Skip to content

cp: Revert packed seq extra checks (2180) into r0.3.0 #2196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ko3n1g merged 1 commit into from
Feb 3, 2026
Merged

cp: Revert packed seq extra checks (2180) into r0.3.0 #2196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 4 additions & 8 deletions src/megatron/bridge/training/utils/packed_seq_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,18 +44,14 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
cu_seqlens_unpadded_argmin = batch.get("cu_seqlens_unpadded_argmin")

if cu_seqlens_argmin is not None:
argmin_idx = cu_seqlens_argmin.item()
assert argmin_idx == 0 or cu_seqlens_padded[argmin_idx] == -1 # cu_seqlens padding is -1
cu_seqlens_padded = cu_seqlens_padded[:argmin_idx]
elif torch.min(cu_seqlens_padded) == -1:
cu_seqlens_padded = cu_seqlens_padded[: cu_seqlens_argmin.item()]
else:
cu_seqlens_padded = cu_seqlens_padded[: torch.argmin(cu_seqlens_padded)]

if cu_seqlens_unpadded is not None:
if cu_seqlens_unpadded_argmin is not None:
argmin_idx = cu_seqlens_unpadded_argmin.item()
assert argmin_idx == 0 or cu_seqlens_unpadded[argmin_idx] == -1 # cu_seqlens padding is -1
cu_seqlens_unpadded = cu_seqlens_unpadded[:argmin_idx]
elif torch.min(cu_seqlens_unpadded) == -1:
cu_seqlens_unpadded = cu_seqlens_unpadded[: cu_seqlens_unpadded_argmin.item()]
else:
cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)]
Comment on lines 46 to 55

@coderabbitai coderabbitai Bot Feb 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🌐 Web query:

PyTorch tensor slicing with 0-dimensional tensor indices CUDA CPU compatibility

💡 Result:

In PyTorch, a 0‑dim tensor used as an index is still a “tensor index” (advanced indexing), so PyTorch applies the device/dtype rules for index tensors:

  • Index tensors must be on CPU or on the same device as the tensor being indexed; otherwise you’ll get errors like: “indices should be either on cpu or on the same device as the indexed tensor”. [1]
  • On CUDA, some invalid tensor-indexing cases (e.g., out-of-bounds when the index is a tensor) may surface as a device-side assert rather than a clean IndexError, which can be harder to catch/debug. [2]

Practical patterns for CPU/CUDA compatibility

1) If you want “scalar slicing”, convert the 0‑dim tensor to a Python int

i = int(i0.item())          # i0 can be CPU or CUDA 0-dim
y = x[i]                    # uses basic indexing; avoids tensor-index device rules

This is the most portable fix. (Note: .item() on a CUDA tensor synchronizes.)

2) If you want to keep it as a tensor index, keep devices consistent

idx = idx.to(x.device)      # and ensure integer dtype (usually torch.long)
y = x[idx]

3) Debugging CUDA “device-side assert” from indexing
If you hit a device-side assert during indexing, reproduce with CUDA_LAUNCH_BLOCKING=1 to get a more accurate stack trace; this class of issue is commonly triggered by invalid CUDA indexing with tensor indices. [2]

Sources:
[1] PyTorch Forums: “Indices should be either on cpu or on the same device as the indexed tensor” (discussion of the device rule for index tensors). [1]
[2] PyTorch GitHub issue discussing CUDA tensor indexing leading to device-side asserts instead of IndexError. [2]

🏁 Script executed:

# First, let's examine the actual file to understand the context
cat -n src/megatron/bridge/training/utils/packed_seq_utils.py | head -70

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 3436

🏁 Script executed:

# Let's verify the full context and see if there's any processing of these cumulative seqlens
cat -n src/megatron/bridge/training/utils/packed_seq_utils.py | sed -n '36,70p'

Repository: NVIDIA-NeMo/Megatron-Bridge

Length of output: 1818

Convert torch.argmin() results to Python int for CUDA compatibility and consistent slicing behavior.

Lines 49 and 55 use torch.argmin() directly as a slice bound, which returns a 0-dimensional tensor. On CUDA, this can trigger device-side assertion errors because PyTorch requires index tensors to match the device of the indexed tensor. Lines 47 and 53 already use .item() correctly; apply the same pattern to the fallback paths.

Proposed fix 
     if cu_seqlens_argmin is not None:
         cu_seqlens_padded = cu_seqlens_padded[: cu_seqlens_argmin.item()]
     else:
-        cu_seqlens_padded = cu_seqlens_padded[: torch.argmin(cu_seqlens_padded)]
+        cu_seqlens_padded = cu_seqlens_padded[: torch.argmin(cu_seqlens_padded).item()]
 
     if cu_seqlens_unpadded is not None:
         if cu_seqlens_unpadded_argmin is not None:
             cu_seqlens_unpadded = cu_seqlens_unpadded[: cu_seqlens_unpadded_argmin.item()]
         else:
-            cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)]
+            cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded).item()]
🤖 Prompt for AI Agents 
In `@src/megatron/bridge/training/utils/packed_seq_utils.py` around lines 46 - 55,
The slice bounds use torch.argmin() (which returns a 0-d tensor) causing
CUDA/device-index issues; update the fallback branches to convert the argmin
results to Python ints (use .item()) before slicing cu_seqlens_padded and
cu_seqlens_unpadded so that cu_seqlens_padded = cu_seqlens_padded[:
torch.argmin(cu_seqlens_padded).item()] and cu_seqlens_unpadded =
cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded).item()], matching the
existing .item() usage in the other branches for consistent, device-safe
indexing of cu_seqlens_padded and cu_seqlens_unpadded.


max_seqlen = batch["max_seqlen"].squeeze() if "max_seqlen" in batch else None
Expand Down
48 changes: 0 additions & 48 deletions tests/unit_tests/training/test_gpt_step.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,23 +103,6 @@ def test_packed_seq_params_with_cu_seqlens_argmin(self):
assert torch.equal(result.max_seqlen_q, expected_max_seqlen)
assert torch.equal(result.max_seqlen_kv, expected_max_seqlen)

def test_packed_seq_params_no_padding(self):
"""Test functionality when cu_seqlens has no padding (-1 values)."""
batch = {
"cu_seqlens": torch.tensor([[0, 7, 14]], dtype=torch.int32),
"max_seqlen": torch.tensor([[10]], dtype=torch.int32),
}

result = get_packed_seq_params(batch)

# Verify the result is a PackedSeqParams object
assert isinstance(result, PackedSeqParams)

# When there's no -1 padding, the tensor is returned unchanged
expected_cu_seqlens = torch.tensor([0, 7, 14], dtype=torch.int32)
assert torch.equal(result.cu_seqlens_q, expected_cu_seqlens)
assert torch.equal(result.cu_seqlens_kv, expected_cu_seqlens)

def test_packed_seq_params_with_cu_seqlens_argmin_zero(self):
"""Test edge case when cu_seqlens_argmin is 0."""
batch = {
Expand Down Expand Up @@ -246,21 +229,6 @@ def test_packed_seq_params_without_unpadded_fallback(self):
assert result.cu_seqlens_q_padded is None
assert result.cu_seqlens_kv_padded is None

def test_packed_seq_params_no_padding_in_cu_seqlens(self):
"""Test when cu_seqlens has no -1 padding markers."""
batch = {
"cu_seqlens": torch.tensor([[0, 5, 10]], dtype=torch.int32), # No -1 padding
"max_seqlen": torch.tensor([[7]], dtype=torch.int32),
}

result = get_packed_seq_params(batch)

# When no -1 present and min != -1, the tensor should remain as-is
expected = torch.tensor([0, 5, 10], dtype=torch.int32)
assert torch.equal(result.cu_seqlens_q, expected)
# Padded fields are None when cu_seqlens_unpadded is not provided
assert result.cu_seqlens_q_padded is None

def test_packed_seq_params_qkv_format_is_thd(self):
"""Test that qkv_format is always set to 'thd'."""
batch = {
Expand All @@ -271,22 +239,6 @@ def test_packed_seq_params_qkv_format_is_thd(self):

assert result.qkv_format == "thd"

def test_packed_seq_params_cu_seqlens_unpadded_no_padding(self):
"""Test cu_seqlens_unpadded with no padding markers."""
batch = {
"cu_seqlens": torch.tensor([[0, 6, 12]], dtype=torch.int32),
"cu_seqlens_unpadded": torch.tensor([[0, 5, 10]], dtype=torch.int32), # No -1
}

result = get_packed_seq_params(batch)

# Unpadded should be used as-is since no -1 and min != -1
expected_unpadded = torch.tensor([0, 5, 10], dtype=torch.int32)
expected_padded = torch.tensor([0, 6, 12], dtype=torch.int32)

assert torch.equal(result.cu_seqlens_q, expected_unpadded)
assert torch.equal(result.cu_seqlens_q_padded, expected_padded)


class TestCreateLossFunction:
"""Tests for the _create_loss_function helper function."""
Expand Down
Loading