PyTorch on A100: setup, memory, and performance tuning

Running PyTorch on an NVIDIA A100 starts with a compatible CUDA build, a working driver, and a check of the GPU memory actually assigned to your process. After that, measure your workload before changing precision, batch size, or the number of GPUs. A faster chip cannot fix a blocked data loader or make an oversized model fit by itself.

This guide is for developers and machine learning engineers setting up an A100 for training or inference. It covers installation checks, mixed precision, memory limits, and common CUDA errors, then explains where a smaller GPU configuration may be enough.

Know which NVIDIA A100 you are using

The NVIDIA A100 is an Ampere data center GPU, available in 40 GB and 80 GB variants and different PCIe or SXM configurations. It supports Tensor Core operations including FP16, BF16, and TF32. Those specifications describe capabilities, not the speed of an arbitrary PyTorch model.

Check the provider's exact instance specification. Memory capacity, GPU count, CPU allocation, storage, and connections between GPUs can all affect the result. An A100 label alone does not establish an uptime commitment or a particular network topology.

Also check whether you have a full GPU or a Multi-Instance GPU (MIG) partition. A partition exposes only its assigned resources. If PyTorch reports less memory than the physical board contains, that can be the expected allocation rather than a broken installation.

Install PyTorch with a compatible CUDA build

Start in a separate virtual environment or an approved GPU container. Keep a record of the existing package versions before changing a working project. Installing into a new environment makes it easier to distinguish an installation problem from a dependency conflict in the application.

1. Run nvidia-smi in the environment where the job will execute. Confirm that the assigned GPU is visible and record the driver version.
2. Use the official PyTorch installation selector for your operating system, package manager, and supported CUDA build. Follow the requirements for that release rather than copying an old CUDA 11.x command from a forum.
3. Check that the host driver supports the selected CUDA runtime, using NVIDIA's driver compatibility guidance. A container still depends on the host driver and GPU access.
4. Run the diagnostic below using the same Python interpreter or notebook kernel as your training code.

The CUDA version shown by nvidia-smi, the build reported by torch.version.cuda, and the toolkit reported by nvcc --version describe different parts of the software stack. They need to be compatible; their version strings do not have to be identical.

Prebuilt PyTorch packages supply their CUDA runtime dependencies. A separate local CUDA toolkit is generally needed when building PyTorch or compiling custom CUDA extensions, not merely to run ordinary tensor operations. Check each extension's build instructions before adding another toolkit.

Verify your current PyTorch installation

This small diagnostic checks the visible device and executes a matrix multiplication. It uses GPU 0, which means the first GPU visible to this process, not necessarily the first physical GPU in the server.

import sys
import torch

print("Python:", sys.executable)
print("PyTorch:", torch.__version__)
print("CUDA build:", torch.version.cuda)
if not torch.cuda.is_available():
    raise RuntimeError("CUDA is unavailable in this Python environment")

device = torch.device("cuda:0")
properties = torch.cuda.get_device_properties(device)
print("GPU:", properties.name)
print("Capability:", torch.cuda.get_device_capability(device))
print("Visible memory (GiB):", round(properties.total_memory / 2**30, 1))
print("Build architectures:", torch.cuda.get_arch_list())

matrix = torch.randn(256, 256, device=device)
result = matrix @ matrix
torch.cuda.synchronize(device)
print("Matrix test:", tuple(result.shape))

An A100 has compute capability 8.0, commonly written as sm_80. The architecture list describes the PyTorch build; it does not certify every third-party extension. NVIDIA's Ampere compatibility guide explains how native GPU binaries and forward-compatible PTX affect whether a kernel can run.

If the diagnostic succeeds but your application fails, preserve that distinction. Test the failing operation or extension separately instead of repeatedly reinstalling the full environment. Save the complete error message, package versions, GPU allocation, and smallest input that reproduces the problem.

Use mixed precision deliberately

Automatic mixed precision selects lower precision for eligible operations while retaining higher precision where appropriate. For an A100 workload that supports BF16, torch.autocast provides a practical starting point. The current PyTorch AMP documentation replaces the older torch.cuda.amp interfaces.

The following is a self-contained example of one BF16 training step with synthetic inputs. It demonstrates the API pattern, not a useful model or a performance benchmark.

import torch

if not torch.cuda.is_available():
    raise RuntimeError("This example requires a CUDA GPU")
if not torch.cuda.is_bf16_supported():
    raise RuntimeError("This example requires BF16 support")

device = torch.device("cuda:0")
model = torch.nn.Linear(32, 8).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
inputs = torch.randn(16, 32, device=device)
targets = torch.randn(16, 8, device=device)

optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    predictions = model(inputs)
    loss = torch.nn.functional.mse_loss(predictions, targets)
loss.backward()
optimizer.step()
print("Loss:", loss.item())

BF16 generally does not need gradient scaling, while FP16 training commonly uses torch.amp.GradScaler("cuda"). Follow the AMP training examples when adapting an FP16 loop, especially with gradient accumulation or clipping. Neither format guarantees stable training: compare loss, validation quality, and nonfinite values against your baseline.

Check TF32 settings separately

TF32 can accelerate eligible operations on FP32 tensors, with a precision trade-off. It is a separate choice from BF16 autocast. Current PyTorch CUDA guidance recommends the newer fp32_precision controls for supported releases and warns against mixing them with older settings. Record the setting used in each experiment rather than assuming a default.

Plan for the whole memory footprint

A model's weight file is only a starting estimate. Training also needs gradients, optimizer state, activations, and temporary buffers. Inference has its own costs, including cached attention state for language models. Sequence length, image resolution, concurrency, and batch size can change the peak substantially.

Measure a representative step, including the optimizer update, before committing to a long run. Leave room for the largest expected input and periodic evaluation. There is no universal rule that 70% or 80% memory occupancy is the correct target for every model.

If you run out of memory, reduce the per-device batch size or input size first. Gradient accumulation can preserve a larger effective batch across several smaller steps, but it does not reduce the size of the model itself. Activation checkpointing saves memory by recomputing selected activations during backward, at an additional compute cost.

Distinguish replication from sharding

DistributedDataParallel (DDP) replicates a model across processes and synchronizes gradients. Adding GPUs with DDP does not turn their VRAM into one pool or make an oversized replica fit. Each process still needs enough memory for its local work.

Fully Sharded Data Parallel distributes model state across devices and introduces communication and configuration costs. Evaluate sharding when model state is the constraint. Confirm the actual interconnect and test scaling before assuming that twice the GPU count will halve the job time.

Measure PyTorch A100 performance end to end

Choose a metric that answers your question: examples per second for training, request latency at a stated concurrency for serving, or time to reach an agreed validation result. Keep the model, input sizes, precision, and software versions consistent when comparing configurations.

Warm up the workload before timing it. CUDA operations execute asynchronously, so a CPU timer around a launch can undercount the work. Use CUDA events with completion synchronization for device timing, or synchronize before and after a wall-clock measurement. Report separately whether loading data, compilation, and checkpoint writes are included.

Profile the input pipeline as well as GPU kernels. PyTorch's performance tuning guide covers data-loader workers, pinned memory, and gradient handling. Tune workers to the available CPU, storage, and transformations. A fixed worker count based only on GPU count can waste resources or slow the job.

For inference, put the model into evaluation mode and disable gradient recording where appropriate. If you test torch.compile, include its startup cost in short-job comparisons and check output quality. Change one variable at a time so an apparent improvement has an identifiable cause.

Troubleshoot common CUDA errors

Symptom	What to check	Useful next step
CUDA is unavailable	Driver access, container GPU allocation, CPU-only packages, and the active Python interpreter.	Compare the diagnostic with nvidia-smi inside the same job environment.
CUDA capability sm_80 is not compatible, or no kernel image is available	The installed PyTorch build and any custom extension used by the failing operation.	Install a supported build or rebuild the affected extension for compatible GPU targets.
CUDA out of memory	Actual full-GPU or MIG allocation, other processes, batch size, and retained tensors.	Measure the peak, reduce the per-device workload, then assess checkpointing or sharding.
Low throughput despite a working GPU	Data loading, CPU preprocessing, transfers, synchronization, and the chosen precision.	Profile a representative step and compare one change against the same baseline.

A CUDA error can surface after the operation that caused it. For a reproducible failure, use the debugging guidance for your PyTorch release and reduce the case before changing drivers. Treat a restart as a way to clear a failed process, not proof that the underlying issue is solved.

When an A100 is the right choice

An A100 is worth evaluating when your workload needs its memory capacity, supported precision modes, or the communication characteristics of a particular server. For smaller models, compare a configuration that already fits the job before paying for unused capacity. Published TOPS and FLOPS figures need precision and workload context; they cannot establish your application's training speed.

Compute with Hivenet is a separate option to evaluate. Its current GPU reference lists RTX 5090 instances with 32 GB of VRAM per GPU. Do not assume A100 availability. Check live presets, capacity, and pricing in the console, then compare cost per completed job with your measured runtime.

For a purchasing comparison, the GPU rental cost guide explains why memory fit and total running cost matter alongside the hourly rate. Budget for setup, persistent storage, data movement, and failed runs where those costs apply.

Before starting a long training run

Keep a short run record with the exact GPU allocation, environment versions, input shape, precision setting, peak memory, measured throughput, and validation result. Confirm that checkpoints can be saved and restored. A small reproducible run gives you a stronger basis for scaling than a hardware specification or an optimization score.