Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Gen ddp support (#1961) 

* updated docs

* added mixed

* added mixed
This commit is contained in:
William Falcon 2020-05-26 19:02:30 -04:00 committed by GitHub
parent c967b88fc8
commit 460ab5485e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
pytorch_lightning/trainer/trainer.py
View File

@ -856,17 +856,21 @@ class Trainer(
if self.use_ddp2:
if self.is_slurm_managing_tasks:
task = int(os.environ['SLURM_LOCALID'])
elif 'WORLD_SIZE' in os.environ and 'GROUP_RANK' in os.environ:
# torchelastic or general non_slurm ddp2
elif 'WORLD_SIZE' in os.environ and ('GROUP_RANK' in os.environ or 'NODE_RANK' in os.environ):
task = int(os.environ['LOCAL_RANK'])
self.ddp_train(task, model)
elif self.use_ddp:
if self.is_slurm_managing_tasks:
task = int(os.environ['SLURM_LOCALID'])
self.ddp_train(task, model)
# torchelastic
elif 'WORLD_SIZE' in os.environ and 'GROUP_RANK' in os.environ:
# torchelastic or general non_slurm ddp
elif 'WORLD_SIZE' in os.environ and ('GROUP_RANK' in os.environ or 'NODE_RANK' in os.environ):
task = int(os.environ['LOCAL_RANK'])
self.ddp_train(task, model)
else:
self.__set_random_port()
# track for predict