51.0K tok/s
ML systems performance
LLM Training Performance Lab
Reproducible benchmark lab
Optimized GPT-style Transformer training on an RTX 3060, increasing throughput by 3.02x while reducing peak GPU memory by approximately 52.4%.
Role and scope
Implemented the benchmark harness, configurable GPT-style Transformer, profiling scripts, experiment outputs, optimization matrix, and results documentation.
PythonPyTorchCUDAGPT-style TransformersBF16FP16SDPAPyTorch Profilertorch.compileWSL2
154.2K tok/s
Optimized throughput
3.02x
Throughput lift
52.4%
Peak memory reduction
Case study
What this shows
This lab makes LLM training performance work reproducible on a consumer RTX 3060 Laptop GPU by combining controlled benchmark scripts, profiler traces, precision experiments, and documented tradeoffs.
Problem
Training-performance work requires more than final loss. The lab measures throughput, step latency, GPU memory, profiler traces, precision modes, and optimization tradeoffs on constrained consumer GPU hardware.
Hardware constraints
The primary constraint was useful GPT-style training measurement on local NVIDIA RTX 3060 Laptop GPU hardware rather than a large training cluster.
Baseline
The WSL2 gpt_small baseline measured 50,997.73 tokens/sec and 1,119.76 MB peak GPU memory with FP32 manual attention.
Profiling methodology
Runs captured average tokens/sec, step latency, GPU memory, warmup steps, benchmark steps, validation loss, perplexity, and profiler artifacts where appropriate.
Optimizations tested
The optimization matrix covered TF32, scaled dot-product attention, FP16, BF16, fused AdamW, and torch.compile under WSL2 after environment validation.
Final results
The best WSL2 result was BF16 + fused AdamW + torch.compile at 154,197.54 tokens/sec, 26.56 ms/step, 532.91 MB max GPU memory, and a 3.02x speedup.
Tradeoffs
Mixed precision reduced memory, but small-model throughput did not always improve. The lab documents where optimization helped and where framework overhead or environment support limited results.
Reproducibility
The repository includes training, evaluation, benchmark, profiling, configuration, result, CSV, and JSON artifacts for repeatable experiments.
Future work
Next steps include native Linux benchmarking, distributed data parallel tests, activation checkpointing comparison, custom CUDA kernels, and Nsight Systems profiling.
Performance visualization
Baseline to optimized
A simple visual summary of the WSL2 validation run: throughput increased while peak GPU memory dropped.
Throughput
Baseline51.0K tok/s
Optimized154.2K tok/s
Peak GPU memory
Baseline1,120 MB
Optimized533 MB
Interested in the implementation?