Built a 380M param LLM from scratch (no PyTorch/TF/JAX) and found a nasty residual connection bug via a custom debugger (that works across all major ml frameworks)

Since December, I’ve been building a 380M parameter LLM entirely from scratch, using only NumPy and CuPy (no PyTorch, TensorFlow, or JAX). What started as an exercise to deeply understand transformers by implementing “Attention Is All You Need” grew into a full training setup with AdamW, gradient accumulation, and batching—all written by hand to run on consumer hardware.

Recently, I hit a wall: after increasing the vocab size to 50,000 tokens, the loss plateaued hard at 6.5 and refused to budge. Because a single loss curve gives you zero insight into millions of possible failure points, I built a lightweight live ML training debugger (Pulse) to capture live tensor statistics and heatmaps.

The tool caught the root cause: my residual connection was incorrect:

currentX = currentX + hidden2 / math.sqrt(num_layers)

instead of

currentX = currentX + hidden2 / num_layers

This was silently crushing the signal. Once fixed, the loss finally dropped past 4.0.

Quick Installation & Usage

Step 1:

If you want to try out the debugging tool, here is how to get started:

git clone https://github.com/codeyash09/PulseML.git
cd PulseML
pip install numpy matplotlib pillow litellm --break-system-packages

Add your project files to the new PulseML folder

- or -

Copy and paste the PULSE folder directly into your project.

Step 2:

from PULSE.pulse import auto_track
#Pass your training function for shape discovery (optional), or call directly
if __name__ == '__main__':
      auto_track()
      #Your training loop

Repositories & Links

LLM Repo (NumPy/CuPy): GitHub - codeyash09/Yash-GPT: Generally Productive Toy · GitHub

The Debugger (Pulse): GitHub - GitHub - codeyash09/PulseML: Pulse debugs ml. · GitHub

Happy to answer any questions about the from-scratch implementation or the debugging workflow!