CUDA for ML Training

CUDA isn't a separate "AI mode." Take a production-shaped access-ticket batch with shape (32, 128, 768): 32 tickets, 128 tokens per ticket, and 768 features per token. The previous NumPy shapes lesson used a tiny (2, 3, 4) teaching tensor; the same (B, T, D) contract still applies here at a larger scale. CUDA adds a second contract to that shape contract: where the batch, weights, activations, and gradients live while training runs.

CUDA is NVIDIA's parallel computing platform and programming model. A CPU is built to run a small number of complicated threads with low latency; a GPU is built to run huge amounts of similar arithmetic at high throughput. For dense tensor work like matrix multiplies and attention, that difference can be dramatic once work becomes a large GPU kernel.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/} The flip side matters too: tiny tensors and repeated copies can spend more time on launch and transfer overhead than on useful math.

The arrays are the same ones you already learned to reason about. Now device placement becomes part of the meaning of every tensor. You'll check an environment, move a training batch, catch a placement failure before a forward pass, budget memory, and measure asynchronous work honestly.^{[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html}

CPU orchestration vs GPU execution

A training loop has two different jobs:

1. The CPU side handles Python control flow, dataloading, launching kernels, logging, checkpointing, and filesystem work.
2. The GPU side handles the heavy tensor math: matrix multiplies, attention kernels, layer norms, optimizer updates, and other parallel operations.

That split matters because the CPU and GPU don't share one flat memory space in the way beginners often imagine. On the standard discrete-GPU path, the GPU has its own device memory. If your tensors live on the CPU, GPU kernels can't use them until you copy them over.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Keep this practical comparison in your head:

Think about a CPU coordinator and a GPU worker pool. The CPU schedules work, but the GPU performs the bulk tensor math. Making the CPU path faster doesn't remove the GPU bottleneck when tensor operations dominate the work.

What CUDA is

CUDA is NVIDIA's GPU computing platform and programming model.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/} In practice, for most AI engineers, that means five related ideas:

• Kernels: GPU functions the host launches. From PyTorch's point of view, x @ w or loss.backward() queues one or more kernels; you rarely write the kernel body yourself.
• Thread hierarchy: each kernel launch runs a grid of blocks; each block is a scheduling unit of threads that can share fast on-chip memory. The runtime maps blocks onto the GPU's streaming multiprocessors.
• Warps: on NVIDIA GPUs, threads execute in lockstep groups of 32. If threads in the same warp take different branches, inactive lanes idle until the paths reconverge, which is why branch-heavy kernels lose throughput.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}
• Device memory: in the standard discrete-GPU setup, the GPU has a device-memory pool separate from host RAM.
• Asynchronous launch: the CPU often queues GPU work and continues running until something forces synchronization.

You don't need to write custom CUDA kernels on day one. You do need to understand that model layers, loss computation, backward passes, and optimizer updates launch GPU work once their tensors are on a CUDA device.

Select a compatible PyTorch build

The NVIDIA System Management Interface command, nvidia-smi, reports devices visible to the NVIDIA driver. PyTorch's torch.cuda.is_available() reports whether this Python process can use CUDA. One can succeed while the other fails: for example, the driver may see a GPU while your environment has a CPU-only PyTorch installation.^{[3]Reference 3nvidia-smi documentationhttps://docs.nvidia.com/deploy/nvidia-smi/index.html[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html}

Run these checks before changing code:

If the process can't access CUDA, use PyTorch's official installation selector for the current operating system, package manager, and supported CUDA option.^{[4]Reference 4Get Started.https://pytorch.org/get-started/locally/} Wheel tags and supported runtimes change; a hard-coded installation command in an article ages badly.

First device checks in PyTorch

Before you worry about throughput, make sure tensors land where you think they do. This script is intentionally device-agnostic: it runs on a CUDA machine and remains executable on a CPU-only laptop.

The output above is the happy path from a configured NVIDIA machine. If your local run reports no accessible CUDA device, that doesn't automatically mean your code is wrong. It means one of these is true:

• you're on a machine without an NVIDIA GPU
• the driver is missing or mismatched
• the environment isn't linked to a CUDA-enabled PyTorch build
• the process can't access the GPU

The first quick checks are usually:

nvidia-smi tells you whether the driver sees the device. PyTorch tells you whether the framework can use it.

Host memory vs device memory

Next comes memory placement.

An out-of-memory (OOM) failure is local to the device running your job. A model that loads can still fail on its first training batch because weights are only one part of the budget. For a simplified full-precision Adam optimizer floor, count weights, gradients, and Adam's two running statistics. This still excludes activations, temporary buffers, and allocator overhead, so it's a lower bound rather than a capacity promise.

Three beginner rules cover most cases:

1. Model and inputs must be on compatible devices.
2. Every host-device copy costs time.
3. Training failures often come from memory, not math alone.

A standard PyTorch training loop usually does both:

GPU index can vary; model and batch still need matching CUDA devices, and the shape contract stays stable. CUDA_VISIBLE_DEVICES (covered in the Git chapter) can hide cards from the process: after a mask of 1, PyTorch's cuda:0 is the host's physical GPU 1. Always read placement from next(model.parameters()).device rather than assuming a fixed physical index.

Real batches often contain inputs, labels, and masks. Move every tensor that participates in device work:

If one input stays on CPU while model parameters are on CUDA, the forward pass fails. A small preflight check makes that failure readable before a long training run begins:

A small training example, step by step

Make the example concrete. Suppose you're training an access-ticket model that predicts whether a request should be answered, escalated, or blocked.

1. The dataloader reads a batch of token IDs on the CPU.
2. The batch is copied to device memory.
3. The model weights already live on the GPU.
4. PyTorch launches matmul, attention, and loss kernels on the GPU.
5. Backward pass produces gradients on the GPU.
6. The optimizer updates weights on the GPU.
7. Only when you log a scalar or save results back to disk does the CPU need some of that state again.

That's why CUDA bugs often look strange at first. The Python code line you wrote and the GPU work it triggered are related, but they don't run in one shared place or finish at the same instant.

This complete stochastic-gradient-descent training step keeps the model, features, labels, logits, loss, and gradients on device until the final scalar is brought back for logging:

The rendered output shows the configured NVIDIA path. The training-step structure stays the same across environments, but this lesson's output should model the accelerator run you are aiming for.

The same idea as a small table:

Those five rows are the core CUDA path: assemble on CPU, copy once, compute on GPU, and synchronize only at an intentional boundary.

Asynchronous execution and hidden sync points

One reason CUDA feels confusing is that the CPU usually launches GPU work asynchronously.^{[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html} That means:

• Python may continue before the GPU finishes the queued kernels.
• timing a block with a naive host timer can under-report real GPU time
• operations that need a CPU value force the host to wait for completion

Common sync points include:

• loss.item() when loss is a CUDA tensor
• tensor.cpu(), including the tensor.cpu().numpy() path used for NumPy analysis
• logging or printing that materializes a CUDA value on the CPU
• explicit torch.cuda.synchronize()
• tensor.tolist() and booleanization that forces a CUDA value onto the host (for example using a CUDA tensor in a Python if)
• some distributed collectives that wait for peer ranks

Calling .numpy() directly on a CUDA tensor isn't the route back to NumPy: move it to CPU first. This explicit boundary is a useful place to control logging frequency:

On a CUDA device, the .cpu() call above waits until data needed for the copy is ready. That's why a loop can look fast until you add "just one print."

A timing trap you should recognize by hand

Suppose one forward pass queues 40 ms of GPU work, but the CPU finishes launching it in 2 ms.

• A naive timer wrapped only around the Python call might report about 2 ms.
• A synchronized timer reports the real end-to-end GPU time: about 40 ms.

That mismatch isn't a rounding error. It changes the engineering conclusion.

• If you believe the 2 ms number, you may think the GPU is extremely fast and the bottleneck must be elsewhere.
• If you measure the real 40 ms number, you may correctly conclude that sequence length, batch size, or kernel efficiency still need work.

Beginner CUDA debugging should always ask: did the measurement include synchronization, or did it only measure kernel launch overhead?

For real CUDA measurements, PyTorch recommends CUDA events or explicit synchronization around host timers.^{[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html} Warm up the operation before recording steady-state work because first execution can include one-time setup costs. This script uses events when CUDA is available and keeps a runnable CPU fallback:

Reading nvidia-smi without over-trusting it

nvidia-smi is useful, but it isn't a full profiler. PyTorch also uses a caching allocator, so memory visible in nvidia-smi can include reserved memory that's not currently occupied by live tensors.^{[3]Reference 3nvidia-smi documentationhttps://docs.nvidia.com/deploy/nvidia-smi/index.html[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html}

Use it for:

• checking whether the process attached to the GPU
• checking rough process and device memory footprint
• spotting obvious OOM pressure
• seeing rough utilization snapshots

Don't use it as your only answer for:

• kernel-level bottlenecks
• whether dataloading is the issue
• whether synchronization is killing throughput
• whether the GPU is compute-bound or memory-bound

For code-level memory checks, separate live tensor bytes from allocator reservations. memory_allocated() tracks memory occupied by tensors. memory_reserved() tracks the larger pool managed by PyTorch's caching allocator. That pool can include unused memory kept for fast reuse, which is why nvidia-smi can report more memory than your live tensors occupy.^{[2]Reference 2CUDA semantics.https://docs.pytorch.org/docs/2.12/notes/cuda.html}

At the beginning, nvidia-smi, correct device placement, and these counters catch a large share of broken setups. Detailed profiling comes later.

First CUDA mistakes in training loops

1. Model on GPU, batch on CPU

• Symptom: device mismatch error on forward pass.

• Cause: weights and input tensors are on different devices.

• Fix: move the whole batch, not a single field.

2. OOM on the first real batch

• Symptom: the script starts, maybe even builds the model, then fails on the forward or backward pass.

• Cause: activations and optimizer state push total memory over the card limit. Parameters alone aren't the full bill.

• Fix: shrink per-step batch size first. If you need to preserve effective batch size, accumulate gradients across several smaller steps. Reduce sequence length or enable mixed precision when the task allows it. Activation checkpointing (recompute activations in backward instead of storing them) is the next named lever on transformer workloads. Read the failing allocation size in the CUDA OOM message before guessing; empty_cache() only returns unused reserved blocks to the driver and rarely fixes a true peak that still exceeds free memory.

Memory lever: The first two knobs reduce the number of token positions in a batch. Halving batch size halves that count; halving sequence length does too. Attention score tensors can drop faster when sequence length shrinks because they have two token axes. Parameter + Adam state is a useful floor (see the memory example later), but activations often dominate transformers: attention scores alone scale as O(B · H · T²) in the score tensor before values are mixed.

3. Slow loop despite high GPU memory usage

• Symptom: the GPU memory is full enough to look "active," but throughput is poor.

• Cause: the bottleneck may be dataloading, synchronization, small batch size, or repeated host-device copies.

• Fix: check whether the data pipeline feeds the GPU fast enough before assuming the math kernels are the problem. The standard first host→device throughput habit is pinned host memory plus non-blocking copies:

4. Timing without synchronization

• Symptom: a kernel appears to take almost no time.

• Cause: the timer stopped before queued CUDA work completed.

• Fix: warm up the operation, then synchronize before starting and after enqueueing the measured work, or use CUDA events.

What to understand before writing custom kernels

You don't need Triton or CUDA C++ to start training models, but you should already understand:

• why GPUs help matrix-heavy workloads
• why tensor placement is explicit
• why device memory is limited and precious
• why copies and sync points can dominate step time
• why "GPU utilization" alone isn't a diagnosis

That foundation makes later topics less mysterious:

• mixed precision
• FlashAttention
• FSDP and ZeRO
• tensor parallelism
• custom kernels

Self-check before bigger training runs

Answer these before moving on.

Why can model(x) return before GPU work is done?

Answer

CUDA launch is asynchronous, so Python can continue while kernels are still queued or running.

Why can a training loop OOM even when weights fit?

Answer

Because activations, gradients, optimizer state, and buffers also consume device memory.

Why can one innocent print(loss.item()) change timing?

Answer

Because .item() forces synchronization, so the CPU waits for the GPU result before printing.

What must match before a forward pass works?

Answer

Model parameters and input batch must be on compatible devices.

Your batch tensor is on cpu, your model weights are on cuda:0, and the first forward pass crashes. What is the first fix?

Answer

Move the batch to the same device as the model before the forward pass, typically with batch = batch.to(device). The strong answer should also note that device placement must stay consistent for every tensor the model touches.

What to remember

• CUDA is an execution and memory model, not a speed checkbox.
• The CPU orchestrates. The GPU executes dense parallel math.
• Host RAM and device memory are different places with real transfer costs.
• PyTorch queues CUDA work asynchronously, so .item() and .cpu() can stall the host.
• OOM errors usually mean the full training footprint doesn't fit, not model weights alone.

If that picture feels solid, you're ready to reason about training loops on accelerators instead of treating the GPU as an opaque speed device.

Use this checklist as the handoff artifact: run the device check, measure one operation with explicit synchronization, and write which memory terms can trigger OOM.