Data-driven software training: LLM error isolation and profiling steps

Goals and toolset — in a training context

In this article our goals are: to reproduce errors with isolation, measure GPU/CPU bottlenecks, and apply model/training pipeline optimizations validated by verifiable measurements; concrete tools used: PyTorch, Hugging Face Transformers, PEFT, torch.profiler, NVIDIA Nsight Systems (nsys), nvidia-smi and Weights & Biases (wandb). Example commands will be given at each step, for instance to install required packages:

pip install torch transformers accelerate peft torch-profiler nvidia-ml-py3 wandb

This installation line provides a directly reproducible starting point.

End-to-end error isolation: reproducer and system profile

The first rule of error isolation is to build a deterministic, small reproducer; write a script that runs the training loop with a single batch and collect GPU/CPU measurements. Example minimal reproducer:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained('gpt2')
model = AutoModelForCausalLM.from_pretrained('gpt2').cuda()
inputs = tok('def foo():\n    return 1', return_tensors='pt').to('cuda')
for i in range(10):
    out = model(**inputs, labels=inputs['input_ids'])
    out.loss.backward()
    torch.cuda.synchronize()
print('done')

With this script you can first capture a system profile with nsys:

nsys profile --output=repro.prof python reproducible_run.py

Inspect the nsys output with 'nsys stats repro.prof' or the Nsight Systems GUI; host-to-device transfers, kernel launch latency and memcpy times will be visible directly there.

Detailed profiling: layer-level timing and memory with torch.profiler

Use torch.profiler to capture GPU kernel timings, operator-level FLOP estimates, and memory snapshots; add the following wrapper into your training loop:

from torch.profiler import profile, record_function, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True,
             with_stack=True,
             profile_memory=True) as prof:
    for step in range(20):
        with record_function('step'):
            loss = model(**inputs, labels=inputs['input_ids']).loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=20))

The output will show the operators consuming the most time (e.g. aten::mma, aten::conv2d) and memory hotspots; this data technically justifies which layers will benefit from mixed precision or kernel tuning.

Before-and-after comparison and measurable KPIs

Define clear KPIs to validate optimizations: tokens/sec, p99 latency, peak GPU memory (MiB), and host-to-device bandwidth (GB/s). To measure token throughput add a simple counter to the training loop and record a baseline:

import time
steps=100
start=time.time()
for i in range(steps):
    loss = model(**inputs, labels=inputs['input_ids']).loss
    loss.backward()
    optimizer.step(); optimizer.zero_grad()
    torch.cuda.synchronize()
end=time.time()
print('tokens/sec', (inputs['input_ids'].numel()*steps)/(end-start))

After the baseline, re-run each change you make (e.g. gradient_checkpointing, amp, batch size) as a separate run and report percentage change; use pandas to automatically produce comparison tables and collect them in CI with wandb.log.

Concrete optimizations: configuration examples and mechanisms

Below are commonly used changes, application examples and the mechanical explanations of why they work: 1) Mixed precision (AMP/BF16) — mechanism: FP16 reduces arithmetic bandwidth and register usage so kernels complete faster; example usage:

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data in loader:
    with autocast():
        loss = model(**data).loss
    scaler.scale(loss).backward()
    scaler.step(optimizer); scaler.update()

2) Gradient checkpointing — mechanism: reduce memory by not storing intermediate activations and recomputing them in the backward pass; for Hugging Face Transformers:

model.gradient_checkpointing_enable()

3) Data pipeline: tune num_workers, pin_memory and prefetch for the DataLoader; example:

loader = DataLoader(ds, batch_size=8, num_workers=8, pin_memory=True, prefetch_factor=2)

Relate these changes to torch.profiler outputs for memory and memcpy times to verify why they are effective.

CI, regression tests and replay: ensuring with automation

To prevent performance regressions run small, deterministic perf tests on every PR; example GitHub Actions step:

- name: Run perf smoke
  run: |
    pip install -r requirements.txt
    python perf_repro.py --steps 10 --log perf.json
    python perf_compare.py --baseline baseline.json --current perf.json

For metric logging use wandb or MLflow:

import wandb
wandb.init(project='perf-ci')
wandb.log({'tokens/sec': tokens_per_sec, 'peak_gpu_mem': peak_mem})

Also set tolerances for regression detection (e.g. tokens/sec -3%) and add a CI check that blocks PRs exceeding the tolerance.

Frequently Asked Questions

How do you quickly write a perf reproducer in a training context?

For a quick perf reproducer use a minimal batch, a single batch-level seed and a single GPU; example command:

python reproducible_run.py --seed 42 --steps 20

In the script call torch.manual_seed(42), run a short loop in train mode instead of eval, and use torch.cuda.synchronize() for timing. This approach produces results that are directly comparable with Nsight or torch.profiler outputs.

Which metrics best summarize training performance and how are they measured?

KPI set: tokens/sec, p99 latency, peak GPU memory (MiB) and host-to-device throughput (GB/s). Measure tokens/sec from the training loops as processed tokens per second; capture memcpy times with nsys or nvidia-smi to compute host-to-device throughput. Example nvidia-smi sampling:

nvidia-smi --query-gpu=utilization.gpu,utilization.memory,memory.used --format=csv -l 1

How to automate regression tests during training and which tools to use?

Add small perf smoke tests into CI, send results to wandb or MLflow and write a tolerance-checking script. Example tolerance check:

import json
b=json.load(open('baseline.json'))
 c=json.load(open('perf.json'))
 assert c['tokens/sec'] >= b['tokens/sec']*0.97

If this assert fails in CI the PR is rejected, catching regressions early.