yatnmn-softplus-d22-chinchilla-1B-pytorch

YatNMN-Softplus d=22 Chinchilla (1.08B) — PyTorch / HuggingFace Transformers

A 1.08B-parameter nanochat-architecture GPT with YatNMN-Softplus MLP, trained on English C4 to Chinchilla-optimal token budget (20× params ≈ 21.5B tokens) on a single TPU v6e-8 using FSDP + gradient accumulation.

This is the 1B-scale version of the d=12 (261M) YatNMN-Softplus model that achieved 2.98 loss — testing whether the YatNMN advantage scales.

Results

The loss improved by 0.15 nats going from 261M → 1.08B at Chinchilla-optimal compute — consistent with standard scaling-law predictions.

Quick start

pip install torch transformers safetensors

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "mlnomad/yatnmn-softplus-d22-chinchilla-1B-pytorch",
    trust_remote_code=True,
    dtype=torch.float32,
).eval()

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

prompt = "The meaning of life is"
ids = tokenizer(prompt, return_tensors="pt").input_ids
with torch.no_grad():
    out = model.generate(
        ids, max_new_tokens=50,
        do_sample=True, temperature=0.8, top_p=0.9,
        use_cache=True, pad_token_id=tokenizer.eos_token_id or 0,
    )
print(tokenizer.decode(out[0], skip_special_tokens=True))

Model details

YatNMN-Softplus MLP

y = α · (x · W + softplus(b))² / (||x − W||² + softplus(ε))

Per-neuron (ff,) bias, scalar learnable ε, scalar learnable α. Same config as the 261M model — architecture identical except depth (22 vs 12) and width (1408 vs 768).

Training setup (1B on single v6e-8)

Training a 1B model on a single v6e-8 required:

• FSDP: model params sharded 8-way across chips (first-dim partitioning on the fsdp mesh axis)
• Gradient checkpointing (remat): dots_saveable policy on all blocks
• Gradient accumulation: 8 micro-steps per optimizer apply (effective batch 512 samples = 524K tokens)
• Batch per device: 8 (reduced from 32 at d=12 for HBM headroom)

Files

├── config.json                       # HF config with auto_map
├── model.safetensors                 # ~4.3 GB, fp32
├── yatnmn_gpt.py                     # pure PyTorch Yat_GPT + YatNMN layer
├── torch_gpt.py                      # shared building blocks
├── configuration_yatnmn_gpt.py       # PretrainedConfig subclass
├── modeling_yatnmn_gpt.py            # PreTrainedModel + KV cache + GenerationMixin
└── README.md

Related

• mlnomad/yatnmn-softplus-d12-chinchilla-261M-pytorch — same architecture at 261M (d=12), loss 2.98
• mlnomad/gelu-d12-chinchilla-261M-pytorch — GELU baseline at 261M, loss 3.12
• flaxchat — JAX/Flax training harness
• nmn — YatNMN layer

License

Apache 2.0.