cp: fix: Fix Fp8 sequence padding for PP>1 case (1579) into r0.5.0

[chtruong814]

beep boop [🤖]: Hi @guyueh1 👋,

we've cherry picked #1579 into  for you! 🚀

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Improved validation of sequence packing parameters with clearer error messages for invalid configurations.

• Refactor

  • Enhanced sequence padding normalization to ensure proper alignment in packed sequences, improving consistency in sequence length calculations during training.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR modifies Megatron sequence packing validation and padding logic by enforcing divisibility constraints on pad_packed_seq_to and computing it as a rounded-up multiple of pad_packed_seq_to_multiple_of, then updates the policy worker to respect pad_full_seq_to when overriding sequence dimensions during training.

Changes

Cohort / File(s)	Summary
Megatron sequence packing padding enforcement nemo_rl/models/megatron/common.py	Modified _pack_sequences_for_megatron to validate that pad_packed_seq_to is a multiple of pad_packed_seq_to_multiple_of with an assertion and descriptive error message, replacing prior rounding. Updated _get_pack_sequence_parameters_for_megatron to round pad_packed_seq_to up to the next multiple of pad_packed_seq_to_multiple_of.
Policy worker sequence dimension override nemo_rl/models/policy/workers/megatron_policy_worker.py	Modified MegatronPolicyWorker.train to override seq_dim_size to pad_full_seq_to when sequence packing is enabled and pad_full_seq_to is set, ensuring padded sequence length is reflected in downstream Megatron forward/backward operations.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Understanding of Megatron padding parameter semantics and their propagation through the worker training path

Possibly related PRs

• fix: Fix the sequence padding for FP8 case #1569 — Modifies Megatron sequence-packing padding logic and uses updated _get_pack_sequence_parameters_for_megatron behavior for rounding/enforcement and worker sequence length handling.
• fix: Fix Fp8 sequence padding for PP>1 case #1579 — Changes the same Megatron packing functions in common.py and applies similar seq_dim_size override logic in the policy worker.

Suggested labels

CI:L1, r0.5.0

Suggested reviewers

• guyueh1
• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR modifies FP8 sequence padding and memory layout for distributed training but lacks test results, regression evidence, or performance metrics in description.	Add documentation including test results from original PR #1579, evidence of no numeric/convergence regression, and relevant performance metrics to PR description.
Title check	❓ Inconclusive	The title is a cherry-pick notation that doesn't clearly convey the main change. While it references fixing FP8 sequence padding for PP>1, the 'cp:' prefix and branch notation obscure the primary change from a developer scanning PR history.	Consider using a clearer title format like 'Fix FP8 sequence padding for PP>1 case' without the cherry-pick prefix, or keep the current format if it aligns with project conventions for cherry-pick PRs.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
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
• Commit unit tests in branch cherry-pick-1579-r0.5.0

---

📜 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 d488064 and 66ce259.

📒 Files selected for processing (2)

• nemo_rl/models/megatron/common.py (2 hunks)
• 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/megatron/common.py
• 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/megatron/common.py
• 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/megatron/common.py
• 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/megatron/common.py
• 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). (5)

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

🔇 Additional comments (3)

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

283-287: LGTM! Normalization ensures FP8 compatibility.

The normalization step correctly rounds up pad_packed_seq_to to satisfy the divisibility requirements for FP8 operations, which is essential when pipeline parallelism is enabled (PP > 1). This works in conjunction with the assertion at lines 92-94 to enforce correct padding.

---

92-94: All callers properly use _get_pack_sequence_parameters_for_megatron before calling _pack_sequences_for_megatron.

Verification confirms that every discoverable caller of _pack_sequences_for_megatron obtains parameters via _get_pack_sequence_parameters_for_megatron, which normalizes pad_packed_seq_to to be a multiple of pad_packed_seq_to_multiple_of (lines 283-287). This includes callers in megatron_policy_worker.py and internal calls within common.py. The assertion is valid and safe; it documents the precondition and will catch any future misuse.

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

1150-1151: LGTM! Sequence dimension override aligns with padding normalization.

The override correctly ensures that when pad_full_seq_to is computed (for PP > 1 cases), the sequence dimension used in forward/backward passes matches the normalized padded length from _get_pack_sequence_parameters_for_megatron. This is consistent with similar logic in get_logprobs (line 1351) and get_topk_logits (line 1612).

---

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