Nano-Megatron TinyStories 311M
This repository is a reproducibility release for a 310,821,888-parameter GPT
trained from scratch on the full TinyStories corpus with Nano-Megatron. It
contains the native single-device inference artifact, the complete distributed
training checkpoint for the best validation step, the exact training metrics
and console log, and a replayable W&B run. The exact implementation is
published as a normal GitHub commit and is not duplicated in this repository.
This is a native Nano-Megatron model, not a Hugging Face Transformers model.
AutoModelForCausalLM.from_pretrained() is not supported.
Model
| Field | Value |
|---|
| Parameters | 310,821,888 |
| Layers | 24 |
| Hidden size | 1,024 |
| FFN hidden size | 2,736 |
| Attention heads / KV heads | 16 / 16 |
| Context length | 512 |
| Vocabulary | 8,192 byte-level BPE tokens |
| Architecture | RMSNorm, RoPE, SwiGLU, tied embeddings |
| Weight dtype | BF16 |
| Inference source | best validation checkpoint at step 10,250 |
| Tokenizer fingerprint | dac8b8de30b96fe3076c3af55ea501da06810ec175ede5dd47ca39bb4be373a1 |
The inference artifact is stored under model/. Its model.pt contains 219
BF16 tensors and has SHA256
b3aea969505eec52f26d9051c601524a77b625b1045a49a9c8977cd60d8bab7c.
No quantization or additional precision conversion was applied during export.
Training
| Field | Value |
|---|
| Hardware | 8 × NVIDIA L40S (46,068 MiB visible per GPU) |
| Parallel topology | TP2 × PP2 × DP2 |
| Pipeline schedule | 1F1B with P2P overlap |
| Sequence parallel | enabled |
| Micro batch / accumulation | 8 / 16 |
| Global batch | 256 sequences |
| Sequence length | 512 |
| Tokens per optimizer step | 131,072 |
| Optimizer | AdamW, LR 3e-4, betas (0.9, 0.95), weight decay 0.1 |
| Scheduler | 210-step warmup, cosine decay to 3e-5 |
| Training horizon | 10,488 steps, 3 corpus epochs |
| Consumed tokens | 1,374,683,136 |
| Metrics wall time | 3.190 hours |
| Steady tokens/s median / p10 | 126,367.6 / 125,375.6 |
| Peak CUDA allocated / reserved | 3.352 / 3.883 GiB per process |
The specific RunPod host used for this run had a broken default NCCL P2P
transport, so the run used NCCL_P2P_DISABLE=1. This is a host-specific
workaround and is not a general recommendation for L40S systems.
Results
| Metric | Value |
|---|
| First logged train loss (step 10) | 9.16173 |
| Final train loss (step 10,488) | 1.07856 |
| First validation loss (step 250) | 2.70606 |
| Best validation loss | 1.179513 |
| Best validation perplexity | 3.25279 |
| Best validation step | 10,250 |
The full run ended at step 10,488, while the exported inference model and the
published recoverable checkpoint come from the best validation step, 10,250.
training/metrics.jsonl and training/train.log continue through step 10,488.
Repository contents
1model/ Native single-device BF16 inference artifact
2checkpoint/step_00010250/ Complete TP2×PP2×DP2 recoverable checkpoint
3training/config.yaml Formal training configuration
4training/metrics.jsonl Complete platform-independent metric history
5training/train.log Complete console log
6training/best_validation.json Best-checkpoint pointer
7training/wandb/ W&B config, summary, requirements, and raw run
8SHA256SUMS Checksums for every published file
Download and integrity check
1hf download FishBoard/nano-megatron-tinystories-311m \
2 --local-dir nano-megatron-tinystories-311m
3
4cd nano-megatron-tinystories-311m
5sha256sum -c SHA256SUMS
Inference
Use the exact Nano-Megatron Git commit associated with this release:
1git clone https://github.com/pyy233/Nano-Megatron.git
2cd Nano-Megatron
3git checkout 757756cc7f163c85da8082c34a01bd6a0d1849c2
4uv sync
5
6CUDA_VISIBLE_DEVICES=0 uv run nano-megatron-generate \
7 --model /path/to/nano-megatron-tinystories-311m/model \
8 --prompt "Once upon a time, there was a little girl named Lily." \
9 --device cuda \
10 --dtype bfloat16 \
11 --max-new-tokens 128 \
12 --temperature 0.8 \
13 --top-p 0.9 \
14 --seed 20260718
The artifact was validated on an NVIDIA L40S and independently downloaded and
validated on an RTX 4060 Laptop GPU. Five fixed prompts, UTF-8 input,
same-device greedy determinism, and EOS termination were tested.
Training checkpoint
checkpoint/step_00010250/ contains model and optimizer shards, trainer and
scheduler state, RNG state, data progress, rank runtime state, manifest
metadata, and a .complete marker. It is intended for Nano-Megatron's native
checkpoint loader.
Exact continuation also requires the matching TinyStories mmap artifacts and
tokenizer fingerprint, plus the original TP2×PP2×DP2 world-size layout. The
dataset itself is not redistributed in this repository.
W&B replay
The original project could not be made public, so the raw local run is included
for replay into another W&B account:
1wandb login
2wandb sync \
3 --entity YOUR_ENTITY \
4 --project nano-megatron-tinystories-reproduction \
5 --include-online \
6 --include-synced \
7 training/wandb/run-kk7ntk3g.wandb
metrics.jsonl remains the authoritative, vendor-independent metric history.
Security note
model.pt and the distributed checkpoint use PyTorch serialization. Load them
only from a trusted copy of this repository and verify SHA256SUMS first.
The release does not currently provide Safetensors weights.
Source and data provenance
The exact source implementation is GitHub commit
757756cc7f163c85da8082c34a01bd6a0d1849c2
from
https://github.com/pyy233/Nano-Megatron. Source code is maintained in
the project repository rather than copied into this model artifact repository.
Training data came from
roneneldan/TinyStories,
whose dataset card declares the CDLA Sharing 1.0 license. The Nano-Megatron
source repository did not contain a standalone license file at release time,
so this artifact repository does not assert an additional source or model
license. Review the upstream terms before redistribution or commercial use.
Limitations
- The model is trained for short, simple English stories and is not a general
assistant or factual language model.
- Context length is limited to 512 tokens.
- The native generator currently recomputes the full prefix and has no KV
cache.
- The checkpoint format is Nano-Megatron-specific and is not directly loadable
by Transformers, vLLM, or Megatron-LM.
