remove duplicated RMSNorm and use LlamaRMSNorm from transformers

[yeyu-nvidia]

What does this PR do?

Code cleanup

Overview:
Remove RMSNorm which is identical to LlamaRMSNorm from transformers.

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

Summary by CodeRabbit

• Refactor
  • Updated the normalization implementation in the Eagle speculative module.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Removed RMSNorm class from Eagle utilities module and replaced its usage with LlamaRMSNorm in the transformers plugin. The substitution maintains the same constructor signature while switching the normalization implementation source.

Changes

Cohort / File(s)	Summary
Eagle utilities removal modelopt/torch/speculative/eagle/utils.py	Removed RMSNorm class definition (19 lines) and torch.nn import; retained make_causal_mask and expand_mask functions
Transformers plugin substitution modelopt/torch/speculative/plugins/transformers.py	Replaced RMSNorm import from Eagle utilities with LlamaRMSNorm usage in EagleModule initialization when use_last_layernorm is enabled

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main changes: removing a duplicated RMSNorm class and replacing it with LlamaRMSNorm from transformers library.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

modelopt/torch/speculative/plugins/transformers.py (1)

221-222: Use keyword argument for consistency with existing code.

Lines 263-266 already use LlamaRMSNorm with the keyword argument style (eps=config.rms_norm_eps). Using the same style here improves consistency and makes the parameter's purpose clearer.

Suggested change

         if config.use_last_layernorm:
-            self.norm = LlamaRMSNorm(config.hidden_size, config.rms_norm_eps)
+            self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5e0d365 and 6f61827.

📒 Files selected for processing (2)

• modelopt/torch/speculative/eagle/utils.py
• modelopt/torch/speculative/plugins/transformers.py

💤 Files with no reviewable changes (1)

• modelopt/torch/speculative/eagle/utils.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)

• GitHub Check: linux
• GitHub Check: wait-checks / wait
• GitHub Check: wait-checks / wait
• GitHub Check: build-docs
• GitHub Check: code-quality

🔇 Additional comments (1)

modelopt/torch/speculative/plugins/transformers.py (1)

52-52: LGTM!

The import cleanup correctly removes the now-unused RMSNorm from the local utils module.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.61%. Comparing base (5e0d365) to head (6f61827).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #774      +/-   ##
==========================================
- Coverage   74.62%   74.61%   -0.02%     
==========================================
  Files         192      192              
  Lines       18989    18977      -12     
==========================================
- Hits        14171    14159      -12     
  Misses       4818     4818              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.