File size: 6,161 Bytes
8de65a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | ---
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.
|