cp: fix: Fix DTensor slice crash after PyTorch 2.9 bump (1689) into r0.5.0

[chtruong814]

beep boop [🤖]: Hi @zpqiu 👋,

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

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

• Refactor
  • Optimized internal sequence-packing slicing logic in loss computation.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactors the sequence-packing logits slicing logic in SequencePackingLossWrapper by replacing slice-object indexing with explicit start, end, and length indices computed separately, then applied via tensor.narrow(). The functional behavior remains unchanged—only the implementation approach for selecting next-token logits is modified.

Changes

Cohort / File(s)	Summary
Sequence-Packing Logits Slicing Refactoring nemo_rl/algorithms/loss_functions.py	Replaced logit_slice_idxs-based slice object indexing with explicit slice index computation (start, end, length) and tensor.narrow() method for selecting next-token logits in sequence-packed scenarios. No functional change.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• fix: Fix DTensor slice crash after PyTorch 2.9 bump #1689: Directly related; both PRs refactor the same file to replace slice-object indexing of next-token logits with explicit indices and tensor.narrow(), avoiding __getitem__ with slice/view semantics.

Suggested labels

CI:L1, r0.5.0

Suggested reviewers

• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title describes a cherry-pick operation, but the actual changeset modifies sequence-packing slicing logic in loss_functions.py—a functional change unrelated to the described DTensor slice crash.	Clarify whether this is a cherry-pick of PR #1689 (update title to reflect that) or ensure the title accurately describes the actual sequence-packing refactoring changes in the diff.

✅ Passed checks (3 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%.
Test Results For Major Changes	✅ Passed	Localized bug fix for DTensor slice crash without altering functional behavior, classified as minor change and cherry-pick of already-reviewed PR #1689.

✨ 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)

nemo_rl/algorithms/loss_functions.py (1)

925-930: The narrow() refactoring is correct and preserves intended behavior.

The logic properly computes slice indices accounting for context parallelism:

• logit_start and logit_end apply integer division by cp_size to account for distributed sharding
• logit_length is correctly derived as the difference
• narrow(dim=1, start, length) properly slices the tensor with the computed bounds

Consider adding a brief inline comment (e.g., # Use narrow() for DTensor compatibility with PyTorch 2.9+) to document why this approach is necessary and prevent future refactoring back to slice syntax.

📜 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 e883ac4 and 359ef83.

📒 Files selected for processing (1)

• nemo_rl/algorithms/loss_functions.py

🧰 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/algorithms/loss_functions.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/algorithms/loss_functions.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/algorithms/loss_functions.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/algorithms/loss_functions.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