Research archive

Training recipe and evidence

Archived reference. Read the current results →

This repository targets validation loss for 20M–92M Llama-style models on English C4, trained from scratch for 20 prediction targets per unique trainable parameter. The synchronous 20M preset uses the winner of the three-seed 128K-token sweep. The 50M/90M presets follow the earlier single-seed campaign. These selections do not establish optimality outside the tested settings.

Published precedents

The original Llama paper motivates pre-RMSNorm, RoPE, SwiGLU and bias-free projections. Embedding tying, exact small-model dimensions, context 1024, initialization std 0.02 with residual projection scaling, and the effective batch sizes are engineering choices. They are not directly established as optimal by these papers.

Selected recipe

• BF16 autocast for CUDA operations; FP32 parameters, gradients, optimizer moments, normalization reductions, and cross-entropy. No FP16 gradient scaler.
• AdamW: LR 0.0056 and beta2 0.98 for synchronous 20M; LR 0.001 and beta2 0.99 for 50M/90M. All use beta1 0.9, epsilon 1e-8, and weight decay 0.1. Matrix weights, including tied embeddings, receive decay; norm scales do not.
• Global raw-gradient clipping at norm 1.0, after accumulation and before AdamW.
• The reported studies used linear token-based warmup for 5% of training, then cosine decay to 10% of peak.
• Synchronous 20M: 131,072 prediction targets per update, microbatch 128, without accumulation. 50M/90M: 32,768 targets per update, microbatch 32. Shortened batches at virtual epoch boundaries use their actual target count as the gradient denominator.
• Zero dropout; seed 42; runtime.deterministic: false.

Current scheduler configs instead default to 312 warmup optimizer updates, configured by lr_schedule.warmup_steps; see the schedule guide. The schedule is evaluated at the end-token position of each optimizer update. Virtual epoch boundaries do not reset the schedule, optimizer, or data order. The token budget counts tied embeddings once; the report also records the non-embedding count to avoid confusing different papers’ size conventions.

Current training uses loader v2: consecutive groups of cache shards are visited in file order, with independently shuffled sequence indices within each group. data.shuffle_group_size defaults to 2 shards; one prefetched group overlaps I/O with model computation. The stream is independent of the total training budget. The audited campaign in runs/campaign used loader v1 for all twelve runs: sequence-aligned 64 MiB ranges were shuffled globally, then shuffled internally. Those results remain tied to their original implementation.

Historical basis for the 50M/90M recipes

1. 20M, six runs: LR {0.0003, 0.001, 0.003} × decay {0, 0.1}; beta2=0.95.
2. 50M, four runs: best two LR/decay pairs × beta2 {0.95, 0.99}.
3. 90M, two runs: best two complete recipes from the 50M stage.

Every completed run trains for 20 targets per parameter and evaluates the full official validation split at the end. Promote by final-checkpoint full-validation loss, breaking exact ties by lower LR, then decay, then beta2. Intermediate subset-validation best checkpoints are recorded separately. Architecture, data, seed, effective batch, schedule, clipping, and beta1 remain fixed in the sweep.

This is a single-seed comparison. It does not measure seed robustness, establish batch-size optimality, or provide an unbiased test-set performance estimate. Validation data is explicitly used for recipe selection.

Precision, reproducibility, and performance

The synchronous 20M recipe uses microbatch 128 and a 131,072-target batch, validated in the 128K-token sweep on GH200. The 50M/90M recipes use microbatch 32 and a 32,768-target batch. All three use runtime.compile=true, runtime.compile_mode=default, automatic SDPA, and eight CPU threads.

Validation defaults to 128 sequences per forward pass for all three sizes, without gradient accumulation. Both subset and full validation sum losses over valid tokens and divide by their total count, including partial final batches and excluding padding. Changing the batch size preserves this weighting but can introduce small floating-point differences. This larger validation batch is a subsequent default change; the recorded GH200 performance measurements used evaluation batch size 8.

deterministic=false retains seeded initialization, seeded block ordering, and saved RNG states while allowing fast kernels. It does not promise bitwise replay. deterministic=true enables strict PyTorch deterministic algorithms, configures cuBLAS before CUDA initialization, disables cuDNN benchmarking and compilation, uses reference attention, and disables fused/foreach AdamW. Unsupported operations raise errors. Cross-hardware or cross-version bitwise equality is not guaranteed.

The benchmark runs each candidate in a separate process, using full optimizer updates at the same effective batch size. It measures compiled/uncompiled paths and microbatches 4, 8, 16, and 32, rejecting candidates above 90% GPU memory usage. Candidates warm up for 20 updates and measure three 100-update windows. The selected configuration is saved rather than recomputed during training. Real-C4 benchmarks and a bounded GH200 tuner are also available; see GH200 measurements. Synthetic throughput estimates exclude data loading, validation, compilation, and checkpoint overhead; training JSONL records observed throughput.

The fast path uses PyTorch SDPA with automatic kernel selection (including FlashAttention and cuDNN), or an explicitly requested backend. The reference path uses explicit matmul, softmax, and masking and retains FP64 for second-order analysis. Copy weights with strict load_state_dict; disable AMP and compilation for derivative calculations. These tests validate derivatives of model losses, not differentiation through the optimizer’s training history.

Packed decentralized simulation

The default packed-4 preset uses the three-seed 128K-token sweep winner: LR 0.008, beta1 0.95, beta2 0.99, four local models, microbatch 32 per model, and one_peer_exponential topology. Its global batch is 131,072 targets, without accumulation. The adaptive-consensus preset retains its separate 32K-token recipe; adaptive consensus was disabled in this sweep.

The opt-in decentralized path keeps the global token budget based on one local model’s parameter count. It executes N workers together and assigns every Nth sample from the buffered loader’s shuffled stream to each worker. This preserves disjoint worker partitions across buffer and epoch boundaries, without separate per-worker file reads. The same seed initializes every worker exactly as the ordinary Llama path, without consuming additional initialization RNG.

There is no gradient accumulation: batch_tokens = N * micro_batch_size * context_length. For worker i, the gradient is the derivative of its local microbatch’s mean valid-token loss. Backpropagating the sum of worker means preserves this normalization. Each step clips gradients independently, mixes parameters, and applies local AdamW updates using gradients evaluated at the pre-mixing parameters. Weight decay acts on mixed weights; first/second moments and counters stay local. Shortened epoch-ending steps use equal smaller local batches and their actual token counts. Incompatible epoch boundaries are rejected.

Parameters, first moments, and second moments each occupy an aligned contiguous [N, D_padded] arena. Local AdamW tensors alias these arenas, and checkpoints serialize the three blocks once plus counters and layout metadata. Alignment padding is excluded from parameters and optimization. The moment arenas are exposed for future experiments but are not mixed by the trainer. Packed buffered checkpoints use format v3 and preserve loader v2’s committed cursor; ordinary buffered checkpoints retain upstream format v2. Prefetch may change on resume, but shard-group size may not. Historical loader checkpoints require their original source version.

Every training step performs one topology event: complete averaging, alternating left/right one-peer ring averaging, or one-peer exponential averaging with offsets 1, 2, 4, … below N. Peers are incoming, so row i receives row (i - offset) % N. One-peer weights are half self and half peer. For exponential graphs, powers of two admit exact cycle averaging; arbitrary N remains supported without that guarantee. See the exponential-graph paper.

Evaluation first snapshots the global parameter mean into an ordinary Llama. The snapshot uses FP32 averaging and the global evaluation batch size. It never changes local training states. Selection and reported validation losses concern this averaged model, whose ordinary-format weights are exported each epoch only with training.checkpoint_policy: interval. The default policy final saves only final.pt.

See the SLURM launcher setup for submission on the aarch64 GPU nodes, automatic CPU allocation, and console logs. benchmark-packed compares N=4/8 packed and sequential workers with matched global batches, reports end-to-end and separate component timings, and saves profiler traces.