Lightning-AI / Lightning-AI/pytorch-lightning
Add support for converting `RMSNorm` when using `transformer-engine`
Open
Nobody has claimed this yet.
feature
precision: te
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Recently, models (e.g., Qwen3/Llama) have replaced the `LayerNorm` layers with the `RMSNorm` layers because they perform identically but are slightly more efficient.
To support it, a new branch almost identical to the `LayerNorm` would have to be added to the conversion function.
cc @lantiga @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Locate the conversion function and its existing LayerNorm branch, then inspect how RMSNorm appears in Qwen3 or Llama models. Confirm that transformer-engine conversion handles RMSNorm alongside LayerNorm and run the relevant conversion tests, if available.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100