How to use Huggingface Trainer with multiple GPUs?

Huggingface Trainer

multi-GPU training

machine learning

parallel processing

deep learning

How to use Huggingface Trainer with multiple GPUs?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning Machine Learning Concurrency

Introduction

The Hugging Face Trainer class supports multi-GPU training out of the box. For data parallelism on a single machine, launch with torchrun or accelerate launch and the Trainer automatically distributes data across GPUs using PyTorch's DistributedDataParallel (DDP). For model parallelism or multi-node training, use DeepSpeed or FSDP integration via TrainingArguments.

Single Command Multi-GPU (Simplest)

If you have a standard training script using Trainer, multi-GPU works by changing only the launch command:

python

1# train.py
2from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
3from datasets import load_dataset
4
5model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
6tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
7dataset = load_dataset("glue", "mrpc")
8
9def tokenize(batch):
10    return tokenizer(batch["sentence1"], batch["sentence2"], padding="max_length", truncation=True)
11
12dataset = dataset.map(tokenize, batched=True)
13
14training_args = TrainingArguments(
15    output_dir="./results",
16    per_device_train_batch_size=16,  # Per GPU batch size
17    num_train_epochs=3,
18    logging_steps=100,
19)
20
21trainer = Trainer(
22    model=model,
23    args=training_args,
24    train_dataset=dataset["train"],
25    eval_dataset=dataset["validation"],
26)
27
28trainer.train()

bash

1# Launch with torchrun (PyTorch 1.10+)
2torchrun --nproc_per_node=4 train.py
3
4# Or with accelerate
5accelerate launch --num_processes=4 train.py
6
7# Or with the older torch.distributed.launch
8python -m torch.distributed.launch --nproc_per_node=4 train.py

The per_device_train_batch_size is per GPU. With 4 GPUs and per_device_train_batch_size=16, the effective batch size is 64.

Controlling GPU Selection

bash

1# Use specific GPUs
2CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 train.py
3
4# Check available GPUs in Python
5import torch
6print(torch.cuda.device_count())  # Number of visible GPUs

Gradient Accumulation

When GPU memory is limited, accumulate gradients over multiple steps to simulate a larger batch:

python

1training_args = TrainingArguments(
2    output_dir="./results",
3    per_device_train_batch_size=4,          # Small batch per GPU
4    gradient_accumulation_steps=8,           # Accumulate 8 steps
5    # Effective batch size = 4 * 8 * num_gpus
6    # With 4 GPUs: 4 * 8 * 4 = 128
7)

DeepSpeed Integration

For large models that do not fit on a single GPU, use DeepSpeed ZeRO:

json

1// ds_config.json
2{
3    "zero_optimization": {
4        "stage": 2,
5        "offload_optimizer": {
6            "device": "cpu"
7        }
8    },
9    "train_batch_size": "auto",
10    "train_micro_batch_size_per_gpu": "auto",
11    "gradient_accumulation_steps": "auto",
12    "fp16": {
13        "enabled": true
14    }
15}

python

1training_args = TrainingArguments(
2    output_dir="./results",
3    per_device_train_batch_size=8,
4    deepspeed="ds_config.json",   # Enable DeepSpeed
5    fp16=True,
6)

bash

1# Launch with DeepSpeed
2deepspeed --num_gpus=4 train.py --deepspeed ds_config.json
3
4# Or with torchrun
5torchrun --nproc_per_node=4 train.py

DeepSpeed ZeRO stages:

Stage 1: Partitions optimizer states across GPUs
Stage 2: Also partitions gradients
Stage 3: Also partitions model parameters (enables training models larger than single GPU memory)

FSDP (Fully Sharded Data Parallel)

PyTorch-native alternative to DeepSpeed:

python

1training_args = TrainingArguments(
2    output_dir="./results",
3    per_device_train_batch_size=8,
4    fsdp="full_shard auto_wrap",
5    fsdp_config={
6        "fsdp_min_num_params": 1e6,
7        "fsdp_transformer_layer_cls_to_wrap": "BertLayer",
8    },
9)

bash

torchrun --nproc_per_node=4 train.py

Mixed Precision Training

Reduce memory usage and speed up training with fp16 or bf16:

python

1training_args = TrainingArguments(
2    output_dir="./results",
3    per_device_train_batch_size=16,
4    fp16=True,            # Use float16 (NVIDIA GPUs with Tensor Cores)
5    # bf16=True,          # Use bfloat16 (A100, H100, or newer)
6)

Multi-Node Training

For training across multiple machines:

bash

1# On node 0 (master)
2torchrun \
3    --nproc_per_node=4 \
4    --nnodes=2 \
5    --node_rank=0 \
6    --master_addr=192.168.1.1 \
7    --master_port=29500 \
8    train.py
9
10# On node 1
11torchrun \
12    --nproc_per_node=4 \
13    --nnodes=2 \
14    --node_rank=1 \
15    --master_addr=192.168.1.1 \
16    --master_port=29500 \
17    train.py

Common Pitfalls

Running with python train.py instead of torchrun: The Trainer detects multi-GPU only when launched with torchrun or accelerate launch. Running with plain python uses a single GPU regardless of how many are available.
Confusing per-device and total batch size: per_device_train_batch_size is per GPU. The effective batch size is per_device * gradient_accumulation_steps * num_gpus. Accidentally setting it to the total desired batch size results in OOM errors.
Not setting CUDA_VISIBLE_DEVICES: Without this, all GPUs are used. On shared machines, this conflicts with other users. Always set CUDA_VISIBLE_DEVICES to restrict which GPUs your job uses.
Saving/loading checkpoints in multi-GPU: The Trainer handles this correctly by default (only the main process saves). But custom saving code must check trainer.is_world_process_zero() to avoid duplicate writes from all processes.
DeepSpeed config conflicting with TrainingArguments: When using DeepSpeed, set batch size and gradient accumulation to "auto" in the DeepSpeed config to let the Trainer control them. Hardcoded values that conflict with TrainingArguments cause silent misconfiguration.

Summary

Launch with torchrun --nproc_per_node=N train.py for multi-GPU training
per_device_train_batch_size is per GPU — effective batch = per_device * accumulation * num_gpus
Use gradient_accumulation_steps to simulate larger batches when GPU memory is limited
Use DeepSpeed ZeRO (stage 2 or 3) for models too large for a single GPU
Use fp16=True or bf16=True for mixed precision training to reduce memory and increase speed
Set CUDA_VISIBLE_DEVICES on shared machines to avoid GPU conflicts

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

How to use Huggingface Trainer with multiple GPUs?

ML System Design practice on Codemia

Introduction

Single Command Multi-GPU (Simplest)

Controlling GPU Selection

Gradient Accumulation

DeepSpeed Integration

FSDP (Fully Sharded Data Parallel)

Mixed Precision Training

Multi-Node Training

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia