Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
This builds on the NumPy shape contracts. An access-ticket classifier labels each ticket answer, escalate, or block. A training-sized batch looks like (32, 128, 768). CUDA indexing is easier to see on a miniature of that same task: four tickets, eight token positions, 16 features, shape (4, 8, 16). Axis names stay the same as in the NumPy lesson. Residence and scale change.
The new constraint is physical. Those numbers live in CPU memory or GPU memory, and an operation fails if its tensors disagree about which.
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 kernels for tensor operations whose data lives on a CUDA device.[1]
If you train on a Mac with Apple silicon, continue through this lesson for the shared accelerator concepts, then use MPS & Metal for ML on Mac for its backend and memory differences.
Eight tokens onto two blocks
A CPU is tuned for low-latency control flow and varied work. A GPU is tuned for high throughput when many data items need similar arithmetic. Large matrix multiplications, convolutions, and attention scores are good GPU work. A tiny tensor, a branch-heavy Python loop, or a copy that happens every line may be faster on the CPU, because launch and transfer overhead can exceed the useful math.
| Part of a training step | CPU host job | GPU device job |
|---|---|---|
| data preparation | read, tokenize, pad, and assemble a batch | no work yet |
| device transfer | request a copy | receive tensor data in device memory |
| forward and backward | launch PyTorch operations | execute tensor kernels |
| reporting | request a Python number or file write | finish and return needed values |
That split shows up inside one kernel launch. A thread handles one logical slice of work. Threads are grouped into thread blocks, and all blocks launched 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]
Take the eight token positions in the miniature ticket batch. Two blocks with four threads each is a convenient teaching size. Each thread computes a unique global index:
For block 1, thread 2, the global index is 1 × 4 + 2 = 6, so that thread handles token position T6. This Python trace checks all eight assignments before the CUDA vocabulary grows:
1threads_per_block = 4
2blocks = 2
3
4for block_id in range(blocks):
5 positions = []
6 for thread_id in range(threads_per_block):
7 global_id = block_id * threads_per_block + thread_id
8 positions.append(f"T{global_id}")
9 print(f"block {block_id}: {positions}")1block 0: ['T0', 'T1', 'T2', 'T3']
2block 1: ['T4', 'T5', 'T6', 'T7']
An NVIDIA 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]
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.
A launch uses this stack, from the Python call down to one thread:
| Level | Meaning | Beginner question |
|---|---|---|
| PyTorch operation | a request such as x * 2 or x @ w | Is the tensor on CUDA? |
| kernel | GPU function launched for that operation | Is the work large enough to justify a launch? |
| grid | every block in one launch | How much work exists? |
| block | threads that run on one SM and can cooperate | Which slice of work stays together? |
| warp | 32 threads scheduled together | Are branches, or a short last warp, leaving lanes idle? |
| thread | one logical worker with its own index | Which data item does it handle? |
PyTorch's matrix-multiply kernels use optimized tiling that is more complex than one thread per token. The eight-position trace teaches CUDA's indexing and scheduling vocabulary without pretending to describe an optimized library kernel exactly.
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
Four software layers often get collapsed into "my CUDA version." Keep them separate. The NVIDIA System Management Interface command, nvidia-smi, reads driver and GPU state; nvcc is NVIDIA's CUDA compiler:
| Signal | What it tells you | What it doesn't prove |
|---|---|---|
nvidia-smi | NVIDIA driver can see a GPU. The CUDA Version header (documented as CUDA UMD Version; the older CUDA Version label is deprecated) is the latest CUDA version that driver supports | which CUDA toolkit is installed, or whether this Python has a CUDA-enabled PyTorch build |
torch.__version__ | installed PyTorch version and often its build suffix | whether GPU access works |
torch.version.cuda | CUDA runtime version used to build this PyTorch package, or None for a CPU-only build | whether driver and device are reachable |
torch.cuda.is_available() | current process can initialize and use CUDA | whether a particular workload fits or runs fast |
nvcc --version | version of local CUDA toolkit compiler, when installed | which runtime a prebuilt PyTorch package uses |
NVIDIA's nvidia-smi manual says that CUDA UMD Version is the latest CUDA version the driver supports, and that this is usually, but not always, the installed toolkit version.[2] It never proves 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] Pick the current command from that selector instead of copying an old wheel URL from a tutorial.
For ordinary prebuilt PyTorch use, start with a compatible NVIDIA driver and the selected PyTorch package. A local toolkit and nvcc become relevant when you build PyTorch from source or compile custom CUDA extensions. Version strings don't need to be identical; the driver must support the runtime used by the package.
Run the driver check first, then inspect the Python environment:
1nvidia-smi
2python3 - <<'PY'
3import torch
4
5print("PyTorch:", torch.__version__)
6print("PyTorch CUDA runtime:", torch.version.cuda)
7print("CUDA available:", torch.cuda.is_available())
8
9if torch.cuda.is_available():
10 print("device:", torch.cuda.get_device_name(0))
11 print("compute capability:", torch.cuda.get_device_capability(0))
12 print("build architectures:", torch.cuda.get_arch_list())
13PYThe compute capability pair describes hardware features supported by a GPU generation. get_arch_list() reports architectures included in the PyTorch build. Read failures from the first layer that disagrees:
| Observed state | Likely boundary | Next check |
|---|---|---|
nvidia-smi fails | driver, hardware, or container access | driver installation and device exposure |
nvidia-smi works, torch.version.cuda is None | CPU-only PyTorch package | reinstall from official CUDA selector |
PyTorch has a CUDA runtime, but availability is False | driver compatibility or process visibility | driver support, container GPU access, CUDA_VISIBLE_DEVICES |
availability is True, but kernel reports unsupported architecture | PyTorch build is too old for GPU | install a build that includes current GPU architecture |
availability is True and device name is correct | basic stack works | test placement, memory, and timing |
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.
Move the ticket batch once
Return to the (4, 8, 16) ticket tensor. Shape stays constant during a device transfer. Residence changes from host memory to CUDA device memory.
On a discrete NVIDIA GPU, host-to-device (H2D) means copying 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.
This runnable step stays small enough for any CUDA card, and it still exercises the full hot path. It averages eight token vectors into one 16-feature vector per ticket, predicts the three ticket classes, computes loss, runs backward, and updates weights. Labels are 0 = answer, 1 = escalate, 2 = block.
1import math
2
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6
7torch.manual_seed(7)
8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
10model = nn.Linear(16, 3).to(device)
11optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
12
13cpu_batch = {
14 "token_features": torch.randn(4, 8, 16),
15 "labels": torch.tensor([0, 2, 1, 0]),
16}
17batch = {name: tensor.to(device) for name, tensor in cpu_batch.items()}
18
19optimizer.zero_grad()
20ticket_vectors = batch["token_features"].mean(dim=1)
21logits = model(ticket_vectors)
22loss = F.cross_entropy(logits, batch["labels"])
23loss.backward()
24optimizer.step()
25
26model_device = next(model.parameters()).device
27reported_loss = loss.detach().cpu().item()
28print("selected accelerator:", device.type == "cuda")
29print("model and batch agree:", model_device == batch["token_features"].device)
30print("ticket vectors:", tuple(ticket_vectors.shape))
31print("logits:", tuple(logits.shape))
32print("finite loss:", math.isfinite(reported_loss))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 the tensor and training logic. In either case, ticket vectors must be (4, 16) and logits must be (4, 3).
Device placement applies to every tensor an operation touches. Moving features but leaving labels on CPU still fails during a CUDA loss calculation. This deliberate failure uses a real mixed-device forward pass when CUDA exists and explains the validation limit otherwise:
1import torch
2import torch.nn as nn
3
4if torch.cuda.is_available():
5 cuda_model = nn.Linear(16, 3).to("cuda")
6 cpu_features = torch.randn(4, 16)
7 try:
8 cuda_model(cpu_features)
9 except RuntimeError:
10 print("caught a real CPU/CUDA device mismatch")
11 else:
12 raise AssertionError("expected a mixed-device forward pass to fail")
13else:
14 print("CUDA unavailable: real mixed-device failure wasn't executed")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.
The miniature step can succeed and a training-sized batch can still fail. The missing piece is how much extra memory a real forward and backward pass asks for.
Count memory before the first big batch
CUDA exposes several storage levels. They aren't interchangeable:
| Storage | Scope | Typical training data |
|---|---|---|
| host RAM | CPU process | dataset objects, decoded examples, CPU batches |
| device global memory, often called video RAM (VRAM) or high-bandwidth memory (HBM) | all SMs on one GPU | weights, activations, gradients, optimizer state |
| L2 and L1 caches | GPU hardware | recently accessed device data |
| shared memory | threads in one block | tiles reused inside a kernel |
| registers | one thread | counters, addresses, and small working values |
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. The training-sized (32, 128, 768) float32 batch that the Mac lesson continues is 32 × 128 × 768 × 4 = 12,582,912 bytes, exactly 12.0 MiB of input features. Training memory still grows quickly because input is one term among several:
For a simplified full-precision Adam budget, one parameter may need 4 bytes for its weight, 4 for its gradient, and 8 for Adam's two running statistics. One billion parameters therefore have a parameter-related floor of 16 billion bytes. Activations, temporary buffers, allocator overhead, and any extra parameter copies still sit outside that floor.
1parameters = 1_000_000_000
2bytes_per_parameter = 4 + 4 + 8
3total_bytes = parameters * bytes_per_parameter
4feature_bytes = 32 * 128 * 768 * 4
5
6print(f"bytes per parameter: {bytes_per_parameter}")
7print(f"parameter-related floor: {total_bytes / (1024 ** 3):.2f} GiB")
8print(f"training-sized ticket features: {feature_bytes / (1024 ** 2):.1f} MiB")
9print("activations and temporary work: add more memory")1bytes per parameter: 16
2parameter-related floor: 14.90 GiB
3training-sized ticket features: 12.0 MiB
4activations and temporary work: add more memoryA model can load and still hit an 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.
For the running text batch, two first-line levers reduce activation work:
1batch_size = 32
2sequence_length = 128
3baseline = batch_size * sequence_length
4
5for name, batch, tokens in [
6 ("baseline", 32, 128),
7 ("half batch", 16, 128),
8 ("half length", 32, 64),
9]:
10 share = batch * tokens / baseline
11 print(f"{name:11s}: {share:.0%} of token positions")1baseline : 100% of token positions
2half batch : 50% of token positions
3half length: 50% of token positionsHalving sequence length can reduce attention-score storage faster because a score tensor has two token axes, commonly shaped (B, H, T, T). Halving batch size cuts its B term in half; halving T cuts its T × T term to one quarter. Later attention lessons derive that shape in full. Feature width D = 768 doesn't change that 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.[4] Confirm finite loss, finite gradients, and expected model quality. The dedicated mixed-precision lesson covers that workflow.
Ampere and newer NVIDIA GPUs can run some matmuls in TensorFloat-32 (TF32). Since PyTorch 1.12, torch.backends.cuda.matmul.allow_tf32 defaults to False, so a new install doesn't silently trade IEEE FP32 matmul accuracy for TF32 throughput. Enable it only after you've measured the speedup and checked quality.[5]
Memory pressure answers "does it fit?" It doesn't answer "how long did the GPU actually work?" CUDA's default async launch makes that second question 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.[5]
Suppose the ticket classifier's forward pass takes 40 ms on the GPU, but Python needs only 2 ms to queue it. A host timer around the call can report about 2 ms. That number measures launch time, not completed work.
Warm up first because initial calls may include library setup or kernel selection. Then use CUDA events for device time, or synchronize around a host timer. This script uses events on CUDA and a normal completed CPU timer as a fallback:
1import time
2
3import torch
4
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6x = torch.randn(256, 256, device=device)
7
8for _ in range(3):
9 result = x @ x
10
11if device.type == "cuda":
12 torch.cuda.synchronize()
13 start = torch.cuda.Event(enable_timing=True)
14 end = torch.cuda.Event(enable_timing=True)
15 start.record()
16 for _ in range(10):
17 result = x @ x
18 end.record()
19 torch.cuda.synchronize()
20 elapsed_ms = start.elapsed_time(end)
21else:
22 start_time = time.perf_counter()
23 for _ in range(10):
24 result = x @ x
25 elapsed_ms = (time.perf_counter() - start_time) * 1000
26
27print("result shape:", tuple(result.shape))
28print("completed timing is positive:", elapsed_ms > 0)Common host waits include:
loss.item()for a Python scalartensor.cpu()before CPU analysistensor.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.
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.[5]
When the ticket step is slow, empty, or dead, which boundary failed first?
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. PyTorch's caching allocator can hold unused blocks for reuse, so nvidia-smi may show more memory than live tensors occupy.[2][5]
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 increase memory available to the same PyTorch job, because that job could already reuse its cached blocks.[5]
Use symptom, cause, and next action together:
| Symptom | Likely cause | First action |
|---|---|---|
torch.cuda.is_available() is False | driver, visibility, or PyTorch build mismatch | compare nvidia-smi, torch.version.cuda, and process visibility |
| forward says tensors are on different devices | model and one batch field disagree | move every tensor used by model or loss to model device |
| model loads, backward OOMs | activations, gradients, optimizer state, or workspace exceed free memory | lower per-step batch size, then sequence length; read allocation size in OOM message |
| effective batch must stay large | smaller steps change optimization batch | accumulate gradients across several smaller steps and scale loss correctly |
loss or gradients become not-a-number (NaN) under FP16 | numerical overflow or invalid mixed-precision path | disable AMP to reproduce, then use autocast and scaling with finite-value checks |
| host timer says a kernel took almost zero time | CPU timed enqueue only | warm up and use CUDA events or explicit synchronization |
| memory bar is high but examples per second are low | allocation isn't utilization; data, copies, sync, or tiny kernels may stall | inspect dataloading and host waits before blaming matrix kernels |
| GPU utilization repeatedly falls to zero | input pipeline can't feed device steadily | profile data loading, preprocessing, and transfer cadence |
| CUDA error points at an innocent later line | asynchronous error surfaced at next wait | reproduce once with CUDA_LAUNCH_BLOCKING=1 |
Pinned host memory can speed H2D copies. DataLoader(..., pin_memory=True) returns batches in page-locked memory, and .to(device, non_blocking=True) lets the host continue without waiting for each transfer. Pinning uses a limited host-RAM pool, and real copy/compute overlap also depends on stream scheduling.[5]
The next snippet is a transfer-boundary fragment, not a full script. It assumes dataset and device already exist:
1from torch.utils.data import DataLoader
2
3loader = DataLoader(dataset, batch_size=32, pin_memory=True)
4
5for features, labels in loader:
6 features = features.to(device, non_blocking=True)
7 labels = labels.to(device, non_blocking=True)
8 assert features.device == device and labels.device == device
9 # Forward, backward, and optimizer work stay on device.Measure before and after. Pinned memory and non-blocking copies help a transfer bottleneck; they won't fix a kernel that is already compute-bound or an OOM caused by live tensors.
Preflight one training step
Use the complete training step as a small accelerator artifact. Record evidence for each boundary instead of writing "GPU works."
- Predict shapes before running: batch
(4, 8, 16), pooled vectors(4, 16), logits(4, 3). - Run
nvidia-smiand the PyTorch environment check. Record driver, PyTorch runtime, device name, and availability separately. - Run the training step. Confirm model, features, labels, logits, and loss use one device until reporting.
- Deliberately leave features on CPU while model is on CUDA. Capture the device-mismatch symptom, then restore
.to(device). - Time ten matrix multiplies with CUDA events. Compare that result with a naive host timer around the same loop.
- Write the terms that would make a larger run OOM: weights, activations, gradients, optimizer state, workspaces, and allocator overhead.
Compare your expected output with this table, then correct the first boundary that fails:
| Evidence | Pass condition | If it fails |
|---|---|---|
| shape trace | (4, 8, 16) → (4, 16) → (4, 3) | revisit reduction axis or linear-layer input width |
| device trace | every training tensor matches model device | move missing batch field before operation that uses it |
| timing trace | event reports completed GPU work | add warmup and synchronization after end event |
| memory trace | full training footprint is named | add activations, gradients, optimizer state, and temporary work |
Check the reasoning without running code:
Why can model(x) return to Python before its CUDA kernels finish?
Answer
PyTorch queues CUDA work asynchronously. Python can continue after enqueueing while the device executes operations in stream order.
Why can a model fit in device memory and then OOM during backward?
Answer
Model loading accounts for current state, but training also needs activations, gradients, optimizer state, temporary workspaces, and allocator overhead.
Why can nvidia-smi show more memory than memory_allocated()?
Answer
PyTorch's caching allocator can reserve unused blocks for fast reuse. memory_allocated() counts live tensors, while nvidia-smi can reflect the larger reserved pool.
What does the CUDA value printed by nvidia-smi mean?
Answer
It's the latest CUDA version the installed driver supports (CUDA UMD Version; the older CUDA Version label is deprecated). It isn't proof that the same toolkit is installed or that this PyTorch environment can use CUDA.
Three contracts for the next lesson
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 the 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.
The Mac lesson keeps those contracts on the same access-ticket classifier, then changes the hardware picture: Apple silicon doesn't give you a separate VRAM pool copied over PCIe.