[bug] equal-length batches of 3D mRoPE position_ids produce an inconsistent jagged layout → split_with_sizes crash in V1 trainer

Summary

In the V1 trainer, per-sample 3D (mRoPE) position_ids of shape (num_components, seq_len) are carried through the pipeline as jagged NestedTensors. Two interacting issues crash training whenever all sequences in a batch have the same length (which includes the trivial case batch_size == 1):

1. Write side. torch.nested.as_nested_tensor(list_of_2d, layout=torch.jagged) mis-takes the component dim as the jagged dim when all input samples are equal-length: instead of the canonical ragged@2 layout (lengths=[L_1..L_B], values=(C, ΣL_i), _ragged_idx=2) it produces a ragged@1 layout (lengths=[C]×B, values=(B·C, L), _ragged_idx=1).
2. Consumer side. maybe_fix_3d_position_ids() (verl/utils/tensordict_utils.py) unconditionally sets _ragged_idx = 2 on any 3D nested position_ids without validating the actual layout. Applied to the ragged@1 tensor above, this produces a metadata-inconsistent tensor whose next unbind() raises:

RuntimeError: split_with_sizes expects split_sizes to sum exactly to 34135
(input tensor's size at dimension 1), but got split_sizes=[4, 4, 4, ...×128]

We hit exactly this in a real 35B MoE multi-turn RL run of ours — 81 healthy steps, then this crash. Details below. The quirk is not torch-2.9-specific — the minimal repro runs on CPU with torch 2.13.0.

Minimal repro (CPU, torch 2.13.0)

import torch
C, L, B = 4, 34135, 128   # 4 mRoPE components; L = a real batch's sequence length
eq = [torch.arange(C * L).reshape(C, L) % 100000 for _ in range(B)]

nt = torch.nested.as_nested_tensor(eq, layout=torch.jagged)
print(nt.shape, nt._ragged_idx, nt.values().shape, nt.offsets()[:5].tolist())
# -> (128, j1, 34135) 1 (512, 34135) [0, 4, 8, 12, 16]     <- ragged@1: component dim treated as jagged

var = [torch.arange(C * l).reshape(C, l) for l in (100, 200, 300)]
nt2 = torch.nested.as_nested_tensor(var, layout=torch.jagged)
print(nt2.shape, nt2._ragged_idx, nt2.values().shape)
# -> (3, 4, j2) 2 (4, 600)                                  <- variable-length: canonical ragged@2

nt._ragged_idx = 2                                          # what maybe_fix_3d_position_ids() does
nt.unbind(0)
# RuntimeError: split_with_sizes expects split_sizes to sum exactly to 34135
# (input tensor's size at dimension 1), but got split_sizes=[4, 4, 4, ...]

Expected: equal-length batches build canonical ragged@2, so unbind() returns 128 tensors of shape (4, 34135).

Real-world occurrence in our RL training

We hit this crash in a real RL training run: a 35B MoE model with 4-component mRoPE, multi-turn agentic RL, fully-async V1 trainer with TransferQueue, dynamic batch size, max_response_length=32768. Training ran healthy for 81 steps and crashed at step 82 in the dataloader with exactly the error above (34135 = 1367 prompt + 32768 response cap).

The trigger: at that point the policy had degenerated into "repeat until truncation" (validation truncated ratio ≈ 0.97), so every rollout in the batch was capped at the same length — the batch became fully equal-length, as_nested_tensor flipped to the ragged@1 layout, and the blind _ragged_idx = 2 retag detonated it at the next unbind().

We reproduced both the malformed layout production and the blind-retag crash at the exact real-batch scale (128 × (4, 34135)) on the training environment; the minimal script above shows the same on CPU.

Note the delayed detonation pattern: the malformed layout is produced at the write site many steps before the crash, and the crash site (unbind() in the dataloader) is far from where the wrong layout was built — which makes the bug look like a transient hardware/communication failure and very hard to attribute.

Corroborating report

#6851 reports the same crash with all batch sizes = 1 — a single-sample batch is trivially "equal-length", so the quirk fires deterministically there as well. Its traceback goes through the TransferQueue path (transferqueue_utils.py), confirming the bug is reachable with TQ enabled: TQ faithfully preserves whatever layout the write site produced, so the malformation needs no pickle/consolidation step to survive.

Affected code paths on master (cb21203a)

Write sites that build jagged NT from per-sample 2D mRoPE slices via as_nested_tensor (quirk-prone for equal-length inputs):

• verl/workers/utils/padding.py — left_right_2_no_padding() rebuilds 3D position_ids (~line 67)
• verl/utils/seqlen_balancing.py — restore_dynamic_batch() (~line 583)
• verl/protocol.py — deserialize_tensordict() (~line 287); serialize_tensordict() also does not record _ragged_idx, so a canonical ragged@2 tensor round-trips into an as_nested_tensor rebuild that re-rolls the quirk
• verl/utils/tensordict_utils.py — list_of_dict_to_tensordict() fallback (~line 940)

Consumer side:

• verl/utils/tensordict_utils.py — maybe_fix_3d_position_ids() (~line 910): unconditional _ragged_idx = 2; called from verl/workers/engine/base.py (train_batch / infer_batch), engine_workers.py, engine_workers_tinker.py
• verl/utils/tensordict_utils.py — index_select_tensor_dict() (~line 458): unbind() with no layout validation, so the inconsistency surfaces as the cryptic split_with_sizes error
• concat_nested_tensors(): uses the first tensor's _ragged_idx for all inputs — mixed-layout inputs are silently rebuilt with the wrong tag

Relation to #6851 and #7767

• #6851 (see Corroborating report above): the same crash with all batch sizes = 1 — multiple users have reported hitting it, confirming this is a common failure mode rather than specific to our setup.
• #7767 (closed, unmerged) attempted to rebuild the tensor inside maybe_fix_3d_position_ids(). Its closure rationale ("with TransferQueue enabled by default, we no longer consolidate/pickle TensorDict") removes only the serialization door. The write-site doors listed above build the malformed layout directly in verl code — no pickle involved — and the blind _ragged_idx = 2 is still applied to whatever arrives. Also, its detection used a 3-or-4-rows heuristic for the component count, which misses other component counts.
• Related: #7586 (min/max seqlen for jagged rebuilds, open), pytorch/pytorch#159380 (to_padded_tensor without min/max seqlen).

Proposed fix direction

Layered defense:

1. Write sites: build 3D position_ids with explicit ragged semantics — verl already has nested_tensor_from_tensor_list(..., ragged_idx=...); for 3D position_ids use ragged_idx=2 (sequence dim jagged, component dim preserved). For serialize_tensordict, record _ragged_idx and restore it on deserialize instead of re-rolling as_nested_tensor.
2. Replace the blind tag with a validating normalize step that recognizes non-canonical states by layout signature (offsets/values shape relationships, component-count agnostic — no C in (3,4) hardcode), rebuilds to canonical ragged@2, and is a zero-cost no-op for canonical input.
3. Fail-fast: validate sum(lengths) == values.shape[ragged_idx - 1] before unbind() in index_select_tensor_dict() with an actionable error message (key name, ragged_idx, sums); reject mixed-_ragged_idx concat instead of silently using the first tensor's tag.