CUDA for ML Training

CUDA isn't a separate "AI mode." Suppose the access-ticket batch from the previous lesson has shape (32, 128, 768): 32 tickets, 128 tokens per ticket, and 768 features per token. 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: functions that run on the GPU across many threads in parallel.
• Thread hierarchy: threads are grouped into blocks, and blocks are grouped into a grid.
• Warps: on NVIDIA GPUs, threads execute in groups of 32 called warps, so branch-heavy code can waste throughput when lanes in a warp diverge.^{[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.

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:

If you can explain those five rows in your own words, you already understand more CUDA than many people who only know the slogan "GPUs are parallel."

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()

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.

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.

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.

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.

Expected output from your own explanation

At this point, explain without code:

1. where the tensor starts
2. when it moves
3. where the heavy math runs
4. which operations force the CPU to wait
5. which memory terms can trigger OOM

If one of those five is fuzzy, re-read the step-by-step table and the timing trap section before moving on. Later training chapters assume this picture is stable.

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.

Mastery check

Key concepts

• CPU vs GPU execution roles
• host memory vs device memory
• kernels, thread blocks, and warps
• PyTorch device placement
• asynchronous CUDA execution
• common synchronization points
• CUDA OOM debugging basics
• nvidia-smi and runtime sanity checks

Evaluation rubric

• Foundational: Explains why training work moves from CPU orchestration to GPU kernels and device memory
• Intermediate: Uses PyTorch device placement correctly and names when host-device transfers become a bottleneck
• Advanced: Diagnoses first-line CUDA issues such as missing device placement, accidental synchronization, and out-of-memory failures

Follow-up questions

Common pitfalls

• Treating CUDA as a speed flag instead of a different execution and memory model.
• Moving the model to GPU but forgetting the inputs, causing device-mismatch errors.
• Reading nvidia-smi as if it were a profiler. It shows memory and utilization snapshots, not full kernel timelines.
• Calling .item(), .cpu(), or print() inside tight loops without realizing they can force synchronization.