micm_nlp.training.batching

Token-budget batching for eval/test: variable-size batches whose total token count is bounded by an auto-detected per-GPU budget.

Two pieces:

  • calibrate_token_budget — binary search on sorted lengths. Finds the largest k such that (k, padded(L_k)) (the widest shape the sampler can yield at budget = k × L_k) fits in VRAM. Real-shape probe, deterministic.

  • TokenBudgetBatchSampler — length-sorted BatchSampler that yields variable-size index batches bounded by that budget.

Designed for eval/test only: variable batch size affects optimizer dynamics in training but is semantically invisible for forward-only inference. training_args.group_by_length is ignored on this path — the sampler always length-sorts internally.

Classes

TokenBudgetBatchSampler

Length-sorted batch sampler with a per-batch token cap.

Functions

calibrate_token_budget(→ int)

Find the largest token budget such that the sampler's widest batch

Module Contents

class micm_nlp.training.batching.TokenBudgetBatchSampler(lengths: collections.abc.Sequence[int], token_budget: int, pad_multiple: int = 1)

Bases: torch.utils.data.Sampler[list[int]]

Length-sorted batch sampler with a per-batch token cap.

Parameters:

  • lengths – per-sample sequence lengths (ints), same length as the dataset.

  • token_budget – max (batch_size * padded_max_length) per batch.

  • pad_multiple – alignment for padded length (matches data_collator’s pad_to_multiple_of); 1 disables rounding.

Notes

  • A sample longer than token_budget is still yielded — as a singleton batch. We don’t silently drop test samples; the caller must either raise the budget or accept the OOM risk for that batch.

  • __len__ is exact (computed once by simulating one iteration).

Parameters:

  • lengths – per-sample sequence lengths, one per dataset row.

  • token_budget – cap on batch_size * padded_max_length per batch.

  • pad_multiple – alignment for the padded length; 1 disables rounding.

Raises:

ValueError – if token_budget is not positive or pad_multiple is below 1.

property order: list[int]

The order samples are emitted in, length-ascending.

Public because batches are not in dataset order: a consumer that needs to line predictions back up with dataset rows has to know the permutation.

micm_nlp.training.batching.calibrate_token_budget(*, model, lengths: collections.abc.Sequence[int], pad_multiple: int = 1, floor: int = 256, tolerance: int | None = None) → int

Find the largest token budget such that the sampler’s widest batch fits in VRAM.

The sampler yields length-sorted batches packed to budget total tokens. For budget B and sorted lengths, the first (widest) batch has shape (k, padded(L_k)) where k samples fit before the budget is exceeded. We binary-search k to find the largest fitting shape.

Because the sampler’s batches all have total_tokens ≤ B, and Aya’s sdpa→flash attention is linear in seq_len (not quadratic), and the LayerNorm-fp32-cast is linear in total_tokens, every batch at this budget will fit if the calibration shape fits — assuming linear-cost attention (sdpa/flash) and a LayerNorm-style activation cost that scales with total tokens. Quadratic-attention models (eager) may violate this for narrow-tall batches.

Parameters:

  • model – callable model on its target device.

  • lengths – per-sample sequence lengths (the dataset’s length column).

  • pad_multiple – collator’s pad_to_multiple_of (1 = no rounding).

  • floor – raise if no fitting shape produces a budget at or above this.

  • tolerance – deprecated and ignored. The search now always runs to convergence (exact largest fitting k). Kept for call-site compatibility.

Returns:

token budget = (largest fitting k) × padded(L_k) × _HEADROOM.

Return type:

int

Raises:

  • ValueError – empty lengths.

  • RuntimeError – not even shape (1, padded(L_1)) fits, or returned budget < floor.