How do we do multi-GPU fine-tuning with DeepSpeed or FSDP?

Multi-GPU fine-tuning becomes necessary once a model's parameters, gradients and optimizer states no longer fit in a single GPU's memory, which happens for full fine-tuning of models above roughly ten billion parameters, and both DeepSpeed and PyTorch's native FSDP solve this by sharding those components across GPUs instead of replicating them on every device. DeepSpeed's ZeRO optimizer offers staged sharding, with ZeRO-1 sharding only optimizer states, ZeRO-2 adding gradient sharding, and ZeRO-3 sharding the model parameters themselves, letting teams pick the level of memory savings against the added communication overhead each stage introduces. FSDP achieves a similar effect natively within PyTorch without a separate library dependency, and it has become the more common default for teams already standardized on PyTorch tooling. Both approaches benefit substantially from fast GPU interconnects, since sharding introduces frequent cross-GPU communication, which is why multi-node full fine-tuning runs are typically paired with InfiniBand networking rather than standard Ethernet. Gradient checkpointing and mixed precision training are usually combined with either framework to push memory savings further. Nanobase AI, an NVIDIA Inception Program member, configures DeepSpeed and FSDP training clusters, including GPU Operator and InfiniBand setup, for clients running full fine-tuning at scale.

Why single-GPU memory runs out before compute does

Full fine-tuning of a model requires holding the model's parameters, gradients, and optimizer states (for Adam, typically two additional copies of each parameter for momentum and variance) simultaneously in memory, which for a 70B-class model in mixed precision adds up to well beyond what any single GPU, including an H200 at 141 GB, can hold once activations are added. This is a memory problem before it is a compute problem: the GPUs involved could easily handle the arithmetic for a 70B model, but no single card has enough HBM to hold the full training state, which is exactly what sharding across multiple GPUs solves. LoRA and QLoRA sidestep this by training a tiny fraction of parameters, which is why they remain single-GPU feasible at sizes where full fine-tuning requires a cluster.

Comparing ZeRO stages and FSDP

ZeRO-3 and FSDP's full-shard mode achieve similar memory savings through a similar mechanism, sharding every major component of training state across GPUs and reconstructing full parameters only transiently during each layer's forward and backward pass, which is why the choice between them often comes down to which integrates more smoothly with your existing training framework rather than a fundamental capability difference.

A concrete setup sequence

Sizing memory need before choosing a sharding stage prevents both under-provisioning the run and over-sharding it into unnecessary communication overhead.

1. Calculate total training memory need first: roughly 2 bytes per parameter for weights in BF16, plus 2 bytes for gradients, plus 8-12 bytes per parameter for Adam optimizer states, before any sharding is applied.
2. Choose the lowest ZeRO stage (or FSDP shard granularity) that fits your available GPU count and per-GPU memory, since more aggressive sharding adds communication overhead that can reduce training throughput even as it enables the run.
3. Configure high-bandwidth interconnect (NVLink within a node, InfiniBand across nodes) as a prerequisite, not an afterthought, since ZeRO-3 and full FSDP sharding are communication-heavy and will bottleneck badly on standard Ethernet at any meaningful scale.
4. Enable gradient checkpointing alongside sharding for large models, trading recomputation compute for further activation memory savings, which is often necessary even after sharding parameters and optimizer states.
5. Benchmark actual throughput at your chosen configuration before committing to a full training run, since the right ZeRO stage or shard granularity for a specific model and GPU count is easier to determine empirically than to calculate precisely in advance.

A minimal DeepSpeed ZeRO-3 config sketch

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"},
    "overlap_comm": true
  },
  "bf16": {"enabled": true},
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8
}

CPU offloading for optimizer states, shown above, trades additional GPU-to-CPU transfer time for further GPU memory headroom, which is worth enabling when a training run does not fit even after full ZeRO-3 sharding across the available GPU count.

When to reach for this versus staying single-GPU

Multi-GPU sharded training is necessary specifically for full fine-tuning of large models or for LoRA/QLoRA fine-tuning of very large models (a 70B-class LoRA run can still exceed single high-memory GPU capacity depending on batch size and sequence length). For most LoRA and QLoRA fine-tuning of models up to roughly 13B parameters, a single high-memory GPU is sufficient, and reaching for DeepSpeed or FSDP unnecessarily adds real setup and debugging complexity without a corresponding benefit. Understanding which GPU is needed to fine-tune a 70B model helps calibrate whether a given project genuinely needs multi-GPU sharding or can stay on a simpler single-GPU setup.

Frequently asked questions

Is DeepSpeed or FSDP better for fine-tuning?

Both achieve similar memory savings through similar sharding mechanisms; DeepSpeed offers more granular staged control (ZeRO-1 through ZeRO-3) and mature CPU/NVMe offloading options, while FSDP is natively integrated into PyTorch with less external dependency overhead. Many teams choose based on which is already integrated into their training framework of choice.

Do we need InfiniBand for multi-GPU fine-tuning, or is Ethernet enough?

For ZeRO-3 or full FSDP sharding across multiple nodes, high-bandwidth interconnect like InfiniBand or NVLink matters significantly, since these approaches communicate parameter shards frequently during training. Within a single node with NVLink-connected GPUs, standard networking between nodes becomes less critical.

Can we use ZeRO or FSDP with LoRA instead of full fine-tuning?

Yes, and it can still help for very large base models even when only training LoRA adapters, since the frozen base model's parameters still need to be held in memory during the forward pass, though the memory savings are less dramatic than for full fine-tuning since gradients and optimizer states apply only to the much smaller adapter parameters.

How Nanobase AI helps

Nanobase AI configures and operates multi-GPU fine-tuning clusters using DeepSpeed and FSDP on Kubernetes GPU Operator or Slurm-managed infrastructure, sizing the GPU count and interconnect to the specific model and training method rather than over- or under-provisioning.

Ready to discuss your project? Contact Nanobase AI or email hello@bumu.tech.