Fix CodeGen inference, make it fast, make training train

[Ayyankhan101]

Three passes over this repo. Each found the same shape of defect: code that runs, prints
plausible numbers, and does the wrong thing or nothing at all. Each is now covered by a
check that did not exist before.

1. CodeGen inference was incorrect

The model could not produce valid output before this.

• The fused qkv_proj tensor was split wrong. The code assumed
  [all q | all v | all k]. CodeGen keeps the sharded layout of the original TPU
  implementation: four interleaved groups of [q | v | k] (mp_num = 4 upstream). Every
  attention head was reading a different slice of the projection than it should.
• Sampling was always greedy. The hand-rolled LCG divided a u64 state by u32::MAX,
  producing values around 6e9, so the selection loop always fell through to argmax —
  temperature, top_k and top_p did nothing.
• complete, repl and POST /generate sliced off the wrong tokens, stripping
  prompt.len() entries from a generated-only vector: dropped output, or a panic when the
  prompt was longer than the completion.
• The HTTP server ignored config.json and --f16, loading with
  CodeGenConfig::default() — rotary_dim 64 rather than the checkpoint's 32.
• Missing checkpoint tensors were silently ignored, leaving zeroed layers that produce
  nonsense rather than an error.
• F16 models panicked on the first forward: new_blank hardcoded F32 regardless of
  config.dtype.

tests/codegen_parity.rs now checks logits against the HuggingFace reference using
committed tiny-model fixtures, so it needs no weight download. It fails at a max logit
difference of 0.66 against the old QKV split and passes at 4e-7 against the new one.

2. Performance

Measured on the benchmark model:

	before	after
decode, per token	58.4 ms	5.6 ms
prefill, 32 tokens	88.5 ms	36.2 ms

KVCache::append used slice_assign, which is not a targeted write — candle zero-pads the
source to the full buffer shape, allocates a full-size mask, pads that too, then runs
where_cond across every element, for k and v, in every layer, for every token.
slice_set copies only the new tokens in place.

lm_head no longer runs over the whole prompt when only the final position is read: 6.0 ms
of a 35.8 ms prefill.

benches/transformer.rs now benchmarks this crate. Every benchmark used to be a
hand-inlined copy of the model, carrying the QKV bug fixed above, so the suite could report
neither a speedup nor a regression.

The CI lint gate was already failing. main.rs re-declared the whole module tree
instead of using the library, so the binary compiled a second copy and every pub item the
CLI did not call was reported as dead code — 34 warnings from one cause.
cargo clippy --all-targets --all-features -- -D warnings goes from exit 101 to exit 0.

Removed src/codegen/quantized.rs: never constructed outside its own tests, and its
forward dequantized the entire weight matrix per call, making it strictly slower than the
F32 path it would have replaced. README, docs and the comparison table no longer claim it.
Added safetensors loading, so the converter's output is finally loadable.

3. glm-train did not train

Both cross-entropy paths summed into an f32 and returned Tensor::new(total, device) — a
fresh leaf with no autograd history. Backprop from a constant yields an empty gradient
store, so AdamW updated nothing and every step was a no-op. The Var wiring was correct;
only the loss was detached.

A real 80-step run now falls from loss 10.80 to 4.99. Before the fix, 60 optimizer steps
left the loss at 4.1934047 both before and after, and 0 of 19 parameters received gradients.

Found while verifying that:

• Every checkpoint save failed — save_safetensors called to_vec1::<u8>(), which errors
  on rank > 1, so runs ended in unexpected rank, expected: 1, got: 2.
• configs/train.yaml did not deserialize: eval_steps sat under a separate evaluation:
  section and tokenizer_path was missing. Nothing had ever loaded the file because
  glm-train had no --config flag. Added the flag, plus #[serde(default)] so partial
  configs fall back field by field.
• gradient_accumulation_steps and max_grad_norm were declared and claimed as delivered,
  but read nowhere. Both are implemented now.
• The learning rate was applied after optimizer.step, so each step ran on the previous
  step's rate.
• Deleted src/training/checkpoint.rs: 192 lines never declared in mod.rs, so never
  compiled, duplicating the checkpoint code in train.rs.

The README and docs/train-code-infill.md carried a fabricated training log — invented
losses and per-step timings, in a format the code never emitted. Replaced with real output.

Verification

98 tests pass
cargo clippy --all-targets --all-features -- -D warnings   exit 0 (was 101)
cargo fmt --all -- --check                                  clean
cargo build --release --features server                     ok

Known gap

Nothing here has run against the real 350M checkpoint —
codegen_weights/pytorch_model.bin has never been present, so both tests in
tests/codegen_integration.rs self-skip. The parity fixtures share CodeGen's architecture
but exercise 2 layers and a 256-token vocabulary. That validation is the next pass.

[git-mind-pr-guardian]

GitMind PR Review

{'pr_url': '#15', 'repo': 'Ayyankhan101/Transformer-In-Rust', 'pr_num': '15', 'title': 'Fix CodeGen inference, make it fast, make training train', 'author': 'Ayyankhan101', 'additions': 1701, 'deletions': 1509, 'changed_files': 48, 'review': '## Summary\nAutomatic review could not be generated — please review the diff manually.\n\n## Issues\nNo significant issues found.\n\n## Verdict\n💬 NEEDS DISCUSSION\nThe automated review failed; human review required.'}

---

Auto-generated by GitMind — AI-powered code analysis.

[Ayyankhan101]

Validated against the real 350M checkpoint

Downloaded Salesforce/codegen-350M-multi and ran it. The inference fixes hold — greedy decoding produces correct Python:

def fibonacci(n):
    if n == 0:
        return 0
    elif n == 1:
        return 1
    else:
        return fibonacci(n-1) + fibonacci(n-2)

That confirms the QKV model-parallel split against Salesforce's actual weights, not just the parity fixtures. Both integration tests that had always self-skipped now run, plus two new ones.

Four more bugs, all invisible without weights

• codegen download reported success without downloading anything. huggingface-cli has been renamed to hf; the old name prints a deprecation notice, downloads nothing, and exits 0. The command trusted that exit code and printed "✓ Download complete!" over an empty directory.
• complete printed nothing in its default streaming mode. with_tokenizer existed but was never called, so decode_token returned "" for every token.
• chat sliced off the wrong tokens — the same generated-only slicing bug fixed in complete, repl and the server earlier in this PR, missed in the fourth caller.
• -t was bound to both --temperature and --template, and --stream had default_value = "true" so it could never be false — the non-streaming branch was unreachable.

Real numbers

Apple M1 Pro, 64 tokens from def quicksort(arr):, greedy:

	F32	F16
Per token	68.8 ms	20.3 ms
64 tokens	4.4 s	1.3 s
Peak RSS	1.63 GB	1.08 GB

F16 is 3.4× faster, not the 23% the README claimed. README and docs now carry measured numbers.

One hypothesis that did not survive testing

CodeGen pads vocab_size to 51200 while the tokenizer stops at 50294, so 905 untrained rows sit on top of lm_head. I expected those to be a sampling hazard and wrote logit masking for it — then tested at temperature 2.0 with no top-k or nucleus cut, and they never win. The masking was removed and the test kept, so a regression would surface instead of being silently suppressed.

Also verified: the safetensors converter's output loads and gives byte-identical greedy output to the pickle path.

102 tests pass, clippy clean.

[git-mind-pr-guardian]

GitMind PR Review

{'pr_url': '#15', 'repo': 'Ayyankhan101/Transformer-In-Rust', 'pr_num': '15', 'title': 'Fix CodeGen inference, make it fast, make training train', 'author': 'Ayyankhan101', 'additions': 2060, 'deletions': 1603, 'changed_files': 51, 'review': '## Summary\nAutomatic review could not be generated — please review the diff manually.\n\n## Issues\nNo significant issues found.\n\n## Verdict\n💬 NEEDS DISCUSSION\nThe automated review failed; human review required.'}

---

Auto-generated by GitMind — AI-powered code analysis.