Optimizer Model Consistency: pretrained AdamW and Muon checkpoints
The two original pretrained bases used to initialize SFT in the Optimizer Model
Consistency experiments. Both checkpoints were saved at step 31,000. Their
source directory names contain iters-32000, which is the configured training
budget, not the actual saved checkpoint step.
| Checkpoint | Optimizer name in experiment code | Bytes | Saved validation loss |
|---|---|---|---|
adamw/ckpt.pt |
adamw |
1,492,568,524 | 3.0113711357 |
muon/ckpt.pt |
matrix-p-sgd |
1,152,806,764 | 2.9760797024 |
These are byte-for-byte original PyTorch checkpoints, including model,
optimizer, model_args, iter_num, best_val_loss, and config.
Muon optimizer state is a list; AdamW optimizer state is a dictionary.
The validation losses above are values stored during pretraining, not a new evaluation.
Architecture and format
Custom nanoGPT / GPT-2-sized model pretrained on OpenWebText: 12 layers,
12 attention heads, embedding dimension 768, context length 1024, SwiGLU,
no linear/layer-norm biases, dropout 0, learned absolute position embeddings,
and padded vocabulary size 50304. Tokenization uses tiktoken.get_encoding("gpt2").
This is a custom nanoGPT checkpoint format. Use the supplied model code and loader;
AutoModelForCausalLM.from_pretrained() does not directly support these files.
The matching model implementation and its optimizer import dependencies are in
code/. The upstream code license is preserved in code/LICENSE.
Download
Log in with a Hugging Face account/token that can read this private repository:
python -m pip install -U huggingface_hub
hf auth login
hf download yuxing6/optimizer-model-consistency-pretrained --local-dir ./omc-pretrained
cd omc-pretrained
python -m pip install -r requirements.txt
sha256sum -c SHA256SUMS
To download only one checkpoint, pass its filename, e.g. hf download yuxing6/optimizer-model-consistency-pretrained adamw/ckpt.pt.
Load and use
python load_checkpoint.py --checkpoint adamw/ckpt.pt --smoke-test
python load_checkpoint.py --checkpoint muon/ckpt.pt --device cuda --prompt "The purpose of this experiment is"
For Python:
from load_checkpoint import load_checkpoint
model, checkpoint = load_checkpoint("adamw/ckpt.pt", device="cpu")
print(checkpoint["iter_num"]) # 31000
For SFT in the original experimental code, set init_from = 'pretrained_ckpt'
and pretrained_ckpt_path to the downloaded adamw/ckpt.pt or muon/ckpt.pt.
Keep model_args from the checkpoint. Restoring optimizer state additionally
requires matching the original optimizer implementation and parameter-group order;
simply loading the model does not restore the optimizer or training schedule.
Private sharing
A repository URL alone does not grant access. For account-based collaboration,
place/transfer the private repository into a dedicated Hugging Face organization
and invite your collaborator with the read role. They can then log in with their
own account and download using their own read token. The standard organization
read role can read that organization's repositories, so use a dedicated
organization if you want to isolate this project.
See the official organization access control and access token documentation. No collaborators are automatically invited by uploading this repository.