ai infrastructure · gpu llm

From-scratch GPU LLM training.

A complete GPU forward + backward stack for Qwen3-32B, written bottom-up, parity-validated against a CPU reference, and exercised on a real, vendor-neutral mixed fleet — datacenter-class and local GPUs. Not a notebook demo; runs on the metal.

Why from-scratch

Most LLM training stacks are vendored — PyTorch + Megatron + DeepSpeed + a cloud abstraction layer. Each link is a black box from the math perspective and the deployment story is "rent a hyperscaler".

The KeiLab stack is written bottom-up: six gradient kernels, hand-rolled, finite-difference verified, parity-checked against a CPU reference layer-backward. The same code runs across the fleet — datacenter-class and local GPUs alike — no vendor lock, no hidden randomness in the toolchain.

Hardware actually exercised

  • Datacenter-class GPU tiers— sized to the run: a heaviest tier for full-depth runs where activations and optimizer state would squeeze a smaller card, a single-GPU tier for domain-injection fine-tunes with QLoRA, and a multi-GPU DDP tier reserved for when wall-clock pressure justifies the cost multiplier (per the lab's own hardware-right-sizing rule — more GPUs is never the reflexive default).
  • Local unified-memory ARM GPU workstations — used for inference, KV-streaming, GPU kernel development, and zero-cost experimentation outside the cloud cost budget. The same kernels run here unmodified.

What is built

  • Six gradient kernels on GPU — matrix-multiply, RMSNorm, SiLU / SwiGLU, RoPE, softmax cross-entropy, scaled-dot-product attention.
  • GPU layer-backward orchestrating all six — parity-validated against the CPU reference.
  • GPU forward path with paged key-value cache, saved-activations for backward, and FlashAttention-2 wired for long context.
  • GPU lm-head — softmax cross-entropy + argmax fused in a single launch.
  • Layer-streaming forward — full 64-layer Qwen3-32B runs by holding one layer at a time. Depth-independent peak RAM, lossless against bulk forward.
  • End-to-end smoke on real Qwen3-32B weights with a real next-token CE objective — loss decreasing monotonically under gradient clipping across the fleet.

Operating discipline

Every cloud launch goes through a pre-flight memory and cost gate before the meter starts — measured per-GPU budget, dry-run loss, single-GPU-first sizing — under the lab's standing rules. Launches are orchestrated end-to-end (spawn → upload → train → poll → download → terminate) so there are no orphan billed pods.

Sensitive training artifacts are transferred directly between the lab's own machines, never staged through a public model hub. The audit trail is preserved per run.

Honesty

Combat-scale training at full depth still requires streaming offload and grad-clip discipline. The kernels are proven; integration into a continuous training loop is bounded engineering — not a missing capability. Partners can run the kernels in their own environment to verify before committing.