Plan a reproducible PyTorch run

Prove one complete step before scaling

Lock the model and dataset revisions, dependency versions, training configuration, random seeds and container digest. Run model loading, a complete training step, evaluation, checkpoint save and checkpoint restore on one accelerator.

Use the same sequence length and micro-batch settings intended for the longer run. A tiny pilot that omits optimizer state or uses shorter sequences can hide the memory problem you are trying to measure.

Measure a representative optimization step

Inside your existing training program, surround a representative step with a peak-memory measurement. The example assumes model, batch and optimizer have already been initialized on the GPU. Include the optimizer step because some state is created lazily.

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

This reports peak tensor allocation, not the complete device footprint. Also inspect reserved memory, runtime overhead and the system’s device-memory report.

Choose DDP or sharding for the right reason

DistributedDataParallel keeps a model replica on each process and synchronizes gradients. It is useful when the full training state fits on each GPU and you want to process more data in parallel; it does not merge VRAM into one contiguous pool.

Fully Sharded Data Parallel can distribute parameters, gradients and optimizer state. Use it when memory pressure justifies sharding, and account for extra communication and distributed-checkpoint complexity.

Mixed precision and activation checkpointing are separate choices. Validate numerical behavior for lower precision; measure the compute cost of recomputing activations when checkpointing saves memory.

Launch a single-node pilot

The training script must initialize distributed execution and assign one process to each GPU using the local-rank information. Match the process count to the visible GPUs. The following is a launch example for a distributed-aware script, not a complete training implementation.

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

Run this only on a provisioned node. The local AnchorGPU demo does not execute training or allocate remote GPUs.

Select hardware from the pilot result

Use A100 as a lower-priced 80 GB CUDA baseline, and test H100 if the workload has a useful Hopper-specific path. H200 offers a larger CUDA memory configuration. MI300X offers more per-accelerator capacity in this catalog, provided the full application is qualified on ROCm.

For AMD, check custom extensions, kernels, package versions and the ROCm-tested PyTorch image as a complete set. Framework-level compatibility does not validate every optional operation.

Test restore, not just save

Record model, optimizer, scheduler, training step, configuration and mixed-precision scaler state where relevant. Save a checkpoint, terminate the process, then restore and resume before committing to the long run.

Use a checkpoint method appropriate to distributed state. Test that the saved artifacts can be read independently of the original process, and that the evaluation output remains within the expected tolerance.

Budget for the whole pipeline

Include downloads, preprocessing, data loading, validation, checkpointing, export and copying outputs off the node. Measure throughput together with memory, data-loading stalls, communication time, checkpoint duration and validation quality.

A 7-day term is useful for a bounded compatibility and profiling pass. Choose a 30-day reservation only after you can estimate the useful work and the operational buffer. More GPUs help only when additional computation outweighs communication and input-pipeline limits.