fix: Fix Fp8 sequence padding for PP>1 case

[guyueh1]

What does this PR do ?

This is a follow-up after #1569 , to fix the sequence length for PP>1 case.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

Summary by CodeRabbit

• Bug Fixes
  • Improved sequence packing parameter validation and alignment during training with packed sequences to ensure consistent padding enforcement.
  • Fixed sequence dimension calculation to use explicitly padded sequence lengths, enhancing correctness in forward and backward passes.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes enforce stricter padding constraints for packed sequences in Megatron models by asserting divisibility requirements, normalizing pad_packed_seq_to to multiples of a specified value, and using the explicitly padded sequence length for forward/backward computations.

Changes

Cohort / File(s)	Summary
Megatron sequence packing enforcement nemo_rl/models/megatron/common.py	Enforces pad_packed_seq_to divisibility via assertion in _pack_sequences_for_megatron; adds post-computation rounding in _get_pack_sequence_parameters_for_megatron to normalize pad_packed_seq_to to multiples of pad_packed_seq_to_multiple_of.
Policy worker sequence length alignment nemo_rl/models/policy/megatron_policy_worker.py	In train method, overrides seq_dim_size with pad_full_seq_to when active, aligning sequence length used in forward/backward calls and related parameters.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• Attention areas:
  • Verify that assertion in _pack_sequences_for_megatron doesn't break backward compatibility or existing call sites
  • Confirm the rounding logic in _get_pack_sequence_parameters_for_megatron correctly handles edge cases and preserves intended padding semantics
  • Ensure seq_dim_size override in megatron_policy_worker.py doesn't inadvertently affect downstream computations or loss calculations

Possibly related PRs

• fix: Fix the sequence padding for FP8 case #1569: Modifies the same Megatron sequence-packing padding functions and parameters (pad_packed_seq_to, pad_packed_seq_to_multiple_of, _get_pack_sequence_parameters_for_megatron).

Suggested reviewers

• youngeunkwon0405
• zpqiu
• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR contains medium-level changes to sequence padding logic for FP8 sequences with Pipeline Parallelism > 1, but lacks test results, performance metrics, or regression validation evidence.	Add test results validating the PP>1 FP8 sequence padding fix, numeric validation demonstrating no convergence regressions, performance benchmarks (before/after), and details of executed unit/functional tests.

✅ 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 'fix: Fix Fp8 sequence padding for PP>1 case' is partially related to the changeset. It mentions FP8 sequence padding and PP>1, but the actual changes involve sequence padding validation and normalization logic in common.py and policy worker updates. The title accurately describes the domain but lacks specificity about the actual implementation changes (asserting divisibility, rounding up padding parameters).
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

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

---

📜 Recent 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 25ff3f6 and f76efcd.

📒 Files selected for processing (2)

• nemo_rl/models/megatron/common.py (2 hunks)
• nemo_rl/models/policy/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/megatron_policy_worker.py
• nemo_rl/models/megatron/common.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/megatron_policy_worker.py
• nemo_rl/models/megatron/common.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/megatron_policy_worker.py
• nemo_rl/models/megatron/common.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/megatron_policy_worker.py
• nemo_rl/models/megatron/common.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: sphinx-build / Build docs
• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Post automodel integration comment / Comment on PR
• GitHub Check: Post submodule check comment / Comment on PR

🔇 Additional comments (3)

nemo_rl/models/megatron/common.py (2)

87-89: LGTM! Proper contract validation.

The assertion enforces that pad_packed_seq_to must be a multiple of pad_packed_seq_to_multiple_of, which is essential for FP8 correctness with pipeline parallelism. This validates the contract between _get_pack_sequence_parameters_for_megatron (which normalizes the value) and this function.

---

278-282: LGTM! Core fix for FP8 padding with PP > 1.

This normalization step ensures that when pipeline parallelism is used (PP > 1), the packed sequence length is rounded up to satisfy FP8 alignment requirements (multiples of 128 for blockwise, 16 for others). This is essential because all pipeline stages must see identical sequence lengths.

nemo_rl/models/policy/megatron_policy_worker.py (1)

1048-1049: LGTM! Correct sequence length override for PP > 1.

When sequence packing is enabled and pad_full_seq_to is not None (which occurs when PP > 1), this correctly overrides seq_dim_size with the normalized padded sequence length. This ensures that forward_backward_func (lines 1062-1082) receives consistent sequence lengths across all pipeline stages, which is required for pipeline parallelism correctness with FP8.

The pattern is also consistently applied in get_logprobs (line 1235) and get_topk_logits (line 1517).

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions:
   | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context.
   Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

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.}

[terrykong]

@yaoyu-33 can you review?

[guyueh1]

@terrykong this PR is passing L2 but I think the fp8 code pass is not tested due to CI being on A100