How to Use NVIDIA GPUs for AI Training

Fundamentals of NVIDIA GPU Architecture for Deep Learning

NVIDIA Graphics Processing Units (GPUs) serve as the standard computational workhorse for artificial intelligence (AI) and deep learning workloads. While central processing units (CPUs) are optimized for low-latency, sequential instruction execution across a modest number of cores, GPUs feature massively parallel architectures designed to process thousands of mathematical operations concurrently. Because artificial neural network training relies fundamentally on dense linear algebra—specifically general matrix multiplications (GEMM) and tensor convolutions—GPUs drastically reduce the training time of modern deep learning models from months to hours.

Modern NVIDIA GPU architectures, ranging from Volta and Ampere to Hopper and Blackwell, incorporate specialized hardware units alongside conventional Streaming Multiprocessors (SMs):

• Streaming Multiprocessors (SMs): The core compute clusters of the GPU. Each SM contains sub-units for handling arithmetic logic, dispatching warps (collections of 32 parallel threads), managing local registers, and caching shared data.
• CUDA Cores: Standard execution units dedicated to single-instruction, multiple-data (SIMD) floating-point (FP32, FP64) and integer (INT32) calculations.
• Tensor Cores: Specialized execution units engineered specifically to accelerate matrix multiply-accumulate operations in hardware. Unlike CUDA cores that execute one arithmetic operation per cycle per thread, Tensor Cores perform fused matrix operations ($D = A \times B + C$) over small matrix tiles in a single instruction, supporting mixed-precision formats such as FP16, BF16, TF32, and FP8. NVIDIA Tensor Cores
• High-Bandwidth Memory (HBM) and GDDR: Physical onboard memory. Data-center architectures utilize HBM (such as HBM3 or HBM3e) to achieve memory bandwidth exceeding several terabytes per second, mitigating the memory-wall bottleneck during massive weight and activation exchanges.
• High-Speed Interconnects (NVLink and NVSwitch): Direct GPU-to-GPU interconnect technology providing bidirectional bandwidth vastly superior to standard PCIe lanes, minimizing communication overhead during distributed, multi-GPU synchronization.

The NVIDIA Software Stack

Executing deep learning workloads on NVIDIA hardware relies on an integrated, multi-tier software ecosystem collectively known as NVIDIA CUDA-X AI. Deep Learning Software Understanding each tier ensures proper system installation and performance optimization.

+-------------------------------------------------------------+
|    Application / Framework Tier (PyTorch, JAX, TensorFlow)   |
+-------------------------------------------------------------+
|    Domain Libraries Tier (cuDNN, cuBLAS, NCCL, TensorRT)     |
+-------------------------------------------------------------+
|    Compute Driver & Platform Tier (CUDA Runtime, Driver)    |
+-------------------------------------------------------------+
|    Hardware Tier (NVIDIA GPUs: H100, A100, RTX Series)       |
+-------------------------------------------------------------+

Display Drivers and CUDA Driver

The base layer consists of the kernel-mode NVIDIA GPU driver, which communicates directly with the PCI bus and the physical GPU. It provides both low-level device control and the CUDA Driver API. System administrators must ensure that the installed driver version matches or exceeds the minimum driver version mandated by the CUDA Toolkit release in use.

CUDA Toolkit and Runtime

The CUDA (Compute Unified Device Architecture) platform provides the parallel computing platform and programming model. It exposes APIs that allow developers to write C, C++, or Fortran code targeted directly at the GPU. For high-level practitioners, the CUDA Toolkit delivers essential runtime components, dynamic link libraries, and development tools like the NVIDIA System Management Interface (nvidia-smi) and the NVIDIA CUDA Compiler (nvcc).

Acceleration Libraries

Rather than requiring practitioners to write manual matrix multiplication kernels in raw CUDA, NVIDIA maintains specialized, highly optimized mathematical libraries:

• cuBLAS: A GPU-accelerated implementation of the Basic Linear Algebra Subprograms (BLAS), handling vector and matrix multiplications.
• cuDNN: The CUDA Deep Neural Network library, containing tuned routines for common deep learning primitives such as forward and backward 2D/3D convolutions, pooling, activation functions, normalization (Batch Normalization, LayerNorm), and attention mechanisms. NVIDIA cuDNN: Powering Faster Deep Learning Models
• NCCL (NVIDIA Collective Communications Library): Highly optimized multi-GPU and multi-node collective communication routines (such as AllReduce, AllGather, Broadcast, and ReduceScatter) that automatically maximize throughput across NVLink and InfiniBand fabrics.

Environment Setup and Driver Verification

A properly configured environment prevents software mismatch issues between frameworks and low-level drivers. Most enterprise and academic training workflows standardize on Linux environments (such as Ubuntu 22.04/24.04 LTS or enterprise distributions).

Driver Installation and Verification

NVIDIA drivers can be installed via the distribution's package manager or via official NVIDIA runfiles. Once installed, verify that the driver and GPU are recognized by executing:

nvidia-smi

The output displays the installed driver version, supported CUDA runtime ceiling, attached GPUs, temperature, power draw, volatile memory usage, and running compute processes.

Framework Setup

Modern frameworks such as PyTorch bundle pre-compiled CUDA runtime libraries and cuDNN binaries directly inside their wheel packages, eliminating the need to install a standalone CUDA Toolkit system-wide unless compiling custom C++/CUDA extensions.

To install PyTorch with CUDA acceleration using pip:

# Example targeting CUDA 12.1 or CUDA 12.4
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

Verify GPU visibility and device access inside a Python shell:

import torch

print(f"CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"Device Count: {torch.cuda.device_count()}")
    print(f"Device Name: {torch.cuda.get_device_name(0)}")
    print(f"Current Device Index: {torch.cuda.current_device()}")

Single-GPU Training Pipeline

To harness the GPU for model training, data and operations must be explicitly allocated to GPU memory. In PyTorch, devices are designated via torch.device("cuda") or torch.device("cuda:0").

Core Workflow

A standard GPU training loop executes four steps:

1. Transferring Model Parameters: Moving network weights to the GPU device memory via .to(device).
2. Transferring Input Batches: Loading data from the host CPU memory into GPU memory before the forward pass.
3. Forward and Backward Passes: Computing loss on the GPU and invoking backpropagation (loss.backward()) to calculate gradients using cuDNN/cuBLAS kernels.
4. Parameter Update: Stepping the optimizer (optimizer.step()) to update the weights directly inside VRAM.

Implementation Example

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 1. Device selection
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. Synthetic dataset and DataLoader
X_data = torch.randn(10000, 128)
y_data = torch.randint(0, 2, (10000,))
dataset = TensorDataset(X_data, y_data)

# pin_memory speeds up CPU-to-GPU page-locked memory transfers
loader = DataLoader(
    dataset, 
    batch_size=64, 
    shuffle=True, 
    num_workers=4, 
    pin_memory=True
)

# 3. Model definition and device placement
class Classifier(nn.Module):
    def __init__(self, in_features, classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(),
            nn.Linear(256, classes)
        )

    def forward(self, x):
        return self.net(x)

model = Classifier(in_features=128, classes=2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

# 4. Training loop
model.train()
for epoch in range(5):
    for batch_idx, (inputs, targets) in enumerate(loader):
        # Asynchronously transfer data tensors to GPU
        inputs = inputs.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

Accelerating Workloads: Mixed Precision and Tensor Cores

By default, neural networks historically trained using 32-bit single-precision floating-point numbers (FP32). However, modern deep learning algorithms are robust against small numerical variations, enabling the use of reduced-precision representations.

Precision Formats

Automatic Mixed Precision (AMP)

Mixed precision training runs performance-critical operations (like convolutions and GEMMs) in lower precision (FP16 or BF16) to leverage Tensor Cores, while maintaining numerical stability by accumulating intermediate calculations and storing master weights in FP32. NVIDIA Tensor Cores

Because the dynamic range of FP16 is narrow, small gradients can encounter numerical underflow (rounding down to zero). PyTorch resolves this via a GradScaler, which multiplies the loss by a dynamic scale factor before backward propagation, shifting gradients into the representable range of FP16 before scaling them back down prior to optimizer updates. Automatic Mixed Precision package - torch.amp Is GradScaler necessary with Mixed precision training ... When training on architectures supporting BF16 (Ampere and later), gradient scaling is typically unnecessary because BF16 preserves the full dynamic range of FP32. Automatic Mixed Precision package - torch.amp

from torch.amp import autocast, GradScaler

scaler = GradScaler("cuda")

for inputs, targets in loader:
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)

    # Cast forward operations to FP16/BF16
    with autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # Scales loss, executes backward pass, unscales gradients, and steps optimizer
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Multi-GPU and Distributed Training Paradigms

When models or datasets exceed the capacity of a single GPU, training must be partitioned across multiple devices.

       Data Parallelism (DDP)                     Model Parallelism (FSDP / Megatron)
+-----------------------------------+    +-----------------------------------------------+
|  Full Model       Full Model      |    |  Layers 1-4         Layers 5-8        Layers 9-12 |
|   (GPU 0)          (GPU 1)        |    |   (GPU 0)             (GPU 1)           (GPU 2)   |
|      |                |           |    |      |                   |                 |      |
| Mini-batch 0    Mini-batch 1      |    |            Same Mini-batch Passed             |
|      +-------AllReduce------+     |    |            Across Partitioned Model           |
|         (Sync Gradients)          |    +-----------------------------------------------+
+-----------------------------------+

Distributed Data Parallel (DDP)

In Data Parallelism, the model is replicated across every GPU. Each device processes an independent slice of the global batch. At the completion of the backward pass, devices synchronize their gradients across the network using NCCL AllReduce before the optimizer updates weights simultaneously on all devices.

In PyTorch, torch.nn.parallel.DistributedDataParallel (DDP) is preferred over the older torch.nn.DataParallel because DDP spawns separate operating system processes per GPU, avoiding Python's Global Interpreter Lock (GIL) and maximizing NVLink throughput.

Fully Sharded Data Parallel (FSDP) and ZeRO

When a model's parameters, optimizer states, and activations exceed the VRAM of a single GPU, standard DDP causes out-of-memory (OOM) errors. Fully Sharded Data Parallel (inspired by the Zero Redundancy Optimizer, or ZeRO) eliminates redundancy across GPUs:

• ZeRO-Stage 1: Shards optimizer states across the GPU cluster.
• ZeRO-Stage 2: Shards both optimizer states and gradients.
• ZeRO-Stage 3 (Full FSDP): Shards optimizer states, gradients, and model parameters. Layers are gathered dynamically via all-gather operations immediately before computation and released immediately afterward.

Pipeline and Tensor Parallelism

For models containing tens or hundreds of billions of parameters:

• Tensor Parallelism (TP): Individual linear layers or attention projection matrices are mathematically split across multiple GPUs (e.g., Megatron-LM). Requires extreme interconnect speeds (NVLink).
• Pipeline Parallelism (PP): Sequential layers of a network are placed across sequential GPUs, passing activation tensors downstream and gradient tensors upstream.

Profiling, Diagnostics, and Bottleneck Elimination

Maximum hardware utilization requires ensuring that the GPU's SMs are not left idle waiting for data transfers or host-side computation.

Identifying Bottlenecks

A standard GPU execution cycle alternates between compute and wait phases:

• Input/Output (I/O) Bottlenecks: The GPU spends significant cycles idle (% GPU Utilization drops below 70%) because the CPU DataLoader cannot fetch, decode, and augment data quickly enough. Mitigate this by increasing num_workers, utilizing shared memory, using hardware-accelerated augmentation pipelines (such as NVIDIA DALI), and specifying pin_memory=True.
• Host-to-Device Latency: Transferring tensors from host RAM to device VRAM sequentially over PCIe introduces latency. Setting non_blocking=True on .to(device) allows memory copies to overlap asynchronously with CUDA compute streams.
• Memory Fragmentation: Frequent dynamic memory allocations cause fragmentation, triggering unexpected out-of-memory errors despite free reported VRAM. PyTorch manages this via an internal caching allocator, configurable via the environment variable PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.

Diagnostic Tools

1. nvidia-smi: Inspects point-in-time GPU utilization, VRAM allocation, and thermal throttling. Run continuously via watch -n 1 nvidia-smi.
2. PyTorch Profiler: Generates Chrome-compatible traces (chrome://tracing or Perfetto) detailing CPU execution, memory allocation, and CUDA kernel runtimes:
   pythonCopy

   with torch.profiler.profile(
       activities=[
           torch.profiler.ProfilerActivity.CPU,
           torch.profiler.ProfilerActivity.CUDA,
       ],
       record_shapes=True,
       profile_memory=True,
   ) as prof:
       # Run training step here
       pass
   prof.export_chrome_trace("trace.json")

3. NVIDIA Nsight Systems (nsys): Enterprise-grade system-wide profiling that visualizes CUDA kernel execution, CPU thread dispatching, OS runtime calls, and multi-GPU NCCL communication patterns.

Memory Management and Best Practices

To extract peak performance and prevent memory exhaustion during deep learning training, adhere to established architectural patterns:

• Set Gradients to None: Use optimizer.zero_grad(set_to_none=True) instead of writing zeros to memory. This avoids an explicit read-modify-write operation across parameter gradient buffers and frees allocation memory earlier during the backward pass.
• Use Activation Checkpointing (Gradient Checkpointing): Trading compute for memory, activation checkpointing discards intermediate activation tensors during the forward pass and recalculates them on demand during the backward pass. This reduces activation memory consumption from $O(N)$ (proportional to network depth) down to $O(\sqrt{N})$.
• Enable cuDNN Auto-Tuner: For networks with static input tensor dimensions, set torch.backends.cudnn.benchmark = True. This benchmarks several cuDNN convolution algorithms on the host architecture during the initial iteration and selects the fastest hardware-specific implementation.
• Align Tensor Dimensions: Tensor Core micro-architectures operate most efficiently when matrix dimensions ($M, N, K$) and layer batch sizes are multiples of 8 or 16 (for FP16/BF16) or 32 (for INT8/FP8). Non-aligned dimensions require padding, decreasing operational efficiency. NVIDIA Tensor Cores

What "Using an NVIDIA GPU for AI Training" Actually Means

Training a neural network on an NVIDIA GPU comes down to four things: a GPU with enough memory for your model, a driver and CUDA runtime that the framework can talk to, code that places tensors and model weights on the GPU device, and a numerical precision setting that lets the GPU's Tensor Cores do the matrix multiplications at full speed. In practice, the modern workflow for learning how to use NVIDIA GPUs for AI training is less about writing CUDA kernels and more about configuration: install a matching driver, install a CUDA-enabled build of PyTorch or JAX, call .to("cuda"), enable mixed precision, then make sure the GPU is not sitting idle waiting for data. Everything beyond that — multi-GPU sharding, FP8, interconnect topology — is scaling work that matters only once a single GPU is saturated.

The reason GPUs dominate this workload is structural. Training is dominated by dense matrix multiplication and elementwise operations over large tensors, which map cleanly onto thousands of parallel arithmetic units fed by very high-bandwidth memory. Since the Volta generation, NVIDIA GPUs have also carried Tensor Cores: fixed-function units that compute small matrix-multiply-accumulate operations in reduced precision. A GPU's advertised "AI TFLOPS" figure almost always refers to Tensor Core throughput at a reduced precision, not to general FP32 math, which is why a model trained in plain FP32 may leave most of the hardware's theoretical performance untouched.

Picking the Right Class of GPU

Memory capacity, not raw compute, is usually the binding constraint. A model that does not fit in VRAM cannot be trained at all without sharding or offloading, whereas a slow model merely takes longer. Roughly speaking there are three tiers:

Within the data center tier the generational jumps are large and matter for planning. The H100 SXM carries 80 GB of HBM3 at about 3.35 TB/s, the H200 moves to 141 GB of HBM3e, and the Blackwell-generation B200 is a dual-die part of 208 billion transistors whose published memory figure is given as 180 GB usable or 192 GB installed depending on the source, at roughly 7.7–8 TB/s. Its fifth-generation Tensor Cores are rated near 4,500 dense TFLOPS at both BF16 and FP8 — versus roughly 1,979 FP8 TFLOPS on H100 SXM — with 9,000 dense TFLOPS at FP4, a format Hopper has no native support for. Blackwell also supports Multi-Instance GPU partitioning, which splits one physical GPU into isolated slices, useful when many small training jobs must share expensive hardware.

At the top end, GPUs stop being individual devices. The GB200 Grace Blackwell Superchip pairs a 72-core Arm Neoverse V2 CPU with two Blackwell GPUs over a coherent link, and the GB200 NVL72 rack ties 72 GPUs together through NVSwitch so they present roughly 13.4 TB of pooled HBM3e and operate much like one very large accelerator. For a reader planning training work, the practical takeaway is that scale-up bandwidth inside a node or rack (NVLink, reported at 1.8 TB/s per GPU on Blackwell) is typically an order of magnitude higher than scale-out bandwidth between nodes, and that asymmetry drives nearly every parallelism decision later.

Getting the Software Stack Right

The stack has four layers, and version mismatches between them are the single most common reason a GPU "isn't detected."

1. Kernel driver — the only component that must be installed on the host.
2. CUDA runtime and toolkit — normally bundled inside the framework wheel or container; you rarely need a separate toolkit unless you compile custom kernels.
3. Libraries — cuDNN for convolutions and normalization, cuBLAS/cuBLASLt for GEMMs, NCCL for multi-GPU collectives.
4. Framework — PyTorch, JAX, or TensorFlow built against a specific CUDA major version.

The compatibility rules are more forgiving than they used to be. The CUDA driver is backward compatible, so applications built against an older toolkit keep working on newer drivers, and minor version compatibility lets an application built with any toolkit in a major family run on an older driver within a published range: CUDA 13.x requires driver release 580 or newer, CUDA 12.x requires 525 or newer, and CUDA 11.x requires 450 or newer. Note also that NVIDIA no longer bundles the driver with the toolkit — on Windows from CUDA 13.1 and on Linux from CUDA 13.4 — so the driver is now a separate install. These specifics shift between releases, so verify against the release notes for the version you are deploying.

For anything beyond a single laptop, containers are the better default. With the NVIDIA Container Toolkit installed, a container gets access to the host driver while carrying its own CUDA userspace, so you can run a CUDA 12 image and a CUDA 13 image on the same machine. The toolkit expects a Linux host, a reasonably modern container runtime, and a supported driver branch; recent versions also inspect libcuda.so metadata to decide when to activate in-container compatibility libraries, which reduces a class of confusing failures. One practical note from NVIDIA's own images: the default 64 MB shared-memory allocation in Docker is often too small for PyTorch DataLoader workers, so raising --shm-size is routinely necessary.

A sanity check before writing any training code:

nvidia-smi                     # driver version, GPU model, memory, running processes
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

Mixed Precision: The Highest-Leverage Setting

Running a model in pure FP32 is the most common way to leave performance on the table. Mixed precision keeps numerically sensitive operations in FP32 while running the heavy linear algebra in a 16-bit format, which cuts both runtime and memory footprint. In PyTorch this is handled by torch.amp, which maintains an op-level policy: matrix multiplications and convolutions autocast to the low-precision type, while reductions and loss functions that need FP32's dynamic range stay in FP32. You wrap only the forward pass and loss computation in the autocast context — running the backward pass inside it is