GPT-2 small with warm-target OLMo 2 OT initialization

Final checkpoint of gpt2-olmo2-1b-warm120-ot-muon-lr4e-3-tok401k-seq1024-cos-wu5-485335. This is a GPT-2 small base language model trained on C4 English, with selected attention weights initialized through activation-based optimal transport from OLMo 2 1B. It is a GPT-2 architecture with 124,475,904 parameters, 12 layers, width 768, 12 attention heads, tied input/output embeddings and a 1024-token context. It uses the GPT-2 tokenizer (50,257 tokens), with a padded model vocabulary of 50,304.

Initialization and training

1. Start GPT-2 from random Gaussian initialization and train it briefly on C4. The shared warm stage completed 171 updates, 68,640,768 input tokens and 115.05 seconds of measured training-loop time.
2. Align this warm target with frozen pretrained OLMo 2 1B using matched text prefixes, activation correlations and feature/layer optimal transport. Donor revision: a1847dff35000b4271fa70afc5db10fd29fedbdf. For selected attention channels, blend 90% target weights with 10% mapped, RMS-calibrated donor weights. Select 128 channels per projection and cap donor contribution scaling at 10. Embeddings, MLPs and layer norms retain their warm-target values. initialization_report.json records the exact masks, sampled rows, maps and diagnostics.
3. Reset optimizers and train all target weights normally with Muon plus AdamW for embeddings, norms and biases. Peak LR is 4e-3, with elapsed-time cosine decay and 5% warmup. Use two H200 GPUs, batch98 per GPU and accumulation2: 401,408 input tokens per optimizer update, all sequences length1024.

The downstream job completed 3,806 updates and processed 1,527,758,848 input tokens. Training-loop time was 1710.03 seconds, within a nominal30-minute Slurm allocation. Warm training and OT preparation are additional compute.

Final validation cross-entropy: 3.325869. Final validation perplexity: 27.823174. Evaluation uses 2,048 held-out C4 English documents truncated to1024 tokens, with padding masked and token-weighted cross-entropy measured in BF16. W&B training record. Training uses BF16 autocast with FP32 master parameters; this checkpoint stores FP32 weights.

Experimental scope

This is the warm-target OT v1 experiment, an adaptation inspired by Transport and Merge. It is an exploratory extension to a briefly trained GPT-2 target, rather than a faithful reproduction of the paper's pretrained-target setting. OT initialization changed calibration CE from 5.501624 to 5.505724. A continuation without donation from the same warm checkpoint reached validation CE3.324907 in this single-seed comparison; this run's result does not demonstrate a donation benefit. OLMo's rotary positions, Q/K normalization and gated MLP architecture are not transplanted as complete computations into GPT-2.

Load

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "burnedinthesky/gpt2-olmo2-1b-warm120-ot-muon-lr4e-3-tok401k-seq1024-cos-wu5-485335"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
inputs = tokenizer("The next step in this experiment is", return_tensors="pt")
outputs = model.generate(
    **inputs, max_new_tokens=60,
    pad_token_id=tokenizer.eos_token_id,
    suppress_tokens=list(range(len(tokenizer), model.config.vocab_size)),
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

The example suppresses the 47 unused padding IDs during generation.

License and provenance

The donor model is released under Apache2.0. This derivative includes LICENSE and a provenance NOTICE. GPT-2 target attention weights were modified through OT transport/blending and then all target parameters were trained on C4. training_metadata.json, metrics.json, warm_stage_metrics.json and initialization_report.json record the initialization, compute and results. The final checkpoint was reloaded with standard Transformers and checked for finite inference, tied embeddings, and agreement with all12 logged norm metrics.