Revert packed seq extra checks

[cuichenx]

What does this PR do ?

Packed sequence changes in #1997 resulted in perf degradation.

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

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 add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

• Refactor
  • Simplified internal sequence handling logic by streamlining the slicing pattern, reducing dependency on padding assumptions and improving code clarity.

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactors slicing logic for cu_seqlens sequences in packed sequence utilities to use argmin-based approach instead of index-based slicing. Replaces conditional branches checking for -1 padding with unified argmin pattern, reducing dependency on padding assumptions and simplifying control flow.

Changes

Cohort / File(s)	Summary
Packed Sequence Slicing Refactor src/megatron/bridge/training/utils/packed_seq_utils.py	Reworks slicing logic for cu_seqlens_padded and cu_seqlens_unpadded to use argmin-based approach directly when indicators are present; removes conditional min==-1 branches and explicit assertions; unifies control flow to consistent argmin-based pattern.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR reverts changes due to performance degradation but lacks required performance benchmarking results, before-and-after metrics, test validation, and incomplete pre-check checklist.	Add performance benchmarking with before-and-after metrics, include test results validating packed sequence slicing logic, and complete pre-check testing checklist.

✅ 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 'Revert packed seq extra checks' accurately describes the main change—reverting recent packed sequence modifications from PR #1997 that caused performance issues.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 chcui/revert_packed_seq

---

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

🤖 Fix all issues with AI agents

In `@src/megatron/bridge/training/utils/packed_seq_utils.py`:
- Around line 46-55: The current slicing using torch.argmin on cu_seqlens_padded
and cu_seqlens_unpadded can produce [:0] when there is no -1 marker (min at
index 0); modify the logic in the block handling cu_seqlens_padded and
cu_seqlens_unpadded so that when the corresponding argmin indicator is None you
first check the minimum value (e.g. via torch.min(...).item() or tensor.min())
and only perform the slice when that minimum is negative, otherwise leave the
tensor unchanged; update the branches that reference cu_seqlens_argmin,
cu_seqlens_unpadded_argmin, torch.argmin, cu_seqlens_padded and
cu_seqlens_unpadded accordingly.

[cuichenx]

/ok to test 87c2d06

[cuichenx]

/ok to test 87c2d06