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

> Learn multi-GPU pre-training with MiniMind using PyTorch DDP. Effortlessly shard data and synchronize checkpoints across devices with torchrun for efficient model training.

- Repository: [jingyaogong/minimind](https://github.com/jingyaogong/minimind)
- Tags: how-to-guide
- Published: 2026-03-24

---

**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`](https://github.com/jingyaogong/minimind/blob/main/trainer/train_pretrain.py) delegates environment setup to `init_distributed_mode()` inside [`trainer/trainer_utils.py`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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:

```bash
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:

```bash
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:

```bash
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`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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`](https://github.com/jingyaogong/minimind/blob/main/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.