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

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

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

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.