feat: Add attention_backend config support for Megatron policy

[sahgerlad]

What does this PR do ?

Enable configuring the attention backend (flash, fused, unfused, local, auto) via megatron_cfg.attention_backend in the YAML configuration.

• Adds support for configuring the attention backend via policy.megatron_cfg.attention_backend in YAML configuration
• Enables users to choose between flash, fused, unfused, local, or auto attention implementations

Usage

• You can potentially add a usage example below

policy:
  megatron_cfg:
    enabled: true
    attention_backend: "flash"  # Options: flash, fused, unfused, local, auto

Summary by CodeRabbit

• New Features
  • Megatron policy worker now supports configurable attention backend settings during initialization when specified in the configuration.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds optional wiring to configure Megatron's attention backend in the policy worker initialization. When attention_backend is provided in megatron_cfg, the code imports AttnBackend from Megatron and sets it on the model configuration.

Changes

Cohort / File(s)	Summary
Megatron attention backend configuration nemo_rl/models/policy/workers/megatron_policy_worker.py	Adds conditional logic to import and configure AttnBackend enum from Megatron when attention_backend is specified in megatron_cfg; maps the string value to the corresponding enum entry and assigns it to model_cfg.attention_backend.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~15–20 minutes

• Verify the AttnBackend enum mapping is correct and aligns with Megatron's current API
• Confirm the conditional logic correctly handles the presence/absence of attention_backend in megatron_cfg
• Check for any potential initialization order dependencies or side effects

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The pull request title accurately describes the main change: adding support for configurable attention_backend in the Megatron policy initialization.
Test Results For Major Changes	✅ Passed	The PR includes test coverage for the attention_backend feature with configuration in conftest.py and active test functions marked @pytest.mark.mcore, plus extensive integration tests.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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: 1

📜 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 e3cfb11 and c23bafb.

📒 Files selected for processing (1)

• nemo_rl/models/policy/workers/megatron_policy_worker.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (4)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Conform code to Python 3.12+
Indent code with 4 spaces. Do not use tabs
Use snake_case for file names
Use PascalCase for class names
Use snake_case for function and method names
Use snake_case for local variables
Prefix variable names that start with a number with 'k' (e.g., k_99th_percentile)
Use upper snake_case with 'G' prefix for global variables (e.g., G_MY_GLOBAL)
Use upper snake_case for constants
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
Prefer docstrings over comments for interfaces that may be used outside a file
Reserve comments for code within a function or interfaces that are local to a file
If a piece of code is commented out, include a comment describing its usage and why it's commented out. Remove debug comments before merging
Use Google style docstrings for classes and functions in Python, which can be parsed by Sphinx
Avoid using reflection when functionality can be easily achieved without reflection
When using try-except blocks, limit the except clause to the smallest set of specific errors possible
When using try-except blocks for duck-typing, keep the body of the try as small as possible and use the else block for logic
YAML is the single source of truth for configuration defaults. Do not set non-None defaults in code for configuration values
For required configuration attributes, access config directly and expect presence (e.g., policy_cfg['precision']) without hidden defaults
Use typing.NotRequired to mark optional attributes in TypedDict for configuration
When adding a new config key to a TypedDict subclass, document the key's purpose, valid values/types, and recommended default, and reflect the default in exemplar YAMLs under examples/configs/*.yaml
Follow the Google Python Style Guide for Python code

Files:

• nemo_rl/models/policy/workers/megatron_policy_worker.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

For any source file under nemo_rl/*.py that defines a class or function decorated with @ray.remote, add a coverage pragma (# pragma: no cover) because these run in separate Ray processes

Files:

• nemo_rl/models/policy/workers/megatron_policy_worker.py

!(**/tests/**|**/test_*.py|**/test_*.sh)

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Add the NVIDIA copyright header to all Python files and shell scripts (excluding tests). The header should include the current year

Files:

• nemo_rl/models/policy/workers/megatron_policy_worker.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

The NVIDIA copyright header should appear at the top of all Python files and shell scripts (excluding tests)

Files:

• nemo_rl/models/policy/workers/megatron_policy_worker.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). (2)

• GitHub Check: Post submodule check comment / Comment on PR
• GitHub Check: Post automodel integration comment / Comment on PR

[sahgerlad]

Hi @terrykong — just checking in on this PR. I rebased it onto the latest main, so it should apply cleanly.

This builds on my previously merged PRs (#1610, #1611) and addresses MoE scalability. Today, the default (dtensor) path doesn’t support expert parallelism and with the Megatron-LM backend the current PyTorch version fails to compile on B300. This PR fixes that by propagating the relevant parameter so we can select a compatible attention implementation; it also enables using flash, which is typically more optimized than the default (inductor) path.

Could you take a look when you have a moment? Thanks!

[terrykong]

thanks for the contribution @sahgerlad . sorry this one slipped. @yaoyu-33 could you review

[cuichenx]

minor comment

[guyueh1]

LGTM, except for one comment; I'll approve

[guyueh1]

@sahgerlad there has been a refactor on megatron policy worker causing merge conflict. Can you resolve it? Also fix the lint please.

[terrykong]

@sahgerlad looks like the new FA test you added fails. TE can't find flash. did it work for you locally?

[sahgerlad]

@sahgerlad looks like the new FA test you added fails. TE can't find flash. did it work for you locally?

FA doesn't support FP32. Specified BF16. Should be okay now

[sahgerlad]

On my end I ran: pytest unit/models/policy/test_megatron_worker.py::test_megatron_policy_training

=========================== short test summary info ============================
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_dp2_llama]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_tp2_llama]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_dp2_qwen2]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_tp2_qwen2]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_dp2_llama_bf16]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_dp2_llama_ac]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_tp2_llama_sp]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_tp2_llama_fp8]
PASSED unit/models/policy/test_megatron_worker.py::test_megatron_policy_training[2gpu_dp2_llama_attention_backend_flash]
================== 9 passed, 27 warnings in 713.36s (0:11:53) ==================

[yuki-97]

/ok to test 86f92f8

[yuki-97]

hi @sahgerlad , overall lgtm, could you take a look at the remaining comments?

[sahgerlad]

hi @sahgerlad , overall lgtm, could you take a look at the remaining comments?

@yuki-97 I believe all comments have been addressed. Let me know if there is anything remaining

[yuki-97]

hi @sahgerlad , overall lgtm, could you take a look at the remaining comments?

@yuki-97 I believe all comments have been addressed. Let me know if there is anything remaining

hi @sahgerlad , just a small Q here #1628 (comment)

[sahgerlad]

hi @sahgerlad , overall lgtm, could you take a look at the remaining comments?

@yuki-97 I believe all comments have been addressed. Let me know if there is anything remaining

hi @sahgerlad , just a small Q here #1628 (comment)

Sorry, had a few comments sitting in review. Just posted. Let me know if there are remaining open items

[yuki-97]

thanks for the contribution again! lgtm

[yuki-97]

/ok to test b7c8ccd