How to Perform Multi-GPU Pre-Training with MiniMind: A Complete DDP Guide

MiniMind supports multi-GPU pre-training through native PyTorch Distributed Data Parallel (DDP), automatically initializing process groups, sharding data via DistributedSampler, and synchronizing checkpoints across devices when launched with torchrun.

The MiniMind repository (jingyaogong/minimind) provides a production-ready pre-training pipeline that scales seamlessly from single-GPU debugging to multi-node clusters. By leveraging PyTorch's distributed training utilities, the framework handles complex coordination tasks—environment setup, gradient synchronization, and dataset partitioning—without requiring manual code changes.

Understanding MiniMind's Distributed Architecture

MiniMind's pre-training system is built on standard PyTorch DDP primitives. The architecture ensures that model parameters, gradients, and batch normalization buffers stay synchronized across all GPUs during the forward-backward pass.

Distributed Environment Initialization

The entry point trainer/train_pretrain.py delegates environment setup to init_distributed_mode() inside trainer/trainer_utils.py. This function reads the RANK environment variable injected by the launcher, initializes an NCCL backend process group, and binds each process to its designated CUDA device using local_rank【/cache/repos/github.com/jingyaogong/minimind/master/trainer/trainer_utils.py#L44-L51】.

The script automatically detects whether it is running in distributed mode by checking for the presence of these environment variables, falling back to single-GPU training if they are absent.

Data Parallelism and Model Wrapping

Once the process group is established, the training script creates a DistributedSampler for the PretrainDataset class defined in dataset/lm_dataset.py. This ensures each GPU receives a disjoint, non-overlapping slice of the pre-training corpus【/cache/repos/github.com/jingyaogong/minimind/master/trainer/train_pretrain.py#L29-L30】.

The core model—MiniMindForCausalLM wrapping MiniMindModel from model/model_minimind.py—is then wrapped in PyTorch's DistributedDataParallel container. All model parameters are registered within standard nn.Module instances, ensuring full compatibility with DDP's gradient synchronization hooks【/cache/repos/github.com/jingyaogong/minimind/master/model/model_minimind.py#L27-L31】【/cache/repos/github.com/jingyaogong/minimind/master/train_pretrain.py#L42-L46】.

Launching Multi-GPU Training Jobs

MiniMind requires no internal configuration changes to scale across GPUs. You control parallelism entirely through the launch command, using either torchrun (recommended) or the legacy torch.distributed.launch wrapper.

Single-Node Multi-GPU Setup

For a single machine with multiple GPUs, use torchrun with --nproc_per_node set to the number of available devices. This command automatically sets WORLD_SIZE, RANK, and LOCAL_RANK for each spawned process:

torchrun \
  --nproc_per_node=4 \
  trainer/train_pretrain.py \
  --save_dir ../out \
  --save_weight pretrain \
  --epochs 2 \
  --batch_size 32 \
  --learning_rate 5e-4 \
  --dtype bfloat16 \
  --num_workers 8 \
  --accumulation_steps 8 \
  --grad_clip 1.0 \
  --max_seq_len 340 \
  --use_moe 0 \
  --data_path ../dataset/pretrain_hq.jsonl \
  --use_wandb \
  --wandb_project MiniMind-Pretrain

Each process loads the full model but trains only on its data partition, with gradients averaged across all four GPUs during the backward pass.

Multi-Node Cluster Configuration

For distributed clusters, specify the total number of nodes (--nnodes), the current node's rank (--node_rank), and the master node's address and port. Run the identical command on every node in the cluster:

torchrun \
  --nnodes=2 \
  --nproc_per_node=4 \
  --node_rank=$NODE_RANK \
  --master_addr=$MASTER_ADDR \
  --master_port=$MASTER_PORT \
  trainer/train_pretrain.py \
  --save_dir ../out \
  --epochs 2 \
  --batch_size 32 \
  --learning_rate 5e-4

The RANK variable ranges from 0 to WORLD_SIZE-1 (0-7 in this two-node example), while LOCAL_RANK always maps to the specific GPU index on the current machine (0-3).

Legacy Launch Method

If your infrastructure requires the deprecated launcher, use python -m torch.distributed.launch with identical argument syntax:

python -m torch.distributed.launch \
  --nproc_per_node=4 \
  trainer/train_pretrain.py \
  --epochs 2 \
  --batch_size 32 \
  --learning_rate 5e-4

Note that torchrun is preferred as it provides better error handling and automatic rendezvous backend selection.

Checkpoint Handling Across GPUs

MiniMind's checkpointing logic in trainer/trainer_utils.py is designed for elastic distributed training. The lm_checkpoint() function stores a shared checkpoint file that records the current world_size alongside model weights and optimizer states【/cache/repos/github.com/jingyaogong/minimind/master/trainer/trainer_utils.py#L63-L66】.

When resuming training, the script automatically detects if the number of available GPUs has changed from the saved configuration. It rescales the global step count accordingly to maintain consistent learning rate schedules and logging intervals across different hardware topologies【/cache/repos/github.com/jingyaogong/minimind/master/trainer/trainer_utils.py#L11-L13】.

This allows you to seamlessly transition between single-GPU debugging and multi-GPU production runs without manual step recalculation or learning rate warm-up adjustments.

Summary

• Automatic DDP Integration: MiniMind's trainer/train_pretrain.py detects distributed environments via init_distributed_mode() and configures NCCL communication automatically.
• Data Sharding: The DistributedSampler ensures non-overlapping data partitions per GPU without manual dataset splitting.
• Flexible Launching: Use torchrun with --nproc_per_node for single-node training or add --nnodes and --node_rank for multi-node clusters.
• Elastic Checkpoints: The lm_checkpoint() system handles world-size changes, allowing training resumption on different GPU counts.
• Model Compatibility: The MiniMindForCausalLM architecture in model/model_minimind.py uses standard PyTorch modules, ensuring full DDP compatibility.

Frequently Asked Questions

What launch command should I use for multi-GPU pre-training with MiniMind?

Use torchrun --nproc_per_node=N trainer/train_pretrain.py where N is the number of GPUs. This injects the required RANK, WORLD_SIZE, and LOCAL_RANK environment variables that MiniMind's init_distributed_mode() function expects. The script then automatically wraps the model in DistributedDataParallel and configures the DistributedSampler.

How does MiniMind handle data sharding across multiple GPUs?

The framework instantiates a DistributedSampler for the PretrainDataset class inside trainer/train_pretrain.py【/cache/repos/github.com/jingyaogong/minimind/master/trainer/train_pretrain.py#L29-L30】. This sampler ensures each process receives a unique subset of the JSONL pre-training data based on its RANK, eliminating data duplication while maintaining deterministic epoch boundaries across all GPUs.

Can I resume training on a different number of GPUs than originally used?

Yes. The checkpointing utility lm_checkpoint() saves the original world_size with each checkpoint. On resume, MiniMind automatically rescales the global step counter to account for the new GPU count, ensuring learning rate schedules and logging remain consistent even when switching between single-GPU debugging and multi-GPU production runs【/cache/repos/github.com/jingyaogong/minimind/master/trainer/trainer_utils.py#L63-L66】.

Is the MiniMind model architecture compatible with PyTorch DDP?

Yes. The MiniMindForCausalLM and underlying MiniMindModel classes in model/model_minimind.py inherit from standard nn.Module and register all trainable parameters within PyTorch's module hierarchy【/cache/repos/github.com/jingyaogong/minimind/master/model/model_minimind.py#L27-L31】. This design ensures that DistributedDataParallel can correctly trace all gradients and synchronize them across processes during the backward pass.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}