[delock]

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).
• Did you read the contributor guideline and the
  Pull Request checks?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes according to the guidelines?
• Did you write any new necessary tests?

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

[delock]

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

[delock]

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.

[github-actions]

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 for the exact violations.

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

[sywangyi]

self.args.train_batch_size is the batch_size per rank

[sywangyi]

The optimization may be valid in isolation, but the current Trainer integration changes distributed sampling semantics. In DDP it appears to double-shard rank-local batches through Accelerate, reduces the configured per-device batch size by dividing args.train_batch_size by world_size, and ignores dataloader_drop_last=False, causing silent sample drops. These are correctness regressions in common Trainer usage, even if the sampler can improve performance in the narrow direct/single-process case.

[sywangyi]

test example for your reference “accelerate launch --num_processes 2 example.py”

import os
import sys

import torch
from torch.utils.data import DataLoader, Dataset

from accelerate import Accelerator

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

from transformers.trainer_pt_utils import BatchRebalanceSampler


class IndexDataset(Dataset):
    def __init__(self, size=16):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        return index


def flatten(batches):
    return [sample for batch in batches for sample in batch]


def demonstrate_drop_last_regression():
    dataset_size = 18
    effective_batch_size = 16
    lengths = list(range(100, 100 + dataset_size))

    baseline = DataLoader(IndexDataset(dataset_size), batch_size=effective_batch_size, drop_last=False)
    baseline_batches = [batch.tolist() for batch in baseline]

    sampler = BatchRebalanceSampler(
        lengths=lengths,
        effective_batch_size=effective_batch_size,
        dp_size=1,
        grad_accum=1,
        shuffle=False,
        rank=0,
        drop_last=False,
    )
    sampler_batches = [list(batch) for batch in sampler]

    baseline_samples = set(flatten(baseline_batches))
    sampler_samples = set(flatten(sampler_batches))
    dropped_samples = baseline_samples - sampler_samples

    print("\nDrop-last regression")
    print(f"  dataset size = {dataset_size}")
    print(f"  effective batch size = {effective_batch_size}")
    print(f"  regular DataLoader(drop_last=False) batches = {baseline_batches}")
    print(f"  BatchRebalanceSampler(drop_last=False) batches = {sampler_batches}")
    print(f"  dropped despite drop_last=False = {sorted(dropped_samples)}")

    assert baseline_samples == set(range(dataset_size))
    assert sampler_samples < baseline_samples
    assert dropped_samples == {16, 17}


def main():
    accelerator = Accelerator()
    world_size = 2
    gradient_accumulation_steps = 2
    per_device_train_batch_size = 4
    expected_effective_batch_size = per_device_train_batch_size * world_size * gradient_accumulation_steps
    lengths = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1600]

    if accelerator.num_processes != world_size:
        raise RuntimeError("Run this script with: accelerate launch --num_processes 2 example.py")

    rank = accelerator.process_index
    sampler = BatchRebalanceSampler(
        lengths=lengths,
        effective_batch_size=expected_effective_batch_size,
        dp_size=world_size,
        grad_accum=gradient_accumulation_steps,
        shuffle=False,
        rank=rank,
    )
    dataloader = DataLoader(IndexDataset(), batch_sampler=sampler)

    before_prepare = [list(batch) for batch in sampler]
    prepared_dataloader = accelerator.prepare(dataloader)
    after_prepare = [batch.tolist() if isinstance(batch, torch.Tensor) else list(batch) for batch in prepared_dataloader]

    local_result = {
        "rank": rank,
        "before_prepare": before_prepare,
        "after_prepare": after_prepare,
    }

    gathered = [None for _ in range(world_size)]
    torch.distributed.all_gather_object(gathered, local_result)

    gathered = sorted(gathered, key=lambda item: item["rank"])
    before_by_rank = {item["rank"]: item["before_prepare"] for item in gathered}
    after_by_rank = {item["rank"]: item["after_prepare"] for item in gathered}

    before_samples = set(flatten(flatten(before_by_rank.values())))
    after_samples = set(flatten(flatten(after_by_rank.values())))
    dropped_samples = before_samples - after_samples

    if accelerator.is_main_process:
        print("Setup")
        print(f"  world_size = {world_size}")
        print(f"  gradient_accumulation_steps = {gradient_accumulation_steps}")
        print(f"  per_device_train_batch_size = {per_device_train_batch_size}")
        print(f"  expected effective batch size = {expected_effective_batch_size}")

        print("\nBefore accelerator.prepare: sampler is already rank-aware")
        for item in gathered:
            print(f"  rank {item['rank']}: {item['before_prepare']}")
        print(f"  samples covered before prepare = {sorted(before_samples)}")

        print("\nAfter accelerator.prepare(DataLoader(..., batch_sampler=sampler))")
        for item in gathered:
            print(f"  rank {item['rank']}: {item['after_prepare']}")
        print(f"  samples covered after prepare = {sorted(after_samples)}")
        print(f"  dropped by double-sharding = {sorted(dropped_samples)}")
        demonstrate_drop_last_regression()

    accelerator.wait_for_everyone()

    assert len(before_samples) == expected_effective_batch_size
    assert len(after_samples) < len(before_samples)
    assert dropped_samples, "accelerator.prepare unexpectedly preserved all rank-aware sampler batches"


if __name__ == "__main__":
    main()

[delock]

@sywangyi thanks for your comments. The latest commit has the fix for three issues you mentioned, please kindly review, thanks!

BTW, I have verified that the experiment environment actually contains the patches for the first and second issue in your comments, but they are omitted in the original commit. So the performance data is real. I reran the performance test with the patch I submitted and its consistent with previous result, thanks!

[sywangyi]

should be effective_batch_size = self.args.train_batch_size * grad_accum * world_size to cover DP case?

[delock]

Fixed. Thanks for the catch!

[sywangyi]

The current implementation preserves a fixed effective_batch_size even for the final partial batch by replaying up to effective_batch_size - remainder samples. That is not the usual drop_last=False behavior; it can massively overweight duplicated examples. If fixed-size final steps are intentional, this should be documented and probably opt-in. Otherwise the tail should be padded only to the minimum size needed for rank/gradient-accumulation alignment, accepting a smaller final optimizer step.

[delock]

That's right, we don't need to pad so many samples, fixed.

[sywangyi]

The current regression test constructs BatchSamplerShard directly and then mutates its attributes, which verifies the wrapper behavior in isolation. It does not exercise the actual Trainer._get_dataloader -> accelerator.prepare(DataLoader(...)) -> patch prepared_dataloader.batch_sampler path. Since the fix depends on the exact structure returned by accelerator.prepare, could we add an integration test that builds the Trainer dataloader and verifies the prepared dataloader preserves the rank-aware BatchRebalanceSampler coverage?

[SunMarc]

Thanks ! Will try to review this shortly

[SunMarc]

Thanks for your work @delock ! Happy to have this new sampler. Left a bunch of comments but I think most of them are easy fix and maybe a bit more iteration on the sampler to make it easier to understand, this will help a lot when it comes to maintaining it in the future ! Also, it would be nice to have some docs so that users learn about this feature !

[SunMarc]

Suggested change

	sampler: torch.utils.data.Sampler | None = None
	sampler = None

[delock]

Will fix as suggestd

[SunMarc]

can you explain better why batch_sampler vs sampler here ?

[delock]

BatchRebalanceSampler decide how many samples a batch (or microbatch) will have based on a cost model. It tries to maintain close cost of the batch among ranks in the same iteration, so number of samples in each batch may change. sampler cannot achieve this goal because sampler return one sample per iter and there is no hint how many iters needs to be called to form the next batch. batch_sampler returns a list of samples per iter which can be used to form next batch.

[SunMarc]

can you try to make this comment shorter and more comprehensible ?

[delock]

Will do

[SunMarc]

Suggested change

	# `self.args.train_batch_size` is the per-process batch (per_device_train_batch_size *
	# #GPUs in this process) and already excludes world_size. Multiplying it by grad_accum
	# and world_size yields the true global effective batch size across all processes and
	# accumulation steps, covering both DDP (world_size > 1, n_gpu == 1) and DP
	# (world_size == 1, n_gpu > 1).

[delock]

Will remove.

[SunMarc]

we should update here for batch rebalance no ?

[delock]

Yes, we need to update for batch rebalance too for its current implementation.

[SunMarc]

I think that until some users request them, we can remove some args that I think nobody will really use. I think we can remove both of them but happy to keep one if you think it is really important to have to tune. But in your bench, I didn't see you run some ablation on top of them.

[delock]

Yes true, they should be added when there's real request. I'll remove then from this PR. Thanks!

[SunMarc]

note that this will break when we add tp in it no ? maybe make a note that it is only tested with dp only so far

[SunMarc]

not used

[SunMarc]

is it needed ?

[delock]

Yes, it is needed so each epoch would reshuffle. Good catch that it was never called. I'll add a fix to call it at epoch boundary.

[SunMarc]

Can you try to make this implementation a bit simpler to read ? There are a bunch of methods that doesn't have explanation or don't really need to be a method . The simpler the code is implemented or read, the better it is !

[delock]

Yes, the original code is too long. I'll break them into smaller methods and add comments for better understanding, thanks!

[delock]

Hi @SunMarc , thanks for your review comments. I have updated the PR followup the comments. The balancing algorithm had been breakdown and commented to improve readability. Unused args without ablation is removed from arg list. Excessive UTs are removed and e2e test enhanced.

[SunMarc]

Thanks a lot for updating ! LGTM, can you just fix one last thing in the test ?

[SunMarc]

instead of recreating a new train script, can we kind of reuse the one we used for ddp/fsdp/deepspeed training (train.py in trainer/distributed/scripts folder )? If this requires a lot of changes, we can just copy the file and update it. This way we test on a real model.

[delock]

Thanks, I will reused the existing train.py, and add var length dataset capability to it (currently it will pad each sample to 128 tokens so rebalance won't take effect).

[delock]

For test_batch_rebalance_no_double_sharding, the training test alone cannot cover this because double-sharding will silently drop samples from the batch and current train.py cannot catch this situation which leads to a silent failure. Do you think we need a UT to catch this?

[SunMarc]

no it's fine !

[SunMarc]

up, just test that training works

[delock]

Thanks a lot for updating ! LGTM, can you just fix one last thing in the test ?

Hi SunMarc, thanks for your comments. I have update the code accordingly. Thanks!

[SunMarc]

Thanks for iterating ! Last thing is it would be great to have some minor docs around this feature if possible.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[delock]

@SunMarc @sywangyi thank you for your kindly review! I added a small section in docs/source/en/training.md to introduce training_sampling_strategy="batch_rebalance" as an optimization.

[SunMarc]

Awesome work, thanks !

[stevhliu]

Thanks for adding!

I think it may be better to add this as a short subsection to trainer_recipes.md (next to ## Dataloader performance) rather than training.md which is more of a basic fine-tuning guide where the defaults are expected to work reasonably well.

Some other little notes:

• brief sentence about the problem it solves would be nice (rank idle time from padding on variable-length data) and that its DP only
• maybe a snippet showing the config for it
• per_device_train_batch_size is an average in this case and not a per-step count (may cause confusion if users see per-step batch sizes that don't match their config)
• something about the length being precomputed and skipping the full-dataset length scan at startup may also be helpful

[delock]

Hi @stevhliu , thanks for the suggestion. I have moved the doc to trainer_recipes.md, added a small sub section after Dataloader performance. The points in your comment are covered, thanks!

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33409906363:2
Result: success | Jobs: 16 | Tests: 180,127 | Failures: 0 | Duration: 14h 38m