fix: pad bshd samples to their own length, not the rollout-global max

[EazyReal]

Summary

bshd uses a rectangular [batch, sequence, heads, head-dim] tensor. Miles previously chose that sequence dimension from the longest sample in the entire rollout, so one long outlier forced even unrelated short samples in other microbatches to carry the same padding. This changes max_seq_lens from a repeated rollout-global maximum to each sample's own aligned padded length.

Symptom

Measured on a 128x H100 DeepSeek-V4-Flash run at 262,144-token context: a 14,010-token sample padded to 132,096 OOMed the expert MLP at 79 GB. Pad tokens still dispatch through MoE. With per-sample padded lengths, the same step peaks at 38.5 GB.

Root Cause

get_rollout_data computed max(total_lengths), rounded it for the existing TP/CP/compression alignment requirements, and copied that value into every entry of max_seq_lens. Because padding was sized across the rollout rather than at the microbatch that consumes a sample, a single outlier inflated every sample's token count (about 9.4x in the measured case).

max_seq_lens is also layout metadata, not only a memory hint: token tensors, rollout log probabilities, replay tensors, and response-logit offsets must use the same padded length when applying context-parallel slicing.

Fix

• In get_rollout_data, round each sample's own total_length up with the existing pad_size rule. This preserves all current alignment constraints while avoiding padding across unrelated samples.
• Keep max_seq_lens as the single source of BSHD padding geometry. Existing per-sample consumers continue to read the corresponding entry, including build-time CP slicing of rollout log probabilities.
• In get_batch and fill_replay_data, assert that every sample in one BSHD microbatch has the same padded length. BSHD is rectangular, and these paths currently form the batch from max_seq_lens[0]; accepting different padded lengths would make token, log-probability, and replay layouts disagree silently.
• Warn during argument validation for BSHD with --micro-batch-size > 1, with remediation to use one sample or bucket samples by padded length.

Why This Shape

The change is intentionally made at the existing max_seq_lens producer, before rollout-derived tensors are CP-sliced. Computing a larger value only later in get_batch would be incorrect because those tensors would already have been sliced using different per-sample geometry. The guards therefore fail closed at the two remaining rectangular-batch consumers instead of hiding a layout mismatch.

This is the smallest change that fixes the production configuration without introducing a second padding owner or reorganizing scheduling. It supports:

• BSHD with one sample per microbatch, used by the checked-in BSHD recipes.
• Multi-sample BSHD microbatches when the samples round to the same padded length.

Multi-sample microbatches with different padded lengths are rejected with a targeted error. Fully general support requires establishing the final microbatch membership before CP slicing, assigning one shared maximum to that group, and then slicing every dependent field with that same value. The TODO at the producer records that separate scheduling/data-flow change; taking max() only at batch consumption would not be sufficient.

Performance Trade-off

Per-sample padding can produce more BSHD sequence shapes than one rollout-global target. That reduces dense token work and memory, but may trade away some kernel, execution-plan, or compilation-cache reuse. This PR therefore claims the measured memory result above, not a universal step-time improvement. --data-pad-size-multiplier remains the existing granularity control: increasing it coarsens aligned-length buckets, trading more padding for fewer distinct shapes. A separate global-max mode is not added without evidence that its cache benefit outweighs the observed OOM risk in a Miles workload.

Verification

New CPU coverage in tests/fast/backends/training_utils/test_bshd_per_sample_pad.py verifies:

• Per-sample aligned values are produced.
• Different true lengths that share one padded length form the expected rectangular batch.
• Different padded lengths are rejected by both get_batch and fill_replay_data.

The focused regression file and the complete argument-validation test file pass 146/146 locally on CPU-only torch. The 38.5 GB result above is the same change validated on the H100 run.

Remaining Boundary

max_seq_lens is still computed before multimodal token expansion, matching the previous ordering rather than introducing a regression. No checked-in BSHD recipe uses multimodal token expansion; supporting that combination should recompute the padded lengths after expansion.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.