CUDA for ML Training

The NumPy lesson tracked what each tensor axis means: examples, token positions, and features. Reuse that shape reasoning for a small access-ticket classifier: four tickets, eight token positions, and 16 features give shape (4, 8, 16). Each ticket is a request that the classifier will label answer, escalate, or block.

Now ask a physical question: where does that same tensor live? An operation can have the right shape and still fail when its tensors disagree about CPU memory versus GPU memory. Follow one miniature batch through indexing, placement, memory pressure, timing, and diagnosis.

CUDA is NVIDIA's platform and programming model for running parallel work on NVIDIA GPUs. It doesn't turn every Python line into GPU code. Python still runs on the CPU, while PyTorch launches GPU functions called CUDA kernels for tensor operations whose data lives on a CUDA device.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Eight tokens onto two blocks

Would the four-ticket miniature automatically run faster on a GPU? Not necessarily. A GPU earns its keep when many items perform similar arithmetic; a tiny tensor, branch-heavy Python loop, or copy on every line can lose to launch and transfer overhead.

A CPU is tuned for low-latency control flow and varied work. A GPU is tuned for high throughput. Large matrix multiplications, convolutions, and attention scores usually offer enough parallel arithmetic to justify the handoff.

The CPU prepares and launches; the GPU executes. Inside one kernel launch, a thread handles one logical slice of work. Threads are grouped into thread blocks, and all blocks for one kernel form a grid. CUDA schedules each whole block onto one streaming multiprocessor (SM), a hardware processor inside the GPU. Several blocks may be active on one SM, and CUDA doesn't promise which block runs first.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Take just the eight token positions of the first ticket, not all four tickets or all their features. Two blocks with four threads each is a convenient teaching size. Before running the trace, predict which block and thread own T6. For a one-dimensional grid of one-dimensional blocks, each thread computes a unique global index:

$$\text{global\_id} = \text{block\_id} \times \text{threads\_per\_block} + \text{thread\_id}$$

For block 1, thread 2, the global index is 1 × 4 + 2 = 6, so that thread handles token position T6. The Python trace checks all eight assignments before the CUDA vocabulary grows. It takes block and thread counts as input and prints each block's token positions:

The trace has eight useful threads, but hardware groups threads into a larger unit. An NVIDIA CUDA warp contains 32 threads from one block. Warp lanes execute through a single-instruction, multiple-thread (SIMT) model. If lanes take different branches, CUDA masks inactive lanes while it executes each required path, so divergence can waste throughput.^{[1]Reference 1CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

The four-thread teaching block is legal, and it still occupies a full warp. CUDA fills that warp with consecutive thread ids from the same block, so 4 lanes do ticket work and 28 lanes sit unused for the whole launch. That's a teaching cost, not a production default.

Read the hierarchy from a kernel launch downward: its grid contains blocks; each block is scheduled on one SM and advances in warps; each thread owns its index. A PyTorch operation can launch several kernels, and a shape-only view can launch none. Use the table to name what each level contributes:

PyTorch's matrix-multiply kernels use optimized tiling that is more complex than one thread per token. The eight-position trace teaches indexing and scheduling vocabulary without pretending to describe an optimized library kernel exactly.

What if there are ten positions instead of eight? Two four-thread blocks miss the last two positions; three blocks create twelve threads. The final two threads need a bounds check before reading or writing an array. This CPU simulation checks both coverage and the unused tail:

Removing the guard would attempt indexes 10 and 11, beyond the ten-element array. These are ordinary Python simulations of CUDA indexing, not GPU kernels. PyTorch supplies the indexing and bounds handling for the tensor operations used below.

Those launches only happen if this process can see a CUDA device. The next check is whether the driver, the PyTorch build, and the GPU agree.

Check driver, PyTorch build, and device

The phrase "my CUDA version" hides several boundaries. Before running any command, predict what each tool can actually prove. The NVIDIA System Management Interface command, nvidia-smi, reads driver and GPU state; nvcc is NVIDIA's CUDA compiler:

The CUDA UMD Version describes driver capability, not an installed compiler. A machine can run a prebuilt CUDA-enabled PyTorch package without having nvcc installed.^{[2]Reference 2nvidia-smi documentationhttps://docs.nvidia.com/deploy/nvidia-smi/index.html} It can't prove that this Python process loaded a CUDA-enabled PyTorch build.

PyTorch's installer offers builds for supported compute platforms, and its verification step uses torch.cuda.is_available().^{[3]Reference 3Get Started.https://pytorch.org/get-started/locally/} For ordinary prebuilt use, a local toolkit and nvcc matter when you build PyTorch from source or compile custom CUDA extensions. Version strings don't need to match; the driver must support the runtime used by the package.

Run the driver check first, then inspect the Python environment. The shell command produces driver output, PyTorch version and runtime, availability, device name, capability, and architectures when CUDA is reachable:

The compute capability pair describes hardware features supported by a GPU generation. get_arch_list() reports architectures included in the PyTorch build. Read the first disagreement, not the most alarming line:

CUDA_VISIBLE_DEVICES changes the GPU indexes visible inside the process. If physical GPU 1 is the only visible device, PyTorch may call it cuda:0. Inspect tensor.device and device names inside the running process instead of assuming host indexes survive a mask.

Once the stack answers "yes," the miniature ticket batch still has to move. A visible GPU is only permission to try work; it isn't proof that tensors, memory, or timing are correct.

Move the ticket batch once

Return to the (4, 8, 16) ticket tensor. What changes during .to(device): its shape or its residence? Shape stays constant; residence changes from host memory to CUDA device memory.

On a discrete NVIDIA GPU, a host-to-device (H2D) transfer copies data from CPU memory to GPU memory. Device-to-host (D2H) is the return path. Keep model weights, batch tensors, activations, gradients, and optimizer state on the GPU through the hot part of a training step. Bring back the small value needed for reporting.

The runnable step uses random feature values so it needs no dataset download. It tests execution, not classifier quality. It averages eight token vectors into one vector per ticket, then scores three classes. Labels are 0 = answer, 1 = escalate, 2 = block.

nn.Linear(16, 3) multiplies each 16-feature vector by learned weights and adds a bias, producing three scores called logits. Cross-entropy turns those scores and the correct label into one error value, the loss. backward() computes how each weight affects that loss; those derivatives are gradients. Stochastic gradient descent (SGD) updates weights opposite the gradient, aiming to reduce the loss.^{[4]Reference 4Optimizing Model Parameters.https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html} The later PyTorch training lesson derives the loop; here, watch where its tensors live.

Predict its shape and device checks before running: (4, 8, 16) should become ticket vectors (4, 16) and logits (4, 3). On CUDA, model and batch should agree; on CPU, the same shape and loss checks still run.

On an accessible NVIDIA GPU, the first two lines should report True. A CPU-only machine reports False for the accelerator line but still validates tensor and training logic. If a shape fails, debug the reduction or linear layer first; a device error needs a different branch.

Device placement applies to every tensor an operation touches. Moving features but leaving labels on CPU still fails during a CUDA loss calculation. Deliberately leave the features on CPU in the next fragment. Predict the error on a CUDA machine, and the fallback message on a CPU-only machine:

The symptom is an error saying tensors or arguments aren't on the same device. Cause is placement, not shape. Move every participating tensor to the model's device before the forward or loss call. Synchronizing won't repair a tensor stored on the wrong device.

The miniature step can succeed while a training-sized batch still fails. Shapes and device placement answer whether an operation is legal; memory pressure decides whether the full forward and backward pass can coexist.

Count memory before the first big batch

The small step fits. Before increasing batch size, ask which values must survive between steps and which exist only while one kernel runs. CUDA exposes several storage levels, and they aren't interchangeable:

PyTorch allocates model tensors in device global memory. Optimized kernels decide when to reuse tiles through caches, shared memory, or registers. Calling .to("cuda") doesn't place a whole tensor in a register or in shared memory. An optimizer adds its own long-lived state during training.

The miniature batch itself contains 4 × 8 × 16 = 512 float32 values, only 2,048 bytes. A larger (32, 128, 768) float32 batch contains 32 × 128 × 768 × 4 = 12,582,912 bytes, exactly 12.0 MiB of input features. That input is only one term in a training step:

$$\text{training memory} \approx \text{weights} + \text{activations} + \text{gradients} + \text{optimizer state} + \text{temporary workspaces}$$

The example uses SGD without momentum, which doesn't keep per-parameter running statistics. To see why the optimizer choice changes memory, consider Adam, an optimizer that tracks two statistics for each parameter: a running average of gradients and a running average of squared gradients.^{[5]Reference 5Adamhttps://docs.pytorch.org/docs/2.13/generated/torch.optim.Adam.html}

For a simplified full-precision Adam budget, count 4 bytes for each weight, 4 for its gradient, and 8 for its two statistics. One billion parameters therefore require 16 billion bytes for these tensors when they're all present. Activations, temporary buffers, allocator overhead, and any extra parameter copies still sit outside this accounting.

Compute those two scales before touching a GPU. The first value is a batch input; the second is a parameter-related floor. The script takes the parameter count and feature shape as constants and prints both in human-readable units:

A model can load and still hit an out-of-memory (OOM) failure on its first forward or backward pass. Loading proves that current state fits. It doesn't prove that activations and gradients for a real batch fit. Adam's moment tensors are commonly created on the first optimizer step, so that step can fail even after backward succeeds.

For the running text batch, predict which change halves token positions: halve batch size, or halve sequence length? Both halve the simple B × T count, but attention storage responds differently. This small calculation makes that distinction visible:

The output counts token positions, not full memory. Our mean-pooling classifier has no attention matrix. A later transformer compares each token with every other token; storing those scores can require a tensor shaped (B, H, T, T), where B is batch size, H is the number of attention heads, and T is sequence length. Halving B halves that tensor; halving T quarters it. Memory-efficient attention can avoid storing the full score matrix, so this isn't a universal prediction of peak memory. Feature width D = 768 doesn't change the token-position count.

Mixed precision is another memory and throughput tool, but it isn't a safe synonym for calling .half() on everything. PyTorch's automatic mixed precision (AMP) uses torch.autocast(device_type="cuda") to choose an operation-specific dtype. FP16 (16-bit floating point) training may need torch.amp.GradScaler("cuda") to prevent small gradients from underflowing, and some models overflow in FP16.^{[6]Reference 6Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html} Confirm finite loss, finite gradients, and expected model quality. The dedicated mixed-precision lesson covers that workflow.

TensorFloat-32 (TF32) is a separate choice on Ampere and newer NVIDIA GPUs: eligible operations keep float32 tensor storage but use reduced-precision inputs internally, with FP32 accumulation. It can change numerical results without halving tensor memory. Current PyTorch documentation recommends fp32_precision controls such as torch.backends.cuda.matmul.fp32_precision = "ieee" or "tf32"; mixing these with legacy allow_tf32 controls isn't supported. Record the precision setting when comparing timings, and check accuracy before accepting a faster setting.^{[7]Reference 7CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

Memory pressure answers "does it fit?" It doesn't answer whether the GPU is moving data efficiently while it works. Access order supplies that missing piece.

Let access order decide bandwidth

Two kernels can read the same number of float32 values and still ask the memory system for very different work. On GPUs with compute capability 6.0 or newer, a warp's global-memory requests are serviced in 32-byte segments. Reading adjacent words lets the lanes share segments, an access pattern called coalescing. Reading a column in a row-major matrix can instead make neighboring lanes jump by a stride of 32 words, so each lane touches a different segment.

Keep eight lanes in the illustration so the addresses stay readable. For a full warp, the adjacent pattern needs four 32-byte transactions when the first word is aligned. The stride-32 pattern can need one transaction per lane. That's why a mathematically identical operation can land on opposite sides of a bandwidth limit.^{[8]Reference 8CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Shared memory can repair the global access pattern, but it introduces banks, independently accessible storage partitions. In the standard 32-bank layout for 32-bit words, the bank index is the word index modulo 32. For column zero of a 32 × 32 tile, lanes 0, 1, and 2 access word indexes 0, 32, and 64, all in bank 0. These are different addresses, so the accesses serialize. With a 32 × 33 allocation, the indexes become 0, 33, and 66, which map to banks 0, 1, and 2. Padding changes storage layout, not the logical matrix result. Multiple lanes reading the same address are a broadcast case, not this bank-conflict case.^{[8]Reference 8CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Use this compact table after you've inspected the access pattern. It names the first change worth testing, not a promise that one trick wins every kernel:

The table explains why "put it in shared memory" isn't a complete optimization. Shared memory helps when it removes global traffic or fixes an access pattern; bank conflicts can give that shortcut a new bottleneck.^{[8]Reference 8CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Memory layout answers where time can go. The next question is whether arithmetic or bytes set the upper bound.

Choose the next optimization with a bound

Suppose a kernel performs 1,000 floating-point operations (FLOPs) and transfers 4,000 bytes to or from device memory. That's 0.25 FLOPs/byte: little arithmetic for each byte fetched. This ratio is called arithmetic intensity:

$$I = \frac{\text{floating-point operations}}{\text{bytes moved}}$$

Double the bytes without changing the math and intensity falls to 0.125. The Roofline performance model combines this ratio with two hardware limits: arithmetic per second and bytes per second. Its upper bound is the smaller of the compute ceiling and bandwidth multiplied by intensity. Count bytes and bandwidth at the same memory level; a device-memory Roofline uses traffic beyond the caches, not every load instruction.^{[9]Reference 9Roofline: An Insightful Visual Performance Model for Multicore Architectureshttps://doi.org/10.1145/1498765.1498785[10]Reference 10Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/}

For a hypothetical device with a 100 GFLOP/s compute ceiling and 40 GB/s memory bandwidth, intensity 0.25 gives a bandwidth ceiling of 40 × 0.25 = 10 GFLOP/s. The bound is therefore 10, not 100 GFLOP/s. These invented numbers teach the units; they aren't a GPU benchmark.

Predict before tuning: when intensity is low and measured bandwidth is high, will adding more arithmetic make the kernel faster? Usually not. Reuse data, reduce bytes, or make requests coalesced first. When intensity is high but achieved floating-point rate stays far below the compute ceiling, inspect instruction choice, layout, and launch resources instead.

Occupancy is the fraction of an SM's maximum resident warps that are currently resident. A resident warp holds execution resources even while waiting for data. Having other ready warps helps the scheduler hide that wait, but 100% occupancy isn't a target by itself. Registers and shared memory are finite; using more of either can reduce the number of resident blocks. A kernel with lower occupancy can still win if its tiles reuse data or its threads do more useful work. Treat occupancy as a constraint and a clue, then let measured throughput and stall reasons decide.^{[8]Reference 8CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

NVIDIA's Assess, Parallelize, Optimize, Deploy (APOD) cycle gives the work a safe rhythm. At this level, combine its Parallelize and Optimize stages into one controlled change:

1. Measure (Assess): profile a representative shape and record end-to-end time, kernel time, output, and memory behavior.
2. Change (Parallelize, then Optimize): alter access order, reuse, dtype, launch shape, or transfer behavior because the evidence points there.
3. Verify, Deploy, and repeat: compare output with a trusted reference, rerun the same measurement, keep the change only when the whole step improves, then start the next pass from a fresh assessment.

Pick the tool from the question. Nsight Systems answers whether the CPU feeds the GPU and where transfers, kernels, streams, or host waits leave gaps. Nsight Compute answers what one kernel does with its launch resources, memory hierarchy, and compute or bandwidth ceilings. Compute Sanitizer answers whether custom CUDA work reads or writes illegal memory, races on shared memory, or violates synchronization. A speedup without the last check is only a faster way to trust a bad result.^{[11]Reference 11Nsight Systems User Guidehttps://docs.nvidia.com/nsight-systems/UserGuide/[10]Reference 10Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/[12]Reference 12Compute Sanitizerhttps://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/}

Run correctness checks before comparing numbers. Once the bound and the failure surface are named, the diagnosis table can point to a narrow next action instead of turning every slowdown into a memory tweak.

CUDA's default async launch makes timing another boundary that is easy to get wrong.

Measure asynchronous work honestly

PyTorch queues CUDA operations asynchronously by default. The CPU can finish a Python call while its GPU kernels are still waiting or running. Operations in the same CUDA stream keep their order, so later GPU work sees correct results without a host wait. A host-visible value or explicit synchronization makes the CPU wait.^{[7]Reference 7CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

Suppose a larger classifier's forward pass takes 40 ms on the GPU, while Python needs only 2 ms to queue it. These are illustrative numbers, not measurements of our miniature classifier. Predict what a host timer around the call reports: it measures the caller's elapsed time, not necessarily completed device work.

The naive timer stops at 2 ms while the GPU continues. Events record timestamps on a CUDA stream, an ordered queue of device work. Their interval includes work and any gaps between the events on that stream; it isn't automatically the sum of individual kernel durations. A synchronized host timer can also include Python overhead, so it needn't match event time exactly.

Warm up first because initial calls may include library setup or kernel selection. This script times ten matrix multiplications with CUDA events, or a host timer on CPU, and checks the result against a CPU reference. The timed region excludes input allocation, warmup, transfers, and the correctness check:

Common host waits include:

• loss.item() for a Python scalar
• tensor.cpu() before CPU analysis
• tensor.cpu().numpy() before NumPy work
• printing a CUDA tensor's values
• torch.cuda.synchronize()

Those operations aren't bugs. They become performance bugs when they sit inside a hot loop more often than reporting or correctness requires. A scalar needed once per full pass is different from a scalar pulled back every step.

Asynchrony also affects error location. A bad kernel may report its error on a later Python line that finally waits for the device. For one debugging reproduction, run with CUDA_LAUNCH_BLOCKING=1 to make CUDA calls synchronous and recover a more useful stack trace. Remove it before performance measurement because it changes execution behavior.^{[7]Reference 7CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

When the ticket step is slow, empty, or dead, start with one question: did the CPU fail to feed the GPU, did the GPU run work that doesn't fit, or did the measurement hide the wait? The diagnosis table turns that question into first actions.

Diagnose setup, memory, and throughput failures

nvidia-smi is a useful first observation, not a kernel profiler. Use it to confirm device visibility, process attachment, rough memory pressure, and utilization samples. Before comparing its memory number with PyTorch, predict which one includes reserved-but-unused allocator blocks. PyTorch's caching allocator can hold blocks for reuse, so nvidia-smi may show more memory than live tensors occupy.^{[2]Reference 2nvidia-smi documentationhttps://docs.nvidia.com/deploy/nvidia-smi/index.html[7]Reference 7CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

PyTorch separates two allocator views:

• torch.cuda.memory_allocated() counts memory occupied by live tensors.
• torch.cuda.memory_reserved() counts the larger pool managed by PyTorch's caching allocator.

torch.cuda.empty_cache() releases unused cached blocks for other applications. It doesn't free live tensors or shrink the workload's required memory. Cached blocks are normally reusable by the same job already; repeatedly clearing them can add allocation overhead. nvidia-smi also includes CUDA context and other allocations outside PyTorch's tensor allocator, so it needn't equal memory_reserved().^{[7]Reference 7CUDA semanticshttps://docs.pytorch.org/docs/2.14/notes/cuda.html}

Use symptom, cause, and next action together. Choose the first row that matches observed evidence instead of tuning every knob at once:

Pinned host memory can speed H2D copies. DataLoader(..., pin_memory=True) prepares CPU tensors in page-locked memory, and .to(device, non_blocking=True) can let the host continue without waiting for each transfer. This doesn't by itself overlap the copy with a later kernel in the same stream: their order is preserved. Actual copy/compute overlap requires suitable hardware, separate streams, and correct dependencies. Don't modify a pinned source buffer until its asynchronous copy finishes.^{[13]Reference 13A guide on good usage of non_blocking and pin_memory() in PyTorchhttps://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html}

The next snippet is a transfer-boundary fragment, not a full script. It assumes a CPU dataset yielding (features, labels) and a model already exist. Deriving the device from the model gives an explicit index such as cuda:0, not an unindexed cuda specification:

Measure before and after. Pinned memory and non-blocking copies help a transfer bottleneck; they won't fix a kernel that's already compute-bound or an OOM caused by live tensors. A faster copy doesn't change the amount of model state that must fit.

Trace the training step end-to-end

Use the complete training step as a small accelerator artifact. A note saying "GPU works" collapses setup, shape, placement, memory, and timing into one unverifiable claim. Record one piece of evidence for each boundary.

1. Predict shapes before running: batch (4, 8, 16), pooled vectors (4, 16), logits (4, 3).
2. Run nvidia-smi and the PyTorch environment check. Record driver, PyTorch runtime, device name, and availability separately.
3. Run the training step. Confirm model, features, labels, logits, and loss use one device until reporting.
4. Deliberately leave features on CPU while model is on CUDA. Capture the device-mismatch symptom, then restore .to(device).
5. Time ten matrix multiplies with CUDA events. Compare that result with a naive host timer around the same loop.
6. Write the terms that would make a larger run OOM: weights, activations, gradients, optimizer state, workspaces, and allocator overhead.

Compare expected output with this table, then correct the first boundary that fails. A later symptom can be downstream noise:

Check the reasoning without running code. Each answer should name a mechanism, not only repeat a command:

Carry shape, device, and timing forward

A CUDA training step has three contracts:

• Shape: each axis still means what the model expects.
• Device: every tensor used together lives on a compatible device.
• Time: measurements include the GPU work whose latency you intend to report.

The CPU prepares data and launches work. CUDA maps kernels onto grids, blocks, warps, and SMs. Device memory holds long-lived training state, while kernels use caches, shared memory, and registers for smaller working sets. Copies and synchronization are explicit costs, so placement, OOM diagnosis, and timing all follow from the same execution path.

If a run fails, ask which contract broke first: shape, device, or time. That question narrows a large CUDA surface into an observable next check.

The Mac lesson keeps those checks but changes the hardware picture: Apple silicon doesn't give you a separate VRAM pool copied over PCIe.