Balance Batch: Token-Budget Micro-Batching

[WentseChen]

Sequences are grouped into fixed-size micro-batches (N sequences each), padded to the longest in each batch. With mixed-length workloads (64–32768 tokens), this wastes significant compute on pad tokens.

Add max_tokens_per_microbatch config: sort by length, greedily fill micro-batches to a token budget so batch_size * max_seq_len ≤ max_tokens.

Phase 1: Core Helper

Add make_token_balanced_microbatches to hosted_tinker/_utils.py:

def make_token_balanced_microbatches(
    sorted_indices: list[int],
    all_input_ids: list[list[int]],
    max_tokens: int,
) -> list[list[int]]:
    """Greedy bin-pack sorted indices so each micro-batch has ≤ max_tokens padded tokens."""
    micro_batches, current, current_max = [], [], 0
    for idx in sorted_indices:
        seq_len = len(all_input_ids[idx])
        new_max = max(current_max, seq_len)
        if current and (len(current) + 1) * new_max > max_tokens:
            micro_batches.append(current)
            current, current_max = [idx], seq_len
        else:
            current.append(idx)
            current_max = new_max
    if current:
        micro_batches.append(current)
    return micro_batches

Single oversized sequences are never dropped — they get their own micro-batch.

Phase 2: FSDP2 Backend

hosted_tinker/fsdp2_backend.py

• Add max_tokens_per_microbatch: int | None = None to FSDP2BackendConfig
• Add _max_tokens_per_microbatch state + _refresh_max_tokens_per_microbatch() reading /dev/shm/hosted_tinker_mtpm + set_max_tokens_per_microbatch() writing same file — mirror of existing _micro_batch_size pattern
• Pass "max_tokens_per_microbatch" in command dict to worker

hosted_tinker/fsdp2_worker.py

FSDP requires all ranks to call model() the same number of times.

• Read max_tokens_per_microbatch from command dict
• If set, use make_token_balanced_microbatches(my_indices, all_input_ids, max_tokens) to form variable-size micro-batches
• Else fall back to fixed micro_batch_size slices (existing behavior)
• dist.all_reduce(n_t, op=MAX) on NCCL (not obj_group) to sync micro-batch count across ranks
• Pad schedule with None dummy sentinels for ranks with fewer micro-batches
• Loop body: iterate mb_schedule, use None sentinel for dummies (generalized from existing empty-rank pattern)

Phase 3: Megatron Backend

hosted_tinker/megatron_backend.py

• Add max_tokens_per_microbatch: int | None = None to MegatronBackendConfig
• Pass in command dict to worker

hosted_tinker/megatron_worker.py

• Read max_tokens_per_microbatch from command dict
• Replace fixed loop with token-balanced version when set
• No FSDP sync needed — ranks process independently; dist.all_gather_object at end handles result collection

Phase 4: PyTorch Backend + API Endpoint

hosted_tinker/pytorch_backend.py

• Add max_tokens_per_microbatch: int | None = None to PyTorchBackendConfig
• Add length-sorting before batching (currently absent)
• Switch result storage from append to index-by-original-position ([None]*n_examples) since request_batch_slices indexes by original position
• Use helper when set, else fall back to fixed micro_batch_size slices over sorted indices

hosted_tinker/api.py

Add /admin/set_max_tokens_per_microbatch runtime endpoint (writes /dev/shm/hosted_tinker_mtpm). FSDP2 picks up changes immediately; other backends need restart.

Phase 5: Benchmark & README

Sweep max_tokens_per_microbatch values on both FSDP2 and Megatron DDP (4× B200, Qwen3.5-35B-A3B, gc=on). Same 128-example mixed-length workload as existing README table.

Values to test

max_tokens_per_microbatch	Rationale
16384	Conservative — ~1 long seq or many short seqs per micro-batch
32768	Medium — roughly 1× max_seq_len
65536	Aggressive — 2× max_seq_len, packs many short seqs
131072	Very aggressive — may OOM on fwd+bwd

Procedure per backend

For each mtpm value:

1. Start server with max_tokens_per_microbatch=<value>, gc=on
2. Run bench_gpu_throughput.py (128 examples, 64–32768 tok, 1 warmup + 3 repeats)
3. Record: fwd tok/s, fwd GPU util%, fwd+bwd tok/s, fwd+bwd GPU util%, GPU mem%

Existing baseline (from README)

backend	GPUs	mbs	fwd tok/s	GPU util (fwd)	fwd+bwd tok/s	GPU util (fwd+bwd)
FSDP2	4	1	23,403	64%	2,550	76%
FSDP2	4	2	27,032	73%	2,631	87%
Megatron DDP	4	1	23,276	56%	2,788	64%
Megatron DDP	4	2	23,429	57%	2,798	64%

Add balanced-batch results to README table alongside existing mbs rows.

Phase 6: Pad-Token Efficiency in Benchmark

Add pad-efficiency reporting to benchmarks/bench_gpu_throughput.py only (not in production). Simulate micro-batch grouping client-side from known sequence lengths to compute real_tokens / (batch_size * max_len) per micro-batch. Print alongside tok/s and GPU util.