NSC64 trained a 109M PSSA in PyTorch

NSC64 ported the architecture to PyTorch and trained the result from scratch: 4B tokens in two phases, 50/50 TinyStories and FineWeb-Edu, then an 85/15 web/story fine-tune. The port keeps the same scan and memory-bank design as the Rust reference; the two departures it needed at this scale are backprop truncated at 256 tokens and the scan running as a fused Triton kernel with a hand-written backward that matches a plain PyTorch loop to about 1e-7.

Held-out scores, 1M unseen tokens per corpus, batch 8, sequence length 512: TinyStories 1.312 cross-entropy (3.71 perplexity, 64.8% next-token accuracy) and FineWeb-Edu web 2.866 (17.6 perplexity, 44.2% next-token accuracy). Web loss was still falling at 4B, so more tokens would help.

The weights are public: bf16 model.safetensors, a 16k byte-level BPE tokenizer and a self-contained loader that runs a plain scan on CPU and the Triton kernel on CUDA. See the model card for samples and the full write-up.