Stabilize training step time by reducing data-dependent micro-batch shape variance in sequence packing during SFT training
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Is your feature request related to a problem? Please describe.**
When using sequence packing, training step time varies significantly because micro-batch shapes (packed sequence length and number of micro-batches) are data-dependent and change every step.
In our SFT training on Nemotron-3-Nano-30B-A3B (MoE, 64 experts, TP=4, EP=8, PP=2, DP=8, 8 nodes/64 GPUs, `algorithm=modified_first_fit_decreasing`, `train_mb_tokens=32768`), we observed over 30,000 steps:
| Metric | Value |
|--------|-------|
| Mean step time | 2.807s |
| Std | 0.416s (14.8%) |
| P5–P95 range | 1.98s–3.48s |
| Min–Max | 1.44s–5.93s |
The variance comes from multiple data-dependent factors:
1. **DataLoader** produces a random batch of samples with varying lengths each step (`shuffle=True`, no length bucketing). The length distribution changes every step.
2. **MFFD bin packing** packs them into bins with `bin_capacity=train_mb_tokens`. Each bin's actual packed token count varies depending on how well sequences fit. The number of bins also varies per step, meaning the gradient accumulation step count is not fixed.
3. **When PP > 1**, all micro-batches are padded to `max(micro_batch_lengths)`, amplifying the variance. But **even with PP=1**, each micro-batch has a different packed sequence length, and the number of micro-batches varies per step, both causing step time variance.
4. `padded_seq_length` is computed **independently per DP rank**, so different ranks may have different computation loads. The slowest rank determines overall step time due to gradient synchronization.
5. The core issue is the design of **fixed sample count with variable tensor shapes**, rather than **fixed tensor shapes with variable sample count**.
6. **Total training steps is inflated.** The current design fixes sample count per step (`global_batch_size`), so total steps = `N / global_batch_size` regardless of sequence lengths. In a fixed-shape packing mode where each step consumes a constant token budget, total steps would be `total_tokens / tokens_per_step`, where `tokens_per_step == global_batch_size * max_sequence_length` , which is significantly fewer when most samples are shorter than `max_sequence_length`. The current approach means each step processes a variable (and often suboptimal) number of tokens, leading to both more total steps and unstable per-step time.
Additionally, `sft_processor` discards any sample with `length >= max_seq_length` entirely by setting `loss_multiplier=0.0`, rather than truncating. A sample with 65,536 tokens could be truncated to 32,768 and still have thousands of useful tokens for loss computation. There is an existing TODO acknowledging this (`TODO @sahilj handle too-long prompts`).
**Describe the solution you'd like**
1. **Fixed-shape packing mode (concatenate-and-chunk):** Support a mode where a variable number of samples are drawn to fill a fixed tensor shape (constant `padded_seq_length` and constant micro-batch count), guaranteeing constant computation per step. This would also reduce total training steps from `N / global_batch_size` to `total_tokens / tokens_per_step`, since each step consumes a full token budget instead of a fixed sample count with variable fill. Both packing styles can use `cu_seqlens` for attention isolation, so this doesn't require accepting cross-sequence attention.
2. **Truncate over-length samples instead of discarding:** Modify `sft_processor` to truncate at `max_seq_length` with valid loss computation on the retained tokens, rather than wasting the entire sample.
**Additional context**
Step time distribution over 30,000 steps:
| Bucket | Count | Percentage |
|--------|-------|------------|
| 1.6s | 484 | 1.6% |
| 1.8s | 775 | 2.6% |
| 2.0s | 520 | 1.7% |
| 2.4s | 1,869 | 6.2% |
| 2.6s | 7,226 | 24.1% |
| 2.8s | 8,553 | 28.6% |
| 3.0s | 4,822 | 16.1% |
| 3.2s | 2,741 | 9.2% |
| 3.4s | 1,268 | 4.2% |
| 3.6s+ | 1,681 | 5.6% |
82% of steps fall in 2.5–3.5s, but the tail (1.4s–5.9s) indicates large per-step compute variation. This variance also means `estimated_tokens_per_sec` reported in logs is not a reliable throughput metric, since the denominator (effective tokens) changes every step.
Contributor guide
Assessment
This issue has not been assessed yet.