Add support for distributed training (multi-node and multi-GPU) via PyTorch DDP #2018
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
AdeelH
force-pushed
the
multi-gpu2
branch
2 times, most recently
from
609c778
to
355b299
Compare
4 tasks
AdeelH
force-pushed
the
multi-gpu2
branch
3 times, most recently
from
fa80f93
to
b60f8ed
Compare
AdeelH
changed the title
4 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Overview
This PR adds support for distributed training (multi-node and multi-GPU) via PyTorch DDP.
It can be used in the following ways:
Learner.train()is called.torchrunCLI command. For example, to run on a single machine with 4 GPUs:torchrun --standalone --nnodes=1 --nproc-per-node=4 --no-python \ rastervision run local rastervision_pytorch_backend/rastervision/pytorch_backend/examples/tiny_spacenet.pyChecklist
needs-backportlabel if the change should be back-ported to the previous releaseNotes
Learner.eval_model()has been renamed toLearner.validate().RASTERVISION_USE_DDP:YESby default. Set toNOto disable distributed training.RASTERVISION_DDP_BACKEND:ncclby default. This is the recommended backend for CUDA GPUs.RASTERVISION_DDP_START_METHOD: One ofspawn,fork, orforkserver. Default:spawn.spawnis what PyTorch documentation recommends (in fact, it doesn't even mention the alternatives), but it has the disadvantage that it requires everything to be pickleable, which rasterio dataset objects are not. This is also true forforkserver, which needs to spawn a server process. However,forkdoes not have the same limitation.fork, we avoid building the dataset in the base process and instead delay it until the worker processes are created.forkorforkserver, the CUDA runtime must not be initialized before the fork happens; otherwise, aRuntimeError: Cannot re-initialize CUDA in forked subprocess.error will be raised. We avoid this by not calling anytorch.cudafunctions or creating tensors on the GPU.TORCH_HOMEto$TORCH_HOME/<local rank>. And only the master process copies the downloaded files to the training directory.WORLD_SIZEenv var.Testing Instructions
torchrun --standalone --nnodes=1 --nproc-per-node=1 --no-python \ rastervision run local rastervision_pytorch_backend/rastervision/pytorch_backend/examples/tiny_spacenet.py# pragma: no cover) since the CI machine does not have a GPU.Sample output from running
isprs_potsdam.pyon a p3.8xlarge instance.Closes #1352