fix: flatten multi-component position_ids to 1D for nested tensor compatibility

[yifannnwu]

Summary

Fixes a crash when training models with multi-dimensional RoPE position_ids (e.g. Qwen3.5's 4-component (text, height, width, temporal), Qwen2-VL) on the FSDP engine_workers.py path (use_legacy_worker_impl=disable).

Related: #5772 (same class of bug on the Megatron/mbridge path)

Problem

left_right_2_no_padding() in padding.py creates per-sample position_ids of shape (num_components, valid_len) and passes them directly to torch.nested.as_nested_tensor(..., layout=torch.jagged).

Jagged layout treats dim 0 as the ragged dimension, so it interprets num_components (e.g. 4) as ragged instead of seq_len. The resulting nested tensor has shape [batch, *(ragged=4), seq_len] instead of the intended [batch, *(ragged=seq_len), num_components].

Downstream in transformer_impl.py, to_padded_tensor() produces incorrect shapes → crash in apply_rotary_pos_emb.

Transposing to (valid_len, num_components) before nesting doesn't work either, because 3D jagged nested tensors have broken unbind() and to_padded_tensor() in PyTorch (see pytorch/pytorch#153238).

Fix

Flatten multi-component position_ids to 1D (seq_len * num_components,) before creating the nested tensor, keeping it purely 2D jagged. Store num_pos_components as non-tensor metadata. Reshape back to (num_components, batch, seq_len) in prepare_model_inputs.

Changes

verl/workers/utils/padding.py — left_right_2_no_padding():

• Flatten (num_components, valid_len) → (valid_len * num_components,) via .T.contiguous().flatten()
• Store num_pos_components in non-tensor metadata via tu.assign_non_tensor_data

verl/workers/engine/fsdp/transformer_impl.py — prepare_model_inputs():

• NO_PADDING path: .view(total_nnz, num_components).T.unsqueeze(1) → (num_components, 1, total_nnz)
• Padded path: to_padded_tensor(..., output_size=(batch, max_seq_len * num_components)).view(batch, max_seq_len, num_components).permute(2, 0, 1) → (num_components, batch, max_seq_len)

Affected models

Any model with multi-dimensional RoPE on the FSDP engine_workers.py path:

• Qwen3.5 (4-component: text + 3D spatial)
• Qwen2-VL (3-component mRoPE)
• Other VLMs with multi-component position_ids

Test plan

• Qwen3.5-0.8B, FSDP2, use_legacy_worker_impl=disable: 3 GRPO training steps + validation pass. Previously crashed with RoPE shape mismatch.
• Qwen2-VL (not tested, same position_ids format — should work)

[gemini-code-assist]

Code Review

This pull request introduces support for multi-component position_ids by flattening them into 1D tensors to circumvent PyTorch nested tensor limitations, while adding logic to reconstruct the original shapes during model input preparation. Review feedback suggests optimizing the flattening process in padding.py for better performance and using inferred dimensions in view calls within transformer_impl.py for more idiomatic code.

[gemini-code-assist]

The current implementation performs multiple intermediate operations and an explicit contiguous copy. Using curr_pos_ids.T[curr_mask].flatten() is more efficient and idiomatic. Since boolean indexing with a mask already creates a new contiguous tensor, this approach avoids the extra overhead of contiguous() and simplifies the operation in this hot path.

Suggested change

	valid_ids = curr_pos_ids[:, curr_mask].T.contiguous().flatten() # (valid_len * num_components,)
	valid_ids = curr_pos_ids.T[curr_mask].flatten() # (valid_len * num_components,)

[gemini-code-assist]

The explicit calculation of total_nnz is redundant here. Using -1 in the view method is more idiomatic and robust, as it allows PyTorch to automatically infer the dimension size while ensuring the total number of elements is compatible with num_pos_components. This also makes the code slightly cleaner by removing an unnecessary intermediate variable.

Suggested change

	total_nnz = flat_pos.shape[0] // num_pos_components
	position_ids_rmpad = flat_pos.view(total_nnz, num_pos_components).T.unsqueeze(1)
	position_ids_rmpad = flat_pos.view(-1, num_pos_components).T.unsqueeze(1)

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}