FishBoard/nano-megatron-tinystories-311m

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

model/                         Native single-device BF16 inference artifact
checkpoint/step_00010250/      Complete TP2×PP2×DP2 recoverable checkpoint
training/config.yaml           Formal training configuration
training/metrics.jsonl         Complete platform-independent metric history
training/train.log             Complete console log
training/best_validation.json  Best-checkpoint pointer
training/wandb/                W&B config, summary, requirements, and raw run
SHA256SUMS                     Checksums for every published file

Download and integrity check

hf download FishBoard/nano-megatron-tinystories-311m \
  --local-dir nano-megatron-tinystories-311m

cd nano-megatron-tinystories-311m
sha256sum -c SHA256SUMS

Inference

Use the exact Nano-Megatron Git commit associated with this release:

git clone https://github.com/pyy233/Nano-Megatron.git
cd Nano-Megatron
git checkout 757756cc7f163c85da8082c34a01bd6a0d1849c2
uv sync

CUDA_VISIBLE_DEVICES=0 uv run nano-megatron-generate \
  --model /path/to/nano-megatron-tinystories-311m/model \
  --prompt "Once upon a time, there was a little girl named Lily." \
  --device cuda \
  --dtype bfloat16 \
  --max-new-tokens 128 \
  --temperature 0.8 \
  --top-p 0.9 \
  --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:

wandb login
wandb sync \
  --entity YOUR_ENTITY \
  --project nano-megatron-tinystories-reproduction \
  --include-online \
  --include-synced \
  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.