adamskrodzki/batch-size-experiments

Tokens per optimizer step: fixed versus scheduled

This repository studies whether changing the number of training tokens used for each optimizer update can improve language-model loss.

The model, model context, dataset, token order, optimizer family, and hardware configuration are held fixed. The experimental variable is the schedule of tokens_per_step (T):

• Constant-T baseline: the same T is used for every optimizer update.
• Scheduled-T treatment: T changes according to a predefined schedule during one run.

The two budget comparisons are independent experiments:

1. Fixed total tokens: every run consumes the same number of input tokens; compare final validation loss.
2. Fixed wall-clock time: every run is measured for the same elapsed time; compare validation loss at the deadline and report how many tokens each run processed.

The question is not whether a larger model context is better. The context is fixed. The question is whether changing the effective batch size and optimizer-update frequency during training leads to a better model than keeping them constant.

Hypotheses

Primary hypothesis

With model context and all other training conditions fixed, a scheduled tokens_per_step can produce a different final validation loss from a constant-T baseline.

Fixed-token hypothesis

When runs process the same total number of input tokens, at least one scheduled-T treatment will achieve lower final validation loss than the constant-T baseline.

Fixed-time hypothesis

When runs train for the same wall-clock duration, at least one scheduled-T treatment will achieve lower validation loss at the deadline than the constant-T baseline.

The direction and useful shape of the schedule are not assumed in advance. A scheduled T changes both gradient-noise scale and the number of optimizer updates per trained token, so the number of optimizer updates and the learning-rate schedule must be reported with every result.

Related work

This experiment is closely related to Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws (Wang et al., 2026). That paper studies batch-size schedules under a fixed data budget and reports that using smaller batches for most of training before switching to larger batches late can outperform constant-batch baselines. With fixed seq_len, its batch-size variable corresponds to this experiment's tokens_per_step; this repository tests the same general idea on a small Polish language model and adds an independent fixed-wall-clock comparison.

A second reference is Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling (Meterez et al., 2026). It proposes coupling batch-size increases with learning-rate decay, arguing that a larger batch can replace part of the learning-rate reduction while reducing serial optimizer steps. This motivates a possible extension beyond the primary experiment: compare constant T with scheduled T under unchanged cosine decay, then test a Seesaw-inspired schedule that changes T and the learning rate together.

Bonus research direction: joint T and learning-rate scheduling

The primary experiment isolates the effect of tokens_per_step while keeping the current cosine learning-rate schedule fixed. A bonus study can test whether jointly scheduling T and the learning rate produces a better loss-throughput trade-off than changing either schedule alone. Candidate arms are:

1. Constant T with cosine decay.
2. Scheduled T with the same cosine decay policy.
3. Scheduled T with a coupled learning-rate schedule inspired by Seesaw.

This extension must define the learning-rate schedule against a common coordinate—preferably cumulative input tokens—so changing T does not silently change when learning-rate decay occurs. It should report validation loss, cumulative tokens, optimizer updates, wall-clock time, and the exact T/learning-rate trajectory.

Experimental variables

Fixed controls

The following remain fixed within a comparison:

• MiniGPT architecture and initialization seed.
• seq_len, the fixed context length supplied to the model. The default is 1024 and it does not change during the intended experiment.
• Dataset, tokenizer, packed token stream, block permutation, and validation split.
• Optimizer, learning-rate policy, weight decay, gradient clipping, precision, compilation mode, and hardware. The current learning-rate schedule is linear warmup, cosine decay from lr to min_lr, then a constant min_lr after the horizon.
• device_token_capacity, the hardware micro-batch capacity.

Experimental variable: tokens_per_step

T is the number of input tokens accumulated into one optimizer update. With a fixed seq_len:

effective_batch_size = T / seq_len
device_batch_size    = device_token_capacity / seq_len
accumulation_steps   = T / device_token_capacity

Changing T therefore changes the effective batch size, gradient accumulation count, and optimizer updates per input token. It does not change the model context or the compiled input shape.

Each candidate value of T must satisfy the divisibility constraints required by the data loader and hardware micro-batching. The implementation currently requires T to be divisible by seq_len and device_token_capacity.

The experiment uses one fixed seq_len for every run. It does not use a 256 → 512 → 1024 → 2048 growing-context schedule.

Tokens-per-step calculation and learning-rate scaling

The choice of tokens_per_step is informed by How Does Critical Batch Size Scale in Pre-training, which studies how batch size, data size, optimization steps, and learning-rate choices interact during language-model pre-training. For these runs, the baseline value is calculated as:

tokens_per_step = device_token_capacity × accumulation_steps
                 = 20,480 × 12
                 = 245,760 tokens per optimizer update

The scheduled runs change tokens_per_step by a divider and adjust the learning rate with the square root of the relative token count:

T_stage = T_base / divider
lr_stage = lr_base × sqrt(T_stage / T_base)
         = lr_base / sqrt(divider)

For the 12,4,1 schedule, the resolved token counts are 20,480, 61,440, and 245,760, with learning-rate multipliers 1 / sqrt(12), 1 / sqrt(4), and 1, respectively. The square-root rule is the experiment's scaling choice; the paper motivates treating batch size and learning rate as coupled hyperparameters but does not prescribe this exact schedule for these runs.

Budget comparisons

Fixed total-token comparison

All runs process the same cumulative number of input tokens. The primary result is final validation loss. Results also include:

• Number of optimizer updates.
• Learning-rate state at the end of training.
• Training loss, final validation loss, and optional periodic validation diagnostics.
• Wall-clock time and tokens per second.

A scheduled treatment will generally perform a different number of optimizer updates than a constant-T baseline for the same total token count. That is part of the treatment, not a hidden equivalence.

Fixed wall-clock comparison

All runs use the same wall-clock budget. Initialization, one warm-up optimizer update, and the final validation pass are excluded from the measured interval. The clock starts after the first update completes and any lazy compilation has been synchronized; the run then stops after the first completed update at or beyond the deadline. If the training stream is exhausted first, the run ends early; its recorded step and token counts show that it did not reach the requested time. The primary result is the validation loss measured immediately after training. Results also include:

• Total input tokens processed.
• Number of optimizer updates.
• Tokens per second.
• Training loss, final validation loss, and optional periodic validation diagnostics.
• The exact elapsed time and any final-step overshoot.

A time-budget comparison therefore measures the combined effect of optimization behavior and hardware throughput.

Dataset and tokenization

The dataset is SlayerLab/polish-dynaword, using:

data/european_hplt_v3_pl/european_hplt_v3_pl.parquet

prepare_data.py trains a 16,384-token byte-level BPE tokenizer, inserts <|eos|> between documents, and writes one flat uint16 token stream. Documents are packed without padding, so a training window may cross a document boundary.

The stream is independent of the experiment's tokens_per_step schedule. Training cuts the packed stream into fixed-length rows of seq_len; the stream hash and shuffle metadata are recorded with each run.

Prepare data with:

python prepare_data.py --out data/

For a small data-pipeline check:

python prepare_data.py --out data/ --limit 1000

Model

The model is an in-repository decoder-only GPT-style transformer with RoPE:

property	value
parameters	31,433,920
vocabulary	16,384
hidden size	448
layers	10
attention heads	7, head dimension 64
MLP size	1,792
normalization	pre-LN LayerNorm, no bias
positional encoding	RoPE, theta 10,000
weight tying	token embedding and LM head
dropout	0.0
training context	seq_len, default 1024

The model uses causal self-attention, GELU MLPs, fp32 parameters and optimizer states, and optional bf16 autocasting. The language-model loss predicts the next token, so each row of length seq_len contributes seq_len - 1 prediction targets.

Current implementation

train.py implements the constant-T baseline:

• One fixed seq_len per run.
• One fixed scalar tokens_per_step for the whole run.
• Gradient accumulation to realize that scalar T.
• A fixed total-token budget or a fixed wall-clock budget.
• Seeded block order, automatic final validation, optional periodic validation for debugging, checkpoint saving, and run metadata.

scheduled_train.py implements the scheduled treatment without changing train.py:

• --batch-dividers D1,D2,...,Dn creates n stages. Stage i requests tokens_per_step / Di input tokens per optimizer update.
• --stage-length L1,L2,...,L(n-1) gives the relative lengths of the first n - 1 stages; the last stage receives the remaining fraction.
• Each requested stage size is rounded to the nearest positive multiple of device_token_capacity, with ties rounded upward. The resolved stage sizes are recorded in run_meta.json.
• Token-budget stages use cumulative input tokens as their coordinate. Time-budget stages use measured elapsed seconds after the warm-up update.
• Stage changes happen between complete optimizer updates. A token-budget run does not start an update that would exceed its budget; a time-budget run stops after the first completed update at or beyond its deadline.
• The global warmup/cosine learning-rate curve is multiplied by 1 / sqrt(divider) for the active stage.

For example, --tokens-per-step 98304 --device-token-capacity 8192 --batch-dividers 12,4,1 resolves to 8192, 24576, and 98304 tokens per update.

Running the current baseline

A fixed-total-token baseline:

python train.py \
  --seq-len 1024 \
  --tokens-per-step 65536 \
  --token-budget 20000000 \
  --out runs/constant_t_tokens

A fixed-wall-clock baseline:

python train.py \
  --seq-len 1024 \
  --tokens-per-step 65536 \
  --time-budget 120 \
  --out runs/constant_t_time

A three-stage scheduled run with approximately equal one-minute stages:

python scheduled_train.py \
  --seq-len 1024 \
  --tokens-per-step 98304 \
  --device-token-capacity 8192 \
  --batch-dividers 12,4,1 \
  --stage-length 0.333333,0.333333 \
  --time-budget 180 \
  --out runs/scheduled_12x_4x_3min

Exactly one of --token-budget and --time-budget is required. Token-budget runs measure the learning-rate schedule by cumulative input tokens, while time-budget runs measure it by elapsed wall-clock time after the first warm-up update. Initialization and warm-up are outside the measured interval, while subsequent training is inside it. Time-budget runs do not require a planned step count. run_meta.json records total steps and tokens separately from measured_steps, measured_tokens, and warmup_elapsed_seconds.

Outputs and comparison data

Each run writes:

runs/<name>/
  model/
  run.log
  run_meta.json

run_meta.json records the resolved scalar or scheduled tokens_per_step values, fixed seq_len, derived batch and accumulation values, budget, total and measured steps/tokens, warm-up and training timing, loss values, the budget-specific lr_schedule coordinate and horizon, stream hash, stream range, shuffle metadata, and the skip_validation setting. Scheduled runs additionally record requested dividers, relative stage lengths, normalized boundaries, per-stage resolved batch values, learning-rate multipliers, actual stage coordinates, and termination reasons. Every run performs one full validation pass after training and stores its result as validation_loss; this pass is outside the token or wall-clock training budget. If eval_every is positive, periodic validation is also run during training for debugging and stored separately as periodic_validation_loss. Pass --skip-validation only when validation is explicitly unnecessary; it disables both final and periodic validation and records skip_validation: true.

For the final experiment comparison, validation_loss is the primary model-quality metric. Training loss, processed tokens, optimizer updates, throughput, and timing are supporting measurements. The final validation pass is automatic by default and requires no special settings. All console output, warnings, and tracebacks are also written automatically to run.log.

Extracting and plotting training curves

parse_run_logs.py extracts one row for each logged training metric. It records the run name, step, cumulative input tokens, training loss, and elapsed wall-clock time in seconds. Pass only complete runs to the parser and exclude runs/capacity_search/ when it is not part of the comparison:

python parse_run_logs.py \
  runs/chinchila_20260917_063456_462446/run.log \
  runs/chinchila_scheduled_12x_4x_20260917_153350_360982/run.log \
  --output runs_metrics.csv

plot_run_metrics.py reads the CSV and creates separate loss curves against wall-clock time and cumulative tokens:

python plot_run_metrics.py runs_metrics.csv --output-dir runs/plots

The generated files are runs/plots/loss_vs_time.png and runs/plots/loss_vs_tokens.png. The time axis is in hours and the token axis is in millions; each run in the CSV is plotted as a separate colored line.

Loss versus wall-clock time

Loss versus tokens processed

Repository structure

• prepare_data.py — creates the tokenizer and packed token stream.
• train.py — trains one constant-T run and saves a checkpoint plus metadata.
• scheduled_train.py — trains a scheduled-T run and saves per-stage metadata.
• test_scheduled_train.py — tests schedule parsing, rounding, boundaries, and LR scaling.
• parse_run_logs.py — extracts plotting metrics from training logs.
• plot_run_metrics.py — creates loss-versus-time and loss-versus-token plots.
• configuration_minigpt.py — model configuration.
• modeling_minigpt.py — model implementation.
• config.json — shared training defaults.
• check_model.py — model shape, parameter-count, causality, and checkpoint checks.
• technical_debt.md — known differences between the intended experiment and the current implementation.