fix: Fix Fp8 sequence padding for PP>1 case

Open guyueh1 opened this issue 4 weeks ago • 1 comments

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

Nov 29 '25 19:11 guyueh1

📝 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

#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 25ff3f67be84be1f7a616bccaf235510023f2af8 and f76efcdbc02ed607d89b596aae540454be3da904.

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

📝 Description — Summarize the main change in 50–60 words, explaining what was done.

📓 References — List relevant issues, discussions, documentation, or related PRs.

📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.

📊 Contributor Summary — Include a Markdown table showing contributions: | Contributor | Lines Added | Lines Removed | Files Changed |

✔️ 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

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

Nov 29 '25 19:11 coderabbitai[bot]