Enable dynamic batching

[jiosephlee]

What does this PR do?

Padding-free packing with FA2 already solves the issue of padding + packing, but not all models support FA2 (e.g. GPT-OSS uses FlexAttention).

With sample-based batching, the token footprint of each microbatch varies with sequence length. Thus, users often have to choose a conservatively small microbatch size so an unlucky group of long sequences doesn't OOM. This is inefficient and unreliable.

This PR enables dynamic batching. It balances sequences across data-parallel ranks, groups similar-length sequences, and constructs rank-local microbatches under --train.max_tokens_per_gpu. Each GPU can therefore use its available memory according to a token-level budget instead of a fixed number of samples.

Length-aware grouping spends a controlled amount of padding to replace several small forwards with fewer, larger, more GPU-efficient forwards. Each rank may use a different local padded length while executing the same number of microbatches, keeping FSDP collectives aligned.

The scheduler preserves optimizer windows, sample coverage, and global-token loss normalization. Non-singleton batches remain within the configured padded-footprint budget. A sequence that individually exceeds the budget remains a warned singleton rather than being dropped, so the configured budget must still account for the longest supported sequence and other model memory.

--train.dynamic_batch_pad_to_multiple controls optional shape bucketing and defaults to 1, meaning exact padded lengths. This is dynamic microbatch sizing, not sequence packing or padding-free training.

Existing fixed microbatching and packed dynamic batching remain unchanged, and dynamic batching stays opt-in.

GPU validation

The fixed and dynamic paths were compared on two NVIDIA B200 GPUs using frozen replay. Each arm ran in three fresh processes for 23 optimizer steps; the first three steps were excluded as warm-up.

Both arms used padded training with --train.micro_batch_size 1. The dynamic arm additionally used:

--train.dynamic_batch_enable
--train.max_tokens_per_gpu 8192
--train.dynamic_batch_pad_to_multiple 1

Model	Fixed median step	Dynamic median step	Paired speedup	Median peak-memory reduction
Qwen3-4B	3.604 s	1.506 s	2.36×	14.14 GiB
GPT-OSS-20B Flex/MoE	8.579 s	5.406 s	1.59×	4.32 GiB

All 12 runs consumed 184/184 samples and completed 23/23 optimizer steps without OOMs, hangs, missing samples, or collective failures.

The correctness comparisons used the same checkpoints, samples, rewards, and optimizer windows:

• Qwen3-4B produced identical policy loss, gradient cosine 0.998696, final-weight cosine 0.999999999999, and maximum parameter difference 3.815e-6.
• The Qwen first-step Adam update was more sensitive to small BF16 gradient differences: update cosine 0.972842 and relative update difference 23.305%.
• GPT-OSS-20B produced identical policy loss, update cosine 0.999956, relative update difference 0.942%, and maximum parameter difference 3.815e-6.
• A direct CUDA FlexAttention test passed for valid outputs and Q/K/V gradients across exact and right-padded sequence lengths.

These results establish numerical training equivalence within expected BF16 reduction differences, not bitwise-identical gradients or updates.

Validation

• python -m compileall -q molt examples/python tests
• python -m pytest -q tests/unit/test_dynamic_batch.py: 11 passed, 1 skipped
• python -m pytest -q: 202 passed, 3 skipped
• Direct CUDA FlexAttention test: 1 passed, 11 deselected

Caveats

• Dynamic batching remains disabled by default.
• Exact-length padding (pad_to_multiple=1) is the recommended configuration.
• A separate 1,024-token shape-bucket experiment reduced FlexAttention recompilations but regressed median step time by 1.4%, throughput by 3.0%, and padding efficiency from 71.2% to 18.2%; it remains opt-in.
• On the longest GPT-OSS step, dynamic batching was 4.0% slower and maximum allocated memory was effectively flat (+0.059 GiB) because the new shape triggered a FlexAttention compile (issues with FlexAttention should be investigated deeper and addressed by a different PR)
• GPU validation covered one node, two B200 GPUs, Qwen3-4B, and GPT-OSS-20B with CP=1. Multi-node, live context-parallel, and VLM configurations were not tested.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hijkzzz]

/claude review

[claude]

LGTM

Reviewed the padded dynamic-batch feature end to end (CLI validation → replay-buffer partitioning → rollout forward-batch scheduling → worker forward_batch → result cropping) plus the new make_experience_batch padding and README switch. The forward-batch schedule alignment across actor/reference/critic groups, the rank-then-batch output reassembly in _dispatch_forward, the right-pad-then-crop of results, and the split-to-target termination all check out, and the new test file gives comprehensive coverage of the padded path (pricing, CP alignment, shape bucketing, cross-rank count matching, over-budget warnings, worker reload with routing/VLM). No critical bugs, typos, coverage gaps, or doc inaccuracies found.

[hijkzzz]

Automodel team will implement this via VLM seq packing~

[jiosephlee]

@hijkzzz sounds good! I'll leave the PR open for now in case the team faces blockers; I would like to see this feature sooner than later as it's a rather large speed-up for non-FA users

(was a big fan of your previous repo OpenRLHF; happy to see this project get started)