neox-ckpt-pythia-14m-seed0

This repository contains GPT-NeoX format checkpoints for a 14M parameter Pythia model. Checkpoints are stored as branches (e.g., step0, step1000, step143000), with 154 checkpoints total matching the standard Pythia checkpoint schedule.

Model Details

Parameter	Value
Parameters	14M
Layers	6
Hidden size	128
Attention heads	4
Sequence length	2048
Vocab size	50304
Training steps	143,000
Batch size	32 per GPU
Learning rate	1e-3 (cosine decay, 10% warmup)
Optimizer	Adam (betas 0.9, 0.95)
Precision	FP16
Init method	small_init (weights), wang_init (output layer)

Training Data

The training config specifies /weka/pile/pile_20B_tokenizer_text_document as the training data path, which is the standard (non-deduplicated) Pile.

Relationship to Other Pythia Models

This model was trained as part of an "extra seeds" project (wandb_project: pythia-extra-seeds). The config specifies seed: 1234, which is the same seed used for the original Pythia models.

Despite sharing the same seed value and init method, the step 0 (pre-training) weights in this repository do not match the step 0 weights of any published Pythia model. We investigated this by comparing against EleutherAI/pythia-14m-deduped, which was trained on the deduplicated Pile but uses the same architecture and seed.

Step 0 comparison with pythia-14m-deduped

At step 0, before any training occurs, the two models should reflect only their random initialization. We found:

• 50 of 76 weight tensors are bitwise identical. These are all biases (initialized to zero) and layer norm parameters (initialized to ones/zeros).
• 26 weight tensors differ. These are the randomly initialized weight matrices: embeddings, attention projections, and MLP projections.

To distinguish between "different init method" and "same init, different RNG state," we compared the standard deviations of corresponding weight tensors:

Layer type	HF std (typical)	NeoX std (typical)	Ratio
Embedding	0.0559	0.0559	~1.000
QKV projection	0.0558	0.0558	~1.000
Attention dense	0.0294	0.0295	~1.000
MLP up-proj	0.0560	0.0560	~1.000
MLP down-proj	0.0295	0.0294	~1.000

The standard deviations match to within ~1% across all layers, confirming the same initialization distribution was used. A different init method would produce a systematic scale difference. This means the same small_init + wang_init scheme was applied, but the RNG produced different values.

Step 143000 (final checkpoint) comparison

At the final checkpoint, all 76 weight tensors differ substantially (as expected, since training data also differs). Max differences range from 0.08 to 136, confirming these are fully distinct models.

Conclusion

We are not sure what happened. The training config specifies the same seed (1234) and the same init methods as the original Pythia models, and the weight statistics confirm the same init distribution was used, but the actual initialized values differ. Possible explanations include differences in the software version, tensor parallelism configuration, or other factors that affect RNG state. The training data also differs (standard Pile vs deduplicated Pile for the published pythia-14m-deduped).

Checkpoint Format

Each step branch contains:

• mp_rank_00_model_states.pt — Full NeoX checkpoint (model weights, optimizer state, LR scheduler, RNG states)
• configs/ — Training configuration YAML files

These are not HuggingFace Transformers format checkpoints. To use them for inference, convert to HF format first using the NeoX conversion scripts.

License

Apache 2.0