Batch Rebalance Data Sampler
Author: delock
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47340)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47340)
<!-- ci-dashboard-badge:end -->

# What does this PR do?
This PR introduce a training data sampler using a batch rebalance technique that reduces padding, balance cost between ranks and grad accumulation steps, which improves training efficiency.

Different from packed sequence data sampling, batch rebalance data sampling does not require varlen support in attention implementation.  This makes it idea for researching or hardware that does not have varlen attention kernel implementation yet.

Different from group length sampler, this sampler does not require a global sort, so the sampler maintains true randomness in each effective batch.

## How it works
The basic idea of batch rebalance data sampler is like this:

In distributed training, each rank processes a subset of samples called a micro-batch. When samples have different sequence lengths, the rank that happens to get longer sequences takes more time to compute, while other ranks finish early and wait. This idle waiting wastes GPU time.

Batch rebalance solves this by looking at all samples in a global batch, sorting them by length, and distributing them across ranks so that each rank gets a roughly equal computational cost — not an equal number of samples, but an equal workload. Short-sequence ranks get more samples; long-sequence ranks get fewer. The assignment is recomputed for every global batch using a lightweight greedy heuristic, with no communication between ranks.

For example, say you have 2 ranks, gradient accumulation of 1, and a global batch of 8 samples with token lengths [128, 256, 512, 768, 1024, 1536, 2048, 4096]. A random split assigns 4 samples per rank. If rank 0 happens to get [128, 512, 1024, 4096], every sample in that micro-batch is padded to 4096 — even the 128-token one — so rank 0 computes 4×4096 = 16,384 padded tokens. Meanwhile rank 1 gets [256, 768, 1536, 2048], padding to only 2048, computing 4×2048 = 8,192 tokens. Rank 0 takes twice as long, and rank 1 sits idle.

Batch rebalance sorts all 8 samples and assigns them in unequal counts: rank 0 gets just the 2 longest samples [4096, 2048] (2×4096 = 8,192 tokens), while rank 1 gets the remaining 6 shorter samples [1536, 1024, 768, 512, 256, 128] (6×1536 = 9,216 tokens). The workload is now nearly equal — no idle waiting.

Batch rebalance sampler uses a cost model to estimate the compute of each group of samples, then balances the total cost across ranks within the same synchronization step.

Specifically, for each global batch, the sampler sorts samples by length, partitions them into dp_size × grad_accum groups, and assigns groups to ranks so that ranks sharing the same gradient accumulation slot receive similar costs. Longer-sequence groups contain fewer samples; shorter-sequence groups contain more — but the estimated compute is even.

The default cost model is intercept + bs × max_len + alpha × bs × max_len², where bs is the number of samples in the group and max_len is the longest sequence. The quadratic term captures the super-linear scaling of attention. Users can plug in their own cost function via the cost_fn parameter to match their model's actual compute characteristics.

## How to use
Set `--train_sampling_strategy batch_rebalance` in your training command. No other changes to your dataset or model code are needed.

python train.py \
    --train_sampling_strategy batch_rebalance \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 2

Optional tuning parameters:

| Parameter | Default | Description |
|---|---|---|
| --batch_rebalance_alpha | 0.001 | Quadratic cost weight. Increase if your model's attention cost dominates (e.g. long context). Set to 0.0 to use a linear-only cost model (bs × max_len). |
| --batch_rebalance_max_tokens | 0 | Safety cap on padded tokens (bs × max_len) per micro-batch to prevent OOM. 0 = no cap. |

I confirm that this PR description and code is written with aid of LLM, and reviewed by a human being prior to PR submission.

## Before submitting
[ ] (First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://huggingface.co/docs/transformers/contributing) and the
      [Pull Request](https://huggingface.co/docs/transformers/pr_checks) checks?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes according to the [guidelines](https://github.com/huggingface/transformers/tree/main/docs)?
- [x] Did you write any new necessary [tests](https://huggingface.co/docs/transformers/testing)?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
- trainer: @SunMarc
- distributed: @3outeille @ArthurZucker

Comments:
<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47340)
**Latest run:** [29406918422:2](https://github.com/huggingface/transformers/actions/runs/29410765820)
**Result:** `failure` | **Jobs:** 15 | **Tests:** 171,078 | **Failures:** 0 | **Duration:** 13h 10m

<!-- ci-dashboard-recap:end -->

### GPU DDP benchmark

Ran the sampler end-to-end through the `Trainer` on a real multi-GPU setup to validate the distributed path.

**Setup**
- Hardware: 2× NVIDIA RTX 5090 (real DDP, `world_size=2`)
- Model: `Qwen/Qwen2.5-0.5B-Instruct` (bf16, gradient checkpointing on so all three strategies fit at the chosen batch size)
- Data: 1600 synthetic sequences, lengths ~ `lognormvariate(mu=5.0, sigma=1.0)` clamped to `[16, 1024]`
- Config: `per_device_train_batch_size=8`, `gradient_accumulation_steps=2` → **effective batch size = 32**, 50 optimizer steps
- Methodology: the benchmark sizes the dataset to exactly `max_steps × eff_bs`, so every strategy processes the **identical set of samples** (375,880 real tokens) once — an apples-to-apples comparison. It also verifies per-epoch coverage (`compute_loss` samples seen, all-reduced) and reports **real-tokens/s** (= real tokens / end-to-end wall time, padding excluded), not the config-based `train_samples_per_second`.

**Results**

| sampler | real tok/s | speedup vs random | step time median / mean / p90 / max (s) |
|---|---|---|---|
| random (distributed) | 15,873 | 1.00× | 0.437 / 0.440 / 0.540 / 0.665 |
| group_by_length | 24,485 | 1.54× | 0.232 / 0.270 / 0.395 / 0.640 |
| **batch_rebalance** | **25,085** | **1.58×** | 0.256 / 0.269 / **0.287** / 0.613 |

**Step 5 sample distribution** (step 5 = `global_batch=32`; 2 ranks × 2 grad-accum slots = 4 micro-batches; per-rank/slot `bs` and per-element lengths shown descending; each micro-batch is padded to its max):

```
random                         group_by_length                batch_rebalance
r0/s0 bs=8                     r0/s0 bs=8                     r0/s0 bs=2
 [495,354,343,257,187,121,      [244,242,242,241,239,239,      [806,753]
        92, 91]                        238,237]               r1/s0 bs=8
r1/s0 bs=8                     r1/s0 bs=8                      [249,246,243,215,203,
 [409,383,363,275,272, 76,      [235,234,234,232,232,226,            197,186,143]
        49, 41]                        225,225]               r0/s1 bs=3
r1/s1 bs=8                     r0/s1 bs=8                      [537,495,346]
 [735,613,388,257,205,184,      [223,222,215,213,210,208,     r1/s1 bs=19
       173,134]                        208,207]                [103,90,82,79,72,67,65,
r0/s1 bs=8                     r1/s1 bs=8                            60,57,57,54,54,50,
 [278,179,173,171,118, 80,      [206,205,201,200,196,195,            40,35,34,31,28,16]
        65, 34]                        194,193]
pad target (max) varies wildly  pad target tight per batch     long groups shrink (bs 2-3),
(91 -> 735); ranks imbalanced   but slot0 ~240 vs slot1 ~200   short groups grow (bs 8,19);
                                                               cost balanced across ranks/slots
```

**Reproduce:** benchmark script is at https://gist.github.com/delock/516efc09d9a4217e209729e91b9294de (`bench_sampler.py`). Run one strategy per invocation:

```
torchrun --nproc_per_node=2 bench_sampler.py --strategy batch_rebalance --gradient_checkpointing
```


### `group_by_length` vs `batch_rebalance`: peak memory and max-batch throughput

A closer look at the trade-off between the two samplers. Same setup as the comment above, but **no gradient checkpointing** — so the memory limit each sampler hits is its own.

**Setup:** 2× NVIDIA RTX 5090 (32 GB), `Qwen/Qwen2.5-0.5B-Instruct` (bf16), real DDP (`world_size=2`), `gradient_accumulation_steps=2`, lengths ~ `lognormvariate(mu=5.0, sigma=1.0)` clamped to `[16, 1024]`. Peak VRAM = max across ranks of `torch.cuda.max_memory_allocated`. Reproduce: https://gist.github.com/delock/516efc09d9a4217e209729e91b9294de (`bench_sampler.py`).

#### 1) Peak VRAM vs `per_device_train_batch_size`

| per_device_bs | group_by_length | batch_rebalance |
|---|---|---|
| 4  | 17.7 GB | 12.3 GB |
| 6  | 24.5 GB | 15.4 GB |
| 8  | 31.3 GB | 20.4 GB |
| 10 | OOM     | 21.7 GB |
| 12 | OOM     | 24.5 GB |
| 14 | OOM     | 27.1 GB |
| 16 | OOM     | ~30 GB (borderline) |
| 18 | OOM     | OOM |

- `group_by_length` OOMs at `per_device_bs ≥ 10` → largest reliable batch = **8** (peak ~31 GB, right at the edge).
- `batch_rebalance` fits reliably up to **14** (peak ~28–30 GB); **16** is borderline (fits on short runs, OOMs over a full epoch); **≥18** OOMs.

→ `batch_rebalance` tolerates roughly **2× the per-device batch** before OOM (8 → 14 reliable, 16 borderline; effective batch 8×2×2 = 32 vs 14×2×2 = 56), and at every batch size its peak is ~30–40% lower.

**Why:** `group_by_length` globally sorts by length, so the longest sequences land in a single micro-batch (`per_device × 1024`), and that one batch sets peak VRAM (the LM-head logits + activations dominate). `batch_rebalance` balances cost within each global batch, so the long-sequence group shrinks to just 2–3 samples → padded-tokens-per-micro-batch is bounded → no single batch spikes.

#### 2) Throughput when each runs at its own max batch

| sampler | max per_device_bs | real tok/s | peak VRAM |
|---|---|---|---|
| group_by_length | 8  | 30,174 | 32.3 GB |
| **batch_rebalance** | **14** | **37,957** | 29.6 GB |

→ `batch_rebalance` delivers **~1.26× higher real-tokens/s** (37,957 vs 30,174) — and does so at **lower** peak VRAM (29.6 vs 32.3 GB).

The gain is not from per-step efficiency (at equal `per_device_bs` the two are close). It comes from `batch_rebalance` being able to use a larger batch on the same GPU → better GPU occupancy → more real tokens/s. `group_by_length` is memory-capped at `per_device_bs=8` by its long-batch spike; `batch_rebalance` converts the extra memory headroom into throughput.

**Net:** compared to `group_by_length`, `batch_rebalance` has lower peak VRAM at every batch size, tolerates ~2× the batch before OOM, and when each is run at its max batch it's ~1.26× faster in real tokens/s.


Thank you for your contribution 🤗!

## CI Security Gate — automatic approval blocked

This PR was **not** automatically approved for CI because the security gate failed.

**Possible reasons:**
- The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
- A changed file is outside the allowed directories (`src/`, `tests/`, `docs/`, `utils/`), has a disallowed extension (only `.py`, `.txt`, `.md` permitted outside `tests/` and `docs/`), or is not `.md`/`.yml` inside `docs/`
- A new high-severity security issue was detected in the changed Python files (Bandit check)

See the [workflow run](https://github.com/huggingface/transformers/actions/runs/29479726367) for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.