chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,54 @@
|
||||
# isort: off
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"PyTorch isn't installed. To install PyTorch, run 'pip install torch'"
|
||||
)
|
||||
# isort: on
|
||||
|
||||
from ray.train.torch.config import TorchConfig
|
||||
from ray.train.torch.torch_checkpoint import TorchCheckpoint
|
||||
from ray.train.torch.torch_trainer import TorchTrainer
|
||||
from ray.train.torch.train_loop_utils import (
|
||||
accelerate,
|
||||
backward,
|
||||
enable_reproducibility,
|
||||
get_device,
|
||||
get_devices,
|
||||
prepare_data_loader,
|
||||
prepare_model,
|
||||
prepare_optimizer,
|
||||
)
|
||||
from ray.train.v2._internal.constants import is_v2_enabled
|
||||
|
||||
if is_v2_enabled():
|
||||
from ray.train.v2.torch.torch_trainer import TorchTrainer # noqa: F811
|
||||
from ray.train.v2.torch.train_loop_utils import ( # noqa: F811
|
||||
accelerate,
|
||||
backward,
|
||||
enable_reproducibility,
|
||||
get_device,
|
||||
get_devices,
|
||||
prepare_data_loader,
|
||||
prepare_model,
|
||||
prepare_optimizer,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"TorchTrainer",
|
||||
"TorchCheckpoint",
|
||||
"TorchConfig",
|
||||
"accelerate",
|
||||
"get_device",
|
||||
"get_devices",
|
||||
"prepare_model",
|
||||
"prepare_optimizer",
|
||||
"prepare_data_loader",
|
||||
"backward",
|
||||
"enable_reproducibility",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,256 @@
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from packaging.version import Version
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._private import ray_constants
|
||||
from ray.air._internal.device_manager import register_custom_torch_dist_backend
|
||||
from ray.exceptions import GetTimeoutError
|
||||
from ray.train._internal.base_worker_group import BaseWorkerGroup
|
||||
from ray.train._internal.utils import get_address_and_port
|
||||
from ray.train.backend import Backend, BackendConfig
|
||||
from ray.train.constants import (
|
||||
DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
|
||||
TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
from ray.util import PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TorchConfigContextManager:
|
||||
def __enter__(self):
|
||||
# Set default cuda device
|
||||
if torch.cuda.is_available():
|
||||
device = ray.train.torch.get_device()
|
||||
if device.type == "cuda":
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
# Propagate exceptions if any
|
||||
return False
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@dataclass
|
||||
class TorchConfig(BackendConfig):
|
||||
"""Configuration for torch process group setup.
|
||||
|
||||
See https://pytorch.org/docs/stable/distributed.html for more info.
|
||||
|
||||
Args:
|
||||
backend: The backend to use for training.
|
||||
See ``torch.distributed.init_process_group`` for more info and
|
||||
valid values.
|
||||
If set to None, nccl will be used if GPUs are requested, else gloo
|
||||
will be used.
|
||||
init_method: The initialization method to use. Either "env"
|
||||
for environment variable initialization or "tcp" for TCP
|
||||
initialization. Defaults to "env".
|
||||
timeout_s: Seconds for process group operations to timeout.
|
||||
"""
|
||||
|
||||
backend: Optional[str] = None
|
||||
init_method: str = "env"
|
||||
timeout_s: int = 1800
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _TorchBackend
|
||||
|
||||
@property
|
||||
def train_func_context(self):
|
||||
return TorchConfigContextManager
|
||||
|
||||
@property
|
||||
def framework(self):
|
||||
return TrainingFramework.TORCH
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
config_dict = {
|
||||
"backend": self.backend,
|
||||
"init_method": self.init_method,
|
||||
"timeout_s": self.timeout_s,
|
||||
}
|
||||
return config_dict
|
||||
|
||||
|
||||
def _is_backend_nccl(backend: str) -> bool:
|
||||
# Check containment because comma separated lists of backends like cpu:gloo,cuda:nccl are supported.
|
||||
return backend == "nccl" or any(
|
||||
item.split(":")[1] == "nccl"
|
||||
for item in backend.split(",")
|
||||
if item.startswith("cuda:")
|
||||
)
|
||||
|
||||
|
||||
def _setup_torch_process_group(
|
||||
backend: str,
|
||||
world_rank: int,
|
||||
world_size: int,
|
||||
init_method: str,
|
||||
timeout_s: int = 1800,
|
||||
):
|
||||
"""Connects the distributed PyTorch backend.
|
||||
|
||||
Args:
|
||||
backend: The backend (nccl, gloo, etc.) to use for training.
|
||||
world_rank: Rank of the current worker.
|
||||
world_size: Number of workers participating in the job.
|
||||
init_method: URL specifying how to initialize the process group.
|
||||
timeout_s: Seconds for process group operations to timeout.
|
||||
"""
|
||||
if world_rank == 0:
|
||||
logger.info(
|
||||
f"Setting up process group for: {init_method} [rank={world_rank}, "
|
||||
f"world_size={world_size}]"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"Setting up process group for: {init_method} [rank={world_rank}, "
|
||||
f"world_size={world_size}]"
|
||||
)
|
||||
logger.debug(f"using {backend}")
|
||||
|
||||
if _is_backend_nccl(backend):
|
||||
# See https://github.com/pytorch/pytorch/blob/c263bd43e8e8502d4726643bc6fd046f0130ac0e/torch/distributed/distributed_c10d.py#L803-L823 # noqa: E501
|
||||
# We do not use TORCH_NCCL_BLOCKING_WAIT due to performance overhead.
|
||||
if Version(torch.__version__) < Version("2.2.0"):
|
||||
TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR = "NCCL_ASYNC_ERROR_HANDLING"
|
||||
TORCH_NCCL_BLOCKING_WAIT_ENV_VAR = "NCCL_BLOCKING_WAIT"
|
||||
else:
|
||||
TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR = "TORCH_NCCL_ASYNC_ERROR_HANDLING"
|
||||
TORCH_NCCL_BLOCKING_WAIT_ENV_VAR = "TORCH_NCCL_BLOCKING_WAIT"
|
||||
if (
|
||||
TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR not in os.environ
|
||||
and TORCH_NCCL_BLOCKING_WAIT_ENV_VAR not in os.environ
|
||||
):
|
||||
logger.debug(
|
||||
f"Setting {TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR}=1 to fail if NCCL collective communication operations are timing out. " # noqa: E501
|
||||
f"To override this behavior, you can set {TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR}=0." # noqa: E501
|
||||
)
|
||||
os.environ[TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR] = "1"
|
||||
elif backend == "hccl":
|
||||
register_custom_torch_dist_backend(backend)
|
||||
|
||||
dist.init_process_group(
|
||||
backend=backend,
|
||||
init_method=init_method,
|
||||
rank=world_rank,
|
||||
world_size=world_size,
|
||||
timeout=timedelta(seconds=timeout_s),
|
||||
)
|
||||
|
||||
|
||||
def _shutdown_torch(destroy_process_group=False):
|
||||
from ray.air._internal.torch_utils import get_devices
|
||||
|
||||
devices = get_devices()
|
||||
if destroy_process_group and dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
if torch.cuda.is_available():
|
||||
for device in devices:
|
||||
if device.type == "cuda":
|
||||
with torch.cuda.device(device):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _set_torch_distributed_env_vars():
|
||||
# Same env vars as in
|
||||
# https://pytorch.org/docs/stable/elastic/run.html#environment-variables
|
||||
from ray.train.torch import get_device
|
||||
|
||||
context = ray.train.get_context()
|
||||
os.environ["LOCAL_RANK"] = str(context.get_local_rank())
|
||||
os.environ["LOCAL_WORLD_SIZE"] = str(context.get_local_world_size())
|
||||
os.environ["NODE_RANK"] = str(context.get_node_rank())
|
||||
os.environ["RANK"] = str(context.get_world_rank())
|
||||
os.environ["WORLD_SIZE"] = str(context.get_world_size())
|
||||
|
||||
# Makes sure Hugging Face Accelerate uses the correct device
|
||||
device = get_device()
|
||||
os.environ["ACCELERATE_TORCH_DEVICE"] = str(device)
|
||||
|
||||
|
||||
class _TorchBackend(Backend):
|
||||
share_cuda_visible_devices: bool = True
|
||||
|
||||
def on_start(self, worker_group: BaseWorkerGroup, backend_config: TorchConfig):
|
||||
if dist.is_available():
|
||||
# Set the appropriate training backend.
|
||||
if backend_config.backend is None:
|
||||
resources = worker_group.get_resources_per_worker()
|
||||
num_gpus_per_worker = resources.get("GPU", 0)
|
||||
|
||||
if num_gpus_per_worker > 0:
|
||||
backend = "nccl"
|
||||
else:
|
||||
backend = "gloo"
|
||||
else:
|
||||
backend = backend_config.backend
|
||||
|
||||
master_addr, master_port = worker_group.execute_single(
|
||||
0, get_address_and_port
|
||||
)
|
||||
if backend_config.init_method == "env":
|
||||
|
||||
def set_env_vars(addr, port):
|
||||
os.environ["MASTER_ADDR"] = addr
|
||||
os.environ["MASTER_PORT"] = str(port)
|
||||
|
||||
worker_group.execute(set_env_vars, addr=master_addr, port=master_port)
|
||||
url = "env://"
|
||||
elif backend_config.init_method == "tcp":
|
||||
url = f"tcp://{build_address(master_addr, master_port)}"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The provided init_method ("
|
||||
f"{backend_config.init_method}) is not supported. Must "
|
||||
f"be either 'env' or 'tcp'."
|
||||
)
|
||||
|
||||
setup_futures = []
|
||||
for i in range(len(worker_group)):
|
||||
setup_futures.append(
|
||||
worker_group.execute_single_async(
|
||||
i,
|
||||
_setup_torch_process_group,
|
||||
backend=backend,
|
||||
world_rank=i,
|
||||
world_size=len(worker_group),
|
||||
init_method=url,
|
||||
timeout_s=backend_config.timeout_s,
|
||||
)
|
||||
)
|
||||
ray.get(setup_futures)
|
||||
else:
|
||||
raise RuntimeError("Distributed torch is not available.")
|
||||
|
||||
def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config):
|
||||
futures = worker_group.execute_async(
|
||||
_shutdown_torch,
|
||||
destroy_process_group=len(worker_group) > 1,
|
||||
)
|
||||
timeout_s = ray_constants.env_integer(
|
||||
TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
|
||||
DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
try:
|
||||
ray.get(futures, timeout=timeout_s)
|
||||
except GetTimeoutError:
|
||||
logger.warning(
|
||||
f"Torch process group shutdown timed out after {timeout_s} seconds"
|
||||
)
|
||||
|
||||
def on_training_start(
|
||||
self, worker_group: BaseWorkerGroup, backend_config: BackendConfig
|
||||
):
|
||||
worker_group.execute(_set_torch_distributed_env_vars)
|
||||
@@ -0,0 +1,196 @@
|
||||
import os
|
||||
import tempfile
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from ray.air._internal.torch_utils import (
|
||||
consume_prefix_in_state_dict_if_present_not_in_place,
|
||||
load_torch_model,
|
||||
)
|
||||
from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.preprocessor import Preprocessor
|
||||
|
||||
ENCODED_DATA_KEY = "torch_encoded_data"
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class TorchCheckpoint(FrameworkCheckpoint):
|
||||
"""A :class:`~ray.train.Checkpoint` with Torch-specific functionality."""
|
||||
|
||||
MODEL_FILENAME = "model.pt"
|
||||
|
||||
@classmethod
|
||||
def from_state_dict(
|
||||
cls,
|
||||
state_dict: Dict[str, Any],
|
||||
*,
|
||||
preprocessor: Optional["Preprocessor"] = None,
|
||||
) -> "TorchCheckpoint":
|
||||
"""Create a :class:`~ray.train.Checkpoint` that stores a model state dictionary.
|
||||
|
||||
.. tip::
|
||||
|
||||
This is the recommended method for creating
|
||||
:class:`TorchCheckpoints<TorchCheckpoint>`.
|
||||
|
||||
Args:
|
||||
state_dict: The model state dictionary to store in the checkpoint.
|
||||
preprocessor: A fitted preprocessor to be applied before inference.
|
||||
|
||||
Returns:
|
||||
A :class:`TorchCheckpoint` containing the specified state dictionary.
|
||||
|
||||
Examples:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from ray.train.torch import TorchCheckpoint
|
||||
|
||||
# Set manual seed
|
||||
torch.manual_seed(42)
|
||||
|
||||
# Function to create a NN model
|
||||
def create_model() -> nn.Module:
|
||||
model = nn.Sequential(nn.Linear(1, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10,1))
|
||||
return model
|
||||
|
||||
# Create a TorchCheckpoint from our model's state_dict
|
||||
model = create_model()
|
||||
checkpoint = TorchCheckpoint.from_state_dict(model.state_dict())
|
||||
|
||||
# Now load the model from the TorchCheckpoint by providing the
|
||||
# model architecture
|
||||
model_from_chkpt = checkpoint.get_model(create_model())
|
||||
|
||||
# Assert they have the same state dict
|
||||
assert str(model.state_dict()) == str(model_from_chkpt.state_dict())
|
||||
print("worked")
|
||||
|
||||
.. testoutput::
|
||||
:hide:
|
||||
|
||||
...
|
||||
"""
|
||||
tempdir = tempfile.mkdtemp()
|
||||
|
||||
model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix()
|
||||
stripped_state_dict = consume_prefix_in_state_dict_if_present_not_in_place(
|
||||
state_dict, "module."
|
||||
)
|
||||
torch.save(stripped_state_dict, model_path)
|
||||
|
||||
checkpoint = cls.from_directory(tempdir)
|
||||
if preprocessor:
|
||||
checkpoint.set_preprocessor(preprocessor)
|
||||
return checkpoint
|
||||
|
||||
@classmethod
|
||||
def from_model(
|
||||
cls,
|
||||
model: torch.nn.Module,
|
||||
*,
|
||||
preprocessor: Optional["Preprocessor"] = None,
|
||||
) -> "TorchCheckpoint":
|
||||
"""Create a :class:`~ray.train.Checkpoint` that stores a Torch model.
|
||||
|
||||
.. note::
|
||||
|
||||
PyTorch recommends storing state dictionaries. To create a
|
||||
:class:`TorchCheckpoint` from a state dictionary, call
|
||||
:meth:`~ray.train.torch.TorchCheckpoint.from_state_dict`. To learn more
|
||||
about state dictionaries, read
|
||||
`Saving and Loading Models <https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict>`_. # noqa: E501
|
||||
|
||||
Args:
|
||||
model: The Torch model to store in the checkpoint.
|
||||
preprocessor: A fitted preprocessor to be applied before inference.
|
||||
|
||||
Returns:
|
||||
A :class:`TorchCheckpoint` containing the specified model.
|
||||
|
||||
Examples:
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray.train.torch import TorchCheckpoint
|
||||
import torch
|
||||
|
||||
# Create model identity and send a random tensor to it
|
||||
model = torch.nn.Identity()
|
||||
input = torch.randn(2, 2)
|
||||
output = model(input)
|
||||
|
||||
# Create a checkpoint
|
||||
checkpoint = TorchCheckpoint.from_model(model)
|
||||
print(checkpoint)
|
||||
|
||||
.. testoutput::
|
||||
:hide:
|
||||
|
||||
...
|
||||
"""
|
||||
tempdir = tempfile.mkdtemp()
|
||||
|
||||
model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix()
|
||||
torch.save(model, model_path)
|
||||
|
||||
checkpoint = cls.from_directory(tempdir)
|
||||
if preprocessor:
|
||||
checkpoint.set_preprocessor(preprocessor)
|
||||
return checkpoint
|
||||
|
||||
def get_model(self, model: Optional[torch.nn.Module] = None) -> torch.nn.Module:
|
||||
"""Retrieve the model stored in this checkpoint.
|
||||
|
||||
.. warning::
|
||||
|
||||
The checkpoint path must point to a **trusted** source.
|
||||
Checkpoints created with
|
||||
:meth:`~ray.train.torch.TorchCheckpoint.from_model` store the entire
|
||||
``nn.Module`` via pickle serialization. Loading such a checkpoint from an
|
||||
untrusted path (shared storage, downloaded artifact, checkpoint produced by
|
||||
a different party) is equivalent to executing arbitrary Python code. Prefer
|
||||
checkpoints created with
|
||||
:meth:`~ray.train.torch.TorchCheckpoint.from_state_dict`, which stores
|
||||
only model weights and is safe to load from untrusted sources.
|
||||
|
||||
Args:
|
||||
model: If the checkpoint contains a model state dict, and not
|
||||
the model itself, then the state dict will be loaded to this
|
||||
``model``. Otherwise, the model will be discarded.
|
||||
|
||||
Returns:
|
||||
The loaded ``torch.nn.Module``.
|
||||
"""
|
||||
with self.as_directory() as tempdir:
|
||||
model_path = Path(tempdir, self.MODEL_FILENAME).as_posix()
|
||||
if not os.path.exists(model_path):
|
||||
raise RuntimeError(
|
||||
"`model.pt` not found within this checkpoint. Make sure you "
|
||||
"created this `TorchCheckpoint` from one of its public "
|
||||
"constructors (`from_state_dict` or `from_model`)."
|
||||
)
|
||||
model_or_state_dict = torch.load(
|
||||
model_path, map_location="cpu", weights_only=False
|
||||
)
|
||||
|
||||
if isinstance(model_or_state_dict, torch.nn.Module):
|
||||
if model:
|
||||
warnings.warn(
|
||||
"TorchCheckpoint already contains all information needed. "
|
||||
"Discarding provided `model` argument."
|
||||
)
|
||||
model = load_torch_model(
|
||||
saved_model=model_or_state_dict, model_definition=model
|
||||
)
|
||||
return model
|
||||
@@ -0,0 +1,200 @@
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.train import Checkpoint, DataConfig, RunConfig, ScalingConfig
|
||||
from ray.train.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.torch.config import TorchConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.util import PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class TorchTrainer(DataParallelTrainer):
|
||||
"""A Trainer for data parallel PyTorch training.
|
||||
|
||||
At a high level, this Trainer does the following:
|
||||
|
||||
1. Launches multiple workers as defined by the ``scaling_config``.
|
||||
2. Sets up a distributed PyTorch environment
|
||||
on these workers as defined by the ``torch_config``.
|
||||
3. Ingests the input ``datasets`` based on the ``dataset_config``.
|
||||
4. Runs the input ``train_loop_per_worker(train_loop_config)``
|
||||
on all workers.
|
||||
|
||||
For more details, see:
|
||||
|
||||
* :ref:`PyTorch Guide <train-pytorch>`
|
||||
* :ref:`PyTorch Lightning Guide <train-pytorch-lightning>`
|
||||
* :ref:`Hugging Face Transformers Guide <train-pytorch-transformers>`
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
import ray
|
||||
from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
# If using GPUs, set this to True.
|
||||
use_gpu = False
|
||||
# Number of processes to run training on.
|
||||
num_workers = 4
|
||||
|
||||
# Define your network structure.
|
||||
class NeuralNetwork(nn.Module):
|
||||
def __init__(self):
|
||||
super(NeuralNetwork, self).__init__()
|
||||
self.layer1 = nn.Linear(1, 32)
|
||||
self.relu = nn.ReLU()
|
||||
self.layer2 = nn.Linear(32, 1)
|
||||
|
||||
def forward(self, input):
|
||||
return self.layer2(self.relu(self.layer1(input)))
|
||||
|
||||
# Training loop.
|
||||
def train_loop_per_worker(config):
|
||||
|
||||
# Read configurations.
|
||||
lr = config["lr"]
|
||||
batch_size = config["batch_size"]
|
||||
num_epochs = config["num_epochs"]
|
||||
|
||||
# Fetch training dataset.
|
||||
train_dataset_shard = ray.train.get_dataset_shard("train")
|
||||
|
||||
# Instantiate and prepare model for training.
|
||||
model = NeuralNetwork()
|
||||
model = ray.train.torch.prepare_model(model)
|
||||
|
||||
# Define loss and optimizer.
|
||||
loss_fn = nn.MSELoss()
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
|
||||
|
||||
# Create data loader.
|
||||
dataloader = train_dataset_shard.iter_torch_batches(
|
||||
batch_size=batch_size, dtypes=torch.float
|
||||
)
|
||||
|
||||
# Train multiple epochs.
|
||||
for epoch in range(num_epochs):
|
||||
|
||||
# Train epoch.
|
||||
for batch in dataloader:
|
||||
output = model(batch["input"])
|
||||
loss = loss_fn(output, batch["label"])
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# Create checkpoint.
|
||||
base_model = (model.module
|
||||
if isinstance(model, DistributedDataParallel) else model)
|
||||
checkpoint_dir = tempfile.mkdtemp()
|
||||
torch.save(
|
||||
{"model_state_dict": base_model.state_dict()},
|
||||
os.path.join(checkpoint_dir, "model.pt"),
|
||||
)
|
||||
checkpoint = Checkpoint.from_directory(checkpoint_dir)
|
||||
|
||||
# Report metrics and checkpoint.
|
||||
ray.train.report({"loss": loss.item()}, checkpoint=checkpoint)
|
||||
|
||||
|
||||
# Define configurations.
|
||||
train_loop_config = {"num_epochs": 20, "lr": 0.01, "batch_size": 32}
|
||||
scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
|
||||
run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1))
|
||||
|
||||
# Define datasets.
|
||||
train_dataset = ray.data.from_items(
|
||||
[{"input": [x], "label": [2 * x + 1]} for x in range(2000)]
|
||||
)
|
||||
datasets = {"train": train_dataset}
|
||||
|
||||
# Initialize the Trainer.
|
||||
trainer = TorchTrainer(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
scaling_config=scaling_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets
|
||||
)
|
||||
|
||||
# Train the model.
|
||||
result = trainer.fit()
|
||||
|
||||
# Inspect the results.
|
||||
final_loss = result.metrics["loss"]
|
||||
|
||||
Args:
|
||||
|
||||
train_loop_per_worker: The training function to execute on each worker.
|
||||
This function can either take in zero arguments or a single ``Dict``
|
||||
argument which is set by defining ``train_loop_config``.
|
||||
Within this function you can use any of the
|
||||
:ref:`Ray Train Loop utilities <train-loop-api>`.
|
||||
train_loop_config: A configuration ``Dict`` to pass in as an argument to
|
||||
``train_loop_per_worker``.
|
||||
This is typically used for specifying hyperparameters. Passing large
|
||||
datasets via `train_loop_config` is not recommended and may introduce
|
||||
large overhead and unknown issues with serialization and deserialization.
|
||||
torch_config: The configuration for setting up the PyTorch Distributed backend.
|
||||
If set to None, a default configuration will be used in which
|
||||
GPU training uses NCCL and CPU training uses Gloo.
|
||||
scaling_config: The configuration for how to scale data parallel training.
|
||||
``num_workers`` determines how many Python processes are used for training,
|
||||
and ``use_gpu`` determines whether or not each process should use GPUs.
|
||||
See :class:`~ray.train.ScalingConfig` for more info.
|
||||
run_config: The configuration for the execution of the training run.
|
||||
See :class:`~ray.train.RunConfig` for more info.
|
||||
datasets: The Ray Datasets to ingest for training.
|
||||
Datasets are keyed by name (``{name: dataset}``).
|
||||
Each dataset can be accessed from within the ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_dataset_shard(name)``.
|
||||
Sharding and additional configuration can be done by
|
||||
passing in a ``dataset_config``.
|
||||
dataset_config: The configuration for ingesting the input ``datasets``.
|
||||
By default, all the Ray Dataset are split equally across workers.
|
||||
See :class:`~ray.train.DataConfig` for more details.
|
||||
metadata: Dict that should be made available via
|
||||
`ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
|
||||
for checkpoints saved from this Trainer. Must be JSON-serializable.
|
||||
resume_from_checkpoint: A checkpoint to resume training from.
|
||||
This checkpoint can be accessed from within ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_checkpoint()``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
torch_config: Optional[TorchConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[DataConfig] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
if not torch_config:
|
||||
torch_config = TorchConfig()
|
||||
|
||||
super(TorchTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=torch_config,
|
||||
scaling_config=scaling_config,
|
||||
dataset_config=dataset_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
resume_from_checkpoint=resume_from_checkpoint,
|
||||
metadata=metadata,
|
||||
)
|
||||
@@ -0,0 +1,793 @@
|
||||
import collections
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import types
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
from torch.cuda.amp import GradScaler, autocast
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.optim import Optimizer
|
||||
from torch.utils.data import (
|
||||
DataLoader,
|
||||
DistributedSampler,
|
||||
IterableDataset,
|
||||
RandomSampler,
|
||||
SequentialSampler,
|
||||
)
|
||||
|
||||
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
||||
from ray.air._internal.device_manager import (
|
||||
get_torch_device_manager_by_context,
|
||||
get_torch_device_manager_by_device_type,
|
||||
)
|
||||
from ray.train._internal import session
|
||||
from ray.train._internal.accelerator import Accelerator
|
||||
from ray.train._internal.session import get_accelerator, set_accelerator
|
||||
from ray.train.utils import _log_deprecation_warning
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
def get_device() -> torch.device:
|
||||
"""Gets the correct torch device configured for this process.
|
||||
|
||||
Returns the torch device for the current worker. If more than 1 GPU is
|
||||
requested per worker, returns the device with the minimal device index.
|
||||
|
||||
.. note::
|
||||
|
||||
If you requested multiple GPUs per worker, and want to get
|
||||
the full list of torch devices, please use
|
||||
:meth:`~ray.train.torch.get_devices`.
|
||||
|
||||
Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
|
||||
superset of the `ray.get_gpu_ids()`.
|
||||
|
||||
Examples:
|
||||
|
||||
Example: Launched 2 workers on the current node, each with 1 GPU
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
|
||||
ray.get_gpu_ids() == [2]
|
||||
torch.cuda.is_available() == True
|
||||
get_device() == torch.device("cuda:0")
|
||||
|
||||
Example: Launched 4 workers on the current node, each with 1 GPU
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
|
||||
ray.get_gpu_ids() == [2]
|
||||
torch.cuda.is_available() == True
|
||||
get_device() == torch.device("cuda:2")
|
||||
|
||||
Example: Launched 2 workers on the current node, each with 2 GPUs
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
|
||||
ray.get_gpu_ids() == [2,3]
|
||||
torch.cuda.is_available() == True
|
||||
get_device() == torch.device("cuda:2")
|
||||
|
||||
|
||||
You can move a model to device by:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
model.to(ray.train.torch.get_device())
|
||||
|
||||
Instead of manually checking the device type:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
model.to("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
Returns:
|
||||
The torch device for the current worker.
|
||||
"""
|
||||
from ray.air._internal import torch_utils
|
||||
|
||||
record_extra_usage_tag(TagKey.TRAIN_TORCH_GET_DEVICE, "1")
|
||||
return torch_utils.get_devices()[0]
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
def get_devices() -> List[torch.device]:
|
||||
"""Gets the correct torch device list configured for this process.
|
||||
|
||||
Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
|
||||
superset of the `ray.get_gpu_ids()`.
|
||||
|
||||
|
||||
Examples:
|
||||
|
||||
Example: Launched 2 workers on the current node, each with 1 GPU
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] == "2,3"
|
||||
ray.get_gpu_ids() == [2]
|
||||
torch.cuda.is_available() == True
|
||||
get_devices() == [torch.device("cuda:0")]
|
||||
|
||||
Example: Launched 4 workers on the current node, each with 1 GPU
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
|
||||
ray.get_gpu_ids() == [2]
|
||||
torch.cuda.is_available() == True
|
||||
get_devices() == [torch.device("cuda:2")]
|
||||
|
||||
Example: Launched 2 workers on the current node, each with 2 GPUs
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3"
|
||||
ray.get_gpu_ids() == [2,3]
|
||||
torch.cuda.is_available() == True
|
||||
get_devices() == [torch.device("cuda:2"), torch.device("cuda:3")]
|
||||
|
||||
Returns:
|
||||
A list of torch devices for the current worker.
|
||||
"""
|
||||
|
||||
from ray.air._internal import torch_utils
|
||||
|
||||
record_extra_usage_tag(TagKey.TRAIN_TORCH_GET_DEVICES, "1")
|
||||
return torch_utils.get_devices()
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
def prepare_model(
|
||||
model: torch.nn.Module,
|
||||
move_to_device: Union[bool, torch.device] = True,
|
||||
parallel_strategy: Optional[str] = "ddp",
|
||||
parallel_strategy_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> torch.nn.Module:
|
||||
"""Prepares the model for distributed execution.
|
||||
|
||||
This allows you to use the same exact code regardless of number of
|
||||
workers or the device type being used (CPU, GPU).
|
||||
|
||||
Args:
|
||||
model: A torch model to prepare.
|
||||
move_to_device: Either a boolean indiciating whether to move
|
||||
the model to the correct device or an actual device to
|
||||
move the model to. If set to False, the model needs
|
||||
to manually be moved to the correct device.
|
||||
parallel_strategy: Whether to wrap models in
|
||||
``DistributedDataParallel``, ``FullyShardedDataParallel``,
|
||||
or neither. Must be one of ``"ddp"``, ``"fsdp"``, or ``None``.
|
||||
parallel_strategy_kwargs: Args to pass into
|
||||
``DistributedDataParallel`` or ``FullyShardedDataParallel``
|
||||
initialization if ``parallel_strategy`` is set to "ddp"
|
||||
or "fsdp", respectively.
|
||||
|
||||
Returns:
|
||||
The prepared model, wrapped according to ``parallel_strategy``.
|
||||
"""
|
||||
if parallel_strategy == "fsdp" and Version(torch.__version__) < Version("1.11.0"):
|
||||
raise ImportError(
|
||||
"FullyShardedDataParallel requires torch>=1.11.0. "
|
||||
"Run `pip install 'torch>=1.11.0'` to use FullyShardedDataParallel."
|
||||
)
|
||||
|
||||
record_extra_usage_tag(TagKey.TRAIN_TORCH_PREPARE_MODEL, "1")
|
||||
return get_accelerator(_TorchAccelerator).prepare_model(
|
||||
model,
|
||||
move_to_device=move_to_device,
|
||||
parallel_strategy=parallel_strategy,
|
||||
parallel_strategy_kwargs=parallel_strategy_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
def prepare_data_loader(
|
||||
data_loader: torch.utils.data.DataLoader,
|
||||
add_dist_sampler: bool = True,
|
||||
move_to_device: bool = True,
|
||||
auto_transfer: bool = True,
|
||||
) -> torch.utils.data.DataLoader:
|
||||
"""Prepares :class:`~torch.utils.data.DataLoader` for distributed execution.
|
||||
|
||||
This allows you to use the same exact code regardless of number of
|
||||
workers or the device type being used (CPU, GPU).
|
||||
|
||||
.. note::
|
||||
|
||||
This method adds a `DistributedSampler` to the `DataLoader` if the
|
||||
number of training workers is greater than 1. If shuffling is
|
||||
enabled on the original `DataLoader`, then `shuffle=True` will also
|
||||
be passed into the `DistributedSampler` constructor. `shuffle=False`
|
||||
on the original `DataLoader` also means that shuffling is disabled
|
||||
on the sampler.
|
||||
|
||||
With more than 1 worker, calling the `DistributedSampler.set_epoch` method
|
||||
at the beginning of each epoch before creating the DataLoader iterator
|
||||
is necessary to make shuffling work properly across multiple epochs.
|
||||
Otherwise, the same ordering will be always used.
|
||||
See: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler # noqa: E501
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode:
|
||||
:skipif: True
|
||||
|
||||
import torch
|
||||
|
||||
import ray.train.torch
|
||||
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
..., batch_size=..., shuffle=True
|
||||
)
|
||||
train_dataloader = ray.train.torch.prepare_data_loader(train_loader)
|
||||
|
||||
for epoch in range(10):
|
||||
if ray.train.get_context().get_world_size() > 1:
|
||||
# Required for the distributed sampler to shuffle properly across epochs
|
||||
train_dataloader.sampler.set_epoch(epoch)
|
||||
|
||||
for X, y in train_loader:
|
||||
# No need to move data to GPU, this is done by `prepare_data_loader`!
|
||||
# X, y = X.to("cuda"), y.to("cuda")
|
||||
...
|
||||
|
||||
Args:
|
||||
data_loader: The DataLoader to prepare.
|
||||
add_dist_sampler: Whether to add a DistributedSampler to
|
||||
the provided DataLoader.
|
||||
move_to_device: If set, automatically move the data
|
||||
returned by the data loader to the correct device.
|
||||
auto_transfer: If set and device is GPU, another CUDA stream
|
||||
is created to automatically copy data from host (CPU) memory
|
||||
to device (GPU) memory (the default CUDA stream still runs the
|
||||
training procedure). If device is CPU, it will be disabled
|
||||
regardless of the setting. This configuration will be ignored
|
||||
if ``move_to_device`` is False.
|
||||
|
||||
Returns:
|
||||
The prepared DataLoader.
|
||||
"""
|
||||
record_extra_usage_tag(TagKey.TRAIN_TORCH_PREPARE_DATALOADER, "1")
|
||||
return get_accelerator(_TorchAccelerator).prepare_data_loader(
|
||||
data_loader,
|
||||
add_dist_sampler=add_dist_sampler,
|
||||
move_to_device=move_to_device,
|
||||
auto_transfer=auto_transfer,
|
||||
)
|
||||
|
||||
|
||||
def _log_amp_deprecation_warning():
|
||||
# Keep V2 imports out of top-level V1 imports.
|
||||
from ray.train.v2.torch.train_loop_utils import _TORCH_AMP_DEPRECATION_MESSAGE
|
||||
|
||||
_log_deprecation_warning(_TORCH_AMP_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@Deprecated
|
||||
def accelerate(amp: bool = False) -> None:
|
||||
"""[Deprecated] Enables training optimizations.
|
||||
|
||||
Arguments:
|
||||
amp: If true, perform training with automatic mixed precision.
|
||||
Otherwise, use full precision.
|
||||
|
||||
.. warning:: ``train.torch.accelerate`` cannot be called more than once, and it
|
||||
must be called before any other ``train.torch`` utility function.
|
||||
"""
|
||||
_log_amp_deprecation_warning()
|
||||
try:
|
||||
set_accelerator(_TorchAccelerator(amp=amp))
|
||||
except RuntimeError:
|
||||
raise RuntimeError(
|
||||
"An accelerator has already been set. Make sure "
|
||||
"`train.torch.accelerate()` is not called multiple times, and is called "
|
||||
"before any of the prepare methods."
|
||||
)
|
||||
|
||||
|
||||
@Deprecated
|
||||
def prepare_optimizer(optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
|
||||
"""[Deprecated] Wraps optimizer to support automatic mixed precision.
|
||||
|
||||
Args:
|
||||
optimizer: The DataLoader to prepare.
|
||||
|
||||
Returns:
|
||||
A wrapped optimizer.
|
||||
"""
|
||||
_log_amp_deprecation_warning()
|
||||
return get_accelerator(_TorchAccelerator).prepare_optimizer(optimizer)
|
||||
|
||||
|
||||
@Deprecated
|
||||
def backward(tensor: torch.Tensor) -> None:
|
||||
"""[Deprecated] Computes the gradient of the specified tensor w.r.t. graph leaves.
|
||||
|
||||
Args:
|
||||
tensor: Tensor of which the derivative will be computed.
|
||||
"""
|
||||
_log_amp_deprecation_warning()
|
||||
get_accelerator(_TorchAccelerator).backward(tensor)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
def enable_reproducibility(seed: int = 0) -> None:
|
||||
"""Limits sources of nondeterministic behavior.
|
||||
|
||||
This function:
|
||||
|
||||
* Seeds PyTorch, Python, and NumPy.
|
||||
* Disables CUDA convolution benchmarking.
|
||||
* Configures PyTorch to use determinstic algorithms.
|
||||
* Seeds workers spawned for multi-process data loading.
|
||||
|
||||
Args:
|
||||
seed: The number to seed libraries and data workers with.
|
||||
|
||||
.. warning:: ``train.torch.enable_reproducibility()`` can't guarantee
|
||||
completely reproducible results across executions. To learn more, read
|
||||
the `PyTorch notes on randomness
|
||||
<https://pytorch.org/docs/stable/notes/randomness.html>`_.
|
||||
"""
|
||||
get_accelerator(_TorchAccelerator).enable_reproducibility(seed)
|
||||
|
||||
|
||||
@Deprecated
|
||||
class TorchWorkerProfiler:
|
||||
"""Utility class for running PyTorch Profiler on a Train worker.
|
||||
|
||||
Args:
|
||||
trace_dir: The directory to store traces on the worker node.
|
||||
If ``None``, this will use a default temporary dir.
|
||||
"""
|
||||
|
||||
WORKER_TRACE_DIR_NAME = "pytorch_profiler_worker_traces"
|
||||
|
||||
def __init__(self, trace_dir: Optional[str] = None):
|
||||
raise DeprecationWarning(
|
||||
"The `ray.train.torch.TorchWorkerProfiler` API is deprecated in Ray 2.0.",
|
||||
)
|
||||
|
||||
|
||||
class _TorchAccelerator(Accelerator):
|
||||
"""A utility that implements methods to accelerate PyTorch training.
|
||||
|
||||
Arguments:
|
||||
amp: If true, perform training with automatic mixed precision.
|
||||
Otherwise, use full precision.
|
||||
"""
|
||||
|
||||
def __init__(self, amp: bool = False):
|
||||
self.amp_is_enabled = amp
|
||||
self.scaler = GradScaler() if amp else None
|
||||
self._seed = None
|
||||
self.device_manager = get_torch_device_manager_by_context()
|
||||
|
||||
def prepare_model(
|
||||
self,
|
||||
model: torch.nn.Module,
|
||||
move_to_device: bool = True,
|
||||
parallel_strategy: Optional[str] = "ddp",
|
||||
parallel_strategy_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> torch.nn.Module:
|
||||
"""Prepares the model for distributed execution.
|
||||
|
||||
This allows you to use the same exact code regardless of number of
|
||||
workers or the device type being used (CPU, GPU).
|
||||
|
||||
Args:
|
||||
model: A torch model to prepare.
|
||||
move_to_device: Whether to move the model to the correct
|
||||
device. If set to False, the model needs to manually be moved
|
||||
to the correct device.
|
||||
parallel_strategy: Whether to wrap models in
|
||||
``DistributedDataParallel``, ``FullyShardedDataParallel`` (
|
||||
Experimental), or neither. Must be one of ``"ddp"``, ``"fsdp"``,
|
||||
or ``None``.
|
||||
parallel_strategy_kwargs: Args to pass into
|
||||
``DistributedDataParallel`` or ``FullyShardedDataParallel``
|
||||
initialization if ``parallel_strategy`` is set to "ddp"
|
||||
or "fsdp", respectively.
|
||||
|
||||
Returns:
|
||||
The prepared model, wrapped according to ``parallel_strategy``.
|
||||
"""
|
||||
parallel_strategy_kwargs = parallel_strategy_kwargs or {}
|
||||
|
||||
rank = session.get_local_rank()
|
||||
|
||||
if isinstance(move_to_device, torch.device):
|
||||
device = move_to_device
|
||||
else:
|
||||
device = get_device()
|
||||
if isinstance(device, list):
|
||||
device = device[0]
|
||||
|
||||
if self.device_manager.is_available():
|
||||
self.device_manager.set_device(device)
|
||||
|
||||
if move_to_device:
|
||||
if rank == 0:
|
||||
logger.info(f"Moving model to device: {device}")
|
||||
else:
|
||||
logger.debug(f"Moving model to device: {device}")
|
||||
model = model.to(device)
|
||||
|
||||
def model_get_state(self):
|
||||
# `__getstate__` is an special method that informs pickle which attributes
|
||||
# to serialize. This custom implementation ensures that the wrapped forward
|
||||
# method and custom `__getstate__` method aren't serialized.
|
||||
if hasattr(self, "_original_get_state"):
|
||||
state = self._original_get_state()
|
||||
state["__getstate__"] = state["_original_get_state"]
|
||||
del state["_original_get_state"]
|
||||
else:
|
||||
# If model does not have a `__getstate__` already defined, use default
|
||||
# implementation.
|
||||
state = self.__dict__.copy()
|
||||
del state["__getstate__"]
|
||||
state["forward"] = state["_unwrapped_forward"]
|
||||
del state["_unwrapped_forward"]
|
||||
|
||||
return state
|
||||
|
||||
if self.amp_is_enabled:
|
||||
# Pickle cannot serialize the wrapped forward method. As a workaround,
|
||||
# define a custom `__getstate__` method that unwraps the forward method.
|
||||
model._unwrapped_forward = model.forward
|
||||
model.forward = autocast()(model.forward)
|
||||
|
||||
# TODO(amogkam): Replace below logic with a generic "unpack model" method.
|
||||
# Replacing the `model.forward` method makes the model no longer
|
||||
# serializable. When serializing the model, we have to override the
|
||||
# `__getstate__` method to set back the original forward method.
|
||||
if hasattr(model, "__getstate__"):
|
||||
model._original_get_state = model.__getstate__
|
||||
# `__getstate__` must be a bound method rather than an callable attribute.
|
||||
# See https://stackoverflow.com/questions/972/adding-a-method-to-an-existing-object-instance. # noqa: E501
|
||||
model.__getstate__ = types.MethodType(model_get_state, model)
|
||||
|
||||
world_size = session.get_world_size()
|
||||
|
||||
if parallel_strategy and world_size > 1:
|
||||
if parallel_strategy == "ddp":
|
||||
DataParallel = DistributedDataParallel
|
||||
if self.device_manager.is_available() and device.type != "cpu":
|
||||
parallel_strategy_kwargs = {
|
||||
"device_ids": [device],
|
||||
"output_device": device,
|
||||
**parallel_strategy_kwargs,
|
||||
}
|
||||
else:
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
"FSDP is only available with GPU-enabled "
|
||||
"training. Set "
|
||||
"`use_gpu=True` in your Trainer to train with "
|
||||
"GPUs."
|
||||
)
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel
|
||||
|
||||
DataParallel = FullyShardedDataParallel
|
||||
if rank == 0:
|
||||
logger.info(f"Wrapping provided model in {DataParallel.__name__}.")
|
||||
else:
|
||||
logger.debug(f"Wrapping provided model in {DataParallel.__name__}.")
|
||||
model = DataParallel(model, **parallel_strategy_kwargs)
|
||||
|
||||
return model
|
||||
|
||||
def prepare_data_loader(
|
||||
self,
|
||||
data_loader: torch.utils.data.DataLoader,
|
||||
add_dist_sampler: bool = True,
|
||||
move_to_device: bool = True,
|
||||
auto_transfer: bool = False,
|
||||
) -> torch.utils.data.DataLoader:
|
||||
"""Prepares DataLoader for distributed execution.
|
||||
|
||||
This allows you to use the same exact code regardless of number of
|
||||
workers or the device type being used (CPU, GPU).
|
||||
|
||||
Args:
|
||||
data_loader: The DataLoader to prepare.
|
||||
add_dist_sampler: Whether to add a DistributedSampler to
|
||||
the provided DataLoader.
|
||||
move_to_device: If set, automatically move the data
|
||||
returned by the data loader to the correct device.
|
||||
auto_transfer: (Experimental) If set and device is GPU, another CUDA stream
|
||||
is created to automatically copy data from host (CPU) memory
|
||||
to device (GPU) memory (the default CUDA stream still runs the
|
||||
training procedure). If device is CPU, it will be disabled
|
||||
regardless of the setting. This configuration will be ignored
|
||||
if ``move_to_device`` is False.
|
||||
|
||||
Returns:
|
||||
The prepared DataLoader.
|
||||
"""
|
||||
|
||||
world_size = session.get_world_size()
|
||||
world_rank = session.get_world_rank()
|
||||
|
||||
# Only add Distributed Sampler if the following conditions hold:
|
||||
# 1. More than one training worker is being used.
|
||||
# 2. A DistributedSampler has not already been added by the user.
|
||||
# 3. The dataset is not an IterableDataset. Samplers do not worker with
|
||||
# IterableDatasets.
|
||||
if (
|
||||
world_size > 1
|
||||
and not isinstance(data_loader.sampler, DistributedSampler)
|
||||
and not (
|
||||
hasattr(data_loader, "dataset")
|
||||
and isinstance(data_loader.dataset, IterableDataset)
|
||||
)
|
||||
and add_dist_sampler
|
||||
):
|
||||
|
||||
def with_sampler(loader):
|
||||
# Automatically set the DistributedSampler
|
||||
|
||||
# If you're using a sampler, the DataLoader shuffle flag must be set to
|
||||
# False. Shuffling is instead determined by the shuffle argument passed
|
||||
# to the DistributedSampler constructor.
|
||||
|
||||
# If no sampler is passed to the DataLoader constructor, Torch
|
||||
# constructs a default sampler. The default sampler is a RandomSampler
|
||||
# if shuffling is enabled and a SequentialSampler otherwise. DataLoader
|
||||
# does not have a shuffle attribute, so we instead identify whether
|
||||
# shuffling is enabled by checking the default sampler type.
|
||||
shuffle = not isinstance(loader.sampler, SequentialSampler)
|
||||
|
||||
def seeded_worker_init_fn(
|
||||
worker_init_fn: Optional[Callable[[int], None]],
|
||||
):
|
||||
def wrapper(worker_id: int):
|
||||
worker_seed = torch.initial_seed() % 2**32
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
if worker_init_fn:
|
||||
worker_init_fn(worker_id)
|
||||
|
||||
return wrapper
|
||||
|
||||
worker_init_fn: Optional[Callable[[int], None]] = loader.worker_init_fn
|
||||
generator: Optional[torch.Generator] = loader.generator
|
||||
if self._seed is not None:
|
||||
worker_init_fn = seeded_worker_init_fn(worker_init_fn)
|
||||
generator = torch.Generator()
|
||||
generator.manual_seed(self._seed)
|
||||
|
||||
using_default_sampler = isinstance(
|
||||
loader.sampler, (SequentialSampler, RandomSampler)
|
||||
)
|
||||
if not using_default_sampler and world_rank == 0:
|
||||
logger.warning(
|
||||
f"The {loader.sampler.__class__.__name__} will be overwritten "
|
||||
"with a DistributedSampler. You can disable this by setting "
|
||||
"`with_sampler` to False in `prepare_data_loader`."
|
||||
)
|
||||
|
||||
data_loader_args = {
|
||||
"dataset": loader.dataset,
|
||||
"batch_size": loader.batch_size,
|
||||
"shuffle": False,
|
||||
"num_workers": loader.num_workers,
|
||||
"collate_fn": loader.collate_fn,
|
||||
"pin_memory": loader.pin_memory,
|
||||
"drop_last": loader.drop_last,
|
||||
"timeout": loader.timeout,
|
||||
"worker_init_fn": worker_init_fn,
|
||||
"generator": generator,
|
||||
"sampler": DistributedSampler(loader.dataset, shuffle=shuffle),
|
||||
}
|
||||
return DataLoader(**data_loader_args)
|
||||
|
||||
data_loader = with_sampler(data_loader)
|
||||
|
||||
if move_to_device:
|
||||
device = get_device()
|
||||
data_loader = _WrappedDataLoader(data_loader, device, auto_transfer)
|
||||
|
||||
return data_loader
|
||||
|
||||
def prepare_optimizer(self, optimizer: Optimizer) -> Optimizer:
|
||||
"""Wraps optimizer to support automatic mixed precision.
|
||||
|
||||
Args:
|
||||
optimizer: The DataLoader to prepare.
|
||||
|
||||
Returns:
|
||||
A wrapped optimizer.
|
||||
"""
|
||||
return _WrappedOptimizer(optimizer, scaler=self.scaler)
|
||||
|
||||
def backward(self, tensor: torch.Tensor) -> None:
|
||||
"""Computes the gradient of the specified tensor w.r.t. graph leaves.
|
||||
|
||||
Args:
|
||||
tensor: Tensor of which the derivative will be computed.
|
||||
"""
|
||||
if self.amp_is_enabled:
|
||||
self.scaler.scale(tensor).backward()
|
||||
else:
|
||||
tensor.backward()
|
||||
|
||||
def enable_reproducibility(self, seed: int = 0) -> None:
|
||||
"""Limits sources of nondeterministic behavior."""
|
||||
self._seed = seed
|
||||
|
||||
torch.manual_seed(seed)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
torch.use_deterministic_algorithms(True)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
# If you want to use deterministic algorithms with CUDA, then you need to set
|
||||
# the CUBLAS_WORKSPACE_CONFIG environment variable; otherwise, Torch errors.
|
||||
# See https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility.
|
||||
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
||||
|
||||
|
||||
class _WrappedDataLoader(DataLoader):
|
||||
def __init__(
|
||||
self, base_dataloader: DataLoader, device: torch.device, auto_transfer: bool
|
||||
):
|
||||
self.__dict__.update(getattr(base_dataloader, "__dict__", {}))
|
||||
self._dataloader = base_dataloader
|
||||
self.dataloader_iter = None
|
||||
self.device = device
|
||||
|
||||
self.device_manager = get_torch_device_manager_by_device_type(device.type)
|
||||
|
||||
# disable auto transfer (host->device) if cpu is used
|
||||
if device.type != "cpu" and self.device_manager.supports_stream():
|
||||
self._auto_transfer = auto_transfer
|
||||
else:
|
||||
self._auto_transfer = False
|
||||
# create a new device stream to move data from host to device concurrently
|
||||
self._memcpy_stream = (
|
||||
self.device_manager.create_stream(device)
|
||||
if device.type != "cpu" and self._auto_transfer
|
||||
else None
|
||||
)
|
||||
self.next_batch = None
|
||||
|
||||
def _move_to_device(self, item):
|
||||
if item is None:
|
||||
return None
|
||||
|
||||
def try_move_device(i):
|
||||
try:
|
||||
i = i.to(self.device, non_blocking=self._auto_transfer)
|
||||
except AttributeError:
|
||||
logger.debug(f"Item {i} cannot be moved to device " f"{self.device}.")
|
||||
return i
|
||||
|
||||
with self.device_manager.get_stream_context(self._memcpy_stream):
|
||||
if isinstance(item, collections.abc.Mapping):
|
||||
item_on_device = {k: self._move_to_device(v) for k, v in item.items()}
|
||||
elif isinstance(item, tuple):
|
||||
item_on_device = tuple(self._move_to_device(i) for i in item)
|
||||
elif isinstance(item, list):
|
||||
item_on_device = [self._move_to_device(i) for i in item]
|
||||
elif isinstance(item, torch.Tensor):
|
||||
item_on_device = try_move_device(item)
|
||||
else:
|
||||
logger.debug(
|
||||
f"Data type {type(item)} doesn't support being moved to device."
|
||||
)
|
||||
item_on_device = item
|
||||
|
||||
return item_on_device
|
||||
|
||||
def _wait_for_batch(self, item):
|
||||
if self._memcpy_stream is None:
|
||||
return
|
||||
# Reference:
|
||||
# https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html
|
||||
# The training stream (current) needs to wait until
|
||||
# the memory copy stream finishes.
|
||||
curr_stream = self.device_manager.get_current_stream()
|
||||
curr_stream.wait_stream(self._memcpy_stream)
|
||||
# When a tensor is used by CUDA streams different from
|
||||
# its original allocator, we need to call ``record_stream``
|
||||
# to inform the allocator of all these streams. Otherwise,
|
||||
# the tensor might be freed once it is no longer used by
|
||||
# the creator stream.
|
||||
for i in item:
|
||||
# The Pytorch DataLoader has no restrictions on what is outputted for
|
||||
# each batch. We should only ``record_stream`` if the item has the
|
||||
# ability to do so.
|
||||
try:
|
||||
i.record_stream(curr_stream)
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
def __len__(self):
|
||||
return len(self._dataloader)
|
||||
|
||||
def _prefetch_next_batch(self):
|
||||
next_batch = next(self.dataloader_iter, None)
|
||||
self.next_batch = self._move_to_device(next_batch)
|
||||
|
||||
def __iter__(self):
|
||||
self.dataloader_iter = iter(self._dataloader)
|
||||
self._prefetch_next_batch()
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
next_batch = self.next_batch
|
||||
if next_batch is None:
|
||||
raise StopIteration
|
||||
self._wait_for_batch(next_batch)
|
||||
self._prefetch_next_batch()
|
||||
return next_batch
|
||||
|
||||
|
||||
class _WrappedOptimizer(Optimizer):
|
||||
def __init__(self, optimizer: Optimizer, scaler: Optional[GradScaler] = None):
|
||||
self.optimizer = optimizer
|
||||
self.scaler = scaler
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self.optimizer.state
|
||||
|
||||
@state.setter
|
||||
def state(self, state):
|
||||
self.optimizer.state = state
|
||||
|
||||
@property
|
||||
def param_groups(self):
|
||||
return self.optimizer.param_groups
|
||||
|
||||
@param_groups.setter
|
||||
def param_groups(self, param_groups):
|
||||
self.optimizer.param_groups = param_groups
|
||||
|
||||
@property
|
||||
def defaults(self):
|
||||
return self.optimizer.defaults
|
||||
|
||||
@defaults.setter
|
||||
def defaults(self, defaults):
|
||||
self.optimizer.defaults = defaults
|
||||
|
||||
def add_param_group(self, param_group):
|
||||
self.optimizer.add_param_group(param_group)
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
self.optimizer.load_state_dict(state_dict)
|
||||
|
||||
def state_dict(self):
|
||||
return self.optimizer.state_dict()
|
||||
|
||||
def zero_grad(self):
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
def step(self, closure=None):
|
||||
if self.scaler is not None:
|
||||
self.scaler.step(self.optimizer, closure)
|
||||
self.scaler.update()
|
||||
else:
|
||||
self.optimizer.step(closure)
|
||||
@@ -0,0 +1,5 @@
|
||||
from ray.train.torch.xla.config import TorchXLAConfig
|
||||
|
||||
__all__ = [
|
||||
"TorchXLAConfig",
|
||||
]
|
||||
@@ -0,0 +1,171 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
|
||||
import ray
|
||||
from ray.train._internal.base_worker_group import BaseWorkerGroup
|
||||
from ray.train._internal.utils import get_address_and_port
|
||||
from ray.train.backend import Backend
|
||||
from ray.train.torch import TorchConfig
|
||||
from ray.util import PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@dataclass
|
||||
class TorchXLAConfig(TorchConfig):
|
||||
"""
|
||||
Configuration for torch XLA setup.
|
||||
See https://pytorch.org/xla/release/1.13/index.html for more info.
|
||||
Currently, only "neuron_cores" accelerator (AwsNeuronXLABackend)
|
||||
is supported with xrt runtime.
|
||||
"""
|
||||
|
||||
neuron_parallel_compile: bool = False
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _TorchAwsNeuronXLABackend
|
||||
|
||||
|
||||
def _kill_xrt_server():
|
||||
import subprocess
|
||||
|
||||
subprocess.call(["pkill", "-f", "xrt_run_server"])
|
||||
|
||||
|
||||
def _set_xla_env_vars():
|
||||
# https://pytorch.org/docs/1.13/elastic/run.html#environment-variables
|
||||
context = ray.train.get_context()
|
||||
|
||||
os.environ["LOCAL_RANK"] = str(context.get_local_rank())
|
||||
os.environ["RANK"] = str(context.get_world_rank())
|
||||
os.environ["LOCAL_WORLD_SIZE"] = str(context.get_local_world_size())
|
||||
os.environ["WORLD_SIZE"] = str(context.get_world_size())
|
||||
os.environ["GROUP_RANK"] = str(context.get_node_rank())
|
||||
os.environ["GROUP_WORLD_SIZE"] = str(
|
||||
context.get_world_size() / context.get_local_world_size()
|
||||
)
|
||||
os.environ["ROLE_RANK"] = str(context.get_world_rank())
|
||||
os.environ["ROLE_WORLD_RANK"] = str(context.get_world_rank())
|
||||
os.environ["ROLE_WORLD_SIZE"] = str(context.get_world_size())
|
||||
|
||||
# EFA and XLA setup
|
||||
# https://github.com/aws/libfabric/blob/master/prov/efa/src/rxr/rxr_init.c
|
||||
# https://github.com/aws-neuron/aws-neuron-samples/blob/master/torch-neuronx/training/dp_bert_hf_pretrain/run_dp_bert_large_hf_pretrain_bf16_s128.sh # noqa
|
||||
os.environ["FI_PROVIDER"] = "efa"
|
||||
os.environ["FI_EFA_USE_DEVICE_RDMA"] = "1"
|
||||
os.environ["FI_EFA_FORK_SAFE"] = "1"
|
||||
os.environ["XLA_TRANSFER_SEED_ASYNC"] = "1"
|
||||
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1"
|
||||
|
||||
|
||||
def _setup_xla_torch_process_group():
|
||||
try:
|
||||
import torch.distributed as dist
|
||||
import torch_xla.core.xla_model as xm # noqa F401
|
||||
import torch_xla.distributed.xla_backend # noqa F401
|
||||
|
||||
dist.init_process_group("xla")
|
||||
except ImportError:
|
||||
raise ImportError("torch_xla must be installed to use torch_xla backend.")
|
||||
|
||||
|
||||
# The following env vars enable Neuron graph extraction for parallel compilation
|
||||
# Note: model outputs are invalid and should be ignored while these env vars are set
|
||||
def _set_neuron_parallel_compile_env_vars():
|
||||
os.environ["NEURON_PARALLEL_COMPILE"] = "1"
|
||||
os.environ["NEURON_EXTRACT_GRAPHS_ONLY"] = "1"
|
||||
os.environ["NEURON_FALL_BACK_TO_NULL_NEFF"] = "1"
|
||||
|
||||
|
||||
# Compile previously extracted Neuron graphs
|
||||
def _neuron_compile_extracted_graphs():
|
||||
try:
|
||||
from libneuronxla.neuron_cc_cache import CacheUrl
|
||||
from libneuronxla.neuron_parallel_compile import parallel_compile
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"libneuronxla must be installed to use Neuron parallel compilation."
|
||||
)
|
||||
|
||||
# Only 1 worker per node should run parallel_compile()
|
||||
if os.environ.get("LOCAL_RANK") == "0":
|
||||
logger.info("Compiling extracted graphs on local rank0 worker")
|
||||
|
||||
parallel_compile_workdir = (
|
||||
f"/tmp/{os.environ.get('USER','no-user')}/parallel_compile_workdir/"
|
||||
)
|
||||
if os.path.exists(parallel_compile_workdir):
|
||||
shutil.rmtree(parallel_compile_workdir)
|
||||
os.makedirs(parallel_compile_workdir, exist_ok=True)
|
||||
|
||||
# Users can set the cache directory using --cache_dir in NEURON_CC_FLAGS or by
|
||||
# using NEURON_COMPILE_CACHE_URL. --cache_dir takes precedence.
|
||||
explicit_cache_dir = None
|
||||
if neuron_cc_flags := os.environ.get("NEURON_CC_FLAGS"):
|
||||
if s := re.search(r"--cache_dir[= ](\S+)", neuron_cc_flags):
|
||||
explicit_cache_dir = s.group(1)
|
||||
|
||||
parallel_compile(
|
||||
parallel_compile_workdir,
|
||||
CacheUrl.get_cache_url(explicit_cache_dir),
|
||||
)
|
||||
|
||||
|
||||
class _TorchAwsNeuronXLABackend(Backend):
|
||||
unique_run_id: str = str(uuid.uuid4())
|
||||
|
||||
def on_start(self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig):
|
||||
"""Logic ran right before training is started."""
|
||||
|
||||
# On previous worker failure, we don't run graceful shutdown on workers.
|
||||
# This would leak any running xrt server.
|
||||
worker_group.execute(_kill_xrt_server)
|
||||
|
||||
# Get master address and port from the first worker.
|
||||
master_addr, master_port = worker_group.execute_single(0, get_address_and_port)
|
||||
|
||||
def set_env_vars(addr, port):
|
||||
os.environ["MASTER_ADDR"] = addr
|
||||
os.environ["MASTER_PORT"] = str(port)
|
||||
# To trigger the xrt server
|
||||
os.environ["TORCHELASTIC_RUN_ID"] = self.unique_run_id
|
||||
|
||||
# Set the env vars on all workers.
|
||||
worker_group.execute(set_env_vars, addr=master_addr, port=master_port)
|
||||
|
||||
# Set up env vars for neuron parallel compilation graph extraction
|
||||
if backend_config.neuron_parallel_compile:
|
||||
logger.info("Extracting graphs for Neuron parallel compilation")
|
||||
worker_group.execute(_set_neuron_parallel_compile_env_vars)
|
||||
|
||||
def on_training_start(
|
||||
self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig
|
||||
):
|
||||
"""
|
||||
Configure the environment variables for the worker group.
|
||||
And initialize the xla distributed process group.
|
||||
TODO: Current setup only supports homogenous cluster with
|
||||
neuron_cores accelerator and xrt runtime.
|
||||
"""
|
||||
worker_group.execute(_set_xla_env_vars)
|
||||
worker_group.execute(_setup_xla_torch_process_group)
|
||||
|
||||
def on_shutdown(
|
||||
self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig
|
||||
):
|
||||
"""
|
||||
Logic ran right after training is finished.
|
||||
This is a sanity cleanup to kill xrt server, and to optionally
|
||||
run neuron parallel graph compilation
|
||||
"""
|
||||
worker_group.execute(_kill_xrt_server)
|
||||
|
||||
# Compile the extracted graphs. This must run at end of training.
|
||||
if backend_config.neuron_parallel_compile:
|
||||
worker_group.execute(_neuron_compile_extracted_graphs)
|
||||
Reference in New Issue
Block a user