docs: update model card for 200M CoT training run

README.md

@@ -1,136 +1,222 @@
 ---
-license: apache-2.0
+language:
+  - en
 tags:
   - hyperbolic
   - lorentz
   - geometric-deep-learning
   - language-model
-  - pretraining
-datasets:
-  - wikimedia/wikipedia
-language:
-  - en
-base_model:
-  - Graph-and-Geometric-Learning/helm
+  - chain-of-thought
+  - reasoning
 pipeline_tag: text-generation
+license: mit
+datasets:
+  - open-thoughts/OpenThoughts-114k
+  - HuggingFaceTB/smollm-corpus
 ---
 
-# HELM-D 130M: Hyperbolic Efficient Language Model
+# HELM-D: Hyperbolic Chain-of-Thought Reasoning Engine
+
+> Fork of [Graph-and-Geometric-Learning/helm](https://github.com/Graph-and-Geometric-Learning/helm) — a **200M parameter** fully hyperbolic transformer trained on NVIDIA H200 for structured reasoning.
+>
+> **Checkpoints**: [datasysdev/helm-d-130m-hyperbolic](https://huggingface.co/datasysdev/helm-d-130m-hyperbolic) on HuggingFace
+
+All computations live on the [Lorentz manifold](https://en.wikipedia.org/wiki/Hyperboloid_model): $-x_0^2 + x_1^2 + \dots + x_d^2 = -1$. The model uses hyperbolic embeddings, Lorentzian attention, and Riemannian optimization — making it natively suited for hierarchical data like code ASTs, dependency trees, and chain-of-thought reasoning traces.
 
-A 130M parameter language model that operates entirely on the **Lorentz manifold** (hyperboloid model of hyperbolic space). All embeddings, attention, and optimization live in hyperbolic space — the model is geometrically native, not a Euclidean model with hyperbolic post-hoc modifications.
+---
 
-Pretrained on NVIDIA H200 at **193K tokens/sec** using Flash Attention 2, selective BF16, and torch.compile optimizations.
+## Current Training Run
 
-## Architecture
+Training a **200M parameter** HELM-D from scratch on a multi-domain reasoning corpus:
 
 | Parameter | Value |
 |---|---|
-| Architecture | L6W384A6 (6 layers, width 384, 6 heads) |
-| Parameters | 130M |
-| Manifold | Lorentz (hyperboloid, curvature K=1) |
-| Tokenizer | Qwen3-30B-A3B (151,669 vocab) |
-| Context length | 2048 |
-| Attention | Flash Attention 2 (spatial-only with time reconstruction) |
-| Optimizer | RiemannianAdam (geoopt) |
+| Architecture | `L16W768A12` (16 layers, 768 width, 12 heads) |
+| Parameters | **200M** (175.8M Euclidean + 24.6M Hyperbolic) |
+| Tokenizer | TinyLlama 32K (dense coverage, no dead tokens) |
+| Context | 4096 tokens (full CoT traces fit in one pass) |
+| Throughput | **130K tok/s** on single H200 |
+| Optimizer | Dual-group RiemannianAdam (see below) |
+| Learning Rate | 3e-4, cosine decay with 500-step warmup |
+| Gradient Clip | 0.5 |
+| Manifold | Lorentz $-x_0^2 + \|x\|^2 = -1$, verified at 1.0000±0.0000 |
+
+### Training Data (60/20/20 Mix)
+
+| Domain | Weight | Source | Purpose |
+|---|---|---|---|
+| CoT Reasoning | 60% | [OpenThoughts-114k](https://huggingface.co/datasets/open-thoughts/OpenThoughts-114k) | Math, code, science reasoning with `<think>` traces |
+| Python Code | 20% | [SmolLM-Corpus python-edu](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus) | Educational Python |
+| Text | 20% | [SmolLM-Corpus cosmopedia-v2](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus) | General knowledge |
 
-## Training
+Streamed via `interleave_datasets` with a **512-chunk shuffle buffer** to prevent domain clustering (see Architecture Decisions below).
 
-Pretrained on 100K English Wikipedia articles + 100K Python source files (~221M unique tokens, ~4 epochs). This is a **proof-of-concept checkpoint** — it validates the hyperbolic training pipeline but does not produce coherent text generation due to the small dataset size.
+---
 
-### Performance (H200)
+## Key Changes from Upstream HELM
 
-| Configuration | ms/step | tok/s | Speedup |
-|---|---|---|---|
-| Original FP32 | 5,966 | 43,917 | 1.0× |
-| + BF16 logits | 3,601 | 72,770 | 1.7× |
-| + FA2 (width=384) | 1,875 | 140,025 | 3.2× |
-| **+ torch.compile + python -O** | **1,357** | **193,000** | **4.4×** |
+### 1. Tokenizer: Llama-3.1 → TinyLlama 32K
 
-### Training Curve
+The original HELM uses the Llama-3.1 tokenizer (128K vocab). We switched to **TinyLlama's 32K tokenizer** for the CoT training run:
 
-Loss stabilized around 6.5-7.0 after exhausting the 221M-token dataset (4+ epochs).
+- **Dense coverage**: No dead tokens — every token gets trained
+- **Smaller embedding matrix**: 32K × 768 vs 128K × 768 — significant VRAM savings
+- **Better for small models**: 200M params can't support 128K vocab efficiently
 
-## Checkpoints
+### 2. Architecture: L6W384A6 → L16W768A12
 
-| File | Step | Description |
+Scaled up from the original 31M parameter toy model to a **200M parameter** engine:
+
+| | Original | Ours |
 |---|---|---|
-| `h200_step2400.pt` | 2400 | End of first torch.compile run (stable, loss ~7.0) |
-| `h200_step4100.pt` | 4100 | Final checkpoint with all optimizations (-O flag, geoopt patch) |
+| Layers | 6 | **16** |
+| Width | 390 | **768** |
+| Heads | 6 | **12** |
+| Head dim | 65 | **64** (Tensor Core aligned) |
+| Parameters | 31M | **200M** |
+
+### 3. Dual-Group Optimizer (Matching Original Authors)
 
-Each checkpoint contains:
-- `model_state_dict`: Full model weights (FP32, Lorentz manifold)
-- `optimizer_state_dict`: RiemannianAdam state
-- `global_step`: Training step counter
+The original HELM repo uses **two separate optimizers**: AdamW for Euclidean params and RiemannianAdam for hyperbolic params, with `weight_decay=0.0` on manifold parameters.
 
-### Loading
+We implement this as a single RiemannianAdam with dual parameter groups:
 
 ```python
-import torch
-from helm.hypercore.manifolds import Lorentz
-from helm.modules.helm_d import LTransformerDecoder
-
-model = LTransformerDecoder(
-    manifold_in=Lorentz(1.0),
-    manifold_hidden=Lorentz(1.0),
-    manifold_out=Lorentz(1.0),
-    arch="L6W384A6",
-    vocab_size=151669,
-    context_length=2048,
-)
-ckpt = torch.load("h200_step4100.pt", map_location="cpu", weights_only=False)
-model.load_state_dict(ckpt["model_state_dict"], strict=False)
+optimizer = RiemannianAdam([
+    {"params": euclidean_params, "weight_decay": 0.01},   # 175.8M params
+    {"params": hyperbolic_params, "weight_decay": 0.0},   # 24.6M params
+], lr=3e-4)
 ```
 
-## Tokenizer Surgery
+**Why**: Standard L2 weight decay pulls parameters toward the Euclidean origin `[0,0,...,0]`, which is **not on the Lorentz manifold**. Applying decay to manifold parameters causes the optimizer to constantly drag embeddings off the $-1$ surface, then the `expmap` projection violently snaps them back — destabilizing training.
 
-The original HELM uses Llama-3.1 tokenizer (128K vocab). We transferred embeddings to the Qwen3-30B-A3B tokenizer (151K vocab) using **Lorentzian Fréchet Mean** — computing the geometric centroid on the hyperboloid for novel tokens by decomposing them into Llama sub-tokens and projecting via the Einstein midpoint.
+### 4. Shuffle Buffer Dataloader
 
-## Key Optimizations
+The streaming `interleave_datasets` interleaves at the **document** level. Since OpenThoughts reasoning traces can be 4,000-16,000 tokens (1-4 consecutive 4096-token chunks), the model receives bursts of pure math followed by bursts of pure code — causing catastrophic loss spikes.
 
-- **Flash Attention 2**: Runs on spatial dimensions only (strips Lorentz time coordinate), reconstructs via manifold projection after attention.
-- **Selective BF16**: Only the output projection (Euclidean) uses BF16. All Lorentz operations remain FP32.
-- **python -O**: Strips 30+ `assert torch.isnan()` checks from the manifold code, eliminating GPU→CPU synchronization stalls.
-- **geoopt patch**: `torch.norm(p=2)` → `torch.linalg.vector_norm(ord=2)` for torch.compile compatibility.
-- **Width 384**: Aligned to 64-wide Tensor Core tiles (original was 390).
+**Fix**: A 512-chunk shuffle buffer accumulates tokenized chunks before yielding, ensuring every batch is a representative mix of all 3 domains:
 
-## Intended Use
+```
+Documents → Tokenize → Pack into 4096-token chunks → Buffer (512) → Shuffle → Yield to GPU
+```
 
-This checkpoint serves as a **seed for Network Morphism** — upscaling to 1B+ parameters by zero-padding Lorentz spatial dimensions and cloning transformer layers. The learned manifold geometry, token distributions, and attention patterns transfer to the larger model.
+This eliminated gradient spikes of 46+ and stabilized the loss descent.
 
-## Geometric Compromises
+### 5. TF32 Tensor Core Acceleration
 
-- FA2 computes Euclidean dot products instead of Minkowski inner products (drops the -q₀k₀ term)
-- Periodic re-projection of embeddings onto the manifold every 100 steps
-- Einstein midpoint used instead of iterative Karcher mean for tokenizer surgery
+```python
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision("high")
+```
 
-## Citation
+Throughput: **40K → 130K tok/s** (3.25× speedup). All upstream Lorentz operations remain in FP32 — only matmul operations use TF32's 10-bit mantissa through the Tensor Cores.
 
-Based on:
-```bibtex
-@article{helm2024,
-  title={Hyperbolic Efficient Language Models},
-  author={Graph and Geometric Learning Lab},
-  year={2024},
-  url={https://github.com/Graph-and-Geometric-Learning/helm}
-}
+### 6. LR Override on Checkpoint Resume
+
+PyTorch's `optimizer.load_state_dict()` restores the learning rate from the checkpoint, silently overriding CLI arguments. We force the LR after restore:
+
+```python
+for pg in optimizer.param_groups:
+    pg["lr"] = args.lr
+    pg["initial_lr"] = args.lr
+```
+
+---
+
+## Quick Start
+
+### Requirements
+
+```bash
+pip install torch flash-attn --no-build-isolation
+pip install geoopt transformers datasets
 ```
 
-## Source Code
+### Training on H200
 
-[unixsysdev/helm (h200-optimizations branch)](https://github.com/unixsysdev/helm/tree/h200-optimizations)
+```bash
+export PYTHONPATH=/path/to/helm-src:$PYTHONPATH
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
-## Roadmap: 1.37B Pretraining
+# Fresh training
+python3 -O train_cot.py \
+    --batch_size 16 --grad_accum 8 \
+    --lr 3e-4 --seq_len 4096 \
+    --save_dir /tmp/checkpoints/cot \
+    --log_every 1
 
-The 130M checkpoints in this repo are seeds for the **1.37B HELM-D** model (L24W1536A24), upscaled via Network Morphism:
+# Resume from checkpoint
+python3 -O train_cot.py \
+    --batch_size 16 --grad_accum 8 \
+    --lr 3e-4 --save_dir /tmp/checkpoints/cot \
+    --log_every 1 --resume
+```
+
+### Generation Test
+
+```bash
+python3 test_gen.py --checkpoint /tmp/checkpoints/cot/cot_step5000.pt
+```
+
+---
+
+## Architecture Decisions
+
+### Gradient Clipping: 1.0 → 0.5
 
-1. **Width 384→1536**: Zero-pad Lorentz spatial dims (manifold constraint preserved exactly)
-2. **Depth 6→24 layers**: Interleaved cloning — repeats the full 6-layer pipeline 4× with residual scaling
-3. **All linear weights**: Top-left corner placement in expanded matrices, remainder N(0, 0.001)
+The original authors use `grad_clip=1.0` on a 6-layer model. At 16 layers, gradient variance compounds across 10 additional layers. Clip of 0.5 on 16 layers is physically equivalent to 1.0 on 6 layers.
+
+### LR Scaling: 4e-4 → 3e-4
+
+The original authors use `lr=4e-4` on a 31M model. As parameter count and depth scale, optimal learning rates must decrease. 3e-4 is the correct scaling for 200M parameters.
+
+### Flash Attention 2
+
+FA2 computes Euclidean dot products, but hyperbolic attention requires the Minkowski inner product $\langle x, y \rangle_{\mathcal{L}} = -x_0 y_0 + \sum x_i y_i$. We run FA2 on **spatial dimensions only** (strip the time coordinate), then reconstruct via manifold projection: $x_0 = \sqrt{\|x_{1:d}\|^2 + 1}$.
+
+### Periodic Re-projection
+
+Embeddings are snapped back to $-x_0^2 + \|x\|^2 = -1$ every 100 steps to correct constraint drift from mixed-precision gradient updates.
+
+---
 
-The 1.37B model is currently training on **2B tokens from FineWeb-Edu** on a single NVIDIA H200.
+## Files
+
+| File | Description |
+|---|---|
+| `train_cot.py` | **Main training script** — 200M HELM-D with streaming 60/20/20 mix, shuffle buffer, dual optimizer |
+| `test_gen.py` | Temperature sweep generation test with repetition penalty grid |
+| `train_h200.py` | H200 pretraining with FA2, BF16, torch.compile (130M seed model) |
+| `train_h200_130m.py` | 130M config (L6W384A6) for seed training |
+| `tokenizer_surgery.py` | Llama→Qwen3 embedding transfer via Lorentzian Fréchet Mean |
+| `upscale_130m_to_1b.py` | Network Morphism: 130M→1.37B (Lorentz zero-pad + layer cloning) |
+| `setup_h200.sh` | H200 environment setup (CUDA, PyTorch, Flash Attention) |
+| `helm/modules/helm_d.py` | HELM-D decoder with RoPE odd-dim fix, BF16 output projection |
+| `helm/hypercore/` | Lorentz manifold operations, Riemannian optimizers |
+
+---
+
+## Known Issues
+
+- **torch.compile modes**: `max-autotune` and `reduce-overhead` crash with CUDAGraphs in LorentzEmbeddings. Only default mode works.
+- **geoopt + torch.compile**: Requires patching `torch.norm` → `torch.linalg.vector_norm` in geoopt's `lorentz/math.py`.
+- **Tokenizer max length warnings**: TinyLlama tokenizer reports max_length=2048 but we use 4096 seq_len — this is harmless (we handle truncation ourselves).
+
+---
+
+## Citation
+
+Based on:
+```bibtex
+@article{he2025helm,
+  title={HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Experts},
+  author={He, Neil and Anand, Rishabh and Madhu, Hiren and Maatouk, Ali and Krishnaswamy, Smita and Tassiulas, Leandros and Yang, Menglin and Ying, Rex},
+  journal={arXiv preprint arXiv:2505.24722},
+  year={2025},
+}
+```
 
-### Next Steps
+## License
 
-- **KL divergence distillation** from Qwen3-30B using Nebius SWE-agent trajectories (80K agentic tool-use sequences)
-- **Context extension** to 128K via NTK-RoPE scaling
-- **Fine-tuning** on agentic coding trajectories for downstream tool-use tasks
+MIT — see [LICENSE](LICENSE).