Update README.md

README.md

@@ -1,3 +1,82 @@
 ---
 license: apache-2.0
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+tags:
+- reinforcement-learning
+- grpo
+- lora
+- jax
+- countdown
+- reasoning
+language:
+- en
+pipeline_tag: text-generation
 ---
+
+# NanoZero Countdown LoRA — GRPO from scratch in JAX, on one free T4
+
+A LoRA adapter for [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct), RL-trained with **GRPO implemented from scratch in pure JAX** on the [Countdown task](https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4) — a tinker-free reproduction of [rLLM](https://github.com/rllm-org/rllm)'s Countdown RL recipe that fits on a single free Colab T4 (16 GB).
+
+**Code:** [github.com/Zayed024/nanozero-jax](https://github.com/Zayed024/nanozero-jax) — single-file model + GRPO, no training framework.
+
+## Results
+
+| | pass@1 (128 held-out Countdown problems, greedy) |
+|---|---|
+| Qwen2.5-0.5B-Instruct (baseline) | **0.00%** |
+| + this adapter (200 GRPO steps) | **15.62%** |
+
+The task: given numbers like `[3, 7, 11]` and a target like `28`, produce an arithmetic
+equation using each number exactly once, in `<answer>...</answer>` tags. Reward is rLLM's
+exact `compute_score` (1.0 correct / 0.1 valid-format / 0.0 no answer), vendored verbatim.
+Training showed the TinyZero-style two-phase dynamic: format acquisition first (mean reward
+0.01 → 0.10), then actual solving (solved% climbing from 0).
+
+## Training setup
+
+- **Algorithm:** GRPO — group-relative advantages (8 rollouts/prompt, z-scored), PPO-style
+  clipped policy gradient, k3 KL penalty to the frozen reference (β=0.001).
+- **Adapter:** LoRA rank 16 on attention projections (q/k/v/o). The **reference model is
+  the same frozen base with adapters off** — one weight copy serves policy and reference;
+  B is zero-init so the policy starts exactly at the reference.
+- **Budget:** 200 steps × (8 prompts × group 8) × 256 new tokens, temperature 1.0,
+  AdamW lr 1e-4, global-norm clip 1.0. ~40 s/step on a T4 after KV-caching.
+- **Fitting 16 GB** (the point of the exercise): never materialize the full `[B, T, vocab]`
+  logits tensor (chunked LM head, ≈12 GB avoided), per-layer gradient checkpointing on the
+  differentiated pass, bit-exact KV-cache decoding, degenerate-group skip.
+- **Forward-pass fidelity:** the from-scratch JAX Qwen2 matches HF logits to max
+  |diff| = 2.6e-4 (argmax agreement 1.000) on the same weights.
+
+## Usage
+
+The adapter is a plain `.npz` of LoRA A/B matrices keyed `"{layer}__{proj}__{A|B}"`
+(projections `wq/wk/wv/wo`, applied as `W·x + (x·A)·B`). With the NanoZero code:
+
+```python
+import nanozero as nz
+from huggingface_hub import hf_hub_download
+
+params, cfg, path = nz.load_params("Qwen/Qwen2.5-0.5B-Instruct")
+lora = nz.load_lora(hf_hub_download("Zayed024/nanozero-countdown-lora",
+                                    "nanozero_countdown_lora.npz"))
+# generate with the adapter:
+ids, mask, resp, lp = nz.generate(params, prompt_ids, prompt_mask, cfg,
+                                  max_new=256, key=key, eos_id=eos, pad_id=pad,
+                                  temperature=0.0, lora=lora)
+```
+
+(It is **not** a PEFT-format adapter; it pairs with the NanoZero codebase. Conversion to
+PEFT is straightforward from the key layout above if you need it.)
+
+## Intended use & limitations
+
+Research/educational artifact demonstrating a minimal, verified GRPO pipeline. Trained
+only on Countdown arithmetic — it improves equation-writing under the `<answer>` format
+and nothing else; expect no gains (and possible format quirks) outside that task. Base
+model license and usage terms apply.
+
+## Acknowledgements
+
+- Reward and recipe: [rLLM](https://github.com/rllm-org/rllm) (Berkeley Sky Computing Lab), `countdown_reward.py` vendored under Apache-2.0.
+- Task data: [Jiayi-Pan/Countdown-Tasks-3to4](https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4) (TinyZero).
+- Base model: Qwen2.5-0.5B-Instruct.