chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,2 @@
|
||||
# Note: This import is to avoid circular import errors
|
||||
import ray.train.torch # noqa: F401
|
||||
@@ -0,0 +1,213 @@
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.train import Checkpoint, DataConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.v2._internal.execution.local_mode.torch import LocalTorchController
|
||||
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
||||
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.util import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# NOTE: `ray.train.torch` module imports in this file will break
|
||||
# with a circular import error if the TorchTrainer class is captured
|
||||
# in the scope of a Ray task.
|
||||
from ray.train.torch.config import TorchConfig
|
||||
|
||||
|
||||
@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::
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
import ray.train
|
||||
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 = 2
|
||||
|
||||
# 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_fn_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
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
torch.save(
|
||||
{"model_state_dict": base_model.state_dict()},
|
||||
os.path.join(temp_checkpoint_dir, "model.pt"),
|
||||
)
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
# Report metrics and checkpoint.
|
||||
ray.train.report({"loss": loss.item()}, checkpoint=checkpoint)
|
||||
|
||||
# Define datasets.
|
||||
train_dataset = ray.data.from_items(
|
||||
[{"input": [x], "label": [2 * x + 1]} for x in range(128)]
|
||||
)
|
||||
|
||||
# Initialize the Trainer.
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
train_loop_config={"num_epochs": 1, "lr": 0.01, "batch_size": 32},
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
|
||||
datasets={"train": train_dataset},
|
||||
)
|
||||
|
||||
# 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.
|
||||
validation_config: [Alpha] Configuration for checkpoint validation.
|
||||
If provided and ``ray.train.report`` is called with the ``validation``
|
||||
argument, Ray Train will validate the reported checkpoint using
|
||||
the validation function specified in this config.
|
||||
metadata: [Deprecated]
|
||||
resume_from_checkpoint: [Deprecated]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
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,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
# TODO: [Deprecated]
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
from ray.train.torch.config import TorchConfig
|
||||
|
||||
torch_config = torch_config or TorchConfig()
|
||||
if not torch_config.backend:
|
||||
is_gpu_training = scaling_config and scaling_config.use_gpu
|
||||
torch_config.backend = "nccl" if is_gpu_training else "gloo"
|
||||
|
||||
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,
|
||||
run_config=run_config,
|
||||
dataset_config=dataset_config,
|
||||
datasets=datasets,
|
||||
resume_from_checkpoint=resume_from_checkpoint,
|
||||
metadata=metadata,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
|
||||
def _get_local_controller(self) -> LocalTorchController:
|
||||
return LocalTorchController(
|
||||
experiment_name=self.run_config.name,
|
||||
datasets=self.datasets,
|
||||
)
|
||||
@@ -0,0 +1,117 @@
|
||||
import logging
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Dict
|
||||
|
||||
import ray
|
||||
from ray.train.torch.config import (
|
||||
TorchConfig,
|
||||
_TorchBackend,
|
||||
)
|
||||
from ray.train.v2._internal.constants import TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR
|
||||
from ray.train.v2._internal.execution.worker_group import ReplicaGroup, WorkerGroup
|
||||
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TorchftConfig(TorchConfig):
|
||||
"""Configuration for torchft-based fault tolerant training.
|
||||
|
||||
See https://github.com/meta-pytorch/torchft for more info.
|
||||
|
||||
Args:
|
||||
lighthouse_kwargs: Keyword arguments to pass to the torchft.Lighthouse constructor.
|
||||
**kwargs: Additional keyword arguments to pass to the TorchConfig constructor.
|
||||
"""
|
||||
|
||||
def __init__(self, lighthouse_kwargs: Dict[str, Any], **kwargs):
|
||||
self.lighthouse_kwargs = lighthouse_kwargs
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _TorchftBackend
|
||||
|
||||
|
||||
@ray.remote
|
||||
class LighthouseServerActor:
|
||||
"""Actor that runs the torchft.Lighthouse server.
|
||||
|
||||
ray.remote(LighthouseServer) does not work because it is a PyO3 type.
|
||||
"""
|
||||
|
||||
def __init__(self, lighthouse_kwargs: Dict[str, Any]):
|
||||
from torchft.coordination import LighthouseServer
|
||||
|
||||
self.lighthouse = LighthouseServer(**lighthouse_kwargs)
|
||||
|
||||
def address(self) -> str:
|
||||
return self.lighthouse.address()
|
||||
|
||||
|
||||
class _TorchftBackend(_TorchBackend):
|
||||
"""Backend for torchft-based fault-tolerant training with replica groups.
|
||||
|
||||
Creates a separate process group per replica group by calling
|
||||
the parent _TorchBackend.on_start() once per replica group.
|
||||
"""
|
||||
|
||||
has_replica_groups: bool = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.lighthouse_actor = None
|
||||
|
||||
def _maybe_create_lighthouse_actor(self, backend_config: TorchftConfig) -> str:
|
||||
"""Create lighthouse actor if it doesn't exist and return its address."""
|
||||
if self.lighthouse_actor is not None:
|
||||
# Intentionally read address from actor in case it was restarted.
|
||||
return ray.get(self.lighthouse_actor.address.remote())
|
||||
|
||||
# Let the OS pick a free port by default
|
||||
if "bind" in backend_config.lighthouse_kwargs:
|
||||
lighthouse_kwargs = backend_config.lighthouse_kwargs
|
||||
else:
|
||||
lighthouse_kwargs = {"bind": "[::]:0"} | backend_config.lighthouse_kwargs
|
||||
|
||||
# Store reference so the actor lives as long as the backend/controller.
|
||||
self.lighthouse_actor = LighthouseServerActor.options(
|
||||
# Schedule lightweight lighthouse actor on head node
|
||||
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
||||
node_id=ray.get_runtime_context().get_node_id(),
|
||||
soft=False,
|
||||
)
|
||||
).remote(lighthouse_kwargs=lighthouse_kwargs)
|
||||
lighthouse_address = ray.get(self.lighthouse_actor.address.remote())
|
||||
logger.info(f"Created torchft lighthouse at {lighthouse_address}")
|
||||
return lighthouse_address
|
||||
|
||||
def on_start(self, worker_group, backend_config: TorchftConfig):
|
||||
lighthouse_address = self._maybe_create_lighthouse_actor(backend_config)
|
||||
|
||||
# Push the lighthouse address to all workers in this group.
|
||||
# Necessary because workers were already started before on_start runs.
|
||||
def _set_lighthouse_address(addr: str):
|
||||
os.environ[TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR] = addr
|
||||
|
||||
worker_group.execute(_set_lighthouse_address, lighthouse_address)
|
||||
|
||||
# Bind super() eagerly — the zero-arg form relies on a __class__ cell
|
||||
# that doesn't transfer correctly into ThreadPoolExecutor submissions.
|
||||
parent_on_start = super().on_start
|
||||
|
||||
if isinstance(worker_group, ReplicaGroup):
|
||||
# Single replica group replacement — just start this one.
|
||||
parent_on_start(worker_group, backend_config)
|
||||
else:
|
||||
# Full worker group startup — start all replica groups in parallel.
|
||||
assert isinstance(worker_group, WorkerGroup)
|
||||
replica_groups = worker_group.get_replica_groups()
|
||||
with ThreadPoolExecutor() as executor:
|
||||
futures = [
|
||||
executor.submit(parent_on_start, rg, backend_config)
|
||||
for rg in replica_groups
|
||||
]
|
||||
for f in futures:
|
||||
f.result()
|
||||
@@ -0,0 +1,435 @@
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import (
|
||||
DataLoader,
|
||||
DistributedSampler,
|
||||
IterableDataset,
|
||||
RandomSampler,
|
||||
SequentialSampler,
|
||||
)
|
||||
|
||||
import ray.train.torch
|
||||
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
||||
from ray.train.torch.train_loop_utils import (
|
||||
_WrappedDataLoader,
|
||||
get_devices as get_devices_distributed,
|
||||
)
|
||||
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
||||
from ray.train.v2._internal.util import requires_train_worker
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_TORCH_AMP_DEPRECATION_MESSAGE = (
|
||||
"The `accelerate`, `backward`, and `prepare_optimizer` utility methods "
|
||||
"in the `ray.train.torch` module are deprecated and will be removed in a "
|
||||
"future release. "
|
||||
"Please use the native PyTorch mixed precision API directly, or "
|
||||
"a library such as Lightning or HuggingFace Transformers/Accelerate. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker()
|
||||
def get_device() -> torch.device:
|
||||
"""Gets the correct torch device configured for the current worker.
|
||||
|
||||
Returns the torch device for the current worker. If more than 1 GPU is
|
||||
requested per worker, returns the device with the lowest 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()`.
|
||||
|
||||
Returns:
|
||||
The torch device assigned to the current worker.
|
||||
|
||||
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")
|
||||
"""
|
||||
return get_devices()[0]
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
@requires_train_worker()
|
||||
def get_devices() -> List[torch.device]:
|
||||
"""Gets the list of torch devices configured for the current worker.
|
||||
|
||||
Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
|
||||
superset of the `ray.get_gpu_ids()`.
|
||||
|
||||
Returns:
|
||||
The list of torch devices assigned to the current worker.
|
||||
|
||||
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")]
|
||||
"""
|
||||
if get_train_fn_utils().is_distributed():
|
||||
return get_devices_distributed()
|
||||
else:
|
||||
# Local mode, we defer to torch.cuda
|
||||
# TODO(xgui): Use `ScalingConfig.use_gpu` instead
|
||||
if torch.cuda.is_available():
|
||||
return [torch.device(f"cuda:{torch.cuda.current_device()}")]
|
||||
else:
|
||||
return [torch.device("cpu")]
|
||||
|
||||
|
||||
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")
|
||||
|
||||
parallel_strategy_kwargs = parallel_strategy_kwargs or {}
|
||||
|
||||
rank = ray.train.get_context().get_local_rank()
|
||||
|
||||
if isinstance(move_to_device, torch.device):
|
||||
device = move_to_device
|
||||
else:
|
||||
device = ray.train.torch.get_device()
|
||||
if isinstance(device, list):
|
||||
device = device[0]
|
||||
|
||||
if torch.cuda.is_available() and device.type == "cuda":
|
||||
torch.cuda.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)
|
||||
|
||||
world_size = ray.train.get_context().get_world_size()
|
||||
|
||||
if parallel_strategy and world_size > 1:
|
||||
if parallel_strategy == "ddp":
|
||||
DataParallel = DistributedDataParallel
|
||||
if torch.cuda.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
|
||||
|
||||
|
||||
@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")
|
||||
|
||||
world_size = ray.train.get_context().get_world_size()
|
||||
world_rank = ray.train.get_context().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)
|
||||
worker_init_fn: Optional[Callable[[int], None]] = loader.worker_init_fn
|
||||
generator: Optional[torch.Generator] = loader.generator
|
||||
|
||||
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 = ray.train.torch.get_device()
|
||||
data_loader = _WrappedDataLoader(data_loader, device, auto_transfer)
|
||||
|
||||
return data_loader
|
||||
|
||||
|
||||
@Deprecated
|
||||
def accelerate(amp: bool = False) -> None:
|
||||
raise DeprecationWarning(_TORCH_AMP_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@Deprecated
|
||||
def prepare_optimizer(optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
|
||||
raise DeprecationWarning(_TORCH_AMP_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@Deprecated
|
||||
def backward(tensor: torch.Tensor) -> None:
|
||||
raise DeprecationWarning(_TORCH_AMP_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@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>`_.
|
||||
"""
|
||||
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"
|
||||
Reference in New Issue
Block a user