Update README for 350M PG-19 byte-level SEDD checkpoint

README.md

@@ -1,3 +1,70 @@
----
-license: mit
----
+---
+license: mit
+tags:
+  - discrete-diffusion
+  - text-generation
+  - byte-level
+  - sedd
+datasets:
+  - pg19
+---
+
+# GPT Diffusion (SEDD) — 350M, PG-19 Byte-Level
+
+A **350M parameter discrete diffusion transformer** trained on raw UTF-8 bytes from [PG-19](https://github.com/google-deepmind/pg19) (Project Gutenberg books), using the **Score Entropy Discrete Diffusion (SEDD)** framework with a nanoGPT-style architecture.
+
+## Model Details
+
+| Parameter | Value |
+|-----------|-------|
+| Parameters | ~355M |
+| Architecture | Transformer (DDiT) with adaptive LayerNorm |
+| Layers | 25 |
+| Heads | 16 |
+| Embedding dim | 1024 |
+| Conditioning dim | 256 |
+| Context length | 8192 bytes |
+| Vocab size | 256 (raw bytes) |
+| Activation | SwiGLU |
+| Position embedding | RoPE |
+| Normalization | RMSNorm |
+| Training steps | 5,000 |
+| Noise schedule | Geometric |
+
+## Checkpoint Contents
+
+The `.pth` file contains:
+
+- `model` — main model state dict
+- `ema_model` — exponential moving average weights
+- `optimizer` — AdamW optimizer state
+- `adamw_optimizer` — secondary optimizer state
+- `step` — training step (5000)
+- `config` — model config dict
+- `chunk_idx` — data chunk index for resuming
+
+## Usage
+
+```python
+import torch
+from model import GPT, GPTConfig
+
+ckpt = torch.load("pg19_step_5000.pth", map_location="cpu")
+config = GPTConfig(**ckpt["config"])
+model = GPT(config)
+model.load_state_dict(ckpt["ema_model"])  # use EMA weights for best quality
+```
+
+## Training
+
+Trained with the `train_pg19.py` script:
+
+```bash
+torchrun --nproc_per_node=2 train_pg19.py
+```
+
+Effective batch size: micro_batch=4 × grad_accum=4 × 8192 context = ~131K tokens/step.
+
+## License
+
+MIT