add README.md

README.md

@@ -0,0 +1,108 @@
+---
+license: mit
+language:
+  - code
+tags:
+  - code-generation
+  - python
+  - from-scratch
+  - pretrained
+  - wsm
+library_name: pytorch
+---
+
+# code-llm-1b (base) — a 1.1B Python code-completion model trained from scratch on one RTX 5090
+
+A 1,095,845,888-parameter decoder-only Transformer for Python code completion. Data pipeline,
+tokenizer, architecture, training loop, checkpoint merging and evaluation were all written and run by
+one person on one consumer GPU. This folder holds the weights; the code, design record, ablation
+report and failure log live at https://github.com/TightXu/Simple_Code_LLM.
+
+> **This folder is the pretrained base.** Its weights are the *merged* (weight-averaged) product of the
+> pretraining run: under this training scheme the merged checkpoint, not the last optimizer step, is
+> the finished pretrained model — hence the name. The instruction-tuned variant is in [`../sft`](../sft).
+
+## Architecture
+
+| | |
+|---|---|
+| Parameters | **1,095,845,888** (bf16), 143 tensors |
+| Layers | 20 |
+| d_model / heads | 2048 / 16 (head_dim 128) |
+| FFN | SwiGLU, d_ff 5120 (gate / up / down) |
+| Position | RoPE, theta 500000 |
+| Norm | RMSNorm, learned `scale`, eps 1e-6 (pre-norm) |
+| Attention | causal SDPA, fused QKV projection, no biases anywhere |
+| Embeddings | untied (input `embed` + output `lm_head`) |
+| Vocabulary | 32,000 byte-level BPE (same tokenizer as the 435M release, not retrained) |
+| Context | 1024 tokens |
+| Precision | bf16 (released file; trained with FP8 matmuls on fp32 master weights) |
+
+Parameter count is the sum of the 143 released tensors, and equals the count of the instantiated
+`CodeLLM` class. It is deliberately not the value of `ModelConfig.num_params`, whose closed form
+double-counts the two per-block norms and omits `final_norm` (that property returns 1,095,925,760).
+
+## Training
+
+- **Data**: the filtered Python corpus (`data/all/train.bin`), one continuous token stream with no
+  dataset boundaries. This checkpoint stopped at **30,977,679,360 tokens (31.0 B)** after
+  **630,000 steps**.
+- **Schedule**: WSM — constant learning rate, then a cooldown; the released weights are the
+  **equal-weight average of the last 21 checkpoints (saved every 500 steps, a 10,500-step window)** (`merge_checkpoints_weighted.py`),
+  not the last step.
+- **Hardware**: a single RTX 5090, 32 GB, Blackwell sm_120.
+- **SFT**: the instruction-tuned release starts from exactly these weights.
+
+## What it can and cannot do
+
+**Can**: complete short Python functions when given a signature and the opening indentation.
+On a human-graded benchmark — 11 tasks, five seeds, scored 0/1/2, maximum 110 — this model scores
+**83/110 on the natural-language prompt column** and **61/110 on the signature-only column**. Same
+ruler, same fine-tuned sibling (`../sft`): 97/110 and 91/110, i.e. the SFT pass is worth +14 on the
+prose column and +30 on the signature-only column, mostly on multi-step tasks.
+
+**Cannot**: follow instructions — this is a completion model, not a chat model. Use [`../sft`](../sft)
+for the instruction-tuned variant. Failure modes are structural rather than random: recursive sorting
+(`mergesort`, and `quicksort` on the prose column) and primality testing never pass, several tasks
+produce a correct skeleton with one wrong operator, and name/edge-case discipline is weak. It is also
+blind to docstrings: at this scale a from-scratch model reads a docstring as the end of the function,
+so standard HumanEval docstring prompts score near zero and docstring-free prompts are used instead.
+
+## Files
+
+| File | What |
+|---|---|
+| `model.safetensors` | bf16 weights, 143 tensors, 2,191,706,272 bytes — verified bit-identical to `checkpoints_wsm_1b/merged/checkpoint.pt` after the fp32→bf16 cast |
+| `config.json` | architecture config taken from the checkpoint itself, plus `torch_dtype: bf16` |
+| `tokenizer.json` | byte-level BPE, 32,000 tokens |
+| `SHA256SUMS.txt` | `6b89f4eebe5bb79bd3f98c5dfaae303d80661c8a3498336ae2dcbf2920184081  model.safetensors` |
+| `arch_facts.json` | provenance: source checkpoint, step, token count, source config |
+
+Verification artifacts (the hash, key/shape, bitwise-equality, strict-load and CPU-forward logs) stay in the training workspace and are **not** part of this release; `SHA256SUMS.txt` is what lets anyone re-check the weights independently.
+
+`config.json` carries `use_fp8: false`. The training config said `true`, but that selects FP8 training
+kernels (`FP8Linear`); it is not a property of these weights, and leaving it set makes
+`ModelConfig(**config)` build layers that cannot run outside CUDA.
+
+Loading requires the model class from the training repository (`pretrain_train_wsm.py`, class
+`CodeLLM`, `ModelConfig(**config.json)`): `embed` → scaled by `sqrt(d_model)` → 20 pre-norm blocks →
+`final_norm` → `lm_head`. This is not a `transformers` architecture and has no `AutoModel` mapping.
+
+One loading gotcha, verified the hard way on this folder: `RMSNorm` computes its mean of squares in
+fp32, so its output is fp32 even when the weights are bf16, and a bare bf16 CPU forward raises
+`expected m1 and m2 to have the same dtype` inside `nn.Linear`. Training never saw this because it
+ran under CUDA autocast. Run inference under `torch.autocast("cpu", dtype=torch.bfloat16)` (what the
+verification used, 0.1 s for a 13-token prompt) or load the model in fp32.
+
+## Intended use and limits
+
+Research and education: what a 1.1B model actually learns when trained end to end on real data with
+one GPU. Not for production code generation, not for instruction following, not a substitute for a
+competent developer. Trained only on code; no personal data. Outputs may reproduce patterns (and
+The evaluation is 11 tasks × 5 seeds per checkpoint, every sample graded 0/1/2 by hand (110 points per column); treat differences smaller than the per-seed spread as noise.
+
+## Links
+
+Everything else — `DESIGN.md` (decisions and why), `CHALLENGES.md` (every bug and contaminant), the
+ablation report, the fine-tuning log with corrections, and every evaluated generation — is at
+https://github.com/TightXu/Simple_Code_LLM. MIT licensed.