launch.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  1. r"""
  2. ``torch.distributed.launch`` is a module that spawns up multiple distributed
  3. training processes on each of the training nodes.
  4. .. warning::
  5. This module is going to be deprecated in favor of :ref:`torchrun <launcher-api>`.
  6. The utility can be used for single-node distributed training, in which one or
  7. more processes per node will be spawned. The utility can be used for either
  8. CPU training or GPU training. If the utility is used for GPU training,
  9. each distributed process will be operating on a single GPU. This can achieve
  10. well-improved single-node training performance. It can also be used in
  11. multi-node distributed training, by spawning up multiple processes on each node
  12. for well-improved multi-node distributed training performance as well.
  13. This will especially be benefitial for systems with multiple Infiniband
  14. interfaces that have direct-GPU support, since all of them can be utilized for
  15. aggregated communication bandwidth.
  16. In both cases of single-node distributed training or multi-node distributed
  17. training, this utility will launch the given number of processes per node
  18. (``--nproc_per_node``). If used for GPU training, this number needs to be less
  19. or equal to the number of GPUs on the current system (``nproc_per_node``),
  20. and each process will be operating on a single GPU from *GPU 0 to
  21. GPU (nproc_per_node - 1)*.
  22. **How to use this module:**
  23. 1. Single-Node multi-process distributed training
  24. ::
  25. >>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
  26. YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
  27. arguments of your training script)
  28. 2. Multi-Node multi-process distributed training: (e.g. two nodes)
  29. Node 1: *(IP: 192.168.1.1, and has a free port: 1234)*
  30. ::
  31. >>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
  32. --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
  33. --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
  34. and all other arguments of your training script)
  35. Node 2:
  36. ::
  37. >>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
  38. --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
  39. --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
  40. and all other arguments of your training script)
  41. 3. To look up what optional arguments this module offers:
  42. ::
  43. >>> python -m torch.distributed.launch --help
  44. **Important Notices:**
  45. 1. This utility and multi-process distributed (single-node or
  46. multi-node) GPU training currently only achieves the best performance using
  47. the NCCL distributed backend. Thus NCCL backend is the recommended backend to
  48. use for GPU training.
  49. 2. In your training program, you must parse the command-line argument:
  50. ``--local_rank=LOCAL_PROCESS_RANK``, which will be provided by this module.
  51. If your training program uses GPUs, you should ensure that your code only
  52. runs on the GPU device of LOCAL_PROCESS_RANK. This can be done by:
  53. Parsing the local_rank argument
  54. ::
  55. >>> import argparse
  56. >>> parser = argparse.ArgumentParser()
  57. >>> parser.add_argument("--local_rank", type=int)
  58. >>> args = parser.parse_args()
  59. Set your device to local rank using either
  60. ::
  61. >>> torch.cuda.set_device(args.local_rank) # before your code runs
  62. or
  63. ::
  64. >>> with torch.cuda.device(args.local_rank):
  65. >>> # your code to run
  66. 3. In your training program, you are supposed to call the following function
  67. at the beginning to start the distributed backend. It is strongly recommended
  68. that ``init_method=env://``. Other init methods (e.g. ``tcp://``) may work,
  69. but ``env://`` is the one that is officially supported by this module.
  70. ::
  71. torch.distributed.init_process_group(backend='YOUR BACKEND',
  72. init_method='env://')
  73. 4. In your training program, you can either use regular distributed functions
  74. or use :func:`torch.nn.parallel.DistributedDataParallel` module. If your
  75. training program uses GPUs for training and you would like to use
  76. :func:`torch.nn.parallel.DistributedDataParallel` module,
  77. here is how to configure it.
  78. ::
  79. model = torch.nn.parallel.DistributedDataParallel(model,
  80. device_ids=[args.local_rank],
  81. output_device=args.local_rank)
  82. Please ensure that ``device_ids`` argument is set to be the only GPU device id
  83. that your code will be operating on. This is generally the local rank of the
  84. process. In other words, the ``device_ids`` needs to be ``[args.local_rank]``,
  85. and ``output_device`` needs to be ``args.local_rank`` in order to use this
  86. utility
  87. 5. Another way to pass ``local_rank`` to the subprocesses via environment variable
  88. ``LOCAL_RANK``. This behavior is enabled when you launch the script with
  89. ``--use_env=True``. You must adjust the subprocess example above to replace
  90. ``args.local_rank`` with ``os.environ['LOCAL_RANK']``; the launcher
  91. will not pass ``--local_rank`` when you specify this flag.
  92. .. warning::
  93. ``local_rank`` is NOT globally unique: it is only unique per process
  94. on a machine. Thus, don't use it to decide if you should, e.g.,
  95. write to a networked filesystem. See
  96. https://github.com/pytorch/pytorch/issues/12042 for an example of
  97. how things can go wrong if you don't do this correctly.
  98. """
  99. import logging
  100. import warnings
  101. from torch.distributed.run import get_args_parser, run
  102. logger = logging.getLogger(__name__)
  103. def parse_args(args):
  104. parser = get_args_parser()
  105. parser.add_argument(
  106. "--use_env",
  107. default=False,
  108. action="store_true",
  109. help="Use environment variable to pass "
  110. "'local rank'. For legacy reasons, the default value is False. "
  111. "If set to True, the script will not pass "
  112. "--local_rank as argument, and will instead set LOCAL_RANK.",
  113. )
  114. return parser.parse_args(args)
  115. def launch(args):
  116. if args.no_python and not args.use_env:
  117. raise ValueError(
  118. "When using the '--no_python' flag,"
  119. " you must also set the '--use_env' flag."
  120. )
  121. run(args)
  122. def main(args=None):
  123. warnings.warn(
  124. "The module torch.distributed.launch is deprecated\n"
  125. "and will be removed in future. Use torchrun.\n"
  126. "Note that --use_env is set by default in torchrun.\n"
  127. "If your script expects `--local_rank` argument to be set, please\n"
  128. "change it to read from `os.environ['LOCAL_RANK']` instead. See \n"
  129. "https://pytorch.org/docs/stable/distributed.html#launch-utility for \n"
  130. "further instructions\n",
  131. FutureWarning,
  132. )
  133. args = parse_args(args)
  134. launch(args)
  135. if __name__ == "__main__":
  136. main()