TightX/code-llm-1b-base

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.

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