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
- 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.
- 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.jsonrecords the exact masks, sampled rows, maps and diagnostics. - 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.