run.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765
  1. #!/usr/bin/env python3
  2. # Copyright (c) Facebook, Inc. and its affiliates.
  3. # All rights reserved.
  4. #
  5. # This source code is licensed under the BSD-style license found in the
  6. # LICENSE file in the root directory of this source tree.
  7. """
  8. ``torchrun`` provides a superset of the functionality as ``torch.distributed.launch``
  9. with the following additional functionalities:
  10. 1. Worker failures are handled gracefully by restarting all workers.
  11. 2. Worker ``RANK`` and ``WORLD_SIZE`` are assigned automatically.
  12. 3. Number of nodes is allowed to change between minimum and maximum sizes (elasticity).
  13. .. note:: ``torchrun`` is a python
  14. `console script <https://packaging.python.org/en/latest/specifications/entry-points/#use-for-scripts>`_
  15. to the main module
  16. `torch.distributed.run <https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py>`_
  17. declared in the ``entry_points`` configuration in
  18. `setup.py <https://github.com/pytorch/pytorch/blob/master/setup.py>`_.
  19. It is equivalent to invoking ``python -m torch.distributed.run``.
  20. Transitioning from torch.distributed.launch to torchrun
  21. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  22. ``torchrun`` supports the same arguments as ``torch.distributed.launch`` **except**
  23. for ``--use_env`` which is now deprecated. To migrate from ``torch.distributed.launch``
  24. to ``torchrun`` follow these steps:
  25. 1. If your training script is already reading ``local_rank`` from the ``LOCAL_RANK`` environment variable.
  26. Then you need simply omit the ``--use_env`` flag, e.g.:
  27. +--------------------------------------------------------------------+--------------------------------------------+
  28. | ``torch.distributed.launch`` | ``torchrun`` |
  29. +====================================================================+============================================+
  30. | | |
  31. | .. code-block:: shell-session | .. code-block:: shell-session |
  32. | | |
  33. | $ python -m torch.distributed.launch --use_env train_script.py | $ torchrun train_script.py |
  34. | | |
  35. +--------------------------------------------------------------------+--------------------------------------------+
  36. 2. If your training script reads local rank from a ``--local_rank`` cmd argument.
  37. Change your training script to read from the ``LOCAL_RANK`` environment variable as
  38. demonstrated by the following code snippet:
  39. +-------------------------------------------------------+----------------------------------------------------+
  40. | ``torch.distributed.launch`` | ``torchrun`` |
  41. +=======================================================+====================================================+
  42. | | |
  43. | .. code-block:: python | .. code-block:: python |
  44. | | |
  45. | | |
  46. | import argparse | import os |
  47. | parser = argparse.ArgumentParser() | local_rank = int(os.environ["LOCAL_RANK"]) |
  48. | parser.add_argument("--local_rank", type=int) | |
  49. | args = parser.parse_args() | |
  50. | | |
  51. | local_rank = args.local_rank | |
  52. | | |
  53. +-------------------------------------------------------+----------------------------------------------------+
  54. The aformentioned changes suffice to migrate from ``torch.distributed.launch`` to ``torchrun``.
  55. To take advantage of new features such as elasticity, fault-tolerance, and error reporting of ``torchrun``
  56. please refer to:
  57. * :ref:`elastic_train_script` for more information on authoring training scripts that are ``torchrun`` compliant.
  58. * the rest of this page for more information on the features of ``torchrun``.
  59. Usage
  60. --------
  61. Single-node multi-worker
  62. ++++++++++++++++++++++++++++++
  63. ::
  64. >>> torchrun
  65. --standalone
  66. --nnodes=1
  67. --nproc_per_node=$NUM_TRAINERS
  68. YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
  69. Stacked single-node multi-worker
  70. +++++++++++++++++++++++++++++++++++
  71. To run multiple instances (separate jobs) of single-node, multi-worker on the
  72. same host, we need to make sure that each instance (job) is
  73. setup on different ports to avoid port conflicts (or worse, two jobs being merged
  74. as a single job). To do this you have to run with ``--rdzv_backend=c10d``
  75. and specify a different port by setting ``--rdzv_endpoint=localhost:$PORT_k``.
  76. For ``--nodes=1``, its often convenient to let ``torchrun`` pick a free random
  77. port automatically instead of manually assgining different ports for each run.
  78. ::
  79. >>> torchrun
  80. --rdzv_backend=c10d
  81. --rdzv_endpoint=localhost:0
  82. --nnodes=1
  83. --nproc_per_node=$NUM_TRAINERS
  84. YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
  85. Fault tolerant (fixed sized number of workers, no elasticity, tolerates 3 failures)
  86. ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  87. ::
  88. >>> torchrun
  89. --nnodes=$NUM_NODES
  90. --nproc_per_node=$NUM_TRAINERS
  91. --max_restarts=3
  92. --rdzv_id=$JOB_ID
  93. --rdzv_backend=c10d
  94. --rdzv_endpoint=$HOST_NODE_ADDR
  95. YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
  96. ``HOST_NODE_ADDR``, in form <host>[:<port>] (e.g. node1.example.com:29400), specifies the node and
  97. the port on which the C10d rendezvous backend should be instantiated and hosted. It can be any
  98. node in your training cluster, but ideally you should pick a node that has a high bandwidth.
  99. .. note::
  100. If no port number is specified ``HOST_NODE_ADDR`` defaults to 29400.
  101. Elastic (``min=1``, ``max=4``, tolerates up to 3 membership changes or failures)
  102. +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  103. ::
  104. >>> torchrun
  105. --nnodes=1:4
  106. --nproc_per_node=$NUM_TRAINERS
  107. --max_restarts=3
  108. --rdzv_id=$JOB_ID
  109. --rdzv_backend=c10d
  110. --rdzv_endpoint=$HOST_NODE_ADDR
  111. YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
  112. ``HOST_NODE_ADDR``, in form <host>[:<port>] (e.g. node1.example.com:29400), specifies the node and
  113. the port on which the C10d rendezvous backend should be instantiated and hosted. It can be any
  114. node in your training cluster, but ideally you should pick a node that has a high bandwidth.
  115. .. note::
  116. If no port number is specified ``HOST_NODE_ADDR`` defaults to 29400.
  117. Note on rendezvous backend
  118. ------------------------------
  119. For multi-node training you need to specify:
  120. 1. ``--rdzv_id``: A unique job id (shared by all nodes participating in the job)
  121. 2. ``--rdzv_backend``: An implementation of
  122. :py:class:`torch.distributed.elastic.rendezvous.RendezvousHandler`
  123. 3. ``--rdzv_endpoint``: The endpoint where the rendezvous backend is running; usually in form
  124. ``host:port``.
  125. Currently ``c10d`` (recommended), ``etcd-v2``, and ``etcd`` (legacy) rendezvous backends are
  126. supported out of the box. To use ``etcd-v2`` or ``etcd``, setup an etcd server with the ``v2`` api
  127. enabled (e.g. ``--enable-v2``).
  128. .. warning::
  129. ``etcd-v2`` and ``etcd`` rendezvous use etcd API v2. You MUST enable the v2 API on the etcd
  130. server. Our tests use etcd v3.4.3.
  131. .. warning::
  132. For etcd-based rendezvous we recommend using ``etcd-v2`` over ``etcd`` which is functionally
  133. equivalent, but uses a revised implementation. ``etcd`` is in maintenance mode and will be
  134. removed in a future version.
  135. Definitions
  136. --------------
  137. 1. ``Node`` - A physical instance or a container; maps to the unit that the job manager works with.
  138. 2. ``Worker`` - A worker in the context of distributed training.
  139. 3. ``WorkerGroup`` - The set of workers that execute the same function (e.g. trainers).
  140. 4. ``LocalWorkerGroup`` - A subset of the workers in the worker group running on the same node.
  141. 5. ``RANK`` - The rank of the worker within a worker group.
  142. 6. ``WORLD_SIZE`` - The total number of workers in a worker group.
  143. 7. ``LOCAL_RANK`` - The rank of the worker within a local worker group.
  144. 8. ``LOCAL_WORLD_SIZE`` - The size of the local worker group.
  145. 9. ``rdzv_id`` - A user-defined id that uniquely identifies the worker group for a job. This id is
  146. used by each node to join as a member of a particular worker group.
  147. 9. ``rdzv_backend`` - The backend of the rendezvous (e.g. ``c10d``). This is typically a strongly
  148. consistent key-value store.
  149. 10. ``rdzv_endpoint`` - The rendezvous backend endpoint; usually in form ``<host>:<port>``.
  150. A ``Node`` runs ``LOCAL_WORLD_SIZE`` workers which comprise a ``LocalWorkerGroup``. The union of
  151. all ``LocalWorkerGroups`` in the nodes in the job comprise the ``WorkerGroup``.
  152. Environment Variables
  153. ----------------------
  154. The following environment variables are made available to you in your script:
  155. 1. ``LOCAL_RANK`` - The local rank.
  156. 2. ``RANK`` - The global rank.
  157. 3. ``GROUP_RANK`` - The rank of the worker group. A number between 0 and ``max_nnodes``. When
  158. running a single worker group per node, this is the rank of the node.
  159. 4. ``ROLE_RANK`` - The rank of the worker across all the workers that have the same role. The role
  160. of the worker is specified in the ``WorkerSpec``.
  161. 5. ``LOCAL_WORLD_SIZE`` - The local world size (e.g. number of workers running locally); equals to
  162. ``--nproc_per_node`` specified on ``torchrun``.
  163. 6. ``WORLD_SIZE`` - The world size (total number of workers in the job).
  164. 7. ``ROLE_WORLD_SIZE`` - The total number of workers that was launched with the same role specified
  165. in ``WorkerSpec``.
  166. 8. ``MASTER_ADDR`` - The FQDN of the host that is running worker with rank 0; used to initialize
  167. the Torch Distributed backend.
  168. 9. ``MASTER_PORT`` - The port on the ``MASTER_ADDR`` that can be used to host the C10d TCP store.
  169. 10. ``TORCHELASTIC_RESTART_COUNT`` - The number of worker group restarts so far.
  170. 11. ``TORCHELASTIC_MAX_RESTARTS`` - The configured maximum number of restarts.
  171. 12. ``TORCHELASTIC_RUN_ID`` - Equal to the rendezvous ``run_id`` (e.g. unique job id).
  172. 13. ``PYTHON_EXEC`` - System executable override. If provided, the python user script will
  173. use the value of ``PYTHON_EXEC`` as executable. The `sys.executable` is used by default.
  174. Deployment
  175. ------------
  176. 1. (Not needed for the C10d backend) Start the rendezvous backend server and get the endpoint (to be
  177. passed as ``--rdzv_endpoint`` to the launcher script)
  178. 2. Single-node multi-worker: Start the launcher on the host to start the agent process which
  179. creates and monitors a local worker group.
  180. 3. Multi-node multi-worker: Start the launcher with the same arguments on all the nodes
  181. participating in training.
  182. When using a job/cluster manager the entry point command to the multi-node job should be this
  183. launcher.
  184. Failure Modes
  185. ---------------
  186. 1. Worker failure: For a training job with ``n`` workers, if ``k<=n`` workers fail all workers
  187. are stopped and restarted up to ``max_restarts``.
  188. 2. Agent failure: An agent failure results in a local worker group failure. It is up to the job
  189. manager to fail the entire job (gang semantics) or attempt to replace the node. Both behaviors
  190. are supported by the agent.
  191. 3. Node failure: Same as agent failure.
  192. Membership Changes
  193. --------------------
  194. 1. Node departure (scale-down): The agent is notified of the departure, all existing workers are
  195. stopped, a new ``WorkerGroup`` is formed, and all workers are started with a new ``RANK`` and
  196. ``WORLD_SIZE``.
  197. 2. Node arrival (scale-up): The new node is admitted to the job, all existing workers are stopped,
  198. a new ``WorkerGroup`` is formed, and all workers are started with a new ``RANK`` and
  199. ``WORLD_SIZE``.
  200. Important Notices
  201. --------------------
  202. 1. This utility and multi-process distributed (single-node or
  203. multi-node) GPU training currently only achieves the best performance using
  204. the NCCL distributed backend. Thus NCCL backend is the recommended backend to
  205. use for GPU training.
  206. 2. The environment variables necessary to initialize a Torch process group are provided to you by
  207. this module, no need for you to pass ``RANK`` manually. To initialize a process group in your
  208. training script, simply run:
  209. ::
  210. >>> import torch.distributed as dist
  211. >>> dist.init_process_group(backend="gloo|nccl")
  212. 3. In your training program, you can either use regular distributed functions
  213. or use :func:`torch.nn.parallel.DistributedDataParallel` module. If your
  214. training program uses GPUs for training and you would like to use
  215. :func:`torch.nn.parallel.DistributedDataParallel` module,
  216. here is how to configure it.
  217. ::
  218. local_rank = int(os.environ["LOCAL_RANK"])
  219. model = torch.nn.parallel.DistributedDataParallel(model,
  220. device_ids=[local_rank],
  221. output_device=local_rank)
  222. Please ensure that ``device_ids`` argument is set to be the only GPU device id
  223. that your code will be operating on. This is generally the local rank of the
  224. process. In other words, the ``device_ids`` needs to be ``[int(os.environ("LOCAL_RANK"))]``,
  225. and ``output_device`` needs to be ``int(os.environ("LOCAL_RANK"))`` in order to use this
  226. utility
  227. 4. On failures or membership changes ALL surviving workers are killed immediately. Make sure to
  228. checkpoint your progress. The frequency of checkpoints should depend on your job's tolerance
  229. for lost work.
  230. 5. This module only supports homogeneous ``LOCAL_WORLD_SIZE``. That is, it is assumed that all
  231. nodes run the same number of local workers (per role).
  232. 6. ``RANK`` is NOT stable. Between restarts, the local workers on a node can be assgined a
  233. different range of ranks than before. NEVER hard code any assumptions about the stable-ness of
  234. ranks or some correlation between ``RANK`` and ``LOCAL_RANK``.
  235. 7. When using elasticity (``min_size!=max_size``) DO NOT hard code assumptions about
  236. ``WORLD_SIZE`` as the world size can change as nodes are allowed to leave and join.
  237. 8. It is recommended for your script to have the following structure:
  238. ::
  239. def main():
  240. load_checkpoint(checkpoint_path)
  241. initialize()
  242. train()
  243. def train():
  244. for batch in iter(dataset):
  245. train_step(batch)
  246. if should_checkpoint:
  247. save_checkpoint(checkpoint_path)
  248. 9. (Recommended) On worker errors, this tool will summarize the details of the error
  249. (e.g. time, rank, host, pid, traceback, etc). On each node, the first error (by timestamp)
  250. is heuristically reported as the "Root Cause" error. To get tracebacks as part of this
  251. error summary print out, you must decorate your main entrypoint function in your
  252. training script as shown in the example below. If not decorated, then the summary
  253. will not include the traceback of the exception and will only contain the exitcode.
  254. For details on torchelastic error handling see: https://pytorch.org/docs/stable/elastic/errors.html
  255. ::
  256. from torch.distributed.elastic.multiprocessing.errors import record
  257. @record
  258. def main():
  259. # do train
  260. pass
  261. if __name__ == "__main__":
  262. main()
  263. """
  264. import logging
  265. import os
  266. import sys
  267. import uuid
  268. from argparse import REMAINDER, ArgumentParser
  269. from typing import Callable, List, Tuple, Union
  270. import torch
  271. from torch.distributed.argparse_util import check_env, env
  272. from torch.distributed.elastic.multiprocessing import Std
  273. from torch.distributed.elastic.multiprocessing.errors import record
  274. from torch.distributed.elastic.rendezvous.utils import _parse_rendezvous_config
  275. from torch.distributed.elastic.utils import macros
  276. from torch.distributed.elastic.utils.logging import get_logger
  277. from torch.distributed.launcher.api import LaunchConfig, elastic_launch
  278. log = get_logger()
  279. def get_args_parser() -> ArgumentParser:
  280. """Helper function parsing the command line options."""
  281. parser = ArgumentParser(description="Torch Distributed Elastic Training Launcher")
  282. #
  283. # Worker/node size related arguments.
  284. #
  285. parser.add_argument(
  286. "--nnodes",
  287. action=env,
  288. type=str,
  289. default="1:1",
  290. help="Number of nodes, or the range of nodes in form <minimum_nodes>:<maximum_nodes>.",
  291. )
  292. parser.add_argument(
  293. "--nproc_per_node",
  294. action=env,
  295. type=str,
  296. default="1",
  297. help="Number of workers per node; supported values: [auto, cpu, gpu, int].",
  298. )
  299. #
  300. # Rendezvous related arguments
  301. #
  302. parser.add_argument(
  303. "--rdzv_backend",
  304. action=env,
  305. type=str,
  306. default="static",
  307. help="Rendezvous backend.",
  308. )
  309. parser.add_argument(
  310. "--rdzv_endpoint",
  311. action=env,
  312. type=str,
  313. default="",
  314. help="Rendezvous backend endpoint; usually in form <host>:<port>.",
  315. )
  316. parser.add_argument(
  317. "--rdzv_id",
  318. action=env,
  319. type=str,
  320. default="none",
  321. help="User-defined group id.",
  322. )
  323. parser.add_argument(
  324. "--rdzv_conf",
  325. action=env,
  326. type=str,
  327. default="",
  328. help="Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).",
  329. )
  330. parser.add_argument(
  331. "--standalone",
  332. action=check_env,
  333. help="Start a local standalone rendezvous backend that is represented by a C10d TCP store "
  334. "on port 29400. Useful when launching single-node, multi-worker job. If specified "
  335. "--rdzv_backend, --rdzv_endpoint, --rdzv_id are auto-assigned; any explicitly set values "
  336. "are ignored.",
  337. )
  338. #
  339. # User-code launch related arguments.
  340. #
  341. parser.add_argument(
  342. "--max_restarts",
  343. action=env,
  344. type=int,
  345. default=0,
  346. help="Maximum number of worker group restarts before failing.",
  347. )
  348. parser.add_argument(
  349. "--monitor_interval",
  350. action=env,
  351. type=float,
  352. default=5,
  353. help="Interval, in seconds, to monitor the state of workers.",
  354. )
  355. parser.add_argument(
  356. "--start_method",
  357. action=env,
  358. type=str,
  359. default="spawn",
  360. choices=["spawn", "fork", "forkserver"],
  361. help="Multiprocessing start method to use when creating workers.",
  362. )
  363. parser.add_argument(
  364. "--role",
  365. action=env,
  366. type=str,
  367. default="default",
  368. help="User-defined role for the workers.",
  369. )
  370. parser.add_argument(
  371. "-m",
  372. "--module",
  373. action=check_env,
  374. help="Change each process to interpret the launch script as a Python module, executing "
  375. "with the same behavior as 'python -m'.",
  376. )
  377. parser.add_argument(
  378. "--no_python",
  379. action=check_env,
  380. help="Skip prepending the training script with 'python' - just execute it directly. Useful "
  381. "when the script is not a Python script.",
  382. )
  383. parser.add_argument(
  384. "--run_path",
  385. action=check_env,
  386. help="Run the training script with runpy.run_path in the same interpreter."
  387. " Script must be provided as an abs path (e.g. /abs/path/script.py)."
  388. " Takes precedence over --no_python.",
  389. )
  390. parser.add_argument(
  391. "--log_dir",
  392. action=env,
  393. type=str,
  394. default=None,
  395. help="Base directory to use for log files (e.g. /var/log/torch/elastic). The same "
  396. "directory is re-used for multiple runs (a unique job-level sub-directory is created with "
  397. "rdzv_id as the prefix).",
  398. )
  399. parser.add_argument(
  400. "-r",
  401. "--redirects",
  402. action=env,
  403. type=str,
  404. default="0",
  405. help="Redirect std streams into a log file in the log directory (e.g. [-r 3] redirects "
  406. "both stdout+stderr for all workers, [-r 0:1,1:2] redirects stdout for local rank 0 and "
  407. "stderr for local rank 1).",
  408. )
  409. parser.add_argument(
  410. "-t",
  411. "--tee",
  412. action=env,
  413. type=str,
  414. default="0",
  415. help="Tee std streams into a log file and also to console (see --redirects for format).",
  416. )
  417. #
  418. # Backwards compatible parameters with caffe2.distributed.launch.
  419. #
  420. parser.add_argument(
  421. "--node_rank",
  422. type=int,
  423. action=env,
  424. default=0,
  425. help="Rank of the node for multi-node distributed training.",
  426. )
  427. parser.add_argument(
  428. "--master_addr",
  429. default="127.0.0.1",
  430. type=str,
  431. action=env,
  432. help="Address of the master node (rank 0). It should be either the IP address or the "
  433. "hostname of rank 0. For single node multi-proc training the --master_addr can simply be "
  434. "127.0.0.1; IPv6 should have the pattern `[0:0:0:0:0:0:0:1]`.",
  435. )
  436. parser.add_argument(
  437. "--master_port",
  438. default=29500,
  439. type=int,
  440. action=env,
  441. help="Port on the master node (rank 0) to be used for communication during distributed "
  442. "training.",
  443. )
  444. #
  445. # Positional arguments.
  446. #
  447. parser.add_argument(
  448. "training_script",
  449. type=str,
  450. help="Full path to the (single GPU) training program/script to be launched in parallel, "
  451. "followed by all the arguments for the training script.",
  452. )
  453. # Rest from the training program.
  454. parser.add_argument("training_script_args", nargs=REMAINDER)
  455. return parser
  456. def parse_args(args):
  457. parser = get_args_parser()
  458. return parser.parse_args(args)
  459. def parse_min_max_nnodes(nnodes: str):
  460. arr = nnodes.split(":")
  461. if len(arr) == 1:
  462. min_nodes = max_nodes = int(arr[0])
  463. elif len(arr) == 2:
  464. min_nodes = int(arr[0])
  465. max_nodes = int(arr[1])
  466. else:
  467. raise RuntimeError(f'nnodes={nnodes} is not in "MIN:MAX" format')
  468. return min_nodes, max_nodes
  469. def determine_local_world_size(nproc_per_node: str):
  470. try:
  471. logging.info(f"Using nproc_per_node={nproc_per_node}.")
  472. return int(nproc_per_node)
  473. except ValueError:
  474. if nproc_per_node == "cpu":
  475. num_proc = os.cpu_count()
  476. device_type = "cpu"
  477. elif nproc_per_node == "gpu":
  478. if not torch.cuda.is_available():
  479. raise ValueError("Cuda is not available.")
  480. device_type = "gpu"
  481. num_proc = torch.cuda.device_count()
  482. elif nproc_per_node == "auto":
  483. if torch.cuda.is_available():
  484. num_proc = torch.cuda.device_count()
  485. device_type = "gpu"
  486. else:
  487. num_proc = os.cpu_count()
  488. device_type = "cpu"
  489. else:
  490. raise ValueError(f"Unsupported nproc_per_node value: {nproc_per_node}")
  491. log.info(
  492. f"Using nproc_per_node={nproc_per_node},"
  493. f" seting to {num_proc} since the instance "
  494. f"has {os.cpu_count()} {device_type}"
  495. )
  496. return num_proc
  497. def get_rdzv_endpoint(args):
  498. if args.rdzv_backend == "static" and not args.rdzv_endpoint:
  499. return f"{args.master_addr}:{args.master_port}"
  500. return args.rdzv_endpoint
  501. def get_use_env(args) -> bool:
  502. """
  503. Retrieves ``use_env`` from the args.
  504. ``use_env`` is a legacy argument, if ``use_env`` is False, the
  505. ``--node_rank`` argument will be transferred to all worker processes.
  506. ``use_env`` is only used by the ``torch.distributed.launch`` and will
  507. be deprecated in future releases.
  508. """
  509. if not hasattr(args, "use_env"):
  510. return True
  511. return args.use_env
  512. def config_from_args(args) -> Tuple[LaunchConfig, Union[Callable, str], List[str]]:
  513. # If ``args`` not passed, defaults to ``sys.argv[:1]``
  514. min_nodes, max_nodes = parse_min_max_nnodes(args.nnodes)
  515. assert 0 < min_nodes <= max_nodes
  516. assert args.max_restarts >= 0
  517. nproc_per_node = determine_local_world_size(args.nproc_per_node)
  518. if "OMP_NUM_THREADS" not in os.environ and nproc_per_node > 1:
  519. omp_num_threads = 1
  520. log.warning(
  521. f"\n*****************************************\n"
  522. f"Setting OMP_NUM_THREADS environment variable for each process to be "
  523. f"{omp_num_threads} in default, to avoid your system being overloaded, "
  524. f"please further tune the variable for optimal performance in "
  525. f"your application as needed. \n"
  526. f"*****************************************"
  527. )
  528. # This env variable will be passed down to the subprocesses
  529. os.environ["OMP_NUM_THREADS"] = str(omp_num_threads)
  530. rdzv_configs = _parse_rendezvous_config(args.rdzv_conf)
  531. if args.rdzv_backend == "static":
  532. rdzv_configs["rank"] = args.node_rank
  533. rdzv_endpoint = get_rdzv_endpoint(args)
  534. config = LaunchConfig(
  535. min_nodes=min_nodes,
  536. max_nodes=max_nodes,
  537. nproc_per_node=nproc_per_node,
  538. run_id=args.rdzv_id,
  539. role=args.role,
  540. rdzv_endpoint=rdzv_endpoint,
  541. rdzv_backend=args.rdzv_backend,
  542. rdzv_configs=rdzv_configs,
  543. max_restarts=args.max_restarts,
  544. monitor_interval=args.monitor_interval,
  545. start_method=args.start_method,
  546. redirects=Std.from_str(args.redirects),
  547. tee=Std.from_str(args.tee),
  548. log_dir=args.log_dir,
  549. )
  550. with_python = not args.no_python
  551. cmd: Union[Callable, str]
  552. cmd_args = []
  553. use_env = get_use_env(args)
  554. if args.run_path:
  555. cmd = run_script_path
  556. cmd_args.append(args.training_script)
  557. else:
  558. if with_python:
  559. cmd = os.getenv("PYTHON_EXEC", sys.executable)
  560. cmd_args.append("-u")
  561. if args.module:
  562. cmd_args.append("-m")
  563. cmd_args.append(args.training_script)
  564. else:
  565. if args.module:
  566. raise ValueError(
  567. "Don't use both the '--no_python' flag"
  568. " and the '--module' flag at the same time."
  569. )
  570. cmd = args.training_script
  571. if not use_env:
  572. cmd_args.append(f"--local_rank={macros.local_rank}")
  573. cmd_args.extend(args.training_script_args)
  574. return config, cmd, cmd_args
  575. def run_script_path(training_script: str, *training_script_args: str):
  576. """
  577. Runs the provided `training_script` from within this interpreter.
  578. Usage: `script_as_function("/abs/path/to/script.py", "--arg1", "val1")`
  579. """
  580. import runpy
  581. import sys
  582. sys.argv = [training_script] + [*training_script_args]
  583. runpy.run_path(sys.argv[0], run_name="__main__")
  584. def run(args):
  585. if args.standalone:
  586. args.rdzv_backend = "c10d"
  587. args.rdzv_endpoint = "localhost:29400"
  588. args.rdzv_id = str(uuid.uuid4())
  589. log.info(
  590. f"\n**************************************\n"
  591. f"Rendezvous info:\n"
  592. f"--rdzv_backend={args.rdzv_backend} "
  593. f"--rdzv_endpoint={args.rdzv_endpoint} "
  594. f"--rdzv_id={args.rdzv_id}\n"
  595. f"**************************************\n"
  596. )
  597. config, cmd, cmd_args = config_from_args(args)
  598. elastic_launch(
  599. config=config,
  600. entrypoint=cmd,
  601. )(*cmd_args)
  602. @record
  603. def main(args=None):
  604. args = parse_args(args)
  605. run(args)
  606. if __name__ == "__main__":
  607. main()