[training_utils, perf] feat: reuse jagged rows across dynamic batches

[zhangxin81]

What does this PR do?

Dynamic batching currently builds each micro-batch by calling index_select_tensor_dict() once per partition. For every jagged NestedTensor field, each call unbinds the entire input batch before selecting its rows.

With F jagged fields and P partitions, this repeats approximately F * P full-batch unbind operations.

This PR adds a multi-partition TensorDict helper that unbinds each jagged field once, caches the resulting row views while constructing the outputs, and reuses the existing field-selection implementation for every partition.

rearrange_micro_batches() uses the new helper. Dense-only and single-partition batches explicitly retain the existing index_select_tensor_dict() path.

Mainstream training stacks avoid this specific repeated-jagged-selection shape in different ways:

• TorchTitan's RL Batcher packs TrainingSample objects into fixed-shape rows before building the micro-batch grid.
• TRL's padding_free collator concatenates per-example fields into flat tensors once.
• OpenRLHF dynamic batching stores partitions as Experience indices and collates each selected micro-batch once.

verl keeps a richer jagged TensorDict online so it can repartition all fields dynamically. Reusing the unbound row views provides the equivalent single-materialization property without replacing verl's data model.

Checklist Before Starting

• Inspected all 483 open PRs and searched for jagged unbind, index_select_tensor_dict, partition_tensor_dict, dynamic batch unbind, and nested tensor unbind performance.
• No open PR implements shared unbind work across dynamic-batch partitions.
• fix: flatten multi-component position_ids to 1D for nested tensor compatibility #5886 is not a duplicate. It fixes multi-component RoPE layout by flattening position_ids; it does not change dynamic-batch partition construction.
• [training_utils] fix: cap micro-batch tokens at max_token_len #6735 modifies rearrange_micro_batches, but addresses token-cap correctness by changing the number of partitions. Its current head still builds partitions one at a time and repeats nested-tensor unbind work. This PR only changes the final partition materialization and can be mechanically rebased if [training_utils] fix: cap micro-batch tokens at max_token_len #6735 lands first.
• The title follows the required format: [training_utils, perf] feat: reuse jagged rows across dynamic batches.

Test

Focused CPU tests:

PYTHONPATH=$PWD python -m pytest -q \
  tests/test_protocol_v2_on_cpu.py \
  tests/utils/test_seqlen_balancing.py \
  -k 'not distributed_params'

Result:

52 passed, 1 deselected

The deselected existing test requires CUDA in this environment and fails unchanged on origin/main.

Coverage includes:

• equality with repeated index_select_tensor_dict() calls;
• regular tensors, jagged tensors, non-last ragged dimensions, NonTensorStack, and scalar NonTensorData;
• output key-order preservation;
• mixed list/tensor partition indices;
• tensor partitions are converted to Python row indices once for jagged tuple selection while retaining the original tensor for dense indexing;
• dense-only and single-partition tensor inputs avoid host conversion;
• one source unbind per jagged field instead of one per partition;
• production rearrange_micro_batches() routing through the batched helper;
• dense-only and one-partition fallback to the existing helper.
• tensor-indexed dense-only and single-jagged fallbacks avoid host conversion.

Additional validation:

• tests/test_protocol_v2_on_cpu.py: 45 passed;
• related dataset tests: 6 passed, with 11 failures caused by unavailable hard-coded local model paths;
• related worker tests: 29 passed, 2 skipped, with one existing 5.72e-06 > 5e-06 tolerance failure;
• complete utils CPU batch: 179 passed, 17 skipped, with failures matching the previously reproduced origin/main environment baseline;
• full pre-commit run --all-files --show-diff-on-failure: all hooks passed.

Reproducible performance benchmark

Benchmark source, raw samples, analysis, hashes, and limitations:

https://gist.github.com/zhangxin81/721fa3c810fe5b04cf1ba07b56d43768

The benchmark times the complete rearrange_micro_batches() function. Baseline and optimized paths use the same partitioning algorithm and validate every output field on every repeat. The only baseline substitution is restoring the old repeated index_select_tensor_dict() partition construction.

Each input has six jagged fields, a dense attention mask, a NonTensorStack, and scalar metadata. Results use 5 warmups, 31 paired repeats, alternating order, 7 independent process launches, and a 100,000-sample bootstrap CI over launch-level log-speedups.

Batch	Partitions	Baseline	Optimized	Savings	Paired speedup	95% CI	Latency reduction
128	9	17.680 ms	8.397 ms	9.283 ms	2.133x	[2.113x, 2.136x]	52.50%
256	17	49.845 ms	17.416 ms	32.429 ms	2.865x	[2.857x, 2.923x]	65.06%
512	17	84.462 ms	24.076 ms	60.386 ms	3.518x	[3.509x, 3.526x]	71.49%

A 201-repeat dense-only negative control measured 1.024x. Dense batches explicitly use the existing helper.

These results cover dynamic-batch CPU partition construction, not complete PPO/GRPO training throughput. End-to-end impact depends on batch size, number of jagged fields, partition count, and the fraction of step time spent preparing micro-batches.

API and Usage Example

There is no user-facing configuration change. Existing callers of rearrange_micro_batches() keep the same API and output partition order.

Internally, callers that need several row partitions may use:

micro_batches = partition_tensor_dict(batch, partitions)

Design & Code Changes

• Keep one shared internal field-selection implementation for single and multi-partition paths.
• Normalize all partition indices once.
• Cache one tuple of row views per jagged field.
• Reuse those views while rebuilding each jagged output with its original ragged dimension.
• Preserve existing behavior for dense tensors, NonTensorStack, NonTensorData, one partition, and dense-only inputs.
• Change rearrange_micro_batches() only at its final partition construction boundary.

Checklist Before Submitting

• Read CONTRIBUTING.md and repository AGENTS.md.
• Duplicate-work checks completed.
• Focused and broader CPU tests run.
• Full pre-commit passed.
• Reproducible benchmark artifacts published.
• No generated config, recipe, or documentation surface is affected.
• CI and review were requested in the PR; fork workflows are awaiting maintainer approval.

AI assistance

OpenAI Codex assisted with repository exploration, implementation, tests, benchmarking, and PR drafting. The human submitter reviewed every changed line, understands the cached-view lifetime and fallback behavior, and personally confirmed the test and benchmark evidence before submission.

[gemini-code-assist]

Code Review

This pull request introduces partition_tensor_dict in tensordict_utils.py to optimize selecting multiple row partitions from a TensorDict by unbinding nested tensors only once. This function is integrated into rearrange_micro_batches in seqlen_balancing.py, and corresponding unit tests are added to verify correctness and performance. Feedback on the changes highlights that iterating directly over a torch.Tensor in partition_tensor_dict is inefficient and will cause a TypeError if the tensor is on GPU; it is recommended to convert the partition tensor to a list of native Python integers using .tolist() first.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Iterating directly over a torch.Tensor in a Python loop or list comprehension (e.g., for idx in partition) is highly inefficient because it creates a new 0D Tensor object for every element, which introduces significant Python overhead. Furthermore, if the partition tensor is on a GPU (CUDA), attempting to index the Python tuple tensor_rows with a 0D CUDA tensor will raise a TypeError (cannot convert CUDA tensor to python integer) because Python's sequence indexing requires host-side integers.

To optimize performance and ensure GPU compatibility, convert the partition tensor to a list of native Python integers using .tolist() before iterating.

Suggested change

	for selected_batch, partition in zip(selected_batches, normalized_partitions, strict=True):
	selected_batch[key] = nested_tensor_from_tensor_list(
	[tensor_rows[idx] for idx in partition], ragged_idx=ragged_idx
	)
	for selected_batch, partition in zip(selected_batches, normalized_partitions, strict=True):
	partition_list = partition.tolist()
	selected_batch[key] = nested_tensor_from_tensor_list(
	[tensor_rows[idx] for idx in partition_list], ragged_idx=ragged_idx
	)

[zhangxin81]

Confirmed: this is a valid issue for tensor partitions. The local follow-up converts each tensor partition to Python row indices once only after entering the multi-jagged fast path, retains the original device tensor for dense TensorDict indexing, and reuses the Python rows across all jagged fields. Dense-only and single-partition fallbacks do zero .tolist() calls, avoiding unnecessary host synchronization. Validation includes 52 focused tests, 200 mixed dense/jagged and list/tensor randomized equivalence trials, full pre-commit, valid-index parity, and a real MPS device-partition check. Updated seven-launch results remain 2.133x-3.518x. The code is intentionally not pushed yet because the repository requires human review of every post-review changed line; the PR body now explicitly marks the current head as pending this reviewed follow-up.

[zhangxin81]

Fixed in 441685b.

The implementation keeps the original tensor partition for dense TensorDict indexing and materializes a Python row list once, outside the jagged-field loop, for cached tuple selection. The conversion only occurs on the multi-partition jagged fast path; dense-only and single-partition fallbacks perform no .tolist() conversion.

Validation after the fix:

• 52 focused tests passed (1 existing CUDA-only test deselected)
• 200 randomized mixed dense/jagged list/tensor equivalence trials passed
• fallback tests verify zero tensor .tolist() calls
• full pre-commit passed
• affected-path benchmark remains 2.133x to 3.518x across the tested cases

The submitter reviewed the follow-up line by line and personally confirmed the test and benchmark results before this commit was pushed.

[zhangxin81]

This PR is ready for upstream CI and maintainer review. All 41 fork-triggered workflows are currently action_required, and the submitter account also lacks RequestReviewsByLogin permission.

Validation at head 441685b0e3d9d5792434496820e4b074a51868a1:

• 52 focused tests passed (one existing CUDA-only test deselected)
• 200 randomized mixed dense/jagged list/tensor equivalence trials passed
• dense-only and single-partition tensor fallbacks perform zero .tolist() host conversions
• full repository pre-commit passed
• full affected-path benchmark shows 52.50% to 71.49% latency reduction

The tensor-partition review finding was fixed in 441685b0, and the submitter reviewed every changed line and personally confirmed the updated tests and benchmark evidence. Could a maintainer approve/run CI and assign/review this PR?