pytorch-lightning/docs/Trainer/Distributed training.md

Lightning makes multi-gpu training and 16 bit training trivial.

Note:
None of the flags below require changing anything about your lightningModel definition.

---

16-bit mixed precision

16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.
First, install apex (if install fails, look here):

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

then set this use_amp to True.

# DEFAULT
trainer = Trainer(amp_level='O2', use_amp=False)

---

Single-gpu

Make sure you're on a GPU machine.

# set these flags
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# DEFAULT
trainer = Trainer(gpus=[0])

---

multi-gpu

Make sure you're on a GPU machine. You can set as many GPUs as you want. In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.

# set these flags
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"


trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])

---

Multi-node

Multi-node training is easily done by specifying these flags.

# train on 12*8 GPUs
trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)

In addition, make sure to set up your SLURM job correctly via the SlurmClusterObject. In particular, specify the number of tasks per node correctly.

cluster = SlurmCluster(
    hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
    log_path='/some/path/to/save',
)

# configure cluster
cluster.per_experiment_nb_nodes = 12 
cluster.per_experiment_nb_gpus = 8

cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')

---

Self-balancing architecture

Here lightning distributes parts of your module across available GPUs to optimize for speed and memory.

COMING SOON.