mirror of
https://github.com/twitter/the-algorithm-ml.git
synced 2024-06-28 04:06:09 +02:00
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
from typing import Optional
|
|
|
|
from tml.ml_logging.torch_logging import logging # type: ignore[attr-defined]
|
|
from twitter.ml.tensorflow.experimental.distributed import utils
|
|
|
|
import torch
|
|
import torch.distributed.run
|
|
|
|
|
|
def is_distributed_worker():
|
|
"""
|
|
Checks if the current process is a distributed worker.
|
|
|
|
Returns:
|
|
bool: True if the necessary distributed PyTorch environment variables (WORLD_SIZE, RANK) are set, else False.
|
|
"""
|
|
world_size = os.environ.get("WORLD_SIZE", None)
|
|
rank = os.environ.get("RANK", None)
|
|
return world_size is not None and rank is not None
|
|
|
|
|
|
def maybe_run_training(
|
|
train_fn,
|
|
module_name,
|
|
nproc_per_node: Optional[int] = None,
|
|
num_nodes: Optional[int] = None,
|
|
set_python_path_in_subprocess: bool = False,
|
|
is_chief: Optional[bool] = False,
|
|
**training_kwargs,
|
|
):
|
|
"""
|
|
Wrapper function for single node, multi-GPU PyTorch training.
|
|
|
|
If the necessary distributed PyTorch environment variables (WORLD_SIZE, RANK) have been set, then this function executes
|
|
`train_fn(**training_kwargs)`.
|
|
|
|
Otherwise, this function calls torchrun and points at the calling module
|
|
`module_name`. After this call, the necessary environment variables are set
|
|
and training will commence.
|
|
|
|
Args:
|
|
train_fn (callable): The function responsible for training.
|
|
module_name (str): The name of the module that this function was called from; used to indicate torchrun entrypoint.
|
|
nproc_per_node (int, optional): Number of workers per node. Defaults to None.
|
|
num_nodes (int, optional): Number of nodes. Defaults to None.
|
|
is_chief (bool, optional): If the process is running on the chief node. Defaults to False.
|
|
set_python_path_in_subprocess (bool, optional): Whether to set PYTHONPATH in the subprocess. Defaults to False.
|
|
**training_kwargs: Additional keyword arguments to pass to the `train_fn`.
|
|
|
|
Note:
|
|
This function checks if the current process is a distributed worker by examining the environment variables.
|
|
If it is a worker, it directly calls `train_fn(**training_kwargs)`. Otherwise, it sets up the necessary
|
|
environment variables and launches the training process using torchrun.
|
|
|
|
Example:
|
|
To run training on a single node with 4 GPUs, you can use:
|
|
```
|
|
maybe_run_training(train_function, __name__, nproc_per_node=4)
|
|
```
|
|
"""
|
|
|
|
machines = utils.machine_from_env()
|
|
if num_nodes is None:
|
|
num_nodes = 1
|
|
if machines.num_workers:
|
|
num_nodes += machines.num_workers
|
|
|
|
if is_distributed_worker():
|
|
# world_size, rank, etc are set; assuming any other env vars are set (checks to come)
|
|
# start the actual training!
|
|
train_fn(**training_kwargs)
|
|
else:
|
|
if nproc_per_node is None:
|
|
if torch.cuda.is_available():
|
|
nproc_per_node = torch.cuda.device_count()
|
|
else:
|
|
nproc_per_node = machines.chief.num_accelerators
|
|
|
|
# Rejoin all arguments to send back through torchrec
|
|
# this is a temporary measure, will replace the os.system call
|
|
# with torchrun API calls
|
|
args = list(f"--{key}={val}" for key, val in training_kwargs.items())
|
|
|
|
cmd = [
|
|
"--nnodes",
|
|
str(num_nodes),
|
|
]
|
|
if nproc_per_node:
|
|
cmd.extend(["--nproc_per_node", str(nproc_per_node)])
|
|
if num_nodes > 1:
|
|
cluster_resolver = utils.cluster_resolver()
|
|
backend_address = cluster_resolver.cluster_spec().task_address("chief", 0)
|
|
cmd.extend(
|
|
[
|
|
"--rdzv_backend",
|
|
"c10d",
|
|
"--rdzv_id",
|
|
backend_address,
|
|
]
|
|
)
|
|
# Set localhost on chief because of https://github.com/pytorch/pytorch/issues/79388
|
|
if is_chief:
|
|
cmd.extend(["--rdzv_endpoint", "localhost:2222"])
|
|
else:
|
|
cmd.extend(["--rdzv_endpoint", backend_address])
|
|
else:
|
|
cmd.append("--standalone")
|
|
|
|
cmd.extend(
|
|
[
|
|
str(module_name),
|
|
*args,
|
|
]
|
|
)
|
|
logging.info(f"""Distributed running with cmd: '{" ".join(cmd)}'""")
|
|
|
|
# Call torchrun on this module; will spawn new processes and re-run this
|
|
# function, eventually calling "train_fn". The following line sets the PYTHONPATH to accommodate
|
|
# bazel stubbing for the main binary.
|
|
if set_python_path_in_subprocess:
|
|
subprocess.run(["torchrun"] + cmd, env={**os.environ, "PYTHONPATH": ":".join(sys.path)})
|
|
else:
|
|
torch.distributed.run.main(cmd)
|