Remove Transformer Engine Triton 3.5 source patch
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
NeMo RL still carries `apply_transformer_engine_patch()` in `nemo_rl/models/policy/workers/patches.py`. The patch rewrites `transformer_engine/pytorch/triton/permutation.py` to work around NVIDIA/TransformerEngine#2286.
That TE issue appears fixed upstream, and TE release notes say v2.10 fixed the Triton 3.5 permute-kernel crash. Current NeMo RL upstream also pins newer TE versions, e.g. `release_v2.15` so the workaround may no longer be needed.
Could we remove this local source-rewriting patch, or at least gate it to only run for vulnerable TE versions?
**Steps/Code to reproduce bug**
N/A
**Expected behavior**
Remove `apply_transformer_engine_patch()` and its call sites, or add a TE version guard.
**Additional context**
Contributor guide
Assessment
This issue has not been assessed yet.