---
license: mit
language:
- en
library_name: transformers
tags:
- text-generation
- causal-lm
- gpt2
- tinystories
- llm-kittens
- thunderkittens
- bf16
- cuda
- sd-prelu
- custom-activation
- custom_code
datasets:
- roneneldan/TinyStories
---

# llm.kittens TinyStories 124M BF16 — SD-PReLU

This is a 124M-parameter GPT-2-style causal language model trained from scratch on TinyStories with the `llm.kittens` C++/CUDA trainer, which is a fork of Karpathy's llm.c with some optimisations for SM120, and multi-stack kernel optimisations.

Unlike the GELU baseline ([`tinystories-5090`](https://huggingface.co/adamroberts/tinystories-5090)), this checkpoint replaces the MLP's GELU nonlinearity with a learnable **SD-PReLU** activation (the llm.kittens `-af sd-prelu` activation). Because SD-PReLU is not part of stock Transformers or llama.cpp, the repo ships a small custom module (`modeling_gpt_sdprelu.py`) and must be loaded with `trust_remote_code=True`. See [Activation: SD-PReLU](#activation-sd-prelu) below.

The model is published as a Hugging Face Transformers checkpoint with BF16 `safetensors` weights plus custom modeling code.
It was trained on a single RTX 5090 in roughly 14 hours (2595.7 ms average iteration over 20,000 steps).

## Result

- Model weights: `model.safetensors`
- Training step: `20000 / 20000`
- Final train loss: `0.781594`
- Final validation loss: `0.870315`
- Final throughput: `198871 tokens/s`
- Final step time: `2642.72 ms`
- Final reported BF16 MFU: `38.0%`
- Average iteration time: `2595.669013 ms`
- Safetensors size: `248,896,984` bytes
- Parameter count: `124,475,904` base + `24` learnable SD-PReLU scalars (`theta_a`/`theta_b`, 2 per layer × 12 layers)

For reference, the GELU baseline reached `0.785740` train / `0.875080` validation loss on the same setup; SD-PReLU lands marginally lower on both (`0.781594` / `0.870315`) at near-identical throughput, with the activation adding only 24 scalar parameters.

The TinyStories paper reports eval losses of `1.33` to `1.58` for the 768-hidden-size 1- and 2-layer attention-head ablations in Figure 24. This run's `0.870315` validation loss is lower, but the comparison is not apples-to-apples: this model is a 12-layer GPT-2-style model using GPT-2 tokenization, a 1024-token context, and a different implementation/training setup.

## Activation: SD-PReLU

This is the key difference from the GELU baseline. The MLP's GELU nonlinearity (applied between the `c_fc` up-projection and the `c_proj` down-projection) is replaced by **SD-PReLU**, a self-gated, damped PReLU. Each transformer block learns two scalars (`theta_a`, `theta_b`) that are mapped through a bounded reparameterization into `a` and `b`:

```text
a = alpha_max * sigmoid(theta_a)      # in [0, alpha_max),  alpha_max = 0.30
b = beta_min  + softplus(theta_b)     # in (beta_min, inf), beta_min  = 0.50
phi(x) = x * (a + (1 - a) * sigmoid(b * x))
```

- `a` is a learnable leak/floor (PReLU-like): the gate output never drops below `a`, so negative inputs are not fully zeroed.
- `b` controls the sharpness of the sigmoid gate. With `a = 0`, `phi` reduces to a Swish/SiLU-style `x * sigmoid(b * x)`.
- For numerical parity with the CUDA kernel, the activation is computed in float32 and the reparameterization/gate arguments are clamped to `[-20, 20]`; inputs and outputs stay in the surrounding model dtype (BF16).

The activation is configured via two config fields, `sdprelu_alpha_max` (`0.30`) and `sdprelu_beta_min` (`0.50`). Only 24 extra scalar parameters are introduced over the GELU baseline, so parameter count and on-disk size are essentially unchanged.

Note: the config still carries `"activation_function": "gelu_new"` for GPT-2 compatibility, but it is inert — the custom MLP overrides the activation with SD-PReLU regardless of that field.

## Custom activation inference code

The repo ships `modeling_gpt_sdprelu.py`, which defines:

- `GPTSDPReLUConfig` (`model_type = "gpt-sdprelu"`), a `GPT2Config` with the two extra `sdprelu_*` fields.
- `GPTSDPReLUMLP`, which reuses GPT-2's `c_fc` / `c_proj` / dropout submodules (so weight names stay `mlp.c_fc.*` / `mlp.c_proj.*`) and swaps GELU for SD-PReLU, adding `theta_a` / `theta_b` per layer.
- `GPTSDPReLULMHeadModel`, a `GPT2LMHeadModel` that installs the SD-PReLU MLP into every block.

These are wired into `config.json` via `auto_map`, so `trust_remote_code=True` is required to load the model. Everything else (attention, layernorms, embeddings, tied head, tokenizer) is standard GPT-2.

## Architecture

- Family: GPT-2-style decoder-only Transformer
- Descriptor: `d12`
- Layers: `12`
- Attention heads: `12`
- Hidden size: `768`
- Context length: `1024`
- Vocabulary size: `50,257`
- MLP activation: **SD-PReLU** (learnable, per-layer) — replaces GELU
- `model_type`: `gpt-sdprelu` (custom code, `trust_remote_code=True`)
- Precision: BF16 weights

## Training

The run used the TinyStories GPT-2 dataset files generated by `dev/data/tinystories.py` in `llm.kittens`. The only change from the GELU baseline is the `-af sd-prelu` activation flag.

```bash
./train_gpt2cu \
    -i "dev/data/tinystories/TinyStories_train.bin" \
    -j "dev/data/tinystories/TinyStories_val.bin" \
    -o "log124M/5090_S" \
    -v 250 -s 20000 -g 144 \
    -h 0 \
    -b 64 -t 1024 -d 524288 \
    -r 0 \
    -z 1 \
    -c 0.1 \
    -l 0.0006 -q 0.0 -u 700 -n 5000 \
    -y 0 \
    -e "d12" \
    -af sd-prelu \
    -x 20000
```

Key settings:

- Hardware target: RTX 5090 / SM120
- MLP activation: `sd-prelu` (`-af sd-prelu`)
- Micro batch: `64`
- Sequence length: `1024`
- Total desired batch size: `524,288` tokens
- Max steps: `20,000`
- Optimizer: AdamW as implemented in `llm.kittens`
- Peak learning rate: `6e-4`
- Scheduler: cosine
- Warmup: `700` steps
- Final LR fraction: `0.0`
- Weight decay: `0.1`
- Recompute: off
- ZeRO stage: `1`
- Checkpoint interval: `5000` steps

## Sample

Prompt/sample emitted at the final checkpoint (step 20000):

```text
Once upon a time, there was a little girl named Lily. She loved to play in the park with her friends. One day, they saw a big, dark cloud in the sky. Lily's friend, Timmy, said, "I think it's going to rain soon."
Suddenly, they heard a loud noise. It was a big, scary dog! Lily felt very scared and her skin started to shake. But then, a brave man came and scared the dog away. "Thank you," said Lily. "You're welcome," said the man.
After the storm passed, Lily and her friends went to play on the swings. They saw a beautiful rainbow in the sky. "Look at the pretty colors!" said Lily. "It's so bright and colorful!" Her friends agreed and they all felt happy.
```

## Files

- `model.safetensors`: BF16 Transformers weights (including the per-layer `theta_a` / `theta_b` SD-PReLU scalars).
- `modeling_gpt_sdprelu.py`: custom SD-PReLU model/config code (required, loaded via `trust_remote_code=True`).
- `config.json`: model configuration, including `sdprelu_alpha_max` / `sdprelu_beta_min` and the `auto_map` wiring.
- `generation_config.json`: default generation settings.
- `tokenizer.json`: GPT-2 tokenizer.
- `vocab.json` and `merges.txt`: GPT-2 BPE vocabulary files.

## Loading

Because the SD-PReLU activation lives in custom code, you must pass `trust_remote_code=True`:

```python

```

## GGUF / llama.cpp

GGUF export is **not available** for this checkpoint. The SD-PReLU activation is a custom, learnable nonlinearity with no equivalent in llama.cpp's GPT-2 graph, so the model cannot be quantized to GGUF or run with `llama.cpp` / LM Studio without implementing the activation there. Use the Transformers loading path above instead. The GELU baseline ([`tinystories-5090`](https://huggingface.co/adamroberts/tinystories-5090)) is available if you need a llama.cpp-compatible variant.

Recommended sampling settings (Transformers `generate`):

- Temperature: `0.8`
- Top-p: `0.95`
- Top-k: `50`
- Repetition penalty: `1.05`
- Stop/EOS token: `<|endoftext|>` / token id `50256`

This is a completion model, not a chat/instruction model: prompt it with the start of a story and always set a finite `max_new_tokens`, since it was trained for continuation and may not emit `<|endoftext|>` during normal generation.

Source implementation: `https://github.com/adamdroberts/llm.kittens` (SD-PReLU is in an unreleased branch)

TinyStories reference paper: `https://arxiv.org/abs/2305.07759`