| --- |
| 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. |
|
|
