Downloads · 30 days
6
15% of all-time downloads
yatnmn-softplus-d22-chinchilla-1B
mlnomad/yatnmn-softplus-d22-chinchilla-1B
yatnmn-softplus-d22-chinchilla-1B is a machine learning model from mlnomad. Use it for the machine learning task on the model card, and read the license before you ship it in a product. It is set up for flax. The card lists the license as apache-2.0.
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.
All-time downloads
40
Public
Repo size
14.7 GB
Likes
0
Public
Repo makeup
Click a slice to open those files.
Other14.7 GB · 100%
At a glance
- Library
- flax
- License
- apache-2.0
- Access
- Public
- Created
- Apr 17, 2026
- Updated
- Apr 18, 2026
- SHA
- d90bebc0
Datasets
Model card
From the Hugging Face model README
YatNMN-Softplus d=22 Chinchilla (1.08B) — JAX/Flax
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 repository contains the full Orbax checkpoint (model + AdamW optimizer state) — resumable for continued pretraining or fine-tuning.
Final result
| Metric | Value |
|---|---|
| Final loss | 2.8479 |
| Smooth final loss | 2.8325 |
| Tokens | 21.5 B |
| Steps | 328,718 micro (41,089 optimizer applies) |
| Wall time | 47.85 h |
| Throughput | 125 K tok/s |
Scaling comparison
| Scale | Params | Tokens | C4 smooth loss |
|---|---|---|---|
| d=12 | 261M | 5.22B | 2.98 |
| d=22 | 1.08B | 21.5B | 2.83 |
Architecture
Full nanochat stack with YatNMN-Softplus MLP:
- YatNMN MLP:
y = α · (x·W + softplus(b))² / (||x−W||² + softplus(ε))— per-neuron(ff,)bias, scalar learnable α and ε - d=22, n_embd=1408, n_head=22, n_kv_head=22, seq_len=1024
- RoPE (base 100K), GQA, QK-norm 1.2×, SSSL sliding window
- Tied embeddings, parameterless RMSNorm, value embeddings (alternating layers)
- Per-layer residual scalars, smear, backout, logit soft-cap 15
Training
| Optimizer | plain AdamW, peak LR 0.03, warmup-cosine (500 warmup, 41K decay) |
| FSDP | 8-way param sharding on v6e-8 |
| Remat | gradient checkpointing (dots_saveable policy) |
| Grad accum | 8 micro-steps per apply (effective batch 512 = 524K tok/apply) |
| Hardware | TPU v6e-8 (TRC), europe-west4-a |
Contents
├── 328718/ # final Orbax checkpoint (~14 GB)
│ ├── model/ # nnx.Param state
│ ├── optimizer/ # optax AdamW + MultiSteps state (resumable)
│ └── metadata/ # JSON: final_loss, smooth
├── config.json
├── README.md
└── code/ # source snapshots
Loading
git clone https://github.com/mlnomadpy/flaxchat && cd flaxchat
pixi install && pip install 'nmn>=0.2.29'
huggingface-cli download mlnomad/yatnmn-softplus-d22-chinchilla-1B --local-dir ./hf_model
from flax import nnx
from nmn.nnx.layers import YatNMN
from flaxchat.gpt import GPT, GPTConfig, Block
from flaxchat.checkpoint import restore_model_from_checkpoint
# Patch Block for YatNMN-Softplus MLP (same as training)
# ... (see code/train_d12_chinchilla.py --mlp yatnmn-softplus --depth 22)
config = GPTConfig(
sequence_len=1024, vocab_size=32768,
n_layer=22, n_head=22, n_kv_head=22, n_embd=1408,
window_pattern="SSSL", tie_embeddings=True,
)
model = GPT(config, rngs=nnx.Rngs(0))
meta = restore_model_from_checkpoint(model, "./hf_model/328718")
PyTorch version
For PyTorch/HuggingFace Transformers inference (no JAX needed): mlnomad/yatnmn-softplus-d22-chinchilla-1B-pytorch
License
Apache 2.0.
- Library
- flax
- License
- apache-2.0
- Languages
- en
- Created
- Apr 17, 2026
- Updated
- Apr 18, 2026