How to Use Multi-GPUs to Train Your Model¶
DI-engine supports data-parallel training with multi-GPUs.
During data-parallel training, each device handles a portion of total input. Large training batch significantly accelerate the training process.
About data-parallel training in DI-engine, we support two types which are DataParallel(DP) and DataDistributedParallel(DDP).
The experimental environment referred to here is a single machine with multi-GPUs.
DataParallel(DP) Mode¶
DP is mainly used for single-machine multi-GPUs, single-process control multi-GPUs. There is 1 master node, the default is device[0]. The gradients are aggregated to the master, and the parameters are updated through backpropagation, and then the parameters are synchronized with other GPUs.
In DI-engine, we define ding.torch_utils.DataParallel, which inherits Torch.nn.DataParallel. And at the same time, we rewrite the parameters() method. please refer to
ding/torch_utils/dataparallel.py
import torch
import torch.nn as nn
class DataParallel(nn.DataParallel):
def __init__(self, module, device_ids=None, output_device=None, dim=0):
super().__init__(module, device_ids=None, output_device=None, dim=0)
self.module = module
def parameters(self, recurse: bool = True):
return self.module.parameters(recurse = True)
Training
from ding.entry import serial_pipeline
from ding.model.template.q_learning import DQN
from ding.torch_utils import DataParallel
model = DataParallel(DQN(obs_shape=[4, 84, 84],action_shape=6))
serial_pipeline((main_config, create_config), seed=0, model=model)
We don’t need to change any other code, just simply encapsulate the policy. Please refer to dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_dp.py
For DP, the runnable script demo is demonstrated as follows.
CUDA_VISIBLE_DEVICES=0,1 python -u spaceinvaders_dqn_config_multi_gpu_ddp.py
or (on cluster managed by Slurm)
srun -p PARTITION_NAME --mpi=pmi2 --gres=gpu:2 -n1 --ntasks-per-node=1 python -u spaceinvaders_dqn_config_multi_gpu_ddp.py
DataDistributedParallel(DDP) Mode¶
DataDistributedParallel(DDP) is mainly used for single-machine multi-GPUs and multi-machine multi-GPUs. It adopts multi-process to control multi-GPUs and adopts ring allreduce to synchronize gradient.
In DataDistributedParallel(DDP) Mode, we should simply set config.policy.learn.multi_gpu
as True
in the config file under dizoo/atari/config/serial/spaceinvaders/spaceinvaders_dqn_config_multi_gpu_ddp.py
.
Principle¶
We re-implement the data-parallel training module with APIs in torch.distributed
for high scalability. The detailed principle is shonw as follows:
Parameters on Rank-0 GPU are broadcasted to all devices, so that models on different devices share the same initialization.
def _init_multi_gpu_setting(self, model: torch.nn.Module) -> None:
for name, param in model.state_dict().items():
assert isinstance(param.data, torch.Tensor), type(param.data)
broadcast(param.data, 0)
for name, param in model.named_parameters():
setattr(param, 'grad', torch.zeros_like(param))
Gradients on different devices should be synchronized after the backward procedure.
self._optimizer.zero_grad()
loss.backward()
if self._cfg.learn.multi_gpu:
self.sync_gradients(self._learn_model) # sync gradients
self._optimizer.step()
def sync_gradients(self, model: torch.nn.Module) -> None:
for name, param in model.named_parameters():
if param.requires_grad:
allreduce(param.grad.data)
3. Information including loss and reward should be aggregated among devices when applying data-parallel training. DI-engine achieves this with allreduce operator in learner and evaluator, and only saves log files on process with rank 0.
For more related functions, please refer to ding/utils/pytorch_ddp_dist_helper.py
Usage¶
To enable DDP training in DI-engine existing codes, you just need to add modifications by following steps:
Set
config.policy.learn.multi_gpu
asTrue
Add DDP training context liks this:
from ding.utils import DistContext
from ding.entry import serial_pipeline
# define main_config and create_config
main_config = (...)
create_config = (...)
if __name__ == "__main__":
# call serial_pipeline with DDP
with DistContext():
serial_pipeline(main_config, create_config)
Tip
The whole example is located in dizoo/atari/entry/spaceinvaders_dqn_config_multi_gpu_ddp.py
Execute launch shell script
For DDP, the runnable script demo is demonstrated as follows.
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=2 spaceinvaders_dqn_config_multi_gpu_ddp.py
Or on cluster managed by Slurm
srun -p PARTITION_NAME --mpi=pmi2 --gres=gpu:2 -n2 --ntasks-per-node=2 python -u spaceinvaders_dqn_config_multi_gpu_ddp.py