chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,15 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
try:
|
||||
import jax # noqa: F401
|
||||
except ModuleNotFoundError as exception:
|
||||
raise ModuleNotFoundError(
|
||||
"Jax isn't installed. To install Jax, please check"
|
||||
" `https://github.com/google/jax#installation` for the instructions."
|
||||
) from exception
|
||||
|
||||
from ray.train.v2.jax.config import JaxConfig
|
||||
from ray.train.v2.jax.jax_trainer import JaxTrainer
|
||||
|
||||
__all__ = ["JaxConfig", "JaxTrainer"]
|
||||
@@ -0,0 +1,227 @@
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray._private import ray_constants
|
||||
from ray.train._internal.utils import get_address_and_port
|
||||
from ray.train._internal.worker_group import WorkerGroup
|
||||
from ray.train.backend import Backend, BackendConfig
|
||||
from ray.train.constants import (
|
||||
DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
from ray.util import PublicAPI
|
||||
from ray.util.tpu import get_tpu_coordinator_env_vars, get_tpu_worker_resources
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# JAX multi-slice (megascale) coordination env vars whose values are
|
||||
# authoritatively computed by the controller for the current worker group.
|
||||
# These must always be overwritten on each worker (even when the underlying
|
||||
# TPU node provider has already set them in the pod environment), because
|
||||
# stale values from a previous worker group configuration -- e.g. after a
|
||||
# slice was preempted and replaced -- would otherwise cause libtpu's
|
||||
# megascale topology coordinator to wait for slices that no longer exist
|
||||
# and hang TPU initialization.
|
||||
_JAX_MULTISLICE_OVERRIDE_KEYS = frozenset(
|
||||
{
|
||||
"MEGASCALE_COORDINATOR_ADDRESS",
|
||||
"MEGASCALE_NUM_SLICES",
|
||||
"MEGASCALE_SLICE_ID",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@dataclass
|
||||
class JaxConfig(BackendConfig):
|
||||
use_tpu: bool = False
|
||||
use_gpu: bool = False
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _JaxBackend
|
||||
|
||||
@property
|
||||
def framework(self):
|
||||
return TrainingFramework.JAX
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
config_dict = {
|
||||
"use_tpu": self.use_tpu,
|
||||
"use_gpu": self.use_gpu,
|
||||
}
|
||||
return config_dict
|
||||
|
||||
|
||||
def _setup_jax_distributed_environment(
|
||||
master_addr_with_port: str,
|
||||
num_workers: int,
|
||||
index: int,
|
||||
use_tpu: bool,
|
||||
use_gpu: bool,
|
||||
resources_per_worker: dict,
|
||||
jax_env_vars: Optional[dict] = None,
|
||||
):
|
||||
"""Set up distributed Jax training information.
|
||||
|
||||
This function should be called on each worker. It sets JAX environment
|
||||
variables and initializes JAX distributed training.
|
||||
|
||||
Args:
|
||||
master_addr_with_port: The master address with port for coordination.
|
||||
num_workers: Total number of workers.
|
||||
index: Index of this worker.
|
||||
use_tpu: Whether to configure for TPU. If True and JAX_PLATFORMS is not
|
||||
already set, it will be set to "tpu".
|
||||
use_gpu: Whether to configure for GPU. If True and JAX_PLATFORMS is not
|
||||
already set, it will be set to "cuda".
|
||||
resources_per_worker: The resources per worker.
|
||||
jax_env_vars: The JAX coordinator env vars to inject for multi-slice.
|
||||
Multi-slice coordination keys (``MEGASCALE_*``) always overwrite
|
||||
any pre-existing value in the environment; all other keys are
|
||||
only set if they are not already present.
|
||||
"""
|
||||
# Get JAX_PLATFORMS from environment if already set
|
||||
jax_platforms = os.environ.get("JAX_PLATFORMS", "").lower()
|
||||
|
||||
if not jax_platforms and use_tpu:
|
||||
os.environ["JAX_PLATFORMS"] = "tpu"
|
||||
jax_platforms = "tpu"
|
||||
|
||||
if jax_env_vars:
|
||||
for k, v in jax_env_vars.items():
|
||||
# For multi-slice coordination keys, always override -- the
|
||||
# controller's freshly computed value is the source of truth and
|
||||
# may differ from what the TPU node provider baked into the pod
|
||||
# environment (e.g. after a slice replacement following preemption).
|
||||
# For all other keys, respect any pre-existing value.
|
||||
if k in _JAX_MULTISLICE_OVERRIDE_KEYS or k not in os.environ:
|
||||
os.environ[k] = v
|
||||
|
||||
if not jax_platforms and use_gpu:
|
||||
os.environ["JAX_PLATFORMS"] = "cuda"
|
||||
jax_platforms = "cuda"
|
||||
|
||||
if "cuda" in jax_platforms.split(","):
|
||||
num_gpus_per_worker = resources_per_worker.get("GPU", 0)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
|
||||
str(i) for i in range(num_gpus_per_worker)
|
||||
)
|
||||
|
||||
import jax
|
||||
|
||||
jax_platforms_list = jax_platforms.split(",")
|
||||
|
||||
if "tpu" in jax_platforms_list:
|
||||
jax.distributed.initialize(master_addr_with_port, num_workers, index)
|
||||
logger.info("Initialized JAX distributed on TPU.")
|
||||
elif "cuda" in jax_platforms_list:
|
||||
if num_gpus_per_worker > 0:
|
||||
local_device_ids = list(range(num_gpus_per_worker))
|
||||
else:
|
||||
local_device_ids = 0
|
||||
jax.distributed.initialize(
|
||||
master_addr_with_port, num_workers, index, local_device_ids
|
||||
)
|
||||
logger.info("Initialized JAX distributed on CUDA.")
|
||||
elif "cpu" in jax_platforms_list:
|
||||
jax.distributed.initialize(master_addr_with_port, num_workers, index)
|
||||
logger.info("Initialized JAX distributed on CPU.")
|
||||
|
||||
|
||||
def _shutdown_jax_distributed():
|
||||
"""Shutdown JAX distributed environment.
|
||||
|
||||
This function should be called on each worker during cleanup.
|
||||
If JAX distributed was not initialized, this is a no-op.
|
||||
"""
|
||||
try:
|
||||
import jax
|
||||
|
||||
jax.distributed.shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during JAX distributed shutdown: {e}")
|
||||
|
||||
|
||||
class _JaxBackend(Backend):
|
||||
def on_start(self, worker_group: WorkerGroup, backend_config: JaxConfig):
|
||||
if not backend_config.use_tpu and not backend_config.use_gpu:
|
||||
return
|
||||
|
||||
master_addr, master_port = worker_group.execute_single(0, get_address_and_port)
|
||||
master_addr_with_port = f"{master_addr}:{master_port}"
|
||||
|
||||
if backend_config.use_tpu and hasattr(worker_group, "get_worker_group_context"):
|
||||
num_slices = worker_group.get_worker_group_context().num_slices
|
||||
else:
|
||||
num_slices = 1
|
||||
|
||||
# Calculate the number of workers per slice for multi-slice env setup.
|
||||
if backend_config.use_tpu and num_slices > 1:
|
||||
# Handle the case where a user requests less workers than the total
|
||||
# capacity of the TPU slice.
|
||||
scaling_config = worker_group._train_run_context.scaling_config
|
||||
workers_per_slice, _ = get_tpu_worker_resources(
|
||||
topology=scaling_config.topology,
|
||||
accelerator_type=scaling_config.accelerator_type,
|
||||
resources_per_unit=scaling_config.resources_per_worker,
|
||||
num_slices=1,
|
||||
)
|
||||
else:
|
||||
# Assume even distribution based on the requested number of workers.
|
||||
workers_per_slice = max(1, len(worker_group) // num_slices)
|
||||
|
||||
# Set up JAX distributed environment on all workers
|
||||
num_workers_total = len(worker_group)
|
||||
setup_futures = []
|
||||
for i in range(num_workers_total):
|
||||
env_vars = {}
|
||||
if num_slices > 1:
|
||||
slice_id = min(i // workers_per_slice, num_slices - 1)
|
||||
env_vars = get_tpu_coordinator_env_vars(
|
||||
coordinator_address=master_addr,
|
||||
num_slices=num_slices,
|
||||
slice_id=slice_id,
|
||||
)
|
||||
|
||||
setup_futures.append(
|
||||
worker_group.execute_single_async(
|
||||
i,
|
||||
_setup_jax_distributed_environment,
|
||||
master_addr_with_port=master_addr_with_port,
|
||||
num_workers=len(worker_group),
|
||||
index=i,
|
||||
use_tpu=backend_config.use_tpu,
|
||||
use_gpu=backend_config.use_gpu,
|
||||
resources_per_worker=worker_group.get_resources_per_worker(),
|
||||
jax_env_vars=env_vars,
|
||||
)
|
||||
)
|
||||
ray.get(setup_futures)
|
||||
|
||||
def on_shutdown(self, worker_group: WorkerGroup, backend_config: JaxConfig):
|
||||
"""Cleanup JAX distributed resources when shutting down worker group."""
|
||||
if not backend_config.use_tpu and not backend_config.use_gpu:
|
||||
return
|
||||
|
||||
# Shutdown JAX distributed on all workers
|
||||
shutdown_futures = worker_group.execute_async(_shutdown_jax_distributed)
|
||||
|
||||
timeout_s = ray_constants.env_integer(
|
||||
JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
try:
|
||||
ray.get(shutdown_futures, timeout=timeout_s)
|
||||
logger.debug("JAX distributed shutdown completed")
|
||||
except ray.exceptions.GetTimeoutError:
|
||||
logger.warning(
|
||||
f"JAX distributed shutdown timed out after {timeout_s} seconds. "
|
||||
"This may indicate workers are hung or unresponsive."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during JAX distributed shutdown: {e}")
|
||||
@@ -0,0 +1,169 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
|
||||
from ray.train import DataConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
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.train.v2.jax.config import JaxConfig
|
||||
from ray.util import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class JaxTrainer(DataParallelTrainer):
|
||||
"""A Trainer for Single-Program Multi-Data (SPMD) JAX 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 JAX environment for TPUs or GPUs
|
||||
on these workers as defined by the ``jax_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:`Jax Guide <train-jax>`
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import os
|
||||
from absl import app
|
||||
import logging
|
||||
from typing import Sequence
|
||||
|
||||
import ray
|
||||
from ray.train import ScalingConfig, RunConfig
|
||||
from ray.train.v2.jax import JaxTrainer
|
||||
from MaxText.train import main as maxtext_main
|
||||
|
||||
def train_loop_per_worker(config):
|
||||
argv = config["argv"]
|
||||
maxtext_main(argv)
|
||||
|
||||
def main(argv: Sequence[str]):
|
||||
ray.init()
|
||||
|
||||
# If you want to use TPUs, specify the TPU topology and accelerator type.
|
||||
tpu_scaling_config = ScalingConfig(
|
||||
use_tpu=True,
|
||||
num_workers=4,
|
||||
topology="4x4",
|
||||
accelerator_type="TPU-V6E",
|
||||
placement_strategy="SPREAD",
|
||||
resources_per_worker={"TPU": 4},
|
||||
)
|
||||
|
||||
# If you want to use GPUs, specify the GPU scaling config like below.
|
||||
# gpu_scaling_config = ScalingConfig(
|
||||
# use_gpu=True,
|
||||
# num_workers=4,
|
||||
# resources_per_worker={"GPU": 1},
|
||||
# )
|
||||
|
||||
|
||||
trainer = JaxTrainer(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config={"argv": absolute_argv},
|
||||
scaling_config=tpu_scaling_config,
|
||||
run_config=RunConfig(
|
||||
name="maxtext_jaxtrainer",
|
||||
worker_runtime_env={
|
||||
"env_vars": {
|
||||
"JAX_PLATFORMS": "tpu",
|
||||
# If you want to use GPUs, set the JAX_PLATFORMS to "cuda".
|
||||
# "JAX_PLATFORMS": "cuda",
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
result = trainer.fit()
|
||||
|
||||
If the ``datasets`` dict contains datasets (e.g. "train" and "val"), then it will be split into multiple dataset
|
||||
shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` and ``ray.train.get_dataset_shard("val")``.
|
||||
|
||||
Note:
|
||||
* If you are using TPUs, importing `jax` should occur within `train_loop_per_worker` to
|
||||
avoid driver-side TPU lock issues.
|
||||
|
||||
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.
|
||||
jax_config: The configuration for setting up the JAX backend.
|
||||
If set to None, a default configuration will be used based on the ``scaling_config`` and ``JAX_PLATFORMS`` environment variable.
|
||||
scaling_config: Configuration for how to scale data parallel training
|
||||
with SPMD. ``num_workers`` should be set to the number of TPU hosts or GPU workers.
|
||||
If using TPUs, ``topology`` should be set to the TPU topology.
|
||||
See :class:`~ray.train.ScalingConfig` for more info.
|
||||
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.
|
||||
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``.
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
jax_config: Optional[JaxConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
dataset_config: Optional[Dict[str, DataConfig]] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
if not jax_config:
|
||||
jax_config = JaxConfig(
|
||||
use_tpu=scaling_config.use_tpu,
|
||||
use_gpu=scaling_config.use_gpu,
|
||||
)
|
||||
super(JaxTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=jax_config,
|
||||
scaling_config=scaling_config,
|
||||
dataset_config=dataset_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig:
|
||||
"""Return scaling config dataclass after validating updated keys."""
|
||||
ensure_only_allowed_dataclass_keys_updated(
|
||||
dataclass=scaling_config,
|
||||
allowed_keys=cls._scaling_config_allowed_keys,
|
||||
)
|
||||
|
||||
return scaling_config
|
||||
Reference in New Issue
Block a user