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

LoRA fine-tuning on a single 8-GPU node fits in per-device memory, so plain DDP is simpler and faster than tensor or pipeline parallelism. When a DDP job fails, though, it usually hangs until the NCCL watchdog times out or corrupts weights across ranks. Fine-tuning Gemma 4 12B reliably meant calculating VRAM usage upfront and making sure every rank runs the exact same collectives in lockstep.

Sudhanva wrote the reference training script, memory analysis, and failure reproductions independently, building on his open-source Megatron Parallelism Visualizer.

• Standalone torchrun script covering process-group init, regex-scoped LoRA targets, DistributedSampler sharding, no_sync() gradient accumulation, and rank-0 checkpointing.
• Per-GPU memory and NCCL bus-bandwidth budget derived from the Gemma 4 12B config and safetensors header.
• Minimal reproductions of common DDP deadlocks and unused-parameter crashes.
• The open-source Megatron Parallelism Visualizer that models the data-parallel execution loop.

• Checked per-GPU memory first. Because one BF16 replica with LoRA fits on a 48 GB GPU, plain DDP avoids the extra communication of FSDP, ZeRO, or tensor parallelism.
• Scoped LoRA target_modules via regex to language_model projections so vision and audio towers never receive adapters that go untouched during text-only batches.
• Enabled non-reentrant gradient checkpointing (use_reentrant=False) and PEFT before wrapping the model in DDP with gradient_as_bucket_view=True.
• Sharded batches via DistributedSampler(drop_last=True) and computed step counts from len(loader) so no rank ever runs an extra iteration.
• Wrapped micro-batches 1..N-1 in model.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.
• Restricted checkpoint writes to rank 0 followed by a dist.barrier() so workers do not overwrite shared storage or race ahead.

• DDP keeps model code simple and communicates once per optimizer step, at the cost of duplicating frozen base weights on every GPU.
• Full fine-tuning of Gemma 4 12B requires ~191 GB of mixed-precision Adam state before activations, which exceeds single-GPU memory and requires FSDP or ZeRO.
• Scoping LoRA with a strict regex avoids the per-step autograd graph traversal overhead of find_unused_parameters=True.
• Freezing embed_tokens and lm_head keeps the per-step gradient all-reduce around 0.46 GB per GPU instead of ~4.6 GB.

• A standalone torchrun training script for 8-GPU single-node and multi-node runs.
• About 35–40 GB per GPU for Gemma 4 12B with LoRA at 2,048 tokens, which fits a 48 GB GPU.
• A 191 GB state calculation showing why full 12B fine-tuning requires FSDP or ZeRO.
• Seven documented DDP failure modes, from uneven batch counts to closure pickling, each with a reproduction and fix.

• Verified that LoRA adapter parameter checksums stay identical across ranks after training.
• Reproduced the second-step unused-parameter crash caused by target_modules="all-linear" on multimodal checkpoints.
• Replaced closure-based data collators with a top-level PadCollate class so multi-worker DataLoaders pickle cleanly under spawn.
• Used NCCL_DEBUG, TORCH_DISTRIBUTED_DEBUG=DETAIL, and the NCCL flight recorder to trace hung collectives.

• Write out the VRAM budget for weights, gradients, optimizer state, and activations before choosing between DDP, FSDP, and tensor parallelism.
• With a 262K vocabulary, the fp32 logits and cross-entropy buffers take 6–9 GB at 2,048 tokens and are usually what causes OOMs on longer sequences.
• Moving from 1 GPU to 8 multiplies the effective batch size by 8, so fix the global batch size first and adjust gradient accumulation steps to match.
• Check the hand calculation against torch.cuda.max_memory_allocated() after the third training step, once AdamW buffers and allocator caches are initialized.