PyTorch DDP for Fine-Tuning Gemma 4 on Multiple GPUs

Some time ago I built a small Megatron Parallelism Visualizer. It shows the three classic ways to split a training job across GPUs: tensor parallelism, pipeline parallelism, and data parallelism. Tensor and pipeline parallelism get most of the attention, because they make it possible to train a 70B-parameter model at all. Most fine-tuning jobs that I see do not need either of them. They need the plain one.

This post is about the plain one, PyTorch DistributedDataParallel (DDP). The target is concrete: LoRA fine-tuning of Gemma 4 on one multi-GPU machine. The post explains how DDP moves gradients and gives a complete torchrun script that you can adapt. Then it calculates the memory budget for the Gemma 4 12B model. Last, it lists the failure modes that turn a multi-GPU job into a silent hang.

Table of contents

• When DDP is the right first tool
• Gemma 4 facts that matter for training
• How DDP works
• Batch size and learning rate
• Memory budget for Gemma 4 12B with LoRA
• The training script
• Launch with torchrun
• Run it on Modal
• Failure modes that hang or corrupt a DDP job
• NCCL and debugging environment variables
• Measure throughput and scaling

When DDP is the right first tool

The visualizer’s README gives one line to each of the three strategies:

• Tensor parallelism (TP) splits individual weight matrices across GPUs. Each GPU computes a partial matmul. Then the GPUs communicate to put the result together again.
• Pipeline parallelism (PP) puts different layers on different GPUs and streams micro-batches through them. The cost is pipeline “bubble” time.
• Data parallelism (DP) gives every GPU a full copy of the model and a different slice of the batch. Then an all-reduce averages the gradients.

Megatron-Core combines all three (“3D parallelism”), because at frontier scale no single strategy is enough. To fine-tune one model on one eight-GPU node is a different problem. The decision depends on whether one replica of the training state fits on one GPU.

flowchart TD
  Fit{"One replica fits<br/>on one GPU?"}
  Fit -->|Yes| DDP["DDP<br/>copy model, split data"]
  Fit -->|No| Shard{"Only optimizer state<br/>and grads too big?"}
  Shard -->|Yes| FSDP["FSDP / ZeRO<br/>shard the state"]
  Shard -->|No| TPPP["Tensor / pipeline<br/>parallelism + DP"]
  DDP -.-> Shrink["Almost fits? LoRA, QLoRA,<br/>checkpointing, shorter seqs"]

With LoRA, the answer for most Gemma 4 sizes is yes. The frozen base weights stay in bf16, and the trainable adapter has tens of millions of parameters. Gradient checkpointing keeps the activations bounded. When a replica fits, DDP is the simplest correct way to use more GPUs. Nothing inside the model changes, the only collective is a gradient all-reduce, and every rank runs the same code.

When a replica does not fit, you have other options. Each option solves a different problem:

For full fine-tuning of Gemma 4 12B, the table points to FSDP, not DDP. The memory section shows why. For LoRA on the same model, DDP is the right tool.

Gemma 4 facts that matter for training

Google released Gemma 4 in four sizes on March 31, 2026, and added a 12B model on June 3, 2026. The table shows the details that matter for a training plan. The data comes from the Gemma 4 model card and the Gemma 4 model overview:

The primary sources also give these points:

• License. Google releases the weights under Apache 2.0. The model card, the Gemma 4 launch post, and the Hugging Face repository for gemma-4-12B-it all state this.
• Model IDs. The base checkpoints are google/gemma-4-E2B, google/gemma-4-E4B, google/gemma-4-12B, google/gemma-4-26B-A4B, and google/gemma-4-31B. The instruction-tuned variants add -it to the end.
• Tooling. The launch post lists day-one support in Hugging Face Transformers, TRL, and many other tools. Google’s own text fine-tuning guide for Gemma uses Transformers, PEFT, and TRL. The guide installs transformers>=5.10.1 and peft>=0.19.0. It trains a LoRA adapter (QLoRA in that guide) on a single 16 GB T4.
• Architecture. All sizes use a 262K vocabulary. They interleave sliding-window attention layers with global attention layers. The 12B model is encoder-free. Image and audio inputs go directly into the decoder, not through separate encoders.

I also read the published checkpoint headers and configs for google/gemma-4-12B-it, because a few details there change the training plan:

• 11.96B parameters, all stored in BF16. Of these, 11.91B are under model.language_model. The input embedding alone is 262144 × 3840, about 1.0B parameters. It is tied to the output head.
• Eight of the 48 layers are global attention (every sixth layer). Those layers set attention_k_eq_v, so they have a k_proj but no v_proj. A LoRA config that lists v_proj still works. In those eight layers, it has nothing to attach to.
• LoRA at rank 16 on every attention and MLP projection of the text decoder gives 65.6M trainable parameters. That is about 0.55% of the model. I counted it from the tensor shapes.

For this post I use google/gemma-4-12B-it. It is big enough for the memory budget to matter. It is also small enough for a bf16 LoRA replica to fit on a single 48 GB or 80 GB GPU. In that case, DDP is the right tool.

How DDP works

The visualizer’s data-parallel simulation, simulate_data_parallel in simulator.py, gives the mental model. Every GPU runs forward and backward on its own slice of the batch, for all gradient-accumulation steps. Then one all-reduce averages the gradients, so every replica applies the same update. The dataParallel client simulator in simulator.js animates the same sequence.

Real DDP has the same shape, with one refinement that the simulation leaves out. The steps below show the life of a DDP job. They follow the PyTorch DDP design notes and the DistributedDataParallel API reference.

1. One process per GPU

torchrun starts one Python process per GPU and sets RANK, LOCAL_RANK, and WORLD_SIZE in each process’s environment. Each process pins itself to its GPU and joins a process group. With the NCCL backend, the collectives run on the GPUs. In a node, they use NVLink or PCIe. Across nodes, they use the network.

2. Replicas start identical

When you wrap a model in DistributedDataParallel, the constructor broadcasts parameters and buffers from rank 0 to every other rank. All replicas start from the same weights, including newly initialized LoRA matrices. DDP also registers an autograd hook on every parameter that requires a gradient. Frozen parameters get no hook and never cross the wire. That is why LoRA and DDP combine so cheaply.

3. The sampler shards the data

DistributedSampler gives each rank a disjoint slice of the dataset indices. It shuffles with a seed plus the epoch number. Call sampler.set_epoch(epoch) at the start of every epoch. If you do not, every epoch replays the same order.

By default, the sampler pads the index list, so every rank gets the same number of samples. With drop_last=True, it trims the list instead. In both cases, every rank sees the same number of batches. An unequal count is one of the ways that a job hangs (see failure modes).

4. Backward overlaps with communication

The simulation runs forward, backward, and all-reduce as separate phases. Real DDP puts parameters into buckets (the default bucket_cap_mb is 25 MB). The bucket order is approximately the reverse of the model’s parameter order, because gradients become ready in about that order during backward. When all gradients in a bucket are ready, DDP starts an asynchronous all-reduce for that bucket. At the same time, backward continues on the earlier layers. When autograd gets to the first layer, most of the communication is already complete.

sequenceDiagram
  participant G0 as GPU 0
  participant G1 as GPU 1
  participant NCCL as NCCL
  Note over G0,G1: micro-batches 1..N-1 in no_sync()
  G0->>G0: forward (last micro-batch)
  G1->>G1: forward (last micro-batch)
  G0->>G0: backward: last layers
  G1->>G1: backward: last layers
  G0->>NCCL: bucket k ready
  G1->>NCCL: bucket k ready
  Note over NCCL: all-reduce bucket k
  G0->>G0: backward: earlier layers
  G1->>G1: backward: earlier layers
  G0->>NCCL: bucket 1 ready
  G1->>NCCL: bucket 1 ready
  NCCL-->>G0: averaged gradients
  NCCL-->>G1: averaged gradients
  Note over G0,G1: clip + optimizer.step()

DDP averages the gradients across ranks (it does not sum them). Then the optimizer step runs independently on every rank, from identical gradients and identical starting weights. Thus the replicas stay identical, and DDP never needs to broadcast the weights again.

5. no_sync() for gradient accumulation

With gradient accumulation, you want one all-reduce per optimizer step, not one per micro-batch. model.no_sync() is a context manager that skips the reduction. Gradients accumulate locally in .grad. The first forward and backward pass outside the context then synchronizes all the gradients accumulated up to that point. The API docs warn that the forward pass must also run inside the context. If it does not, DDP still synchronizes the gradients.

The visualizer shows the same thing, with one all-reduce after all the accumulation steps. Its step comment even reads # All-reduce after accumulation.

6. What DDP doesn’t do

DDP only averages gradients. It does not average your loss for logs or coordinate checkpoint writes. It also does not make evaluation distributed or protect you from ranks that take different code paths. All of that is on you, and most of the script below handles it.

Batch size and learning rate

The effective (global) batch size is:

global_batch = micro_batch_per_gpu × world_size × accumulation_steps

The simulator gives the same identity in slightly different terms. Its batch_size argument is the per-step batch across all GPUs. Thus micro_batch_size = batch_size // num_gpus and effective_batch_size = batch_size × gradient_accumulation_steps. The data-parallel tests pin it down: 128 samples on 2 GPUs with 4 accumulation steps give an effective batch of 512.

There is one nit in my own code. If a batch does not divide evenly, the integer division drops the remainder. But effective_batch_size still reports the requested size. In a real training script, always calculate the global batch from the micro-batch. The script below prints that value.

This has two consequences for fine-tuning:

• More GPUs change the optimization problem. Go from 1 GPU to 8 with the same per-GPU micro-batch and accumulation, and the global batch becomes 8 times larger. The number of optimizer steps per epoch becomes 8 times smaller. If you tuned a learning rate on one GPU, that tuning does not carry over automatically. The DDP docs say the same thing. Because DDP averages gradients, a DDP model behaves like a single-GPU model trained on the global batch.
• Pick the global batch first. Then solve for accumulation. If the recipe that you trust used a global batch of 64, keep 64. On 8 GPUs with a micro-batch of 1, that is 8 accumulation steps. Then DDP changes only the wall-clock time.

The linear scaling rule (scale the learning rate with the batch, and add warmup) came from large-batch SGD on ImageNet. For Adam-style optimizers on LoRA fine-tunes, it is only a starting heuristic. Sweep the learning rate at the global batch that you actually run.

Language models add one subtlety. The Hugging Face causal-LM loss averages over the tokens in each micro-batch. Then DDP averages equally across ranks. Thus micro-batches with very different token counts get equal weight. For chat fine-tuning with variable-length samples, this is usually fine. If it matters for your data, do these steps:

1. Sum the per-token loss.
2. All-reduce the token count.
3. Divide by the global token count.

Memory budget for Gemma 4 12B with LoRA

Use this budget to decide which GPUs to request. To get the real peak on your hardware, call torch.cuda.max_memory_allocated() after the first few steps. GB means 10⁹ bytes.

The budget uses these assumptions:

• google/gemma-4-12B-it
• bf16 base weights
• LoRA rank 16 on all text-decoder projections (65.6M trainable parameters)
• Adapter weights in fp32 (PEFT’s default when the base model is bf16)
• AdamW
• Gradient checkpointing on
• Micro-batch 1
• Sequence length 2,048

The trainable state is small. The adapter weights, gradients, and optimizer moments together come to about 1 GB. This makes DDP viable. Every GPU holds a full copy of them, and it barely matters.

The vocabulary is the activation hog. With 262K entries, the logits tensor at 2,048 tokens has 537M values, and cross-entropy wants it in fp32. Gemma 4 also applies a final logit softcap (tanh scaling), which makes extra temporaries. If you double the sequence length to 4,096, this line approximately doubles, to around 12–18 GB. If you run out of memory at long sequence lengths, look here first, before the decoder layers.

The 16-bytes-per-parameter rule rules out DDP for full fine-tuning. Mixed-precision Adam holds about 16 bytes per trainable parameter. The ZeRO paper gives the breakdown: bf16 weights and gradients, plus fp32 master weights and two fp32 moments. For 11.96B parameters, that is about 191 GB before any activations. No single GPU can hold that. To fully fine-tune the 12B model, shard that state with FSDP or ZeRO. Do not replicate it with DDP.

This table applies the same analysis to the other sizes (weights only, bf16):

For the MoE model, the 3.8B active parameters set the compute per token. But every expert must still be resident on every replica. Its module layout is also different from the dense models. Before you train it under DDP, check it for unused parameters.

Communication per optimizer step

DDP moves gradients. Thus LoRA decreases the network cost by the same factor that it decreases the trainable state. A ring all-reduce sends about 2 × (N − 1) / N times the buffer size per GPU. This is the bus-bandwidth factor that nccl-tests uses in its performance notes. For N = 8, that factor is 1.75.

Google’s fine-tuning guide sets modules_to_save=["lm_head", "embed_tokens"] in its LoRA config. On a single GPU, that is a memory question. Under DDP, it also makes the all-reduce about ten times larger. If you do not need to train the embeddings, keep them frozen.

To change bytes into time, divide by the bus bandwidth that nccl-tests measures on your own machine. At 100 GB/s of bus bandwidth, the LoRA all-reduce takes about 5 ms. It occurs once per optimizer step, and most of it hides behind backward.

The training script

This is a complete, minimal DDP training script. It uses plain PyTorch for the distributed parts, Transformers to load Gemma 4, and PEFT for LoRA. I left out the Hugging Face Trainer and TRL on purpose, so that you can see every distributed concept. In production, I usually reach for TRL’s SFTTrainer, launched under torchrun. It does the same things underneath.

The dataset uses the same format as Google’s guide, with one JSON object per line:

{
	"messages": [
		{ "role": "user", "content": "..." },
		{ "role": "assistant", "content": "..." }
	]
}

# train_ddp.py: LoRA fine-tuning of Gemma 4 with plain PyTorch DDP.
# Launch: torchrun --standalone --nproc_per_node=8 train_ddp.py
import contextlib
import json
import os
import time

import torch
import torch.distributed as dist
from peft import LoraConfig, get_peft_model
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    get_cosine_schedule_with_warmup,
)

MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-12B-it")
DATA_PATH = os.environ.get("DATA_PATH", "train.jsonl")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "checkpoints")

MAX_LEN = 2048  # tokens per sample after truncation
MICRO_BATCH = 1  # samples per GPU per forward pass
ACCUM_STEPS = 8  # micro-batches per optimizer step
EPOCHS = 1
LR = 1e-4
WARMUP_STEPS = 20
MAX_GRAD_NORM = 1.0
LOG_EVERY = 10
SEED = 42

# Only the text decoder's projections. Keeps LoRA off the vision/audio paths,
# which a text-only batch never touches.
LORA_TARGETS = (
    r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)"
)


class ChatJsonl(Dataset):
    """One JSON object per line: {"messages": [{"role": ..., "content": ...}]}."""

    def __init__(self, path, tokenizer, max_len):
        with open(path) as f:
            self.rows = [json.loads(line) for line in f if line.strip()]
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, idx):
        text = self.tokenizer.apply_chat_template(
            self.rows[idx]["messages"], tokenize=False
        )
        # The chat template already emits <bos>, so don't add it twice.
        ids = self.tokenizer(
            text, add_special_tokens=False, truncation=True, max_length=self.max_len
        )["input_ids"]
        return torch.tensor(ids, dtype=torch.long)


class PadCollate:
    """Right-pads a batch and masks padding out of the loss.

    A top-level class, not a closure, so DataLoader workers can pickle it.
    """

    def __init__(self, pad_id):
        self.pad_id = pad_id

    def __call__(self, batch):
        longest = max(len(x) for x in batch)
        input_ids = torch.full((len(batch), longest), self.pad_id, dtype=torch.long)
        attention_mask = torch.zeros((len(batch), longest), dtype=torch.long)
        for i, ids in enumerate(batch):
            input_ids[i, : len(ids)] = ids
            attention_mask[i, : len(ids)] = 1
        labels = input_ids.masked_fill(attention_mask == 0, -100)
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }


def main():
    # torchrun sets RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR and MASTER_PORT.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl", device_id=device)
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    is_main = rank == 0
    torch.manual_seed(SEED)  # same LoRA init on every rank

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, dtype=torch.bfloat16, device_map={"": local_rank}
    )
    model.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False}
    )
    model = get_peft_model(
        model,
        LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            target_modules=LORA_TARGETS,
            task_type="CAUSAL_LM",
        ),
    )
    if is_main:
        model.print_trainable_parameters()

    model = DDP(model, device_ids=[local_rank], gradient_as_bucket_view=True)
    torch.manual_seed(SEED + rank)  # but different dropout masks per rank

    dataset = ChatJsonl(DATA_PATH, tokenizer, MAX_LEN)
    sampler = DistributedSampler(dataset, shuffle=True, seed=SEED, drop_last=True)
    loader = DataLoader(
        dataset,
        batch_size=MICRO_BATCH,
        sampler=sampler,
        collate_fn=PadCollate(tokenizer.pad_token_id),
        drop_last=True,
        num_workers=2,
        pin_memory=True,
    )

    trainable = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable, lr=LR, weight_decay=0.0, fused=True)
    steps_per_epoch = len(loader) // ACCUM_STEPS  # identical on every rank
    scheduler = get_cosine_schedule_with_warmup(
        optimizer, WARMUP_STEPS, steps_per_epoch * EPOCHS
    )
    if is_main:
        global_batch = MICRO_BATCH * world_size * ACCUM_STEPS
        print(f"world={world_size} global_batch={global_batch} steps={steps_per_epoch}")

    for epoch in range(EPOCHS):
        sampler.set_epoch(epoch)  # new shuffle each epoch, same on all ranks
        model.train()
        batches = iter(loader)
        for step in range(steps_per_epoch):
            t0 = time.perf_counter()
            loss_sum = torch.zeros((), device=device)
            tokens = torch.zeros((), device=device)
            for micro in range(ACCUM_STEPS):
                batch = {
                    k: v.to(device, non_blocking=True) for k, v in next(batches).items()
                }
                last = micro == ACCUM_STEPS - 1
                # Skip the gradient all-reduce on all but the last micro-batch.
                sync = contextlib.nullcontext() if last else model.no_sync()
                with sync:
                    with torch.autocast("cuda", dtype=torch.bfloat16):
                        out = model(**batch, use_cache=False)
                    loss = out.loss / ACCUM_STEPS
                    loss.backward()
                loss_sum += loss.detach()
                tokens += batch["attention_mask"].sum()

            grad_norm = torch.nn.utils.clip_grad_norm_(trainable, MAX_GRAD_NORM)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad(set_to_none=True)

            if step % LOG_EVERY == 0:
                # Logging only: training already synced gradients, not losses.
                stats = torch.stack([loss_sum, tokens])
                dist.all_reduce(stats, op=dist.ReduceOp.SUM)
                loss_avg = stats[0].item() / world_size
                total_tokens = stats[1].item()
                dt = time.perf_counter() - t0  # .item() above waited for the GPU
                if is_main:
                    print(
                        f"epoch={epoch} step={step} loss={loss_avg:.4f} "
                        f"grad_norm={grad_norm.item():.3f} "
                        f"lr={scheduler.get_last_lr()[0]:.2e} "
                        f"tok/s={total_tokens / dt:,.0f}"
                    )

        if is_main:
            ckpt = os.path.join(OUTPUT_DIR, f"epoch-{epoch}")
            model.module.save_pretrained(ckpt)  # adapter weights only
            torch.save(
                {
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                    "epoch": epoch,
                },
                os.path.join(ckpt, "trainer_state.pt"),
            )
        dist.barrier()  # nobody races ahead while rank 0 writes

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Why each piece is there

Process group and device. torch.cuda.set_device(local_rank) runs before any other code uses CUDA, so each process allocates memory only on its own GPU. The device_id argument to init_process_group binds the process group to that device at the start. LOCAL_RANK selects the GPU on this machine. RANK (via dist.get_rank()) is the global identity for “only rank 0 does this.”

How the model loads. device_map={"": local_rank} loads the whole model directly onto the GPU of this rank. Every rank loads its own copy, which is the point of DDP. Download the checkpoint once before you start the job (hf download google/gemma-4-12B-it). If you do not, eight processes race to fill the same cache. With current Transformers, AutoModelForCausalLM resolves the 12B checkpoint to its full multimodal class. The model card uses AutoModelForMultimodalLM, and both classes are fine for text-only training.

LoRA targets. LORA_TARGETS is a regex that matches only the projections under language_model. (PEFT treats a string target_modules as a regex match.) With a text-only dataset, adapters on image or audio paths never receive gradients. DDP raises an error when a parameter that requires a gradient does not get one. A narrow adapter scope is cleaner than find_unused_parameters=True, which adds an extra graph traversal on every step.

Gradient checkpointing. Non-reentrant checkpointing (use_reentrant=False) is the variant that works cleanly with DDP and with frozen base weights. The reentrant variant has sharp edges with DDP. Checkpointing keeps the activation lines in the memory table small.

Seeds. The same seed before get_peft_model gives every rank the same adapter initialization. (The rank-0 broadcast in DDP also enforces this.) A new seed of SEED + rank after the wrap gives each rank different LoRA dropout masks, which is what you want.

gradient_as_bucket_view=True. With this option, gradients become views into the DDP communication buckets. Without it, they are separate tensors that DDP copies in and out. The option saves one gradient-sized buffer and one copy per step. The saving is small for LoRA and large for anything bigger.

Sampler and loader. With DistributedSampler(drop_last=True) and DataLoader(drop_last=True), every rank gets the same number of full micro-batches. The script calculates steps_per_epoch from len(loader), which is identical on every rank. The loop runs exactly that many optimizer steps. The script drops a partial accumulation window at the end, so different ranks cannot handle it differently.

Accumulation with no_sync(). The first ACCUM_STEPS − 1 micro-batches run forward and backward inside model.no_sync(). Then the last micro-batch runs outside it, and its backward starts the bucketed all-reduce. The script divides the loss by ACCUM_STEPS, so the accumulated gradient is a mean, not a sum.

bf16 autocast. The base weights are already bf16. The LoRA weights are fp32, because PEFT upcasts bf16 adapters by default for stable training. Autocast runs the matmuls in bf16, and the optimizer updates the fp32 adapter weights. You do not need a GradScaler. It is for fp16, which has a narrow exponent range that underflows. bf16 has the same exponent range as fp32.

Clipping. Clipping runs after the all-reduce, on identical averaged gradients, so every rank computes the same norm and applies the same scale.

Logs. The logged loss is an explicit all_reduce of the local losses, and every rank runs it. LOG_EVERY does not depend on the rank. The .item() calls force a sync with the GPU, so the step time is honest. Without the sync, the step time shows only how long it took to queue kernels.

Checkpointing. The replicas are identical, so only rank 0 writes. model.module unwraps DDP, and save_pretrained on the PEFT model writes only the adapter (about 262 MB for this config in fp32). The optimizer state is also identical across ranks, so one copy is enough to resume. The barrier() makes the other ranks wait while rank 0 writes, so they cannot race ahead or exit.

Teardown. destroy_process_group() shuts down NCCL cleanly, so the job exits. It does not hang on shutdown or print warnings.

What I’d add before a real run

• Loss masking to assistant tokens only. The script trains on the whole rendered conversation, including the prompts. Most chat fine-tunes mask everything except the assistant turns. They set those label positions to -100.
• Resume. Load the adapter with PEFT. Load the optimizer and scheduler state from trainer_state.pt. Then fast-forward the sampler to the saved epoch.
• Evaluation. Never call the forward of the DDP-wrapped model on one rank only. You have two options. Run evaluation on every rank over a DistributedSampler, and all-reduce the metrics. Or run it on rank 0 with the unwrapped model.module under torch.no_grad(), while the other ranks wait at a barrier.
• Token-weighted loss. Use it if your sample lengths vary wildly (see batch size and learning rate).

Launch with torchrun

On one node with eight GPUs:

pip install torch "transformers>=5.10.1" "peft>=0.19.0" accelerate
hf download google/gemma-4-12B-it

torchrun --standalone --nproc_per_node=8 train_ddp.py

--standalone sets up a local rendezvous, so you do not have to pick a master address or port. For two nodes, run the same command on each node with a shared rendezvous endpoint:

torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=node0.example.internal:29400 \
  --rdzv_id=gemma4-lora-001 \
  train_ddp.py

The script does not change. RANK and WORLD_SIZE now span 16 processes. LOCAL_RANK still goes from 0 to 7 on each node. Rank 0 (the only writer) is on one node. Point OUTPUT_DIR at storage that the node of rank 0 can write to. The torchrun documentation covers elastic options such as --max_restarts.

Across nodes, the all-reduce leaves NVLink and goes over the network. For the LoRA buffer of approximately a quarter gigabyte, that is usually tolerable. For full fine-tuning, it is the reason people buy InfiniBand.

Run it on Modal

The visualizer’s GPU benchmarks run on Modal. There, the 2×T4 tensor-parallel endpoint is only @app.function(gpu="T4:2"). The same pattern works for this script. Modal’s GPU guide documents the "H100:8" syntax for eight GPUs in one container. It also recommends that you run multi-process training as a subprocess, which is what torchrun wants:

# modal_train.py
import os
import subprocess

import modal

app = modal.App("gemma4-ddp")
volume = modal.Volume.from_name("gemma4-ddp", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch", "transformers>=5.10.1", "peft>=0.19.0", "accelerate")
    .env({"HF_HOME": "/vol/hf"})
    .add_local_file("train_ddp.py", "/root/train_ddp.py")
    .add_local_file("train.jsonl", "/root/train.jsonl")
)


@app.function(image=image, gpu="H100:8", volumes={"/vol": volume}, timeout=6 * 60 * 60)
def train():
    env = {
        **os.environ,
        "DATA_PATH": "/root/train.jsonl",
        "OUTPUT_DIR": "/vol/checkpoints",
    }
    subprocess.run(
        ["torchrun", "--standalone", "--nproc_per_node=8", "/root/train_ddp.py"],
        check=True,
        env=env,
    )
    volume.commit()


@app.local_entrypoint()
def main():
    train.remote()

Run it with modal run modal_train.py. The volume holds the Hugging Face cache and the adapter checkpoints. Because the cache stays on the volume, later runs skip the download.

Failure modes that hang or corrupt a DDP job

Most DDP bugs do not crash. They hang. One rank waits in a collective that another rank never enters, until the NCCL watchdog times out. One rule applies to every item below. Every rank must issue the same collectives, in the same order, with the same shapes.

Unequal numbers of batches

If rank 3 gets one more batch than rank 5, rank 3 enters an all-reduce that rank 5 never joins. Some causes are:

• Custom samplers that do not pad.
• An IterableDataset sharded by file, where the files have different lengths.
• Sample filters inside the training loop on some ranks.

Fix:

• Use DistributedSampler with drop_last=True.
• Compute the step count from len(loader), as the script does.
• Never continue past a batch on one rank only.

If uneven inputs are part of your data, use PyTorch’s Join context manager for uneven inputs. With it, ranks that finish early shadow the collectives of the ranks that still run.

Parameters that don’t receive gradients

Symptom: A runtime error on the second iteration. The error says that DDP expected to finish reduction in the prior iteration. It lists the parameters that did not receive gradients.

Cause: The forward pass did not use a parameter with requires_grad=True. For Gemma 4, the usual culprit is target_modules="all-linear" on a multimodal checkpoint. This setting attaches LoRA to vision and audio projections that a text-only batch never runs.

I reproduced this with a tiny test model built from the 12B config. With all-linear, the first step succeeds. The second step fails with Expected to have finished reduction in the prior iteration before starting a new one. The error names six parameter indices that never received a gradient.

Fix: Limit the adapters to the modules that you actually use, as the regex in the script does. Or freeze the unused modules. find_unused_parameters=True also fixes it, but it adds a full graph traversal after every forward pass. Use it when the unused parameters really depend on the data. That is plausible with a mixture-of-experts model, if your adapters are on expert weights that some micro-batches never route to.

Rank-divergent code paths

Some examples are if rank == 0: dist.all_reduce(...), evaluation through the DDP wrapper on rank 0 only, and a break on one rank because its loss went NaN. All of these desynchronize the collectives.

Fix: Keep collectives unconditional. For decisions such as early stopping, do these steps:

1. Make the decision on rank 0.
2. Broadcast it as a tensor.
3. Make every rank act on the broadcast value.

Model changes after the DDP wrap

DDP registers its hooks when you construct it. If you add adapters, unfreeze layers, or swap modules after DDP(...), the hooks no longer match the parameters. The API docs warn against this. Apply PEFT, gradient checkpointing, and any freezing before you wrap the model. The script uses this order.

Doubled or missing synchronization with accumulation

If forward runs outside no_sync() and only backward runs inside it, DDP silently synchronizes every micro-batch. Nothing breaks, but the job is slower. The opposite mistake is to run the last micro-batch inside no_sync() too. Then DDP skips the all-reduce completely, and the replicas quietly drift apart. For a cheap check, compare a checksum of the trainable parameters across ranks after a few steps.

DataLoader workers and pickling

With the spawn or forkserver multiprocessing start methods, everything that a worker receives must be picklable. (Spawn is the default on macOS, and the DDP docs recommend forkserver or spawn in some NCCL setups.) My first draft used a closure as collate_fn. It failed immediately under spawn with Can't get local object 'make_collate.<locals>.collate'. That is why the script uses a top-level PadCollate class.

Every rank writes checkpoints

On shared filesystems, if eight processes write the same file, the output is corrupt or interleaved. Write on rank 0. Then call barrier().

NCCL and debugging environment variables

I set these environment variables before I dig into a distributed problem. For details, see NVIDIA’s NCCL environment variable reference and the PyTorch distributed documentation.

Two more habits help. Run nvidia-smi topo -m once per machine type. It shows whether the GPUs talk over NVLink or PCIe, and that changes the throughput to expect. When you debug, pass a shorter timeout to init_process_group. Then a hang becomes a stack trace in minutes, not after the default wait.

Measure throughput and scaling

Monitor tokens per second, across all ranks and per GPU. The script logs non-padding tokens per second for each logged step. To measure scaling, do these steps:

1. Run the same config on 1, 2, 4, and 8 GPUs, with the same micro-batch per GPU.
2. Discard the first steps. CUDA context setup, allocator warmup, and NCCL communicator setup all land there.
3. Average over a steady window.
4. Compute scaling efficiency as tokens_per_s(N) / (N × tokens_per_s(1)).

The visualizer’s 2×T4 benchmark reports the same formula. api_multi_gpu_benchmark in app.py computes parallel_efficiency = speedup / ideal_speedup × 100 for a real column-parallel matmul split across two T4s. Its breakdown field separates compute time from GPU-to-GPU transfer time.

The README does not publish a fixed result, because the endpoint measures live on each call. The benchmark shows that a split of work across GPUs makes the work faster, and that communication eats part of that speedup. This lesson carries over directly to DDP.

Qualitatively, expect these results from DDP with LoRA:

• The all-reduce is small and mostly hidden. The LoRA gradients are about 0.26 GB in fp32. DDP reduces them once per optimizer step and overlaps the reduction with backward. On NVLink-connected GPUs, communication should be a small fraction of step time. Scaling efficiency in a node should be high.
• Accumulation amortizes it further. With ACCUM_STEPS=8, seven of every eight backward passes do no communication.
• The bottleneck moves to the input pipeline and padding. When communication is cheap, padding waste (long and short samples in the same batch) and slow tokenization become visible. Sort or bucket samples by length, or pack them. That usually buys more than NCCL tuning.
• Across nodes it depends on the network. The same 0.46 GB per GPU per step can dominate on a slow Ethernet link.

To see the overlap directly, capture a few steps with torch.profiler. Then look at the trace. The NCCL all-reduce kernels should run at the same time as the backward compute kernels, not after them. If they do not overlap, the usual cause is buckets that fill too late. bucket_cap_mb is the knob to try.

---

The visualizer’s interactive data-parallel tab shows the same loop at toy scale. For the production side, see the multi-GPU Kubernetes delivery platform case study. It shows how this kind of multi-GPU work fits into a delivery platform.