Multi-GPU LoRA Fine-Tuning of Gemma 4 with PyTorch DDP

Most fine-tuning jobs on one multi-GPU machine need neither tensor nor pipeline parallelism, so DDP is usually the right tool. When a DDP job goes wrong, it tends to hang or let the replicas drift apart without an error. LoRA fine-tuning Gemma 4 12B needed a rule for picking the parallelism, a per-GPU memory budget before paying for GPUs, and a training script in which every rank issues the same collectives in the same order.

Sudhanva did the study as personal research, built on his open-source Megatron Parallelism Visualizer, and wrote the training script.

• The torchrun training script: process group setup, LoRA targeting, DistributedSampler sharding, no_sync() gradient accumulation, logging, and rank-0 checkpoints.
• The memory and communication budget for Gemma 4 12B, built from the published checkpoint config and safetensors header.
• The reproductions of two failure modes.
• The Megatron Parallelism Visualizer, whose data-parallel simulation is the mental model for the study.

• Picked the parallelism by asking whether one replica fits on one GPU: DDP if it does, FSDP or ZeRO when optimizer state is the problem, and tensor or pipeline parallelism when layers or depth do not fit.
• Scoped LoRA with a regex to the projections under language_model, so no adapter lands on a vision or audio path that a text-only batch never runs.
• Applied PEFT and non-reentrant gradient checkpointing before wrapping the model in DDP with gradient_as_bucket_view=True.
• Sharded the data with DistributedSampler and drop_last=True, and derived the step count from len(loader) so every rank runs the same number of optimizer steps.
• Ran all but the last micro-batch inside no_sync(), so the bucketed all-reduce of LoRA gradients runs once per optimizer step instead of once per micro-batch, and overlaps with backward.
• Wrote checkpoints from rank 0 only, followed by a barrier.

• DDP leaves the model code unchanged and needs one gradient all-reduce per optimizer step, at the cost of a full copy of the training state on every GPU.
• Full fine-tuning Gemma 4 12B holds about 191 GB of mixed-precision Adam state before activations. That fits on no single GPU, so full fine-tuning calls for FSDP or ZeRO.
• Scoping the LoRA targets avoids find_unused_parameters=True and the graph traversal it adds after every forward pass.
• Training embed_tokens and lm_head makes the all-reduce about ten times larger. Leaving them frozen keeps it near 0.46 GB per GPU per step.

• A complete torchrun training script for one eight-GPU node or two nodes.
• About 35–40 GB per GPU for Gemma 4 12B with LoRA at 2,048 tokens, which fits a 48 GB GPU.
• About 191 GB for full fine-tuning, which rules out DDP and points to FSDP.
• Documented seven failure modes that hang or corrupt DDP jobs, each with a fix.

• LoRA parameter checksums match across ranks after training.
• Switching the LoRA targets to all-linear reproduced the unused-parameter error on the second step.
• A closure collate_fn fails to pickle under spawn, so the script uses a top-level PadCollate class.
• NCCL_DEBUG, TORCH_DISTRIBUTED_DEBUG=DETAIL, and the PyTorch NCCL flight recorder show which rank stopped in which collective.

• Check whether one replica fits on one GPU before choosing a parallelism strategy.
• With a 262K vocabulary, the logits tensor is usually the first place to look when long sequences run out of memory.
• Adding GPUs multiplies the global batch, so pick the global batch first and solve for gradient accumulation.
• torch.cuda.max_memory_allocated() after the first few steps gives the real per-GPU peak.