wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/train/api/api.rst
+244
View File
@@ -0,0 +1,244 @@
.. _train-api:
Ray Train API
=============
.. currentmodule:: ray
.. important::
These API references are for the revamped Ray Train V2 implementation that is available starting from Ray 2.43
by enabling the environment variable ``RAY_TRAIN_V2_ENABLED=1``. These APIs assume that the environment variable has been enabled.
See :ref:`train-deprecated-api` for the old API references and the `Ray Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_.
PyTorch Ecosystem
-----------------
.. autosummary::
:nosignatures:
:toctree: doc/
~train.torch.TorchTrainer
~train.torch.TorchConfig
~train.torch.xla.TorchXLAConfig
.. _train-pytorch-integration:
PyTorch
~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.torch.get_device
~train.torch.get_devices
~train.torch.prepare_model
~train.torch.prepare_data_loader
~train.torch.enable_reproducibility
.. _train-lightning-integration:
PyTorch Lightning
~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.lightning.prepare_trainer
~train.lightning.RayLightningEnvironment
~train.lightning.RayDDPStrategy
~train.lightning.RayFSDPStrategy
~train.lightning.RayDeepSpeedStrategy
~train.lightning.RayTrainReportCallback
.. _train-transformers-integration:
Hugging Face Transformers
~~~~~~~~~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.huggingface.transformers.prepare_trainer
~train.huggingface.transformers.RayTrainReportCallback
More Frameworks
---------------
TensorFlow/Keras
~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.tensorflow.TensorflowTrainer
~train.tensorflow.TensorflowConfig
~train.tensorflow.prepare_dataset_shard
~train.tensorflow.keras.ReportCheckpointCallback
XGBoost
~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.xgboost.XGBoostTrainer
~train.xgboost.RayTrainReportCallback
LightGBM
~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.lightgbm.LightGBMTrainer
~train.lightgbm.get_network_params
~train.lightgbm.RayTrainReportCallback
~train.lightgbm.normalize_pandas_for_lightgbm
JAX
~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.v2.jax.JaxTrainer
.. _ray-train-configs-api:
Ray Train Configuration
-----------------------
.. autosummary::
:nosignatures:
:toctree: doc/
~train.CheckpointConfig
~train.DataConfig
~train.FailureConfig
~train.LoggingConfig
~train.RunConfig
~train.ScalingConfig
~train.ValidationConfig
.. _train-loop-api:
Ray Train Utilities
-------------------
**Classes**
.. autosummary::
:nosignatures:
:toctree: doc/
~train.Checkpoint
~train.CheckpointUploadMode
~train.CheckpointConsistencyMode
~train.TrainContext
~train.ValidationFn
~train.ValidationTaskConfig
**Functions**
.. autosummary::
:nosignatures:
:toctree: doc/
~train.get_all_reported_checkpoints
~train.get_checkpoint
~train.get_context
~train.get_dataset_shard
~train.report
**Collective**
.. autosummary::
:nosignatures:
:toctree: doc/
~train.collective.barrier
~train.collective.broadcast_from_rank_zero
Ray Train Output
----------------
.. autosummary::
:nosignatures:
:template: autosummary/class_without_autosummary.rst
:toctree: doc/
~train.ReportedCheckpoint
~train.ReportedCheckpointStatus
~train.Result
Ray Train Errors
----------------
.. autosummary::
:nosignatures:
:template: autosummary/class_without_autosummary.rst
:toctree: doc/
~train.ControllerError
~train.WorkerGroupError
~train.TrainingFailedError
Ray Tune Integration Utilities
------------------------------
.. autosummary::
:nosignatures:
:toctree: doc/
tune.integration.ray_train.TuneReportCallback
Ray Train Developer APIs
------------------------
Trainer Base Class
~~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.v2.api.data_parallel_trainer.DataParallelTrainer
Train Backend Base Classes
~~~~~~~~~~~~~~~~~~~~~~~~~~
.. _train-backend:
.. _train-backend-config:
.. autosummary::
:nosignatures:
:toctree: doc/
:template: autosummary/class_without_autosummary.rst
~train.backend.Backend
~train.backend.BackendConfig
Trainer Callbacks
~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.UserCallback
doc/source/train/api/deprecated.rst
+197
View File
@@ -0,0 +1,197 @@
:orphan:
.. _train-deprecated-api:
Ray Train V1 API
================
.. currentmodule:: ray
.. important::
Ray Train V2 is an overhaul of Ray Train's implementation and select APIs, which can be enabled by setting the environment variable ``RAY_TRAIN_V2_ENABLED=1`` starting in Ray 2.43.
This page contains the deprecated V1 API references. See :ref:`train-api` for the new V2 API references and the `Ray Train V2 Migration Guide <https://github.com/ray-project/ray/issues/49454>`_.
PyTorch Ecosystem
-----------------
.. autosummary::
:nosignatures:
:toctree: doc/
~train.torch.torch_trainer.TorchTrainer
~train.torch.TorchConfig
~train.torch.xla.TorchXLAConfig
PyTorch
~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.torch.get_device
~train.torch.get_devices
~train.torch.prepare_model
~train.torch.prepare_data_loader
~train.torch.enable_reproducibility
PyTorch Lightning
~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.lightning.prepare_trainer
~train.lightning.RayLightningEnvironment
~train.lightning.RayDDPStrategy
~train.lightning.RayFSDPStrategy
~train.lightning.RayDeepSpeedStrategy
~train.lightning.RayTrainReportCallback
Hugging Face Transformers
~~~~~~~~~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.huggingface.transformers.prepare_trainer
~train.huggingface.transformers.RayTrainReportCallback
More Frameworks
---------------
TensorFlow/Keras
~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.tensorflow.tensorflow_trainer.TensorflowTrainer
~train.tensorflow.TensorflowConfig
~train.tensorflow.prepare_dataset_shard
~train.tensorflow.keras.ReportCheckpointCallback
Horovod
~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.horovod.HorovodTrainer
~train.horovod.HorovodConfig
XGBoost
~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.xgboost.xgboost_trainer.XGBoostTrainer
~train.xgboost.RayTrainReportCallback
LightGBM
~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.lightgbm.lightgbm_trainer.LightGBMTrainer
~train.lightgbm.RayTrainReportCallback
Ray Train Configuration
-----------------------
.. autosummary::
:nosignatures:
:toctree: doc/
~air.config.ScalingConfig
~air.config.RunConfig
~air.config.FailureConfig
~train.CheckpointConfig
~train.DataConfig
~train.SyncConfig
Ray Train Utilities
-------------------
**Classes**
.. autosummary::
:nosignatures:
:toctree: doc/
~train.Checkpoint
~train.context.TrainContext
**Functions**
.. autosummary::
:nosignatures:
:toctree: doc/
~train._internal.session.get_checkpoint
~train.context.get_context
~train._internal.session.get_dataset_shard
~train._internal.session.report
Ray Train Output
----------------
.. autosummary::
:nosignatures:
:toctree: doc/
~train.Result
Ray Train Errors
----------------
.. autosummary::
:nosignatures:
:template: autosummary/class_without_autosummary.rst
:toctree: doc/
~train.error.SessionMisuseError
~train.base_trainer.TrainingFailedError
Ray Train Developer APIs
------------------------
Trainer Base Classes
~~~~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
~train.trainer.BaseTrainer
~train.data_parallel_trainer.DataParallelTrainer
Train Backend Base Classes
~~~~~~~~~~~~~~~~~~~~~~~~~~
.. autosummary::
:nosignatures:
:toctree: doc/
:template: autosummary/class_without_autosummary.rst
~train.backend.Backend
~train.backend.BackendConfig
doc/source/train/benchmarks.rst
+167
View File
@@ -0,0 +1,167 @@
.. _train-benchmarks:
Ray Train Benchmarks
====================
Below we document key performance benchmarks for common Ray Train tasks and workflows.
.. _pytorch_gpu_training_benchmark:
GPU image training
------------------
This task uses the TorchTrainer module to train different amounts of data
using a PyTorch ResNet model.
We test out the performance across different cluster sizes and data sizes.
- `GPU image training script`_
- `GPU training small cluster configuration`_
- `GPU training large cluster configuration`_
.. note::
For multi-host distributed training, on AWS we need to ensure ec2 instances are in the same VPC and
all ports are open in the security group.
.. list-table::
* - **Cluster Setup**
- **Data Size**
- **Performance**
- **Command**
* - 1 g3.8xlarge node (1 worker)
- 1 GB (1623 images)
- 79.76 s (2 epochs, 40.7 images/sec)
- `python pytorch_training_e2e.py --data-size-gb=1`
* - 1 g3.8xlarge node (1 worker)
- 20 GB (32460 images)
- 1388.33 s (2 epochs, 46.76 images/sec)
- `python pytorch_training_e2e.py --data-size-gb=20`
* - 4 g3.16xlarge nodes (16 workers)
- 100 GB (162300 images)
- 434.95 s (2 epochs, 746.29 images/sec)
- `python pytorch_training_e2e.py --data-size-gb=100 --num-workers=16`
.. _pytorch-training-parity:
PyTorch training parity
-----------------------
This task checks the performance parity between native PyTorch Distributed and
Ray Train's distributed TorchTrainer.
We demonstrate that the performance is similar (within 2.5\%) between the two frameworks.
Performance may vary greatly across different model, hardware, and cluster configurations.
The reported times are for the raw training times. There is an unreported constant setup
overhead of a few seconds for both methods that is negligible for longer training runs.
- `PyTorch comparison training script`_
- `PyTorch comparison CPU cluster configuration`_
- `PyTorch comparison GPU cluster configuration`_
.. list-table::
* - **Cluster Setup**
- **Dataset**
- **Performance**
- **Command**
* - 4 m5.2xlarge nodes (4 workers)
- FashionMNIST
- 196.64 s (vs 194.90 s PyTorch)
- `python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 4 --cpus-per-worker 8`
* - 4 m5.2xlarge nodes (16 workers)
- FashionMNIST
- 430.88 s (vs 475.97 s PyTorch)
- `python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 2`
* - 4 g4dn.12xlarge nodes (16 workers)
- FashionMNIST
- 149.80 s (vs 146.46 s PyTorch)
- `python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 4 --use-gpu`
.. _tf-training-parity:
TensorFlow training parity
--------------------------
This task checks the performance parity between native TensorFlow Distributed and
Ray Train's distributed TensorflowTrainer.
We demonstrate that the performance is similar (within 1\%) between the two frameworks.
Performance may vary greatly across different model, hardware, and cluster configurations.
The reported times are for the raw training times. There is an unreported constant setup
overhead of a few seconds for both methods that is negligible for longer training runs.
.. note:: The batch size and number of epochs is different for the GPU benchmark, resulting in a longer runtime.
- `TensorFlow comparison training script`_
- `TensorFlow comparison CPU cluster configuration`_
- `TensorFlow comparison GPU cluster configuration`_
.. list-table::
* - **Cluster Setup**
- **Dataset**
- **Performance**
- **Command**
* - 4 m5.2xlarge nodes (4 workers)
- FashionMNIST
- 78.81 s (versus 79.67 s TensorFlow)
- `python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 4 --cpus-per-worker 8`
* - 4 m5.2xlarge nodes (16 workers)
- FashionMNIST
- 64.57 s (versus 67.45 s TensorFlow)
- `python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 2`
* - 4 g4dn.12xlarge nodes (16 workers)
- FashionMNIST
- 465.16 s (versus 461.74 s TensorFlow)
- `python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 200 --num-workers 16 --cpus-per-worker 4 --batch-size 64 --use-gpu`
.. _xgboost-benchmark:
XGBoost training
----------------
This task uses the XGBoostTrainer module to train on different sizes of data
with different amounts of parallelism to show near-linear scaling from distributed
data parallelism.
XGBoost parameters were kept as defaults for ``xgboost==1.7.6`` this task.
- `XGBoost Training Script`_
- `XGBoost Cluster Configuration`_
.. list-table::
* - **Cluster Setup**
- **Number of distributed training workers**
- **Data Size**
- **Performance**
- **Command**
* - 1 m5.4xlarge node with 16 CPUs
- 1 training worker using 12 CPUs, leaving 4 CPUs for Ray Data tasks
- 10 GB (26M rows)
- 310.22 s
- `python train_batch_inference_benchmark.py "xgboost" --size=10GB`
* - 10 m5.4xlarge nodes
- 10 training workers (one per node), using 10x12 CPUs, leaving 10x4 CPUs for Ray Data tasks
- 100 GB (260M rows)
- 326.86 s
- `python train_batch_inference_benchmark.py "xgboost" --size=100GB`
.. _`GPU image training script`: https://github.com/ray-project/ray/blob/cec82a1ced631525a4d115e4dc0c283fa4275a7f/release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py#L95-L106
.. _`GPU training small cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_1_aws.yaml#L6-L24
.. _`GPU training large cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_4x4_aws.yaml#L5-L25
.. _`PyTorch comparison training script`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/workloads/torch_benchmark.py
.. _`PyTorch comparison CPU cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_cpu_4_aws.yaml
.. _`PyTorch comparison GPU cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_4x4_aws.yaml
.. _`TensorFlow comparison training script`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/workloads/tensorflow_benchmark.py
.. _`TensorFlow comparison CPU cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_cpu_4_aws.yaml
.. _`TensorFlow comparison GPU cluster configuration`: https://github.com/ray-project/ray/blob/master/release/air_tests/air_benchmarks/compute_gpu_4x4_aws.yaml
.. _`XGBoost Training Script`: https://github.com/ray-project/ray/blob/9ac58f4efc83253fe63e280106f959fe317b1104/release/train_tests/xgboost_lightgbm/train_batch_inference_benchmark.py
.. _`XGBoost Cluster Configuration`: https://github.com/ray-project/ray/tree/9ac58f4efc83253fe63e280106f959fe317b1104/release/train_tests/xgboost_lightgbm
doc/source/train/common/torch-configure-run.rst
+86
View File
@@ -0,0 +1,86 @@
Configure scale and GPUs
------------------------
Outside of your training function, create a :class:`~ray.train.ScalingConfig` object to configure:
1. :class:`num_workers <ray.train.ScalingConfig>` - The number of distributed training worker processes.
2. :class:`use_gpu <ray.train.ScalingConfig>` - Whether each worker should use a GPU (or CPU).
.. testcode::
from ray.train import ScalingConfig
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
For more details, see :ref:`train_scaling_config`.
Configure persistent storage
----------------------------
Create a :class:`~ray.train.RunConfig` object to specify the path where results
(including checkpoints and artifacts) will be saved.
.. testcode::
from ray.train import RunConfig
# Local path (/some/local/path/unique_run_name)
run_config = RunConfig(storage_path="/some/local/path", name="unique_run_name")
# Shared cloud storage URI (s3://bucket/unique_run_name)
run_config = RunConfig(storage_path="s3://bucket", name="unique_run_name")
# Shared NFS path (/mnt/nfs/unique_run_name)
run_config = RunConfig(storage_path="/mnt/nfs", name="unique_run_name")
.. warning::
Specifying a *shared storage location* (such as cloud storage or NFS) is
*optional* for single-node clusters, but it is **required for multi-node clusters.**
Using a local path will :ref:`raise an error <multinode-local-storage-warning>`
during checkpointing for multi-node clusters.
For more details, see :ref:`persistent-storage-guide`.
Launch a training job
---------------------
Tying this all together, you can now launch a distributed training job
with a :class:`~ray.train.torch.TorchTrainer`.
.. testcode::
:hide:
from ray.train import ScalingConfig
train_func = lambda: None
scaling_config = ScalingConfig(num_workers=1)
run_config = None
.. testcode::
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
train_func, scaling_config=scaling_config, run_config=run_config
)
result = trainer.fit()
Access training results
-----------------------
After training completes, a :class:`~ray.train.Result` object is returned which contains
information about the training run, including the metrics and checkpoints reported during training.
.. testcode::
result.metrics # The metrics reported during training.
result.checkpoint # The latest checkpoint reported during training.
result.path # The path where logs are stored.
result.error # The exception that was raised, if training failed.
For more usage examples, see :ref:`train-inspect-results`.
doc/source/train/common/torch-configure-train_func.rst
+51
View File
@@ -0,0 +1,51 @@
First, update your training code to support distributed training.
Begin by wrapping your code in a :ref:`training function <train-overview-training-function>`:
.. testcode::
:skipif: True
def train_func():
# Your model training code here.
...
Each distributed training worker executes this function.
You can also specify the input argument for `train_func` as a dictionary via the Trainer's `train_loop_config`. For example:
.. testcode:: python
:skipif: True
def train_func(config):
lr = config["lr"]
num_epochs = config["num_epochs"]
config = {"lr": 1e-4, "num_epochs": 10}
trainer = ray.train.torch.TorchTrainer(train_func, train_loop_config=config, ...)
.. warning::
Avoid passing large data objects through `train_loop_config` to reduce the
serialization and deserialization overhead. Instead, it's preferred to
initialize large objects (e.g. datasets, models) directly in `train_func`.
.. code-block:: diff
def load_dataset():
# Return a large in-memory dataset
...
def load_model():
# Return a large in-memory model instance
...
-config = {"data": load_dataset(), "model": load_model()}
def train_func(config):
- data = config["data"]
- model = config["model"]
+ data = load_dataset()
+ model = load_model()
...
trainer = ray.train.torch.TorchTrainer(train_func, train_loop_config=config, ...)
doc/source/train/deepspeed.rst
+136
View File
@@ -0,0 +1,136 @@
.. _train-deepspeed:
Get Started with DeepSpeed
==========================
The :class:`~ray.train.torch.TorchTrainer` can help you easily launch your `DeepSpeed <https://www.deepspeed.ai/>`_ training across a distributed Ray cluster.
DeepSpeed is an optimization library that enables efficient large-scale model training through techniques like ZeRO (Zero Redundancy Optimizer).
Benefits of Using Ray Train with DeepSpeed
------------------------------------------
- **Simplified Distributed Setup**: Ray Train handles all the distributed environment setup for you
- **Multi-Node Scaling**: Easily scale to multiple nodes with minimal code changes
- **Checkpoint Management**: Built-in checkpoint saving and loading across distributed workers
- **Seamless Integration**: Works with your existing DeepSpeed code
Code example
------------
You can use your existing DeepSpeed training code with Ray Train's TorchTrainer. The integration is minimal and preserves your familiar DeepSpeed workflow:
.. testcode::
:skipif: True
import deepspeed
from deepspeed.accelerator import get_accelerator
def train_func():
# Instantiate your model and dataset
model = ...
train_dataset = ...
eval_dataset = ...
deepspeed_config = {...} # Your DeepSpeed config
# Prepare everything for distributed training
model, optimizer, train_dataloader, lr_scheduler = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
training_data=tokenized_datasets["train"],
collate_fn=collate_fn,
config=deepspeed_config,
)
# Define the GPU device for the current worker
device = get_accelerator().device_name(model.local_rank)
# Start training
for epoch in range(num_epochs):
# Training logic
...
# Report metrics to Ray Train
ray.train.report(metrics={"loss": loss})
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(...),
# If running in a multi-node cluster, this is where you
# should configure the run's persistent storage that is accessible
# across all worker nodes.
# run_config=ray.train.RunConfig(storage_path="s3://..."),
...
)
result = trainer.fit()
Complete Examples
-----------------
Below are complete examples of ZeRO-3 training with DeepSpeed. Each example shows a full implementation of fine-tuning
a Bidirectional Encoder Representations from Transformers (BERT) model on the Microsoft Research Paraphrase Corpus (MRPC) dataset.
Install the requirements:
.. code-block:: bash
pip install deepspeed torch datasets transformers torchmetrics "ray[train]"
.. tab-set::
.. tab-item:: Example with Ray Data
.. dropdown:: Show Code
.. literalinclude:: /../../python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py
:language: python
:start-after: __deepspeed_torch_basic_example_start__
:end-before: __deepspeed_torch_basic_example_end__
.. tab-item:: Example with PyTorch DataLoader
.. dropdown:: Show Code
.. literalinclude:: /../../python/ray/train/examples/deepspeed/deepspeed_torch_trainer_no_raydata.py
:language: python
:start-after: __deepspeed_torch_basic_example_no_raydata_start__
:end-before: __deepspeed_torch_basic_example_no_raydata_end__
.. tip::
To run DeepSpeed with pure PyTorch, you **don't need to** provide any additional Ray Train utilities
like :meth:`~ray.train.torch.prepare_model` or :meth:`~ray.train.torch.prepare_data_loader` in your training function. Instead,
keep using `deepspeed.initialize() <https://deepspeed.readthedocs.io/en/latest/initialize.html>`_ as usual to prepare everything
for distributed training.
Fine-tune LLMs with DeepSpeed
-----------------------------
See this step-by-step guide for how to fine-tune large language models (LLMs) with Ray Train and DeepSpeed: :doc:`Fine-tune an LLM with Ray Train and DeepSpeed </_collections/train/examples/pytorch/deepspeed_finetune/README>`.
Run DeepSpeed with Other Frameworks
-----------------------------------
Many deep learning frameworks have integrated with DeepSpeed, including Lightning, Transformers, Accelerate, and more. You can run all these combinations in Ray Train.
Check the below examples for more details:
.. list-table::
:header-rows: 1
* - Framework
- Example
* - Accelerate (:ref:`User Guide <train-hf-accelerate>`)
- `Fine-tune Llama-2 series models with DeepSpeed, Accelerate, and Ray Train. <https://github.com/ray-project/ray/tree/master/doc/source/templates/04_finetuning_llms_with_deepspeed>`_
* - Transformers (:ref:`User Guide <train-pytorch-transformers>`)
- :doc:`Fine-tune GPT-J-6b with DeepSpeed and Hugging Face Transformers <examples/deepspeed/gptj_deepspeed_fine_tuning>`
* - Lightning (:ref:`User Guide <train-pytorch-lightning>`)
- :doc:`Fine-tune vicuna-13b with DeepSpeed and PyTorch Lightning <examples/lightning/vicuna_13b_lightning_deepspeed_finetune>`
For more information about DeepSpeed configuration options, refer to the `official DeepSpeed documentation <https://www.deepspeed.ai/docs/config-json/>`_.
doc/source/train/deprecated-user-guides/fault-tolerance-deprecated-api.rst
+136
View File
@@ -0,0 +1,136 @@
:orphan:
.. _train-fault-tolerance-deprecated-api:
Handling Failures and Node Preemption (Deprecated API)
======================================================
.. important::
This user guide covers deprecated fault tolerance APIs. See :ref:`train-fault-tolerance` for the new API user guide.
Please see :ref:`here <train-fault-tolerance-deprecation-info>` for information about the deprecation and migration.
Automatically Recover from Train Worker Failures
------------------------------------------------
Ray Train has built-in fault tolerance to recover from worker failures (i.e.
``RayActorError``\s). When a failure is detected, the workers will be shut
down and new workers will be added in.
The training function will be restarted, but progress from the previous execution can
be resumed through checkpointing.
.. tip::
In order to retain progress when recovery, your training function
**must** implement logic for both :ref:`saving <train-dl-saving-checkpoints>`
*and* :ref:`loading checkpoints <train-dl-loading-checkpoints>`.
Each instance of recovery from a worker failure is considered a retry. The
number of retries is configurable through the ``max_failures`` attribute of the
:class:`~ray.train.FailureConfig` argument set in the :class:`~ray.train.RunConfig`
passed to the ``Trainer``:
.. literalinclude:: ../doc_code/fault_tolerance.py
:language: python
:start-after: __failure_config_start__
:end-before: __failure_config_end__
Which checkpoint will be restored?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ray Train will automatically resume training from the latest available
:ref:`checkpoint reported to Ray Train <train-checkpointing>`.
This will be the last checkpoint passed to :func:`train.report() <ray.train.report>`.
Restore a Ray Train Experiment
------------------------------
At the experiment level, Trainer restoration
allows you to resume a previously interrupted experiment from where it left off.
A Train experiment may be interrupted due to one of the following reasons:
- The experiment was manually interrupted (e.g., Ctrl+C, or pre-empted head node instance).
- The head node crashed (e.g., OOM or some other runtime error).
- The entire cluster went down (e.g., network error affecting all nodes).
Trainer restoration is possible for all of Ray Train's built-in trainers,
but we use ``TorchTrainer`` in the examples for demonstration.
We also use ``<Framework>Trainer`` to refer to methods that are shared across all
built-in trainers.
Let's say your initial Train experiment is configured as follows.
The actual training loop is just for demonstration purposes: the important detail is that
:ref:`saving <train-dl-saving-checkpoints>` *and* :ref:`loading checkpoints <train-dl-loading-checkpoints>`
has been implemented.
.. literalinclude:: ../doc_code/dl_guide.py
:language: python
:start-after: __ft_initial_run_start__
:end-before: __ft_initial_run_end__
The results and checkpoints of the experiment are saved to the path configured by :class:`~ray.train.RunConfig`.
If the experiment has been interrupted due to one of the reasons listed above, use this path to resume:
.. literalinclude:: ../doc_code/dl_guide.py
:language: python
:start-after: __ft_restored_run_start__
:end-before: __ft_restored_run_end__
.. tip::
You can also restore from a remote path (e.g., from an experiment directory stored in a s3 bucket).
.. literalinclude:: ../doc_code/dl_guide.py
:language: python
:dedent:
:start-after: __ft_restore_from_cloud_initial_start__
:end-before: __ft_restore_from_cloud_initial_end__
.. literalinclude:: ../doc_code/dl_guide.py
:language: python
:dedent:
:start-after: __ft_restore_from_cloud_restored_start__
:end-before: __ft_restore_from_cloud_restored_end__
.. note::
Different trainers may allow more parameters to be optionally re-specified on restore.
Only **datasets** are required to be re-specified on restore, if they were supplied originally.
`TorchTrainer.restore`, `TensorflowTrainer.restore`, and `HorovodTrainer.restore`
can take in the same parameters as their parent class's
:meth:`DataParallelTrainer.restore <ray.train.data_parallel_trainer.DataParallelTrainer.restore>`.
Unless otherwise specified, other trainers will accept the same parameters as
:meth:`BaseTrainer.restore <ray.train.trainer.BaseTrainer.restore>`.
Auto-resume
~~~~~~~~~~~
Adding the branching logic below will allow you to run the same script after the interrupt,
picking up training from where you left on the previous run. Notice that we use the
:meth:`<Framework>Trainer.can_restore <ray.train.trainer.BaseTrainer.can_restore>` utility method
to determine the existence and validity of the given experiment directory.
.. literalinclude:: ../doc_code/dl_guide.py
:language: python
:start-after: __ft_autoresume_start__
:end-before: __ft_autoresume_end__
.. seealso::
See the :meth:`BaseTrainer.restore <ray.train.trainer.BaseTrainer.restore>` docstring
for a full example.
.. note::
`<Framework>Trainer.restore` is different from
:class:`<Framework>Trainer(..., resume_from_checkpoint=...) <ray.train.trainer.BaseTrainer>`.
`resume_from_checkpoint` is meant to be used to start a *new* Train experiment,
which writes results to a new directory and starts over from iteration 0.
`<Framework>Trainer.restore` is used to continue an existing experiment, where
new results will continue to be appended to existing logs.
doc/source/train/deprecated-user-guides/hyperparameter-optimization-deprecated.rst
+134
View File
@@ -0,0 +1,134 @@
:orphan:
.. _train-tune-deprecated-api:
Hyperparameter Tuning with Ray Tune (Deprecated API)
====================================================
.. important::
This user guide covers the deprecated Train + Tune integration. See :ref:`train-tune` for the new API user guide.
Please see :ref:`here <train-tune-deprecation>` for information about the deprecation and migration.
Hyperparameter tuning with :ref:`Ray Tune <tune-main>` is natively supported with Ray Train.
.. https://docs.google.com/drawings/d/1yMd12iMkyo6DGrFoET1TIlKfFnXX9dfh2u3GSdTz6W4/edit
.. figure:: ../images/train-tuner.svg
:align: center
The `Tuner` will take in a `Trainer` and execute multiple training runs, each with different hyperparameter configurations.
Key Concepts
------------
There are a number of key concepts when doing hyperparameter optimization with a :class:`~ray.tune.Tuner`:
* A set of hyperparameters you want to tune in a *search space*.
* A *search algorithm* to effectively optimize your parameters and optionally use a
*scheduler* to stop searches early and speed up your experiments.
* The *search space*, *search algorithm*, *scheduler*, and *Trainer* are passed to a Tuner,
which runs the hyperparameter tuning workload by evaluating multiple hyperparameters in parallel.
* Each individual hyperparameter evaluation run is called a *trial*.
* The Tuner returns its results as a :class:`~ray.tune.ResultGrid`.
.. note::
Tuners can also be used to launch hyperparameter tuning without using Ray Train. See
:ref:`the Ray Tune documentation <tune-main>` for more guides and examples.
Basic usage
-----------
You can take an existing :class:`Trainer <ray.train.base_trainer.BaseTrainer>` and simply
pass it into a :class:`~ray.tune.Tuner`.
.. literalinclude:: ../doc_code/tuner.py
:language: python
:start-after: __basic_start__
:end-before: __basic_end__
How to configure a Tuner?
-------------------------
There are two main configuration objects that can be passed into a Tuner: the :class:`TuneConfig <ray.tune.TuneConfig>` and the :class:`ray.tune.RunConfig`.
The :class:`TuneConfig <ray.tune.TuneConfig>` contains tuning specific settings, including:
- the tuning algorithm to use
- the metric and mode to rank results
- the amount of parallelism to use
Here are some common configurations for `TuneConfig`:
.. literalinclude:: ../doc_code/tuner.py
:language: python
:start-after: __tune_config_start__
:end-before: __tune_config_end__
See the :class:`TuneConfig API reference <ray.tune.TuneConfig>` for more details.
The :class:`ray.tune.RunConfig` contains configurations that are more generic than tuning specific settings.
This includes:
- failure/retry configurations
- verbosity levels
- the name of the experiment
- the logging directory
- checkpoint configurations
- custom callbacks
- integration with cloud storage
Below we showcase some common configurations of :class:`ray.tune.RunConfig`.
.. literalinclude:: ../doc_code/tuner.py
:language: python
:start-after: __run_config_start__
:end-before: __run_config_end__
Search Space configuration
--------------------------
A `Tuner` takes in a `param_space` argument where you can define the search space
from which hyperparameter configurations will be sampled.
Depending on the model and dataset, you may want to tune:
- The training batch size
- The learning rate for deep learning training (e.g., image classification)
- The maximum depth for tree-based models (e.g., XGBoost)
You can use a Tuner to tune most arguments and configurations for Ray Train, including but
not limited to:
- Ray :class:`Datasets <ray.data.Dataset>`
- :class:`~ray.train.ScalingConfig`
- and other hyperparameters.
Read more about :ref:`Tune search spaces here <tune-search-space-tutorial>`.
Train - Tune gotchas
--------------------
There are a couple gotchas about parameter specification when using Tuners with Trainers:
- By default, configuration dictionaries and config objects will be deep-merged.
- Parameters that are duplicated in the Trainer and Tuner will be overwritten by the Tuner ``param_space``.
- **Exception:** all arguments of the :class:`ray.tune.RunConfig` and :class:`ray.tune.TuneConfig` are inherently un-tunable.
See :doc:`/tune/tutorials/tune_get_data_in_and_out` for an example.
Advanced Tuning
---------------
Tuners also offer the ability to tune over different data preprocessing steps and
different training/validation datasets, as shown in the following snippet.
.. literalinclude:: ../doc_code/tuner.py
:language: python
:start-after: __tune_dataset_start__
:end-before: __tune_dataset_end__
doc/source/train/distributed-tensorflow-keras.rst
+355
View File
@@ -0,0 +1,355 @@
.. _train-tensorflow-overview:
Get Started with Distributed Training using TensorFlow/Keras
============================================================
Ray Train's `TensorFlow <https://www.tensorflow.org/>`__ integration enables you
to scale your TensorFlow and Keras training functions to many machines and GPUs.
On a technical level, Ray Train schedules your training workers
and configures ``TF_CONFIG`` for you, allowing you to run
your ``MultiWorkerMirroredStrategy`` training script. See `Distributed
training with TensorFlow <https://www.tensorflow.org/guide/distributed_training>`_
for more information.
Most of the examples in this guide use TensorFlow with Keras, but
Ray Train also works with vanilla TensorFlow.
Quickstart
-----------
.. literalinclude:: ./doc_code/tf_starter.py
:language: python
:start-after: __tf_train_start__
:end-before: __tf_train_end__
Update your training function
-----------------------------
First, update your :ref:`training function <train-overview-training-function>` to support distributed
training.
.. note::
The current TensorFlow implementation supports
``MultiWorkerMirroredStrategy`` (and ``MirroredStrategy``). If there are
other strategies you wish to see supported by Ray Train, submit a `feature request on GitHub <https://github.com/ray-project/ray/issues>`_.
These instructions closely follow TensorFlow's `Multi-worker training
with Keras <https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras>`_
tutorial. One key difference is that Ray Train handles the environment
variable set up for you.
**Step 1:** Wrap your model in ``MultiWorkerMirroredStrategy``.
The `MultiWorkerMirroredStrategy <https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/MultiWorkerMirroredStrategy>`_
enables synchronous distributed training. You *must* build and compile the ``Model`` within the scope of the strategy.
.. testcode::
:skipif: True
with tf.distribute.MultiWorkerMirroredStrategy().scope():
model = ... # build model
model.compile()
**Step 2:** Update your ``Dataset`` batch size to the *global* batch
size.
Set ``batch_size`` appropriately because `batch <https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch>`_
splits evenly across worker processes.
.. code-block:: diff
-batch_size = worker_batch_size
+batch_size = worker_batch_size * train.get_context().get_world_size()
.. warning::
Ray doesn't automatically set any environment variables or configuration
related to local parallelism or threading
:ref:`aside from "OMP_NUM_THREADS" <omp-num-thread-note>`.
If you want greater control over TensorFlow threading, use
the ``tf.config.threading`` module (eg.
``tf.config.threading.set_inter_op_parallelism_threads(num_cpus)``)
at the beginning of your ``train_loop_per_worker`` function.
Create a TensorflowTrainer
--------------------------
``Trainer``\s are the primary Ray Train classes for managing state and
execute training. For distributed TensorFlow,
use a :class:`~ray.train.tensorflow.TensorflowTrainer`
that you can setup like this:
.. testcode::
:hide:
train_func = lambda: None
.. testcode::
from ray.train import ScalingConfig
from ray.train.tensorflow import TensorflowTrainer
# For GPU Training, set `use_gpu` to True.
use_gpu = False
trainer = TensorflowTrainer(
train_func,
scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=2)
)
To customize the backend setup, you can pass a
:class:`~ray.train.tensorflow.TensorflowConfig`:
.. testcode::
:skipif: True
from ray.train import ScalingConfig
from ray.train.tensorflow import TensorflowTrainer, TensorflowConfig
trainer = TensorflowTrainer(
train_func,
tensorflow_backend=TensorflowConfig(...),
scaling_config=ScalingConfig(num_workers=2),
)
For more configurability, see the :py:class:`~ray.train.data_parallel_trainer.DataParallelTrainer` API.
Run a training function
-----------------------
With a distributed training function and a Ray Train ``Trainer``, you are now
ready to start training.
.. testcode::
:skipif: True
trainer.fit()
Load and preprocess data
------------------------
TensorFlow by default uses its own internal dataset sharding policy, as described
`in the guide <https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#dataset_sharding>`__.
If your TensorFlow dataset is compatible with distributed loading, you don't need to
change anything.
If you require more advanced preprocessing, you may want to consider using Ray Data
for distributed data ingest. See :ref:`Ray Data with Ray Train <data-ingest-torch>`.
The main difference is that you may want to convert your Ray Data dataset shard to
a TensorFlow dataset in your training function so that you can use the Keras
API for model training.
`See this example <https://github.com/ray-project/ray/blob/master/python/ray/train/examples/tf/tensorflow_autoencoder_example.py>`__
for distributed data loading. The relevant parts are:
.. testcode::
import tensorflow as tf
from ray import train
from ray.train.tensorflow import prepare_dataset_shard
def train_func(config: dict):
# ...
# Get dataset shard from Ray Train
dataset_shard = train.get_context().get_dataset_shard("train")
# Define a helper function to build a TensorFlow dataset
def to_tf_dataset(dataset, batch_size):
def to_tensor_iterator():
for batch in dataset.iter_tf_batches(
batch_size=batch_size, dtypes=tf.float32
):
yield batch["image"], batch["label"]
output_signature = (
tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
)
tf_dataset = tf.data.Dataset.from_generator(
to_tensor_iterator, output_signature=output_signature
)
# Call prepare_dataset_shard to disable automatic sharding
# (since the dataset is already sharded)
return prepare_dataset_shard(tf_dataset)
for epoch in range(epochs):
# Call our helper function to build the dataset
tf_dataset = to_tf_dataset(
dataset=dataset_shard,
batch_size=64,
)
history = multi_worker_model.fit(tf_dataset)
Report results
--------------
During training, the training loop should report intermediate results and checkpoints
to Ray Train. This reporting logs the results to the console output and appends them to
local log files. The logging also triggers :ref:`checkpoint bookkeeping <train-dl-configure-checkpoints>`.
The easiest way to report your results with Keras is by using the
:class:`~ray.train.tensorflow.keras.ReportCheckpointCallback`:
.. testcode::
from ray.train.tensorflow.keras import ReportCheckpointCallback
def train_func(config: dict):
# ...
for epoch in range(epochs):
model.fit(dataset, callbacks=[ReportCheckpointCallback()])
This callback automatically forwards all results and checkpoints from the
Keras training function to Ray Train.
Aggregate results
~~~~~~~~~~~~~~~~~
TensorFlow Keras automatically aggregates metrics from all workers. If you wish to have more
control over that, consider implementing a `custom training loop <https://www.tensorflow.org/tutorials/distribute/custom_training>`__.
Save and load checkpoints
-------------------------
You can save :class:`Checkpoints <ray.train.Checkpoint>` by calling ``train.report(metrics, checkpoint=Checkpoint(...))`` in the
training function. This call saves the checkpoint state from the distributed
workers on the ``Trainer``, where you executed your python script.
You can access the latest saved checkpoint through the ``checkpoint`` attribute of
the :py:class:`~ray.train.Result`, and access the best saved checkpoints with the ``best_checkpoints``
attribute.
These concrete examples demonstrate how Ray Train appropriately saves checkpoints, model weights but not models, in distributed training.
.. testcode::
import json
import os
import tempfile
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.tensorflow import TensorflowTrainer
import numpy as np
def train_func(config):
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import tensorflow as tf
n = 100
# create a toy dataset
# data : X - dim = (n, 4)
# target : Y - dim = (n, 1)
X = np.random.normal(0, 1, size=(n, 4))
Y = np.random.uniform(0, 1, size=(n, 1))
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# toy neural network : 1-layer
model = tf.keras.Sequential([tf.keras.layers.Dense(1, activation="linear", input_shape=(4,))])
model.compile(optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
dataset = tf.data.Dataset.from_tensor_slices((X, Y)).batch(20)
for epoch in range(config["num_epochs"]):
history = model.fit(dataset)
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
model.save(os.path.join(temp_checkpoint_dir, "model.keras"))
checkpoint_dict = os.path.join(temp_checkpoint_dir, "checkpoint.json")
with open(checkpoint_dict, "w") as f:
json.dump({"epoch": epoch}, f)
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report({"loss": history.history["loss"][0]}, checkpoint=checkpoint)
trainer = TensorflowTrainer(
train_func,
train_loop_config={"num_epochs": 5},
scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()
print(result.checkpoint)
By default, checkpoints persist to local disk in the :ref:`log
directory <train-log-dir>` of each run.
Load checkpoints
~~~~~~~~~~~~~~~~
.. testcode::
import os
import tempfile
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.tensorflow import TensorflowTrainer
import numpy as np
def train_func(config):
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import tensorflow as tf
n = 100
# create a toy dataset
# data : X - dim = (n, 4)
# target : Y - dim = (n, 1)
X = np.random.normal(0, 1, size=(n, 4))
Y = np.random.uniform(0, 1, size=(n, 1))
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# toy neural network : 1-layer
checkpoint = train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as checkpoint_dir:
model = tf.keras.models.load_model(
os.path.join(checkpoint_dir, "model.keras")
)
else:
model = tf.keras.Sequential(
[tf.keras.layers.Dense(1, activation="linear", input_shape=(4,))]
)
model.compile(optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
dataset = tf.data.Dataset.from_tensor_slices((X, Y)).batch(20)
for epoch in range(config["num_epochs"]):
history = model.fit(dataset)
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
model.save(os.path.join(temp_checkpoint_dir, "model.keras"))
extra_json = os.path.join(temp_checkpoint_dir, "checkpoint.json")
with open(extra_json, "w") as f:
json.dump({"epoch": epoch}, f)
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report({"loss": history.history["loss"][0]}, checkpoint=checkpoint)
trainer = TensorflowTrainer(
train_func,
train_loop_config={"num_epochs": 5},
scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()
print(result.checkpoint)
Further reading
---------------
See :ref:`User Guides <train-user-guides>` to explore more topics:
- :ref:`Experiment tracking <train-experiment-tracking-native>`
- :ref:`Fault tolerance and training on spot instances <train-fault-tolerance>`
- :ref:`Hyperparameter optimization <train-tune>`
doc/source/train/doc_code/asynchronous_validation.py
+311
View File
@@ -0,0 +1,311 @@
# __validation_fn_simple_start__
import os
import torch
import ray.train
import ray.data
# Define Ray Data validation dataset outside validation function because it is not json serializable
validation_dataset = ...
def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
# Load the checkpoint
model = ...
with checkpoint.as_directory() as checkpoint_dir:
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
model.load_state_dict(model_state_dict)
model.eval()
# Perform validation on the data
total_accuracy = 0
with torch.no_grad():
for batch in validation_dataset.iter_torch_batches(batch_size=128):
images, labels = batch["image"], batch["label"]
outputs = model(images)
total_accuracy += (outputs.argmax(1) == labels).sum().item()
return {"score": total_accuracy / len(validation_dataset)}
# __validation_fn_simple_end__
# __validation_fn_torch_trainer_start__
import torchmetrics
from torch.nn import CrossEntropyLoss
import ray.train.torch
from ray.data import ExecutionOptions
def eval_only_train_fn(config_dict: dict) -> dict:
# Load the checkpoint
model = ...
with config_dict["checkpoint"].as_directory() as checkpoint_dir:
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
model.load_state_dict(model_state_dict)
model.cuda().eval()
# Set up metrics and data loaders
criterion = CrossEntropyLoss()
mean_valid_loss = torchmetrics.MeanMetric().cuda()
test_data_shard = ray.train.get_dataset_shard("validation")
test_dataloader = test_data_shard.iter_torch_batches(batch_size=128)
# Compute metric and return it directly from the train function
with torch.no_grad():
for batch in test_dataloader:
images, labels = batch["image"], batch["label"]
outputs = model(images)
loss = criterion(outputs, labels)
mean_valid_loss(loss)
return {"score": mean_valid_loss.compute().item()}
def validation_fn(checkpoint: ray.train.Checkpoint, train_run_name: str, epoch: int) -> dict:
trainer = ray.train.torch.TorchTrainer(
eval_only_train_fn,
train_loop_config={"checkpoint": checkpoint},
scaling_config=ray.train.ScalingConfig(
num_workers=2, use_gpu=True, accelerator_type="A10G"
),
# Give unique name to validation run so it does not attempt to load placeholder checkpoint.
# Also allows you to better associate training runs with validation runs.
run_config=ray.train.RunConfig(
name=f"{train_run_name}_validation_epoch_{epoch}"
),
# Use weaker GPUs for validation
datasets={"validation": validation_dataset},
# Pin to the "validation" subcluster so it doesn't compete with
# training. See https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
dataset_config=ray.train.DataConfig(
execution_options={
"validation": ExecutionOptions(
label_selector={"ray-subcluster": "validation"}
),
},
),
)
result = trainer.fit()
# return_value holds the value returned by train function of worker 0
return result.return_value
# __validation_fn_torch_trainer_end__
# __validation_fn_map_batches_start__
import ray.data
class Predictor:
def __init__(self, checkpoint: ray.train.Checkpoint):
self.model = ...
with checkpoint.as_directory() as checkpoint_dir:
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
self.model.load_state_dict(model_state_dict)
self.model.cuda().eval()
def __call__(self, batch: dict) -> dict:
image = torch.as_tensor(batch["image"], dtype=torch.float32, device="cuda")
label = torch.as_tensor(batch["label"], dtype=torch.float32, device="cuda")
pred = self.model(image)
return {"res": (pred.argmax(1) == label).cpu().numpy()}
# Construct ``validation_dataset`` under a DataContext copy pinned to the
# "validation" subcluster. ``Dataset.context`` is a deep copy of the
# current context taken at construction, so the selector is baked in and
# every downstream operator (including the ``map_batches`` below) inherits
# it — no in-function mutation needed. See
# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
with ray.data.DataContext.current(ctx):
validation_dataset = ray.data.read_parquet(...)
def validation_fn(checkpoint: ray.train.Checkpoint) -> dict:
# Set name to avoid confusion; default name is "Dataset"
validation_dataset.set_name("validation")
eval_res = validation_dataset.map_batches(
Predictor,
batch_size=128,
num_gpus=1,
fn_constructor_kwargs={"checkpoint": checkpoint},
concurrency=2,
)
mean = eval_res.mean(["res"])
return {
"score": mean,
}
# __validation_fn_map_batches_end__
# __validation_fn_report_start__
import tempfile
from ray.data import ExecutionOptions
from ray.train import ValidationConfig, ValidationTaskConfig
def train_func(config: dict) -> None:
...
epochs = ...
model = ...
rank = ray.train.get_context().get_world_rank()
for epoch in epochs:
... # training step
if rank == 0:
training_metrics = {"loss": ..., "epoch": epoch}
local_checkpoint_dir = tempfile.mkdtemp()
torch.save(
model.module.state_dict(),
os.path.join(local_checkpoint_dir, "model.pt"),
)
ray.train.report(
training_metrics,
checkpoint=ray.train.Checkpoint.from_directory(local_checkpoint_dir),
checkpoint_upload_mode=ray.train.CheckpointUploadMode.ASYNC,
validation=ValidationTaskConfig(fn_kwargs={
"train_run_name": ray.train.get_context().get_experiment_name(),
"epoch": epoch,
}),
)
else:
ray.train.report({}, None)
def run_trainer() -> ray.train.Result:
# 1) Construction-time tasks (parquet schema inference, file listing)
# read the current DataContext. Pin them to "training" with a copy of
# the DataContext applied via the DataContext.current() context
# manager — scoped to the `with` block so it doesn't leak. See
# https://docs.ray.io/en/latest/data/concurrent-dataset-execution.html.
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "training"}
with ray.data.DataContext.current(ctx):
train_dataset = ray.data.read_parquet(...)
trainer = ray.train.torch.TorchTrainer(
train_func,
validation_config=ValidationConfig(fn=validation_fn),
# Pass training dataset in datasets arg to split it across training workers
datasets={"train": train_dataset},
# 2) DataConfig.execution_options REPLACES ds.context.execution_options
# wholesale at training start, dropping anything not re-specified
# (including label_selector). Restate the selector here so per-worker
# ingest stays pinned to "training".
dataset_config=ray.train.DataConfig(
datasets_to_split=["train"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "training"}
),
},
),
scaling_config=ray.train.ScalingConfig(
num_workers=2,
use_gpu=True,
# Use powerful GPUs for training
accelerator_type="A100",
),
)
return trainer.fit()
# __validation_fn_report_end__
# __exp_tracking_same_run_wandb_start__
import wandb
import ray.train
from ray.train import ValidationConfig, ValidationTaskConfig
entity = "my_entity"
project = "my_project"
num_epochs = ...
def validation_fn(checkpoint: ray.train.Checkpoint, wandb_run_id: str, val_step: int) -> dict:
wandb.init(
entity=entity,
project=project,
settings=wandb.Settings(mode="shared", x_primary=False),
id=wandb_run_id,
)
score = ...
wandb.log({"validation/loss": score, "val_step": val_step})
wandb.finish() # flush the metrics
return {"validation/loss": score}
def train_func():
if ray.train.get_context().get_world_rank() == 0:
run = wandb.init(
entity=entity,
project=project,
settings=wandb.Settings(mode="shared", x_primary=True,)
)
wandb.define_metric("val_step", hidden=True)
wandb.define_metric("train_step", hidden=True)
wandb.define_metric("validation/loss", step_metric="val_step")
wandb.define_metric("train/loss", step_metric="train_step")
for epoch in range(num_epochs):
loss = ...
if ray.train.get_context().get_world_rank() == 0:
wandb.log({"train/loss": loss, "train_step": epoch})
checkpoint = ...
ray.train.report(
{"train/loss": loss},
checkpoint=checkpoint,
validation=ValidationTaskConfig(
fn_kwargs={"wandb_run_id": run.id, "val_step": epoch}
),
)
else:
ray.train.report({}, None)
if ray.train.get_context().get_world_rank() == 0:
wandb.finish()
# __exp_tracking_same_run_wandb_end__
# __exp_tracking_same_run_mlflow_start__
import mlflow
from mlflow.tracking import MlflowClient
import ray.train
from ray.train import ValidationConfig, ValidationTaskConfig
tracking_uri = "my_uri"
experiment_name = "my_experiment"
num_epochs = ...
def validation_fn(
checkpoint: ray.train.Checkpoint, mlflow_run_id: str, val_step: int
) -> dict:
client = MlflowClient(tracking_uri=tracking_uri)
score = ...
client.log_metric(mlflow_run_id, "val_score", score, step=val_step)
return {"val_score": score}
def train_func():
if ray.train.get_context().get_world_rank() == 0:
client = MlflowClient(tracking_uri=tracking_uri)
experiment = client.get_experiment_by_name(experiment_name)
run = client.create_run(experiment_id=experiment.experiment_id)
for epoch in range(num_epochs):
loss = ...
if ray.train.get_context().get_world_rank() == 0:
client.log_metric(run.info.run_id, "train_loss", loss, step=epoch)
checkpoint = ...
ray.train.report(
{"train_loss": loss},
checkpoint=checkpoint,
validation=ValidationTaskConfig(
fn_kwargs={"mlflow_run_id": run.info.run_id, "val_step": epoch}
),
)
else:
ray.train.report({}, None)
if ray.train.get_context().get_world_rank() == 0:
client.set_terminated(run.info.run_id)
# __exp_tracking_same_run_mlflow_end__
doc/source/train/doc_code/checkpoints.py
+569
View File
@@ -0,0 +1,569 @@
# flake8: noqa
# isort: skip_file
# __pytorch_save_start__
import os
import tempfile
import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam
import ray.train.torch
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer
def train_func(config):
n = 100
# create a toy dataset
# data : X - dim = (n, 4)
# target : Y - dim = (n, 1)
X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
# toy neural network : 1-layer
# Wrap the model in DDP
model = ray.train.torch.prepare_model(nn.Linear(4, 1))
criterion = nn.MSELoss()
optimizer = Adam(model.parameters(), lr=3e-4)
for epoch in range(config["num_epochs"]):
y = model.forward(X)
loss = criterion(y, Y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
metrics = {"loss": loss.item()}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
checkpoint = None
should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
# In standard DDP training, where the model is the same across all ranks,
# only the global rank 0 worker needs to save and report the checkpoint
if train.get_context().get_world_rank() == 0 and should_checkpoint:
torch.save(
model.module.state_dict(), # NOTE: Unwrap the model.
os.path.join(temp_checkpoint_dir, "model.pt"),
)
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report(metrics, checkpoint=checkpoint)
trainer = TorchTrainer(
train_func,
train_loop_config={"num_epochs": 5},
scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()
# __pytorch_save_end__
# __pytorch_restore_start__
import os
import tempfile
import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam
import ray.train.torch
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer
def train_func(config):
n = 100
# create a toy dataset
# data : X - dim = (n, 4)
# target : Y - dim = (n, 1)
X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
# toy neural network : 1-layer
model = nn.Linear(4, 1)
optimizer = Adam(model.parameters(), lr=3e-4)
criterion = nn.MSELoss()
# Wrap the model in DDP and move it to GPU.
model = ray.train.torch.prepare_model(model)
# ====== Resume training state from the checkpoint. ======
start_epoch = 0
checkpoint = train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as checkpoint_dir:
model_state_dict = torch.load(
os.path.join(checkpoint_dir, "model.pt"),
# map_location=..., # Load onto a different device if needed.
)
model.module.load_state_dict(model_state_dict)
optimizer.load_state_dict(
torch.load(os.path.join(checkpoint_dir, "optimizer.pt"))
)
start_epoch = (
torch.load(os.path.join(checkpoint_dir, "extra_state.pt"))["epoch"] + 1
)
# ========================================================
for epoch in range(start_epoch, config["num_epochs"]):
y = model.forward(X)
loss = criterion(y, Y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
metrics = {"loss": loss.item()}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
checkpoint = None
should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
# In standard DDP training, where the model is the same across all ranks,
# only the global rank 0 worker needs to save and report the checkpoint
if train.get_context().get_world_rank() == 0 and should_checkpoint:
# === Make sure to save all state needed for resuming training ===
torch.save(
model.module.state_dict(), # NOTE: Unwrap the model.
os.path.join(temp_checkpoint_dir, "model.pt"),
)
torch.save(
optimizer.state_dict(),
os.path.join(temp_checkpoint_dir, "optimizer.pt"),
)
torch.save(
{"epoch": epoch},
os.path.join(temp_checkpoint_dir, "extra_state.pt"),
)
# ================================================================
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report(metrics, checkpoint=checkpoint)
if epoch == 1:
raise RuntimeError("Intentional error to showcase restoration!")
trainer = TorchTrainer(
train_func,
train_loop_config={"num_epochs": 5},
scaling_config=ScalingConfig(num_workers=2),
run_config=train.RunConfig(failure_config=train.FailureConfig(max_failures=1)),
)
result = trainer.fit()
# __pytorch_restore_end__
# __checkpoint_from_single_worker_start__
import tempfile
from ray import train
def train_fn(config):
...
metrics = {...}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
checkpoint = None
# Only the global rank 0 worker saves and reports the checkpoint
if train.get_context().get_world_rank() == 0:
... # Save checkpoint to temp_checkpoint_dir
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report(metrics, checkpoint=checkpoint)
# __checkpoint_from_single_worker_end__
# __lightning_save_example_start__
import lightning.pytorch as pl
from ray import train
from ray.train.lightning import RayTrainReportCallback
from ray.train.torch import TorchTrainer
class MyLightningModule(pl.LightningModule):
# ...
def on_validation_epoch_end(self):
...
mean_acc = calculate_accuracy()
self.log("mean_accuracy", mean_acc, sync_dist=True)
def train_func():
...
model = MyLightningModule(...)
datamodule = MyLightningDataModule(...)
trainer = pl.Trainer(
# ...
callbacks=[RayTrainReportCallback()]
)
trainer.fit(model, datamodule=datamodule)
ray_trainer = TorchTrainer(
train_func,
scaling_config=train.ScalingConfig(num_workers=2),
run_config=train.RunConfig(
checkpoint_config=train.CheckpointConfig(
num_to_keep=2,
checkpoint_score_attribute="mean_accuracy",
checkpoint_score_order="max",
),
),
)
# __lightning_save_example_end__
# __lightning_custom_save_example_start__
import os
from tempfile import TemporaryDirectory
from lightning.pytorch.callbacks import Callback
import ray
import ray.train
from ray.train import Checkpoint
class CustomRayTrainReportCallback(Callback):
def on_train_epoch_end(self, trainer, pl_module):
should_checkpoint = trainer.current_epoch % 3 == 0
with TemporaryDirectory() as tmpdir:
# Fetch metrics from `self.log(..)` in the LightningModule
metrics = trainer.callback_metrics
metrics = {k: v.item() for k, v in metrics.items()}
# Add customized metrics
metrics["epoch"] = trainer.current_epoch
metrics["custom_metric"] = 123
checkpoint = None
global_rank = ray.train.get_context().get_world_rank() == 0
if global_rank == 0 and should_checkpoint:
# Save model checkpoint file to tmpdir
ckpt_path = os.path.join(tmpdir, "ckpt.pt")
trainer.save_checkpoint(ckpt_path, weights_only=False)
checkpoint = Checkpoint.from_directory(tmpdir)
# Report to train session
ray.train.report(metrics=metrics, checkpoint=checkpoint)
# __lightning_custom_save_example_end__
# __lightning_restore_example_start__
import os
from ray import train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
from ray.train.lightning import RayTrainReportCallback
def train_func():
model = MyLightningModule(...)
datamodule = MyLightningDataModule(...)
trainer = pl.Trainer(..., callbacks=[RayTrainReportCallback()])
checkpoint = train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as ckpt_dir:
ckpt_path = os.path.join(ckpt_dir, RayTrainReportCallback.CHECKPOINT_NAME)
trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path)
else:
trainer.fit(model, datamodule=datamodule)
ray_trainer = TorchTrainer(
train_func,
scaling_config=train.ScalingConfig(num_workers=2),
run_config=train.RunConfig(
checkpoint_config=train.CheckpointConfig(num_to_keep=2),
),
)
# __lightning_restore_example_end__
# __transformers_save_example_start__
from transformers import TrainingArguments
from ray import train
from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
from ray.train.torch import TorchTrainer
def train_func(config):
...
# Configure logging, saving, evaluation strategies as usual.
args = TrainingArguments(
...,
eval_strategy="epoch",
save_strategy="epoch",
logging_strategy="step",
)
trainer = transformers.Trainer(args, ...)
# Add a report callback to transformers Trainer
# =============================================
trainer.add_callback(RayTrainReportCallback())
trainer = prepare_trainer(trainer)
trainer.train()
ray_trainer = TorchTrainer(
train_func,
run_config=train.RunConfig(
checkpoint_config=train.CheckpointConfig(
num_to_keep=3,
checkpoint_score_attribute="eval_loss", # The monitoring metric
checkpoint_score_order="min",
)
),
)
# __transformers_save_example_end__
# __transformers_custom_save_example_start__
from ray import train
from transformers.trainer_callback import TrainerCallback
class MyTrainReportCallback(TrainerCallback):
def __init__(self):
super().__init__()
self.metrics = {}
def on_log(self, args, state, control, model=None, logs=None, **kwargs):
"""Log is called on evaluation step and logging step."""
self.metrics.update(logs)
def on_save(self, args, state, control, **kwargs):
"""Event called after a checkpoint save."""
checkpoint = None
if train.get_context().get_world_rank() == 0:
# Build a Ray Train Checkpoint from the latest checkpoint
checkpoint_path = transformers.trainer.get_last_checkpoint(args.output_dir)
checkpoint = Checkpoint.from_directory(checkpoint_path)
# Report to Ray Train with up-to-date metrics
ray.train.report(metrics=self.metrics, checkpoint=checkpoint)
# Clear the metrics buffer
self.metrics = {}
# __transformers_custom_save_example_end__
# __distributed_checkpointing_start__
from ray import train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
def train_func(config):
...
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
rank = train.get_context().get_world_rank()
torch.save(
...,
os.path.join(temp_checkpoint_dir, f"model-rank={rank}.pt"),
)
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report(metrics, checkpoint=checkpoint)
trainer = TorchTrainer(
train_func,
scaling_config=train.ScalingConfig(num_workers=2),
run_config=train.RunConfig(storage_path="s3://bucket/"),
)
# The checkpoint in cloud storage will contain: model-rank=0.pt, model-rank=1.pt
# __distributed_checkpointing_end__
# __inspect_checkpoint_example_start__
from pathlib import Path
from ray.train import Checkpoint
# For demonstration, create a locally available directory with a `model.pt` file.
example_checkpoint_dir = Path("/tmp/test-checkpoint")
example_checkpoint_dir.mkdir()
example_checkpoint_dir.joinpath("model.pt").touch()
# Create the checkpoint, which is a reference to the directory.
checkpoint = Checkpoint.from_directory(example_checkpoint_dir)
# Inspect the checkpoint's contents with either `as_directory` or `to_directory`:
with checkpoint.as_directory() as checkpoint_dir:
assert Path(checkpoint_dir).joinpath("model.pt").exists()
checkpoint_dir = checkpoint.to_directory()
assert Path(checkpoint_dir).joinpath("model.pt").exists()
# __inspect_checkpoint_example_end__
# __inspect_transformers_checkpoint_example_start__
# After training finished
checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
hf_checkpoint_path = f"{checkpoint_dir}/checkpoint/"
# __inspect_transformers_checkpoint_example_end__
# __inspect_lightning_checkpoint_example_start__
# After training finished
checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
lightning_checkpoint_path = f"{checkpoint_dir}/checkpoint.ckpt"
# __inspect_lightning_checkpoint_example_end__
# __checkpoint_upload_mode_sync_start__
def train_fn(config):
...
metrics = {...}
with tempfile.TemporaryDirectory() as tmpdir:
... # Save checkpoint to tmpdir
checkpoint = Checkpoint.from_directory(tmpdir)
train.report(
metrics,
checkpoint=checkpoint,
checkpoint_upload_mode=train.CheckpointUploadMode.SYNC,
)
# __checkpoint_upload_mode_sync_end__
# __checkpoint_upload_mode_async_start__
def train_fn(config):
...
metrics = {...}
tmpdir = tempfile.mkdtemp()
... # Save checkpoint to tmpdir
checkpoint = Checkpoint.from_directory(tmpdir)
train.report(
metrics,
checkpoint=checkpoint,
checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
)
# __checkpoint_upload_mode_async_end__
# __checkpoint_upload_mode_no_upload_start__
from s3torchconnector.dcp import S3StorageWriter
from torch.distributed.checkpoint.state_dict_saver import save
from torch.distributed.checkpoint.state_dict import get_state_dict
def train_fn(config):
...
for epoch in range(config["num_epochs"]):
# Directly upload checkpoint to s3 with Torch
model, optimizer = ...
storage_context = ray.train.get_context().get_storage()
checkpoint_path = (
f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
)
storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
save(
{"model": model_dict, "opt": opt_dict},
storage_writer=storage_writer,
)
# Report that checkpoint to Ray Train
metrics = {...}
checkpoint = Checkpoint(checkpoint_path)
train.report(
metrics,
checkpoint=checkpoint,
checkpoint_upload_mode=train.CheckpointUploadMode.NO_UPLOAD,
)
# __checkpoint_upload_mode_no_upload_end__
# __checkpoint_upload_fn_start__
from torch.distributed.checkpoint.state_dict_saver import async_save
from s3torchconnector.dcp import S3StorageWriter
from torch.distributed.checkpoint.state_dict import get_state_dict
from ray import train
from ray.train import Checkpoint
def train_fn(config):
...
for epoch in config["num_epochs"]:
# Start async checkpoint upload to s3 with Torch
model, optimizer = ...
storage_context = train.get_context().get_storage()
checkpoint_path = (
f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
)
storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
ckpt_ref = async_save(
{"model": model_dict, "opt": opt_dict},
storage_writer=storage_writer,
)
def wait_async_save(checkpoint, checkpoint_dir_name):
# This function waits for checkpoint to be finalized before returning it as is
ckpt_ref.result()
return checkpoint
# Ray Train kicks off a thread that waits for the async checkpoint upload to complete
# before reporting the checkpoint
metrics = {...}
checkpoint = Checkpoint(checkpoint_path)
train.report(
metrics=metrics,
checkpoint=checkpoint,
checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
checkpoint_upload_fn=wait_async_save,
# As uploading into the experiment directory then don't delete the checkpoint after upload is complete
delete_local_checkpoint_after_upload=False,
)
trainer = TorchTrainer(
train_fn,
train_loop_config={"num_epochs": 3},
scaling_config=train.ScalingConfig(num_workers=2, use_gpu=True),
# we need a cpu backend for async_save and a gpu backend for training
torch_config=train.torch.TorchConfig(backend="cpu:gloo,cuda:nccl"),
run_config=train.RunConfig(storage_path="s3://bucket/")
)
# __checkpoint_upload_fn_end__
# __get_all_reported_checkpoints_example_start__
import ray.train
from ray.train import CheckpointConsistencyMode
def train_fn():
for epoch in range(2):
metrics = {"train/loss": 0.1}
checkpoint = ...
ray.train.report(
metrics,
checkpoint=checkpoint,
validation=...,
)
# Get committed checkpoints which may still have ongoing validations.
committed_checkpoints = ray.train.get_all_reported_checkpoints(
consistency_mode=CheckpointConsistencyMode.COMMITTED)
# Wait for all pending validations to finish to access reported checkpoints
# with validation metrics attached.
validated_checkpoints = ray.train.get_all_reported_checkpoints()
...
# __get_all_reported_checkpoints_example_end__
doc/source/train/doc_code/collate_utils.py
+118
View File
@@ -0,0 +1,118 @@
from dataclasses import dataclass
from typing import Dict, List, Tuple, Union
import torch
from ray import cloudpickle as pickle
import pyarrow as pa
# (dtype, shape, offset)
FEATURE_TYPE = Tuple[torch.dtype, torch.Size, int]
TORCH_BYTE_ELEMENT_TYPE = torch.uint8
def _create_binary_array_from_buffer(buffer: bytes) -> pa.BinaryArray:
"""Zero-copy create a binary array from a buffer."""
data_buffer = pa.py_buffer(buffer)
return pa.Array.from_buffers(
pa.binary(),
1,
[
None,
pa.array([0, data_buffer.size], type=pa.int32()).buffers()[1],
data_buffer,
],
)
@dataclass
class _Metadata:
features: Dict[str, List[FEATURE_TYPE]]
total_buffer_size: int
@dataclass
class _TensorBatch:
"""Internal class for serializing/deserializing tensor batches."""
buffer: torch.Tensor
metadata: _Metadata
@classmethod
def from_batch(cls, batch: Dict[str, Union[List[torch.Tensor], torch.Tensor]]) -> '_TensorBatch':
"""Serialize a batch of tensors into a single buffer."""
features: Dict[str, List[FEATURE_TYPE]] = {}
flattened_binary_tensors = []
total_buffer_size = 0
for name, tensors in batch.items():
features[name] = []
if not isinstance(tensors, list):
tensors = [tensors]
for tensor in tensors:
flattened_tensor = tensor.flatten().contiguous().view(TORCH_BYTE_ELEMENT_TYPE)
flattened_binary_tensors.append(flattened_tensor)
features[name].append((tensor.dtype, tensor.shape, total_buffer_size))
total_buffer_size += flattened_tensor.shape[0]
buffer = torch.empty(total_buffer_size, dtype=TORCH_BYTE_ELEMENT_TYPE)
cur_offset = 0
for flattened_tensor in flattened_binary_tensors:
buffer[cur_offset:cur_offset + flattened_tensor.shape[0]] = flattened_tensor
cur_offset += flattened_tensor.shape[0]
return _TensorBatch(
buffer=buffer,
metadata=_Metadata(
features=features,
total_buffer_size=total_buffer_size,
),
)
def to_table(self) -> pa.Table:
"""Convert to a single-row PyArrow table."""
buffer_array = _create_binary_array_from_buffer(self.buffer.numpy().data)
metadata_array = _create_binary_array_from_buffer(pickle.dumps(self.metadata))
return pa.Table.from_arrays(
arrays=[buffer_array, metadata_array],
names=["_buffer", "_metadata"],
)
@classmethod
def from_table(cls, table: pa.Table) -> '_TensorBatch':
"""Deserialize from a single-row PyArrow table."""
return _TensorBatch(
buffer=torch.frombuffer(
table["_buffer"].chunks[0].buffers()[2],
dtype=TORCH_BYTE_ELEMENT_TYPE
),
metadata=pickle.loads(table["_metadata"].chunks[0].buffers()[2]),
)
def to_batch(self, pin_memory: bool = False) -> Dict[str, List[torch.Tensor]]:
"""Deserialize back to a batch of tensors."""
batch = {}
storage_buffer = self.buffer.untyped_storage()
offsets = []
for name, features in self.metadata.features.items():
for _, _, offset in features:
offsets.append(offset)
offsets.append(self.metadata.total_buffer_size)
offset_id = 0
for name, features in self.metadata.features.items():
batch[name] = []
for dtype, shape, _ in features:
# Create a zero-copy view of the byte slice.
byte_slice = self.buffer[offsets[offset_id]:offsets[offset_id + 1]]
tensor = torch.frombuffer(
byte_slice.numpy().data, dtype=dtype
).view(shape)
if pin_memory:
tensor = tensor.pin_memory()
batch[name].append(tensor)
offset_id += 1
return batch
# Helper functions for use in your code
def serialize_tensors_to_table(batch: Dict[str, Union[List[torch.Tensor], torch.Tensor]]) -> pa.Table:
"""Serialize a batch of tensors to a PyArrow table."""
return _TensorBatch.from_batch(batch).to_table()
def deserialize_table_to_tensors(table: pa.Table, pin_memory: bool = False) -> Dict[str, List[torch.Tensor]]:
"""Deserialize a PyArrow table back to tensors."""
return _TensorBatch.from_table(table).to_batch(pin_memory=pin_memory)
doc/source/train/doc_code/data_ingest_torch_new.py
+147
View File
@@ -0,0 +1,147 @@
# flake8: noqa
# isort: skip_file
# __basic__
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
import numpy as np
from typing import Dict
# Load the data.
train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")
## Uncomment to randomize the block order each epoch.
# train_ds = train_ds.randomize_block_order()
# Define a preprocessing function.
def normalize_length(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
new_col = batch["sepal.length"] / np.max(batch["sepal.length"])
batch["normalized.sepal.length"] = new_col
del batch["sepal.length"]
return batch
# Preprocess your data any way you want. This will be re-run each epoch.
# You can use Ray Data preprocessors here as well,
# e.g., preprocessor.fit_transform(train_ds)
train_ds = train_ds.map_batches(normalize_length)
def train_loop_per_worker():
# Get an iterator to the dataset we passed in below.
it = train.get_dataset_shard("train")
# Train for 10 epochs over the data. We'll use a shuffle buffer size
# of 10k elements, and prefetch up to 10 batches of size 128 each.
for _ in range(10):
for batch in it.iter_batches(
local_shuffle_buffer_size=10000, batch_size=128, prefetch_batches=10
):
print("Do some training on batch", batch)
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": train_ds},
)
my_trainer.fit()
# __basic_end__
# __custom_split__
dataset_a = ray.data.read_text(
"s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)
dataset_b = ray.data.read_csv("s3://anonymous@ray-example-data/dow_jones.csv")
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"a": dataset_a, "b": dataset_b},
dataset_config=ray.train.DataConfig(
datasets_to_split=["a"],
),
)
# __custom_split_end__
def augment_data(batch):
return batch
# __materialized__
# Load the data.
train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")
# Preprocess the data. Transformations that are made to the materialize call below
# will only be run once.
train_ds = train_ds.map_batches(normalize_length)
# Materialize the dataset in object store memory.
train_ds = train_ds.materialize()
# Add per-epoch preprocessing. Transformations that you want to run per-epoch, such
# as data augmentation, should go after the materialize call.
train_ds = train_ds.map_batches(augment_data)
# __materialized_end__
# __options__
from ray.train import DataConfig
options = DataConfig.default_ingest_options()
options.resource_limits = options.resource_limits.copy(object_store_memory=10e9)
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
dataset_config=ray.train.DataConfig(
execution_options=options,
),
)
# __options_end__
# __custom__
# Note that this example class is doing the same thing as the basic DataConfig
# impl included with Ray Train.
from typing import Optional, Dict, List
from ray.data import Dataset, DataIterator, NodeIdStr
from ray.actor import ActorHandle
class MyCustomDataConfig(DataConfig):
def configure(
self,
datasets: Dict[str, Dataset],
world_size: int,
worker_handles: Optional[List[ActorHandle]],
worker_node_ids: Optional[List[NodeIdStr]],
**kwargs,
) -> List[Dict[str, DataIterator]]:
assert len(datasets) == 1, "This example only handles the simple case"
# Configure Ray Data for ingest.
ctx = ray.data.DataContext.get_current()
ctx.execution_options = DataConfig.default_ingest_options()
# Split the stream into shards.
iterator_shards = datasets["train"].streaming_split(
world_size, equal=True, locality_hints=worker_node_ids
)
# Return the assigned iterators for each worker.
return [{"train": it} for it in iterator_shards]
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": train_ds},
dataset_config=MyCustomDataConfig(),
)
my_trainer.fit()
# __custom_end__
doc/source/train/doc_code/dl_guide.py
+110
View File
@@ -0,0 +1,110 @@
# flake8: noqa
# TODO: [V2] Deprecated doc code to delete.
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
MOCK = True
# __ft_initial_run_start__
import os
import tempfile
from typing import Dict, Optional
import torch
import ray
from ray import train
from ray.train import Checkpoint
from ray.train.torch import TorchTrainer
def get_datasets() -> Dict[str, ray.data.Dataset]:
return {"train": ray.data.from_items([{"x": i, "y": 2 * i} for i in range(10)])}
def train_loop_per_worker(config: dict):
from torchvision.models import resnet18
model = resnet18()
# Checkpoint loading
checkpoint: Optional[Checkpoint] = train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as checkpoint_dir:
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
model.load_state_dict(model_state_dict)
model = train.torch.prepare_model(model)
train_ds = train.get_dataset_shard("train")
for epoch in range(5):
# Do some training...
# Checkpoint saving
with tempfile.TemporaryDirectory() as tmpdir:
torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
train.report({"epoch": epoch}, checkpoint=Checkpoint.from_directory(tmpdir))
trainer = TorchTrainer(
train_loop_per_worker=train_loop_per_worker,
datasets=get_datasets(),
scaling_config=train.ScalingConfig(num_workers=2),
run_config=train.RunConfig(
name="dl_trainer_restore", storage_path=os.path.expanduser("~/ray_results")
),
)
result = trainer.fit()
# __ft_initial_run_end__
# __ft_restored_run_start__
from ray.train.torch import TorchTrainer
restored_trainer = TorchTrainer.restore(
path=os.path.expanduser("~/ray_results/dl_trainer_restore"),
datasets=get_datasets(),
)
# __ft_restored_run_end__
if not MOCK:
# __ft_restore_from_cloud_initial_start__
original_trainer = TorchTrainer(
# ...
run_config=train.RunConfig(
# Configure cloud storage
storage_path="s3://results-bucket",
name="dl_trainer_restore",
),
)
result = trainer.fit()
# __ft_restore_from_cloud_initial_end__
# __ft_restore_from_cloud_restored_start__
restored_trainer = TorchTrainer.restore(
"s3://results-bucket/dl_trainer_restore",
datasets=get_datasets(),
)
# __ft_restore_from_cloud_restored_end__
# __ft_autoresume_start__
experiment_path = os.path.expanduser("~/ray_results/dl_restore_autoresume")
if TorchTrainer.can_restore(experiment_path):
trainer = TorchTrainer.restore(experiment_path, datasets=get_datasets())
result = trainer.fit()
else:
trainer = TorchTrainer(
train_loop_per_worker=train_loop_per_worker,
datasets=get_datasets(),
scaling_config=train.ScalingConfig(num_workers=2),
run_config=train.RunConfig(
storage_path=os.path.expanduser("~/ray_results"),
name="dl_restore_autoresume",
),
)
result = trainer.fit()
# __ft_autoresume_end__
doc/source/train/doc_code/fault_tolerance.py
+107
View File
@@ -0,0 +1,107 @@
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
# __failure_config_start__
import ray.train
# Tries to recover a run up to this many times.
failure_config = ray.train.FailureConfig(max_failures=2)
# No limit on the number of retries.
failure_config = ray.train.FailureConfig(max_failures=-1)
# __failure_config_end__
# __worker_fault_tolerance_start__
import tempfile
import uuid
import ray.train
import ray.train.torch
def train_fn_per_worker(train_loop_config: dict):
# [1] Train worker restoration logic.
checkpoint = ray.train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as temp_checkpoint_dir:
# model.load_state_dict(torch.load(...))
...
# [2] Checkpoint saving and reporting logic.
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
# torch.save(...)
ray.train.report(
{"loss": 0.1},
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
scaling_config=ray.train.ScalingConfig(num_workers=4),
run_config=ray.train.RunConfig(
# (If multi-node, configure S3 / NFS as the storage path.)
# storage_path="s3://...",
name=f"train_run-{uuid.uuid4().hex}",
# [3] Enable worker-level fault tolerance to gracefully handle
# Train worker failures.
failure_config=ray.train.FailureConfig(max_failures=3),
),
)
trainer.fit()
# __worker_fault_tolerance_end__
# Avoid running the code below so that the argument parser is not used.
__name__ = "__dummy__"
# __job_driver_fault_tolerance_start__
# entrypoint.py
import argparse
import tempfile
import uuid
import ray.train
import ray.train.torch
def train_fn_per_worker(train_loop_config: dict):
# [1] Train worker restoration logic.
checkpoint = ray.train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as temp_checkpoint_dir:
# model.load_state_dict(torch.load(...))
...
# [2] Checkpoint saving and reporting logic.
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
# torch.save(...)
ray.train.report(
{"loss": 0.1},
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--storage_path", type=str, required=True)
parser.add_argument("--run_name", type=str, required=True)
args = parser.parse_args()
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
scaling_config=ray.train.ScalingConfig(num_workers=4),
run_config=ray.train.RunConfig(
# [3] Enable worker-level fault tolerance to gracefully handle
# Train worker failures.
failure_config=ray.train.FailureConfig(max_failures=3),
# [4] (Recommendation) The (storage_path, name) pair should be
# determined by the job submitter and passed in as arguments
# to the entrypoint script.
storage_path=args.storage_path,
name=args.run_name,
),
)
trainer.fit()
# __job_driver_fault_tolerance_end__
doc/source/train/doc_code/hvd_trainer.py
+80
View File
@@ -0,0 +1,80 @@
# TODO: [V2] Deprecated doc code to delete.
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
import tempfile
import horovod.torch as hvd
import ray
from ray import train
from ray.train import Checkpoint, ScalingConfig
import ray.train.torch # Need this to use `train.torch.get_device()`
from ray.train.horovod import HorovodTrainer
import torch
import torch.nn as nn
# If using GPUs, set this to True.
use_gpu = False
input_size = 1
layer_size = 15
output_size = 1
num_epochs = 3
class NeuralNetwork(nn.Module):
def __init__(self):
super(NeuralNetwork, self).__init__()
self.layer1 = nn.Linear(input_size, layer_size)
self.relu = nn.ReLU()
self.layer2 = nn.Linear(layer_size, output_size)
def forward(self, input):
return self.layer2(self.relu(self.layer1(input)))
def train_loop_per_worker():
hvd.init()
dataset_shard = train.get_dataset_shard("train")
model = NeuralNetwork()
device = train.torch.get_device()
model.to(device)
loss_fn = nn.MSELoss()
lr_scaler = 1
optimizer = torch.optim.SGD(model.parameters(), lr=0.1 * lr_scaler)
# Horovod: wrap optimizer with DistributedOptimizer.
optimizer = hvd.DistributedOptimizer(
optimizer,
named_parameters=model.named_parameters(),
op=hvd.Average,
)
for epoch in range(num_epochs):
model.train()
for batch in dataset_shard.iter_torch_batches(
batch_size=32, dtypes=torch.float
):
inputs, labels = torch.unsqueeze(batch["x"], 1), batch["y"]
outputs = model(inputs)
loss = loss_fn(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"epoch: {epoch}, loss: {loss.item()}")
with tempfile.TemporaryDirectory() as tmpdir:
torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
train.report(
{"loss": loss.item()}, checkpoint=Checkpoint.from_directory(tmpdir)
)
train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
scaling_config = ScalingConfig(num_workers=3, use_gpu=use_gpu)
trainer = HorovodTrainer(
train_loop_per_worker=train_loop_per_worker,
scaling_config=scaling_config,
datasets={"train": train_dataset},
)
result = trainer.fit()
doc/source/train/doc_code/key_concepts.py
+136
View File
@@ -0,0 +1,136 @@
# flake8: noqa
# isort: skip_file
from pathlib import Path
import tempfile
import ray.train
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
def train_fn(config):
for i in range(3):
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
Path(temp_checkpoint_dir).joinpath("model.pt").touch()
ray.train.report(
{"loss": i},
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
return {"total loss": 3}
trainer = DataParallelTrainer(
train_fn, scaling_config=ray.train.ScalingConfig(num_workers=2)
)
# __run_config_start__
import os
from ray.train import RunConfig
run_config = RunConfig(
# Name of the training run (directory name).
name="my_train_run",
# The experiment results will be saved to: storage_path/name
storage_path=os.path.expanduser("~/ray_results"),
# storage_path="s3://my_bucket/tune_results",
)
# __run_config_end__
# __checkpoint_config_start__
from ray.train import RunConfig, CheckpointConfig
# Example 1: Only keep the 2 *most recent* checkpoints and delete the others.
run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=2))
# Example 2: Only keep the 2 *best* checkpoints and delete the others.
run_config = RunConfig(
checkpoint_config=CheckpointConfig(
num_to_keep=2,
# *Best* checkpoints are determined by these params:
checkpoint_score_attribute="mean_accuracy",
checkpoint_score_order="max",
),
# This will store checkpoints on S3.
storage_path="s3://remote-bucket/location",
)
# __checkpoint_config_end__
# __result_metrics_start__
result = trainer.fit()
print("Observed metrics:", result.metrics)
# __result_metrics_end__
# __result_dataframe_start__
df = result.metrics_dataframe
print("Minimum loss", min(df["loss"]))
# __result_dataframe_end__
# __result_return_value_start__
print("Returned data", result.return_value)
# __result_return_value_end__
# __result_checkpoint_start__
print("Last checkpoint:", result.checkpoint)
with result.checkpoint.as_directory() as tmpdir:
# Load model from directory
...
# __result_checkpoint_end__
# __result_best_checkpoint_start__
# Print available checkpoints
for checkpoint, metrics in result.best_checkpoints:
print("Loss", metrics["loss"], "checkpoint", checkpoint)
# Get checkpoint with minimal loss
best_checkpoint = min(
result.best_checkpoints, key=lambda checkpoint: checkpoint[1]["loss"]
)[0]
with best_checkpoint.as_directory() as tmpdir:
# Load model from directory
...
# __result_best_checkpoint_end__
import pyarrow
# __result_path_start__
result_path: str = result.path
result_filesystem: pyarrow.fs.FileSystem = result.filesystem
print(f"Results location (fs, path) = ({result_filesystem}, {result_path})")
# __result_path_end__
# __result_restore_start__
from ray.train import Result
restored_result = Result.from_path(result_path)
print("Restored loss", restored_result.metrics["loss"])
# __result_restore_end__
def error_train_fn(config):
raise RuntimeError("Simulated training error")
trainer = DataParallelTrainer(
error_train_fn, scaling_config=ray.train.ScalingConfig(num_workers=1)
)
# __result_error_start__
try:
result = trainer.fit()
except ray.train.TrainingFailedError as e:
if isinstance(e, ray.train.WorkerGroupError):
print(e.worker_failures)
# __result_error_end__
doc/source/train/doc_code/lightgbm_quickstart.py
+128
View File
@@ -0,0 +1,128 @@
# flake8: noqa
# isort: skip_file
# __lightgbm_start__
import pandas as pd
import lightgbm as lgb
# 1. Load your data as a `lightgbm.Dataset`.
train_df = pd.read_csv("s3://ray-example-data/iris/train/1.csv")
eval_df = pd.read_csv("s3://ray-example-data/iris/val/1.csv")
train_X = train_df.drop("target", axis=1)
train_y = train_df["target"]
eval_X = eval_df.drop("target", axis=1)
eval_y = eval_df["target"]
train_set = lgb.Dataset(train_X, label=train_y)
eval_set = lgb.Dataset(eval_X, label=eval_y)
# 2. Define your LightGBM model training parameters.
params = {
"objective": "multiclass",
"num_class": 3,
"metric": ["multi_logloss", "multi_error"],
"verbosity": -1,
"boosting_type": "gbdt",
"num_leaves": 31,
"learning_rate": 0.05,
"feature_fraction": 0.9,
"bagging_fraction": 0.8,
"bagging_freq": 5,
}
# 3. Do non-distributed training.
model = lgb.train(
params,
train_set,
valid_sets=[eval_set],
valid_names=["eval"],
num_boost_round=100,
)
# __lightgbm_end__
# __lightgbm_ray_start__
import lightgbm as lgb
import ray.train
from ray.train.lightgbm import LightGBMTrainer, RayTrainReportCallback
# 1. Load your data as a Ray Data Dataset.
train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/train")
eval_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/val")
def train_func():
# 2. Load your data shard as a `lightgbm.Dataset`.
# Get dataset shards for this worker
train_shard = ray.train.get_dataset_shard("train")
eval_shard = ray.train.get_dataset_shard("eval")
# Convert shards to PyArrow tables. LightGBM (>=4.2.0) supports PyArrow
# natively, which avoids a round-trip through pandas.
import pyarrow as pa
train_table = pa.concat_tables(
train_shard.iter_batches(batch_format="pyarrow", batch_size=None)
)
eval_table = pa.concat_tables(
eval_shard.iter_batches(batch_format="pyarrow", batch_size=None)
)
train_X = train_table.drop(["target"])
train_y = train_table.column("target")
eval_X = eval_table.drop(["target"])
eval_y = eval_table.column("target")
train_set = lgb.Dataset(train_X, label=train_y)
eval_set = lgb.Dataset(eval_X, label=eval_y)
# 3. Define your LightGBM model training parameters.
params = {
"objective": "multiclass",
"num_class": 3,
"metric": ["multi_logloss", "multi_error"],
"verbosity": -1,
"boosting_type": "gbdt",
"num_leaves": 31,
"learning_rate": 0.05,
"feature_fraction": 0.9,
"bagging_fraction": 0.8,
"bagging_freq": 5,
# Adding the lines below are the only changes needed
# for your `lgb.train` call!
"tree_learner": "data_parallel",
"pre_partition": True,
**ray.train.lightgbm.get_network_params(),
}
# 4. Do distributed data-parallel training.
# Ray Train sets up the necessary coordinator processes and
# environment variables for your workers to communicate with each other.
model = lgb.train(
params,
train_set,
valid_sets=[eval_set],
valid_names=["eval"],
num_boost_round=100,
# Optional: Use the `RayTrainReportCallback` to save and report checkpoints.
callbacks=[RayTrainReportCallback()],
)
# 5. Configure scaling and resource requirements.
scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 2})
# 6. Launch distributed training job.
trainer = LightGBMTrainer(
train_func,
scaling_config=scaling_config,
datasets={"train": train_dataset, "eval": eval_dataset},
)
result = trainer.fit()
# 7. Load the trained model.
model = RayTrainReportCallback.get_model(result.checkpoint)
# __lightgbm_ray_end__
doc/source/train/doc_code/metric_logging.py
+138
View File
@@ -0,0 +1,138 @@
# flake8: noqa
# isort: skip_file
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
# __torchmetrics_start__
# First, pip install torchmetrics
# This code is tested with torchmetrics==0.7.3 and torch==1.12.1
import os
import tempfile
import ray.train.torch
from ray import train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
import torch
import torch.nn as nn
import torchmetrics
from torch.optim import Adam
import numpy as np
def train_func(config):
n = 100
# create a toy dataset
X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
X_valid = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
Y_valid = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
# toy neural network : 1-layer
# wrap the model in DDP
model = ray.train.torch.prepare_model(nn.Linear(4, 1))
criterion = nn.MSELoss()
mape = torchmetrics.MeanAbsolutePercentageError()
# for averaging loss
mean_valid_loss = torchmetrics.MeanMetric()
optimizer = Adam(model.parameters(), lr=3e-4)
for epoch in range(config["num_epochs"]):
model.train()
y = model.forward(X)
# compute loss
loss = criterion(y, Y)
# back-propagate loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
# evaluate
model.eval()
with torch.no_grad():
pred = model(X_valid)
valid_loss = criterion(pred, Y_valid)
# save loss in aggregator
mean_valid_loss(valid_loss)
mape(pred, Y_valid)
# collect all metrics
# use .item() to obtain a value that can be reported
valid_loss = valid_loss.item()
mape_collected = mape.compute().item()
mean_valid_loss_collected = mean_valid_loss.compute().item()
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
torch.save(
model.state_dict(), os.path.join(temp_checkpoint_dir, "model.pt")
)
train.report(
{
"mape_collected": mape_collected,
"valid_loss": valid_loss,
"mean_valid_loss_collected": mean_valid_loss_collected,
},
checkpoint=train.Checkpoint.from_directory(temp_checkpoint_dir),
)
# reset for next epoch
mape.reset()
mean_valid_loss.reset()
trainer = TorchTrainer(
train_func,
train_loop_config={"num_epochs": 5},
scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()
print(result.metrics["valid_loss"], result.metrics["mean_valid_loss_collected"])
# 0.5109779238700867 0.5512474775314331
# __torchmetrics_end__
# __report_callback_start__
import os
assert os.environ["RAY_TRAIN_V2_ENABLED"] == "1"
from typing import Any, Dict, List, Optional
import ray.train
import ray.train.torch
def train_fn_per_worker(config):
# Free-floating metrics can be accessed from the callback below.
ray.train.report({"rank": ray.train.get_context().get_world_rank()})
class CustomMetricsCallback(ray.train.UserCallback):
def after_report(
self,
run_context,
metrics: List[Dict[str, Any]],
checkpoint: Optional[ray.train.Checkpoint],
):
rank_0_metrics = metrics[0]
print(rank_0_metrics)
# Ex: Write metrics to a file...
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
scaling_config=ray.train.ScalingConfig(num_workers=2),
run_config=ray.train.RunConfig(callbacks=[CustomMetricsCallback()]),
)
trainer.fit()
# __report_callback_end__
doc/source/train/doc_code/random_text_generator.py
+58
View File
@@ -0,0 +1,58 @@
import random
import string
import ray
def random_text(length: int) -> str:
"""Generate random text of specified length."""
if length <= 0:
return ""
if length <= 3:
return "".join(random.choices(string.ascii_lowercase, k=length))
words = []
current_length = 0
while current_length < length:
remaining = length - current_length
if remaining <= 4:
word_length = remaining
word = "".join(random.choices(string.ascii_lowercase, k=word_length))
words.append(word)
break
else:
max_word_length = min(10, remaining - 1)
if max_word_length >= 3:
word_length = random.randint(3, max_word_length)
else:
word_length = remaining
word = "".join(random.choices(string.ascii_lowercase, k=word_length))
words.append(word)
current_length += len(word) + 1
text = " ".join(words)
return text[:length]
def random_label() -> int:
"""Pick a random label."""
labels = [0, 1, 2, 3, 4, 5, 6, 7]
return random.choice(labels)
def create_mock_ray_text_dataset(dataset_size: int = 96, min_len: int = 5, max_len: int = 100):
"""Create a mock Ray dataset with random text and labels."""
numbers = random.choices(range(min_len, max_len + 1), k=dataset_size)
ray_dataset = ray.data.from_items(numbers)
def map_to_text_and_label(item):
length = item['item']
text = random_text(length)
label = random_label()
return {
"length": length,
"text": text,
"label": label
}
text_dataset = ray_dataset.map(map_to_text_and_label)
return text_dataset
doc/source/train/doc_code/tf_starter.py
+84
View File
@@ -0,0 +1,84 @@
# flake8: noqa
# isort: skip_file
# __tf_train_start__
import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import ray
import tensorflow as tf
from ray import train
from ray.train import ScalingConfig
from ray.train.tensorflow import TensorflowTrainer
from ray.train.tensorflow.keras import ReportCheckpointCallback
# If using GPUs, set this to True.
use_gpu = False
a = 5
b = 10
size = 100
def build_model() -> tf.keras.Model:
model = tf.keras.Sequential(
[
tf.keras.layers.InputLayer(input_shape=()),
# Add feature dimension, expanding (batch_size,) to (batch_size, 1).
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(10),
tf.keras.layers.Dense(1),
]
)
return model
def train_func(config: dict):
import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import tensorflow as tf
batch_size = config.get("batch_size", 64)
epochs = config.get("epochs", 3)
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
# Model building/compiling need to be within `strategy.scope()`.
multi_worker_model = build_model()
multi_worker_model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
loss="mean_squared_error",
metrics=["mean_squared_error"],
)
dataset = train.get_dataset_shard("train")
results = []
for _ in range(epochs):
tf_dataset = dataset.to_tf(
feature_columns="x", label_columns="y", batch_size=batch_size
)
history = multi_worker_model.fit(
tf_dataset, callbacks=[ReportCheckpointCallback()]
)
results.append(history.history)
return results
config = {"lr": 1e-3, "batch_size": 32, "epochs": 4}
train_dataset = ray.data.from_items(
[{"x": x / 200, "y": 2 * x / 200} for x in range(200)]
)
scaling_config = ScalingConfig(num_workers=2, use_gpu=use_gpu)
trainer = TensorflowTrainer(
train_loop_per_worker=train_func,
train_loop_config=config,
scaling_config=scaling_config,
datasets={"train": train_dataset},
)
result = trainer.fit()
print(result.metrics)
# __tf_train_end__
doc/source/train/doc_code/train_tune_interop.py
+204
View File
@@ -0,0 +1,204 @@
# flake8: noqa
# isort: skip_file
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "1"
# __quickstart_start__
import random
import tempfile
import uuid
import ray.train
import ray.train.torch
import ray.tune
from ray.tune.integration.ray_train import TuneReportCallback
# [1] Define your Ray Train worker code.
def train_fn_per_worker(train_loop_config: dict):
# Unpack train worker hyperparameters.
# Train feeds in the `train_loop_config` defined below.
lr = train_loop_config["lr"]
# training code here...
print(
ray.train.get_context().get_world_size(),
ray.train.get_context().get_world_rank(),
train_loop_config,
)
# model = ray.train.torch.prepare_model(...) # Wrap model in DDP.
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
ray.train.report(
{"loss": random.random()},
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
# [2] Define a function that launches the Ray Train run.
def train_driver_fn(config: dict):
# Unpack run-level hyperparameters.
# Tune feeds in hyperparameters defined in the `param_space` below.
num_workers = config["num_workers"]
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
train_loop_config=config["train_loop_config"],
scaling_config=ray.train.ScalingConfig(
num_workers=num_workers,
# Uncomment to use GPUs.
# use_gpu=True,
),
run_config=ray.train.RunConfig(
# [3] Assign unique names to each run.
# Recommendation: use the trial id as part of the run name.
name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
# [4] (Optional) Pass in a `TuneReportCallback` to propagate
# reported results to the Tuner.
callbacks=[TuneReportCallback()],
# (If multi-node, configure S3 / NFS as the storage path.)
# storage_path="s3://...",
),
)
trainer.fit()
# Launch a single Train run.
# Note that you can only create a TuneReportCallback in a Ray Tune session.
# train_driver_fn({"num_workers": 4, "train_loop_config": {"lr": 1e-3}})
# Launch a sweep of hyperparameters with Ray Tune.
tuner = ray.tune.Tuner(
train_driver_fn,
param_space={
"num_workers": ray.tune.choice([2, 4]),
"train_loop_config": {
"lr": ray.tune.grid_search([1e-3, 3e-4]),
"batch_size": ray.tune.grid_search([32, 64]),
},
},
run_config=ray.tune.RunConfig(
name=f"tune_train_example-{uuid.uuid4().hex[:6]}",
# (If multi-node, configure S3 / NFS as the storage path.)
# storage_path="s3://...",
),
# [5] (Optional) Set the maximum number of concurrent trials
# in order to prevent too many Train driver processes from
# being launched at once.
tune_config=ray.tune.TuneConfig(max_concurrent_trials=2),
)
results = tuner.fit()
print(results.get_best_result(metric="loss", mode="min"))
# __quickstart_end__
# __max_concurrent_trials_start__
# For a fixed size cluster, calculate this based on the limiting resource (ex: GPUs).
total_cluster_gpus = 8
num_gpu_workers_per_trial = 4
max_concurrent_trials = total_cluster_gpus // num_gpu_workers_per_trial
def train_driver_fn(config: dict):
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
scaling_config=ray.train.ScalingConfig(
num_workers=num_gpu_workers_per_trial, use_gpu=True
),
)
trainer.fit()
tuner = ray.tune.Tuner(
train_driver_fn,
tune_config=ray.tune.TuneConfig(max_concurrent_trials=max_concurrent_trials),
)
# __max_concurrent_trials_end__
# __trainable_resources_start__
# Cluster setup:
# head_node:
# resources:
# CPU: 16.0
# worker_node_cpu:
# resources:
# CPU: 32.0
# TRAIN_DRIVER_RESOURCE: 1.0
# worker_node_gpu:
# resources:
# GPU: 4.0
import ray.tune
def train_driver_fn(config):
# trainer = TorchTrainer(...)
...
tuner = ray.tune.Tuner(
ray.tune.with_resources(
train_driver_fn,
# Note: 0.01 is an arbitrary value to schedule the actor
# onto the `worker_node_cpu` node type.
{"TRAIN_DRIVER_RESOURCE": 0.01},
),
)
# __trainable_resources_end__
# __fault_tolerance_start__
import tempfile
import ray.tune
import ray.train
import ray.train.torch
def train_fn_per_worker(train_loop_config: dict):
# [1] Train worker restoration logic.
checkpoint = ray.train.get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as temp_checkpoint_dir:
# model.load_state_dict(torch.load(...))
...
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
# torch.save(...)
ray.train.report(
{"loss": 0.1},
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
def train_fn_driver(config: dict):
trainer = ray.train.torch.TorchTrainer(
train_fn_per_worker,
run_config=ray.train.RunConfig(
# [2] Train driver restoration is automatic, as long as
# the (storage_path, name) remains the same across trial restarts.
# The easiest way to do this is to attach the trial ID in the name.
# **Do not include any timestamps or random values in the name.**
name=f"train-trial_id={ray.tune.get_context().get_trial_id()}",
# [3] Enable worker-level fault tolerance to gracefully handle
# Train worker failures.
failure_config=ray.train.FailureConfig(max_failures=3),
# (If multi-node, configure S3 / NFS as the storage path.)
# storage_path="s3://...",
),
)
trainer.fit()
tuner = ray.tune.Tuner(
train_fn_driver,
run_config=ray.tune.RunConfig(
# [4] Enable trial-level fault tolerance to gracefully handle
# Train driver process failures.
failure_config=ray.tune.FailureConfig(max_failures=3)
),
)
tuner.fit()
# __fault_tolerance_end__
doc/source/train/doc_code/tuner.py
+245
View File
@@ -0,0 +1,245 @@
# flake8: noqa
# isort: skip_file
# TODO: [V2] Deprecated doc code to delete.
import os
os.environ["RAY_TRAIN_V2_ENABLED"] = "0"
# __basic_start__
import ray
import ray.tune
import ray.train
from ray.tune import Tuner
from ray.train.xgboost import XGBoostTrainer
dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
trainer = XGBoostTrainer(
label_column="target",
params={
"objective": "binary:logistic",
"eval_metric": ["logloss", "error"],
"max_depth": 4,
},
datasets={"train": dataset},
scaling_config=ray.train.ScalingConfig(num_workers=2),
)
# Create Tuner
tuner = Tuner(
trainer,
# Add some parameters to tune
param_space={"params": {"max_depth": ray.tune.choice([4, 5, 6])}},
# Specify tuning behavior
tune_config=ray.tune.TuneConfig(metric="train-logloss", mode="min", num_samples=2),
)
# Run tuning job
tuner.fit()
# __basic_end__
# __xgboost_start__
import ray.data
import ray.train
import ray.tune
from ray.tune import Tuner
from ray.train.xgboost import XGBoostTrainer
dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
# Create an XGBoost trainer
trainer = XGBoostTrainer(
label_column="target",
params={
"objective": "binary:logistic",
"eval_metric": ["logloss", "error"],
"max_depth": 4,
},
num_boost_round=10,
datasets={"train": dataset},
)
param_space = {
# Tune parameters directly passed into the XGBoostTrainer
"num_boost_round": ray.tune.randint(5, 20),
# `params` will be merged with the `params` defined in the above XGBoostTrainer
"params": {
"min_child_weight": ray.tune.uniform(0.8, 1.0),
# Below will overwrite the XGBoostTrainer setting
"max_depth": ray.tune.randint(1, 5),
},
# Tune the number of distributed workers
"scaling_config": ray.train.ScalingConfig(num_workers=ray.tune.grid_search([1, 2])),
}
tuner = Tuner(
trainable=trainer,
run_config=ray.tune.RunConfig(name="test_tuner_xgboost"),
param_space=param_space,
tune_config=ray.tune.TuneConfig(
mode="min", metric="train-logloss", num_samples=2, max_concurrent_trials=2
),
)
result_grid = tuner.fit()
# __xgboost_end__
# __torch_start__
import os
import ray.train
import ray.tune
from ray.tune import Tuner
from ray.train.examples.pytorch.torch_linear_example import (
train_func as linear_train_func,
)
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(
train_loop_per_worker=linear_train_func,
train_loop_config={"lr": 1e-2, "batch_size": 4, "epochs": 10},
scaling_config=ray.train.ScalingConfig(num_workers=1, use_gpu=False),
)
param_space = {
# The params will be merged with the ones defined in the TorchTrainer
"train_loop_config": {
# This is a parameter that hasn't been set in the TorchTrainer
"hidden_size": ray.tune.randint(1, 4),
# This will overwrite whatever was set when TorchTrainer was instantiated
"batch_size": ray.tune.choice([4, 8]),
},
# Tune the number of distributed workers
"scaling_config": ray.train.ScalingConfig(num_workers=ray.tune.grid_search([1, 2])),
}
tuner = Tuner(
trainable=trainer,
run_config=ray.tune.RunConfig(
name="test_tuner", storage_path=os.path.expanduser("~/ray_results")
),
param_space=param_space,
tune_config=ray.tune.TuneConfig(
mode="min", metric="loss", num_samples=2, max_concurrent_trials=2
),
)
result_grid = tuner.fit()
# __torch_end__
# __tune_dataset_start__
import ray.data
import ray.tune
from ray.data.preprocessors import StandardScaler
def get_dataset():
ds1 = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
prep_v1 = StandardScaler(["worst radius", "worst area"])
ds1 = prep_v1.fit_transform(ds1)
return ds1
def get_another_dataset():
ds2 = ray.data.read_csv(
"s3://anonymous@air-example-data/breast_cancer_with_categorical.csv"
)
prep_v2 = StandardScaler(["worst concavity", "worst smoothness"])
ds2 = prep_v2.fit_transform(ds2)
return ds2
dataset_1 = get_dataset()
dataset_2 = get_another_dataset()
tuner = ray.tune.Tuner(
trainer,
param_space={
"datasets": {
"train": ray.tune.grid_search([dataset_1, dataset_2]),
}
# Your other parameters go here
},
)
# __tune_dataset_end__
# __tune_optimization_start__
from ray.tune.search.bayesopt import BayesOptSearch
from ray.tune.schedulers import HyperBandScheduler
from ray.tune import TuneConfig
config = TuneConfig(
# ...
search_alg=BayesOptSearch(),
scheduler=HyperBandScheduler(),
)
# __tune_optimization_end__
# __result_grid_inspection_start__
from ray.tune import Tuner, TuneConfig
tuner = Tuner(
trainable=trainer,
param_space=param_space,
tune_config=TuneConfig(mode="min", metric="loss", num_samples=5),
)
result_grid = tuner.fit()
num_results = len(result_grid)
# Check if there have been errors
if result_grid.errors:
print("At least one trial failed.")
# Get the best result
best_result = result_grid.get_best_result()
# And the best checkpoint
best_checkpoint = best_result.checkpoint
# And the best metrics
best_metric = best_result.metrics
# Or a dataframe for further analysis
results_df = result_grid.get_dataframe()
print("Shortest training time:", results_df["time_total_s"].min())
# Iterate over results
for result in result_grid:
if result.error:
print("The trial had an error:", result.error)
continue
print("The trial finished successfully with the metrics:", result.metrics["loss"])
# __result_grid_inspection_end__
# __run_config_start__
import ray.tune
run_config = ray.tune.RunConfig(
name="MyExperiment",
storage_path="s3://...",
checkpoint_config=ray.tune.CheckpointConfig(checkpoint_frequency=2),
)
# __run_config_end__
# __tune_config_start__
from ray.tune import TuneConfig
from ray.tune.search.bayesopt import BayesOptSearch
tune_config = TuneConfig(
metric="loss",
mode="min",
max_concurrent_trials=10,
num_samples=100,
search_alg=BayesOptSearch(),
)
# __tune_config_end__
# __tune_restore_start__
tuner = Tuner.restore(
path=os.path.expanduser("~/ray_results/test_tuner"),
trainable=trainer,
restart_errored=True,
)
tuner.fit()
# __tune_restore_end__
doc/source/train/doc_code/xgboost_quickstart.py
+113
View File
@@ -0,0 +1,113 @@
# flake8: noqa
# isort: skip_file
# __xgboost_start__
import pandas as pd
import xgboost
# 1. Load your data as an `xgboost.DMatrix`.
train_df = pd.read_csv("s3://ray-example-data/iris/train/1.csv")
eval_df = pd.read_csv("s3://ray-example-data/iris/val/1.csv")
train_X = train_df.drop("target", axis=1)
train_y = train_df["target"]
eval_X = eval_df.drop("target", axis=1)
eval_y = eval_df["target"]
dtrain = xgboost.DMatrix(train_X, label=train_y)
deval = xgboost.DMatrix(eval_X, label=eval_y)
# 2. Define your xgboost model training parameters.
params = {
"tree_method": "approx",
"objective": "reg:squarederror",
"eta": 1e-4,
"subsample": 0.5,
"max_depth": 2,
}
# 3. Do non-distributed training.
bst = xgboost.train(
params,
dtrain=dtrain,
evals=[(deval, "validation")],
num_boost_round=10,
)
# __xgboost_end__
# __xgboost_ray_start__
import xgboost
import ray.train
from ray.train.xgboost import XGBoostTrainer, RayTrainReportCallback
# 1. Load your data as a Ray Data Dataset.
train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/train")
eval_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris/val")
def train_func():
# 2. Load your data shard as an `xgboost.DMatrix`.
# Get dataset shards for this worker
train_shard = ray.train.get_dataset_shard("train")
eval_shard = ray.train.get_dataset_shard("eval")
# Convert shards to pandas DataFrames
train_df = train_shard.materialize().to_pandas()
eval_df = eval_shard.materialize().to_pandas()
train_X = train_df.drop("target", axis=1)
train_y = train_df["target"]
eval_X = eval_df.drop("target", axis=1)
eval_y = eval_df["target"]
dtrain = xgboost.DMatrix(train_X, label=train_y)
deval = xgboost.DMatrix(eval_X, label=eval_y)
# 3. Define your xgboost model training parameters.
params = {
"tree_method": "approx",
"objective": "reg:squarederror",
"eta": 1e-4,
"subsample": 0.5,
"max_depth": 2,
}
# 4. Do distributed data-parallel training.
# Ray Train sets up the necessary coordinator processes and
# environment variables for your workers to communicate with each other.
bst = xgboost.train(
params,
dtrain=dtrain,
evals=[(deval, "validation")],
num_boost_round=10,
# Optional: Use the `RayTrainReportCallback` to save and report checkpoints.
callbacks=[RayTrainReportCallback()],
)
# 5. Configure scaling and resource requirements.
scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 2})
# 6. Launch distributed training job.
trainer = XGBoostTrainer(
train_func,
scaling_config=scaling_config,
datasets={"train": train_dataset, "eval": eval_dataset},
# If running in a multi-node cluster, this is where you
# should configure the run's persistent storage that is accessible
# across all worker nodes.
# run_config=ray.train.RunConfig(storage_path="s3://..."),
)
result = trainer.fit()
# 7. Load the trained model
import os
with result.checkpoint.as_directory() as checkpoint_dir:
model_path = os.path.join(checkpoint_dir, RayTrainReportCallback.CHECKPOINT_NAME)
model = xgboost.Booster()
model.load_model(model_path)
# __xgboost_ray_end__
doc/source/train/examples.yml
+236
View File
@@ -0,0 +1,236 @@
# This file is used to auto-generate the Examples Gallery page.
# Do not edit the generated examples.rst page directly.
# To request formatting changes to the generated page, file an issue with the Ray docs team.
# To reference the generated page, use examples.html.
# When adding a new example, include the skill level and framework, if applicable.
text: Below are examples for using Ray Train with a variety of frameworks and use cases. Ray Train makes it easy to scale out each of these examples to a large cluster of GPUs.
columns_to_show:
- frameworks
groupby: skill_level
examples:
- title: Distributing your PyTorch Training Code with Ray Train and Ray Data
skill_level: beginner
frameworks:
- pytorch
use_cases:
- computer vision
link: ../_collections/train/examples/pytorch/distributing-pytorch/README
- title: Train an image classifier with Lightning
skill_level: beginner
frameworks:
- lightning
use_cases:
- computer vision
link: examples/lightning/lightning_mnist_example
- title: Train a text classifier with Hugging Face Accelerate
frameworks:
- accelerate
- pytorch
- hugging face
skill_level: beginner
use_cases:
- large language models
- natural language processing
link: examples/accelerate/accelerate_example
- title: Train an image classifier with TensorFlow
frameworks:
- tensorflow
skill_level: beginner
use_cases:
- computer vision
link: examples/tf/tensorflow_mnist_example
- title: Train a GPT-2-style Transformer with JAX and Flax
frameworks:
- jax
- flax
skill_level: beginner
use_cases:
- natural language processing
link: examples/jax/intro_to_jax_trainer/README
- title: Train with Horovod and PyTorch
frameworks:
- horovod
skill_level: beginner
link: examples/horovod/horovod_example
- title: "Train ResNet model with Intel Gaudi"
frameworks:
- pytorch
skill_level: beginner
use_cases:
- computer vision
contributor: community
link: examples/intel_gaudi/resnet
- title: "Train BERT model with Intel Gaudi"
frameworks:
- transformers
skill_level: beginner
use_cases:
- natural language processing
contributor: community
link: examples/intel_gaudi/bert
- title: Profiling a Ray Train Workload with PyTorch Profiler
frameworks:
- pytorch
skill_level: beginner
use_cases:
- computer vision
link: ../_collections/train/examples/pytorch/pytorch-profiling/README
- title: Get started with PyTorch Fully Sharded Data Parallel (FSDP2) and Ray Train
skill_level: intermediate
frameworks:
- pytorch
use_cases:
- computer vision
link: ../_collections/train/examples/pytorch/pytorch-fsdp/README
- title: Get started with Tensor Parallelism (DeepSpeed AutoTP) and Ray Train
skill_level: intermediate
frameworks:
- pytorch
- deepspeed
use_cases:
- large language models
- natural language processing
link: ../_collections/train/examples/pytorch/tensor_parallel_autotp/README
- title: Get started with 2D Parallelism (Tensor + Data Parallelism) using FSDP2 and Ray Train
skill_level: intermediate
frameworks:
- pytorch
use_cases:
- large language models
- natural language processing
link: ../_collections/train/examples/pytorch/tensor_parallel_dtensor/README
- title: Fine-tune an LLM with Ray Train and DeepSpeed
skill_level: intermediate
frameworks:
- pytorch
- deepspeed
use_cases:
- large language models
- natural language processing
link: ../_collections/train/examples/pytorch/deepspeed_finetune/README
- title: Train a text classifier with DeepSpeed
frameworks:
- deepspeed
- pytorch
skill_level: intermediate
use_cases:
- large language models
- natural language processing
link: examples/deepspeed/deepspeed_example
- title: Fine-tune a personalized Stable Diffusion model
skill_level: intermediate
frameworks:
- pytorch
use_cases:
- computer vision
- generative ai
link: examples/pytorch/dreambooth_finetuning
- title: Finetune Stable Diffusion and generate images with Intel Gaudi
skill_level: intermediate
frameworks:
- accelerate
- transformers
use_cases:
- computer vision
- generative ai
contributor: community
link: examples/intel_gaudi/sd
- title: Train a text classifier with PyTorch Lightning and Ray Data
frameworks:
- lightning
skill_level: intermediate
use_cases:
- natural language processing
link: examples/lightning/lightning_cola_advanced
- title: Train a text classifier with Hugging Face Transformers
frameworks:
- transformers
skill_level: intermediate
use_cases:
- natural language processing
link: examples/transformers/huggingface_text_classification
- title: RL Post-Train an LLM using HuggingFace TRL with GRPO
frameworks:
- transformers
- trl
skill_level: intermediate
use_cases:
- natural language processing
- reinforcement learning
link: examples/transformers/transformer_reinforcement_learning/README
- title: "Fine-tune Llama-2-7b and Llama-2-70b with Intel Gaudi"
frameworks:
- accelerate
- transformers
skill_level: intermediate
use_cases:
- natural language processing
- large language models
contributor: community
link: examples/intel_gaudi/llama
- title: "Pre-train Llama-2 with Intel Gaudi"
frameworks:
- accelerate
- transformers
- deepspeed
skill_level: intermediate
use_cases:
- natural language processing
- large language models
contributor: community
link: examples/intel_gaudi/llama_pretrain
- title: Fine-tune Llama3.1 with AWS Trainium
frameworks:
- pytorch
- aws neuron
skill_level: advanced
use_cases:
- natural language processing
- large language models
contributor: community
link: examples/aws-trainium/llama3
- title: Fine-tune a Llama-2 text generation model with DeepSpeed and Hugging Face Accelerate
frameworks:
- accelerate
- deepspeed
- hugging face
skill_level: advanced
use_cases:
- natural language processing
- large language models
link: https://github.com/ray-project/ray/tree/master/doc/source/templates/04_finetuning_llms_with_deepspeed
- title: Fine-tune a GPT-J-6B text generation model with DeepSpeed and Hugging Face Transformers
frameworks:
- hugging face
- deepspeed
skill_level: advanced
use_cases:
- natural language processing
- large language models
- generative ai
link: examples/deepspeed/gptj_deepspeed_fine_tuning
- title: Fine-tune a vicuna-13b text generation model with PyTorch Lightning and DeepSpeed
frameworks:
- lightning
- deepspeed
skill_level: advanced
use_cases:
- large language models
- generative ai
link: examples/lightning/vicuna_13b_lightning_deepspeed_finetune
- title: Fine-tune a dolly-v2-7b text generation model with PyTorch Lightning and FSDP
frameworks:
- lightning
skill_level: advanced
use_cases:
- large language models
- generative ai
- natural language processing
link: examples/lightning/dolly_lightning_fsdp_finetuning
- title: Train a tabular model with XGBoost
frameworks:
- xgboost
skill_level: beginner
link: ../_collections/ray-overview/examples/e2e-xgboost/README
doc/source/train/examples/BUILD.bazel
+13
View File
@@ -0,0 +1,13 @@
filegroup(
name = "train_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]) + [
"//doc/source/train/examples/pytorch:pytorch_examples_ci_configs",
"//doc/source/train/examples/transformers:transformers_examples_ci_configs",
"//doc/source/train/examples/lightning:lightning_examples_ci_configs",
],
visibility = ["//doc:__pkg__"],
)
doc/source/train/examples/accelerate/accelerate_example.rst
+29
View File
@@ -0,0 +1,29 @@
:orphan:
Distributed Training with Hugging Face Accelerate
=================================================
.. raw:: html
<a id="try-anyscale-quickstart-accelerate_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=accelerate_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example does distributed data parallel training
with Hugging Face Accelerate, Ray Train, and Ray Data.
It fine-tunes a BERT model and is adapted from
https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/accelerate/accelerate_torch_trainer.py
See also
--------
* :ref:`Get Started with Hugging Face Accelerate <train-hf-accelerate>` for a tutorial on using Ray Train and HF Accelerate
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/aws-trainium/llama3.rst
+109
View File
@@ -0,0 +1,109 @@
:orphan:
Distributed fine-tuning of Llama 3.1 8B on AWS Trainium with Ray and PyTorch Lightning
======================================================================================
.. raw:: html
<a id="try-anyscale-quickstart-aws-trainium-llama3" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=aws-trainium-llama3">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example demonstrates how to fine-tune the `Llama 3.1 8B <https://huggingface.co/NousResearch/Meta-Llama-3.1-8B/>`__ model on `AWS
Trainium <https://aws.amazon.com/ai/machine-learning/trainium/>`__ instances using Ray Train, PyTorch Lightning, and AWS Neuron SDK.
AWS Trainium is the machine learning (ML) chip that AWS built for deep
learning (DL) training of 100B+ parameter models. `AWS Neuron
SDK <https://aws.amazon.com/machine-learning/neuron/>`__ helps
developers train models on Trainium accelerators.
Prepare the environment
-----------------------
See `Setup EKS cluster and tools <https://github.com/aws-neuron/aws-neuron-eks-samples/tree/master/llama3.1_8B_finetune_ray_ptl_neuron#setupeksclusterandtools>`__ for setting up an Amazon EKS cluster leveraging AWS Trainium instances.
Create a Docker image
---------------------
When the EKS cluster is ready, create an Amazon ECR repository for building and uploading the Docker image containing artifacts for fine-tuning a Llama3.1 8B model:
1. Clone the repo.
::
git clone https://github.com/aws-neuron/aws-neuron-eks-samples.git
2. Go to the ``llama3.1_8B_finetune_ray_ptl_neuron`` directory.
::
cd aws-neuron-eks-samples/llama3.1_8B_finetune_ray_ptl_neuron
3. Trigger the script.
::
chmod +x 0-kuberay-trn1-llama3-finetune-build-image.sh
./0-kuberay-trn1-llama3-finetune-build-image.sh
4. Enter the zone your cluster is running in, for example: us-east-2.
5. Verify in the AWS console that the Amazon ECR service has the newly
created ``kuberay_trn1_llama3.1_pytorch2`` repository.
6. Update the ECR image ARN in the manifest file used for creating the Ray cluster.
Replace the <AWS_ACCOUNT_ID> and <REGION> placeholders with actual values in the ``1-llama3-finetune-trn1-create-raycluster.yaml`` file using commands below to reflect the ECR image ARN created above:
::
export AWS_ACCOUNT_ID=<enter_your_aws_account_id> # for ex: 111222333444
export REGION=<enter_your_aws_region> # for ex: us-east-2
sed -i "s/<AWS_ACCOUNT_ID>/$AWS_ACCOUNT_ID/g" 1-llama3-finetune-trn1-create-raycluster.yaml
sed -i "s/<REGION>/$REGION/g" 1-llama3-finetune-trn1-create-raycluster.yaml
Configuring Ray Cluster
-----------------------
The ``llama3.1_8B_finetune_ray_ptl_neuron`` directory in the AWS Neuron samples repository simplifies the
Ray configuration. KubeRay provides a manifest that you can apply
to the cluster to set up the head and worker pods.
Run the following command to set up the Ray cluster:
::
kubectl apply -f 1-llama3-finetune-trn1-create-raycluster.yaml
Accessing Ray Dashboard
-----------------------
Port forward from the cluster to see the state of the Ray dashboard and
then view it on `http://localhost:8265 <http://localhost:8265/>`__.
Run it in the background with the following command:
::
kubectl port-forward service/kuberay-trn1-head-svc 8265:8265 &
Launching Ray Jobs
------------------
The Ray cluster is now ready to handle workloads. Initiate the data preparation and fine-tuning Ray jobs:
1. Launch the Ray job for downloading the dolly-15k dataset and the Llama3.1 8B model artifacts:
::
kubectl apply -f 2-llama3-finetune-trn1-rayjob-create-data.yaml
2. When the job has executed successfully, run the following fine-tuning job:
::
kubectl apply -f 3-llama3-finetune-trn1-rayjob-submit-finetuning-job.yaml
3. Monitor the jobs via the Ray Dashboard
For detailed information on each of the steps above, see the `AWS documentation link <https://github.com/aws-neuron/aws-neuron-eks-samples/blob/master/llama3.1_8B_finetune_ray_ptl_neuron/README.md/>`__.
doc/source/train/examples/deepspeed/deepspeed_example.rst
+29
View File
@@ -0,0 +1,29 @@
:orphan:
Train with DeepSpeed ZeRO-3 and Ray Train
=========================================
.. raw:: html
<a id="try-anyscale-quickstart-deepspeed_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=deepspeed_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This is an intermediate example that shows how to do distributed training with DeepSpeed ZeRO-3 and Ray Train.
It demonstrates how to use :ref:`Ray Data <data>` with DeepSpeed ZeRO-3 and Ray Train.
If you just want to quickly convert your existing TorchTrainer scripts into Ray Train, you can refer to the :ref:`Train with DeepSpeed <train-deepspeed>`.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases.
* :ref:`Get Started with DeepSpeed <train-deepspeed>` for a tutorial.
doc/source/train/examples/deepspeed/gptj_deepspeed_fine_tuning.ipynb
+761
View File
@@ -0,0 +1,761 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# GPT-J-6B Fine-Tuning with Ray Train and DeepSpeed\n",
"\n",
"<a id=\"try-anyscale-quickstart-gptj_deepspeed_fine_tuning\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=gptj_deepspeed_fine_tuning\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example showcases how to use Ray Train for **GPT-J fine-tuning**. GPT-J is a GPT-2-like causal language model trained on the Pile dataset. This particular model has 6 billion parameters. For more information, see [GPT-J](https://huggingface.co/docs/transformers/model_doc/gptj).\n",
"\n",
"This example uses the Ray Train \ud83e\udd17 Transformers integration and a pre-trained model from the Hugging Face Hub. Note that this example is adaptable to other similar models.\n",
"\n",
"This is an advanced example that focuses on the performance and distributed computing aspects of Ray Train. For a beginner-friendly introduction to the Ray Train \ud83e\udd17 Transformers integration, see {ref}`Basic Example for HuggingFace Transformers <transformers_torch_trainer_basic_example>`.\n",
"\n",
"Read [Ray Train Key Concepts](train-key-concepts) and [Ray Data Integration User Guides](data-ingest-torch) before starting this example.\n",
"\n",
"```{note}\n",
"To run this example, make sure your Ray cluster has access to at least one GPU with 16 or more GBs of memory. The required amount of memory depends on the model. This notebook is tested with 16 g4dn.4xlarge instances (including the head node).\n",
"```\n",
"\n",
"This notebook has the following steps:\n",
"1. [Set up Ray](#gptj-setup)\n",
"2. [Load the dataset](#gptj-load)\n",
"3. [Preprocess the dataset with Ray Data](#gptj-preprocess)\n",
"4. [Run the training with Ray Train](#gptj-train)\n",
"5. [Generate text from prompt](#gptj-predict)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Uncomment and run the following line in order to install all the necessary dependencies (this notebook was tested with `accelerate==0.33.0`, `transformers==4.36.2`, `deepspeed==0.17.2`):"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"! pip install -q \"datasets\" \"evaluate\" \"accelerate==0.33.0\" \"transformers==4.36.2\" \"deepspeed==0.17.2\""
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import numpy as np\n",
"import pandas as pd\n",
"import os"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-setup)=\n",
"## Set up Ray\n",
"\n",
"First, let's set some global variables. We will use 16 workers, each being assigned 1 GPU and 8 CPUs."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"model_name = \"EleutherAI/gpt-j-6B\"\n",
"use_gpu = True\n",
"num_workers = 16\n",
"cpus_per_worker = 8"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use `ray.init()` to initialize a local cluster. By default, this cluster will be comprised of only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.\n",
"\n",
"We define a {ref}`runtime environment <runtime-environments>` to ensure that the Ray workers have access to all the necessary packages. You can omit the `runtime_env` argument if you have all of the packages already installed on each node in your cluster."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"\n",
"ray.init(\n",
" runtime_env={\n",
" \"pip\": [\n",
" \"datasets\",\n",
" \"evaluate\",\n",
" \"accelerate==0.33.0\",\n",
" \"transformers==4.36.2\",\n",
" \"deepspeed==0.17.2\",\n",
" ],\n",
" },\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-cell"
]
},
"outputs": [],
"source": [
"# THIS SHOULD BE HIDDEN IN DOCS AND ONLY RAN IN CI\n",
"# Download the model from our S3 mirror as it's faster\n",
"\n",
"import ray\n",
"import subprocess\n",
"import ray.util.scheduling_strategies\n",
"\n",
"\n",
"def force_on_node(node_id: str, remote_func_or_actor_class):\n",
" scheduling_strategy = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n",
" node_id=node_id, soft=False\n",
" )\n",
" options = {\"scheduling_strategy\": scheduling_strategy}\n",
" return remote_func_or_actor_class.options(**options)\n",
"\n",
"\n",
"def run_on_every_node(remote_func_or_actor_class, **remote_kwargs):\n",
" refs = []\n",
" for node in ray.nodes():\n",
" if node[\"Alive\"] and node[\"Resources\"].get(\"GPU\", None):\n",
" refs.append(\n",
" force_on_node(node[\"NodeID\"], remote_func_or_actor_class).remote(\n",
" **remote_kwargs\n",
" )\n",
" )\n",
" return ray.get(refs)\n",
"\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def download_model():\n",
" from transformers.utils.hub import TRANSFORMERS_CACHE\n",
"\n",
" path = os.path.expanduser(\n",
" os.path.join(TRANSFORMERS_CACHE, \"models--EleutherAI--gpt-j-6B\")\n",
" )\n",
" subprocess.run([\"mkdir\", \"-p\", os.path.join(path, \"snapshots\", \"main\")])\n",
" subprocess.run([\"mkdir\", \"-p\", os.path.join(path, \"refs\")])\n",
" if os.path.exists(os.path.join(path, \"refs\", \"main\")):\n",
" return\n",
" subprocess.run(\n",
" [\n",
" \"aws\",\n",
" \"s3\",\n",
" \"sync\",\n",
" \"--no-sign-request\",\n",
" \"s3://large-dl-models-mirror/models--EleutherAI--gpt-j-6B/main/\",\n",
" os.path.join(path, \"snapshots\", \"main\"),\n",
" ]\n",
" )\n",
" with open(os.path.join(path, \"snapshots\", \"main\", \"hash\"), \"r\") as f:\n",
" f_hash = f.read().strip()\n",
" with open(os.path.join(path, \"refs\", \"main\"), \"w\") as f:\n",
" f.write(f_hash)\n",
" os.rename(\n",
" os.path.join(path, \"snapshots\", \"main\"), os.path.join(path, \"snapshots\", f_hash)\n",
" )\n",
"\n",
"\n",
"_ = run_on_every_node(download_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-load)=\n",
"## Loading the dataset\n",
"\n",
"We will be fine-tuning the model on the [`tiny_shakespeare` dataset](https://huggingface.co/datasets/tiny_shakespeare), comprised of 40,000 lines of Shakespeare from a variety of Shakespeare's plays. The aim will be to make the GPT-J model better at generating text in the style of Shakespeare."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"\n",
"print(\"Loading tiny_shakespeare dataset\")\n",
"current_dataset = load_dataset(\"tiny_shakespeare\", trust_remote_code=True)\n",
"current_dataset"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use [Ray Data](https://docs.ray.io/en/latest/data/data.html) for distributed preprocessing and data ingestion. We can easily convert the dataset obtained from Hugging Face Hub to Ray Data by using `ray.data.from_huggingface`."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'train': MaterializedDataset(num_blocks=1, num_rows=1, schema={text: string}),\n",
" 'validation': MaterializedDataset(num_blocks=1, num_rows=1, schema={text: string})}"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import ray.data\n",
"\n",
"ray_datasets = {\n",
" \"train\": ray.data.from_huggingface(current_dataset[\"train\"]),\n",
" \"validation\": ray.data.from_huggingface(current_dataset[\"validation\"]),\n",
"}\n",
"\n",
"ray_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-preprocess)=\n",
"Note that the dataset is represented by a single line of large string, and needs some preprocessing. To do this, use the {meth}`~ray.data.Dataset.map_batches` API to apply transformation functions to batches of data.\n",
"\n",
"The `split_text` function takes the single string and splits it into separate lines, removing empty lines and character names ending with ':' (eg. 'ROMEO:'). The `tokenize` function takes the lines and tokenizes them using the \ud83e\udd17 Tokenizer associated with the model, ensuring each entry has the same length (`block_size`) by padding and truncating. This preprocessing is necessary for training.\n",
"\n",
"```{note}\n",
"This preprocessing can be done in other ways. A common pattern is to tokenize first, and then split the obtained tokens into equally-sized blocks.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"block_size = 512"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'train': MapBatches(tokenize)\n",
" +- MapBatches(split_text)\n",
" +- Dataset(num_blocks=1, num_rows=1, schema={text: string}),\n",
" 'validation': MapBatches(tokenize)\n",
" +- MapBatches(split_text)\n",
" +- Dataset(num_blocks=1, num_rows=1, schema={text: string})}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from transformers import AutoTokenizer\n",
"\n",
"\n",
"def split_text(batch: pd.DataFrame) -> pd.DataFrame:\n",
" text = list(batch[\"text\"])\n",
" flat_text = \"\".join(text)\n",
" split_text = [\n",
" x.strip()\n",
" for x in flat_text.split(\"\\n\")\n",
" if x.strip() and not x.strip()[-1] == \":\"\n",
" ]\n",
" return pd.DataFrame(split_text, columns=[\"text\"])\n",
"\n",
"\n",
"def tokenize(batch: pd.DataFrame) -> dict:\n",
" tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
" ret = tokenizer(\n",
" list(batch[\"text\"]),\n",
" truncation=True,\n",
" max_length=block_size,\n",
" padding=\"max_length\",\n",
" return_tensors=\"np\",\n",
" )\n",
" ret[\"labels\"] = ret[\"input_ids\"].copy()\n",
" return dict(ret)\n",
"\n",
"\n",
"processed_datasets = {\n",
" key: (\n",
" ds.map_batches(split_text, batch_format=\"pandas\")\n",
" .map_batches(tokenize, batch_format=\"pandas\")\n",
" )\n",
" for key, ds in ray_datasets.items()\n",
"}\n",
"processed_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-train)=\n",
"### Fine-tuning the model with Ray Train\n",
"\n",
"Configure Ray Train's {class}`~ray.train.torch.TorchTrainer` to perform distributed fine-tuning of the model. Specify a `train_loop_per_worker` function, which defines the training logic to be distributed by Ray using Distributed Data Parallelism, which uses the PyTorch Distributed backend internally. Each worker has its own copy of the model, but operates on different data. At the end of each step, all the workers sync gradients.\n",
"\n",
"Because GPT-J is a relatively large model, it may not be possible to fit it on smaller GPU types (<=16 GB GRAM). To deal with that issue, this example uses [DeepSpeed](https://github.com/microsoft/DeepSpeed), a library to optimize the training process and to offload and partition optimizer and parameter states, reducing GRAM usage. Furthermore, DeepSpeed ZeRO Stage 3 can load large models without running out of memory.\n",
"\n",
"\ud83e\udd17 Transformers and Ray Train's {ref}`integrations <train-transformers-integration>` allow you to easily configure and use DDP and DeepSpeed. All you need to do is specify the DeepSpeed configuration in the [`TrainingArguments`](https://huggingface.co/docs/transformers/en/main_classes/trainer#transformers.TrainingArguments) object.\n",
"\n",
"```{tip}\n",
"There are many DeepSpeed settings that allow you to trade-off speed for memory usage. The settings used below are tailored to the cluster setup used (16 g4dn.4xlarge nodes) and per device batch size of 16. Some things to keep in mind:\n",
"- If your GPUs support bfloat16, use that instead of float16 mixed precision to get better performance and prevent overflows. Replace `fp16=True` with `bf16=True` in `TrainingArguments`.\n",
"- If you are running out of GRAM: try reducing batch size (defined in the cell below the next one), set `\"overlap_comm\": False` in DeepSpeed config.\n",
"- If you are running out of RAM, add more nodes to your cluster, use nodes with more RAM, set `\"pin_memory\": False` in the DeepSpeed config, reduce the batch size, and remove `\"offload_param\"` from the DeepSpeed config.\n",
"\n",
"For more information on DeepSpeed configuration, refer to [Hugging Face documentation](https://huggingface.co/docs/transformers/main_classes/deepspeed) and [DeepSpeed documentation](https://www.deepspeed.ai/docs/config-json/).\n",
"\n",
"Additionally, if you prefer a lower-level API, the logic below can be expressed as an [Accelerate training loop](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py) distributed by a Ray Train {class}`~ray.train.torch.torch_trainer.TorchTrainer`.\n",
"```\n",
"\n",
"#### Training speed\n",
"\n",
"As this example uses data parallelism, each worker operates on its own shard of the data. The batch size set in `train_ds.iter_torch_batches` is the **per device batch size** (per worker batch size). By changing the number of workers, you can change the **effective batch size** and thus the time needed for training to complete. Calculate the effective batch size as `per device batch size * number of workers * number of gradient accumulation steps`. As you add more workers, the effective batch size rises and thus less time is needed to complete a full epoch. While the speedup is not exactly linear due to extra communication overheads, in many cases it can be close to linear.\n",
"\n",
"The preprocessed dataset has 1348 examples. We have set per device batch size to 16.\n",
"\n",
"* With 16 g4dn.4xlarge nodes, the effective batch size was 256, which equals to 85 steps per epoch. One epoch took **~2440 seconds** (including initialization time).\n",
"\n",
"* With 32 g4dn.4xlarge nodes, the effective batch size was 512, which equals to 43 steps per epoch. One epoch took **~1280 seconds** (including initialization time)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import evaluate\n",
"import torch\n",
"from transformers import (\n",
" Trainer,\n",
" TrainingArguments,\n",
" GPTJForCausalLM,\n",
" AutoTokenizer,\n",
" default_data_collator,\n",
")\n",
"from transformers.utils.logging import disable_progress_bar, enable_progress_bar\n",
"\n",
"from ray import train\n",
"from ray.train.huggingface.transformers import prepare_trainer, RayTrainReportCallback\n",
"\n",
"\n",
"def train_func(config):\n",
" # Use the actual number of CPUs assigned to this worker by Ray\n",
" runtime_ctx = ray.get_runtime_context()\n",
" assigned_cpus = runtime_ctx.get_assigned_resources().get(\"CPU\", 1)\n",
" os.environ[\"OMP_NUM_THREADS\"] = str(int(assigned_cpus))\n",
" # Enable tf32 for better performance\n",
" torch.backends.cuda.matmul.allow_tf32 = True\n",
"\n",
" batch_size = config.get(\"batch_size\", 4)\n",
" epochs = config.get(\"epochs\", 2)\n",
" warmup_steps = config.get(\"warmup_steps\", 0)\n",
" learning_rate = config.get(\"learning_rate\", 0.00002)\n",
" weight_decay = config.get(\"weight_decay\", 0.01)\n",
" steps_per_epoch = config.get(\"steps_per_epoch\")\n",
"\n",
" deepspeed = {\n",
" \"fp16\": {\n",
" \"enabled\": \"auto\",\n",
" \"initial_scale_power\": 8,\n",
" \"hysteresis\": 4,\n",
" \"consecutive_hysteresis\": True,\n",
" },\n",
" \"bf16\": {\"enabled\": \"auto\"},\n",
" \"optimizer\": {\n",
" \"type\": \"AdamW\",\n",
" \"params\": {\n",
" \"lr\": \"auto\",\n",
" \"betas\": \"auto\",\n",
" \"eps\": \"auto\",\n",
" },\n",
" },\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"offload_optimizer\": {\n",
" \"device\": \"cpu\",\n",
" \"pin_memory\": True,\n",
" },\n",
" \"overlap_comm\": True,\n",
" \"contiguous_gradients\": True,\n",
" \"gather_16bit_weights_on_model_save\": True,\n",
" \"round_robin_gradients\": True,\n",
" },\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"gradient_clipping\": \"auto\",\n",
" \"steps_per_print\": 10,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"wall_clock_breakdown\": False,\n",
" }\n",
"\n",
" print(\"Preparing training arguments\")\n",
" training_args = TrainingArguments(\n",
" \"output\",\n",
" logging_steps=1,\n",
" save_strategy=\"steps\",\n",
" save_steps=steps_per_epoch,\n",
" max_steps=steps_per_epoch * epochs,\n",
" per_device_train_batch_size=batch_size,\n",
" gradient_accumulation_steps=1,\n",
" learning_rate=learning_rate,\n",
" weight_decay=weight_decay,\n",
" warmup_steps=warmup_steps,\n",
" label_names=[\"input_ids\", \"attention_mask\"],\n",
" push_to_hub=False,\n",
" report_to=\"none\",\n",
" disable_tqdm=True, # declutter the output a little\n",
" fp16=True,\n",
" gradient_checkpointing=True,\n",
" deepspeed=deepspeed,\n",
" )\n",
" disable_progress_bar()\n",
"\n",
" tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
"\n",
" print(\"Loading model\")\n",
"\n",
" model = GPTJForCausalLM.from_pretrained(model_name, use_cache=False)\n",
" model.resize_token_embeddings(len(tokenizer))\n",
"\n",
" print(\"Model loaded\")\n",
"\n",
" enable_progress_bar()\n",
"\n",
" metric = evaluate.load(\"accuracy\")\n",
"\n",
" train_ds = train.get_dataset_shard(\"train\")\n",
" eval_ds = train.get_dataset_shard(\"validation\")\n",
"\n",
" train_ds_iterable = train_ds.iter_torch_batches(\n",
" batch_size=batch_size,\n",
" local_shuffle_buffer_size=train.get_context().get_world_size() * batch_size,\n",
" )\n",
" eval_ds_iterable = eval_ds.iter_torch_batches(batch_size=batch_size)\n",
"\n",
" def compute_metrics(eval_pred):\n",
" logits, labels = eval_pred\n",
" predictions = np.argmax(logits, axis=-1)\n",
" return metric.compute(predictions=predictions, references=labels)\n",
"\n",
" trainer = Trainer(\n",
" model=model,\n",
" args=training_args,\n",
" train_dataset=train_ds_iterable,\n",
" eval_dataset=eval_ds_iterable,\n",
" compute_metrics=compute_metrics,\n",
" tokenizer=tokenizer,\n",
" data_collator=default_data_collator,\n",
" )\n",
"\n",
" # Add callback to report checkpoints to Ray Train\n",
" trainer.add_callback(RayTrainReportCallback())\n",
" trainer = prepare_trainer(trainer)\n",
" trainer.train()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After defining the training function, instantiate the {class}`~ray.train.torch.TorchTrainer`. Aside from the function, set the `scaling_config` to control the number of workers and amount of resources to use, and `datasets`(the preprocessed Ray Datasets) to use for training and evaluation.\n",
"\n",
"```{note}\n",
"Running with multiple nodes necessitates the persistence of checkpoints\n",
"and other outputs to some external storage for access after training has completed.\n",
"**You should set up cloud storage or NFS, then replace `storage_path` with your own cloud bucket URI or NFS path.**\n",
"\n",
"See {ref}`Configuration and Persistent Storage<persistent-storage-guide>` for more details.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"storage_path = \"s3://your-bucket-here\" # TODO: Set up cloud storage\n",
"# storage_path=\"/mnt/path/to/nfs\" # TODO: Alternatively, set up NFS"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"import os, re\n",
"\n",
"artifact_storage = os.environ.get(\"ANYSCALE_ARTIFACT_STORAGE\", \"artifact_storage\")\n",
"user_name = re.sub(r\"\\s+\", \"__\", os.environ.get(\"ANYSCALE_USERNAME\", \"user\"))\n",
"storage_path = f\"{artifact_storage}/{user_name}/gptj-deepspeed-finetune\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"batch_size = 16\n",
"train_ds_size = processed_datasets[\"train\"].count()\n",
"steps_per_epoch = train_ds_size // (batch_size * num_workers)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# SMOKE TEST SETTINGS FOR CI\n",
"steps_per_epoch = 10\n",
"num_workers = 8\n",
"batch_size = 1"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig\n",
"\n",
"trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func,\n",
" train_loop_config={\n",
" \"epochs\": 1,\n",
" \"batch_size\": batch_size, # per device\n",
" \"steps_per_epoch\": steps_per_epoch,\n",
" },\n",
" scaling_config=ScalingConfig(\n",
" num_workers=num_workers,\n",
" use_gpu=use_gpu,\n",
" resources_per_worker={\"GPU\": 1, \"CPU\": cpus_per_worker},\n",
" ),\n",
" datasets=processed_datasets,\n",
" run_config=RunConfig(storage_path=storage_path),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, call the {meth}`~ray.train.torch.TorchTrainer.fit` method to start training with Ray Train. Save the {class}`~ray.train.Result` object to a variable to access metrics and checkpoints."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"hide-output"
]
},
"outputs": [],
"source": [
"results = trainer.fit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Use the returned {class}`~ray.train.Result` object to access metrics and the Ray Train {class}`~ray.train.Checkpoint` associated with the last iteration."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Checkpoint(filesystem=<pyarrow._s3fs.S3FileSystem object at 0x7f8c59d311b0>, path=anyscale-staging-data-cld-kvedzwag2qa8i5bjxuevf5i7/org_7c1Kalm9WcX2bNIjW53GUT/cld_kvedZWag2qA8i5BjxUevf5i7/artifact_storage/yunxuan__xiao/gptj-deepspeed-finetune/TorchTrainer_2023-08-18_18-09-11/TorchTrainer_01ea5_00000_0_2023-08-18_18-09-12/checkpoint_000000)"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"checkpoint = results.checkpoint\n",
"checkpoint"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(gptj-predict)=\n",
"### Generate text from prompt\n",
"\n",
"First, download the persistent Ray Train checkpoint from a gpu node and load the fine-tuned model weights and tokenizer from the checkpoint. Then use \ud83e\udd17 Transformers [`pipeline`](https://huggingface.co/docs/transformers/en/main_classes/pipelines) to generate predictions from the fine-tuned model.\n",
"\n",
"```{tip}\n",
"For large scale batch inference, see {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`.\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Set the `task` to `\"text-generation\"`, and also set `device_map=\"auto\"` for Ray Train to automatically place the model on the right device. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'generated_text': 'Romeo and Juliet. This very night shall they come. A word with you, sir.'}]\n",
"[{'generated_text': 'Romeo! I know thee not. Lord Mercutio, is it you! Signior Montague.'}]\n",
"[{'generated_text': 'Juliet, look up in the vault, and there shalt find a grave; within the monument there is a table:'}]\n"
]
}
],
"source": [
"from transformers import pipeline, AutoTokenizer, GPTJForCausalLM\n",
"import os\n",
"\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def generate_text():\n",
" # Download the checkpoint\n",
" os.system(f\"aws s3 sync s3://{checkpoint.path} /mnt/local_storage/\")\n",
"\n",
" # Load the model and tokenizer \n",
" model = GPTJForCausalLM.from_pretrained(\"/mnt/local_storage/checkpoint\")\n",
" tokenizer = AutoTokenizer.from_pretrained(\"/mnt/local_storage/checkpoint\")\n",
"\n",
" pipe = pipeline(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" task=\"text-generation\",\n",
" torch_dtype=torch.float16,\n",
" device_map=\"auto\",\n",
" )\n",
"\n",
" # Generate from prompts!\n",
" result = []\n",
" for sentence in pipe(\n",
" [\"Romeo and Juliet\", \"Romeo\", \"Juliet\"], do_sample=True, min_length=20\n",
" ):\n",
" result.append(sentence)\n",
" \n",
" return result\n",
"\n",
"ref = generate_text.remote()\n",
"print(ray.get(ref))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.18"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/horovod/horovod_example.rst
+25
View File
@@ -0,0 +1,25 @@
:orphan:
Run Horovod Distributed Training with PyTorch and Ray Train
===========================================================
.. raw:: html
<a id="try-anyscale-quickstart-horovod_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=horovod_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example demonstrates how to run Horovod distributed training with PyTorch and Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/horovod/horovod_example.py
See also
--------
* :ref:`Get Started with Horovod <train-horovod>` for a tutorial on using Horovod with Ray Train
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/intel_gaudi/bert.ipynb
+416
View File
@@ -0,0 +1,416 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_bert_training)=\n",
"# BERT Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-bert\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-bert\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this notebook, we will train a BERT model for sequence classification using the Yelp review full dataset. We will use the `transformers` and `datasets` libraries from Hugging Face, along with `ray.train` for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install the following dependencies to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook transformers datasets evaluate\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
" warnings.warn(\n"
]
}
],
"source": [
"# Import necessary libraries\n",
"\n",
"import os\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"from torch import nn\n",
"from torch.utils.data import DataLoader\n",
"from tqdm import tqdm\n",
"\n",
"import numpy as np\n",
"import evaluate\n",
"from datasets import load_dataset\n",
"import transformers\n",
"from transformers import (\n",
" Trainer,\n",
" TrainingArguments,\n",
" AutoTokenizer,\n",
" AutoModelForSequenceClassification,\n",
")\n",
"\n",
"import ray.train\n",
"from ray.train import ScalingConfig\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Metrics Setup\n",
"\n",
"We will use accuracy as our evaluation metric. The `compute_metrics` function will calculate the accuracy of our model's predictions."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Metrics\n",
"metric = evaluate.load(\"accuracy\")\n",
"\n",
"def compute_metrics(eval_pred):\n",
" logits, labels = eval_pred\n",
" predictions = np.argmax(logits, axis=-1)\n",
" return metric.compute(predictions=predictions, references=labels)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training. It handles data loading, tokenization, model initialization, and the training loop. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def train_func_per_worker(config: Dict):\n",
" \n",
" # Datasets\n",
" dataset = load_dataset(\"yelp_review_full\")\n",
" tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
" \n",
" def tokenize_function(examples):\n",
" return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
"\n",
" lr = config[\"lr\"]\n",
" epochs = config[\"epochs\"]\n",
" batch_size = config[\"batch_size_per_worker\"]\n",
"\n",
" train_dataset = dataset[\"train\"].select(range(1000)).map(tokenize_function, batched=True)\n",
" eval_dataset = dataset[\"test\"].select(range(1000)).map(tokenize_function, batched=True)\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = {}\n",
" dataloaders[\"train\"] = torch.utils.data.DataLoader(\n",
" train_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
" dataloaders[\"test\"] = torch.utils.data.DataLoader(\n",
" eval_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
"\n",
" # Obtain HPU device automatically\n",
" device = ray.train.torch.get_device()\n",
"\n",
" # Prepare model and optimizer\n",
" model = AutoModelForSequenceClassification.from_pretrained(\n",
" \"bert-base-cased\", num_labels=5\n",
" )\n",
" model = model.to(device)\n",
" \n",
" optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
"\n",
" # Start training loops\n",
" for epoch in range(epochs):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"test\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" # breakpoint()\n",
" for batch in dataloaders[phase]:\n",
" batch = {k: v.to(device) for k, v in batch.items()}\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" \n",
" outputs = model(**batch)\n",
" loss = outputs.loss\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
" print(f\"train epoch:[{epoch}]\\tloss:{loss:.6f}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_bert` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def train_bert(num_workers=2):\n",
" global_batch_size = 8\n",
"\n",
" train_config = {\n",
" \"lr\": 1e-3,\n",
" \"epochs\": 10,\n",
" \"batch_size_per_worker\": global_batch_size // num_workers,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" # start your ray cluster\n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_bert` function to start the training process. You can adjust the number of workers to use.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_bert(num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"Downloading builder script: 100%|██████████| 4.20k/4.20k [00:00<00:00, 27.0MB/s]\n",
"2025-03-03 03:37:08,776 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=75123) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=74734) Started distributed worker processes: \n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75123) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75122) world_rank=1, local_rank=1, node_rank=0\n",
"Generating train split: 0%| | 0/650000 [00:00<?, ? examples/s]\n",
"Generating train split: 7%|▋ | 45000/650000 [00:00<00:01, 435976.18 examples/s]\n",
"Generating train split: 15%|█▍ | 95000/650000 [00:00<00:01, 469481.51 examples/s]\n",
"Generating train split: 23%|██▎ | 150000/650000 [00:00<00:01, 477676.99 examples/s]\n",
"Generating train split: 31%|███ | 203000/650000 [00:00<00:00, 493746.70 examples/s]\n",
"Generating train split: 43%|████▎ | 279000/650000 [00:00<00:00, 499340.09 examples/s]\n",
"Generating train split: 55%|█████▍ | 355000/650000 [00:00<00:00, 498613.65 examples/s]\n",
"Generating train split: 66%|██████▋ | 431000/650000 [00:00<00:00, 497799.19 examples/s]\n",
"Generating train split: 78%|███████▊ | 506000/650000 [00:01<00:00, 495696.93 examples/s]\n",
"Generating train split: 86%|████████▌ | 556000/650000 [00:01<00:00, 494508.05 examples/s]\n",
"Generating train split: 94%|█████████▎| 609000/650000 [00:01<00:00, 490725.53 examples/s]\n",
"Generating train split: 100%|██████████| 650000/650000 [00:01<00:00, 494916.42 examples/s]\n",
"Generating test split: 0%| | 0/50000 [00:00<?, ? examples/s]\n",
"Generating test split: 100%|██████████| 50000/50000 [00:00<00:00, 509619.87 examples/s]\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3998.33 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 4051.80 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3869.20 examples/s]\n",
"(RayTrainWorker pid=75123) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75123) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s] [repeated 3x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3782.66 examples/s] [repeated 2x across cluster]\n",
"(RayTrainWorker pid=75123) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=75123) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=75123) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=75123) Num CPU Cores : 160\n",
"(RayTrainWorker pid=75123) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=75123) ------------------------------------------------------------------------------\n",
"2025-03-03 03:41:04,658 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-37-11' in 0.0020s.\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-37-11\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-37-06_983992_65223/artifacts/2025-03-03_03-37-11/TorchTrainer_2025-03-03_03-37-11/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭─────────────────────────────────────────────────╮\n",
"│ Training config │\n",
"├─────────────────────────────────────────────────┤\n",
"│ train_loop_config/batch_size_per_worker 4 │\n",
"│ train_loop_config/epochs 10 │\n",
"│ train_loop_config/lr 0.001 │\n",
"╰─────────────────────────────────────────────────╯\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.979938\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.756611 [repeated 36x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.643875 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.416416 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.272513 [repeated 107x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:2.086884 [repeated 155x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.426217 [repeated 178x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[1] loss:0.991381 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.294097 [repeated 28x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.386306 [repeated 169x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.190416 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.171733 [repeated 130x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.287821 [repeated 152x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.055692 [repeated 179x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[3] loss:1.677789 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.942071 [repeated 19x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:1.592500 [repeated 167x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.936934 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:2.465384 [repeated 141x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.659170 [repeated 156x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.850438 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[5] loss:1.101623 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:2.125591 [repeated 18x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.612838 [repeated 170x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.759160 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.338552 [repeated 139x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.467959 [repeated 157x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.682137 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.395805 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.527835 [repeated 153x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.672311 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75122) train epoch:[8] loss:1.093186 [repeated 166x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.457587 [repeated 13x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.727377 [repeated 171x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.694001 [repeated 182x across cluster]\n",
"\n",
"Training completed after 0 iterations at 2025-03-03 03:41:04. Total running time: 3min 53s\n",
"\n",
"Training result: Result(\n",
" metrics={},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-37-11/TorchTrainer_ca6cf_00000_0_2025-03-03_03-37-11',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=75122) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75122) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"(RayTrainWorker pid=75122) train epoch:[9] loss:0.417845 [repeated 136x across cluster]\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/llama.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Fine-tuning Llama-2 Model with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will:\n",
"- fine-tuning a [Llama-2-7b](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) model by using Intel Gaudi accelerators with DDP method\n",
"- fine-tuning a [Llama-2-70b](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) model by using Intel Gaudi accelerators with DeepSpeed method\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training. We will use dataset [tatsu-lab/alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca).\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this fine-tuning example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- LoRA training.\n",
"- DDP or DeepSpeed based method.\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- Llama-2-7b/Llama-2-70b model.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import copy\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"\n",
"import datasets\n",
"import transformers\n",
"from transformers import DataCollatorForLanguageModeling\n",
"\n",
"from tqdm import tqdm\n",
"\n",
"import peft\n",
"\n",
"from optimum.habana import GaudiTrainer, GaudiConfig, GaudiTrainingArguments\n",
"from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare Dataset Function\n",
"\n",
"Preprocessing the raw dataset's each line with specified format."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset(raw_datasets):\n",
"\n",
" PROMPT_DICT = {\n",
" \"prompt_with_input\": (\n",
" \"Below is an instruction that describes a task, paired with an input that provides further context. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\"\n",
" ),\n",
" \"prompt_without_input\": (\n",
" \"Below is an instruction that describes a task. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Response:\"\n",
" ),\n",
" }\n",
"\n",
" def create_prompts(examples):\n",
" prompts = {}\n",
" prompts[\"source\"] = []\n",
" prompts[\"target\"] = []\n",
" for example in examples:\n",
" prompt_template = (\n",
" PROMPT_DICT[\"prompt_with_input\"] if example[\"input\"] != \"\" else PROMPT_DICT[\"prompt_without_input\"]\n",
" )\n",
" source = prompt_template.format_map(example)\n",
" prompts[\"source\"].append(source)\n",
" prompts[\"target\"].append(example[\"output\"])\n",
" return prompts\n",
"\n",
" # Preprocessing the datasets.\n",
" for key in raw_datasets:\n",
" prompts = create_prompts(raw_datasets[key])\n",
" columns_to_be_removed = list(raw_datasets[key].features.keys())\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_sources\", prompts[\"source\"])\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_targets\", prompts[\"target\"])\n",
" raw_datasets[key] = raw_datasets[key].remove_columns(columns_to_be_removed)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset to Tokenizer Function\n",
"\n",
"Tokenize each line in dataset by model tokenizer.\n",
"\n",
"In example codes, we concatenate the dataset's line content to accelerate training speed.\n",
"\n",
"All datasets are processed as \"train\" datasets, no evaluation datasets are sampled from raw_datasets."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset_to_tokenizer(raw_datasets, tokenizer):\n",
" max_seq_length = 512\n",
" tokenizer.pad_token_id = 0\n",
" tokenizer.eos_token_id = 1\n",
" tokenizer.bos_token_id = 2\n",
"\n",
" def tokenize(prompt, add_eos_token=True):\n",
" results = tokenizer(\n",
" prompt,\n",
" truncation=True,\n",
" max_length=max_seq_length,\n",
" padding=False,\n",
" return_tensors=None,\n",
" )\n",
" for i in range(len(results[\"input_ids\"])):\n",
" if (\n",
" results[\"input_ids\"][i][-1] != tokenizer.eos_token_id\n",
" and len(results[\"input_ids\"][i]) < max_seq_length\n",
" and add_eos_token\n",
" ):\n",
" results[\"input_ids\"][i].append(tokenizer.eos_token_id)\n",
" results[\"attention_mask\"][i].append(1)\n",
"\n",
" results[\"labels\"] = copy.deepcopy(results[\"input_ids\"])\n",
" results[\"input_id_len\"] = [len(result) for result in results[\"input_ids\"]]\n",
" return results\n",
"\n",
" def preprocess_function(examples):\n",
" keys = list(examples.data.keys())\n",
" if len(keys) != 2:\n",
" raise ValueError(\"Unsupported dataset format\")\n",
"\n",
" st = [s + t for s, t in zip(examples[keys[0]], examples[keys[1]])]\n",
"\n",
" examples_tokenized = tokenize(st)\n",
" input_ids = examples_tokenized[\"input_ids\"]\n",
" labels = examples_tokenized[\"labels\"]\n",
" return {\n",
" \"input_ids\": input_ids,\n",
" \"labels\": labels,\n",
" \"attention_mask\": examples_tokenized[\"attention_mask\"],\n",
" }\n",
"\n",
" tokenized_datasets = raw_datasets.map(\n",
" preprocess_function,\n",
" batched=True,\n",
" load_from_cache_file=True,\n",
" )\n",
"\n",
" def concatenate_data(dataset, max_seq_length):\n",
" concatenated_dataset = {}\n",
" for column in dataset.features:\n",
" concatenated_data = [item for sample in dataset[column] for item in sample]\n",
" reshaped_data = [\n",
" concatenated_data[i * max_seq_length : (i + 1) * max_seq_length]\n",
" for i in range(len(concatenated_data) // max_seq_length)\n",
" ]\n",
" concatenated_dataset[column] = reshaped_data\n",
" return datasets.Dataset.from_dict(concatenated_dataset)\n",
"\n",
" tokenized_datasets_ = tokenized_datasets[\"train\"].remove_columns([\"prompt_sources\", \"prompt_targets\"])\n",
" tokenized_datasets[\"train\"] = concatenate_data(tokenized_datasets_, max_seq_length)\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare training arguments\n",
"\n",
"here some arguments are hard coded, you can pass arguments from `config`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_training_args(config: Dict):\n",
" # prepare execution mode config\n",
" execution_mode = config[\"execution_mode\"]\n",
" use_lazy_mode = True if execution_mode == \"lazy\" else False\n",
" torch_compile_backend = \"hpu_backend\" if execution_mode == \"eager.compile\" else None\n",
"\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
"\n",
" return GaudiTrainingArguments(deepspeed=deepspeed,\n",
" output_dir=config[\"output\"],\n",
" do_train=True,\n",
" do_eval=False,\n",
" per_device_train_batch_size=config[\"batch_size_per_worker\"],\n",
" bf16=True,\n",
" learning_rate=config[\"lr\"],\n",
" save_strategy=\"no\",\n",
" torch_compile_backend=torch_compile_backend,\n",
" evaluation_strategy=\"no\",\n",
" lr_scheduler_type=\"cosine\",\n",
" num_train_epochs=config[\"epochs\"],\n",
" use_lazy_mode=use_lazy_mode,\n",
" use_habana=True,\n",
" pipelining_fwd_bwd=True,\n",
" save_only_model=True,\n",
" gradient_checkpointing=True,\n",
" warmup_ratio=0.03,\n",
" throughput_warmup_steps=3,\n",
" logging_steps=5)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare model\n",
"\n",
"1. download model from huggingface or read model from local directory.\n",
"2. convert model to lora model.\n",
"3. move model to HPU device.\n",
"\n",
"If you doesn't want to fine-tune with LoRA, just remove LoRA conversion step."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_model(config: Dict, device):\n",
" # prepare from pretrained model\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
" if deepspeed is not None:\n",
" auto_config = transformers.AutoConfig.from_pretrained(config[\"model\"], use_cache=False, revision=\"main\", use_auth_token=None, trust_remote_code=None)\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], config=auto_config, **config[\"model_config\"])\n",
" model.generation_config.attn_softmax_bf16 = True\n",
" model.generation_config.use_flash_attention = True\n",
" else:\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], **config[\"model_config\"])\n",
" model.enable_input_require_grads()\n",
"\n",
" # convert to peft model for lora training\n",
" peft_config = peft.LoraConfig(**config[\"lora_config\"])\n",
" model = peft.get_peft_model(model, peft_config)\n",
"\n",
" model.to(dtype=config[\"model_config\"][\"torch_dtype\"], device=device)\n",
"\n",
" return model\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"\n",
"- preparing training args, an instance of `GaudiTrainingArguments`.\n",
"- loading datasets and preprocess datasets, just load the first 4096 item as training datasets.\n",
"- loading pretrained model as tokenizer, and process datasets to tokenizer.\n",
"- loading pretrained model.\n",
"- preparing data collator and gaidu_config.\n",
"- preparing instance of `GaudiTrainer`.\n",
"- calling `train()` to train model.\n",
"- saving model results.\n",
"\n",
"Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"- Detect HPU and set the device.\n",
"- Initialize the habana PyTorch backend.\n",
"- Initialize the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_func_per_worker(config: Dict):\n",
" # adapt transformers to gaudi\n",
" adapt_transformers_to_gaudi()\n",
"\n",
" # prepare training arguments\n",
" training_args = prepare_training_args(config)\n",
"\n",
" # prepare datasets\n",
" # here we use dataset \"tatsu-lab/alpaca\" from huggingface\n",
" raw_datasets = datasets.DatasetDict({\"train\": datasets.load_dataset(\"tatsu-lab/alpaca\", split='train[0:4096]')})\n",
" preprocess_dataset(raw_datasets)\n",
"\n",
" # prepare tokenizer\n",
" tokenizer = transformers.AutoTokenizer.from_pretrained(config[\"model\"])\n",
" tokenized_datasets = preprocess_dataset_to_tokenizer(raw_datasets, tokenizer)\n",
"\n",
" # prepare model\n",
" model = prepare_model(config, training_args.device)\n",
"\n",
" # prepare data collator\n",
" data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors=\"pt\", mlm=False)\n",
"\n",
" # prepare gaudi config\n",
" gaudi_config = GaudiConfig()\n",
" gaudi_config.use_fused_adam = True\n",
" gaudi_config.use_fused_clip_norm = True\n",
"\n",
" # instance GaudiTrainer\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=tokenized_datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=data_collator,\n",
" compute_metrics=None,\n",
" preprocess_logits_for_metrics=None,\n",
" )\n",
"\n",
" train_result = trainer.train()\n",
" print(f\"train_result = {train_result}\")\n",
" trainer.save_model()\n",
"\n",
" return train_result"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"The `train_llama` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Set training method, supported method are:\n",
" - \"ddp\"\n",
" - \"deepspeed\"\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_llama(num_workers, execution_mode, training_method):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" # deepspeed config, can also place it to config file\n",
" deepspeed_config = {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" }\n",
"\n",
" # Preparing train configurations\n",
" train_config = {\n",
" \"execution_mode\": execution_mode,\n",
" \"model\": \"/root/models/models--meta-llama--Llama-2-70b-chat-hf/snapshots/e9149a12809580e8602995856f8098ce973d1080/\",\n",
" \"model_config\": {\"torch_dtype\": torch.bfloat16, \"trust_remote_code\": False, \"use_auth_token\": None},\n",
" \"lora_config\": {\"task_type\": \"CAUSAL_LM\", \"r\": 8, \"lora_alpha\": 32, \"lora_dropout\": 0.1, \"target_modules\": [\"q_proj\", \"v_proj\"]},\n",
" \"lr\": 1e-4,\n",
" \"epochs\": 2,\n",
" \"batch_size_per_worker\": 8,\n",
" \"output\": \"/tmp/ray/\",\n",
" \"deepspeed\": deepspeed_config if training_method == \"deepspeed\" else None,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
"\n",
" # start your ray cluster\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_llama` function to start the training process. You can adjust the number of workers to use, and the execution mode for HPU."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_DEVICES, such as\n",
"# os.environ[\"HABANA_VISIBLE_DEVICES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"# training_method are [\"ddp\", \"deepspeed\"]\n",
"training_method = \"deepspeed\"\n",
"if training_method == \"deepspeed\":\n",
" os.environ[\"PT_HPU_MAX_COMPOUND_OP_SIZE\"] = \"10\"\n",
" os.environ[\"DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED\"] = \"1\"\n",
"\n",
"# here use 4 HPUs\n",
"train_llama(num_workers=4, execution_mode=execution_mode, training_method=training_method)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Final output\n",
"\n",
"### For DDP on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=123181) {'loss': 1.8051, 'grad_norm': 0.6015625, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.6754, 'grad_norm': 0.408203125, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.568, 'grad_norm': 0.4453125, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.4934, 'grad_norm': 0.4609375, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3965, 'grad_norm': 0.3515625, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3461, 'grad_norm': 0.34765625, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2924, 'grad_norm': 0.32421875, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2643, 'grad_norm': 0.33203125, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.263, 'grad_norm': 0.318359375, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2502, 'grad_norm': 0.275390625, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2161, 'grad_norm': 0.2734375, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2517, 'grad_norm': 0.294921875, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"```\n",
"\n",
"### For DeepSpeed on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=50067) {'loss': 1.662, 'grad_norm': 0.36514782905578613, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.46, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.6047, 'grad_norm': 0.396455317735672, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4974, 'grad_norm': 0.49250370264053345, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4078, 'grad_norm': 0.49840453267097473, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.315, 'grad_norm': 0.3432576656341553, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.2651, 'grad_norm': 0.32175061106681824, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1947, 'grad_norm': 0.3646097481250763, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1534, 'grad_norm': 0.4598522186279297, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1404, 'grad_norm': 0.2677183449268341, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 32.75, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1283, 'grad_norm': 0.32087600231170654, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.0877, 'grad_norm': 0.28305548429489136, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1238, 'grad_norm': 0.25713953375816345, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/llama_pretrain.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Llama model pre-training on Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama_pretrain\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama_pretrain\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will pre-train a [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b) model by using Intel Gaudi accelerators.\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this pre-training example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- Pre-training llama model use configuration [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b)\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- DeepSpeed based pre-training.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"# more available docker image can be found here: https://vault.habana.ai/ui/native/gaudi-docker\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#!/usr/bin/env python\n",
"\n",
"import os\n",
"from typing import Any, Dict\n",
"from torch.utils.data import DataLoader\n",
"\n",
"import transformers\n",
"from itertools import chain\n",
"from datasets import load_dataset\n",
"from transformers import default_data_collator\n",
"from transformers.testing_utils import CaptureLogger\n",
"from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments\n",
"from optimum.habana.utils import set_seed"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Build datasets\n",
"\n",
"Download and load dataset from huggingface.co"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_datasets(config):\n",
" dataset_name = config[\"name\"] \n",
" dataset_config_name = config[\"config_name\"]\n",
"\n",
" # Downloading and loading a dataset from the hub.\n",
" raw_datasets = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" if \"validation\" not in raw_datasets.keys():\n",
" raw_datasets[\"validation\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[:{data_args.validation_split_percentage}%]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" raw_datasets[\"train\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[{data_args.validation_split_percentage}%:]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
"\n",
" return raw_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load tokenizer\n",
"\n",
"Download vocabulary from huggingface.co."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_tokenizer(config):\n",
" name = config[\"name\"]\n",
" tokenizer_kwargs = {\n",
" \"cache_dir\": None,\n",
" \"use_fast\": True,\n",
" \"revision\": \"main\",\n",
" \"token\": None,\n",
" \"trust_remote_code\": False,\n",
" }\n",
" return transformers.AutoTokenizer.from_pretrained(name, **tokenizer_kwargs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Tokenize dataset\n",
"\n",
"tokenize word to token ids."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def tokenize_dataset(datasets, tokenizer):\n",
" column_names = list(datasets[\"train\"].features)\n",
" text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
"\n",
" tok_logger = transformers.utils.logging.get_logger(\"transformers.tokenization_utils_base\")\n",
"\n",
" def tokenize_function(examples):\n",
" with CaptureLogger(tok_logger) as cl:\n",
" output = tokenizer(examples[text_column_name])\n",
" # clm input could be much much longer than block_size\n",
" if \"Token indices sequence length is longer than the\" in cl.out:\n",
" tok_logger.warning(\n",
" \"^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits\"\n",
" \" before being passed to the model.\"\n",
" )\n",
" return output\n",
"\n",
" tokenized_datasets = datasets.map(\n",
" tokenize_function,\n",
" batched=True,\n",
" num_proc=None,\n",
" remove_columns=column_names,\n",
" load_from_cache_file=True,\n",
" desc=\"Running tokenizer on dataset\",\n",
" )\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Group dataset\n",
"\n",
"This preprocssing will concatenate all texts from our dataset and generate chunks of block_size, and will pre-train model much faster."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def group_dataset(config, datasets, tokenizer):\n",
" config_name = config[\"name\"]\n",
" auto_config = transformers.AutoConfig.from_pretrained(config_name)\n",
" max_pos_embeddings = auto_config.max_position_embeddings\n",
" block_size = tokenizer.model_max_length\n",
" if block_size > max_pos_embeddings:\n",
" print(\n",
" f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n",
" f\"Using block_size={min(1024, max_pos_embeddings)} instead. You can change that default value by passing --block_size xxx.\"\n",
" )\n",
" if max_pos_embeddings > 0:\n",
" block_size = min(1024, max_pos_embeddings)\n",
" else:\n",
" block_size = 1024\n",
"\n",
" # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n",
" def group_texts(examples):\n",
" # Concatenate all texts.\n",
" concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n",
" total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
" # We drop the small remainder, and if the total_length < block_size we exclude this batch and return an empty dict.\n",
" # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.\n",
" total_length = (total_length // block_size) * block_size\n",
" # Split by chunks of max_len.\n",
" result = {\n",
" k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
" for k, t in concatenated_examples.items()\n",
" }\n",
" result[\"labels\"] = result[\"input_ids\"].copy()\n",
" return result\n",
"\n",
" lm_datasets = datasets.map(\n",
" group_texts,\n",
" batched=True,\n",
" num_proc=None,\n",
" load_from_cache_file=True,\n",
" desc=f\"Grouping texts in chunks of {block_size}\",\n",
" )\n",
" return lm_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model\n",
"\n",
"Download and load pre-configed model from huggingface.co, the detail model configurations in config.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_model(config):\n",
" name = config[\"name\"]\n",
" model_config = config.get(\"config\", {})\n",
" auto_config = transformers.AutoConfig.from_pretrained(\n",
" pretrained_model_name_or_path=name, **model_config\n",
" )\n",
" model = transformers.AutoModelForCausalLM.from_config(auto_config, trust_remote_code=False)\n",
"\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare trainer\n",
"\n",
"Instance Trainer with `model`, `gaudi_config`, `training_args`, `tokenizer`\n",
"\n",
"No evaluation dataset passed, just training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def get_trainer(training_args, datasets, tokenizer, model):\n",
" gaudi_config = GaudiConfig.from_pretrained(\n",
" training_args.gaudi_config_name, revision=\"main\",\n",
" )\n",
"\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=default_data_collator,\n",
" )\n",
" return trainer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"- prepare GaudiTrainingArguments object.\n",
"- load datasets from huggingface.co.\n",
"- load pre-configed tokenizer from huggingface.co.\n",
"- tokenize dataset with loaded model tokenizer.\n",
"- concatenate all texts from our dataset and generate chunks of block_size.\n",
"- instance object of `GaudiTrainer` with training_args, datasets, tokenizer, and model.\n",
"- call `train` of trainer.\n",
"- save model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def pretrain_llama(config: Dict[str, Any]):\n",
"\n",
" training_args = GaudiTrainingArguments(**config[\"training_args\"])\n",
" set_seed(training_args.seed)\n",
"\n",
" raw_datasets = load_datasets(config[\"datasets\"])\n",
"\n",
" tokenizer = load_tokenizer(config[\"tokenizer\"])\n",
"\n",
" tokenized_datasets = tokenize_dataset(raw_datasets, tokenizer)\n",
"\n",
" tokenized_datasets = group_dataset(config[\"model\"], tokenized_datasets, tokenizer)\n",
"\n",
" model = load_model(config[\"model\"])\n",
"\n",
" trainer = get_trainer(training_args, tokenized_datasets, tokenizer, model)\n",
"\n",
" result = trainer.train()\n",
" trainer.save_model()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `main` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def main(num_workers, execution_mode):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" pretrain_config = {\n",
" \"datasets\": {\n",
" \"name\": \"wikitext\",\n",
" \"config_name\": \"wikitext-2-raw-v1\",\n",
" },\n",
" \"tokenizer\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {}\n",
" },\n",
" \"model\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {\n",
" \"torch_dtype\": \"bfloat16\",\n",
" },\n",
" },\n",
" \"training_args\": {\n",
" \"per_device_train_batch_size\": 1,\n",
" \"do_train\": True,\n",
" \"save_strategy\": \"no\",\n",
" \"output_dir\": \"/tmp/ray/pretrain-llama-2\",\n",
" \"logging_steps\": 1,\n",
" \"gaudi_config_name\": \"Habana/llama\",\n",
" \"use_habana\": True,\n",
" \"throughput_warmup_steps\": 3,\n",
" \"use_lazy_mode\": True,\n",
" \"overwrite_output_dir\": True,\n",
" \"seed\": 42,\n",
" \"bf16\": True,\n",
" \"report_to\":'tensorboard',\n",
" \"deepspeed\": {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"reduce_scatter\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" },\n",
" },\n",
" }\n",
"\n",
" # if execution mode is eager with compile, must spcified with a compile backend\n",
" if execution_mode == \"eager.compile\":\n",
" pretrain_config[\"training_args\"].update({\"torch_compile_backend\": \"hpu_backend\"})\n",
"\n",
" scaling_config = ScalingConfig(num_workers=num_workers,\n",
" use_gpu=False,\n",
" resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
"\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend=\"hccl\")\n",
"\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=pretrain_llama,\n",
" train_loop_config=pretrain_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `main` function to start the pre-training process.\n",
"\n",
"Before calling `main` function, you must set some environment variables.\n",
"\n",
"1. The visiable devices. Environment variable `HABANA_VISIBLE_DEVICES` and `HABANA_VISIBLE_MODULES` are used to control the HPU device visiable to application, you must set this two environment variable properly. For more detail usage of `HABANA_VISIBLE_DEVICES`, `HABANA_VISIBLE_MODULES`, please visit [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PT_Multiple_Tenants_on_HPU/Multiple_Dockers_each_with_Single_Workload.html)\n",
"\n",
"2. The execution mode. Different execution mode has different runtime performance. The default execution mode is lazy mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_MODULES, such as\n",
"# os.environ[\"HABANA_VISIBLE_MODULES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"main(num_workers=4, execution_mode=execution_mode)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"\n",
"...\n",
"\n",
"RayTrainWorker pid=36561) Setting up process group for: env:// [rank=0, world_size=4]\n",
"(TorchTrainer pid=36054) Started distributed worker processes: \n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36561) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36562) world_rank=1, local_rank=1, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36563) world_rank=2, local_rank=2, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36564) world_rank=3, local_rank=3, node_rank=0\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=36561) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_ACC_PAR_MODE = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=36561) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=36561) Num CPU Cores : 160\n",
"(RayTrainWorker pid=36561) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=36561) ------------------------------------------------------------------------------\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) {'loss': 4.1052, 'grad_norm': 2.225008249282837, 'learning_rate': 8.26086956521739e-06, 'epoch': 2.5, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0472, 'grad_norm': 2.0701019763946533, 'learning_rate': 8.212560386473431e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.097, 'grad_norm': 2.119075059890747, 'learning_rate': 8.164251207729469e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7035, 'grad_norm': 2.1802899837493896, 'learning_rate': 8.115942028985508e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.242, 'grad_norm': 1.9516953229904175, 'learning_rate': 8.067632850241547e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9594, 'grad_norm': 2.0580222606658936, 'learning_rate': 8.019323671497584e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.3415, 'grad_norm': 2.192605495452881, 'learning_rate': 7.971014492753623e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9739, 'grad_norm': 2.0198025703430176, 'learning_rate': 7.922705314009662e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1624, 'grad_norm': 2.0957565307617188, 'learning_rate': 7.874396135265701e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9744, 'grad_norm': 2.1159448623657227, 'learning_rate': 7.82608695652174e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1127, 'grad_norm': 2.159834623336792, 'learning_rate': 7.777777777777777e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0588, 'grad_norm': 2.106534004211426, 'learning_rate': 7.729468599033817e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8734, 'grad_norm': 2.445814371109009, 'learning_rate': 7.681159420289856e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0278, 'grad_norm': 2.0376927852630615, 'learning_rate': 7.632850241545895e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9643, 'grad_norm': 2.1097891330718994, 'learning_rate': 7.584541062801932e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1384, 'grad_norm': 2.157325267791748, 'learning_rate': 7.536231884057972e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9982, 'grad_norm': 2.230065107345581, 'learning_rate': 7.48792270531401e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0988, 'grad_norm': 2.355875015258789, 'learning_rate': 7.439613526570048e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0514, 'grad_norm': 2.1178295612335205, 'learning_rate': 7.391304347826088e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9858, 'grad_norm': 2.089723825454712, 'learning_rate': 7.342995169082126e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1548, 'grad_norm': 2.308490753173828, 'learning_rate': 7.294685990338164e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0356, 'grad_norm': 1.9994627237319946, 'learning_rate': 7.246376811594203e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7696, 'grad_norm': 1.9719663858413696, 'learning_rate': 7.1980676328502416e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0157, 'grad_norm': 2.1598856449127197, 'learning_rate': 7.1497584541062814e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0113, 'grad_norm': 1.997869849205017, 'learning_rate': 7.10144927536232e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1048, 'grad_norm': 2.099222183227539, 'learning_rate': 7.053140096618358e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0048, 'grad_norm': 2.100231885910034, 'learning_rate': 7.004830917874397e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0302, 'grad_norm': 2.18204402923584, 'learning_rate': 6.956521739130435e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7227, 'grad_norm': 2.190962553024292, 'learning_rate': 6.908212560386473e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1111, 'grad_norm': 2.349518060684204, 'learning_rate': 6.859903381642513e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.024, 'grad_norm': 2.5497331619262695, 'learning_rate': 6.811594202898551e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8844, 'grad_norm': 2.3125178813934326, 'learning_rate': 6.7632850241545894e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.2208, 'grad_norm': 2.1103923320770264, 'learning_rate': 6.7149758454106285e-06, 'epoch': 2.6, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"\n",
"...\n",
"```"
]
}
],
"metadata": {
"language_info": {
"name": "python"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/resnet.ipynb
+610
View File
@@ -0,0 +1,610 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_resnet_training)=\n",
"# ResNet Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-resnet\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-resnet\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will train a ResNet-50 model to classify images of ants and bees using HPU. We will use PyTorch for model training and Ray for distributed training. The dataset will be downloaded and processed using torchvision's datasets and transforms.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install Ray and Jupyter to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from typing import Dict\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"import torch\n",
"from filelock import FileLock\n",
"from torch import nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets, transforms, models\n",
"from tqdm import tqdm\n",
"\n",
"import ray\n",
"import ray.train as train\n",
"from ray.train import ScalingConfig, Checkpoint\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define Data Transforms\n",
"\n",
"We will set up the data transforms for preprocessing images for training and validation. This includes random cropping, flipping, and normalization for the training set, and resizing and normalization for the validation set."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Data augmentation and normalization for training\n",
"# Just normalization for validation\n",
"data_transforms = {\n",
" \"train\": transforms.Compose([\n",
" transforms.RandomResizedCrop(224),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
" \"val\": transforms.Compose([\n",
" transforms.Resize(256),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Download Function\n",
"\n",
"We will define a function to download the Hymenoptera dataset. This dataset contains images of ants and bees for a binary classification problem."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def download_datasets():\n",
" os.system(\"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\")\n",
" os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Preparation Function\n",
"\n",
"After downloading the dataset, we need to build PyTorch datasets for training and validation. The `build_datasets` function will apply the previously defined transforms and create the datasets."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" return torch_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model Initialization Functions\n",
"\n",
"We will define two functions to initialize our model. The `initialize_model` function will load a pre-trained ResNet-50 model and replace the final classification layer for our binary classification task. The `initialize_model_from_checkpoint` function will load a model from a saved checkpoint if available."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def initialize_model():\n",
" # Load pretrained model params\n",
" model = models.resnet50(pretrained=True)\n",
"\n",
" # Replace the original classifier with a new Linear layer\n",
" num_features = model.fc.in_features\n",
" model.fc = nn.Linear(num_features, 2)\n",
"\n",
" # Ensure all params get updated during finetuning\n",
" for param in model.parameters():\n",
" param.requires_grad = True\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluation Function\n",
"\n",
"To assess the performance of our model during training, we define an `evaluate` function. This function computes the number of correct predictions by comparing the predicted labels with the true labels."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def evaluate(logits, labels):\n",
" _, preds = torch.max(logits, 1)\n",
" corrects = torch.sum(preds == labels).item()\n",
" return corrects"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Loop Function\n",
"\n",
"This function defines the training loop that will be executed by each worker. It includes downloading the dataset, preparing data loaders, initializing the model, and running the training and validation phases. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"def train_loop_per_worker(configs):\n",
" import warnings\n",
"\n",
" warnings.filterwarnings(\"ignore\")\n",
"\n",
" # Calculate the batch size for a single worker\n",
" worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
"\n",
" # Download dataset once on local rank 0 worker\n",
" if train.get_context().get_local_rank() == 0:\n",
" download_datasets()\n",
" torch.distributed.barrier()\n",
"\n",
" # Build datasets on each worker\n",
" torch_datasets = build_datasets()\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = dict()\n",
" dataloaders[\"train\"] = DataLoader(\n",
" torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
" )\n",
" dataloaders[\"val\"] = DataLoader(\n",
" torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
" )\n",
"\n",
" # Distribute\n",
" dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
" dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
"\n",
" # Obtain HPU device automatically\n",
" device = train.torch.get_device()\n",
"\n",
" # Prepare DDP Model, optimizer, and loss function\n",
" model = initialize_model()\n",
" model = model.to(device)\n",
"\n",
" optimizer = optim.SGD(\n",
" model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
" )\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" # Start training loops\n",
" for epoch in range(configs[\"num_epochs\"]):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"val\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" for inputs, labels in dataloaders[phase]:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # calculate statistics\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += evaluate(outputs, labels)\n",
"\n",
" size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
" epoch_loss = running_loss / size\n",
" epoch_acc = running_corrects / size\n",
"\n",
" if train.get_context().get_world_rank() == 0:\n",
" print(\n",
" \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
" epoch, phase, epoch_loss, epoch_acc\n",
" )\n",
" )\n",
"\n",
" # Report metrics and checkpoint every epoch\n",
" if phase == \"val\":\n",
" train.report(\n",
" metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_resnet` function sets up the distributed training environment using Ray and starts the training process. It specifies the batch size, number of epochs, learning rate, and momentum for the SGD optimizer. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def train_resnet(num_workers=2):\n",
" global_batch_size = 16\n",
"\n",
" train_loop_config = {\n",
" \"input_size\": 224, # Input image size (224 x 224)\n",
" \"batch_size\": 32, # Batch size for training\n",
" \"num_epochs\": 10, # Number of epochs to train for\n",
" \"lr\": 0.001, # Learning Rate\n",
" \"momentum\": 0.9, # SGD optimizer momentum\n",
" }\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config=train_loop_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_resnet` function to start the training process. You can adjust the number of workers to use. Before running this cell, ensure that Ray is properly set up in your environment to handle distributed training.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"train_resnet(num_workers=2) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"2025-03-03 03:32:12,620 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=63669) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=63280) Started distributed worker processes: \n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63669) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63668) world_rank=1, local_rank=1, node_rank=0\n",
"(RayTrainWorker pid=63669) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=63669) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=63669) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=63669) Num CPU Cores : 160\n",
"(RayTrainWorker pid=63669) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=63669) ------------------------------------------------------------------------------\n",
"(RayTrainWorker pid=63668) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 9%|▊ | 8.38M/97.8M [00:00<00:01, 87.7MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 193MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 203MB/s]\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-32-15\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-32-10_695011_53838/artifacts/2025-03-03_03-32-15/TorchTrainer_2025-03-03_03-32-15/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭──────────────────────────────────────╮\n",
"│ Training config │\n",
"├──────────────────────────────────────┤\n",
"│ train_loop_config/batch_size 32 │\n",
"│ train_loop_config/input_size 224 │\n",
"│ train_loop_config/lr 0.001 │\n",
"│ train_loop_config/momentum 0.9 │\n",
"│ train_loop_config/num_epochs 10 │\n",
"╰──────────────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-train Loss: 0.6574 Acc: 0.6066\n",
"\n",
"Training finished iteration 1 at 2025-03-03 03:32:45. Total running time: 29s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 24.684 │\n",
"│ time_total_s 24.684 │\n",
"│ training_iteration 1 │\n",
"│ acc 0.71053 │\n",
"│ loss 0.51455 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-val Loss: 0.5146 Acc: 0.7105\n",
"(RayTrainWorker pid=63669) Epoch 1-train Loss: 0.5016 Acc: 0.7541\n",
"\n",
"Training finished iteration 2 at 2025-03-03 03:32:46. Total running time: 31s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39649 │\n",
"│ time_total_s 26.0805 │\n",
"│ training_iteration 2 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.30218 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 1-val Loss: 0.3022 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 2-train Loss: 0.3130 Acc: 0.9180\n",
"\n",
"Training finished iteration 3 at 2025-03-03 03:32:47. Total running time: 32s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37042 │\n",
"│ time_total_s 27.4509 │\n",
"│ training_iteration 3 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.22201 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 2-val Loss: 0.2220 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 3-train Loss: 0.2416 Acc: 0.9262\n",
"\n",
"Training finished iteration 4 at 2025-03-03 03:32:49. Total running time: 34s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.38353 │\n",
"│ time_total_s 28.8345 │\n",
"│ training_iteration 4 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.17815 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 3-val Loss: 0.1782 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 4-train Loss: 0.1900 Acc: 0.9508\n",
"\n",
"Training finished iteration 5 at 2025-03-03 03:32:50. Total running time: 35s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37318 │\n",
"│ time_total_s 30.2077 │\n",
"│ training_iteration 5 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.17063 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 4-val Loss: 0.1706 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 5-train Loss: 0.1346 Acc: 0.9672\n",
"\n",
"Training finished iteration 6 at 2025-03-03 03:32:52. Total running time: 36s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37999 │\n",
"│ time_total_s 31.5876 │\n",
"│ training_iteration 6 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.1552 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 5-val Loss: 0.1552 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 6-train Loss: 0.1184 Acc: 0.9672\n",
"\n",
"Training finished iteration 7 at 2025-03-03 03:32:53. Total running time: 38s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39198 │\n",
"│ time_total_s 32.9796 │\n",
"│ training_iteration 7 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14702 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 6-val Loss: 0.1470 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 7-train Loss: 0.0864 Acc: 0.9836\n",
"\n",
"Training finished iteration 8 at 2025-03-03 03:32:54. Total running time: 39s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.3736 │\n",
"│ time_total_s 34.3532 │\n",
"│ training_iteration 8 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14443 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 7-val Loss: 0.1444 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 8-train Loss: 0.1085 Acc: 0.9590\n",
"\n",
"Training finished iteration 9 at 2025-03-03 03:32:56. Total running time: 40s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37868 │\n",
"│ time_total_s 35.7319 │\n",
"│ training_iteration 9 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14194 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 8-val Loss: 0.1419 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 9-train Loss: 0.0829 Acc: 0.9754\n",
"\n",
"2025-03-03 03:32:58,628 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-32-15' in 0.0028s.\n",
"Training finished iteration 10 at 2025-03-03 03:32:57. Total running time: 42s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.36497 │\n",
"│ time_total_s 37.0969 │\n",
"│ training_iteration 10 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.14297 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 9-val Loss: 0.1430 Acc: 0.9605\n",
"\n",
"Training completed after 10 iterations at 2025-03-03 03:32:58. Total running time: 43s\n",
"\n",
"Training result: Result(\n",
" metrics={'loss': 0.1429688463869848, 'acc': 0.9605263157894737},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-32-15/TorchTrainer_19fd8_00000_0_2025-03-03_03-32-15',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=63669) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 68%|██████▊ | 66.1M/97.8M [00:00<00:00, 160MB/s] [repeated 6x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/sd.ipynb
+494
View File
File diff suppressed because one or more lines are too long
doc/source/train/examples/jax/intro_to_jax_trainer/README.ipynb
+878
View File
@@ -0,0 +1,878 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0ebcd4f7",
"metadata": {},
"source": [
"# Train a GPT-2 model with Ray Train `JaxTrainer`\n",
"\n",
"**Time to complete**: 15 min\n",
"\n",
"This template shows you how to distribute a JAX/Flax training loop with [Ray Train](https://docs.ray.io/en/latest/train/train.html)'s `JaxTrainer`. You’ll train a small GPT-2-style Transformer from scratch on the [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/) dataset.\n",
"\n",
"Ray Train lets you keep your training code in a normal Python function, then runs that function on a set of Ray workers. `JaxTrainer` handles the orchestration (starting workers, setting up the distributed context, and collecting metrics/checkpoints) so you can focus on the model and the input pipeline.\n",
"\n",
"This tutorial is inspired by Andrej Karpathy’s [nanoGPT](https://github.com/karpathy/nanoGPT/tree/master) and Google’s [Train a GPT-2 model with JAX on TPU for free](https://developers.googleblog.com/en/train-gpt2-model-with-jax-on-tpu/).\n",
"\n",
"In this tutorial, you will:\n",
"\n",
"1. Prepare and the `OpenWebText` dataset and wrap with Ray Data.\n",
"2. Define a basic GPT2 model and train step in Jax/Flax.\n",
"3. Wrap the training loop in a `train_loop_per_worker` function and scale it out using Ray Train `JaxTrainer` with GPUs or TPUs!\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "a9fe11b6",
"metadata": {},
"source": [
"<div id=\"anyscale-note\" class=\"alert alert-block alert-warning\">\n",
"\n",
" <strong>Anyscale specific configuration</strong>\n",
"\n",
" <p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, additional configuration is required. For example, you need to manually:</p>\n",
"\n",
" <ul>\n",
" <li><strong>Configure your Ray cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li>\n",
" <li><strong>Manage dependencies</strong>: Manually install and manage dependencies on each node.</li>\n",
" <li><strong>Set up storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>\n",
" </ul>\n",
"</div>\n",
"\n",
"<style>\n",
" div#anyscale-note > p,\n",
" div#anyscale-note > ul,\n",
" div#anyscale-note > ul li {\n",
" color: black;\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" background-color: rgb(255, 243, 205);\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" border: 1px solid #ccc; \n",
" border-radius: 8px;\n",
" padding: 15px;\n",
" }\n",
"\n",
"</style>"
]
},
{
"cell_type": "markdown",
"id": "bf58e5fd",
"metadata": {},
"source": [
"## Step 1: Install dependencies and prepare the dataset\n"
]
},
{
"cell_type": "markdown",
"id": "1704e5cd",
"metadata": {},
"source": [
"First, install the required Python packages.\n",
"\n",
"If you’re running on GPUs, install `jax[cuda]`. If you’re running on Google TPUs, install `jax[tpu]`. For platform-specific requirements, see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).\n",
"\n",
"This notebook uses `jax[cuda]` in the examples for simplicity.\n"
]
},
{
"cell_type": "markdown",
"id": "39b71121",
"metadata": {},
"source": [
"Run below if you plan to use GPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f8a6d4fc",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"pip install pandas numpy jax[cuda] flax tiktoken datasets transformers orbax optax\n"
]
},
{
"cell_type": "markdown",
"id": "d862bb66",
"metadata": {},
"source": [
"Run below if you plan to use TPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f5bf2907",
"metadata": {},
"outputs": [],
"source": [
"# %%bash\n",
"# pip install pandas numpy jax[tpu] flax tiktoken datasets transformers orbax optax"
]
},
{
"cell_type": "markdown",
"id": "193fafe8",
"metadata": {},
"source": [
"Next, prepare the data that you’ll feed into the training loop.\n",
"\n",
"This notebook uses [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/), an open reproduction of OpenAI’s (private) WebText. The goal of this step is to tokenize the dataset once and write it to disk as two files:\n",
"\n",
"- `train.bin`: training tokens.\n",
"- `val.bin`: validation tokens.\n",
"\n",
"You can use Karpathy’s nanoGPT prep script ([`prepare.py`](https://github.com/karpathy/nanoGPT/blob/master/data/openwebtext/prepare.py)) or download preprocessed data from Kaggle ([OpenWebText GPT-2](https://www.kaggle.com/datasets/windmaple/openwebtext-gpt2)).\n",
"\n",
"If running on the Anyscale workspace, the following code adapts the nanoGPT approach and writes the output to the shared storage path used in an Anyscale workspace (`/mnt/cluster_storage`). If you already have `train.bin` and `val.bin`, you can skip this step.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90dcb1ba",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from tqdm import tqdm\n",
"import numpy as np\n",
"import tiktoken\n",
"from datasets import load_dataset # huggingface datasets\n",
"\n",
"# number of workers in .map() call\n",
"# good number to use is ~order number of cpu cores // 2\n",
"num_proc = 8\n",
"storage_path = \"/mnt/cluster_storage/openwebtext\"\n",
"\n",
"# number of workers in load_dataset() call\n",
"# best number might be different from num_proc above as it also depends on NW speed.\n",
"# it is better than 1 usually though\n",
"num_proc_load_dataset = num_proc\n",
"\n",
"enc = tiktoken.get_encoding(\"gpt2\")\n",
"\n",
"dataset = load_dataset(\"openwebtext\", num_proc=num_proc_load_dataset)\n",
"\n",
"# owt by default only contains the 'train' split, so create a test split\n",
"split_dataset = dataset[\"train\"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)\n",
"split_dataset['val'] = split_dataset.pop('test') # rename the test split to val\n",
"# we now want to tokenize the dataset. first define the encoding function (gpt2 bpe)\n",
"def process(example):\n",
" ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens\n",
" ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe\n",
" # note: I think eot should be prepended not appended... hmm. it's called \"eot\" though...\n",
" out = {'ids': ids, 'len': len(ids)}\n",
" return out\n",
"\n",
"# tokenize the dataset\n",
"tokenized = split_dataset.map(\n",
" process,\n",
" remove_columns=['text'],\n",
" desc=\"tokenizing the splits\",\n",
" num_proc=num_proc,\n",
")\n",
"\n",
"# concatenate all the ids in each dataset into one large file we can use for training\n",
"for split, dset in tokenized.items():\n",
" arr_len = np.sum(dset['len'], dtype=np.uint64)\n",
" filename = os.path.join(storage_path, f'{split}.bin')\n",
" dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16)\n",
" arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))\n",
" total_batches = 1024\n",
"\n",
" idx = 0\n",
" for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):\n",
" # Batch together samples for faster write\n",
" batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')\n",
" arr_batch = np.concatenate(batch['ids'])\n",
" # Write into mmap\n",
" arr[idx : idx + len(arr_batch)] = arr_batch\n",
" idx += len(arr_batch)\n",
" arr.flush()\n",
"\n",
"# train.bin is ~18GB, val.bin ~8.8MB\n",
"# train has ~9B tokens \n",
"# val has ~4M tokens"
]
},
{
"cell_type": "markdown",
"id": "d89cd3f4",
"metadata": {},
"source": [
"After running the script, you should have two files in shared storage:\n",
"\n",
"1. Training dataset: `/mnt/cluster_storage/openwebtext/train.bin`\n",
"2. Validation dataset: `/mnt/cluster_storage/openwebtext/val.bin`\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "72f1caa8",
"metadata": {},
"source": [
"## Step 2: Load the dataset with Ray Data."
]
},
{
"cell_type": "markdown",
"id": "75c698e5",
"metadata": {},
"source": [
"Next, let's load these files with Ray Data. Ray Data can read and preprocess data in parallel, then shard it across workers so each training process streams its own batches. This keeps the input pipeline scalable as you add more GPUs or TPU hosts.\n",
"\n",
"For more details about Ray Data, check out the [Ray Data documentation](https://docs.ray.io/en/latest/data/data.html).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4f234bb",
"metadata": {},
"outputs": [],
"source": [
"import ray\n",
"import ray.data\n",
"\n",
"def make_bin_xy_dataset(\n",
" bin_path: str,\n",
" seqlen: int,\n",
" *,\n",
" # How many sequences to generate per epoch-like pass.\n",
" # You can make this very large and then re-iterate over batches in the\n",
" # training loop.\n",
" num_sequences: int,\n",
" seed: int = 0,\n",
" dtype=np.uint16,\n",
" concurrency: int = 64,\n",
"):\n",
" \"\"\"\n",
" Build a Ray Dataset of (x,y) sequences sampled randomly from a .bin token file.\n",
"\n",
" Produces rows:\n",
" - x: int32[seqlen]\n",
" - y: int32[seqlen]\n",
" \"\"\"\n",
" if not os.path.exists(bin_path):\n",
" raise FileNotFoundError(bin_path)\n",
"\n",
" # Open memmap on driver just to get length.\n",
" data = np.memmap(bin_path, dtype=dtype, mode=\"r\")\n",
" n = int(len(data))\n",
" if n <= seqlen + 1:\n",
" raise ValueError(f\"{bin_path} too small: len={n}, seqlen={seqlen}\")\n",
"\n",
" rng = np.random.default_rng(seed)\n",
" # Each start index uses [i : i+seqlen+1]\n",
" starts = rng.integers(0, n - (seqlen + 1), size=num_sequences, dtype=np.int64)\n",
"\n",
" # Create a dataset of start indices\n",
" ds = ray.data.from_items([{\"i\": int(i)} for i in starts])\n",
"\n",
" def read_xy(batch):\n",
" # Open memmap inside worker process\n",
" mm = np.memmap(bin_path, dtype=dtype, mode=\"r\")\n",
" idx = batch[\"i\"].astype(np.int64, copy=False)\n",
"\n",
" # Allocate fixed arrays\n",
" bs = idx.shape[0]\n",
" x = np.empty((bs, seqlen), dtype=np.int32)\n",
" y = np.empty((bs, seqlen), dtype=np.int32)\n",
"\n",
" # Slice per row (still Python loop, but runs in parallel across Ray workers)\n",
" for j, start in enumerate(idx):\n",
" window = mm[start : start + seqlen + 1].astype(np.int32, copy=False)\n",
" x[j] = window[:-1]\n",
" y[j] = window[1:]\n",
"\n",
" return {\"x\": x, \"y\": y}\n",
"\n",
" # Batch reading for efficiency\n",
" ds = ds.map_batches(\n",
" read_xy,\n",
" batch_format=\"numpy\",\n",
" batch_size=32,\n",
" compute=ray.data.TaskPoolStrategy(size=concurrency),\n",
" zero_copy_batch=True,\n",
" )\n",
"\n",
" return ds\n",
"\n",
"train_ds = make_bin_xy_dataset(\n",
" \"/mnt/cluster_storage/openwebtext/train.bin\",\n",
" seqlen=1024,\n",
" num_sequences=5_000_000,\n",
" seed=2357,\n",
" )\n",
"val_ds = make_bin_xy_dataset(\n",
" \"/mnt/cluster_storage/openwebtext/val.bin\",\n",
" seqlen=1024,\n",
" num_sequences=5_000_000,\n",
" seed=2357,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "87878c88",
"metadata": {},
"source": [
"## Step 3: Define a JAX/Flax GPT-2-style model\n"
]
},
{
"cell_type": "markdown",
"id": "9d3a679f",
"metadata": {},
"source": [
"Now define the model and the core training step with JAX and Flax.\n",
"\n",
"The JAX ecosystem is modular: JAX provides the array programming and compilation primitives, and libraries such as [Flax](https://github.com/google/flax) (neural network building blocks), [Optax](https://github.com/google-deepmind/optax) (optimizers and losses), and [Orbax](https://github.com/google/orbax) (checkpointing) provide higher-level components. You’ll use all three in this tutorial.\n",
"\n",
"In this section, nothing is Ray-specific yet—you’re building a normal single-process JAX/Flax training step that you’ll scale out with `JaxTrainer` in the next section.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25f1fde5",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import time\n",
"import numpy as np\n",
"from dataclasses import dataclass\n",
"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"from jax.experimental import mesh_utils\n",
"from jax.sharding import Mesh\n",
"\n",
"import flax.nnx as nnx\n",
"import optax\n",
"import orbax.checkpoint as orbax\n",
"\n",
"import tiktoken\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "267cb3bf",
"metadata": {},
"outputs": [],
"source": [
"@dataclass(frozen=True)\n",
"class TrainingConfig:\n",
" # Model config (GPT-2 base model configuration).\n",
" tokenizer = tiktoken.get_encoding(\"gpt2\") # We use gpt2 tokenizer for this tutorial.\n",
" vocab_size = tokenizer.n_vocab\n",
"\n",
" num_transformer_blocks: int = 12\n",
" seqlen: int = 1024\n",
" embed_dim: int = 768\n",
"\n",
" num_heads: int = 12\n",
" dropout_rate: float = 0.1\n",
"\n",
" dtype = jnp.bfloat16 # change to jnp.float32 for older GPUs\n",
" param_dtype = jnp.float32\n",
"\n",
" @property\n",
" def feed_forward_dim(self) -> int:\n",
" return 4 * self.embed_dim\n",
"\n",
" # Optimizer config.\n",
" init_learning_rate: float = 5e-4\n",
" weight_decay: float = 1e-1\n",
"\n",
" # Training loop config.\n",
" global_batch_size: int = 32\n",
" max_steps: int = 10_000\n",
" log_every_n_steps: int = 10\n",
" val_every_n_steps: int = 100\n",
" checkpoint_every_n_steps: int = 100\n",
"\n",
" # Data/config paths.\n",
" openwebtext_root: str = \"/mnt/cluster_storage/openwebtext\""
]
},
{
"cell_type": "markdown",
"id": "df1bdea2",
"metadata": {},
"source": [
"This tutorial provides a GPT-2-style Transformer model implemented with [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html). The model code is standard JAX/Flax—Ray Train doesn’t require any special model wrappers.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d9103cb",
"metadata": {},
"outputs": [],
"source": [
"# --- Model Definitions (Safe to keep global) ---\n",
"# Keep dtype settings in one place so the model code doesn't rely on undefined globals.\n",
"dtype = TrainingConfig.dtype\n",
"param_dtype = TrainingConfig.param_dtype\n",
"\n",
"\n",
"def causal_attention_mask(seq_len):\n",
" return jnp.tril(jnp.ones((seq_len, seq_len)))\n",
"\n",
"class TransformerBlock(nnx.Module):\n",
" def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, dropout_rate: float, rngs: nnx.Rngs):\n",
" self.layer_norm1 = nnx.LayerNorm(epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.mha = nnx.MultiHeadAttention(num_heads=num_heads,\n",
" in_features=embed_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.dropout1 = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
" self.layer_norm2 = nnx.LayerNorm(epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.linear1 = nnx.Linear(in_features=embed_dim,\n",
" out_features=ff_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.linear2 = nnx.Linear(in_features=ff_dim,\n",
" out_features=embed_dim,\n",
" kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs)\n",
" self.dropout2 = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
"\n",
" def __call__(self, inputs, training: bool = False):\n",
" input_shape = inputs.shape\n",
" bs, seq_len, emb_sz = input_shape\n",
" attention_output = self.mha(\n",
" inputs_q=self.layer_norm1(inputs),\n",
" mask=causal_attention_mask(seq_len),\n",
" decode=False,\n",
" )\n",
" x = inputs + self.dropout1(attention_output, deterministic=not training)\n",
" mlp_output = self.linear1(self.layer_norm2(x))\n",
" mlp_output = nnx.gelu(mlp_output)\n",
" mlp_output = self.linear2(mlp_output)\n",
" mlp_output = self.dropout2(mlp_output, deterministic=not training)\n",
" return x + mlp_output\n",
"\n",
"class TokenAndPositionEmbedding(nnx.Module):\n",
" def __init__(self, seqlen: int, vocab_size: int, embed_dim: int, rngs: nnx.Rngs):\n",
" self.token_emb = nnx.Embed(num_embeddings=vocab_size, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)\n",
" self.pos_emb = nnx.Embed(num_embeddings=seqlen, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)\n",
"\n",
" def __call__(self, x):\n",
" positions = jnp.arange(0, x.shape[1])[None, :]\n",
" position_embedding = self.pos_emb(positions)\n",
" token_embedding = self.token_emb(x)\n",
" return self.token_emb, token_embedding+position_embedding\n",
"\n",
"class GPT2(nnx.Module):\n",
" def __init__(\n",
" self,\n",
" seqlen: int,\n",
" vocab_size: int,\n",
" embed_dim: int,\n",
" num_heads: int,\n",
" dropout_rate: float,\n",
" feed_forward_dim: int,\n",
" num_transformer_blocks: int,\n",
" rngs: nnx.Rngs,\n",
" ):\n",
" self.embedding_layer = TokenAndPositionEmbedding(seqlen, vocab_size, embed_dim, rngs=rngs)\n",
" self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n",
" self.transformer_blocks = nnx.List([\n",
" TransformerBlock(embed_dim, num_heads, feed_forward_dim, dropout_rate, rngs=rngs)\n",
" for _ in range(num_transformer_blocks)\n",
" ])\n",
" self.layer_norm = nnx.LayerNorm(\n",
" epsilon=1e-6,\n",
" num_features=embed_dim,\n",
" scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), (\"model\",)),\n",
" bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), (\"model\",)),\n",
" dtype=dtype,\n",
" param_dtype=param_dtype,\n",
" rngs=rngs,\n",
" )\n",
"\n",
" def __call__(self, inputs, training: bool = False):\n",
" token_embedding, x = self.embedding_layer(inputs)\n",
" x = self.dropout(x, deterministic=not training)\n",
" for transformer_block in self.transformer_blocks:\n",
" x = transformer_block(x, training=training)\n",
" x = self.layer_norm(x)\n",
" outputs = token_embedding.attend(x)\n",
" return outputs\n",
"\n",
"\n",
"def create_model(*, rngs: nnx.Rngs, config: TrainingConfig):\n",
" return GPT2(\n",
" seqlen=config.seqlen,\n",
" vocab_size=config.vocab_size,\n",
" embed_dim=config.embed_dim,\n",
" num_heads=config.num_heads,\n",
" dropout_rate=config.dropout_rate,\n",
" feed_forward_dim=config.feed_forward_dim,\n",
" num_transformer_blocks=config.num_transformer_blocks,\n",
" rngs=rngs,\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "3d453706",
"metadata": {},
"source": [
"Next, define the `loss_fn_train`, `loss_fn_eval`, and `train_step` functions.\n",
"\n",
"`train_step()` computes the loss, takes gradients, and updates model parameters through the optimizer. The training loop will call this function repeatedly.\n",
"\n",
"For performance, this notebook JIT-compiles these functions with `@nnx.jit`.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a819d3ba",
"metadata": {},
"outputs": [],
"source": [
"\n",
"@nnx.jit\n",
"def loss_fn_train(model, batch):\n",
" logits = model(batch[0], training=True)\n",
" loss = optax.softmax_cross_entropy_with_integer_labels(\n",
" logits=logits, labels=batch[1]\n",
" ).mean()\n",
" return loss, logits\n",
"\n",
"\n",
"@nnx.jit\n",
"def loss_fn_eval(model, batch):\n",
" logits = model(batch[0], training=False)\n",
" loss = optax.softmax_cross_entropy_with_integer_labels(\n",
" logits=logits, labels=batch[1]\n",
" ).mean()\n",
" return loss, logits\n",
"\n",
"\n",
"@nnx.jit\n",
"def train_step(model, optimizer, metrics, batch):\n",
" grad_fn = nnx.value_and_grad(loss_fn_train, has_aux=True)\n",
" (loss, logits), grads = grad_fn(model, batch)\n",
" metrics.update(loss=loss, logits=logits, labels=batch[1])\n",
" optimizer.update(model, grads)\n",
" return loss\n"
]
},
{
"cell_type": "markdown",
"id": "d9cccbed",
"metadata": {},
"source": [
"Next, let's wrap the JAX training logic in a `train_loop_per_worker` function.\n",
"\n",
"Each Ray Train worker runs the same Python function with a different world rank, and Ray sets device visibility per worker (for example, one GPU per worker). Inside this function, you can:\n",
"\n",
"- Read the distributed context (`world_rank`, `world_size`).\n",
"- Get the per-worker dataset shard (`train.get_dataset_shard(...)`) to stream batches.\n",
"- Stream sharded JAX Arrays directly to training devices using [`Dataset.iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html), which handles device sharding natively.\n",
"- Report metrics and checkpoints back to the trainer with `ray.train.report(...)`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f4543c82",
"metadata": {},
"outputs": [],
"source": [
"\n",
"import ray\n",
"from ray import train\n",
"from ray.train import Checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "57ea38fa",
"metadata": {},
"outputs": [],
"source": [
"def train_loop_per_worker(config_dict: dict) -> None:\n",
"\n",
" config = TrainingConfig(**config_dict)\n",
"\n",
" world_rank = ray.train.get_context().get_world_rank()\n",
" world_size = ray.train.get_context().get_world_size()\n",
" print(f\"Worker rank {world_rank}/{world_size} sees devices: {jax.devices()}\")\n",
"\n",
"\n",
" # Create a mesh per worker process. \n",
" device_mesh = mesh_utils.create_device_mesh((jax.process_count(), 1))\n",
" mesh = Mesh(device_mesh, axis_names=(\"data\", \"model\"))\n",
" jax.set_mesh(mesh)\n",
"\n",
"\n",
" # Initialize the model locally.\n",
" # Include a dropout RNG stream so nnx.Dropout can run when training=True.\n",
" model = create_model(rngs=nnx.Rngs(params=0, dropout=1), config=config)\n",
"\n",
"\n",
" # We use Ray data to load the training and validation datasets.\n",
" train_it = ray.train.get_dataset_shard(\"train\")\n",
" val_it = ray.train.get_dataset_shard(\"val\")\n",
" if train_it is None or val_it is None:\n",
" raise RuntimeError(\"No Ray Train datasets provided. Pass datasets={...} to JaxTrainer.\")\n",
" \n",
" local_batch_size = config.global_batch_size // jax.process_count()\n",
" \n",
" # A generator helper that yields JAX Array batches indefinitely.\n",
" # This removes the need for try/except StopIteration blocks in the training loop.\n",
" def batch_generator(dataset_shard):\n",
" while True:\n",
" # iter_jax_batches yields globally sharded JAX Array batches directly,\n",
" # mapped onto all process JAX-addressable training devices.\n",
" yield from dataset_shard.iter_jax_batches(\n",
" batch_size=local_batch_size,\n",
" prefetch_batches=2,\n",
" drop_last=True,\n",
" )\n",
"\n",
" train_batches = batch_generator(train_it)\n",
" val_batches = batch_generator(val_it)\n",
"\n",
" # Initialize the optimizer.\n",
" schedule = optax.cosine_decay_schedule(\n",
" init_value=config.init_learning_rate,\n",
" decay_steps=config.max_steps,\n",
" )\n",
" optax_chain = optax.chain(optax.adamw(learning_rate=schedule, weight_decay=config.weight_decay))\n",
" optimizer = nnx.Optimizer(model, optax_chain, wrt=nnx.Param)\n",
"\n",
" checkpointer = orbax.PyTreeCheckpointer()\n",
" start_time = time.time()\n",
"\n",
" train_metrics = nnx.metrics.Average('loss')\n",
" val_metrics = nnx.metrics.Average('val_loss')\n",
"\n",
" for step in range(config.max_steps):\n",
" batch = next(train_batches)\n",
" global_x, global_y = batch[\"x\"], batch[\"y\"]\n",
"\n",
" train_loss = train_step(model, optimizer, train_metrics, (global_x, global_y))\n",
"\n",
" if (step + 1) % config.log_every_n_steps == 0:\n",
" elapsed = time.time() - start_time\n",
" # Report metrics through Ray Train.\n",
" ray.train.report({\"step\": step + 1, \"train_loss\": float(train_loss), \"elapsed_s\": elapsed})\n",
" start_time = time.time()\n",
" \n",
" if (step + 1) % config.val_every_n_steps == 0:\n",
" validation_batch = next(val_batches)\n",
" global_val_input, global_val_target = validation_batch[\"x\"], validation_batch[\"y\"]\n",
" \n",
" loss, logits = loss_fn_eval(model, (global_val_input, global_val_target))\n",
" val_metrics.update(val_loss=loss, logits=logits)\n",
" val_loss = float(val_metrics.compute())\n",
" metrics = {\"step\": step + 1, \"train_loss\": float(train_loss),\"val_loss\": float(val_loss)}\n",
" \n",
" checkpoint = None\n",
" if (step + 1) % config.checkpoint_every_n_steps == 0:\n",
" \n",
" # Orbax checkpointing is a barrier.\n",
" train_state = nnx.to_pure_dict(nnx.state(model))\n",
" checkpoint_path = os.path.join(\"/mnt/cluster_storage/checkpoint/jax_gpt2_ray_data\", str(step + 1))\n",
" checkpointer.save(checkpoint_path, train_state)\n",
" \n",
" # Save a checkpoint and report validation metrics through Ray Train.\n",
" # The controller persists the checkpoint to the RunConfig storage path.\n",
" checkpoint = Checkpoint.from_directory(checkpoint_path) \n",
" if world_rank == 0:\n",
" train.report(metrics, checkpoint=checkpoint)\n",
" else:\n",
" train.report(metrics, checkpoint=None)"
]
},
{
"cell_type": "markdown",
"id": "c9b59b73",
"metadata": {},
"source": [
"> [!NOTE]\n",
"> [`iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html) uses an internal 1D mesh for sharding. If your training loop uses a complex multi-dimensional mesh, JAX might perform implicit resharding. Ensure your mesh aligns with the device ordering to minimize overhead.\n"
]
},
{
"cell_type": "markdown",
"id": "36f206e9",
"metadata": {},
"source": [
"## Step 5: Define the `ScalingConfig`\n",
"\n",
"Let's define the `ScalingConfig` that we want to scale the training process. \n",
"\n",
"`JaxTrainer` now supports both GPU training and TPU training. \n",
"\n",
"For a walkthrough on configuring `ScalingConfig`, see [Get Started with Distributed Training using JAX](https://docs.ray.io/en/latest/train/getting-started-jax.html). "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "260fbd24",
"metadata": {},
"outputs": [],
"source": [
"from ray.train import ScalingConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f1a1d22d",
"metadata": {},
"outputs": [],
"source": [
"# In this example, we use 2 GPUs.\n",
"scaling_config = ScalingConfig(\n",
" use_gpu=True,\n",
" num_workers=2, # Change this to match on your GPU cluster setting, by default, each worker uses one GPU.\n",
")\n",
"# If you plan to use TPUs, see an example below.\n",
"# This ScalingConfig requires a KubeRay cluster configured for a TPU v6e 4x4 slice with 4 TPU VMs.\n",
"# For more information about TPU clusters with Ray on Kubernetes, see the\n",
"# KubeRay TPU guide: https://docs.ray.io/en/master/cluster/kubernetes/user-guides/tpu.html#kuberay-tpu\n",
"# scaling_config = ScalingConfig(\n",
"# use_tpu=True,\n",
"# num_workers=4,\n",
"# topology=\"4x4\",\n",
"# accelerator_type=\"TPU-V6E\",\n",
"# resources_per_worker={\"TPU\": 4},\n",
"# )"
]
},
{
"cell_type": "markdown",
"id": "dadca9ae",
"metadata": {},
"source": [
"## Step 6: Launch with `JaxTrainer`\n",
"\n",
"To run `train_loop_per_worker` on a Ray cluster, you'll construct a `JaxTrainer` with:\n",
"\n",
"- `train_loop_per_worker`: the training function you've defined earlier. Each Ray Train worker runs this function.\n",
"- `train_loop_config`: a hyperparameter dictionary passed into the function.\n",
"- `scaling_config`: the `ScalingConfig` you've defined earlier.\n",
"- `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`.\n",
"- `run_config`: runtime configuration including where to write outputs such as checkpoints.\n",
"\n",
"If your workers hit CUDA or XLA library load errors, clear `LD_LIBRARY_PATH` in the runtime env to avoid picking up incompatible system libraries.\n",
"\n",
"`trainer.fit` spawns a controller process to orchestrate the training run and worker processes to execute the JAX training code."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "42b00efd",
"metadata": {},
"outputs": [],
"source": [
"from ray.train import RunConfig\n",
"from ray.train.v2.jax import JaxTrainer"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fc12fdfb",
"metadata": {},
"outputs": [],
"source": [
"storage_path = \"/mnt/cluster_storage\"\n",
"\n",
"trainer = JaxTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config={\n",
" \"global_batch_size\": 32,\n",
" },\n",
" scaling_config=scaling_config,\n",
" run_config=RunConfig(\n",
" name=\"jax_gpt2\",\n",
" storage_path=storage_path,\n",
" # Make sure to unset ``LD_LIBRARY_PATH`` if you're using CUDA devices, \n",
" # since ``LD_LIBRARY_PATH`` can override the CUDA libraries.\n",
" worker_runtime_env={\"env_vars\": {\"LD_LIBRARY_PATH\": \"\"}},\n",
" ),\n",
" datasets={\"train\": train_ds, \"val\": val_ds},\n",
")\n",
"\n",
"result = trainer.fit()\n",
"print(result)\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "4a0d85e1",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"In this notebook, you:\n",
"\n",
"1. Prepared and the `OpenWebText` dataset and wrapped with Ray Data.\n",
"2. Defined a basic GPT2 model and train step in Jax/Flax.\n",
"3. Wrapped the training loop in a `train_loop_per_worker` function and scaled it out using Ray Train `JaxTrainer` with GPUs or TPUs!"
]
},
{
"cell_type": "markdown",
"id": "c37928bd",
"metadata": {},
"source": []
}
],
"metadata": {
"language_info": {
"name": "python"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/train/examples/jax/intro_to_jax_trainer/README.md
+684
View File
@@ -0,0 +1,684 @@
# Train a GPT-2 model with Ray Train `JaxTrainer`
**Time to complete**: 15 min
This template shows you how to distribute a JAX/Flax training loop with [Ray Train](https://docs.ray.io/en/latest/train/train.html)'s `JaxTrainer`. You’ll train a small GPT-2-style Transformer from scratch on the [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/) dataset.
Ray Train lets you keep your training code in a normal Python function, then runs that function on a set of Ray workers. `JaxTrainer` handles the orchestration (starting workers, setting up the distributed context, and collecting metrics/checkpoints) so you can focus on the model and the input pipeline.
This tutorial is inspired by Andrej Karpathy’s [nanoGPT](https://github.com/karpathy/nanoGPT/tree/master) and Google’s [Train a GPT-2 model with JAX on TPU for free](https://developers.googleblog.com/en/train-gpt2-model-with-jax-on-tpu/).
In this tutorial, you will:
1. Prepare and the `OpenWebText` dataset and wrap with Ray Data.
2. Define a basic GPT2 model and train step in Jax/Flax.
3. Wrap the training loop in a `train_loop_per_worker` function and scale it out using Ray Train `JaxTrainer` with GPUs or TPUs!
<div id="anyscale-note" class="alert alert-block alert-warning">
<strong>Anyscale specific configuration</strong>
<p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, additional configuration is required. For example, you need to manually:</p>
<ul>
<li><strong>Configure your Ray cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li> <li><strong>Manage dependencies</strong>: Manually install and manage dependencies on each node.</li> <li><strong>Set up storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>
</ul>
</div>
<style>
div#anyscale-note > p, div#anyscale-note > ul, div#anyscale-note > ul li { color: black; }
div#anyscale-note { background-color: rgb(255, 243, 205); }
div#anyscale-note { border: 1px solid #ccc; border-radius: 8px; padding: 15px; }
</style>
## Step 1: Install dependencies and prepare the dataset
First, install the required Python packages.
If you’re running on GPUs, install `jax[cuda]`. If you’re running on Google TPUs, install `jax[tpu]`. For platform-specific requirements, see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).
This notebook uses `jax[cuda]` in the examples for simplicity.
Run below if you plan to use GPUs.
```bash
%%bash
pip install pandas numpy jax[cuda] flax tiktoken datasets transformers orbax optax
```
Run below if you plan to use TPUs.
```python
# %%bash
# pip install pandas numpy jax[tpu] flax tiktoken datasets transformers orbax optax
```
Next, prepare the data that you’ll feed into the training loop.
This notebook uses [OpenWebText](https://openwebtext2.readthedocs.io/en/latest/), an open reproduction of OpenAI’s (private) WebText. The goal of this step is to tokenize the dataset once and write it to disk as two files:
- `train.bin`: training tokens.
- `val.bin`: validation tokens.
You can use Karpathy’s nanoGPT prep script ([`prepare.py`](https://github.com/karpathy/nanoGPT/blob/master/data/openwebtext/prepare.py)) or download preprocessed data from Kaggle ([OpenWebText GPT-2](https://www.kaggle.com/datasets/windmaple/openwebtext-gpt2)).
If running on the Anyscale workspace, the following code adapts the nanoGPT approach and writes the output to the shared storage path used in an Anyscale workspace (`/mnt/cluster_storage`). If you already have `train.bin` and `val.bin`, you can skip this step.
```python
import os
from tqdm import tqdm
import numpy as np
import tiktoken
from datasets import load_dataset # huggingface datasets
# number of workers in .map() call
# good number to use is ~order number of cpu cores // 2
num_proc = 8
storage_path = "/mnt/cluster_storage/openwebtext"
# number of workers in load_dataset() call
# best number might be different from num_proc above as it also depends on NW speed.
# it is better than 1 usually though
num_proc_load_dataset = num_proc
enc = tiktoken.get_encoding("gpt2")
dataset = load_dataset("openwebtext", num_proc=num_proc_load_dataset)
# owt by default only contains the 'train' split, so create a test split
split_dataset = dataset["train"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)
split_dataset['val'] = split_dataset.pop('test') # rename the test split to val
# we now want to tokenize the dataset. first define the encoding function (gpt2 bpe)
def process(example):
ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens
ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe
# note: I think eot should be prepended not appended... hmm. it's called "eot" though...
out = {'ids': ids, 'len': len(ids)}
return out
# tokenize the dataset
tokenized = split_dataset.map(
process,
remove_columns=['text'],
desc="tokenizing the splits",
num_proc=num_proc,
)
# concatenate all the ids in each dataset into one large file we can use for training
for split, dset in tokenized.items():
arr_len = np.sum(dset['len'], dtype=np.uint64)
filename = os.path.join(storage_path, f'{split}.bin')
dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16)
arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))
total_batches = 1024
idx = 0
for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):
# Batch together samples for faster write
batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')
arr_batch = np.concatenate(batch['ids'])
# Write into mmap
arr[idx : idx + len(arr_batch)] = arr_batch
idx += len(arr_batch)
arr.flush()
# train.bin is ~18GB, val.bin ~8.8MB
# train has ~9B tokens
# val has ~4M tokens
```
After running the script, you should have two files in shared storage:
1. Training dataset: `/mnt/cluster_storage/openwebtext/train.bin`
2. Validation dataset: `/mnt/cluster_storage/openwebtext/val.bin`
## Step 2: Load the dataset with Ray Data.
Next, let's load these files with Ray Data. Ray Data can read and preprocess data in parallel, then shard it across workers so each training process streams its own batches. This keeps the input pipeline scalable as you add more GPUs or TPU hosts.
For more details about Ray Data, check out the [Ray Data documentation](https://docs.ray.io/en/latest/data/data.html).
```python
import ray
import ray.data
def make_bin_xy_dataset(
bin_path: str,
seqlen: int,
*,
# How many sequences to generate per epoch-like pass.
# You can make this very large and then re-iterate over batches in the
# training loop.
num_sequences: int,
seed: int = 0,
dtype=np.uint16,
concurrency: int = 64,
):
"""
Build a Ray Dataset of (x,y) sequences sampled randomly from a .bin token file.
Produces rows:
- x: int32[seqlen]
- y: int32[seqlen]
"""
if not os.path.exists(bin_path):
raise FileNotFoundError(bin_path)
# Open memmap on driver just to get length.
data = np.memmap(bin_path, dtype=dtype, mode="r")
n = int(len(data))
if n <= seqlen + 1:
raise ValueError(f"{bin_path} too small: len={n}, seqlen={seqlen}")
rng = np.random.default_rng(seed)
# Each start index uses [i : i+seqlen+1]
starts = rng.integers(0, n - (seqlen + 1), size=num_sequences, dtype=np.int64)
# Create a dataset of start indices
ds = ray.data.from_items([{"i": int(i)} for i in starts])
def read_xy(batch):
# Open memmap inside worker process
mm = np.memmap(bin_path, dtype=dtype, mode="r")
idx = batch["i"].astype(np.int64, copy=False)
# Allocate fixed arrays
bs = idx.shape[0]
x = np.empty((bs, seqlen), dtype=np.int32)
y = np.empty((bs, seqlen), dtype=np.int32)
# Slice per row (still Python loop, but runs in parallel across Ray workers)
for j, start in enumerate(idx):
window = mm[start : start + seqlen + 1].astype(np.int32, copy=False)
x[j] = window[:-1]
y[j] = window[1:]
return {"x": x, "y": y}
# Batch reading for efficiency
ds = ds.map_batches(
read_xy,
batch_format="numpy",
batch_size=32,
compute=ray.data.TaskPoolStrategy(size=concurrency),
zero_copy_batch=True,
)
return ds
train_ds = make_bin_xy_dataset(
"/mnt/cluster_storage/openwebtext/train.bin",
seqlen=1024,
num_sequences=5_000_000,
seed=2357,
)
val_ds = make_bin_xy_dataset(
"/mnt/cluster_storage/openwebtext/val.bin",
seqlen=1024,
num_sequences=5_000_000,
seed=2357,
)
```
## Step 3: Define a JAX/Flax GPT-2-style model
Now define the model and the core training step with JAX and Flax.
The JAX ecosystem is modular: JAX provides the array programming and compilation primitives, and libraries such as [Flax](https://github.com/google/flax) (neural network building blocks), [Optax](https://github.com/google-deepmind/optax) (optimizers and losses), and [Orbax](https://github.com/google/orbax) (checkpointing) provide higher-level components. You’ll use all three in this tutorial.
In this section, nothing is Ray-specific yet—you’re building a normal single-process JAX/Flax training step that you’ll scale out with `JaxTrainer` in the next section.
```python
import os
import time
import numpy as np
from dataclasses import dataclass
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import Mesh
import flax.nnx as nnx
import optax
import orbax.checkpoint as orbax
import tiktoken
```
```python
@dataclass(frozen=True)
class TrainingConfig:
# Model config (GPT-2 base model configuration).
tokenizer = tiktoken.get_encoding("gpt2") # We use gpt2 tokenizer for this tutorial.
vocab_size = tokenizer.n_vocab
num_transformer_blocks: int = 12
seqlen: int = 1024
embed_dim: int = 768
num_heads: int = 12
dropout_rate: float = 0.1
dtype = jnp.bfloat16 # change to jnp.float32 for older GPUs
param_dtype = jnp.float32
@property
def feed_forward_dim(self) -> int:
return 4 * self.embed_dim
# Optimizer config.
init_learning_rate: float = 5e-4
weight_decay: float = 1e-1
# Training loop config.
global_batch_size: int = 32
max_steps: int = 10_000
log_every_n_steps: int = 10
val_every_n_steps: int = 100
checkpoint_every_n_steps: int = 100
# Data/config paths.
openwebtext_root: str = "/mnt/cluster_storage/openwebtext"
```
This tutorial provides a GPT-2-style Transformer model implemented with [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html). The model code is standard JAX/Flax—Ray Train doesn’t require any special model wrappers.
```python
# --- Model Definitions (Safe to keep global) ---
# Keep dtype settings in one place so the model code doesn't rely on undefined globals.
dtype = TrainingConfig.dtype
param_dtype = TrainingConfig.param_dtype
def causal_attention_mask(seq_len):
return jnp.tril(jnp.ones((seq_len, seq_len)))
class TransformerBlock(nnx.Module):
def __init__(self, embed_dim: int, num_heads: int, ff_dim: int, dropout_rate: float, rngs: nnx.Rngs):
self.layer_norm1 = nnx.LayerNorm(epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.mha = nnx.MultiHeadAttention(num_heads=num_heads,
in_features=embed_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.dropout1 = nnx.Dropout(rate=dropout_rate, rngs=rngs)
self.layer_norm2 = nnx.LayerNorm(epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ('model',)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.linear1 = nnx.Linear(in_features=embed_dim,
out_features=ff_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.linear2 = nnx.Linear(in_features=ff_dim,
out_features=embed_dim,
kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, 'model')),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ('model',)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs)
self.dropout2 = nnx.Dropout(rate=dropout_rate, rngs=rngs)
def __call__(self, inputs, training: bool = False):
input_shape = inputs.shape
bs, seq_len, emb_sz = input_shape
attention_output = self.mha(
inputs_q=self.layer_norm1(inputs),
mask=causal_attention_mask(seq_len),
decode=False,
)
x = inputs + self.dropout1(attention_output, deterministic=not training)
mlp_output = self.linear1(self.layer_norm2(x))
mlp_output = nnx.gelu(mlp_output)
mlp_output = self.linear2(mlp_output)
mlp_output = self.dropout2(mlp_output, deterministic=not training)
return x + mlp_output
class TokenAndPositionEmbedding(nnx.Module):
def __init__(self, seqlen: int, vocab_size: int, embed_dim: int, rngs: nnx.Rngs):
self.token_emb = nnx.Embed(num_embeddings=vocab_size, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)
self.pos_emb = nnx.Embed(num_embeddings=seqlen, features=embed_dim, dtype=dtype, param_dtype=param_dtype, rngs=rngs)
def __call__(self, x):
positions = jnp.arange(0, x.shape[1])[None, :]
position_embedding = self.pos_emb(positions)
token_embedding = self.token_emb(x)
return self.token_emb, token_embedding+position_embedding
class GPT2(nnx.Module):
def __init__(
self,
seqlen: int,
vocab_size: int,
embed_dim: int,
num_heads: int,
dropout_rate: float,
feed_forward_dim: int,
num_transformer_blocks: int,
rngs: nnx.Rngs,
):
self.embedding_layer = TokenAndPositionEmbedding(seqlen, vocab_size, embed_dim, rngs=rngs)
self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)
self.transformer_blocks = nnx.List([
TransformerBlock(embed_dim, num_heads, feed_forward_dim, dropout_rate, rngs=rngs)
for _ in range(num_transformer_blocks)
])
self.layer_norm = nnx.LayerNorm(
epsilon=1e-6,
num_features=embed_dim,
scale_init=nnx.with_partitioning(nnx.initializers.ones_init(), ("model",)),
bias_init=nnx.with_partitioning(nnx.initializers.zeros_init(), ("model",)),
dtype=dtype,
param_dtype=param_dtype,
rngs=rngs,
)
def __call__(self, inputs, training: bool = False):
token_embedding, x = self.embedding_layer(inputs)
x = self.dropout(x, deterministic=not training)
for transformer_block in self.transformer_blocks:
x = transformer_block(x, training=training)
x = self.layer_norm(x)
outputs = token_embedding.attend(x)
return outputs
def create_model(*, rngs: nnx.Rngs, config: TrainingConfig):
return GPT2(
seqlen=config.seqlen,
vocab_size=config.vocab_size,
embed_dim=config.embed_dim,
num_heads=config.num_heads,
dropout_rate=config.dropout_rate,
feed_forward_dim=config.feed_forward_dim,
num_transformer_blocks=config.num_transformer_blocks,
rngs=rngs,
)
```
Next, define the `loss_fn_train`, `loss_fn_eval`, and `train_step` functions.
`train_step()` computes the loss, takes gradients, and updates model parameters through the optimizer. The training loop will call this function repeatedly.
For performance, this notebook JIT-compiles these functions with `@nnx.jit`.
```python
@nnx.jit
def loss_fn_train(model, batch):
logits = model(batch[0], training=True)
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch[1]
).mean()
return loss, logits
@nnx.jit
def loss_fn_eval(model, batch):
logits = model(batch[0], training=False)
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch[1]
).mean()
return loss, logits
@nnx.jit
def train_step(model, optimizer, metrics, batch):
grad_fn = nnx.value_and_grad(loss_fn_train, has_aux=True)
(loss, logits), grads = grad_fn(model, batch)
metrics.update(loss=loss, logits=logits, labels=batch[1])
optimizer.update(model, grads)
return loss
```
## Step 4: Wrap training logic in `train_loop_per_worker`
Next, let's wrap the JAX training logic in a `train_loop_per_worker` function.
Each Ray Train worker runs the same Python function with a different world rank, and Ray sets device visibility per worker (for example, one GPU per worker). Inside this function, you can:
- Read the distributed context (`world_rank`, `world_size`).
- Get the per-worker dataset shard (`train.get_dataset_shard(...)`) to stream batches.
- Stream sharded JAX Arrays directly to training devices using [`Dataset.iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html), which handles device sharding natively.
- Report metrics and checkpoints back to the trainer with `ray.train.report(...)`.
```python
import ray
from ray import train
from ray.train import Checkpoint
```
```python
def train_loop_per_worker(config_dict: dict) -> None:
config = TrainingConfig(**config_dict)
world_rank = ray.train.get_context().get_world_rank()
world_size = ray.train.get_context().get_world_size()
print(f"Worker rank {world_rank}/{world_size} sees devices: {jax.devices()}")
# Create a mesh per worker process.
device_mesh = mesh_utils.create_device_mesh((jax.process_count(), 1))
mesh = Mesh(device_mesh, axis_names=("data", "model"))
jax.set_mesh(mesh)
# Initialize the model locally.
# Include a dropout RNG stream so nnx.Dropout can run when training=True.
model = create_model(rngs=nnx.Rngs(params=0, dropout=1), config=config)
# We use Ray data to load the training and validation datasets.
train_it = ray.train.get_dataset_shard("train")
val_it = ray.train.get_dataset_shard("val")
if train_it is None or val_it is None:
raise RuntimeError("No Ray Train datasets provided. Pass datasets={...} to JaxTrainer.")
local_batch_size = config.global_batch_size // jax.process_count()
# A generator helper that yields JAX Array batches indefinitely.
# This removes the need for try/except StopIteration blocks in the training loop.
def batch_generator(dataset_shard):
while True:
# iter_jax_batches yields globally sharded JAX Array batches directly,
# mapped onto all process JAX-addressable training devices.
yield from dataset_shard.iter_jax_batches(
batch_size=local_batch_size,
prefetch_batches=2,
drop_last=True,
)
train_batches = batch_generator(train_it)
val_batches = batch_generator(val_it)
# Initialize the optimizer.
schedule = optax.cosine_decay_schedule(
init_value=config.init_learning_rate,
decay_steps=config.max_steps,
)
optax_chain = optax.chain(optax.adamw(learning_rate=schedule, weight_decay=config.weight_decay))
optimizer = nnx.Optimizer(model, optax_chain, wrt=nnx.Param)
checkpointer = orbax.PyTreeCheckpointer()
start_time = time.time()
train_metrics = nnx.metrics.Average('loss')
val_metrics = nnx.metrics.Average('val_loss')
for step in range(config.max_steps):
batch = next(train_batches)
global_x, global_y = batch["x"], batch["y"]
train_loss = train_step(model, optimizer, train_metrics, (global_x, global_y))
if (step + 1) % config.log_every_n_steps == 0:
elapsed = time.time() - start_time
# Report metrics through Ray Train.
ray.train.report({"step": step + 1, "train_loss": float(train_loss), "elapsed_s": elapsed})
start_time = time.time()
if (step + 1) % config.val_every_n_steps == 0:
validation_batch = next(val_batches)
global_val_input, global_val_target = validation_batch["x"], validation_batch["y"]
loss, logits = loss_fn_eval(model, (global_val_input, global_val_target))
val_metrics.update(val_loss=loss, logits=logits)
val_loss = float(val_metrics.compute())
metrics = {"step": step + 1, "train_loss": float(train_loss),"val_loss": float(val_loss)}
checkpoint = None
if (step + 1) % config.checkpoint_every_n_steps == 0:
# Orbax checkpointing is a barrier.
train_state = nnx.to_pure_dict(nnx.state(model))
checkpoint_path = os.path.join("/mnt/cluster_storage/checkpoint/jax_gpt2_ray_data", str(step + 1))
checkpointer.save(checkpoint_path, train_state)
# Save a checkpoint and report validation metrics through Ray Train.
# The controller persists the checkpoint to the RunConfig storage path.
checkpoint = Checkpoint.from_directory(checkpoint_path)
if world_rank == 0:
train.report(metrics, checkpoint=checkpoint)
else:
train.report(metrics, checkpoint=None)
```
> [!NOTE]
> [`iter_jax_batches`](https://docs.ray.io/en/latest/data/api/doc/ray.data.Dataset.iter_jax_batches.html) uses an internal 1D mesh for sharding. If your training loop uses a complex multi-dimensional mesh, JAX might perform implicit resharding. Ensure your mesh aligns with the device ordering to minimize overhead.
## Step 5: Define the `ScalingConfig`
Let's define the `ScalingConfig` that we want to scale the training process.
`JaxTrainer` now supports both GPU training and TPU training.
For a walkthrough on configuring `ScalingConfig`, see [Get Started with Distributed Training using JAX](https://docs.ray.io/en/latest/train/getting-started-jax.html).
```python
from ray.train import ScalingConfig
```
```python
# In this example, we use 2 GPUs.
scaling_config = ScalingConfig(
use_gpu=True,
num_workers=2, # Change this to match on your GPU cluster setting, by default, each worker uses one GPU.
)
# If you plan to use TPUs, see an example below.
# This ScalingConfig requires a KubeRay cluster configured for a TPU v6e 4x4 slice with 4 TPU VMs.
# For more information about TPU clusters with Ray on Kubernetes, see the
# KubeRay TPU guide: https://docs.ray.io/en/master/cluster/kubernetes/user-guides/tpu.html#kuberay-tpu
# scaling_config = ScalingConfig(
# use_tpu=True,
# num_workers=4,
# topology="4x4",
# accelerator_type="TPU-V6E",
# resources_per_worker={"TPU": 4},
# )
```
## Step 6: Launch with `JaxTrainer`
To run `train_loop_per_worker` on a Ray cluster, you'll construct a `JaxTrainer` with:
- `train_loop_per_worker`: the training function you've defined earlier. Each Ray Train worker runs this function.
- `train_loop_config`: a hyperparameter dictionary passed into the function.
- `scaling_config`: the `ScalingConfig` you've defined earlier.
- `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`.
- `run_config`: runtime configuration including where to write outputs such as checkpoints.
If your workers hit CUDA or XLA library load errors, clear `LD_LIBRARY_PATH` in the runtime env to avoid picking up incompatible system libraries.
`trainer.fit` spawns a controller process to orchestrate the training run and worker processes to execute the JAX training code.
```python
from ray.train import RunConfig
from ray.train.v2.jax import JaxTrainer
```
```python
storage_path = "/mnt/cluster_storage"
trainer = JaxTrainer(
train_loop_per_worker=train_loop_per_worker,
train_loop_config={
"global_batch_size": 32,
},
scaling_config=scaling_config,
run_config=RunConfig(
name="jax_gpt2",
storage_path=storage_path,
# Make sure to unset ``LD_LIBRARY_PATH`` if you're using CUDA devices,
# since ``LD_LIBRARY_PATH`` can override the CUDA libraries.
worker_runtime_env={"env_vars": {"LD_LIBRARY_PATH": ""}},
),
datasets={"train": train_ds, "val": val_ds},
)
result = trainer.fit()
print(result)
```
## Summary
In this notebook, you:
1. Prepared and the `OpenWebText` dataset and wrapped with Ray Data.
2. Defined a basic GPT2 model and train step in Jax/Flax.
3. Wrapped the training loop in a `train_loop_per_worker` function and scaled it out using Ray Train `JaxTrainer` with GPUs or TPUs!
doc/source/train/examples/lightning/BUILD.bazel
+33
View File
@@ -0,0 +1,33 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "lightning_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/lightning:lightning_examples"],
exclude = [
"dolly_lightning_fsdp_finetuning.ipynb", # Release Test
"vicuna_13b_lightning_deepspeed_finetune.ipynb", # Release Test
],
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
filegroup(
name = "lightning_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/lightning/dolly_lightning_fsdp_finetuning.ipynb
+842
View File
@@ -0,0 +1,842 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Fine-tune `dolly-v2-7b` with Ray Train, PyTorch Lightning and FSDP\n",
"\n",
"<a id=\"try-anyscale-quickstart-dolly_lightning_fsdp_finetuning\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=dolly_lightning_fsdp_finetuning\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this example, we demonstrate how to use Ray Train to fine-tune a [`dolly-v2-7b`](https://huggingface.co/databricks/dolly-v2-7b) model. `dolly-v2-7b` is a 7 billion parameter causal language model created by Databricks, derived from EleutherAI’s [Pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b), and fine-tuned on a [~15K record instruction corpus](https://github.com/databrickslabs/dolly/tree/master/data).\n",
"\n",
"We load the pre-trained model from the HuggingFace model hub into a LightningModule and launch an FSDP fine-tuning job across 16 T4 GPUs with the help of {class}`Ray TorchTrainer <ray.train.torch.TorchTrainer>`. It is also straightforward to fine-tune other similar large language models in a similar manner as shown in this example.\n",
"\n",
"Before starting this example, we highly recommend reading [Ray Train Key Concepts](train-key-concepts) and [Ray Data Quickstart](data_quickstart)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this , we demonstrate how to use Ray Train to fine-tune a [`dolly-v2-7b`](https://huggingface.co/databricks/dolly-v2-7b) model. `dolly-v2-7b` is a 7 billion parameter causal language model created by Databricks, derived from EleutherAI’s [Pythia-6.9b](https://huggingface.co/EleutherAI/pythia-6.9b), and fine-tuned on a [~15K record instruction corpus](https://github.com/databrickslabs/dolly/tree/master/data).\n",
"\n",
"We load the pre-trained model from the HuggingFace model hub into a LightningModule and launch an FSDP fine-tuning job across 16 T4 GPUs with the help of {class}`Ray TorchTrainer <ray.train.torch.TorchTrainer>`. It is also straightforward to fine-tune other similar large language models in a similar manner as shown in this example.\n",
"\n",
"Before starting this example, we highly recommend reading [Ray Train Key Concepts](train-key-concepts) and [Ray Data Quickstart](data_quickstart)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Set up ray cluster \n",
"In this example, we are using a Ray cluster with a `m5.2xlarge` head node and 4 `g4dn.12xlarge` worker nodes. Each `g4dn.12xlarge has four Tesla T4 GPUs. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"ray.init()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We then install the necessary dependencies on each node:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"pip install datasets\n",
"pip install evaluate\n",
"pip install \"transformers>=4.26.0\"\n",
"pip install \"torch>=1.12.0\"\n",
"pip install \"lightning>=2.0\"\n",
"pip install \"pydantic>=2,<3\""
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"MODEL_NAME = \"databricks/dolly-v2-7b\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare your data \n",
"We are using tiny_shakespeare for fine-tuning, which contains 40,000 lines of Shakespeare from a variety of Shakespeare's plays. Featured in Andrej Karpathy's blog post ['The Unreasonable Effectiveness of Recurrent Neural Networks'](http://karpathy.github.io/2015/05/21/rnn-effectiveness/). \n",
"\n",
"Dataset samples:\n",
"```\n",
"BAPTISTA:\n",
"I know him well: you are welcome for his sake.\n",
"\n",
"GREMIO:\n",
"Saving your tale, Petruchio, I pray,\n",
"Let us, that are poor petitioners, speak too:\n",
"Baccare! you are marvellous forward.\n",
"\n",
"PETRUCHIO:\n",
"O, pardon me, Signior Gremio; I would fain be doing.\n",
"```\n",
"\n",
"Here, we have adopted similar pre-processing logic from another demo: {doc}`GPT-J-6B Fine-Tuning with Ray Train and DeepSpeed <../deepspeed/gptj_deepspeed_fine_tuning>`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import ray\n",
"import pandas as pd\n",
"from datasets import load_dataset\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
"\n",
"def split_text(batch: pd.DataFrame) -> pd.DataFrame:\n",
" text = list(batch[\"text\"])\n",
" flat_text = \"\".join(text)\n",
" split_text = [\n",
" x.strip()\n",
" for x in flat_text.split(\"\\n\")\n",
" if x.strip() and not x.strip()[-1] == \":\"\n",
" ]\n",
" return pd.DataFrame(split_text, columns=[\"text\"])\n",
"\n",
"\n",
"def tokenize(batch: pd.DataFrame) -> dict:\n",
" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side=\"left\")\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
" ret = tokenizer(\n",
" list(batch[\"text\"]),\n",
" truncation=True,\n",
" max_length=256,\n",
" padding=\"max_length\",\n",
" return_tensors=\"np\",\n",
" )\n",
" ret[\"labels\"] = ret[\"input_ids\"].copy()\n",
" return dict(ret)\n",
"\n",
"hf_dataset = load_dataset(\"tiny_shakespeare\", trust_remote_code=True)\n",
"train_ds = ray.data.from_huggingface(hf_dataset[\"train\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We first split the original paragraphs into multiple sentences, then tokenize them. Here are some samples:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# First split the dataset into multiple sentences.\n",
"train_ds = train_ds.map_batches(split_text, batch_format=\"pandas\", batch_size=\"auto\")\n",
"train_ds.take(10)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Then tokenize the dataset.\n",
"train_ds = train_ds.map_batches(tokenize, batch_format=\"pandas\", batch_size=\"auto\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define your lightning model\n",
"\n",
"In this example, we use the [dolly-v2-7b](https://huggingface.co/databricks/dolly-v2-7b) model for finetuning. It is an instruction-following large language model trained on the Databricks machine learning platform that is licensed for commercial use. We load the model weights from Huggingface Model Hub and encapsulate it into a `pl.LightningModule`.\n",
"\n",
":::{note}\n",
"Make sure you pass the FSDP wrapped model parameters `self.trainer.model.parameters()` into the optimizer, instead of `self.model.parameters()`. \n",
":::\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import torch\n",
"import lightning.pytorch as pl\n",
"\n",
"class DollyV2Model(pl.LightningModule):\n",
" def __init__(self, lr=2e-5, eps=1e-8):\n",
" super().__init__()\n",
" self.save_hyperparameters()\n",
" self.lr = lr\n",
" self.eps = eps\n",
" self.model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)\n",
"\n",
" def forward(self, batch):\n",
" outputs = self.model(\n",
" batch[\"input_ids\"], \n",
" attention_mask=batch[\"attention_mask\"], \n",
" labels=batch[\"labels\"]\n",
" )\n",
" return outputs.loss\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" loss = self.forward(batch)\n",
" self.log(\"train_loss\", loss, prog_bar=True, on_step=True)\n",
" return loss\n",
"\n",
" def configure_optimizers(self):\n",
" if self.global_rank == 0:\n",
" print(self.trainer.model)\n",
" return torch.optim.AdamW(self.trainer.model.parameters(), lr=self.lr, eps=self.eps)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Configure your FSDP strategy\n",
"As `dolly-v2-7b` is a relatively large model, it cannot be properly fit into a single commercial GPU. In this example, we use the FSDP strategy to shard model parameters across multiple workers. This allows us to avoid GPU out-of-memory issues and support a larger global batch size.\n",
"\n",
"![](https://user-images.githubusercontent.com/26745457/236892936-d4b91751-4689-421e-ac5f-edfd2eeeb635.png)\n",
"Image source: [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"\n",
":::{note}\n",
"FSDP is a type of data parallelism that shards model parameters, optimizer states and gradients across DDP ranks. This was inspired by Xu et al. as well as the ZeRO Stage 3 from DeepSpeed. You may refer to these blogs for more information:\n",
"\n",
"- [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"- [Getting Started with Fully Sharded Data Parallel(FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html#:~:text=FSDP%20is%20a%20type%20of,sizes%20for%20our%20training%20job.)\n",
"- [PyTorch FSDP Tutorial](https://www.youtube.com/watch?v=8_k76AHu__s&list=PL_lsbAsL_o2BT6aerEKgIoufVD_fodnuT)\n",
":::\n",
"\n",
"To start training with Lightning's [FSDPStrategy](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html#lightning.pytorch.strategies.FSDPStrategy), you only need to create a {class}`~ray.train.lightning.RayFSDPStrategy` with the same initialization arguments.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import functools\n",
"import lightning.pytorch as pl \n",
"\n",
"from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy\n",
"from torch.distributed.fsdp import ShardingStrategy, BackwardPrefetch\n",
"from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXLayer\n",
"\n",
"from ray.train.lightning import RayFSDPStrategy\n",
"\n",
"\n",
"# Define the model sharding policy:\n",
"# Wrap every GPTNeoXLayer as its own FSDP instance\n",
"auto_wrap_policy = functools.partial(\n",
" transformer_auto_wrap_policy,\n",
" transformer_layer_cls = {GPTNeoXLayer}\n",
")\n",
"\n",
"fsdp_strategy = RayFSDPStrategy(\n",
" sharding_strategy=ShardingStrategy.FULL_SHARD,\n",
" backward_prefetch=BackwardPrefetch.BACKWARD_PRE,\n",
" forward_prefetch=True,\n",
" auto_wrap_policy=auto_wrap_policy,\n",
" limit_all_gathers=True,\n",
" activation_checkpointing=[GPTNeoXLayer],\n",
" cpu_offload=True\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
":::{tip}\n",
"\n",
"Some tips for FSDP configuration:\n",
"- `sharding_strategy`:\n",
" - `ShardingStrategy.NO_SHARD`: Parameters, gradients, and optimizer states are not sharded. Similar to DDP.\n",
" - `ShardingStrategy.SHARD_GRAD_OP`: Gradients and optimizer states are sharded during computation, and additionally, parameters are sharded outside computation. Similar to ZeRO stage-2.\n",
" - `ShardingStrategy.FULL_SHARD`: Parameters, gradients, and optimizer states are sharded. It has minimal GRAM usage among the 3 options. Similar to ZeRO stage-3.\n",
"- `auto_wrap_policy`:\n",
" - Model layers are often wrapped with FSDP in a layered fashion. This means that only the layers in a single FSDP instance are required to aggregate all parameters to a single device during forwarding or backward calculations.\n",
" - Use `transformer_auto_wrap_policy` to automatically wrap each Transformer Block into a single FSDP instance. \n",
"- `backward_prefetch` and `forward_prefetch`:\n",
" - Overlap the upcoming all-gather while executing the current forward/backward pass. It can improve throughput but may slightly increase peak memory usage.\n",
":::"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Fine-tune with Ray TorchTrainer\n",
"\n",
"Ray TorchTrainer allows you to scale your PyTorch Lightning training workload over multiple nodes. See {ref}`train_scaling_config` for more details."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"num_workers = 16\n",
"batch_size_per_worker = 5"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# To accelerate release tests\n",
"train_ds = train_ds.limit(num_workers * batch_size_per_worker * 5) # each worker has 5 batches"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, remember to define a Lightning callback that saves and reports checkpoints. Ray Train offers a simple implementation, {meth}`~ray.train.lightning.RayTrainReportCallback`, which persists your checkpoint and metrics in remote storage at the end of each training epoch. \n",
"\n",
"Note you can also implement your own report callback with customized logics, such as saving customized checkpoint files or reporting at a different frequency."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"from lightning.pytorch.callbacks import TQDMProgressBar\n",
"\n",
"# Create a customized progress bar for Ray Data Iterable Dataset\n",
"class DollyV2ProgressBar(TQDMProgressBar):\n",
" def __init__(self, num_iters_per_epoch, *args, **kwargs):\n",
" super().__init__(*args, **kwargs)\n",
" self.num_iters_per_epoch = num_iters_per_epoch\n",
" \n",
" def on_train_epoch_start(self, trainer, *_):\n",
" super().on_train_epoch_start(trainer, *_)\n",
" self.train_progress_bar.reset(self.num_iters_per_epoch)\n",
"\n",
"num_iters_per_epoch = train_ds.count() // (num_workers * batch_size_per_worker)\n",
"prog_bar = DollyV2ProgressBar(num_iters_per_epoch)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"from ray.train import Checkpoint\n",
"from ray.train.lightning import RayLightningEnvironment, RayTrainReportCallback, prepare_trainer\n",
"\n",
"# Training function for each worker\n",
"def train_func(config):\n",
" lr = config[\"lr\"]\n",
" eps = config[\"eps\"]\n",
" strategy = config[\"strategy\"]\n",
" batch_size_per_worker = config[\"batch_size_per_worker\"]\n",
"\n",
" # Model\n",
" model = DollyV2Model(lr=lr, eps=eps)\n",
"\n",
" # Ray Data Ingestion\n",
" train_ds = ray.train.get_dataset_shard(\"train\")\n",
" train_dataloader = train_ds.iter_torch_batches(batch_size=batch_size_per_worker)\n",
"\n",
" # Lightning Trainer\n",
" trainer = pl.Trainer(\n",
" max_epochs=1, \n",
" devices=\"auto\",\n",
" accelerator=\"auto\", \n",
" precision=\"16-mixed\",\n",
" strategy=strategy,\n",
" plugins=[RayLightningEnvironment()],\n",
" callbacks=[RayTrainReportCallback()],\n",
" enable_checkpointing=False,\n",
" )\n",
"\n",
" trainer = prepare_trainer(trainer)\n",
"\n",
" trainer.fit(model, train_dataloaders=train_dataloader)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```{note}\n",
"Since this example runs with multiple nodes, we need to persist checkpoints\n",
"and other outputs to some external storage for access after training has completed.\n",
"**You should set up cloud storage or NFS, then replace `storage_path` with your own cloud bucket URI or NFS path.**\n",
"\n",
"See the [storage guide](tune-storage-options) for more details.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"storage_path=\"s3://your-bucket-here\" # TODO: Set up cloud storage\n",
"# storage_path=\"/mnt/path/to/nfs\" # TODO: Alternatively, set up NFS"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"storage_path = \"/mnt/cluster_storage\""
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(TrainController pid=6878)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m [State Transition] INITIALIZING -> SCHEDULING.\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m Attempting to start training worker group of size 16 with the following resources: [{'GPU': 1}] * 16\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4350, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(RayTrainWorker pid=4350, ip=10.0.157.249)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Setting up process group for: env:// [rank=0, world_size=16]\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
"\u001b[36m(RayTrainWorker pid=4349, ip=10.0.157.249)\u001b[0m warnings.warn(\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m Started training worker group of size 16: \n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4088) world_rank=0, local_rank=0, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4089) world_rank=1, local_rank=1, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4090) world_rank=2, local_rank=2, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.163.141, pid=4091) world_rank=3, local_rank=3, node_rank=0\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4338) world_rank=4, local_rank=0, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4337) world_rank=5, local_rank=1, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4340) world_rank=6, local_rank=2, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.166.248, pid=4339) world_rank=7, local_rank=3, node_rank=1\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4090) world_rank=8, local_rank=0, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4248) world_rank=9, local_rank=1, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4246) world_rank=10, local_rank=2, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.191.43, pid=4247) world_rank=11, local_rank=3, node_rank=2\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4349) world_rank=12, local_rank=0, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4350) world_rank=13, local_rank=1, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4347) world_rank=14, local_rank=2, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m - (ip=10.0.157.249, pid=4348) world_rank=15, local_rank=3, node_rank=3\n",
"\u001b[36m(TrainController pid=6878)\u001b[0m [State Transition] SCHEDULING -> RUNNING.\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\u001b[32m [repeated 31x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m _torch_pytree._register_pytree_node(\u001b[32m [repeated 31x across cluster]\u001b[0m\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[33m(raylet)\u001b[0m WARNING: 4 PYTHON worker processes have been started on node: 3ffdd02cd1be6b69a64a97e08f75fc0a80eddcf0caa627f3f4266c95 with address: 10.0.150.21. This could be a result of using a large number of actors, or due to tasks blocked in ray.get() calls (see https://github.com/ray-project/ray/issues/3644 for some discussion of workarounds).\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m GPU available: True (cuda), used: True\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m TPU available: False, using: 0 TPU cores\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m HPU available: False, using: 0 HPUs\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m warnings.warn(\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:07.970624: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.208262: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.208291: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.231782: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:08.277889: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 2025-09-30 14:20:10.134645: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m FullyShardedDataParallel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_fsdp_wrapped_module): DollyV2Model(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (model): GPTNeoXForCausalLM(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (gpt_neox): GPTNeoXModel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (embed_in): Embedding(50280, 4096)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (emb_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (layers): ModuleList(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (0-31): 32 x FullyShardedDataParallel(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_fsdp_wrapped_module): CheckpointWrapper(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (_checkpoint_wrapped_module): GPTNeoXLayer(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (input_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_attention_layernorm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_attention_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (post_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (attention): GPTNeoXAttention(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (rotary_emb): GPTNeoXRotaryEmbedding()\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (query_key_value): Linear(in_features=4096, out_features=12288, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense): Linear(in_features=4096, out_features=4096, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (attention_dropout): Dropout(p=0.0, inplace=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (mlp): GPTNeoXMLP(\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense_h_to_4h): Linear(in_features=4096, out_features=16384, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (dense_4h_to_h): Linear(in_features=16384, out_features=4096, bias=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (act): GELUActivation()\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (final_layer_norm): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m (embed_out): Linear(in_features=4096, out_features=50280, bias=False)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m )\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m /tmp/ray/session_2025-09-30_14-10-46_627006_2395/runtime_resources/pip/72a6e451f55d87eb50ebbf5bc30a4a57ed513d34/virtualenv/lib/python3.10/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py:231: Precision 16-mixed is not supported by the model summary. Estimated model size in MB will not be accurate. Using 32 bits instead.\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m \n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m | Name | Type | Params | Mode\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m ----------------------------------------------------\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 0 | model | GPTNeoXForCausalLM | 428 M | eval\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m ----------------------------------------------------\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 428 M Trainable params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 0 Non-trainable params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 428 M Total params\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 1,714.014 Total estimated model params size (MB)\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 64 Modules in train mode\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m 455 Modules in eval mode\n",
"\u001b[36m(RayTrainWorker pid=4246, ip=10.0.191.43)\u001b[0m LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Registered dataset logger for dataset train_12_0\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Starting execution of Dataset train_12_0. Full logs are in /tmp/ray/session_2025-09-30_14-10-46_627006_2395/logs/ray-data\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m Execution plan of Dataset train_12_0: InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(split_text)->MapBatches(tokenize)] -> LimitOperator[limit=400] -> OutputSplitter[split(16, equal=True)]\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m ⚠️ Ray's object store is configured to use only 28.7% of available memory (229.3GiB out of 800.0GiB total). For optimal Ray Data performance, we recommend setting the object store to at least 50% of available memory. You can do this by setting the 'object_store_memory' parameter when calling ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable.\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ab803bffe2224e6591bc452cac07f2a8",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) Running 0: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8d647925625f49c19e16fddce0fab359",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - MapBatches(split_text)->MapBatches(tokenize) 1: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "30978a8d911a44dcb168b0d5a386a42c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - limit=400 2: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "432c8cb6d3e84b749d9341ff104bb25c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"(pid=7661) - split(16, equal=True) 3: 0.00 row [00:00, ? row/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m _torch_pytree._register_pytree_node(\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m warnings.warn(\n",
"\u001b[36m(MapBatches(split_text)->MapBatches(tokenize) pid=4089, ip=10.0.191.43)\u001b[0m Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m /home/ray/anaconda3/lib/python3.10/site-packages/ray/anyscale/data/_internal/cluster_autoscaler/productivity_calculator.py:174: RuntimeWarning: invalid value encountered in divide\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m gpu_fraction_per_op = (optimal_num_tasks_per_op * num_gpus_per_op) / np.sum(\n",
"\u001b[36m(SplitCoordinator pid=7661)\u001b[0m ✔️ Dataset train_12_0 execution finished in 5.10 seconds\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m /tmp/ray/session_2025-09-30_14-10-46_627006_2395/runtime_resources/pip/72a6e451f55d87eb50ebbf5bc30a4a57ed513d34/virtualenv/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py:527: Found 455 module(s) in eval mode at the start of training. This may lead to unexpected behavior during training. If this is intentional, you can ignore this warning.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: | | 0/? [00:00<?, ?it/s] 141)\u001b[0m \n",
"Epoch 0: | | 1/? [00:20<00:00, 0.05it/s, v_num=0, train_loss=14.60]\n",
"Epoch 0: | | 2/? [00:36<00:00, 0.06it/s, v_num=0, train_loss=13.90]\n",
"Epoch 0: | | 3/? [00:52<00:00, 0.06it/s, v_num=0, train_loss=14.80]\n",
"Epoch 0: | | 4/? [01:09<00:00, 0.06it/s, v_num=0, train_loss=14.10]\n",
"Epoch 0: | | 5/? [01:27<00:00, 0.06it/s, v_num=0, train_loss=14.90]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[36m(RayTrainWorker pid=4090, ip=10.0.191.43)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076)\n",
"\u001b[36m(RayTrainWorker pid=4090, ip=10.0.191.43)\u001b[0m Reporting training result 1: TrainingResult(checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), metrics={'train_loss': 14.9804105758667, 'epoch': 0, 'step': 5})\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076)\u001b[32m [repeated 12x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m Reporting training result 1: TrainingResult(checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), metrics={'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5})\u001b[32m [repeated 12x across cluster]\u001b[0m\n",
"\u001b[36m(RayTrainWorker pid=4088, ip=10.0.163.141)\u001b[0m `Trainer.fit` stopped: `max_epochs=1` reached.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: | | 5/? [06:07<00:00, 0.01it/s, v_num=0, train_loss=14.90]\n"
]
},
{
"data": {
"text/plain": [
"Result(metrics={'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5}, checkpoint=Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), error=None, path='/mnt/cluster_storage/finetune_dolly-v2-7b-trial1', metrics_dataframe= train_loss epoch step\n",
"0 14.879827 0 5, best_checkpoints=[(Checkpoint(filesystem=local, path=/mnt/cluster_storage/finetune_dolly-v2-7b-trial1/checkpoint_2025-09-30_14-23-45.998076), {'train_loss': 14.879826545715332, 'epoch': 0, 'step': 5})], _storage_filesystem=<pyarrow._fs.LocalFileSystem object at 0x737af20de5b0>)"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"\n",
"# Save Ray Train checkpoints according to the performance on validation set\n",
"run_config = RunConfig(\n",
" name=\"finetune_dolly-v2-7b-trial1\",\n",
" storage_path=storage_path,\n",
" checkpoint_config=CheckpointConfig(num_to_keep=1),\n",
")\n",
"\n",
"# Scale the FSDP training workload across 16 GPUs\n",
"# You can change this config based on your compute resources.\n",
"scaling_config = ScalingConfig(\n",
" num_workers=num_workers, use_gpu=True\n",
")\n",
"\n",
"# Configuration to pass into train_func\n",
"train_config = {\n",
" \"lr\": 2e-5,\n",
" \"eps\": 1e-8,\n",
" \"strategy\": fsdp_strategy,\n",
" \"batch_size_per_worker\": batch_size_per_worker\n",
"}\n",
"\n",
"# Define a TorchTrainer and launch you training workload\n",
"ray_trainer = TorchTrainer(\n",
" train_func,\n",
" train_loop_config=train_config,\n",
" run_config=run_config,\n",
" scaling_config=scaling_config,\n",
" datasets={\"train\": train_ds},\n",
")\n",
"result = ray_trainer.fit()\n",
"\n",
"result\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We finished training in 2877s. The price for an on-demand g4dn.4xlarge instance is `$1.204/hour`, while a g4dn.8xlarge instance costs `$2.176/hour`. The total cost would be `($1.204 * 15 + $2.176) * 2877 / 3600 = $16.17`."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Text-generation with HuggingFace Pipeline\n",
"\n",
"We can use the [HuggingFace Pipeline](https://huggingface.co/docs/transformers/main_classes/pipelines) to generate predictions from our fine-tuned model. Let's input some prompts and see if our tuned Dolly can speak like Shakespeare:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import os\n",
"from transformers import pipeline\n",
"\n",
"@ray.remote(num_gpus=1)\n",
"def generate_tokens():\n",
" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side=\"right\")\n",
"\n",
" ckpt_path = os.path.join(result.checkpoint.path, \"checkpoint.ckpt\")\n",
"\n",
" dolly = DollyV2Model.load_from_checkpoint(ckpt_path, map_location=torch.device(\"cpu\"))\n",
"\n",
" nlp_pipeline = pipeline(\n",
" task=\"text-generation\", \n",
" model=dolly.model, \n",
" tokenizer=tokenizer, \n",
" device_map=\"auto\"\n",
" )\n",
"\n",
" tokens = []\n",
" for prompt in [\"This is\", \"I am\", \"Once more\"]:\n",
" tokens.append(nlp_pipeline(prompt, max_new_tokens=20, do_sample=True, pad_token_id=tokenizer.eos_token_id))\n",
"\n",
" return tokens\n",
"\n",
"ref = generate_tokens.remote()\n",
"output = ray.get(ref)"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'generated_text': \"This is more like it:\\n\\nIt's just a guess, but maybe the extra processing power of Intel\"}]\n",
"[{'generated_text': \"I am the biggest fan of your wife's writing, and this novella was fantastic. So interesting to see\"}]\n",
"[{'generated_text': 'Once more I wish I could make sense of it.\" \"My friend, you can leave all this behind you'}]\n"
]
}
],
"source": [
"for generated_tokens in output:\n",
" print(generated_tokens)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"References:\n",
"- [PyTorch FSDP Tutorial](https://www.youtube.com/watch?v=8_k76AHu__s&list=PL_lsbAsL_o2BT6aerEKgIoufVD_fodnuT)\n",
"- [Getting Started with Fully Sharded Data Parallel(FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html#:~:text=FSDP%20is%20a%20type%20of,sizes%20for%20our%20training%20job.)\n",
"- [Fully Sharded Data Parallel: faster AI training with fewer GPUs](https://engineering.fb.com/2021/07/15/open-source/fsdp/)\n",
"- [Hugging Face: dolly-v2-7b Model Card](https://huggingface.co/databricks/dolly-v2-7b)\n",
"- [Hugging Face: Handling big models for inference](https://huggingface.co/docs/accelerate/usage_guides/big_modeling)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.18"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/lightning/lightning_cola_advanced.ipynb
+1186
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/lightning/lightning_mnist_example.ipynb
+947
View File
@@ -0,0 +1,947 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Train a Pytorch Lightning Image Classifier\n",
"\n",
"<a id=\"try-anyscale-quickstart-lightning_mnist_example\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=lightning_mnist_example\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example introduces how to train a Pytorch Lightning Module using Ray Train {class}`TorchTrainer <ray.train.torch.TorchTrainer>`. It demonstrates how to train a basic neural network on the MNIST dataset with distributed data parallelism.\n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"!pip install \"torchmetrics>=0.9\" \"lightning>=2.0\""
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import numpy as np\n",
"import random\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"from filelock import FileLock\n",
"from torch.utils.data import DataLoader, random_split, Subset\n",
"from torchmetrics import Accuracy\n",
"from torchvision.datasets import MNIST\n",
"from torchvision import transforms\n",
"\n",
"try:\n",
" import lightning.pytorch as pl\n",
" from lightning.pytorch.loggers import CSVLogger\n",
"except ModuleNotFoundError:\n",
" import pytorch_lightning as pl\n",
" from pytorch_lightning.loggers import CSVLogger"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare a dataset and module\n",
"\n",
"The Pytorch Lightning Trainer takes either `torch.utils.data.DataLoader` or `pl.LightningDataModule` as data inputs. You can continue using them without any changes with Ray Train. "
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"class MNISTDataModule(pl.LightningDataModule):\n",
" def __init__(self, batch_size=100):\n",
" super().__init__()\n",
" self.data_dir = os.getcwd()\n",
" self.batch_size = batch_size\n",
" self.transform = transforms.Compose(\n",
" [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
" )\n",
"\n",
" def setup(self, stage=None):\n",
" with FileLock(f\"{self.data_dir}.lock\"):\n",
" mnist = MNIST(\n",
" self.data_dir, train=True, download=True, transform=self.transform\n",
" )\n",
"\n",
" # Split data into train and val sets\n",
" self.mnist_train, self.mnist_val = random_split(mnist, [55000, 5000])\n",
"\n",
" def train_dataloader(self):\n",
" return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=4)\n",
"\n",
" def val_dataloader(self):\n",
" return DataLoader(self.mnist_val, batch_size=self.batch_size, num_workers=4)\n",
"\n",
" def test_dataloader(self):\n",
" with FileLock(f\"{self.data_dir}.lock\"):\n",
" self.mnist_test = MNIST(\n",
" self.data_dir, train=False, download=True, transform=self.transform\n",
" )\n",
" return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=4)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, define a simple multi-layer perception as the subclass of `pl.LightningModule`."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"class MNISTClassifier(pl.LightningModule):\n",
" def __init__(self, lr=1e-3, feature_dim=128):\n",
" torch.manual_seed(421)\n",
" super(MNISTClassifier, self).__init__()\n",
" self.save_hyperparameters()\n",
"\n",
" self.linear_relu_stack = nn.Sequential(\n",
" nn.Linear(28 * 28, feature_dim),\n",
" nn.ReLU(),\n",
" nn.Linear(feature_dim, 10),\n",
" nn.ReLU(),\n",
" )\n",
" self.lr = lr\n",
" self.accuracy = Accuracy(task=\"multiclass\", num_classes=10, top_k=1)\n",
" self.eval_loss = []\n",
" self.eval_accuracy = []\n",
" self.test_accuracy = []\n",
" pl.seed_everything(888)\n",
"\n",
" def forward(self, x):\n",
" x = x.view(-1, 28 * 28)\n",
" x = self.linear_relu_stack(x)\n",
" return x\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" x, y = batch\n",
" y_hat = self(x)\n",
" loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
" self.log(\"train_loss\", loss)\n",
" return loss\n",
"\n",
" def validation_step(self, val_batch, batch_idx):\n",
" loss, acc = self._shared_eval(val_batch)\n",
" self.log(\"val_accuracy\", acc)\n",
" self.eval_loss.append(loss)\n",
" self.eval_accuracy.append(acc)\n",
" return {\"val_loss\": loss, \"val_accuracy\": acc}\n",
"\n",
" def test_step(self, test_batch, batch_idx):\n",
" loss, acc = self._shared_eval(test_batch)\n",
" self.test_accuracy.append(acc)\n",
" self.log(\"test_accuracy\", acc, sync_dist=True, on_epoch=True)\n",
" return {\"test_loss\": loss, \"test_accuracy\": acc}\n",
"\n",
" def _shared_eval(self, batch):\n",
" x, y = batch\n",
" logits = self.forward(x)\n",
" loss = F.nll_loss(logits, y)\n",
" acc = self.accuracy(logits, y)\n",
" return loss, acc\n",
"\n",
" def on_validation_epoch_end(self):\n",
" avg_loss = torch.stack(self.eval_loss).mean()\n",
" avg_acc = torch.stack(self.eval_accuracy).mean()\n",
" self.log(\"val_loss\", avg_loss, sync_dist=True)\n",
" self.log(\"val_accuracy\", avg_acc, sync_dist=True)\n",
" self.eval_loss.clear()\n",
" self.eval_accuracy.clear()\n",
" \n",
" def configure_optimizers(self):\n",
" optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)\n",
" return optimizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You don't need to modify the definition of the PyTorch Lightning model or datamodule."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define a training function\n",
"\n",
"This code defines a {ref}`training function <train-overview-training-function>` for each worker. Comparing the training function with the original PyTorch Lightning code, notice three main differences:\n",
"\n",
"- Distributed strategy: Use {class}`RayDDPStrategy <ray.train.lightning.RayDDPStrategy>`.\n",
"- Cluster environment: Use {class}`RayLightningEnvironment <ray.train.lightning.RayLightningEnvironment>`.\n",
"- Parallel devices: Always set to `devices=\"auto\"` to use all available devices configured by ``TorchTrainer``.\n",
"\n",
"See {ref}`Getting Started with PyTorch Lightning <train-pytorch-lightning>` for more information.\n",
"\n",
"\n",
"For checkpoint reporting, Ray Train provides a minimal {class}`RayTrainReportCallback <ray.train.lightning.RayTrainReportCallback>` class that reports metrics and checkpoints at the end of each train epoch. For more complex checkpoint logic, implement custom callbacks. See {ref}`Saving and Loading Checkpoint <train-checkpointing>`."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"use_gpu = True # Set to False if you want to run without GPUs\n",
"num_workers = 4"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"try:\n",
" import lightning.pytorch as pl\n",
"except ModuleNotFoundError:\n",
" import pytorch_lightning as pl\n",
"\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.lightning import (\n",
" RayDDPStrategy,\n",
" RayLightningEnvironment,\n",
" RayTrainReportCallback,\n",
" prepare_trainer,\n",
")\n",
"\n",
"def train_func_per_worker():\n",
" model = MNISTClassifier(lr=1e-3, feature_dim=128)\n",
" datamodule = MNISTDataModule(batch_size=128)\n",
"\n",
" trainer = pl.Trainer(\n",
" devices=\"auto\",\n",
" strategy=RayDDPStrategy(),\n",
" plugins=[RayLightningEnvironment()],\n",
" callbacks=[RayTrainReportCallback()],\n",
" max_epochs=10,\n",
" accelerator=\"gpu\" if use_gpu else \"cpu\",\n",
" log_every_n_steps=100,\n",
" logger=CSVLogger(\"logs\"),\n",
" )\n",
" \n",
" trainer = prepare_trainer(trainer)\n",
" \n",
" # Train model\n",
" trainer.fit(model, datamodule=datamodule)\n",
"\n",
" # Evaluation on the test dataset\n",
" trainer.test(model, datamodule=datamodule)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now put everything together:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)\n",
"\n",
"run_config = RunConfig(\n",
" name=\"ptl-mnist-example\",\n",
" storage_path=\"/tmp/ray_results\",\n",
" checkpoint_config=CheckpointConfig(\n",
" num_to_keep=3,\n",
" checkpoint_score_attribute=\"val_accuracy\",\n",
" checkpoint_score_order=\"max\",\n",
" ),\n",
")\n",
"\n",
"trainer = TorchTrainer(\n",
" train_func_per_worker,\n",
" scaling_config=scaling_config,\n",
" run_config=run_config,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now fit your trainer:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div class=\"tuneStatus\">\n",
" <div style=\"display: flex;flex-direction: row\">\n",
" <div style=\"display: flex;flex-direction: column;\">\n",
" <h3>Tune Status</h3>\n",
" <table>\n",
"<tbody>\n",
"<tr><td>Current time:</td><td>2023-08-07 23:41:11</td></tr>\n",
"<tr><td>Running for: </td><td>00:00:39.80 </td></tr>\n",
"<tr><td>Memory: </td><td>24.2/186.6 GiB </td></tr>\n",
"</tbody>\n",
"</table>\n",
" </div>\n",
" <div class=\"vDivider\"></div>\n",
" <div class=\"systemInfo\">\n",
" <h3>System Info</h3>\n",
" Using FIFO scheduling algorithm.<br>Logical resource usage: 1.0/48 CPUs, 4.0/4 GPUs\n",
" </div>\n",
" \n",
" </div>\n",
" <div class=\"hDivider\"></div>\n",
" <div class=\"trialStatus\">\n",
" <h3>Trial Status</h3>\n",
" <table>\n",
"<thead>\n",
"<tr><th>Trial name </th><th>status </th><th>loc </th><th style=\"text-align: right;\"> iter</th><th style=\"text-align: right;\"> total time (s)</th><th style=\"text-align: right;\"> train_loss</th><th style=\"text-align: right;\"> val_accuracy</th><th style=\"text-align: right;\"> val_loss</th></tr>\n",
"</thead>\n",
"<tbody>\n",
"<tr><td>TorchTrainer_78346_00000</td><td>TERMINATED</td><td>10.0.6.244:120026</td><td style=\"text-align: right;\"> 10</td><td style=\"text-align: right;\"> 29.0221</td><td style=\"text-align: right;\"> 0.0315938</td><td style=\"text-align: right;\"> 0.970002</td><td style=\"text-align: right;\"> -12.3466</td></tr>\n",
"</tbody>\n",
"</table>\n",
" </div>\n",
"</div>\n",
"<style>\n",
".tuneStatus {\n",
" color: var(--jp-ui-font-color1);\n",
"}\n",
".tuneStatus .systemInfo {\n",
" display: flex;\n",
" flex-direction: column;\n",
"}\n",
".tuneStatus td {\n",
" white-space: nowrap;\n",
"}\n",
".tuneStatus .trialStatus {\n",
" display: flex;\n",
" flex-direction: column;\n",
"}\n",
".tuneStatus h3 {\n",
" font-weight: bold;\n",
"}\n",
".tuneStatus .hDivider {\n",
" border-bottom-width: var(--jp-border-width);\n",
" border-bottom-color: var(--jp-border-color0);\n",
" border-bottom-style: solid;\n",
"}\n",
".tuneStatus .vDivider {\n",
" border-left-width: var(--jp-border-width);\n",
" border-left-color: var(--jp-border-color0);\n",
" border-left-style: solid;\n",
" margin: 0.5em 1em 0.5em 1em;\n",
"}\n",
"</style>\n"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(TorchTrainer pid=120026)\u001b[0m Starting distributed worker processes: ['120176 (10.0.6.244)', '120177 (10.0.6.244)', '120178 (10.0.6.244)', '120179 (10.0.6.244)']\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m Setting up process group for: env:// [rank=0, world_size=4]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m [rank: 0] Global seed set to 888\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m GPU available: True (cuda), used: True\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m TPU available: False, using: 0 TPU cores\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m IPU available: False, using: 0 IPUs\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m HPU available: False, using: 0 HPUs\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/train-images-idx3-ubyte.gz\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" 0%| | 0/9912422 [00:00<?, ?it/s]\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 94562894.32it/s]\n",
" 9%|▉ | 917504/9912422 [00:00<00:00, 9166590.91it/s]\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 115619443.32it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m Extracting /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_3/MNIST/raw/train-images-idx3-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_3/MNIST/raw\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120177)\u001b[0m Missing logger folder: logs/lightning_logs\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m | Name | Type | Params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ---------------------------------------------------------\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0 | linear_relu_stack | Sequential | 101 K \n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 1 | accuracy | MulticlassAccuracy | 0 \n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ---------------------------------------------------------\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 101 K Trainable params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0 Non-trainable params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 101 K Total params\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m 0.407 Total estimated model params size (MB)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Sanity Checking: 0it [00:00, ?it/s])\u001b[0m \n",
"Sanity Checking DataLoader 0: 0%| | 0/2 [00:00<?, ?it/s]\n",
"Sanity Checking DataLoader 0: 100%|██████████| 2/2 [00:00<00:00, 2.69it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m /mnt/cluster_storage/pypi/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/logger_connector/result.py:432: PossibleUserWarning: It is recommended to use `self.log('val_accuracy', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m warning_cache.warn(\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m [rank: 3] Global seed set to 888\u001b[32m [repeated 7x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)\u001b[0m\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 0: 0%| | 0/108 [00:00<?, ?it/s] \n",
"Epoch 0: 12%|█▏ | 13/108 [00:00<00:02, 39.35it/s, v_num=0]\n",
"Epoch 0: 25%|██▌ | 27/108 [00:00<00:01, 59.26it/s, v_num=0]\n",
"Epoch 0: 26%|██▌ | 28/108 [00:00<00:01, 61.03it/s, v_num=0]\n",
"Epoch 0: 27%|██▋ | 29/108 [00:00<00:01, 62.76it/s, v_num=0]\n",
"Epoch 0: 42%|████▏ | 45/108 [00:00<00:00, 81.02it/s, v_num=0]\n",
"Epoch 0: 53%|█████▎ | 57/108 [00:00<00:00, 86.01it/s, v_num=0]\n",
"Epoch 0: 64%|██████▍ | 69/108 [00:00<00:00, 88.63it/s, v_num=0]\n",
"Epoch 0: 81%|████████ | 87/108 [00:00<00:00, 98.04it/s, v_num=0]\n",
"Epoch 0: 81%|████████▏ | 88/108 [00:00<00:00, 98.69it/s, v_num=0]\n",
"Epoch 0: 82%|████████▏ | 89/108 [00:00<00:00, 99.34it/s, v_num=0]\n",
"Epoch 0: 96%|█████████▋| 104/108 [00:00<00:00, 104.14it/s, v_num=0]\n",
"Epoch 0: 97%|█████████▋| 105/108 [00:01<00:00, 104.71it/s, v_num=0]\n",
"Epoch 0: 98%|█████████▊| 106/108 [00:01<00:00, 105.22it/s, v_num=0]\n",
"Epoch 0: 100%|██████████| 108/108 [00:01<00:00, 105.79it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 171.69it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 200.99it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 221.66it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 215.50it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 194.14it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 205.63it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 215.27it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 216.26it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 198.67it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 205.79it/s]\u001b[A\n",
"Epoch 0: 100%|██████████| 108/108 [00:01<00:00, 79.84it/s, v_num=0] \n",
" \u001b[A\n",
"Epoch 1: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 1: 11%|█ | 12/108 [00:00<00:02, 32.36it/s, v_num=0]\n",
"Epoch 1: 23%|██▎ | 25/108 [00:00<00:01, 50.16it/s, v_num=0]\n",
"Epoch 1: 37%|███▋ | 40/108 [00:00<00:01, 65.95it/s, v_num=0]\n",
"Epoch 1: 38%|███▊ | 41/108 [00:00<00:00, 67.05it/s, v_num=0]\n",
"Epoch 1: 50%|█████ | 54/108 [00:00<00:00, 75.52it/s, v_num=0]\n",
"Epoch 1: 51%|█████ | 55/108 [00:00<00:00, 76.40it/s, v_num=0]\n",
"Epoch 1: 62%|██████▏ | 67/108 [00:00<00:00, 81.72it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 77%|███████▋ | 83/108 [00:00<00:00, 89.48it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/t10k-labels-idx1-ubyte.gz\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 78%|███████▊ | 84/108 [00:00<00:00, 89.21it/s, v_num=0]\n",
"Epoch 1: 91%|█████████ | 98/108 [00:01<00:00, 93.27it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Extracting /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw/t10k-labels-idx1-ubyte.gz to /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/rank_2/MNIST/raw\u001b[32m [repeated 15x across cluster]\u001b[0m\n",
"Epoch 1: 92%|█████████▏| 99/108 [00:01<00:00, 93.94it/s, v_num=0]\n",
"Epoch 1: 93%|█████████▎| 100/108 [00:01<00:00, 94.57it/s, v_num=0]\n",
"Epoch 1: 100%|██████████| 108/108 [00:01<00:00, 98.06it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 320.27it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 291.99it/s]\u001b[A\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 19x across cluster]\u001b[0m\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 291.61it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 268.90it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 290.07it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 293.52it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 299.70it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 304.80it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 310.16it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 303.63it/s]\u001b[A\n",
"Epoch 1: 100%|██████████| 108/108 [00:01<00:00, 76.12it/s, v_num=0]\n",
" \u001b[A\n",
"Epoch 2: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 2: 7%|▋ | 8/108 [00:00<00:03, 25.23it/s, v_num=0]\n",
"Epoch 2: 16%|█▌ | 17/108 [00:00<00:02, 39.73it/s, v_num=0]\n",
"Epoch 2: 17%|█▋ | 18/108 [00:00<00:02, 41.60it/s, v_num=0]\n",
"Epoch 2: 18%|█▊ | 19/108 [00:00<00:02, 43.49it/s, v_num=0]\n",
"Epoch 2: 18%|█▊ | 19/108 [00:00<00:02, 43.46it/s, v_num=0]\n",
"Epoch 2: 19%|█▊ | 20/108 [00:00<00:01, 45.27it/s, v_num=0]\n",
"Epoch 2: 27%|██▋ | 29/108 [00:00<00:01, 53.08it/s, v_num=0]\n",
"Epoch 2: 42%|████▏ | 45/108 [00:00<00:00, 69.12it/s, v_num=0]\n",
"Epoch 2: 43%|████▎ | 46/108 [00:00<00:00, 70.31it/s, v_num=0]\n",
"Epoch 2: 44%|████▎ | 47/108 [00:00<00:00, 71.51it/s, v_num=0]\n",
"Epoch 2: 44%|████▍ | 48/108 [00:00<00:00, 72.71it/s, v_num=0]\n",
"Epoch 2: 59%|█████▉ | 64/108 [00:00<00:00, 83.97it/s, v_num=0]\n",
"Epoch 2: 75%|███████▌ | 81/108 [00:00<00:00, 93.77it/s, v_num=0]\n",
"Epoch 2: 90%|████████▉ | 97/108 [00:00<00:00, 99.35it/s, v_num=0]\n",
"Epoch 2: 100%|██████████| 108/108 [00:01<00:00, 101.71it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 212.13it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 184.45it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 228.42it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 225.00it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 250.65it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 251.36it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 268.85it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 256.15it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 269.87it/s]\u001b[A\n",
"Epoch 2: 100%|██████████| 108/108 [00:01<00:00, 77.52it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 3: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 3: 8%|▊ | 9/108 [00:00<00:03, 25.68it/s, v_num=0]\n",
"Epoch 3: 9%|▉ | 10/108 [00:00<00:03, 28.26it/s, v_num=0]\n",
"Epoch 3: 20%|██ | 22/108 [00:00<00:01, 48.10it/s, v_num=0]\n",
"Epoch 3: 21%|██▏ | 23/108 [00:00<00:01, 49.73it/s, v_num=0]\n",
"Epoch 3: 22%|██▏ | 24/108 [00:00<00:01, 51.34it/s, v_num=0]\n",
"Epoch 3: 23%|██▎ | 25/108 [00:00<00:01, 52.98it/s, v_num=0]\n",
"Epoch 3: 37%|███▋ | 40/108 [00:00<00:00, 69.67it/s, v_num=0]\n",
"Epoch 3: 51%|█████ | 55/108 [00:00<00:00, 80.93it/s, v_num=0]\n",
"Epoch 3: 64%|██████▍ | 69/108 [00:00<00:00, 87.15it/s, v_num=0]\n",
"Epoch 3: 65%|██████▍ | 70/108 [00:00<00:00, 88.04it/s, v_num=0]\n",
"Epoch 3: 66%|██████▌ | 71/108 [00:00<00:00, 88.92it/s, v_num=0]\n",
"Epoch 3: 77%|███████▋ | 83/108 [00:00<00:00, 92.62it/s, v_num=0]\n",
"Epoch 3: 86%|████████▌ | 93/108 [00:01<00:00, 92.33it/s, v_num=0]\n",
"Epoch 3: 87%|████████▋ | 94/108 [00:01<00:00, 92.93it/s, v_num=0]\n",
"Epoch 3: 88%|████████▊ | 95/108 [00:01<00:00, 93.61it/s, v_num=0]\n",
"Epoch 3: 100%|██████████| 108/108 [00:01<00:00, 97.43it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 308.50it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 344.87it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 375.98it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 335.26it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 327.34it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 317.66it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 332.79it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 188.14it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 201.21it/s]\u001b[A\n",
"Epoch 3: 100%|██████████| 108/108 [00:01<00:00, 75.94it/s, v_num=0]6it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 4: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 4: 10%|█ | 11/108 [00:00<00:03, 30.09it/s, v_num=0]\n",
"Epoch 4: 20%|██ | 22/108 [00:00<00:01, 46.96it/s, v_num=0]\n",
"Epoch 4: 21%|██▏ | 23/108 [00:00<00:01, 47.88it/s, v_num=0]\n",
"Epoch 4: 35%|███▌ | 38/108 [00:00<00:01, 65.26it/s, v_num=0]\n",
"Epoch 4: 36%|███▌ | 39/108 [00:00<00:01, 65.73it/s, v_num=0]\n",
"Epoch 4: 53%|█████▎ | 57/108 [00:00<00:00, 81.51it/s, v_num=0]\n",
"Epoch 4: 54%|█████▎ | 58/108 [00:00<00:00, 82.56it/s, v_num=0]\n",
"Epoch 4: 68%|██████▊ | 73/108 [00:00<00:00, 89.69it/s, v_num=0]\n",
"Epoch 4: 69%|██████▊ | 74/108 [00:00<00:00, 90.53it/s, v_num=0]\n",
"Epoch 4: 83%|████████▎ | 90/108 [00:00<00:00, 98.32it/s, v_num=0]\n",
"Epoch 4: 98%|█████████▊| 106/108 [00:01<00:00, 103.12it/s, v_num=0]\n",
"Epoch 4: 100%|██████████| 108/108 [00:01<00:00, 103.78it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 268.49it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 298.62it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 282.88it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 256.50it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 276.28it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 268.05it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 276.18it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 290.08it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 261.92it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 274.00it/s]\u001b[A\n",
"Epoch 4: 100%|██████████| 108/108 [00:01<00:00, 78.54it/s, v_num=0] \n",
" \u001b[A\n",
"Epoch 5: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 5: 5%|▍ | 5/108 [00:00<00:06, 15.52it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 9x across cluster]\u001b[0m\n",
"Epoch 5: 17%|█▋ | 18/108 [00:00<00:02, 42.53it/s, v_num=0]\n",
"Epoch 5: 26%|██▌ | 28/108 [00:00<00:01, 52.36it/s, v_num=0]\n",
"Epoch 5: 27%|██▋ | 29/108 [00:00<00:01, 53.91it/s, v_num=0]\n",
"Epoch 5: 28%|██▊ | 30/108 [00:00<00:01, 55.45it/s, v_num=0]\n",
"Epoch 5: 29%|██▊ | 31/108 [00:00<00:01, 56.96it/s, v_num=0]\n",
"Epoch 5: 37%|███▋ | 40/108 [00:00<00:01, 61.48it/s, v_num=0]\n",
"Epoch 5: 38%|███▊ | 41/108 [00:00<00:01, 62.61it/s, v_num=0]\n",
"Epoch 5: 39%|███▉ | 42/108 [00:00<00:01, 63.79it/s, v_num=0]\n",
"Epoch 5: 40%|███▉ | 43/108 [00:00<00:01, 64.96it/s, v_num=0]\n",
"Epoch 5: 48%|████▊ | 52/108 [00:00<00:00, 67.96it/s, v_num=0]\n",
"Epoch 5: 49%|████▉ | 53/108 [00:00<00:00, 68.88it/s, v_num=0]\n",
"Epoch 5: 50%|█████ | 54/108 [00:00<00:00, 69.77it/s, v_num=0]\n",
"Epoch 5: 62%|██████▏ | 67/108 [00:00<00:00, 76.43it/s, v_num=0]\n",
"Epoch 5: 78%|███████▊ | 84/108 [00:00<00:00, 85.56it/s, v_num=0]\n",
"Epoch 5: 79%|███████▊ | 85/108 [00:00<00:00, 86.17it/s, v_num=0]\n",
"Epoch 5: 93%|█████████▎| 100/108 [00:01<00:00, 92.27it/s, v_num=0]\n",
"Epoch 5: 100%|██████████| 108/108 [00:01<00:00, 94.81it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 255.91it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 206.50it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 214.91it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 236.35it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 240.05it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 234.90it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 240.78it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 252.00it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 255.18it/s]\u001b[A\n",
"Epoch 5: 100%|██████████| 108/108 [00:01<00:00, 72.36it/s, v_num=0]1it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 6: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 6: 15%|█▍ | 16/108 [00:00<00:02, 44.27it/s, v_num=0]\n",
"Epoch 6: 25%|██▌ | 27/108 [00:00<00:01, 57.21it/s, v_num=0]\n",
"Epoch 6: 38%|███▊ | 41/108 [00:00<00:00, 70.86it/s, v_num=0]\n",
"Epoch 6: 39%|███▉ | 42/108 [00:00<00:00, 71.82it/s, v_num=0]\n",
"Epoch 6: 55%|█████▍ | 59/108 [00:00<00:00, 85.97it/s, v_num=0]\n",
"Epoch 6: 68%|██████▊ | 73/108 [00:00<00:00, 91.53it/s, v_num=0]\n",
"Epoch 6: 81%|████████ | 87/108 [00:00<00:00, 96.88it/s, v_num=0]\n",
"Epoch 6: 92%|█████████▏| 99/108 [00:00<00:00, 99.33it/s, v_num=0]\n",
"Epoch 6: 93%|█████████▎| 100/108 [00:01<00:00, 98.66it/s, v_num=0]\n",
"Epoch 6: 94%|█████████▎| 101/108 [00:01<00:00, 99.34it/s, v_num=0]\n",
"Epoch 6: 100%|██████████| 108/108 [00:01<00:00, 102.79it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 197.51it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 143.68it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 156.17it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 180.52it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 205.25it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 212.20it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 195.64it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 211.21it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 225.13it/s]\u001b[A\n",
"Epoch 6: 100%|██████████| 108/108 [00:01<00:00, 76.04it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 7: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 7: 11%|█ | 12/108 [00:00<00:02, 33.31it/s, v_num=0]\n",
"Epoch 7: 20%|██ | 22/108 [00:00<00:01, 46.90it/s, v_num=0]\n",
"Epoch 7: 22%|██▏ | 24/108 [00:00<00:01, 50.49it/s, v_num=0]\n",
"Epoch 7: 31%|███▏ | 34/108 [00:00<00:01, 58.20it/s, v_num=0]\n",
"Epoch 7: 32%|███▏ | 35/108 [00:00<00:01, 59.59it/s, v_num=0]\n",
"Epoch 7: 33%|███▎ | 36/108 [00:00<00:01, 60.97it/s, v_num=0]\n",
"Epoch 7: 48%|████▊ | 52/108 [00:00<00:00, 74.69it/s, v_num=0]\n",
"Epoch 7: 64%|██████▍ | 69/108 [00:00<00:00, 85.96it/s, v_num=0]\n",
"Epoch 7: 80%|███████▉ | 86/108 [00:00<00:00, 94.41it/s, v_num=0]\n",
"Epoch 7: 81%|████████▏ | 88/108 [00:00<00:00, 95.91it/s, v_num=0]\n",
"Epoch 7: 97%|█████████▋| 105/108 [00:01<00:00, 102.61it/s, v_num=0]\n",
"Epoch 7: 100%|██████████| 108/108 [00:01<00:00, 103.00it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 215.46it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 246.46it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 264.39it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 256.84it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 218.46it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 230.90it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 243.53it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 253.83it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 249.22it/s]\u001b[A\n",
"Epoch 7: 100%|██████████| 108/108 [00:01<00:00, 78.36it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 8: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 8: 7%|▋ | 8/108 [00:00<00:03, 25.72it/s, v_num=0]\n",
"Epoch 8: 19%|█▊ | 20/108 [00:00<00:01, 47.54it/s, v_num=0]\n",
"Epoch 8: 31%|███ | 33/108 [00:00<00:01, 62.61it/s, v_num=0]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 7x across cluster]\u001b[0m\n",
"Epoch 8: 44%|████▎ | 47/108 [00:00<00:00, 74.45it/s, v_num=0]\n",
"Epoch 8: 55%|█████▍ | 59/108 [00:00<00:00, 81.84it/s, v_num=0]\n",
"Epoch 8: 56%|█████▌ | 60/108 [00:00<00:00, 80.82it/s, v_num=0]\n",
"Epoch 8: 57%|█████▋ | 62/108 [00:00<00:00, 82.74it/s, v_num=0]\n",
"Epoch 8: 58%|█████▊ | 63/108 [00:00<00:00, 83.69it/s, v_num=0]\n",
"Epoch 8: 70%|███████ | 76/108 [00:00<00:00, 88.60it/s, v_num=0]\n",
"Epoch 8: 85%|████████▌ | 92/108 [00:00<00:00, 96.53it/s, v_num=0]\n",
"Epoch 8: 86%|████████▌ | 93/108 [00:00<00:00, 96.21it/s, v_num=0]\n",
"Epoch 8: 87%|████████▋ | 94/108 [00:00<00:00, 96.72it/s, v_num=0]\n",
"Epoch 8: 88%|████████▊ | 95/108 [00:00<00:00, 97.32it/s, v_num=0]\n",
"Epoch 8: 89%|████████▉ | 96/108 [00:00<00:00, 98.03it/s, v_num=0]\n",
"Epoch 8: 100%|██████████| 108/108 [00:01<00:00, 102.15it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 228.96it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 220.63it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 220.41it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 208.74it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 221.74it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 243.64it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 253.60it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 254.93it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 207.23it/s]\u001b[A\n",
"Epoch 8: 100%|██████████| 108/108 [00:01<00:00, 78.28it/s, v_num=0] it/s]\u001b[A\n",
" \u001b[A\n",
"Epoch 9: 0%| | 0/108 [00:00<?, ?it/s, v_num=0] \n",
"Epoch 9: 11%|█ | 12/108 [00:00<00:02, 33.03it/s, v_num=0]\n",
"Epoch 9: 21%|██▏ | 23/108 [00:00<00:01, 48.82it/s, v_num=0]\n",
"Epoch 9: 31%|███ | 33/108 [00:00<00:01, 58.62it/s, v_num=0]\n",
"Epoch 9: 31%|███▏ | 34/108 [00:00<00:01, 58.61it/s, v_num=0]\n",
"Epoch 9: 32%|███▏ | 35/108 [00:00<00:01, 59.89it/s, v_num=0]\n",
"Epoch 9: 33%|███▎ | 36/108 [00:00<00:01, 61.11it/s, v_num=0]\n",
"Epoch 9: 46%|████▋ | 50/108 [00:00<00:00, 71.95it/s, v_num=0]\n",
"Epoch 9: 61%|██████ | 66/108 [00:00<00:00, 82.62it/s, v_num=0]\n",
"Epoch 9: 70%|███████ | 76/108 [00:00<00:00, 83.77it/s, v_num=0]\n",
"Epoch 9: 71%|███████▏ | 77/108 [00:00<00:00, 84.54it/s, v_num=0]\n",
"Epoch 9: 72%|███████▏ | 78/108 [00:00<00:00, 85.33it/s, v_num=0]\n",
"Epoch 9: 88%|████████▊ | 95/108 [00:01<00:00, 93.18it/s, v_num=0]\n",
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 98.27it/s, v_num=0]\n",
"Validation: 0it [00:00, ?it/s]\u001b[A76)\u001b[0m \n",
"Validation: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
"Validation DataLoader 0: 10%|█ | 1/10 [00:00<00:00, 305.42it/s]\u001b[A\n",
"Validation DataLoader 0: 20%|██ | 2/10 [00:00<00:00, 337.39it/s]\u001b[A\n",
"Validation DataLoader 0: 30%|███ | 3/10 [00:00<00:00, 368.65it/s]\u001b[A\n",
"Validation DataLoader 0: 40%|████ | 4/10 [00:00<00:00, 361.22it/s]\u001b[A\n",
"Validation DataLoader 0: 50%|█████ | 5/10 [00:00<00:00, 250.96it/s]\u001b[A\n",
"Validation DataLoader 0: 60%|██████ | 6/10 [00:00<00:00, 271.98it/s]\u001b[A\n",
"Validation DataLoader 0: 70%|███████ | 7/10 [00:00<00:00, 289.64it/s]\u001b[A\n",
"Validation DataLoader 0: 80%|████████ | 8/10 [00:00<00:00, 304.16it/s]\u001b[A\n",
"Validation DataLoader 0: 90%|█████████ | 9/10 [00:00<00:00, 184.87it/s]\u001b[A\n",
"Validation DataLoader 0: 100%|██████████| 10/10 [00:00<00:00, 196.99it/s]\u001b[A\n",
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 74.63it/s, v_num=0]\n",
" \u001b[A\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m `Trainer.fit` stopped: `max_epochs=10` reached.\n",
"100%|██████████| 4542/4542 [00:00<00:00, 48474627.91it/s]\u001b[32m [repeated 14x across cluster]\u001b[0m\n",
"100%|██████████| 9912422/9912422 [00:00<00:00, 90032420.31it/s]\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m \u001b[32m [repeated 5x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120178)\u001b[0m Missing logger folder: logs/lightning_logs\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120179)\u001b[0m LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\u001b[32m [repeated 3x across cluster]\u001b[0m\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m [rank: 0] Global seed set to 888\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 9: 100%|██████████| 108/108 [00:01<00:00, 66.61it/s, v_num=0]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m /mnt/cluster_storage/pypi/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:225: PossibleUserWarning: Using `DistributedSampler` with the dataloaders. During `trainer.test()`, it is recommended to use `Trainer(devices=1, num_nodes=1)` to ensure each sample/batch gets evaluated exactly once. Otherwise, multi-device settings use `DistributedSampler` that replicates some samples to make sure all devices have same batch size in case of uneven inputs.\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m rank_zero_warn(\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Testing DataLoader 0: 25%|██▌ | 5/20 [00:00<00:00, 146.57it/s]\n",
"Testing DataLoader 0: 100%|██████████| 20/20 [00:00<00:00, 163.98it/s]\n",
"Testing DataLoader 0: 100%|██████████| 20/20 [00:00<00:00, 125.34it/s]\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┃ Test metric ┃ DataLoader 0 ┃\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m │ test_accuracy │ 0.9740999937057495 │\n",
"\u001b[2m\u001b[36m(RayTrainWorker pid=120176)\u001b[0m └───────────────────────────┴───────────────────────────┘\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"2023-08-07 23:41:11,072\tINFO tune.py:1145 -- Total run time: 39.92 seconds (39.80 seconds for the tuning loop).\n"
]
}
],
"source": [
"result = trainer.fit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Check training results and checkpoints"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Result(\n",
" metrics={'train_loss': 0.03159375861287117, 'val_accuracy': 0.9700015783309937, 'val_loss': -12.346583366394043, 'epoch': 9, 'step': 1080},\n",
" path='/tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31',\n",
" checkpoint=LegacyTorchCheckpoint(local_path=/tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31/checkpoint_000009)\n",
")"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Validation Accuracy: 0.9700015783309937\n",
"Trial Directory: /tmp/ray_results/ptl-mnist-example/TorchTrainer_78346_00000_0_2023-08-07_23-40-31\n",
"['checkpoint_000007', 'checkpoint_000008', 'checkpoint_000009', 'events.out.tfevents.1691476838.ip-10-0-6-244', 'params.json', 'params.pkl', 'progress.csv', 'rank_0', 'rank_0.lock', 'rank_1', 'rank_1.lock', 'rank_2', 'rank_2.lock', 'rank_3', 'rank_3.lock', 'result.json']\n"
]
}
],
"source": [
"print(\"Validation Accuracy: \", result.metrics[\"val_accuracy\"])\n",
"print(\"Trial Directory: \", result.path)\n",
"print(sorted(os.listdir(result.path)))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Ray Train saved three checkpoints(`checkpoint_000007`, `checkpoint_000008`, `checkpoint_000009`) in the trial directory. The following code retrieves the latest checkpoint from the fit results and loads it back into the model.\n",
"\n",
"If you lost the in-memory result object, you can restore the model from the checkpoint file. The checkpoint path is: `/tmp/ray_results/ptl-mnist-example/TorchTrainer_eb925_00000_0_2023-08-07_23-15-06/checkpoint_000009/checkpoint.ckpt`."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Global seed set to 888\n"
]
},
{
"data": {
"text/plain": [
"MNISTClassifier(\n",
" (linear_relu_stack): Sequential(\n",
" (0): Linear(in_features=784, out_features=128, bias=True)\n",
" (1): ReLU()\n",
" (2): Linear(in_features=128, out_features=10, bias=True)\n",
" (3): ReLU()\n",
" )\n",
" (accuracy): MulticlassAccuracy()\n",
")"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"checkpoint = result.checkpoint\n",
"\n",
"with checkpoint.as_directory() as ckpt_dir:\n",
" best_model = MNISTClassifier.load_from_checkpoint(f\"{ckpt_dir}/checkpoint.ckpt\")\n",
"\n",
"best_model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## See also\n",
"\n",
"* {ref}`Getting Started with PyTorch Lightning <train-pytorch-lightning>` for a tutorial on using Ray Train and PyTorch Lightning \n",
"\n",
"* {doc}`Ray Train Examples <../../examples>` for more use cases"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "a8c1140d108077f4faeb76b2438f85e4ed675f93d004359552883616a1acd54c"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/lightning/vicuna_13b_lightning_deepspeed_finetune.ipynb
+1272
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/pytorch/BUILD.bazel
+43
View File
@@ -0,0 +1,43 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "train_pytorch_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
exclude = ["convert_existing_pytorch_code_to_ray_train.ipynb"], # CPU test
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
# CPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["convert_existing_pytorch_code_to_ray_train.ipynb"],
data = ["//doc/source/train/examples/pytorch:train_pytorch_examples"],
exclude = [],
tags = [
"exclusive",
"ray_air",
"team:ml",
],
)
filegroup(
name = "pytorch_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/pytorch/convert_existing_pytorch_code_to_ray_train.ipynb
+1030
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/pytorch/dreambooth_finetuning.rst
+304
View File
@@ -0,0 +1,304 @@
:orphan:
Fine-tuning of Stable Diffusion with DreamBooth and Ray Train
=============================================================
.. raw:: html
<a id="try-anyscale-quickstart-dreambooth_finetuning" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This is an intermediate example that shows how to do DreamBooth fine-tuning of a Stable Diffusion model using Ray Train.
It demonstrates how to use :ref:`Ray Data <data>` with PyTorch Lightning in Ray Train.
See the original `DreamBooth project homepage <https://dreambooth.github.io/>`_ for more details on what this fine-tuning method achieves.
.. image:: https://dreambooth.github.io/DreamBooth_files/high_level.png
:target: https://dreambooth.github.io
:alt: DreamBooth fine-tuning overview
This example builds on `this Hugging Face 🤗 tutorial <https://huggingface.co/docs/diffusers/training/dreambooth>`_.
See the Hugging Face tutorial for useful explanations and suggestions on hyperparameters.
**Adapting this example to Ray Train allows you to easily scale up the fine-tuning to an arbitrary number of distributed training workers.**
**Compute requirements:**
* Because of the large model sizes, you need a machine with at least 1 A10G GPU.
* Each training worker uses 1 GPU. You can use multiple GPUs or workers to leverage data-parallel training to speed up training time.
This example fine-tunes both the ``text_encoder`` and ``unet`` models used in the stable diffusion process, with respect to a prior preserving loss.
.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_example.png
:alt: DreamBooth overview
Find the full code repository at `https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning <https://github.com/ray-project/ray/tree/master/doc/source/templates/05_dreambooth_finetuning>`_
How it works
------------
This example uses Ray Data for data loading and Ray Train for distributed training.
Data loading
^^^^^^^^^^^^
.. note::
Find the latest version of the code at `dataset.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/dataset.py>`_
The latest version might differ slightly from the code presented here.
Use Ray Data for data loading. The code has three interesting parts.
First, load two datasets using :func:`ray.data.read_images`:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-at: instance_dataset = read
:end-at: class_dataset = read
:dedent: 4
Then, tokenize the prompt that generated these images:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-at: tokenizer = AutoTokenizer
:end-at: instance_prompt_ids = _tokenize
:dedent: 4
And lastly, apply a ``torchvision`` preprocessing pipeline to the images:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-after: START: image preprocessing
:end-before: END: image preprocessing
:dedent: 4
Apply all three parts in a final step:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/dataset.py
:language: python
:start-after: START: Apply preprocessing
:end-before: END: Apply preprocessing
:dedent: 4
Distributed training
^^^^^^^^^^^^^^^^^^^^
.. note::
Find the latest version of the code at `train.py <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_
The latest version might differ slightly from the code presented here.
The central part of the training code is the :ref:`training function <train-overview-training-function>`. This function accepts a configuration dict that contains the hyperparameters. It then defines a regular PyTorch training loop.
You interact with the Ray Train API in only a few locations, which follow in-line comments in the snippet below.
Remember that you want to do data-parallel training for all the models.
#. Load the data shard for each worker with `session.get_dataset_shard("train")``
#. Iterate over the dataset with `train_dataset.iter_torch_batches()``
#. Report results to Ray Train with `session.report(results)``
The code is compacted for brevity. The `full code <https://github.com/ray-project/ray/blob/master/doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py>`_ is more thoroughly annotated.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
:language: python
:start-at: def train_fn(config)
:end-before: END: Training loop
You can then run this training function with Ray Train's TorchTrainer:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth/train.py
:language: python
:start-at: args = train_arguments
:end-at: trainer.fit()
:dedent: 4
Configure the scale
^^^^^^^^^^^^^^^^^^^
In the TorchTrainer, you can easily configure the scale.
The preceding example uses the ``num_workers`` argument to specify the number
of workers. This argument defaults to 2 workers with 1 GPU each, totalling to 2 GPUs.
To run the example on 4 GPUs, set the number of workers to 4 using ``--num-workers=4``.
Or you can change the scaling config directly:
.. code-block:: diff
scaling_config=ScalingConfig(
use_gpu=True,
- num_workers=args.num_workers,
+ num_workers=4,
)
If you're running multi-node training, make sure that all nodes have access to a shared
storage like NFS or EFS. In the following example script, you can adjust the location with the
``DATA_PREFIX`` environment variable.
Training throughput
~~~~~~~~~~~~~~~~~~~
Compare throughput of the preceding training runs that used 1, 2, and 4 workers or GPUs.
Consider the following setup:
* 1 GCE g2-standard-48-nvidia-l4-4 instance with 4 GPUs
* Model as configured below
* Data from this example
* 200 regularization images
* Training for 4 epochs (local batch size = 2)
* 3 runs per configuration
You expect that the training time should benefit from scale and decreases when running with
more workers and GPUs.
.. image:: /templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_training.png
:alt: DreamBooth training times
.. list-table::
:header-rows: 1
* - Number of workers/GPUs
- Training time (seconds)
* - 1
- 802.14
* - 2
- 487.82
* - 4
- 313.25
While the training time decreases linearly with the amount of workers/GPUs, you can observe some penalty.
Specifically, with double the amount of workers you don't get half of the training time.
This penalty is most likely due to additional communication between processes and the transfer of large model
weights. You are also only training with a batch size of one because of the GPU memory limitation. On larger
GPUs with higher batch sizes you would expect a greater benefit from scaling out.
Run the example
---------------
First, download the pre-trained Stable Diffusion model as a starting point.
Then train this model with a few images of a subject.
To achieve this, choose a non-word as an identifier, such as ``unqtkn``. When fine-tuning the model with this subject, you teach the model that the prompt is ``A photo of a unqtkn <class>``.
After fine-tuning you can run inference with this specific prompt.
For instance: ``A photo of a unqtkn <class>`` creates an image of the subject.
Similarly, ``A photo of a unqtkn <class> at the beach`` creates an image of the subject at the beach.
Step 0: Preparation
^^^^^^^^^^^^^^^^^^^
Clone the Ray repository, go to the example directory, and install dependencies.
.. code-block:: bash
git clone https://github.com/ray-project/ray.git
cd doc/source/templates/05_dreambooth_finetuning
pip install -Ur dreambooth/requirements.txt
Prepare some directories and environment variables.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __preparation_start__
:end-before: __preparation_end__
Step 1: Download the pre-trained model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Download and cache a pre-trained Stable Diffusion model locally.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __cache_model_start__
:end-before: __cache_model_end__
You can access the downloaded model checkpoint at the ``$ORIG_MODEL_PATH``.
Step 2: Supply images of your subject
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Use one of the sample datasets, like `dog` or `lego car`, or provide your own directory
of images, and specify the directory with the ``$INSTANCE_DIR`` environment variable.
Then, copy these images to ``$IMAGES_OWN_DIR``.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: __supply_own_images_start__
:end-before: __supply_own_images_end__
The ``$CLASS_NAME`` should be the general category of your subject.
The images produced by the prompt ``photo of a unqtkn <class>`` should be diverse images
that are different enough from the subject in order for generated images to clearly
show the effect of fine-tuning.
Step 3: Create the regularization images
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Create a regularization image set for a class of subjects using the pre-trained
Stable Diffusion model. This regularization set ensures that
the model still produces decent images for random images of the same class,
rather than just optimize for producing good images of the subject.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 3: START
:end-before: Step 3: END
Use Ray Data to do batch inference with 4 workers, to generate more images in parallel.
Step 4: Fine-tune the model
^^^^^^^^^^^^^^^^^^^^^^^^^^^
Save a few, like 4 to 5, images of the subject being fine-tuned
in a local directory. Then launch the training job with:
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 4: START
:end-before: Step 4: END
Step 5: Generate images of the subject
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Try your model with the same command line as Step 2, but point
to your own model this time.
.. literalinclude:: /templates/05_dreambooth_finetuning/dreambooth_run.sh
:language: bash
:start-after: Step 5: START
:end-before: Step 5: END
Next, try replacing the prompt with something more interesting.
For example, for the dog subject, you can try:
- "photo of a unqtkn dog in a bucket"
- "photo of a unqtkn dog sleeping"
- "photo of a unqtkn dog in a doghouse"
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases
* :ref:`Ray Train User Guides <train-user-guides>` for how-to guides
doc/source/train/examples/pytorch/pytorch_resnet_finetune.ipynb
+519
View File
@@ -0,0 +1,519 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Finetuning a Pytorch Image Classifier with Ray Train\n",
"\n",
"<a id=\"try-anyscale-quickstart-pytorch_resnet_finetune\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=pytorch_resnet_finetune\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"This example fine-tunes a pre-trained ResNet model with Ray Train. \n",
"\n",
"For this example, the network architecture consists of the intermediate layer output of a pre-trained ResNet model, which feeds into a randomly initialized linear layer that outputs classification logits for our new task.\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load and preprocess finetuning dataset\n",
"This example is adapted from Pytorch's [Transfer Learning for Computer Vision](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html) tutorial.\n",
"We will use *hymenoptera_data* as the finetuning dataset, which contains two classes (bees and ants) and 397 total images (across training and validation). This is a quite small dataset and used only for demonstration purposes. "
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"# To run full example, set SMOKE_TEST as False\n",
"SMOKE_TEST = True\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The dataset is publicly available [here](https://www.kaggle.com/datasets/ajayrana/hymenoptera-data). Note that it is structured with directory names as the labels. Use `torchvision.datasets.ImageFolder()` to load the images and their corresponding labels."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets, models, transforms\n",
"import numpy as np\n",
"\n",
"# Data augmentation and normalization for training\n",
"# Just normalization for validation\n",
"data_transforms = {\n",
" \"train\": transforms.Compose(\n",
" [\n",
" transforms.RandomResizedCrop(224),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]\n",
" ),\n",
" \"val\": transforms.Compose(\n",
" [\n",
" transforms.Resize(224),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]\n",
" ),\n",
"}\n",
"\n",
"def download_datasets():\n",
" os.system(\n",
" \"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\"\n",
" )\n",
" os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")\n",
"\n",
"# Download and build torch datasets\n",
"def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" return torch_datasets\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" from torch.utils.data import Subset\n",
"\n",
" def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" \n",
" # Only take a subset for smoke test\n",
" for split in [\"train\", \"val\"]:\n",
" indices = list(range(100))\n",
" torch_datasets[split] = Subset(torch_datasets[split], indices)\n",
" return torch_datasets\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Initialize Model and Fine-tuning configs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's define the training configuration that will be passed into the training loop function later."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"train_loop_config = {\n",
" \"input_size\": 224, # Input image size (224 x 224)\n",
" \"batch_size\": 32, # Batch size for training\n",
" \"num_epochs\": 10, # Number of epochs to train for\n",
" \"lr\": 0.001, # Learning Rate\n",
" \"momentum\": 0.9, # SGD optimizer momentum\n",
"}\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's define our model. You can either create a model from pre-trained weights or reload the model checkpoint from a previous run."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"from ray.train import Checkpoint\n",
"\n",
"# Option 1: Initialize model with pretrained weights\n",
"def initialize_model():\n",
" # Load pretrained model params\n",
" model = models.resnet50(pretrained=True)\n",
"\n",
" # Replace the original classifier with a new Linear layer\n",
" num_features = model.fc.in_features\n",
" model.fc = nn.Linear(num_features, 2)\n",
"\n",
" # Ensure all params get updated during finetuning\n",
" for param in model.parameters():\n",
" param.requires_grad = True\n",
" return model\n",
"\n",
"\n",
"# Option 2: Initialize model with an Train checkpoint\n",
"# Replace this with your own uri\n",
"CHECKPOINT_FROM_S3 = Checkpoint(\n",
" path=\"s3://air-example-data/finetune-resnet-checkpoint/TorchTrainer_4f69f_00000_0_2023-02-14_14-04-09/checkpoint_000001/\"\n",
")\n",
"\n",
"\n",
"def initialize_model_from_checkpoint(checkpoint: Checkpoint):\n",
" with checkpoint.as_directory() as tmpdir:\n",
" state_dict = torch.load(os.path.join(tmpdir, \"checkpoint.pt\"))\n",
" resnet50 = initialize_model()\n",
" resnet50.load_state_dict(state_dict[\"model\"])\n",
" return resnet50\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define the Training Loop\n",
"\n",
"The `train_loop_per_worker` function defines the fine-tuning procedure for each worker.\n",
"\n",
"**1. Prepare dataloaders for each worker**:\n",
"- This tutorial assumes you are using PyTorch's native `torch.utils.data.Dataset` for data input. {meth}`train.torch.prepare_data_loader() <ray.train.torch.prepare_data_loader>` prepares your dataLoader for distributed execution. You can also use Ray Data for more efficient preprocessing. For more details on using Ray Data for images, see the {doc}`Working with Images </data/working-with-images>` Ray Data user guide.\n",
"\n",
"**2. Prepare your model**:\n",
"- {meth}`train.torch.prepare_model() <ray.train.torch.prepare_model>` prepares the model for distributed training. Under the hood, it converts your torch model to `DistributedDataParallel` model, which synchronize its weights across all workers.\n",
"\n",
"**3. Report metrics and checkpoint**:\n",
"- {meth}`train.report() <ray.train.report>` will report metrics and checkpoints to Ray Train.\n",
"- Saving checkpoints through {meth}`train.report(metrics, checkpoint=...) <ray.train.report>` will automatically [upload checkpoints to cloud storage](tune-cloud-checkpointing) (if configured), and allow you to easily enable Ray Train worker fault tolerance in the future."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"import ray.train as train\n",
"from ray.train import Checkpoint\n",
"\n",
"\n",
"\n",
"def evaluate(logits, labels):\n",
" _, preds = torch.max(logits, 1)\n",
" corrects = torch.sum(preds == labels).item()\n",
" return corrects\n",
"\n",
"\n",
"def train_loop_per_worker(configs):\n",
" import warnings\n",
"\n",
" warnings.filterwarnings(\"ignore\")\n",
"\n",
" # Calculate the batch size for a single worker\n",
" worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
"\n",
" # Download dataset once on local rank 0 worker\n",
" if train.get_context().get_local_rank() == 0:\n",
" download_datasets()\n",
" torch.distributed.barrier()\n",
"\n",
" # Build datasets on each worker\n",
" torch_datasets = build_datasets()\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = dict()\n",
" dataloaders[\"train\"] = DataLoader(\n",
" torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
" )\n",
" dataloaders[\"val\"] = DataLoader(\n",
" torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
" )\n",
"\n",
" # Distribute\n",
" dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
" dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
"\n",
" device = train.torch.get_device()\n",
"\n",
" # Prepare DDP Model, optimizer, and loss function\n",
" model = initialize_model()\n",
" model = train.torch.prepare_model(model)\n",
"\n",
" optimizer = optim.SGD(\n",
" model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
" )\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" # Start training loops\n",
" for epoch in range(configs[\"num_epochs\"]):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"val\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" if train.get_context().get_world_size() > 1:\n",
" dataloaders[phase].sampler.set_epoch(epoch)\n",
"\n",
" for inputs, labels in dataloaders[phase]:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # calculate statistics\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += evaluate(outputs, labels)\n",
"\n",
" size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
" epoch_loss = running_loss / size\n",
" epoch_acc = running_corrects / size\n",
"\n",
" if train.get_context().get_world_rank() == 0:\n",
" print(\n",
" \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
" epoch, phase, epoch_loss, epoch_acc\n",
" )\n",
" )\n",
"\n",
" # Report metrics and checkpoint every epoch\n",
" if phase == \"val\":\n",
" with TemporaryDirectory() as tmpdir:\n",
" state_dict = {\n",
" \"epoch\": epoch,\n",
" \"model\": model.module.state_dict(),\n",
" \"optimizer_state_dict\": optimizer.state_dict(),\n",
" }\n",
" torch.save(state_dict, os.path.join(tmpdir, \"checkpoint.pt\"))\n",
" train.report(\n",
" metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
" checkpoint=Checkpoint.from_directory(tmpdir),\n",
" )\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, setup the TorchTrainer:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import ScalingConfig, RunConfig, CheckpointConfig\n",
"\n",
"# Scale out model training across 4 GPUs.\n",
"scaling_config = ScalingConfig(\n",
" num_workers=4, use_gpu=True, resources_per_worker={\"CPU\": 1, \"GPU\": 1}\n",
")\n",
"\n",
"# Save the latest checkpoint\n",
"checkpoint_config = CheckpointConfig(num_to_keep=1)\n",
"\n",
"# Set experiment name and checkpoint configs\n",
"run_config = RunConfig(\n",
" name=\"finetune-resnet\",\n",
" storage_path=\"/tmp/ray_results\",\n",
" checkpoint_config=checkpoint_config,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" scaling_config = ScalingConfig(\n",
" num_workers=8, use_gpu=False, resources_per_worker={\"CPU\": 1}\n",
" )\n",
" train_loop_config[\"num_epochs\"] = 1\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"trainer = TorchTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config=train_loop_config,\n",
" scaling_config=scaling_config,\n",
" run_config=run_config,\n",
")\n",
"\n",
"result = trainer.fit()\n",
"print(result)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load the checkpoint for prediction:\n",
"\n",
" \n",
" The metadata and checkpoints have already been saved in the `storage_path` specified in TorchTrainer:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We now need to load the trained model and evaluate it on test data. The best model parameters have been saved in `log_dir`. We can load the resulting checkpoint from our fine-tuning run using the previously defined `initialize_model_from_checkpoint()` function."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = initialize_model_from_checkpoint(result.checkpoint)\n",
"device = torch.device(\"cuda\")\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"tags": [
"remove-cell"
]
},
"outputs": [],
"source": [
"if SMOKE_TEST:\n",
" device = torch.device(\"cpu\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, define a simple evaluation loop and check the performance of the checkpoint model."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Accuracy: 0.934640522875817\n"
]
}
],
"source": [
"model = model.to(device)\n",
"model.eval()\n",
"\n",
"download_datasets()\n",
"torch_datasets = build_datasets()\n",
"dataloader = DataLoader(torch_datasets[\"val\"], batch_size=32, num_workers=4)\n",
"corrects = 0\n",
"for inputs, labels in dataloader:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
" preds = model(inputs)\n",
" corrects += evaluate(preds, labels)\n",
"\n",
"print(\"Accuracy: \", corrects / len(dataloader.dataset))\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "a8c1140d108077f4faeb76b2438f85e4ed675f93d004359552883616a1acd54c"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/pytorch/torch_detection.ipynb
+1393
View File
File diff suppressed because one or more lines are too long
doc/source/train/examples/pytorch/torch_fashion_mnist_example.rst
+27
View File
@@ -0,0 +1,27 @@
:orphan:
.. _train-pytorch-fashion-mnist:
Train a PyTorch model on Fashion MNIST
======================================
.. raw:: html
<a id="try-anyscale-quickstart-torch_fashion_mnist_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_fashion_mnist_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This example runs distributed training of a PyTorch model on Fashion MNIST with Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_fashion_mnist_example.py
See also
--------
* :ref:`Get Started with PyTorch <train-pytorch>` for a tutorial on using Ray Train and PyTorch
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/examples/pytorch/torch_regression_example.rst
+14
View File
@@ -0,0 +1,14 @@
:orphan:
torch_regression_example
========================
.. raw:: html
<a id="try-anyscale-quickstart-torch_regression_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=torch_regression_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
.. literalinclude:: /../../python/ray/train/examples/pytorch/torch_regression_example.py
doc/source/train/examples/tf/tensorflow_mnist_example.rst
+26
View File
@@ -0,0 +1,26 @@
:orphan:
Training with TensorFlow and Ray Train
======================================
.. raw:: html
<a id="try-anyscale-quickstart-tensorflow_mnist_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=tensorflow_mnist_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example runs distributed training of a TensorFlow model on MNIST with Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/tf/tensorflow_mnist_example.py
See also
--------
* :doc:`Ray Train Examples <../../examples>` for more use cases.
* :ref:`Distributed Tensorflow & Keras <train-tensorflow-overview>` for a tutorial.
doc/source/train/examples/tf/tensorflow_regression_example.rst
+13
View File
@@ -0,0 +1,13 @@
:orphan:
tensorflow_regression_example
=============================
.. raw:: html
<a id="try-anyscale-quickstart-tensorflow_regression_example" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=tensorflow_regression_example">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
.. literalinclude:: /../../python/ray/train/examples/tf/tensorflow_regression_example.py
doc/source/train/examples/transformers/BUILD.bazel
+30
View File
@@ -0,0 +1,30 @@
load("//bazel:python.bzl", "py_test_run_all_notebooks")
filegroup(
name = "huggingface_transformers_examples",
srcs = glob(["*.ipynb"]),
visibility = ["//doc:__subpackages__"],
)
# GPU Tests
py_test_run_all_notebooks(
size = "large",
include = ["*.ipynb"],
data = ["//doc/source/train/examples/transformers:huggingface_transformers_examples"],
exclude = [],
tags = [
"exclusive",
"gpu",
"ray_air",
"team:ml",
],
)
filegroup(
name = "transformers_examples_ci_configs",
srcs = glob([
"**/ci/aws.yaml",
"**/ci/gce.yaml",
]),
visibility = ["//doc/source/train/examples:__pkg__"],
)
doc/source/train/examples/transformers/huggingface_text_classification.ipynb
+1338
View File
File diff suppressed because it is too large Load Diff
doc/source/train/examples/transformers/transformer_reinforcement_learning/README.ipynb
+417
View File
@@ -0,0 +1,417 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "d31e727a103a6eb7",
"metadata": {},
"source": [
"# RL Post-Training using Hugging Face TRL with GRPO\n",
"\n",
"This notebook builds on the TRL GRPO Example, [GRPO Trainer](https://huggingface.co/docs/trl/en/grpo_trainer), to fine-tune a Qwen2.5 0.5B model to answer math questions using the DeepMath-103k dataset and Group Relative Policy Optimization (GRPO) algorithm.\n",
"\n",
"Ray Train scales this implementation efficiently across multiple GPUs without changing training logic.\n",
"\n",
"This notebook consists of the following steps:\n",
"\n",
"1. Quick Summary of GRPO and the DeepMath dataset\n",
"2. Package and Environment Setup\n",
"3. Running TRL with Ray Train\n",
"4. Scaling to more GPUs\n",
"\n",
"<div id=\"anyscale-note\" class=\"alert alert-block alert-warning\">\n",
"\n",
" <strong>Anyscale Specific Configuration</strong>\n",
"\n",
" <p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, you need additional configuration. For example, you would need to manually:</p>\n",
"\n",
" <ul>\n",
" <li><strong>Configure your Ray Cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li>\n",
" <li><strong>Manage Dependencies</strong>: Manually install and manage dependencies on each node.</li>\n",
" <li><strong>Set Up Storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>\n",
" </ul>\n",
"\n",
" <p>All these configurations are handled automatically through the Anyscale platform.\n",
"</div>\n",
"\n",
"<style>\n",
" div#anyscale-note > p,\n",
" div#anyscale-note > ul,\n",
" div#anyscale-note > ul li {\n",
" color: black;\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" background-color: rgb(255, 243, 205);\n",
" }\n",
"\n",
" div#anyscale-note {\n",
" border: 1px solid #ccc;\n",
" border-radius: 8px;\n",
" padding: 15px;\n",
" }\n",
"\n",
"</style>"
]
},
{
"cell_type": "markdown",
"id": "94faee5d4ef298c5",
"metadata": {},
"source": [
"## 1. Summary of the DeepMath dataset and GRPO algorithm\n",
"\n",
"### DeepMath dataset\n",
"\n",
"This example uses the DeepMath-103k dataset by [He et al., 2025](https://arxiv.org/abs/2504.11456), consisting of 103,000 challenging questions across a wide range of mathematical subjects including Algebra, Calculus, Number Theory, Geometry, Probability, and Discrete Mathematics. These questions are more heavily distributed towards challenging questions, in particular Levels 5 to 9, compared to alternative datasets, making them highly complex and difficult to solve without specialist training.\n",
"\n",
"Each example contains two fields:\n",
"- **`prompt`**: a list of chat-format message dicts; the question text is in `prompt[0][\"content\"]`.\n",
"- **`solution`**: a LaTeX-formatted ground-truth answer, used by the reward function to verify the model's response.\n",
"\n",
"Using a structured LaTeX format ensures the reward function can reliably parse and compare both the model's output and the ground truth. Here are a few examples from the dataset:\n",
"\n",
"| Subject | Prompt | Solution |\n",
"|---|---|---|\n",
"| Number Theory | Determine the number of 19th power residues modulo 229. | $12$ |\n",
"| Functional Analysis | Determine the norm $\\lVert T \\rVert$ of the linear operator $T: \\ell^2 \\rightarrow \\ell^2$ given by $(Tx)_1 = 0$, $(Tx)_n = -x_n + \\alpha x_{n+1}$ for $n \\geq 2$, where $\\alpha \\in \\mathbb{C}$. | $1 + \\lvert\\alpha\\rvert$ |\n",
"| Analysis | Determine whether there exists a Schwartz function $g \\in \\mathcal{S}(\\mathbb{R}^n)$ such that for a given continuous function $f: \\mathbb{R}^n \\to \\mathbb{R}$ with $f \\not\\equiv 0$, the integral $\\int_{\\mathbb{R}^n} g^2 f \\, dx \\neq 0$. | Yes |\n",
"\n",
"### Group Relative Policy Optimization Algorithm\n",
"\n",
"To RL post-train the model to answer math questions, this example uses TRL's GRPO (Group Relative Policy Optimization) implementation introduced by [Shao et al., 2024](https://arxiv.org/abs/2402.03300). GRPO is an online algorithm that iteratively improves by training on data the model generates itself. Each iteration proceeds through four steps:\n",
"\n",
"1. **Generate**: For each prompt in the batch, sample a group of G completions from the current policy.\n",
"2. **Score**: Apply the reward function to each completion, producing a scalar reward per completion.\n",
"3. **Advantage**: Normalize the rewards within each group by subtracting the group mean and dividing by the group standard deviation. This relative score becomes the advantage $\\hat{A}_{i,t}$ where positive means better than peers.\n",
"4. **Loss**: Update the policy to increase the probability of high-advantage completions and decrease it for low-advantage ones, while a KL penalty keeps the updated policy close to the reference.\n",
"\n",
"$$\n",
"\\mathcal{L}_{\\text{GRPO}}(\\theta) = -\\frac{1}{\\sum_{i=1}^{G} |o_i|} \\sum_{i=1}^{G} \\sum_{t=1}^{|o_i|} \\left[ \\frac{\\pi_\\theta(o_{i,t} \\mid q, o_{i,\\mathopen{<} t})}{\\left[\\pi_\\theta(o_{i,t} \\mid q, o_{i,\\mathopen{<} t})\\right]_{\\text{no grad}}} \\hat{A}_{i,t} - \\beta \\mathbb{D}_{\\text{KL}}\\left[\\pi_\\theta \\| \\pi_{\\text{ref}}\\right] \\right],\n",
"$$\n",
"\n",
"For a deeper, more technical, description of the algorithm, see the [TRL GRPO page](https://huggingface.co/docs/trl/en/grpo_trainer#looking-deeper-into-the-grpo-method) and the [original paper](https://arxiv.org/abs/2402.03300)."
]
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"## 2. Package and Environment Setup\n",
"\n",
"Install the required dependencies"
],
"id": "1c71d4c53875d94f"
},
{
"metadata": {},
"cell_type": "code",
"source": "# !pip install \"trl[vllm]\" \"math_verify\" \"transformers==4.57.6\"",
"id": "b78729cb7d9d2730",
"outputs": [],
"execution_count": null
},
{
"metadata": {},
"cell_type": "markdown",
"source": "Use `ray.init()` to initialize a local cluster. By default, this cluster contains only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.",
"id": "a362d29fb3189847"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"import ray\n",
"\n",
"ray.init()"
],
"id": "3455573a2a2218f0"
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"Use `ray.cluster_resources()` to check which resources your cluster has access to.\n",
"If you're running this notebook on your local machine or Google Colab, you should see the number of CPU cores and GPUs available to you."
],
"id": "f7018eb1a753b41d"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"from pprint import pprint\n",
"\n",
"pprint(ray.cluster_resources())"
],
"id": "2690f780c2975613"
},
{
"metadata": {},
"cell_type": "markdown",
"source": "Change these two variables to control whether the training uses CPUs or GPUs, and how many workers to spawn. Each worker claims one CPU or GPU, so make sure not to request more resources than are available. By default, the training runs with four GPU worker.",
"id": "c7296dcf7a48bbaa"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"use_gpu = True # set this to `False` to run on CPUs\n",
"num_workers = 4 # set this to the number of GPUs or CPUs you want to use"
],
"id": "776d49988bfefc48"
},
{
"cell_type": "markdown",
"id": "43ad7b293617547a",
"metadata": {},
"source": [
"## 3. Using Hugging Face TRL (Transformer Reinforcement Learning) with Ray Train\n",
"\n",
"In comparison to supervised pre-training where models optimize to minimize the error for their next token, RL-based post-training aims to maximize their reward from a prompt. Therefore, it's crucial to define a reward function to measure the success and to train a model.\n",
"\n",
"This example uses the `trl.rewards.accuracy_reward` function to check whether the model's answer matches the answer in the dataset. As the answers use the LaTeX format, you must parse responses and solution before comparing them. The default `trl.rewards.accuracy_reward` implementation applies timeouts to the `parse` and `verify` functions, which are incompatible with Ray. The version defined below disables these timeouts by setting `parsing_timeout=0` and `timeout_seconds=0`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "902e2cfaa92b68b1",
"metadata": {},
"outputs": [],
"source": [
"from latex2sympy2_extended import NormalizationConfig\n",
"from math_verify import LatexExtractionConfig, parse, verify\n",
"\n",
"def accuracy_reward(completions, solution, **kwargs):\n",
" \"\"\"Reward function that checks mathematical accuracy.\n",
"\n",
" This is a copy of `trl.rewards.accuracy_reward`.\n",
" The only difference is `parse(..., parsing_timeout=0)` and `verify(..., timeout_seconds=0)`\n",
" to avoid `signal.alarm()` issues with ray.\n",
" \"\"\"\n",
" contents = [completion[0][\"content\"] for completion in completions]\n",
" rewards = []\n",
" for content, sol in zip(contents, solution, strict=True):\n",
" gold_parsed = parse(sol, parsing_timeout=0)\n",
" if len(gold_parsed) != 0:\n",
" # We require the answer to be provided in correct latex (no malformed operators)\n",
" answer_parsed = parse(\n",
" content,\n",
" extraction_config=[\n",
" LatexExtractionConfig(\n",
" normalization_config=NormalizationConfig(units=True),\n",
" # Ensures that boxed is tried first\n",
" boxed_match_priority=0,\n",
" try_extract_without_anchor=False,\n",
" )\n",
" ],\n",
" extraction_mode=\"first_match\",\n",
" parsing_timeout=0,\n",
" )\n",
" reward = float(verify(gold_parsed, answer_parsed, timeout_seconds=0))\n",
" else:\n",
" # If the gold solution cannot be parsed, we assign `None` to skip this example\n",
" reward = None\n",
" rewards.append(reward)\n",
"\n",
" return rewards"
]
},
{
"cell_type": "markdown",
"id": "494cc96ea380247d",
"metadata": {},
"source": [
"Next, build the training function that's distributed across all the workers. The `GRPOTrainer` implements the GRPO algorithm for rolling out the model, computing the advantage and backpropagation the losses.\n",
"\n",
"With `vllm_mode=\"colocate\"`, each worker runs a vLLM instance that consumes about 30% of the GPU memory. The alternative `\"server\"` mode dedicates one worker to generation with vLLM while the others perform learning, but this introduces inter-GPU communication overhead that can reduce throughput. See [this blog post](https://huggingface.co/blog/vllm-colocate) for more details.\n",
"\n",
"Two Ray Train-specific additions integrate the HF Trainer into the Ray Train ecosystem:\n",
"\n",
"- **`RayTrainReportCallback`** — hooks into the HF Trainer's logging to forward metrics and checkpoints to Ray Train after each training step. This is what populates the `Result` object returned by `trainer.fit()` with metrics like `reward` and the best checkpoint.\n",
"- **`prepare_trainer`** — configures the HF Trainer for distributed execution across Ray workers. It disables HF's built-in distributed setup so that Ray Train controls the process group instead, and ensures correct device placement on each worker."
]
},
{
"cell_type": "code",
"id": "92c0cb05a65a9da6",
"metadata": {},
"source": [
"from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer\n",
"from trl import GRPOConfig, GRPOTrainer\n",
"from datasets import load_dataset\n",
"\n",
"def train_func(config):\n",
" # Load the DeepMath dataset with only 100 elements for this example.\n",
" # For larger datasets, use Ray Data.\n",
" dataset = load_dataset(\"trl-lib/DeepMath-103K\", split=\"train\").shuffle(seed=42).select(range(100))\n",
"\n",
" # Use vllm_mode=\"colocate\" to allow easy scaling as each GPU handles their own vLLM instance\n",
" training_args = GRPOConfig(\n",
" # Compute a training batch size of 4 for each prompt\n",
" per_device_train_batch_size=4,\n",
" # Use vLLM, colocated on each GPU which uses 30% of the vRAM.\n",
" use_vllm=True,\n",
" vllm_mode=\"colocate\",\n",
" vllm_gpu_memory_utilization=0.3,\n",
" # Run two training epochs over the whole dataset\n",
" num_train_epochs=2,\n",
" # Number of generations per prompt to sample, equivalent to G in the GRPO loss function.\n",
" num_generations=8,\n",
" # Save checkpoints and log metrics every epoch\n",
" save_strategy=\"epoch\",\n",
" logging_strategy=\"epoch\",\n",
" )\n",
"\n",
" # GRPO Trainer\n",
" trainer = GRPOTrainer(\n",
" model=\"Qwen/Qwen2.5-0.5B\",\n",
" args=training_args,\n",
" reward_funcs=accuracy_reward,\n",
" train_dataset=dataset,\n",
" )\n",
"\n",
" # Report metrics and checkpoints to Ray Train\n",
" trainer.add_callback(RayTrainReportCallback())\n",
"\n",
" # Prepare your trainer for TRL and Huggingface Transformer integration\n",
" trainer = prepare_trainer(trainer)\n",
"\n",
" # Start Training\n",
" trainer.train()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"id": "5758d865e9d5b149",
"metadata": {},
"source": "With your `train_func` complete, you can now instantiate the `TorchTrainer`. Aside from calling the function, set the `scaling_config`, which controls the number of workers and resources used, and the `run_config` to configure checkpointing.\n\nThe `CheckpointConfig` controls how Ray Train saves and selects checkpoints during training:\n\n- **`num_to_keep=1`** — Ray Train retains only the single best checkpoint on disk, saving storage.\n- **`checkpoint_score_attribute=\"reward\"`** — Ray Train ranks checkpoints by the `\"reward\"` metric, which is the mean reward across the batch as reported by `RayTrainReportCallback`.\n- **`checkpoint_score_order=\"max\"`** — higher reward is better, so Ray Train keeps the checkpoint with the highest reward seen across all training steps."
},
{
"cell_type": "code",
"execution_count": null,
"id": "fff7fec396590de9",
"metadata": {},
"outputs": [],
"source": [
"from ray.train.torch import TorchTrainer\n",
"from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
"\n",
"trainer = TorchTrainer(\n",
" train_func,\n",
" scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
" run_config=RunConfig(\n",
" # On Anyscale this path is auto-mounted shared storage.\n",
" # For open-source Ray, use a shared NFS path or s3:// URI accessible from all nodes.\n",
" storage_path=\"/mnt/cluster_storage/\",\n",
" checkpoint_config=CheckpointConfig(\n",
" num_to_keep=1,\n",
" checkpoint_score_attribute=\"reward\",\n",
" checkpoint_score_order=\"max\",\n",
" ),\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"id": "b5b92f7546aef635",
"metadata": {},
"source": [
"Finally, call the `fit` method to start training with Ray Train. Save the `Result` object to a variable so you can access metrics and checkpoints."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8ca549e43acc91b1",
"metadata": {},
"outputs": [],
"source": "results = trainer.fit()"
},
{
"cell_type": "markdown",
"id": "af3535209e06ad45",
"metadata": {},
"source": [
"You can use the returned `Result` object to access metrics and the Ray Train `Checkpoint` associated with the last iteration."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fa7b49328202a7a5",
"metadata": {},
"outputs": [],
"source": "print(results.metrics_dataframe)"
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "print(results.get_best_checkpoint(\"reward\", \"max\"))",
"id": "38431d92cf468f46"
},
{
"cell_type": "markdown",
"id": "ark4cvyvi2a",
"source": [
"## 4. Scaling to more GPUs or a larger model\n",
"\n",
"The preceding example trains with 4 workers on a small 100-sample subset of the dataset. To scale up, adjust `num_workers` in `ScalingConfig` and the `num_workers` variable at the top of the notebook, along with the dataset size and model.\n",
"\n",
"**Scaling workers:**\n",
"Each worker claims one GPU. With `vllm_mode=\"colocate\"`, each worker runs its own vLLM instance for generation and its own training process — generation requires no inter-GPU communication. Adding more workers increases the effective batch size and reduces time-to-convergence without any changes to training logic.\n",
"\n",
"**Scaling the model:**\n",
"Replace `\"Qwen/Qwen2.5-0.5B\"` with a larger checkpoint. Larger models require more GPU memory per worker. If a model doesn't fit on a single GPU, consider:\n",
"- Explore using DeepSpeed ZeRO to distribute the model weights across multiple GPUs.\n",
"- Reducing `vllm_gpu_memory_utilization` from `0.3` to leave more memory for model weights at the cost of a smaller vLLM KV cache.\n",
"\n",
"**Fault Tolerance:**\n",
"For longer runs where worker or node failures are possible, set `failure_config=FailureConfig(max_failures=1)` in `RunConfig`. Ray Train will automatically restart training up to twice on failure. Because this example already saves a checkpoint each epoch, a restarted run resumes from the last checkpoint rather than from scratch."
],
"metadata": {}
},
{
"cell_type": "markdown",
"id": "rnid76xr5u",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## Summary\n\nThis notebook demonstrated how to use Ray Train to scale RL post-training of a Qwen2.5 0.5B model across multiple GPUs using the TRL GRPO implementation and the DeepMath-103k dataset.\n\nThe key Ray Train integration points were:\n\n- **`RayTrainReportCallback`** — forwards HF Trainer metrics and checkpoints to Ray Train after each step.\n- **`prepare_trainer`** — configures the HF Trainer to run distributed training under Ray's process group.\n- **`TorchTrainer`** — launches `train_func` on all workers with the configured resources.\n- **`CheckpointConfig`** — automatically tracks and retains the best checkpoint based on the `\"reward\"` metric.\n\nFrom here, you can:\n- Scale to more GPUs or multiple nodes by increasing `num_workers` in `ScalingConfig`.\n- Swap in a different base model or dataset by updating the `model` argument and `load_dataset` call.\n- Experiment with different reward functions to train for other tasks."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/train/examples/transformers/transformer_reinforcement_learning/README.md
+279
View File
@@ -0,0 +1,279 @@
# RL Post-Training using Hugging Face TRL with GRPO
This notebook builds on the TRL GRPO Example, [GRPO Trainer](https://huggingface.co/docs/trl/en/grpo_trainer), to fine-tune a Qwen2.5 0.5B model to answer math questions using the DeepMath-103k dataset and Group Relative Policy Optimization (GRPO) algorithm.
Ray Train scales this implementation efficiently across multiple GPUs without changing training logic.
This notebook consists of the following steps:
1. Quick Summary of GRPO and the DeepMath dataset
2. Package and Environment Setup
3. Running TRL with Ray Train
4. Scaling to more GPUs
<div id="anyscale-note" class="alert alert-block alert-warning">
<strong>Anyscale Specific Configuration</strong>
<p><strong>Note:</strong> This tutorial is optimized for the Anyscale platform. When running on open source Ray, you need additional configuration. For example, you would need to manually:</p>
<ul>
<li><strong>Configure your Ray Cluster</strong>: Set up your multi-node environment and manage resource allocation without Anyscale's automation.</li> <li><strong>Manage Dependencies</strong>: Manually install and manage dependencies on each node.</li> <li><strong>Set Up Storage</strong>: Configure your own distributed or shared storage system for model checkpointing.</li>
</ul>
<p>All these configurations are handled automatically through the Anyscale platform.
</div>
<style>
div#anyscale-note > p, div#anyscale-note > ul, div#anyscale-note > ul li { color: black; }
div#anyscale-note { background-color: rgb(255, 243, 205); }
div#anyscale-note { border: 1px solid #ccc; border-radius: 8px; padding: 15px; }
</style>
## 1. Summary of the DeepMath dataset and GRPO algorithm
### DeepMath dataset
This example uses the DeepMath-103k dataset by [He et al., 2025](https://arxiv.org/abs/2504.11456), consisting of 103,000 challenging questions across a wide range of mathematical subjects including Algebra, Calculus, Number Theory, Geometry, Probability, and Discrete Mathematics. These questions are more heavily distributed towards challenging questions, in particular Levels 5 to 9, compared to alternative datasets, making them highly complex and difficult to solve without specialist training.
Each example contains two fields:
- **`prompt`**: a list of chat-format message dicts; the question text is in `prompt[0]["content"]`.
- **`solution`**: a LaTeX-formatted ground-truth answer, used by the reward function to verify the model's response.
Using a structured LaTeX format ensures the reward function can reliably parse and compare both the model's output and the ground truth. Here are a few examples from the dataset:
| Subject | Prompt | Solution |
|---|---|---|
| Number Theory | Determine the number of 19th power residues modulo 229. | $12$ |
| Functional Analysis | Determine the norm $\lVert T \rVert$ of the linear operator $T: \ell^2 \rightarrow \ell^2$ given by $(Tx)_1 = 0$, $(Tx)_n = -x_n + \alpha x_{n+1}$ for $n \geq 2$, where $\alpha \in \mathbb{C}$. | $1 + \lvert\alpha\rvert$ |
| Analysis | Determine whether there exists a Schwartz function $g \in \mathcal{S}(\mathbb{R}^n)$ such that for a given continuous function $f: \mathbb{R}^n \to \mathbb{R}$ with $f \not\equiv 0$, the integral $\int_{\mathbb{R}^n} g^2 f \, dx \neq 0$. | Yes |
### Group Relative Policy Optimization Algorithm
To RL post-train the model to answer math questions, this example uses TRL's GRPO (Group Relative Policy Optimization) implementation introduced by [Shao et al., 2024](https://arxiv.org/abs/2402.03300). GRPO is an online algorithm that iteratively improves by training on data the model generates itself. Each iteration proceeds through four steps:
1. **Generate**: For each prompt in the batch, sample a group of G completions from the current policy.
2. **Score**: Apply the reward function to each completion, producing a scalar reward per completion.
3. **Advantage**: Normalize the rewards within each group by subtracting the group mean and dividing by the group standard deviation. This relative score becomes the advantage $\hat{A}_{i,t}$ where positive means better than peers.
4. **Loss**: Update the policy to increase the probability of high-advantage completions and decrease it for low-advantage ones, while a KL penalty keeps the updated policy close to the reference.
$$
\mathcal{L}_{\text{GRPO}}(\theta) = -\frac{1}{\sum_{i=1}^{G} |o_i|} \sum_{i=1}^{G} \sum_{t=1}^{|o_i|} \left[ \frac{\pi_\theta(o_{i,t} \mid q, o_{i,\mathopen{<} t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,\mathopen{<} t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right],
$$
For a deeper, more technical, description of the algorithm, see the [TRL GRPO page](https://huggingface.co/docs/trl/en/grpo_trainer#looking-deeper-into-the-grpo-method) and the [original paper](https://arxiv.org/abs/2402.03300).
## 2. Package and Environment Setup
Install the required dependencies
```python
# !pip install "trl[vllm]" "math_verify" "transformers==4.57.6"
```
Use `ray.init()` to initialize a local cluster. By default, this cluster contains only the machine you are running this notebook on. You can also run this notebook on an Anyscale cluster.
```python
import ray
ray.init()
```
Use `ray.cluster_resources()` to check which resources your cluster has access to. If you're running this notebook on your local machine or Google Colab, you should see the number of CPU cores and GPUs available to you.
```python
from pprint import pprint
pprint(ray.cluster_resources())
```
Change these two variables to control whether the training uses CPUs or GPUs, and how many workers to spawn. Each worker claims one CPU or GPU, so make sure not to request more resources than are available. By default, the training runs with four GPU worker.
```python
use_gpu = True # set this to `False` to run on CPUs
num_workers = 4 # set this to the number of GPUs or CPUs you want to use
```
## 3. Using Hugging Face TRL (Transformer Reinforcement Learning) with Ray Train
In comparison to supervised pre-training where models optimize to minimize the error for their next token, RL-based post-training aims to maximize their reward from a prompt. Therefore, it's crucial to define a reward function to measure the success and to train a model.
This example uses the `trl.rewards.accuracy_reward` function to check whether the model's answer matches the answer in the dataset. As the answers use the LaTeX format, you must parse responses and solution before comparing them. The default `trl.rewards.accuracy_reward` implementation applies timeouts to the `parse` and `verify` functions, which are incompatible with Ray. The version defined below disables these timeouts by setting `parsing_timeout=0` and `timeout_seconds=0`.
```python
from latex2sympy2_extended import NormalizationConfig
from math_verify import LatexExtractionConfig, parse, verify
def accuracy_reward(completions, solution, **kwargs):
"""Reward function that checks mathematical accuracy.
This is a copy of `trl.rewards.accuracy_reward`.
The only difference is `parse(..., parsing_timeout=0)` and `verify(..., timeout_seconds=0)`
to avoid `signal.alarm()` issues with ray.
"""
contents = [completion[0]["content"] for completion in completions]
rewards = []
for content, sol in zip(contents, solution, strict=True):
gold_parsed = parse(sol, parsing_timeout=0)
if len(gold_parsed) != 0:
# We require the answer to be provided in correct latex (no malformed operators)
answer_parsed = parse(
content,
extraction_config=[
LatexExtractionConfig(
normalization_config=NormalizationConfig(units=True),
# Ensures that boxed is tried first
boxed_match_priority=0,
try_extract_without_anchor=False,
)
],
extraction_mode="first_match",
parsing_timeout=0,
)
reward = float(verify(gold_parsed, answer_parsed, timeout_seconds=0))
else:
# If the gold solution cannot be parsed, we assign `None` to skip this example
reward = None
rewards.append(reward)
return rewards
```
Next, build the training function that's distributed across all the workers. The `GRPOTrainer` implements the GRPO algorithm for rolling out the model, computing the advantage and backpropagation the losses.
With `vllm_mode="colocate"`, each worker runs a vLLM instance that consumes about 30% of the GPU memory. The alternative `"server"` mode dedicates one worker to generation with vLLM while the others perform learning, but this introduces inter-GPU communication overhead that can reduce throughput. See [this blog post](https://huggingface.co/blog/vllm-colocate) for more details.
Two Ray Train-specific additions integrate the HF Trainer into the Ray Train ecosystem:
- **`RayTrainReportCallback`** — hooks into the HF Trainer's logging to forward metrics and checkpoints to Ray Train after each training step. This is what populates the `Result` object returned by `trainer.fit()` with metrics like `reward` and the best checkpoint.
- **`prepare_trainer`** — configures the HF Trainer for distributed execution across Ray workers. It disables HF's built-in distributed setup so that Ray Train controls the process group instead, and ensures correct device placement on each worker.
```python
from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset
def train_func(config):
# Load the DeepMath dataset with only 100 elements for this example.
# For larger datasets, use Ray Data.
dataset = load_dataset("trl-lib/DeepMath-103K", split="train").shuffle(seed=42).select(range(100))
# Use vllm_mode="colocate" to allow easy scaling as each GPU handles their own vLLM instance
training_args = GRPOConfig(
# Compute a training batch size of 4 for each prompt
per_device_train_batch_size=4,
# Use vLLM, colocated on each GPU which uses 30% of the vRAM.
use_vllm=True,
vllm_mode="colocate",
vllm_gpu_memory_utilization=0.3,
# Run two training epochs over the whole dataset
num_train_epochs=2,
# Number of generations per prompt to sample, equivalent to G in the GRPO loss function.
num_generations=8,
# Save checkpoints and log metrics every epoch
save_strategy="epoch",
logging_strategy="epoch",
)
# GRPO Trainer
trainer = GRPOTrainer(
model="Qwen/Qwen2.5-0.5B",
args=training_args,
reward_funcs=accuracy_reward,
train_dataset=dataset,
)
# Report metrics and checkpoints to Ray Train
trainer.add_callback(RayTrainReportCallback())
# Prepare your trainer for TRL and Huggingface Transformer integration
trainer = prepare_trainer(trainer)
# Start Training
trainer.train()
```
With your `train_func` complete, you can now instantiate the `TorchTrainer`. Aside from calling the function, set the `scaling_config`, which controls the number of workers and resources used, and the `run_config` to configure checkpointing.
The `CheckpointConfig` controls how Ray Train saves and selects checkpoints during training:
- **`num_to_keep=1`** — Ray Train retains only the single best checkpoint on disk, saving storage.
- **`checkpoint_score_attribute="reward"`** — Ray Train ranks checkpoints by the `"reward"` metric, which is the mean reward across the batch as reported by `RayTrainReportCallback`.
- **`checkpoint_score_order="max"`** — higher reward is better, so Ray Train keeps the checkpoint with the highest reward seen across all training steps.
```python
from ray.train.torch import TorchTrainer
from ray.train import RunConfig, ScalingConfig, CheckpointConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
run_config=RunConfig(
# On Anyscale this path is auto-mounted shared storage.
# For open-source Ray, use a shared NFS path or s3:// URI accessible from all nodes.
storage_path="/mnt/cluster_storage/",
checkpoint_config=CheckpointConfig(
num_to_keep=1,
checkpoint_score_attribute="reward",
checkpoint_score_order="max",
),
),
)
```
Finally, call the `fit` method to start training with Ray Train. Save the `Result` object to a variable so you can access metrics and checkpoints.
```python
results = trainer.fit()
```
You can use the returned `Result` object to access metrics and the Ray Train `Checkpoint` associated with the last iteration.
```python
print(results.metrics_dataframe)
```
```python
print(results.get_best_checkpoint("reward", "max"))
```
## 4. Scaling to more GPUs or a larger model
The preceding example trains with 4 workers on a small 100-sample subset of the dataset. To scale up, adjust `num_workers` in `ScalingConfig` and the `num_workers` variable at the top of the notebook, along with the dataset size and model.
**Scaling workers:** Each worker claims one GPU. With `vllm_mode="colocate"`, each worker runs its own vLLM instance for generation and its own training process — generation requires no inter-GPU communication. Adding more workers increases the effective batch size and reduces time-to-convergence without any changes to training logic.
**Scaling the model:** Replace `"Qwen/Qwen2.5-0.5B"` with a larger checkpoint. Larger models require more GPU memory per worker. If a model doesn't fit on a single GPU, consider:
- Explore using DeepSpeed ZeRO to distribute the model weights across multiple GPUs.
- Reducing `vllm_gpu_memory_utilization` from `0.3` to leave more memory for model weights at the cost of a smaller vLLM KV cache.
**Fault Tolerance:** For longer runs where worker or node failures are possible, set `failure_config=FailureConfig(max_failures=1)` in `RunConfig`. Ray Train will automatically restart training up to twice on failure. Because this example already saves a checkpoint each epoch, a restarted run resumes from the last checkpoint rather than from scratch.
## Summary
This notebook demonstrated how to use Ray Train to scale RL post-training of a Qwen2.5 0.5B model across multiple GPUs using the TRL GRPO implementation and the DeepMath-103k dataset.
The key Ray Train integration points were:
- **`RayTrainReportCallback`** — forwards HF Trainer metrics and checkpoints to Ray Train after each step.
- **`prepare_trainer`** — configures the HF Trainer to run distributed training under Ray's process group.
- **`TorchTrainer`** — launches `train_func` on all workers with the configured resources.
- **`CheckpointConfig`** — automatically tracks and retains the best checkpoint based on the `"reward"` metric.
From here, you can:
- Scale to more GPUs or multiple nodes by increasing `num_workers` in `ScalingConfig`.
- Swap in a different base model or dataset by updating the `model` argument and `load_dataset` call.
- Experiment with different reward functions to train for other tasks.
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/aws.yaml
+12
View File
@@ -0,0 +1,12 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west-2
head_node_type:
name: head_node
instance_type: m5.2xlarge
worker_node_types:
- instance_type: g5.12xlarge # 4x A10G GPUs
name: '4xA10G:48CPU-192GB'
min_workers: 1
max_workers: 1
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/gce.yaml
+13
View File
@@ -0,0 +1,13 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-central1
head_node_type:
name: head
instance_type: n2-standard-8
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-48-nvidia-l4-4 # 4x L4 GPUs
min_workers: 1
max_workers: 1
use_spot: false
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/nb2py.py
+72
View File
@@ -0,0 +1,72 @@
#!/usr/bin/env python3
import argparse
import nbformat
def convert_notebook(input_path: str, output_path: str) -> None:
"""
Read a Jupyter notebook and write a Python script, converting all %%bash
cells and IPython "!" commands into subprocess.run calls that raise on error.
Cells that load or autoreload extensions are ignored.
"""
nb = nbformat.read(input_path, as_version=4)
with open(output_path, "w") as out:
for cell in nb.cells:
# Only process code cells
if cell.cell_type != "code":
continue
lines = cell.source.splitlines()
# Skip cells that load or autoreload extensions
if any(
l.strip().startswith("%load_ext autoreload")
or l.strip().startswith("%autoreload all")
for l in lines
):
continue
# Detect a %%bash cell
if lines and lines[0].strip().startswith("%%bash"):
bash_script = "\n".join(lines[1:]).rstrip()
out.write("import subprocess\n")
out.write(
f"subprocess.run(r'''{bash_script}''',\n"
" shell=True,\n"
" check=True,\n"
" executable='/bin/bash')\n\n"
)
else:
# Detect any IPython '!' shell commands in code lines
has_bang = any(line.lstrip().startswith("!") for line in lines)
if has_bang:
out.write("import subprocess\n")
for line in lines:
stripped = line.lstrip()
if stripped.startswith("!"):
cmd = stripped[1:].lstrip()
out.write(
f"subprocess.run(r'''{cmd}''',\n"
" shell=True,\n"
" check=True,\n"
" executable='/bin/bash')\n"
)
else:
out.write(line.rstrip() + "\n")
out.write("\n")
else:
# Regular Python cell: dump as-is
out.write(cell.source.rstrip() + "\n\n")
def main() -> None:
parser = argparse.ArgumentParser(
description="Convert a Jupyter notebook to a Python script, preserving bash cells and '!' commands as subprocess calls."
)
parser.add_argument("input_nb", help="Path to the input .ipynb file")
parser.add_argument("output_py", help="Path for the output .py script")
args = parser.parse_args()
convert_notebook(args.input_nb, args.output_py)
if __name__ == "__main__":
main()
doc/source/train/examples/transformers/transformer_reinforcement_learning/ci/tests.sh
+5
View File
@@ -0,0 +1,5 @@
#!/bin/bash
set -euo pipefail
python ci/nb2py.py README.ipynb README.py # convert notebook to py script
python README.py # run the converted python script
rm README.py # remove the generated script
doc/source/train/examples/transformers/transformer_reinforcement_learning/configs/aws.yaml
+9
View File
@@ -0,0 +1,9 @@
head_node_type:
name: head_node
instance_type: m5.2xlarge
worker_node_types:
- instance_type: g5.12xlarge # 4x A10G GPUs
name: '4xA10G:48CPU-192GB'
min_workers: 1
max_workers: 1
doc/source/train/examples/transformers/transformer_reinforcement_learning/configs/gce.yaml
+10
View File
@@ -0,0 +1,10 @@
head_node_type:
name: head
instance_type: n2-standard-8
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-48-nvidia-l4-4 # 4x L4 GPUs
min_workers: 1
max_workers: 1
use_spot: false
doc/source/train/examples/transformers/transformers_torch_trainer_basic.rst
+28
View File
@@ -0,0 +1,28 @@
:orphan:
.. _transformers_torch_trainer_basic_example:
Fine-tune a Text Classifier with Hugging Face Transformers
==========================================================
.. raw:: html
<a id="try-anyscale-quickstart-transformers_torch_trainer_basic" target="_blank" href="https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=transformers_torch_trainer_basic">
<img src="../../../_static/img/run-on-anyscale.svg" alt="Run on Anyscale" />
<br/><br/>
</a>
This basic example of distributed training with Ray Train and Hugging Face (HF) Transformers
fine-tunes a text classifier on the Yelp review dataset using HF Transformers and Ray Train.
Code example
------------
.. literalinclude:: /../../python/ray/train/examples/transformers/transformers_torch_trainer_basic.py
See also
--------
* :ref:`Get Started with Hugging Face Transformers <train-pytorch-transformers>` for a tutorial
* :doc:`Ray Train Examples <../../examples>` for more use cases
doc/source/train/getting-started-jax.rst
+376
View File
@@ -0,0 +1,376 @@
.. _train-jax:
Get Started with Distributed Training using JAX
===============================================
This guide provides an overview of the :class:`~ray.train.v2.jax.JaxTrainer` in Ray Train.
What is JAX?
------------
`JAX <https://github.com/jax-ml/jax>`_ is a Python library for accelerator-oriented array computation and
program transformation, designed for high-performance numerical computing and large-scale machine learning.
JAX provides an extensible system for transforming numerical functions such as `jax.grad`, `jax.jit`, and `jax.vmap`,
utilizing the XLA compiler to create highly optimized code that scales efficiently on accelerators like GPUs and TPUs.
The core power of JAX lies in its composability, allowing these transformations to be combined to build complex,
high-performance numerical programs for distributed execution.
JAX and :class:`~ray.train.v2.jax.JaxTrainer` support different accelerators such as GPUs and TPUs. For more details, see `JAX Supported platforms <https://docs.jax.dev/en/latest/installation.html#supported-platforms>`_.
What are TPUs?
--------------
Tensor Processing Units (TPUs), are custom-designed accelerators created by Google to optimize machine learning
workloads. Unlike general-purpose CPUs or parallel-processing GPUs, TPUs are highly specialized for the massive
matrix and tensor computations involved in deep learning, making them exceptionally efficient.
The primary advantage of TPUs is performance at scale because they're designed to be connected into large, multi-host
configurations called “PodSlices” via a high-speed ICI interconnect, making them ideal for training large models
that are unable to fit on a single node.
To learn more about configuring TPUs with KubeRay, see :ref:`kuberay-tpu`.
JaxTrainer API
--------------
The :class:`~ray.train.v2.jax.JaxTrainer` is the core component for orchestrating distributed JAX training in Ray Train.
It follows the Single-Program, Multi-Data (SPMD) paradigm, where your training code is executed simultaneously
across multiple workers.
For TPUs, each worker runs on a separate TPU virtual machine within a TPU slice. Ray Train automatically
handles atomically reserving TPU slices.
For GPUs, Ray automatically sets up the JAX distributed system on CUDA devices.
You initialize the `JaxTrainer` with your training logic, defined in a `train_loop_per_worker` function, and a
`ScalingConfig` that specifies the distributed hardware layout. The `JaxTrainer` supports both **Google Cloud TPUs** and **NVIDIA GPUs**.
Configure scale and accelerators
--------------------------------
TPU scaling configuration
^^^^^^^^^^^^^^^^^^^^^^^^^
For TPU training, use :class:`~ray.train.ScalingConfig` to define your TPU slices configuration. Key fields include:
* :class:`use_tpu <ray.train.ScalingConfig>`: It's a new field added in Ray 2.49.0 to the V2 `ScalingConfig`. This boolean flag tells Ray Train to initialize the JAX backend for TPU execution.
* :class:`topology <ray.train.ScalingConfig>`: It's a new field added in Ray 2.49.0 to the V2 `ScalingConfig`. Topology is a string defining the physical arrangement of the TPU chips (for example, "4x4"). It's required for multi-host training and ensures Ray places workers correctly across the slice. For a list of supported TPU topologies by generation,
see the `GKE documentation <https://cloud.google.com/kubernetes-engine/docs/concepts/plan-tpus#topology>`_.
* :class:`num_workers <ray.train.ScalingConfig>`: Set this to the total number of TPU VMs across all slices. For example, one v4-32 slice with a 2x2x4 topology uses 4 VMs, so set `num_workers` to 4. If you use two v4-32 slices, set `num_workers` to 8.
* :class:`resources_per_worker <ray.train.ScalingConfig>`: A dictionary specifying the resources each worker needs. For TPUs, you typically request the number of chips per VM (for example, `{"TPU": 4}`).
* :class:`accelerator_type <ray.train.ScalingConfig>`: For TPUs, `accelerator_type` specifies the TPU generation you're using (for example, "TPU-V6E"), ensuring your workload is scheduled on the desired TPU slice.
.. testcode::
:skipif: True
from ray.train import ScalingConfig
tpu_scaling_config = ScalingConfig(num_workers=4, use_tpu=True, topology="4x4", accelerator_type="TPU-V6E")
GPU scaling configuration
^^^^^^^^^^^^^^^^^^^^^^^^^
For GPU training, use :class:`~ray.train.ScalingConfig` to define your GPU configuration.
Each worker is one Ray Train process. By default, this will request one GPU per worker.
* :class:`num_workers <ray.train.ScalingConfig>`: The number of distributed training worker processes.
* :class:`use_gpu <ray.train.ScalingConfig>`: Whether each worker should use a GPU.
* :class:`resources_per_worker <ray.train.ScalingConfig>`: A dictionary specifying the resources each worker needs.
.. testcode::
from ray.train import ScalingConfig
gpu_scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
For more details, see :ref:`train_scaling_config`.
Quickstart
----------
For reference, the final code is as follows:
.. testcode::
:skipif: True
from ray.train.v2.jax import JaxTrainer
from ray.train import ScalingConfig
def train_func():
# Your JAX training code here.
# Define the TPU scaling configuration with `use_tpu=True`.
scaling_config = ScalingConfig(num_workers=4, use_tpu=True, topology="4x4", accelerator_type="TPU-V6E")
# Define the GPU scaling configuration with `use_gpu=True`.
# scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
# Choose one scaling config.
trainer = JaxTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
1. `train_func` is the Python code that executes on each distributed training worker.
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use TPUs or GPUs.
3. :class:`~ray.train.v2.jax.JaxTrainer` launches the distributed training job.
Compare a JAX training script with and without Ray Train.
.. tab-set::
.. tab-item:: JAX + Ray Train
.. testcode::
:skipif: True
import jax
import jax.numpy as jnp
import optax
import ray.train
from ray.train.v2.jax import JaxTrainer
from ray.train import ScalingConfig
def train_func():
"""This function is run on each distributed worker."""
key = jax.random.PRNGKey(jax.process_index())
X = jax.random.normal(key, (100, 1))
noise = jax.random.normal(key, (100, 1)) * 0.1
y = 2 * X + 1 + noise
def linear_model(params, x):
return x @ params['w'] + params['b']
def loss_fn(params, x, y):
preds = linear_model(params, x)
return jnp.mean((preds - y) ** 2)
@jax.jit
def train_step(params, opt_state, x, y):
loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return params, opt_state, loss
# Initialize parameters and optimizer.
key, w_key, b_key = jax.random.split(key, 3)
params = {'w': jax.random.normal(w_key, (1, 1)), 'b': jax.random.normal(b_key, (1,))}
optimizer = optax.adam(learning_rate=0.01)
opt_state = optimizer.init(params)
# Training loop
epochs = 100
for epoch in range(epochs):
params, opt_state, loss = train_step(params, opt_state, X, y)
# Report metrics back to Ray Train.
ray.train.report({"loss": float(loss), "epoch": epoch})
# Define the TPU scaling configuration for your distributed job.
scaling_config = ScalingConfig(
num_workers=4,
use_tpu=True,
topology="4x4",
accelerator_type="TPU-V6E",
placement_strategy="SPREAD"
)
# Define the GPU scaling configuration with `use_gpu=True`.
# scaling_config = ScalingConfig(
# num_workers=4,
# use_gpu=True,
# )
# Define and run the JaxTrainer.
trainer = JaxTrainer(
train_loop_per_worker=train_func,
scaling_config=scaling_config,
)
result = trainer.fit()
print(f"Training finished. Final loss: {result.metrics['loss']:.4f}")
.. tab-item:: JAX
.. This snippet isn't tested because it doesn't use any Ray code.
.. testcode::
:skipif: True
import jax
import jax.numpy as jnp
import optax
# In a non-Ray script, you would manually initialize the
# distributed environment for multi-host training.
# import jax.distributed
# jax.distributed.initialize()
# Generate synthetic data.
key = jax.random.PRNGKey(0)
X = jax.random.normal(key, (100, 1))
noise = jax.random.normal(key, (100, 1)) * 0.1
y = 2 * X + 1 + noise
# Model and loss function are standard JAX.
def linear_model(params, x):
return x @ params['w'] + params['b']
def loss_fn(params, x, y):
preds = linear_model(params, x)
return jnp.mean((preds - y) ** 2)
@jax.jit
def train_step(params, opt_state, x, y):
loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return params, opt_state, loss
# Initialize parameters and optimizer.
key, w_key, b_key = jax.random.split(key, 3)
params = {'w': jax.random.normal(w_key, (1, 1)), 'b': jax.random.normal(b_key, (1,))}
optimizer = optax.adam(learning_rate=0.01)
opt_state = optimizer.init(params)
# Training loop
epochs = 100
print("Starting training...")
for epoch in range(epochs):
params, opt_state, loss = train_step(params, opt_state, X, y)
if epoch % 10 == 0:
print(f"Epoch {epoch}, Loss: {loss:.4f}")
print("Training finished.")
print(f"Learned parameters: w={params['w'].item():.4f}, b={params['b'].item():.4f}")
Set up a training function
--------------------------
Ray Train automatically initializes the JAX distributed environment based on the `ScalingConfig` and the `JAX_PLATFORMS` environment variable.
To adapt your existing JAX code, you simply need to wrap your training logic in a Python function
that can be passed to the `JaxTrainer`.
This function is the entry point that Ray will execute on each remote worker.
.. code-block:: diff
+from ray.train.v2.jax import JaxTrainer
+from ray.train import ScalingConfig, report
-def main_logic()
+def train_func():
"""This function is run on each distributed worker."""
# ... (JAX model, data, and training step definitions) ...
# Training loop
for epoch in range(epochs):
params, opt_state, loss = train_step(params, opt_state, X, y)
- print(f"Epoch {epoch}, Loss: {loss:.4f}")
+ # In Ray Train, you can report metrics back to the trainer
+ report({"loss": float(loss), "epoch": epoch})
-if __name__ == "__main__":
- main_logic()
+# Define the hardware configuration for your distributed job.
+scaling_config = ScalingConfig(
+ num_workers=4,
+ use_tpu=True,
+ topology="4x4",
+ accelerator_type="TPU-V6E",
+ placement_strategy="SPREAD"
+)
+
+# Define the GPU scaling configuration with `use_gpu=True`.
+# scaling_config = ScalingConfig(
+# num_workers=4,
+# use_gpu=True,
+# resources_per_worker={"GPU": 1},
+# )
+
+# Define and run the JaxTrainer, which executes `train_func`.
+trainer = JaxTrainer(
+ train_loop_per_worker=train_func,
+ scaling_config=scaling_config
+)
+result = trainer.fit()
Configure persistent storage
----------------------------
Create a :class:`~ray.train.RunConfig` object to specify the path where results
(including checkpoints and artifacts) will be saved.
.. testcode::
from ray.train import RunConfig
# Local path (/some/local/path/unique_run_name)
run_config = RunConfig(storage_path="/some/local/path", name="unique_run_name")
# Shared cloud storage URI (s3://bucket/unique_run_name)
run_config = RunConfig(storage_path="s3://bucket", name="unique_run_name")
# Shared NFS path (/mnt/nfs/unique_run_name)
run_config = RunConfig(storage_path="/mnt/nfs", name="unique_run_name")
.. warning::
Specifying a *shared storage location* (such as cloud storage or NFS) is
*optional* for single-node clusters, but it is **required for multi-node clusters.**
Using a local path will :ref:`raise an error <multinode-local-storage-warning>`
during checkpointing for multi-node clusters.
For more details, see :ref:`persistent-storage-guide`.
Launch a training job
---------------------
Tying it all together, you can now launch a distributed training job with a :class:`~ray.train.v2.jax.JaxTrainer`.
.. testcode::
:skipif: True
from ray.train import ScalingConfig
train_func = lambda: None
# Define the TPU scaling configuration with `use_tpu=True`.
scaling_config = ScalingConfig(num_workers=4, use_tpu=True, topology="4x4", accelerator_type="TPU-V6E")
# Define the GPU scaling configuration with `use_gpu=True`.
# scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
run_config = None
.. testcode::
:skipif: True
from ray.train.v2.jax import JaxTrainer
trainer = JaxTrainer(
train_func, scaling_config=scaling_config, run_config=run_config
)
result = trainer.fit()
Access training results
-----------------------
After training completes, a :class:`~ray.train.Result` object is returned which contains
information about the training run, including the metrics and checkpoints reported during training.
.. testcode::
:skipif: True
result.metrics # The metrics reported during training.
result.checkpoint # The latest checkpoint reported during training.
result.path # The path where logs are stored.
result.error # The exception that was raised, if training failed.
For more usage examples, see :ref:`train-inspect-results`.
Next steps
----------
After you have converted your JAX training script to use Ray Train:
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
* Consult the :ref:`API Reference <train-api>` for more details on the classes and methods from this tutorial.
doc/source/train/getting-started-lightgbm.rst
+395
View File
@@ -0,0 +1,395 @@
.. _train-lightgbm:
Get Started with Distributed Training using LightGBM
====================================================
This tutorial walks through the process of converting an existing LightGBM script to use Ray Train.
Learn how to:
1. Configure a :ref:`training function <train-overview-training-function>` to report metrics and save checkpoints.
2. Configure :ref:`scaling <train-overview-scaling-config>` and CPU or GPU resource requirements for a training job.
3. Launch a distributed training job with a :class:`~ray.train.lightgbm.LightGBMTrainer`.
Quickstart
----------
For reference, the final code will look something like this:
.. testcode::
:skipif: True
import ray.train
from ray.train.lightgbm import LightGBMTrainer
def train_func():
# Your LightGBM training code here.
...
scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 4})
trainer = LightGBMTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
1. `train_func` is the Python code that executes on each distributed training worker.
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use GPUs.
3. :class:`~ray.train.lightgbm.LightGBMTrainer` launches the distributed training job.
Compare a LightGBM training script with and without Ray Train.
.. tab-set::
.. tab-item:: LightGBM + Ray Train
.. literalinclude:: ./doc_code/lightgbm_quickstart.py
:language: python
:start-after: __lightgbm_ray_start__
:end-before: __lightgbm_ray_end__
.. tab-item:: LightGBM
.. literalinclude:: ./doc_code/lightgbm_quickstart.py
:language: python
:start-after: __lightgbm_start__
:end-before: __lightgbm_end__
Set up a training function
--------------------------
First, update your training code to support distributed training.
Begin by wrapping your `native <https://lightgbm.readthedocs.io/en/latest/Python-Intro.html>`_
or `scikit-learn estimator <https://lightgbm.readthedocs.io/en/latest/Python-API.html#scikit-learn-api>`_
LightGBM training code in a :ref:`training function <train-overview-training-function>`:
.. testcode::
:skipif: True
def train_func():
# Your native LightGBM training code here.
train_set = ...
lightgbm.train(...)
Each distributed training worker executes this function.
You can also specify the input argument for `train_func` as a dictionary via the Trainer's `train_loop_config`. For example:
.. testcode:: python
:skipif: True
def train_func(config):
label_column = config["label_column"]
num_boost_round = config["num_boost_round"]
...
config = {"label_column": "target", "num_boost_round": 100}
trainer = ray.train.lightgbm.LightGBMTrainer(train_func, train_loop_config=config, ...)
.. warning::
Avoid passing large data objects through `train_loop_config` to reduce the
serialization and deserialization overhead. Instead,
initialize large objects (e.g. datasets, models) directly in `train_func`.
.. code-block:: diff
def load_dataset():
# Return a large in-memory dataset
...
def load_model():
# Return a large in-memory model instance
...
-config = {"data": load_dataset(), "model": load_model()}
def train_func(config):
- data = config["data"]
- model = config["model"]
+ data = load_dataset()
+ model = load_model()
...
trainer = ray.train.lightgbm.LightGBMTrainer(train_func, train_loop_config=config, ...)
Configure distributed training parameters
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To enable distributed LightGBM training, add network communication parameters to your training configuration using :func:`ray.train.lightgbm.get_network_params`.
This function automatically configures the necessary network settings for worker communication:
.. code-block:: diff
def train_func():
...
params = {
# Your LightGBM training parameters here
...
+ "tree_learner": "data_parallel",
+ "pre_partition": True,
+ **ray.train.lightgbm.get_network_params(),
}
model = lightgbm.train(
params,
...
)
...
.. note::
Make sure to set ``tree_learner`` to enable distributed training. See the `LightGBM documentation <https://lightgbm.readthedocs.io/en/latest/Parallel-Learning-Guide.html#tree-learner>`_ for more details.
You should also set ``pre_partition=True`` if using Ray Data to load and shard your dataset, as shown in the quickstart example.
Report metrics and save checkpoints
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To persist your checkpoints and monitor training progress, add a
:class:`ray.train.lightgbm.RayTrainReportCallback` utility callback to your Trainer:
.. testcode:: python
:skipif: True
import lightgbm
from ray.train.lightgbm import RayTrainReportCallback
def train_func():
...
bst = lightgbm.train(
...,
callbacks=[
RayTrainReportCallback(
metrics=["eval-multi_logloss"], frequency=1
)
],
)
...
Reporting metrics and checkpoints to Ray Train enables :ref:`fault-tolerant training <train-fault-tolerance>` and the integration with Ray Tune.
Loading data
------------
When running distributed LightGBM training, each worker should use a different shard of the dataset.
.. testcode:: python
:skipif: True
def get_train_dataset(world_rank: int) -> lightgbm.Dataset:
# Define logic to get the Dataset shard for this worker rank
...
def get_eval_dataset(world_rank: int) -> lightgbm.Dataset:
# Define logic to get the Dataset for each worker
...
def train_func():
rank = ray.train.get_world_rank()
train_set = get_train_dataset(rank)
eval_set = get_eval_dataset(rank)
...
A common way to do this is to pre-shard the dataset and then assign each worker a different set of files to read.
Pre-sharding the dataset is not very flexible to changes in the number of workers, since some workers may be assigned more data than others. For more flexibility, Ray Data provides a solution for sharding the dataset at runtime.
Use Ray Data to shard the dataset
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
:ref:`Ray Data <data>` is a distributed data processing library that allows you to easily shard and distribute your data across multiple workers.
First, load your **entire** dataset as a Ray Data Dataset.
Reference the :ref:`Ray Data Quickstart <data_quickstart>` for more details on how to load and preprocess data from different sources.
.. testcode:: python
:skipif: True
train_dataset = ray.data.read_parquet("s3://path/to/entire/train/dataset/dir")
eval_dataset = ray.data.read_parquet("s3://path/to/entire/eval/dataset/dir")
In the training function, you can access the dataset shards for this worker using :meth:`ray.train.get_dataset_shard`.
Convert this into a native `lightgbm.Dataset <https://lightgbm.readthedocs.io/en/latest/Python-Intro.html#dataset>`_.
.. testcode:: python
:skipif: True
from ray.train.lightgbm import normalize_pandas_for_lightgbm
def get_dataset(dataset_name: str) -> lightgbm.Dataset:
shard = ray.train.get_dataset_shard(dataset_name)
df = normalize_pandas_for_lightgbm(shard.materialize().to_pandas())
X, y = df.drop("target", axis=1), df["target"]
return lightgbm.Dataset(X, label=y)
def train_func():
train_set = get_dataset("train")
eval_set = get_dataset("eval")
...
.. note::
Starting in Ray 2.56, Ray Data preserves Arrow-backed pandas dtypes when
converting Arrow blocks to pandas, for example ``int64[pyarrow]``. LightGBM's
pandas input validation rejects these dtypes, so a pandas DataFrame coming
from Ray Data must be normalized before being passed to ``lightgbm.Dataset``.
:func:`ray.train.lightgbm.normalize_pandas_for_lightgbm` maps Arrow-backed
numeric/boolean columns to NumPy-nullable equivalents and leaves all other
columns untouched. Prefer it over
``df.convert_dtypes(dtype_backend="numpy_nullable")``, which scans every
value in every column and also rewrites NumPy-backed columns into nullable
equivalents even when no Arrow dtypes are present.
Finally, pass the dataset to the Trainer. This will automatically shard the dataset across the workers. These keys must match the keys used when calling ``get_dataset_shard`` in the training function.
.. testcode:: python
:skipif: True
trainer = LightGBMTrainer(..., datasets={"train": train_dataset, "eval": eval_dataset})
trainer.fit()
For more details, see :ref:`data-ingest-torch`.
Configure scale and GPUs
------------------------
Outside of your training function, create a :class:`~ray.train.ScalingConfig` object to configure:
1. :class:`num_workers <ray.train.ScalingConfig>` - The number of distributed training worker processes.
2. :class:`use_gpu <ray.train.ScalingConfig>` - Whether each worker should use a GPU (or CPU).
3. :class:`resources_per_worker <ray.train.ScalingConfig>` - The number of CPUs or GPUs per worker.
.. testcode::
from ray.train import ScalingConfig
# 4 nodes with 8 CPUs each.
scaling_config = ScalingConfig(num_workers=4, resources_per_worker={"CPU": 8})
.. note::
When using Ray Data with Ray Train, be careful not to request all available CPUs in your cluster with the `resources_per_worker` parameter.
Ray Data needs CPU resources to execute data preprocessing operations in parallel.
If all CPUs are allocated to training workers, Ray Data operations may be bottlenecked, leading to reduced performance.
A good practice is to leave some portion of CPU resources available for Ray Data operations.
For example, if your cluster has 8 CPUs per node, you might allocate 6 CPUs to training workers and leave 2 CPUs for Ray Data:
.. testcode::
# Allocate 6 CPUs per worker, leaving resources for Ray Data operations
scaling_config = ScalingConfig(num_workers=4, resources_per_worker={"CPU": 6})
In order to use GPUs, you will need to set the `use_gpu` parameter to `True` in your :class:`~ray.train.ScalingConfig` object.
This will request and assign a single GPU per worker.
.. testcode::
# 1 node with 8 CPUs and 4 GPUs each.
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
# 4 nodes with 8 CPUs and 4 GPUs each.
scaling_config = ScalingConfig(num_workers=16, use_gpu=True)
When using GPUs, you will also need to update your training function to use the assigned GPU.
This can be done by setting the `"device"` parameter as `"gpu"`.
For more details on LightGBM's GPU support, see the `LightGBM GPU documentation <https://lightgbm.readthedocs.io/en/latest/GPU-Tutorial.html>`__.
.. code-block:: diff
def train_func():
...
params = {
...,
+ "device": "gpu",
}
bst = lightgbm.train(
params,
...
)
Configure persistent storage
----------------------------
Create a :class:`~ray.train.RunConfig` object to specify the path where results
(including checkpoints and artifacts) will be saved.
.. testcode::
from ray.train import RunConfig
# Local path (/some/local/path/unique_run_name)
run_config = RunConfig(storage_path="/some/local/path", name="unique_run_name")
# Shared cloud storage URI (s3://bucket/unique_run_name)
run_config = RunConfig(storage_path="s3://bucket", name="unique_run_name")
# Shared NFS path (/mnt/nfs/unique_run_name)
run_config = RunConfig(storage_path="/mnt/nfs", name="unique_run_name")
.. warning::
Specifying a *shared storage location* (such as cloud storage or NFS) is
*optional* for single-node clusters, but it is **required for multi-node clusters.**
Using a local path will :ref:`raise an error <multinode-local-storage-warning>`
during checkpointing for multi-node clusters.
For more details, see :ref:`persistent-storage-guide`.
Launch a training job
---------------------
Tying it all together, you can now launch a distributed training job with a :class:`~ray.train.lightgbm.LightGBMTrainer`.
.. testcode::
:hide:
from ray.train import ScalingConfig
train_func = lambda: None
scaling_config = ScalingConfig(num_workers=1)
run_config = None
.. testcode::
from ray.train.lightgbm import LightGBMTrainer
trainer = LightGBMTrainer(
train_func, scaling_config=scaling_config, run_config=run_config
)
result = trainer.fit()
Access training results
-----------------------
After training completes, a :class:`~ray.train.Result` object is returned which contains
information about the training run, including the metrics and checkpoints reported during training.
.. testcode::
result.metrics # The metrics reported during training.
result.checkpoint # The latest checkpoint reported during training.
result.path # The path where logs are stored.
result.error # The exception that was raised, if training failed.
For more usage examples, see :ref:`train-inspect-results`.
Next steps
----------
After you have converted your LightGBM training script to use Ray Train:
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
* Consult the :ref:`API Reference <train-api>` for more details on the classes and methods from this tutorial.
doc/source/train/getting-started-pytorch-lightning.rst
+533
View File
@@ -0,0 +1,533 @@
.. _train-pytorch-lightning:
Get Started with Distributed Training using PyTorch Lightning
=============================================================
This tutorial walks through the process of converting an existing PyTorch Lightning script to use Ray Train.
Learn how to:
1. Configure the Lightning Trainer so that it runs distributed with Ray and on the correct CPU or GPU device.
2. Configure :ref:`training function <train-overview-training-function>` to report metrics and save checkpoints.
3. Configure :ref:`scaling <train-overview-scaling-config>` and CPU or GPU resource requirements for a training job.
4. Launch a distributed training job with a :class:`~ray.train.torch.TorchTrainer`.
Quickstart
----------
For reference, the final code is as follows:
.. testcode::
:skipif: True
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
def train_func():
# Your PyTorch Lightning training code here.
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
1. `train_func` is the Python code that executes on each distributed training worker.
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use GPUs.
3. :class:`~ray.train.torch.TorchTrainer` launches the distributed training job.
Compare a PyTorch Lightning training script with and without Ray Train.
.. tab-set::
.. tab-item:: PyTorch Lightning + Ray Train
.. code-block:: python
:emphasize-lines: 11-12, 38, 52-57, 59, 63, 66-73
import os
import tempfile
import torch
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
import lightning.pytorch as pl
import ray.train.lightning
from ray.train.torch import TorchTrainer
# Model, Loss, Optimizer
class ImageClassifier(pl.LightningModule):
def __init__(self):
super(ImageClassifier, self).__init__()
self.model = resnet18(num_classes=10)
self.model.conv1 = torch.nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
self.criterion = torch.nn.CrossEntropyLoss()
def forward(self, x):
return self.model(x)
def training_step(self, batch, batch_idx):
x, y = batch
outputs = self.forward(x)
loss = self.criterion(outputs, y)
self.log("loss", loss, on_step=True, prog_bar=True)
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.model.parameters(), lr=0.001)
def train_func():
# Data
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
data_dir = os.path.join(tempfile.gettempdir(), "data")
train_data = FashionMNIST(root=data_dir, train=True, download=True, transform=transform)
train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
# Training
model = ImageClassifier()
# [1] Configure PyTorch Lightning Trainer.
trainer = pl.Trainer(
max_epochs=10,
devices="auto",
accelerator="auto",
strategy=ray.train.lightning.RayDDPStrategy(),
plugins=[ray.train.lightning.RayLightningEnvironment()],
callbacks=[ray.train.lightning.RayTrainReportCallback()],
# [1a] Optionally, disable the default checkpointing behavior
# in favor of the `RayTrainReportCallback` above.
enable_checkpointing=False,
)
trainer = ray.train.lightning.prepare_trainer(trainer)
trainer.fit(model, train_dataloaders=train_dataloader)
# [2] Configure scaling and resource requirements.
scaling_config = ray.train.ScalingConfig(num_workers=2, use_gpu=True)
# [3] Launch distributed training job.
trainer = TorchTrainer(
train_func,
scaling_config=scaling_config,
# [3a] If running in a multi-node cluster, this is where you
# should configure the run's persistent storage that is accessible
# across all worker nodes.
# run_config=ray.train.RunConfig(storage_path="s3://..."),
)
result: ray.train.Result = trainer.fit()
# [4] Load the trained model.
with result.checkpoint.as_directory() as checkpoint_dir:
model = ImageClassifier.load_from_checkpoint(
os.path.join(
checkpoint_dir,
ray.train.lightning.RayTrainReportCallback.CHECKPOINT_NAME,
),
)
.. tab-item:: PyTorch Lightning
.. This snippet isn't tested because it doesn't use any Ray code.
.. testcode::
:skipif: True
import torch
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
from torch.utils.data import DataLoader
import lightning.pytorch as pl
# Model, Loss, Optimizer
class ImageClassifier(pl.LightningModule):
def __init__(self):
super(ImageClassifier, self).__init__()
self.model = resnet18(num_classes=10)
self.model.conv1 = torch.nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
self.criterion = torch.nn.CrossEntropyLoss()
def forward(self, x):
return self.model(x)
def training_step(self, batch, batch_idx):
x, y = batch
outputs = self.forward(x)
loss = self.criterion(outputs, y)
self.log("loss", loss, on_step=True, prog_bar=True)
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.model.parameters(), lr=0.001)
# Data
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
train_data = FashionMNIST(root='./data', train=True, download=True, transform=transform)
train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
# Training
model = ImageClassifier()
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, train_dataloaders=train_dataloader)
Set up a training function
--------------------------
.. include:: ./common/torch-configure-train_func.rst
Ray Train sets up your distributed process group on each worker. You only need to
make a few changes to your Lightning Trainer definition.
.. code-block:: diff
import lightning.pytorch as pl
-from pl.strategies import DDPStrategy
-from pl.plugins.environments import LightningEnvironment
+import ray.train.lightning
def train_func():
...
model = MyLightningModule(...)
datamodule = MyLightningDataModule(...)
trainer = pl.Trainer(
- devices=[0, 1, 2, 3],
- strategy=DDPStrategy(),
- plugins=[LightningEnvironment()],
+ devices="auto",
+ accelerator="auto",
+ strategy=ray.train.lightning.RayDDPStrategy(),
+ plugins=[ray.train.lightning.RayLightningEnvironment()]
)
+ trainer = ray.train.lightning.prepare_trainer(trainer)
trainer.fit(model, datamodule=datamodule)
The following sections discuss each change.
Configure the distributed strategy
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Ray Train offers several sub-classed distributed strategies for Lightning.
These strategies retain the same argument list as their base strategy classes.
Internally, they configure the root device and the distributed
sampler arguments.
- :class:`~ray.train.lightning.RayDDPStrategy`
- :class:`~ray.train.lightning.RayFSDPStrategy`
- :class:`~ray.train.lightning.RayDeepSpeedStrategy`
.. code-block:: diff
import lightning.pytorch as pl
-from pl.strategies import DDPStrategy
+import ray.train.lightning
def train_func():
...
trainer = pl.Trainer(
...
- strategy=DDPStrategy(),
+ strategy=ray.train.lightning.RayDDPStrategy(),
...
)
...
Configure the Ray cluster environment plugin
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Ray Train also provides a :class:`~ray.train.lightning.RayLightningEnvironment` class
as a specification for the Ray Cluster. This utility class configures the worker's
local, global, and node rank and world size.
.. code-block:: diff
import lightning.pytorch as pl
-from pl.plugins.environments import LightningEnvironment
+import ray.train.lightning
def train_func():
...
trainer = pl.Trainer(
...
- plugins=[LightningEnvironment()],
+ plugins=[ray.train.lightning.RayLightningEnvironment()],
...
)
...
Configure parallel devices
^^^^^^^^^^^^^^^^^^^^^^^^^^
In addition, Ray TorchTrainer has already configured the correct
``CUDA_VISIBLE_DEVICES`` for you. One should always use all available
GPUs by setting ``devices="auto"`` and ``acelerator="auto"``.
.. code-block:: diff
import lightning.pytorch as pl
def train_func():
...
trainer = pl.Trainer(
...
- devices=[0,1,2,3],
+ devices="auto",
+ accelerator="auto",
...
)
...
Report checkpoints and metrics
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To persist your checkpoints and monitor training progress, add a
:class:`ray.train.lightning.RayTrainReportCallback` utility callback to your Trainer.
.. code-block:: diff
import lightning.pytorch as pl
from ray.train.lightning import RayTrainReportCallback
def train_func():
...
trainer = pl.Trainer(
...
- callbacks=[...],
+ callbacks=[..., RayTrainReportCallback()],
)
...
Reporting metrics and checkpoints to Ray Train enables you to support :ref:`fault-tolerant training <train-fault-tolerance>` and :ref:`hyperparameter optimization <train-tune>`.
You can also configure :ref:`asynchronous checkpointing <train-checkpoint-upload-mode-async>` and
:ref:`asynchronous validation <train-validating-checkpoints>` through the callback.
The ``checkpoint_upload_mode`` flag offloads checkpoint uploading to a Ray Train managed background thread
instead of blocking the Lightning training loop. The ``validation`` flag launches an asynchronous
Ray task to validate the checkpoint instead of running ``validation_step`` synchronously in the training workers.
Note that this is incompatible with Lightning's `AsyncCheckpointIO <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.plugins.io.AsyncCheckpointIO.html>`_ plugin
because Ray Train needs to control the upload thread in order to wait for it to finish before committing the checkpoint.
.. code-block:: diff
import lightning.pytorch as pl
from ray.train.lightning import RayTrainReportCallback
+from ray.train import CheckpointUploadMode
+from ray.train import ValidationConfig, ValidationTaskConfig
def train_func():
...
trainer = pl.Trainer(
...
- callbacks=[..., RayTrainReportCallback()],
+ callbacks=[..., RayTrainReportCallback(
+ checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
+ validation=ValidationTaskConfig(fn_kwargs={}),
+ )],
)
...
+def validation_fn(checkpoint):
+ # Load and validate the checkpoint, return metrics.
+ return {"val_score": ...}
+
+trainer = TorchTrainer(
+ train_func,
+ validation_config=ValidationConfig(fn=validation_fn),
+ ...
+)
Note that the :class:`ray.train.lightning.RayTrainReportCallback` class only provides a simple implementation, and can be :ref:`further customized <train-dl-saving-checkpoints>`.
Prepare your Lightning Trainer
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Finally, pass your Lightning Trainer into
:meth:`~ray.train.lightning.prepare_trainer` to validate
your configurations.
.. code-block:: diff
import lightning.pytorch as pl
import ray.train.lightning
def train_func():
...
trainer = pl.Trainer(...)
+ trainer = ray.train.lightning.prepare_trainer(trainer)
...
.. include:: ./common/torch-configure-run.rst
Next steps
----------
After you have converted your PyTorch Lightning training script to use Ray Train:
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
* Consult the :ref:`API Reference <train-api>` for more details on the classes and methods from this tutorial.
Version Compatibility
---------------------
Ray Train is tested with `pytorch_lightning` versions `1.6.5` and `2.1.2`. For full compatibility, use ``pytorch_lightning>=1.6.5`` .
Earlier versions aren't prohibited but may result in unexpected issues. If you run into any compatibility issues, consider upgrading your PyTorch Lightning version or
`file an issue <https://github.com/ray-project/ray/issues>`_.
.. note::
If you are using Lightning 2.x, please use the import path `lightning.pytorch.xxx` instead of `pytorch_lightning.xxx`.
.. _lightning-trainer-migration-guide:
LightningTrainer Migration Guide
--------------------------------
Ray 2.4 introduced the `LightningTrainer`, and exposed a
`LightningConfigBuilder` to define configurations for `pl.LightningModule`
and `pl.Trainer`.
It then instantiates the model and trainer objects and runs a pre-defined
training function in a black box.
This version of the LightningTrainer API was constraining and limited
your ability to manage the training functionality.
Ray 2.7 introduced the newly unified :class:`~ray.train.torch.TorchTrainer` API, which offers
enhanced transparency, flexibility, and simplicity. This API is more aligned
with standard PyTorch Lightning scripts, ensuring users have better
control over their native Lightning code.
.. tab-set::
.. tab-item:: (Deprecating) LightningTrainer
.. This snippet isn't tested because it raises a hard deprecation warning.
.. testcode::
:skipif: True
from ray.train.lightning import LightningConfigBuilder, LightningTrainer
config_builder = LightningConfigBuilder()
# [1] Collect model configs
config_builder.module(cls=MyLightningModule, lr=1e-3, feature_dim=128)
# [2] Collect checkpointing configs
config_builder.checkpointing(monitor="val_accuracy", mode="max", save_top_k=3)
# [3] Collect pl.Trainer configs
config_builder.trainer(
max_epochs=10,
accelerator="gpu",
log_every_n_steps=100,
)
# [4] Build datasets on the head node
datamodule = MyLightningDataModule(batch_size=32)
config_builder.fit_params(datamodule=datamodule)
# [5] Execute the internal training function in a black box
ray_trainer = LightningTrainer(
lightning_config=config_builder.build(),
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
run_config=RunConfig(
checkpoint_config=CheckpointConfig(
num_to_keep=3,
checkpoint_score_attribute="val_accuracy",
checkpoint_score_order="max",
),
)
)
result = ray_trainer.fit()
# [6] Load the trained model from an opaque Lightning-specific checkpoint.
lightning_checkpoint = result.checkpoint
model = lightning_checkpoint.get_model(MyLightningModule)
.. tab-item:: (New API) TorchTrainer
.. This snippet isn't tested because it runs with 4 GPUs, and CI is only run with 1.
.. testcode::
:skipif: True
import os
import lightning.pytorch as pl
import ray.train
from ray.train.torch import TorchTrainer
from ray.train.lightning import (
RayDDPStrategy,
RayLightningEnvironment,
RayTrainReportCallback,
prepare_trainer
)
def train_func():
# [1] Create a Lightning model
model = MyLightningModule(lr=1e-3, feature_dim=128)
# [2] Report Checkpoint with callback
ckpt_report_callback = RayTrainReportCallback()
# [3] Create a Lighting Trainer
trainer = pl.Trainer(
max_epochs=10,
log_every_n_steps=100,
# New configurations below
devices="auto",
accelerator="auto",
strategy=RayDDPStrategy(),
plugins=[RayLightningEnvironment()],
callbacks=[ckpt_report_callback],
)
# Validate your Lightning trainer configuration
trainer = prepare_trainer(trainer)
# [4] Build your datasets on each worker
datamodule = MyLightningDataModule(batch_size=32)
trainer.fit(model, datamodule=datamodule)
# [5] Explicitly define and run the training function
ray_trainer = TorchTrainer(
train_func,
scaling_config=ray.train.ScalingConfig(num_workers=4, use_gpu=True),
run_config=ray.train.RunConfig(
checkpoint_config=ray.train.CheckpointConfig(
num_to_keep=3,
checkpoint_score_attribute="val_accuracy",
checkpoint_score_order="max",
),
)
)
result = ray_trainer.fit()
# [6] Load the trained model from a simplified checkpoint interface.
checkpoint: ray.train.Checkpoint = result.checkpoint
with checkpoint.as_directory() as checkpoint_dir:
print("Checkpoint contents:", os.listdir(checkpoint_dir))
checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.ckpt")
model = MyLightningModule.load_from_checkpoint(checkpoint_path)
doc/source/train/getting-started-pytorch.rst
+319
View File
@@ -0,0 +1,319 @@
.. _train-pytorch:
Get Started with Distributed Training using PyTorch
===================================================
This tutorial walks through the process of converting an existing PyTorch script to use Ray Train.
Learn how to:
1. Configure a model to run distributed and on the correct CPU/GPU device.
2. Configure a dataloader to shard data across the :ref:`workers <train-overview-worker>` and place data on the correct CPU or GPU device.
3. Configure a :ref:`training function <train-overview-training-function>` to report metrics and save checkpoints.
4. Configure :ref:`scaling <train-overview-scaling-config>` and CPU or GPU resource requirements for a training job.
5. Launch a distributed training job with a :class:`~ray.train.torch.TorchTrainer` class.
Quickstart
----------
For reference, the final code will look something like the following:
.. testcode::
:skipif: True
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
def train_func():
# Your PyTorch training code here.
...
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
1. `train_func` is the Python code that executes on each distributed training worker.
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use GPUs.
3. :class:`~ray.train.torch.TorchTrainer` launches the distributed training job.
Compare a PyTorch training script with and without Ray Train.
.. tab-set::
.. tab-item:: PyTorch + Ray Train
.. code-block:: python
:emphasize-lines: 12, 14, 21, 32, 36-37, 55-58, 59, 63, 66-73
import os
import tempfile
import torch
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
import ray.train.torch
def train_func():
# Model, Loss, Optimizer
model = resnet18(num_classes=10)
model.conv1 = torch.nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
# [1] Prepare model.
model = ray.train.torch.prepare_model(model)
# model.to("cuda") # This is done by `prepare_model`
criterion = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)
# Data
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
data_dir = os.path.join(tempfile.gettempdir(), "data")
train_data = FashionMNIST(root=data_dir, train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
# [2] Prepare dataloader.
train_loader = ray.train.torch.prepare_data_loader(train_loader)
# Training
for epoch in range(10):
if ray.train.get_context().get_world_size() > 1:
train_loader.sampler.set_epoch(epoch)
for images, labels in train_loader:
# This is done by `prepare_data_loader`!
# images, labels = images.to("cuda"), labels.to("cuda")
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# [3] Report metrics and checkpoint.
metrics = {"loss": loss.item(), "epoch": epoch}
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
torch.save(
model.module.state_dict(),
os.path.join(temp_checkpoint_dir, "model.pt")
)
ray.train.report(
metrics,
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
)
if ray.train.get_context().get_world_rank() == 0:
print(metrics)
# [4] Configure scaling and resource requirements.
scaling_config = ray.train.ScalingConfig(num_workers=2, use_gpu=True)
# [5] Launch distributed training job.
trainer = ray.train.torch.TorchTrainer(
train_func,
scaling_config=scaling_config,
# [5a] If running in a multi-node cluster, this is where you
# should configure the run's persistent storage that is accessible
# across all worker nodes.
# run_config=ray.train.RunConfig(storage_path="s3://..."),
)
result = trainer.fit()
# [6] Load the trained model.
with result.checkpoint.as_directory() as checkpoint_dir:
model_state_dict = torch.load(os.path.join(checkpoint_dir, "model.pt"))
model = resnet18(num_classes=10)
model.conv1 = torch.nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
model.load_state_dict(model_state_dict)
.. tab-item:: PyTorch
.. This snippet isn't tested because it doesn't use any Ray code.
.. testcode::
:skipif: True
import os
import tempfile
import torch
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
# Model, Loss, Optimizer
model = resnet18(num_classes=10)
model.conv1 = torch.nn.Conv2d(
1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
)
model.to("cuda")
criterion = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)
# Data
transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
train_data = FashionMNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
# Training
for epoch in range(10):
for images, labels in train_loader:
images, labels = images.to("cuda"), labels.to("cuda")
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
metrics = {"loss": loss.item(), "epoch": epoch}
checkpoint_dir = tempfile.mkdtemp()
checkpoint_path = os.path.join(checkpoint_dir, "model.pt")
torch.save(model.state_dict(), checkpoint_path)
print(metrics)
Set up a training function
--------------------------
.. include:: ./common/torch-configure-train_func.rst
Set up a model
^^^^^^^^^^^^^^
Use the :func:`ray.train.torch.prepare_model` utility function to:
1. Move your model to the correct device.
2. Wrap it in ``DistributedDataParallel``.
.. code-block:: diff
-from torch.nn.parallel import DistributedDataParallel
+import ray.train.torch
def train_func():
...
# Create model.
model = ...
# Set up distributed training and device placement.
- device_id = ... # Your logic to get the right device.
- model = model.to(device_id or "cpu")
- model = DistributedDataParallel(model, device_ids=[device_id])
+ model = ray.train.torch.prepare_model(model)
...
Set up a dataset
^^^^^^^^^^^^^^^^
.. TODO: Update this to use Ray Data.
Use the :func:`ray.train.torch.prepare_data_loader` utility function, which:
1. Adds a :class:`~torch.utils.data.distributed.DistributedSampler` to your :class:`~torch.utils.data.DataLoader`.
2. Moves the batches to the right device.
Note that this step isn't necessary if you're passing in Ray Data to your Trainer.
See :ref:`data-ingest-torch`.
.. code-block:: diff
from torch.utils.data import DataLoader
+import ray.train.torch
def train_func():
...
dataset = ...
data_loader = DataLoader(dataset, batch_size=worker_batch_size, shuffle=True)
+ data_loader = ray.train.torch.prepare_data_loader(data_loader)
for epoch in range(10):
+ if ray.train.get_context().get_world_size() > 1:
+ data_loader.sampler.set_epoch(epoch)
for X, y in data_loader:
- X = X.to_device(device)
- y = y.to_device(device)
...
.. tip::
Keep in mind that ``DataLoader`` takes in a ``batch_size`` which is the batch size for each worker.
The global batch size can be calculated from the worker batch size (and vice-versa) with the following equation:
.. testcode::
:skipif: True
global_batch_size = worker_batch_size * ray.train.get_context().get_world_size()
.. note::
If you already manually set up your ``DataLoader`` with a ``DistributedSampler``,
:meth:`~ray.train.torch.prepare_data_loader` will not add another one, and will
respect the configuration of the existing sampler.
.. note::
:class:`~torch.utils.data.distributed.DistributedSampler` does not work with a
``DataLoader`` that wraps :class:`~torch.utils.data.IterableDataset`.
If you want to work with an dataset iterator,
consider using :ref:`Ray Data <data>` instead of PyTorch DataLoader since it
provides performant streaming data ingestion for large scale datasets.
See :ref:`data-ingest-torch` for more details.
Report checkpoints and metrics
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To monitor progress, you can report intermediate metrics and checkpoints using the :func:`ray.train.report` utility function.
.. code-block:: diff
+import os
+import tempfile
+import ray.train
def train_func():
...
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
torch.save(
model.state_dict(), os.path.join(temp_checkpoint_dir, "model.pt")
)
+ metrics = {"loss": loss.item()} # Training/validation metrics.
# Build a Ray Train checkpoint from a directory
+ checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
# Ray Train will automatically save the checkpoint to persistent storage,
# so the local `temp_checkpoint_dir` can be safely cleaned up after.
+ ray.train.report(metrics=metrics, checkpoint=checkpoint)
...
For more details, see :ref:`train-monitoring-and-logging` and :ref:`train-checkpointing`.
.. include:: ./common/torch-configure-run.rst
Next steps
----------
After you have converted your PyTorch training script to use Ray Train:
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
* Dive into the :ref:`API Reference <train-api>` for more details on the classes and methods used in this tutorial.
doc/source/train/getting-started-transformers.rst
+447
View File
@@ -0,0 +1,447 @@
.. _train-pytorch-transformers:
Get Started with Distributed Training using Hugging Face Transformers
=====================================================================
This tutorial shows you how to convert an existing Hugging Face Transformers script to use Ray Train for distributed training.
In this guide, learn how to:
1. Configure a :ref:`training function <train-overview-training-function>` that properly reports metrics and saves checkpoints.
2. Configure :ref:`scaling <train-overview-scaling-config>` and resource requirements for CPUs or GPUs for your distributed training job.
3. Launch a distributed training job with :class:`~ray.train.torch.TorchTrainer`.
Requirements
------------
Install the necessary packages before you begin:
.. code-block:: bash
pip install "ray[train]" torch "transformers[torch]" datasets evaluate numpy scikit-learn
Quickstart
----------
Here's a quick overview of the final code structure:
.. testcode::
:skipif: True
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
def train_func():
# Your Transformers training code here
...
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
The key components are:
1. `train_func`: Python code that runs on each distributed training worker.
2. :class:`~ray.train.ScalingConfig`: Defines the number of distributed training workers and GPU usage.
3. :class:`~ray.train.torch.TorchTrainer`: Launches and manages the distributed training job.
Code Comparison: Hugging Face Transformers vs. Ray Train Integration
--------------------------------------------------------------------
Compare a standard Hugging Face Transformers script with its Ray Train equivalent:
.. tab-set::
.. tab-item:: Hugging Face Transformers + Ray Train
.. code-block:: python
:emphasize-lines: 13-15, 21, 67-68, 72, 80-87
import os
import numpy as np
import evaluate
from datasets import load_dataset
from transformers import (
Trainer,
TrainingArguments,
AutoTokenizer,
AutoModelForSequenceClassification,
)
import ray.train.huggingface.transformers
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
# [1] Encapsulate data preprocessing, training, and evaluation
# logic in a training function
# ============================================================
def train_func():
# Datasets
dataset = load_dataset("yelp_review_full")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)
small_train_dataset = (
dataset["train"].select(range(100)).map(tokenize_function, batched=True)
)
small_eval_dataset = (
dataset["test"].select(range(100)).map(tokenize_function, batched=True)
)
# Model
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-cased", num_labels=5
)
# Evaluation Metrics
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
# Hugging Face Trainer
training_args = TrainingArguments(
output_dir="test_trainer",
evaluation_strategy="epoch",
save_strategy="epoch",
report_to="none",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=small_train_dataset,
eval_dataset=small_eval_dataset,
compute_metrics=compute_metrics,
)
# [2] Report Metrics and Checkpoints to Ray Train
# ===============================================
callback = ray.train.huggingface.transformers.RayTrainReportCallback()
trainer.add_callback(callback)
# [3] Prepare Transformers Trainer
# ================================
trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
# Start Training
trainer.train()
# [4] Define a Ray TorchTrainer to launch `train_func` on all workers
# ===================================================================
ray_trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
# [4a] For multi-node clusters, configure persistent storage that is
# accessible across all worker nodes
# run_config=ray.train.RunConfig(storage_path="s3://..."),
)
result: ray.train.Result = ray_trainer.fit()
# [5] Load the trained model
with result.checkpoint.as_directory() as checkpoint_dir:
checkpoint_path = os.path.join(
checkpoint_dir,
ray.train.huggingface.transformers.RayTrainReportCallback.CHECKPOINT_NAME,
)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint_path)
.. tab-item:: Hugging Face Transformers
.. This snippet isn't tested because it doesn't use any Ray code.
.. testcode::
:skipif: True
# Adapted from Hugging Face tutorial: https://huggingface.co/docs/transformers/training
import numpy as np
import evaluate
from datasets import load_dataset
from transformers import (
Trainer,
TrainingArguments,
AutoTokenizer,
AutoModelForSequenceClassification,
)
# Datasets
dataset = load_dataset("yelp_review_full")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)
small_train_dataset = dataset["train"].select(range(100)).map(tokenize_function, batched=True)
small_eval_dataset = dataset["test"].select(range(100)).map(tokenize_function, batched=True)
# Model
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-cased", num_labels=5
)
# Metrics
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
# Hugging Face Trainer
training_args = TrainingArguments(
output_dir="test_trainer", evaluation_strategy="epoch", report_to="none"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=small_train_dataset,
eval_dataset=small_eval_dataset,
compute_metrics=compute_metrics,
)
# Start Training
trainer.train()
Set up a training function
--------------------------
.. include:: ./common/torch-configure-train_func.rst
Ray Train sets up the distributed process group on each worker before entering the training function.
Put all your logic into this function, including:
- Dataset construction and preprocessing
- Model initialization
- Transformers trainer definition
.. note::
When using Hugging Face Datasets or Evaluate, always call ``datasets.load_dataset`` and ``evaluate.load``
inside the training function. Don't pass loaded datasets and metrics from outside the training
function, as this can cause serialization errors when transferring objects to workers.
Report checkpoints and metrics
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To persist checkpoints and monitor training progress, add a
:class:`ray.train.huggingface.transformers.RayTrainReportCallback` utility callback to your Trainer:
.. code-block:: diff
import transformers
from ray.train.huggingface.transformers import RayTrainReportCallback
def train_func():
...
trainer = transformers.Trainer(...)
+ trainer.add_callback(RayTrainReportCallback())
...
Reporting metrics and checkpoints to Ray Train enables integration with Ray Tune and :ref:`fault-tolerant training <train-fault-tolerance>`.
The :class:`ray.train.huggingface.transformers.RayTrainReportCallback` provides a basic implementation, and you can :ref:`customize it <train-dl-saving-checkpoints>` to fit your needs.
Prepare a Transformers Trainer
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Pass your Transformers Trainer into
:meth:`~ray.train.huggingface.transformers.prepare_trainer` to validate
configurations and enable Ray Data integration:
.. code-block:: diff
import transformers
import ray.train.huggingface.transformers
def train_func():
...
trainer = transformers.Trainer(...)
+ trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
trainer.train()
...
.. include:: ./common/torch-configure-run.rst
Next steps
----------
Now that you've converted your Hugging Face Transformers script to use Ray Train:
* Explore :ref:`User Guides <train-user-guides>` to learn about specific tasks
* Browse the :doc:`Examples <examples>` for end-to-end Ray Train applications
* Consult the :ref:`API Reference <train-api>` for detailed information on the classes and methods
.. _transformers-trainer-migration-guide:
TransformersTrainer Migration Guide
-----------------------------------
Ray 2.1 introduced `TransformersTrainer` with a `trainer_init_per_worker` interface
to define `transformers.Trainer` and execute a pre-defined training function.
Ray 2.7 introduced the unified :class:`~ray.train.torch.TorchTrainer` API,
which offers better transparency, flexibility, and simplicity. This API aligns more closely
with standard Hugging Face Transformers scripts, giving you better control over your
training code.
.. tab-set::
.. tab-item:: (Deprecating) TransformersTrainer
.. This snippet isn't tested because it contains skeleton code.
.. testcode::
:skipif: True
import transformers
from transformers import AutoConfig, AutoModelForCausalLM
from datasets import load_dataset
import ray
from ray.train.huggingface import TransformersTrainer
from ray.train import ScalingConfig
from huggingface_hub import HfFileSystem
# Load datasets using HfFileSystem
path = "hf://datasets/Salesforce/wikitext/wikitext-2-raw-v1/"
fs = HfFileSystem()
# List the parquet files for each split
all_files = [f["name"] for f in fs.ls(path)]
train_files = [f for f in all_files if "train" in f and f.endswith(".parquet")]
validation_files = [f for f in all_files if "validation" in f and f.endswith(".parquet")]
ray_train_ds = ray.data.read_parquet(train_files, filesystem=fs)
ray_eval_ds = ray.data.read_parquet(validation_files, filesystem=fs)
# Define the Trainer generation function
def trainer_init_per_worker(train_dataset, eval_dataset, **config):
MODEL_NAME = "gpt2"
model_config = AutoConfig.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_config(model_config)
args = transformers.TrainingArguments(
output_dir=f"{MODEL_NAME}-wikitext2",
evaluation_strategy="epoch",
save_strategy="epoch",
logging_strategy="epoch",
learning_rate=2e-5,
weight_decay=0.01,
max_steps=100,
)
return transformers.Trainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
)
# Build a Ray TransformersTrainer
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
ray_trainer = TransformersTrainer(
trainer_init_per_worker=trainer_init_per_worker,
scaling_config=scaling_config,
datasets={"train": ray_train_ds, "validation": ray_eval_ds},
)
result = ray_trainer.fit()
.. tab-item:: (New API) TorchTrainer
.. This snippet isn't tested because it contains skeleton code.
.. testcode::
:skipif: True
import transformers
from transformers import AutoConfig, AutoModelForCausalLM
from datasets import load_dataset
import ray
from ray.train.torch import TorchTrainer
from ray.train.huggingface.transformers import (
RayTrainReportCallback,
prepare_trainer,
)
from ray.train import ScalingConfig
from huggingface_hub import HfFileSystem
# Load datasets using HfFileSystem
path = "hf://datasets/Salesforce/wikitext/wikitext-2-raw-v1/"
fs = HfFileSystem()
# List the parquet files for each split
all_files = [f["name"] for f in fs.ls(path)]
train_files = [f for f in all_files if "train" in f and f.endswith(".parquet")]
validation_files = [f for f in all_files if "validation" in f and f.endswith(".parquet")]
ray_train_ds = ray.data.read_parquet(train_files, filesystem=fs)
ray_eval_ds = ray.data.read_parquet(validation_files, filesystem=fs)
# [1] Define the full training function
# =====================================
def train_func():
MODEL_NAME = "gpt2"
model_config = AutoConfig.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_config(model_config)
# [2] Build Ray Data iterables
# ============================
train_dataset = ray.train.get_dataset_shard("train")
eval_dataset = ray.train.get_dataset_shard("validation")
train_iterable_ds = train_dataset.iter_torch_batches(batch_size=8)
eval_iterable_ds = eval_dataset.iter_torch_batches(batch_size=8)
args = transformers.TrainingArguments(
output_dir=f"{MODEL_NAME}-wikitext2",
evaluation_strategy="epoch",
save_strategy="epoch",
logging_strategy="epoch",
learning_rate=2e-5,
weight_decay=0.01,
max_steps=100,
)
trainer = transformers.Trainer(
model=model,
args=args,
train_dataset=train_iterable_ds,
eval_dataset=eval_iterable_ds,
)
# [3] Add Ray Train Report Callback
# =================================
trainer.add_callback(RayTrainReportCallback())
# [4] Prepare your trainer
# ========================
trainer = prepare_trainer(trainer)
trainer.train()
# Build a Ray TorchTrainer
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
ray_trainer = TorchTrainer(
train_func,
scaling_config=scaling_config,
datasets={"train": ray_train_ds, "validation": ray_eval_ds},
)
result = ray_trainer.fit()
doc/source/train/getting-started-xgboost.rst
+357
View File
@@ -0,0 +1,357 @@
.. _train-xgboost:
Get Started with Distributed Training using XGBoost
===================================================
This tutorial walks through the process of converting an existing XGBoost script to use Ray Train.
Learn how to:
1. Configure a :ref:`training function <train-overview-training-function>` to report metrics and save checkpoints.
2. Configure :ref:`scaling <train-overview-scaling-config>` and CPU or GPU resource requirements for a training job.
3. Launch a distributed training job with a :class:`~ray.train.xgboost.XGBoostTrainer`.
Quickstart
----------
For reference, the final code will look something like this:
.. testcode::
:skipif: True
import ray.train
from ray.train.xgboost import XGBoostTrainer
def train_func():
# Your XGBoost training code here.
...
scaling_config = ray.train.ScalingConfig(num_workers=2, resources_per_worker={"CPU": 4})
trainer = XGBoostTrainer(train_func, scaling_config=scaling_config)
result = trainer.fit()
1. `train_func` is the Python code that executes on each distributed training worker.
2. :class:`~ray.train.ScalingConfig` defines the number of distributed training workers and whether to use GPUs.
3. :class:`~ray.train.xgboost.XGBoostTrainer` launches the distributed training job.
Compare a XGBoost training script with and without Ray Train.
.. tab-set::
.. tab-item:: XGBoost + Ray Train
.. literalinclude:: ./doc_code/xgboost_quickstart.py
:emphasize-lines: 3-4, 7-8, 11, 15-16, 19-20, 48, 53, 56-64
:language: python
:start-after: __xgboost_ray_start__
:end-before: __xgboost_ray_end__
.. tab-item:: XGBoost
.. literalinclude:: ./doc_code/xgboost_quickstart.py
:language: python
:start-after: __xgboost_start__
:end-before: __xgboost_end__
Set up a training function
--------------------------
First, update your training code to support distributed training.
Begin by wrapping your `native <https://xgboost.readthedocs.io/en/latest/python/python_intro.html>`_
or `scikit-learn estimator <https://xgboost.readthedocs.io/en/latest/python/sklearn_estimator.html>`_
XGBoost training code in a :ref:`training function <train-overview-training-function>`:
.. testcode::
:skipif: True
def train_func():
# Your native XGBoost training code here.
dmatrix = ...
xgboost.train(...)
Each distributed training worker executes this function.
You can also specify the input argument for `train_func` as a dictionary via the Trainer's `train_loop_config`. For example:
.. testcode:: python
:skipif: True
def train_func(config):
label_column = config["label_column"]
num_boost_round = config["num_boost_round"]
...
config = {"label_column": "y", "num_boost_round": 10}
trainer = ray.train.xgboost.XGBoostTrainer(train_func, train_loop_config=config, ...)
.. warning::
Avoid passing large data objects through `train_loop_config` to reduce the
serialization and deserialization overhead. Instead,
initialize large objects (e.g. datasets, models) directly in `train_func`.
.. code-block:: diff
def load_dataset():
# Return a large in-memory dataset
...
def load_model():
# Return a large in-memory model instance
...
-config = {"data": load_dataset(), "model": load_model()}
def train_func(config):
- data = config["data"]
- model = config["model"]
+ data = load_dataset()
+ model = load_model()
...
trainer = ray.train.xgboost.XGBoostTrainer(train_func, train_loop_config=config, ...)
Ray Train automatically performs the worker communication setup that is needed to do distributed xgboost training.
Report metrics and save checkpoints
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To persist your checkpoints and monitor training progress, add a
:class:`ray.train.xgboost.RayTrainReportCallback` utility callback to your Trainer:
.. testcode:: python
:skipif: True
import xgboost
from ray.train.xgboost import RayTrainReportCallback
def train_func():
...
bst = xgboost.train(
...,
callbacks=[
RayTrainReportCallback(
metrics=["eval-logloss"], frequency=1
)
],
)
...
Reporting metrics and checkpoints to Ray Train enables :ref:`fault-tolerant training <train-fault-tolerance>` and the integration with Ray Tune.
Loading data
------------
When running distributed XGBoost training, each worker should use a different shard of the dataset.
.. testcode:: python
:skipif: True
def get_train_dataset(world_rank: int) -> xgboost.DMatrix:
# Define logic to get the DMatrix shard for this worker rank
...
def get_eval_dataset(world_rank: int) -> xgboost.DMatrix:
# Define logic to get the DMatrix for each worker
...
def train_func():
rank = ray.train.get_world_rank()
dtrain = get_train_dataset(rank)
deval = get_eval_dataset(rank)
...
A common way to do this is to pre-shard the dataset and then assign each worker a different set of files to read.
Pre-sharding the dataset is not very flexible to changes in the number of workers, since some workers may be assigned more data than others. For more flexibility, Ray Data provides a solution for sharding the dataset at runtime.
Use Ray Data to shard the dataset
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
:ref:`Ray Data <data>` is a distributed data processing library that allows you to easily shard and distribute your data across multiple workers.
First, load your **entire** dataset as a Ray Data Dataset.
Reference the :ref:`Ray Data Quickstart <data_quickstart>` for more details on how to load and preprocess data from different sources.
.. testcode:: python
:skipif: True
train_dataset = ray.data.read_parquet("s3://path/to/entire/train/dataset/dir")
eval_dataset = ray.data.read_parquet("s3://path/to/entire/eval/dataset/dir")
In the training function, you can access the dataset shards for this worker using :meth:`ray.train.get_dataset_shard`.
Convert this into a native `xgboost.DMatrix <https://xgboost.readthedocs.io/en/stable/python/python_api.html#xgboost.DMatrix>`_.
.. testcode:: python
:skipif: True
def get_dmatrix(dataset_name: str) -> xgboost.DMatrix:
shard = ray.train.get_dataset_shard(dataset_name)
df = shard.materialize().to_pandas()
X, y = df.drop("target", axis=1), df["target"]
return xgboost.DMatrix(X, label=y)
def train_func():
dtrain = get_dmatrix("train")
deval = get_dmatrix("eval")
...
Finally, pass the dataset to the Trainer. This will automatically shard the dataset across the workers. These keys must match the keys used when calling ``get_dataset_shard`` in the training function.
.. testcode:: python
:skipif: True
trainer = XGBoostTrainer(..., datasets={"train": train_dataset, "eval": eval_dataset})
trainer.fit()
For more details, see :ref:`data-ingest-torch`.
Configure scale and GPUs
------------------------
Outside of your training function, create a :class:`~ray.train.ScalingConfig` object to configure:
1. :class:`num_workers <ray.train.ScalingConfig>` - The number of distributed training worker processes.
2. :class:`use_gpu <ray.train.ScalingConfig>` - Whether each worker should use a GPU (or CPU).
3. :class:`resources_per_worker <ray.train.ScalingConfig>` - The number of CPUs or GPUs per worker.
.. testcode::
from ray.train import ScalingConfig
# 4 nodes with 8 CPUs each.
scaling_config = ScalingConfig(num_workers=4, resources_per_worker={"CPU": 8})
.. note::
When using Ray Data with Ray Train, be careful not to request all available CPUs in your cluster with the `resources_per_worker` parameter.
Ray Data needs CPU resources to execute data preprocessing operations in parallel.
If all CPUs are allocated to training workers, Ray Data operations may be bottlenecked, leading to reduced performance.
A good practice is to leave some portion of CPU resources available for Ray Data operations.
For example, if your cluster has 8 CPUs per node, you might allocate 6 CPUs to training workers and leave 2 CPUs for Ray Data:
.. testcode::
# Allocate 6 CPUs per worker, leaving resources for Ray Data operations
scaling_config = ScalingConfig(num_workers=4, resources_per_worker={"CPU": 6})
In order to use GPUs, you will need to set the `use_gpu` parameter to `True` in your :class:`~ray.train.ScalingConfig` object.
This will request and assign a single GPU per worker.
.. testcode::
# 1 node with 8 CPUs and 4 GPUs each.
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
# 4 nodes with 8 CPUs and 4 GPUs each.
scaling_config = ScalingConfig(num_workers=16, use_gpu=True)
When using GPUs, you will also need to update your training function to use the assigned GPU.
This can be done by setting the `"device"` parameter as `"cuda"`.
For more details on XGBoost's GPU support, see the `XGBoost GPU documentation <https://xgboost.readthedocs.io/en/stable/gpu/index.html>`__.
.. code-block:: diff
def train_func():
...
params = {
...,
+ "device": "cuda",
}
bst = xgboost.train(
params,
...
)
Configure persistent storage
----------------------------
Create a :class:`~ray.train.RunConfig` object to specify the path where results
(including checkpoints and artifacts) will be saved.
.. testcode::
from ray.train import RunConfig
# Local path (/some/local/path/unique_run_name)
run_config = RunConfig(storage_path="/some/local/path", name="unique_run_name")
# Shared cloud storage URI (s3://bucket/unique_run_name)
run_config = RunConfig(storage_path="s3://bucket", name="unique_run_name")
# Shared NFS path (/mnt/nfs/unique_run_name)
run_config = RunConfig(storage_path="/mnt/nfs", name="unique_run_name")
.. warning::
Specifying a *shared storage location* (such as cloud storage or NFS) is
*optional* for single-node clusters, but it is **required for multi-node clusters.**
Using a local path will :ref:`raise an error <multinode-local-storage-warning>`
during checkpointing for multi-node clusters.
For more details, see :ref:`persistent-storage-guide`.
Launch a training job
---------------------
Tying this all together, you can now launch a distributed training job
with a :class:`~ray.train.xgboost.XGBoostTrainer`.
.. testcode::
:hide:
from ray.train import ScalingConfig
train_func = lambda: None
scaling_config = ScalingConfig(num_workers=1)
run_config = None
.. testcode::
from ray.train.xgboost import XGBoostTrainer
trainer = XGBoostTrainer(
train_func, scaling_config=scaling_config, run_config=run_config
)
result = trainer.fit()
Access training results
-----------------------
After training completes, a :class:`~ray.train.Result` object is returned which contains
information about the training run, including the metrics and checkpoints reported during training.
.. testcode::
result.metrics # The metrics reported during training.
result.checkpoint # The latest checkpoint reported during training.
result.path # The path where logs are stored.
result.error # The exception that was raised, if training failed.
For more usage examples, see :ref:`train-inspect-results`.
Next steps
----------
After you have converted your XGBoost training script to use Ray Train:
* See :ref:`User Guides <train-user-guides>` to learn more about how to perform specific tasks.
* Browse the :doc:`Examples <examples>` for end-to-end examples of how to use Ray Train.
* Consult the :ref:`API Reference <train-api>` for more details on the classes and methods from this tutorial.
doc/source/train/horovod.rst
+101
View File
@@ -0,0 +1,101 @@
.. _train-horovod:
Get Started with Distributed Training using Horovod
===================================================
Ray Train configures the Horovod environment and Rendezvous
server for you, allowing you to run your ``DistributedOptimizer`` training
script. See the `Horovod documentation <https://horovod.readthedocs.io/en/stable/index.html>`_
for more information.
Quickstart
-----------
.. literalinclude:: ./doc_code/hvd_trainer.py
:language: python
Update your training function
-----------------------------
First, update your :ref:`training function <train-overview-training-function>` to support distributed
training.
If you have a training function that already runs with the `Horovod Ray
Executor <https://horovod.readthedocs.io/en/stable/ray_include.html#horovod-ray-executor>`_,
you shouldn't need to make any additional changes.
To onboard onto Horovod, visit the `Horovod guide
<https://horovod.readthedocs.io/en/stable/index.html#get-started>`_.
Create a HorovodTrainer
-----------------------
``Trainer``\s are the primary Ray Train classes to use to manage state and
execute training. For Horovod, use a :class:`~ray.train.horovod.HorovodTrainer`
that you can setup like this:
.. testcode::
:hide:
train_func = lambda: None
.. testcode::
from ray.train import ScalingConfig
from ray.train.horovod import HorovodTrainer
# For GPU Training, set `use_gpu` to True.
use_gpu = False
trainer = HorovodTrainer(
train_func,
scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=2)
)
When training with Horovod, always use a HorovodTrainer,
irrespective of the training framework, for example, PyTorch or TensorFlow.
To customize the backend setup, you can pass a
:class:`~ray.train.horovod.HorovodConfig`:
.. testcode::
:skipif: True
from ray.train import ScalingConfig
from ray.train.horovod import HorovodTrainer, HorovodConfig
trainer = HorovodTrainer(
train_func,
tensorflow_backend=HorovodConfig(...),
scaling_config=ScalingConfig(num_workers=2),
)
For more configurability, see the :py:class:`~ray.train.data_parallel_trainer.DataParallelTrainer` API.
Run a training function
-----------------------
With a distributed training function and a Ray Train ``Trainer``, you are now
ready to start training.
.. testcode::
:skipif: True
trainer.fit()
Further reading
---------------
Ray Train's :class:`~ray.train.horovod.HorovodTrainer` replaces the distributed
communication backend of the native libraries with its own implementation.
Thus, the remaining integration points remain the same. If you're using Horovod
with :ref:`PyTorch <train-pytorch>` or :ref:`Tensorflow <train-tensorflow-overview>`,
refer to the respective guides for further configuration
and information.
If you are implementing your own Horovod-based training routine without using any of
the training libraries, read through the
:ref:`User Guides <train-user-guides>`, as you can apply much of the content
to generic use cases and adapt them easily.
doc/source/train/huggingface-accelerate.rst
+224
View File
@@ -0,0 +1,224 @@
.. _train-hf-accelerate:
Get Started with Distributed Training using Hugging Face Accelerate
===================================================================
The :class:`~ray.train.torch.TorchTrainer` can help you easily launch your `Accelerate <https://huggingface.co/docs/accelerate>`_ training across a distributed Ray cluster.
You only need to run your existing training code with a TorchTrainer. You can expect the final code to look like this:
.. testcode::
:skipif: True
from accelerate import Accelerator
def train_func():
# Instantiate the accelerator
accelerator = Accelerator(...)
model = ...
optimizer = ...
train_dataloader = ...
eval_dataloader = ...
lr_scheduler = ...
# Prepare everything for distributed training
(
model,
optimizer,
train_dataloader,
eval_dataloader,
lr_scheduler,
) = accelerator.prepare(
model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
)
# Start training
...
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(...),
# If running in a multi-node cluster, this is where you
# should configure the run's persistent storage that is accessible
# across all worker nodes.
# run_config=ray.train.RunConfig(storage_path="s3://..."),
...
)
trainer.fit()
.. tip::
Model and data preparation for distributed training is completely handled by the `Accelerator <https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator>`_
object and its `Accelerator.prepare() <https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare>`_ method.
Unlike with native PyTorch, **don't** call any additional Ray Train utilities
like :meth:`~ray.train.torch.prepare_model` or :meth:`~ray.train.torch.prepare_data_loader` in your training function.
Configure Accelerate
--------------------
In Ray Train, you can set configurations through the `accelerate.Accelerator <https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator>`_
object in your training function. Below are starter examples for configuring Accelerate.
.. tab-set::
.. tab-item:: DeepSpeed
For example, to run DeepSpeed with Accelerate, create a `DeepSpeedPlugin <https://huggingface.co/docs/accelerate/main/en/package_reference/deepspeed>`_
from a dictionary:
.. testcode::
:skipif: True
from accelerate import Accelerator, DeepSpeedPlugin
DEEPSPEED_CONFIG = {
"fp16": {
"enabled": True
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": False
},
"overlap_comm": True,
"contiguous_gradients": True,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"gather_16bit_weights_on_model_save": True,
"round_robin_gradients": True
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 10,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": False
}
def train_func():
# Create a DeepSpeedPlugin from config dict
ds_plugin = DeepSpeedPlugin(hf_ds_config=DEEPSPEED_CONFIG)
# Initialize Accelerator
accelerator = Accelerator(
...,
deepspeed_plugin=ds_plugin,
)
# Start training
...
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(...),
run_config=ray.train.RunConfig(storage_path="s3://..."),
...
)
trainer.fit()
.. tab-item:: FSDP
:sync: FSDP
For PyTorch FSDP, create a `FullyShardedDataParallelPlugin <https://huggingface.co/docs/accelerate/main/en/package_reference/fsdp>`_
and pass it to the Accelerator.
.. testcode::
:skipif: True
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from accelerate import Accelerator, FullyShardedDataParallelPlugin
def train_func():
fsdp_plugin = FullyShardedDataParallelPlugin(
state_dict_config=FullStateDictConfig(
offload_to_cpu=False,
rank0_only=False
),
optim_state_dict_config=FullOptimStateDictConfig(
offload_to_cpu=False,
rank0_only=False
)
)
# Initialize accelerator
accelerator = Accelerator(
...,
fsdp_plugin=fsdp_plugin,
)
# Start training
...
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(...),
run_config=ray.train.RunConfig(storage_path="s3://..."),
...
)
trainer.fit()
Note that Accelerate also provides a CLI tool, `"accelerate config"`, to generate a configuration and launch your training
job with `"accelerate launch"`. However, it's not necessary here because Ray's `TorchTrainer` already sets up the Torch
distributed environment and launches the training function on all workers.
Next, see these end-to-end examples below for more details:
.. tab-set::
.. tab-item:: Example with Ray Data
.. dropdown:: Show Code
.. literalinclude:: /../../python/ray/train/examples/accelerate/accelerate_torch_trainer.py
:language: python
:start-after: __accelerate_torch_basic_example_start__
:end-before: __accelerate_torch_basic_example_end__
.. tab-item:: Example with PyTorch DataLoader
.. dropdown:: Show Code
.. literalinclude:: /../../python/ray/train/examples/accelerate/accelerate_torch_trainer_no_raydata.py
:language: python
:start-after: __accelerate_torch_basic_example_no_raydata_start__
:end-before: __accelerate_torch_basic_example_no_raydata_end__
.. seealso::
If you're looking for more advanced use cases, check out this Llama-2 fine-tuning example:
- `Fine-tuning Llama-2 series models with Deepspeed, Accelerate, and Ray Train. <https://github.com/ray-project/ray/tree/master/doc/source/templates/04_finetuning_llms_with_deepspeed>`_
You may also find these user guides helpful:
- :ref:`train_scaling_config`
- :ref:`Configuration and Persistent Storage <persistent-storage-guide>`
- :ref:`Saving and Loading Checkpoints <train-checkpointing>`
- :ref:`How to use Ray Data with Ray Train <data-ingest-torch>`
AccelerateTrainer Migration Guide
---------------------------------
Before Ray 2.7, Ray Train's `AccelerateTrainer` API was the
recommended way to run Accelerate code. As a subclass of :class:`TorchTrainer <ray.train.torch.TorchTrainer>`,
the AccelerateTrainer takes in a configuration file generated by ``accelerate config`` and applies it to all workers.
Aside from that, the functionality of ``AccelerateTrainer`` is identical to ``TorchTrainer``.
However, this caused confusion around whether this was the *only* way to run Accelerate code.
Because you can express the full Accelerate functionality with the ``Accelerator`` and ``TorchTrainer`` combination, the plan is to deprecate the ``AccelerateTrainer`` in Ray 2.8,
and it's recommend to run your Accelerate code directly with ``TorchTrainer``.
doc/source/train/images/checkpoint_lifecycle.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 216 KiB

doc/source/train/images/checkpoint_metrics_lifecycle.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 81 KiB

doc/source/train/images/fault_tolerance/cluster_failure.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 162 KiB

doc/source/train/images/fault_tolerance/cluster_failure_start.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 186 KiB

doc/source/train/images/fault_tolerance/cluster_recovery.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 196 KiB

doc/source/train/images/fault_tolerance/head_node_failure.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 167 KiB

doc/source/train/images/fault_tolerance/worker_failure_start.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 184 KiB

doc/source/train/images/fault_tolerance/worker_group_recovery.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 186 KiB

doc/source/train/images/fault_tolerance/worker_node_failure.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 164 KiB

doc/source/train/images/fault_tolerance/worker_node_replacement.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 106 KiB

doc/source/train/images/hyperparameter_optimization/train_tune_interop.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 209 KiB

doc/source/train/images/hyperparameter_optimization/train_without_tune.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 76 KiB

doc/source/train/images/logo.svg
+20
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 8.2 KiB

doc/source/train/images/overview.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 108 KiB

doc/source/train/images/persistent_storage_checkpoint.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 626 KiB

doc/source/train/images/predictors.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 212 KiB

doc/source/train/images/sync_vs_async_checkpointing.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 54 KiB

doc/source/train/images/train-specific.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 558 KiB

doc/source/train/images/train-tuner.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 36 KiB

doc/source/train/images/train.svg
+1
View File
File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 543 KiB

doc/source/train/images/train_ingest.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 46 KiB

doc/source/train/more-frameworks.rst
+58
View File
@@ -0,0 +1,58 @@
.. _train-more-frameworks:
More Frameworks
===============
.. toctree::
:hidden:
Hugging Face Accelerate Guide <huggingface-accelerate>
DeepSpeed Guide <deepspeed>
TensorFlow and Keras Guide <distributed-tensorflow-keras>
LightGBM Guide <getting-started-lightgbm>
Horovod Guide <horovod>
.. grid:: 1 2 3 4
:gutter: 1
:class-container: container pb-3
.. grid-item-card::
:img-top: /images/accelerate_logo.png
:class-img-top: mt-2 w-75 d-block mx-auto fixed-height-img
:link: huggingface-accelerate
:link-type: doc
Hugging Face Accelerate
.. grid-item-card::
:img-top: /images/deepspeed_logo.svg
:class-img-top: mt-2 w-75 d-block mx-auto fixed-height-img
:link: deepspeed
:link-type: doc
DeepSpeed
.. grid-item-card::
:img-top: /images/tf_logo.png
:class-img-top: mt-2 w-75 d-block mx-auto fixed-height-img
:link: distributed-tensorflow-keras
:link-type: doc
TensorFlow and Keras
.. grid-item-card::
:img-top: /images/lightgbm_logo.png
:class-img-top: mt-2 w-75 d-block mx-auto fixed-height-img
:link: getting-started-lightgbm
:link-type: doc
LightGBM
.. grid-item-card::
:img-top: /images/horovod.png
:class-img-top: mt-2 w-75 d-block mx-auto fixed-height-img
:link: horovod
:link-type: doc
Horovod
doc/source/train/overview.rst
+102
View File
@@ -0,0 +1,102 @@
.. _train-key-concepts:
.. _train-overview:
Ray Train Overview
==================
To use Ray Train effectively, you need to understand four main concepts:
#. :ref:`Training function <train-overview-training-function>`: A Python function that contains your model training logic.
#. :ref:`Worker <train-overview-worker>`: A process that runs the training function.
#. :ref:`Scaling configuration: <train-overview-scaling-config>` A configuration of the number of workers and compute resources (for example, CPUs or GPUs).
#. :ref:`Trainer <train-overview-trainers>`: A Python class that ties together the training function, workers, and scaling configuration to execute a distributed training job.
.. figure:: images/overview.png
:align: center
.. _train-overview-training-function:
Training function
-----------------
The training function is a user-defined Python function that contains the end-to-end model training loop logic.
When launching a distributed training job, each worker executes this training function.
Ray Train documentation uses the following conventions:
#. `train_func` is a user-defined function that contains the training code.
#. `train_func` is passed into the Trainer's `train_loop_per_worker` parameter.
.. testcode::
def train_func():
"""User-defined training function that runs on each distributed worker process.
This function typically contains logic for loading the model,
loading the dataset, training the model, saving checkpoints,
and logging metrics.
"""
...
.. _train-overview-worker:
Worker
------
Ray Train distributes model training compute to individual worker processes across the cluster.
Each worker is a process that executes the `train_func`.
The number of workers determines the parallelism of the training job and is configured in the :class:`~ray.train.ScalingConfig`.
.. _train-overview-scaling-config:
Scaling configuration
---------------------
The :class:`~ray.train.ScalingConfig` is the mechanism for defining the scale of the training job.
Specify two basic parameters for worker parallelism and compute resources:
* :class:`num_workers <ray.train.ScalingConfig>`: The number of workers to launch for a distributed training job.
* :class:`use_gpu <ray.train.ScalingConfig>`: Whether each worker should use a GPU.
.. testcode::
from ray.train import ScalingConfig
# Single worker with a CPU
scaling_config = ScalingConfig(num_workers=1, use_gpu=False)
# Single worker with a GPU
scaling_config = ScalingConfig(num_workers=1, use_gpu=True)
# Multiple workers, each with a GPU
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
.. _train-overview-trainers:
Trainer
-------
The Trainer ties the previous three concepts together to launch distributed training jobs.
Ray Train provides :ref:`Trainer classes <train-api>` for different frameworks.
Calling the :meth:`fit() <ray.train.trainer.BaseTrainer.fit>` method executes the training job by:
#. Launching workers as defined by the :ref:`scaling_config <train-overview-scaling-config>`.
#. Setting up the framework's distributed environment on all workers.
#. Running the `train_func` on all workers.
.. testcode::
:hide:
def train_func():
pass
scaling_config = ScalingConfig(num_workers=1, use_gpu=False)
.. testcode::
from ray.train.torch import TorchTrainer
trainer = TorchTrainer(train_func, scaling_config=scaling_config)
trainer.fit()
doc/source/train/train.rst
+230
View File
@@ -0,0 +1,230 @@
.. _train-docs:
Ray Train: Scalable Model Training
==================================
.. toctree::
:hidden:
Overview <overview>
PyTorch Guide <getting-started-pytorch>
PyTorch Lightning Guide <getting-started-pytorch-lightning>
Hugging Face Transformers Guide <getting-started-transformers>
XGBoost Guide <getting-started-xgboost>
JAX Guide <getting-started-jax>
more-frameworks
User Guides <user-guides>
Tutorials </_collections/train/tutorials/README>
Examples <examples>
Benchmarks <benchmarks>
.. div:: sd-d-flex-row sd-align-major-center sd-align-minor-center
.. div:: sd-w-50
.. raw:: html
:file: images/logo.svg
Ray Train is a scalable machine learning library for distributed training and fine-tuning.
Ray Train allows you to scale model training code from a single machine to a cluster of machines in the cloud, and abstracts away the complexities of distributed computing.
Whether you have large models or large datasets, Ray Train is the simplest solution for distributed training.
Ray Train provides support for many frameworks:
.. list-table::
:widths: 1 1
:header-rows: 1
* - PyTorch Ecosystem
- More Frameworks
* - PyTorch
- TensorFlow
* - PyTorch Lightning
- Keras
* - Hugging Face Transformers
- Horovod
* - Hugging Face Accelerate
- XGBoost
* - DeepSpeed
- LightGBM
Install Ray Train
-----------------
To install Ray Train, run:
.. code-block:: console
$ pip install -U "ray[train]"
To learn more about installing Ray and its libraries, see
:ref:`Installing Ray <installation>`.
Get started
-----------
.. grid:: 1 2 2 2
:gutter: 1
:class-container: container pb-6
.. grid-item-card::
**Overview**
^^^
Understand the key concepts for distributed training with Ray Train.
+++
.. button-ref:: train-overview
:color: primary
:outline:
:expand:
Learn the basics
.. grid-item-card::
**PyTorch**
^^^
Get started on distributed model training with Ray Train and PyTorch.
+++
.. button-ref:: train-pytorch
:color: primary
:outline:
:expand:
Try Ray Train with PyTorch
.. grid-item-card::
**PyTorch Lightning**
^^^
Get started on distributed model training with Ray Train and Lightning.
+++
.. button-ref:: train-pytorch-lightning
:color: primary
:outline:
:expand:
Try Ray Train with Lightning
.. grid-item-card::
**Hugging Face Transformers**
^^^
Get started on distributed model training with Ray Train and Transformers.
+++
.. button-ref:: train-pytorch-transformers
:color: primary
:outline:
:expand:
Try Ray Train with Transformers
.. grid-item-card::
**JAX**
^^^
Get started on distributed model training with Ray Train and JAX.
+++
.. button-ref:: train-jax
:color: primary
:outline:
:expand:
Try Ray Train with JAX
Learn more
----------
.. grid:: 1 2 2 2
:gutter: 1
:class-container: container pb-6
.. grid-item-card::
**More Frameworks**
^^^
Don't see your framework? See these guides.
+++
.. button-ref:: train-more-frameworks
:color: primary
:outline:
:expand:
Try Ray Train with other frameworks
.. grid-item-card::
**User Guides**
^^^
Get how-to instructions for common training tasks with Ray Train.
+++
.. button-ref:: train-user-guides
:color: primary
:outline:
:expand:
Read how-to guides
.. grid-item-card::
**Tutorials**
^^^
Hands-on tutorials covering ML workload patterns from vision to recommendation systems.
+++
.. button-ref:: /_collections/train/tutorials/README
:color: primary
:outline:
:expand:
:ref-type: doc
Follow tutorials
.. grid-item-card::
**Examples**
^^^
Browse end-to-end code examples for different use cases.
+++
.. button-ref:: examples
:color: primary
:outline:
:expand:
:ref-type: doc
Learn through examples
.. grid-item-card::
**API**
^^^
Consult the API Reference for full descriptions of the Ray Train API.
+++
.. button-ref:: train-api
:color: primary
:outline:
:expand:
Read the API Reference
doc/source/train/user-guides.rst
+22
View File
@@ -0,0 +1,22 @@
.. _train-user-guides:
Ray Train User Guides
=====================
.. toctree::
:maxdepth: 2
user-guides/data-loading-preprocessing
user-guides/using-accelerators
user-guides/persistent-storage
user-guides/monitoring-logging
user-guides/checkpoints
user-guides/asynchronous-validation
user-guides/experiment-tracking
user-guides/results
user-guides/fault-tolerance
user-guides/elastic-training
user-guides/monitor-your-application
user-guides/local_mode
user-guides/reproducibility
Hyperparameter Optimization <user-guides/hyperparameter-optimization>
doc/source/train/user-guides/_collate_utils.rst
+20
View File
@@ -0,0 +1,20 @@
:orphan:
.. _train-collate-utils:
Collate Utilities
=================
.. literalinclude:: ../doc_code/collate_utils.py
:language: python
.. _random-text-generator:
Random Text Generator
=====================
The following helper functions generate random text samples with labels:
.. literalinclude:: ../doc_code/random_text_generator.py
:language: python
doc/source/train/user-guides/asynchronous-validation.rst
+376
View File
@@ -0,0 +1,376 @@
.. _train-validating-checkpoints:
Validating checkpoints asynchronously
=====================================
During training, you may want to validate the model periodically to monitor training progress.
The standard way to do this is to periodically switch between training and validation within
the training loop. Instead, Ray Train allows you to asynchronously validate the model in a
separate Ray task, which does the following:
* Runs validation in parallel without blocking the training loop
* Runs validation on different, potentially cheaper hardware than training, since validation
doesn't require optimizer states or gradients and can use 2-4x less GPU memory
* Leverages :ref:`autoscaling <vms-autoscaling>` to launch user-specified machines only for the duration of the validation
* Lets training continue immediately after saving a checkpoint with partial metrics (for example, loss)
and then receives validation metrics (for example, accuracy) as soon as they are available. If the initial
and validated metrics share the same key, the validated metrics overwrite the initial metrics.
When to use async validation
----------------------------
Asynchronous validation is preferable to alternating between training and validation within the
same training loop in the following scenarios:
* **Validation takes a large percentage of total training time.** If validation is a significant
fraction of your end-to-end training time, running it asynchronously can substantially reduce
wall clock time by overlapping validation with training.
* **Cheaper GPUs are available for validation.** Validation doesn't require optimizer states or
gradients, so it can use 2-4x less GPU memory than training. If you have a pool of cheaper GPUs
or an autoscaling setup that can provision them, async validation lets you run validation on
those cheaper machines instead of occupying your expensive training GPUs.
* **Training throughput stops scaling linearly with more workers.** As worker count increases,
allreduce overhead grows and limits training speed, so doubling workers no longer doubles
throughput. Validation, however, scales more linearly since it requires no gradient synchronization.
Asynchronous validation can therefore utilize otherwise idle cluster capacity without impacting
training.
The best way to know if async validation helps your workload is to try it. Converting is
straightforward (see the tutorial below), so you can run both approaches and compare.
Tutorial
--------
First, define a ``validation_fn`` that takes a :class:`ray.train.Checkpoint` to validate
and any number of json-serializable keyword arguments. This function should return a dictionary
of metrics from that validation.
The following is a simple example for teaching purposes only. It is impractical
because the validation task always runs on cpu; for a more realistic example, see
:ref:`train-distributed-validate-fn`.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_simple_start__
:end-before: __validation_fn_simple_end__
.. note::
In this example, the validation dataset is a ray.data.Dataset object, which is not
json-serializable. We therefore include it with the validation_fn closure instead of passing
it as a keyword argument.
.. warning::
Don't pass large objects to the ``validation_fn`` because Ray Train runs it as a Ray task and
serializes all captured variables. Instead, package large objects in the ``Checkpoint`` and
access them from shared storage later as explained in :ref:`train-checkpointing`.
Next, register your ``validation_fn`` with your trainer by settings its ``validation_config`` argument to a
:class:`~ray.train.v2.api.report_config.ValidationConfig` object that contains your ``validation_fn``
and any default keyword arguments you want to pass to your ``validation_fn``.
Next, within your rank 0 worker's training loop, call :func:`ray.train.report` with ``validation``
set to True, which will call your ``validation_fn`` with the default keyword arguments you passed to the trainer.
Alternatively, you can set ``validation`` to a :class:`~ray.train.v2.api.report_config.ValidationTaskConfig` object
that contains keyword arguments that will override matching keyword arguments you passed to the trainer. If
``validation`` is False, Ray Train will not run validation.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_report_start__
:end-before: __validation_fn_report_end__
Finally, after training is done, you can access your checkpoints and their associated metrics with the
:class:`ray.train.Result` object. See :ref:`train-inspect-results` for more details.
.. _train-distributed-validate-fn:
Write a distributed validation function
---------------------------------------
The ``validation_fn`` above runs in a single Ray task, but you can improve its performance by spawning
even more Ray tasks or actors. The Ray team recommends doing this with one of the following approaches:
* Creating a :class:`ray.train.torch.TorchTrainer` that only does validation, not training.
* Using :func:`ray.data.Dataset.map_batches` to calculate metrics on a validation set.
Choose an approach
~~~~~~~~~~~~~~~~~~
You should use ``TorchTrainer`` if:
* You want to keep your existing validation logic and avoid migrating to Ray Data.
The training function API lets you fully customize the validation loop to match your current setup.
* Your validation code depends on running within a Torch process group — for example, your
metric aggregation logic uses collective communication calls, or your model parallelism
setup requires cross-GPU communication during the forward pass.
* You want a more consistent training and validation experience. The ``map_batches`` approach involves
running multiple Ray Data Datasets in a single ray cluster; we are currently working on better support
for this.
You should use ``map_batches`` if:
* You care about validation performance. Preliminary benchmarks show that ``map_batches`` is
faster.
* You prefer Ray Data’s native metric aggregation APIs over PyTorch, where you must implement
aggregation manually using low-level collective operations or rely on third-party libraries
such as `torchmetrics <https://lightning.ai/docs/torchmetrics/stable>`_.
Example: validation with Ray Train TorchTrainer
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Here is a ``validation_fn`` that uses a ``TorchTrainer`` to calculate average cross entropy
loss on a validation set. Note the following about this example:
* ``TorchTrainer`` is typically used for training, but you can use it for validation like in this
example allowing different resource requirements for training and validation, for example,
A100 for training and A10G for validation.
* The validation train function returns its metrics directly from worker 0 rather than calling
``ray.train.report`` which is accessible via ``result.return_value``. These values can't be torch
tensors and must be python based like ``ray.train.report``.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_torch_trainer_start__
:end-before: __validation_fn_torch_trainer_end__
Example: validation with Ray Data map_batches
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The following is a ``validation_fn`` that uses :func:`ray.data.Dataset.map_batches` to
calculate average accuracy on a validation set. To learn more about how to use
``map_batches`` for batch inference, see :ref:`batch_inference_home`.
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __validation_fn_map_batches_start__
:end-before: __validation_fn_map_batches_end__
Isolating training and validation with subclusters
---------------------------------------------------
When training and validation run concurrently on the same Ray cluster,
they compete for the same nodes by default. To give each phase its own
slice of the cluster — for example, A100s for training and A10Gs for
validation — label your worker pools with a ``ray-subcluster`` value and
pin each Dataset to its subcluster. See :ref:`data_concurrent_execution`
for the background and compute-config setup.
The pattern differs slightly between the ``TorchTrainer`` validation_fn
and the ``map_batches`` validation_fn, because only the former goes
through ``ray.train.DataConfig``.
**TorchTrainer validation_fn.** Set the validation Dataset's selector
through the sub-trainer's ``dataset_config``:
.. code-block:: python
from ray.data import ExecutionOptions
def validation_fn(checkpoint, ...) -> dict:
trainer = ray.train.torch.TorchTrainer(
...,
datasets={"validation": validation_dataset},
dataset_config=ray.train.DataConfig(
execution_options={
"validation": ExecutionOptions(
label_selector={"ray-subcluster": "validation"}
),
},
),
)
...
**map_batches validation_fn.** The ``map_batches`` path doesn't take a
``DataConfig``. Construct ``validation_dataset`` under a
``DataContext.current()`` block so the selector is baked into the
Dataset at construction — every downstream operator inherits it:
.. code-block:: python
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "validation"}
with ray.data.DataContext.current(ctx):
validation_dataset = ray.data.read_parquet(...)
def validation_fn(checkpoint) -> dict:
eval_res = validation_dataset.map_batches(...)
...
**Training-side configuration.** A Train pipeline needs the selector
specified in two places — they cover different phases and are not
redundant:
1. **At Dataset construction**, via the ``DataContext.current()`` context
manager, so construction-time tasks (parquet schema inference, file
listing) land on training nodes.
2. **In the trainer's** ``dataset_config``, because Train wholesale
replaces ``ds.context.execution_options`` with ``DataConfig``'s
per-dataset entry at training start. Anything not restated in
``DataConfig.execution_options`` — ``label_selector`` included — is
dropped, so per-worker ingest would lose its pinning.
.. code-block:: python
from ray.data import ExecutionOptions
def run_trainer() -> ray.train.Result:
# (1) Pin construction-time tasks.
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "training"}
with ray.data.DataContext.current(ctx):
train_dataset = ray.data.read_parquet(...)
# (2) Pin per-worker ingest — Train replaces ds.context options
# wholesale, so the selector must be restated here.
trainer = ray.train.torch.TorchTrainer(
...,
datasets={"train": train_dataset},
dataset_config=ray.train.DataConfig(
datasets_to_split=["train"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "training"}
),
},
),
)
...
.. note::
For *interleaved* validation — where you reuse the training workers
to validate on a separate "validation" Dataset inside the same
``TorchTrainer`` — pass both Datasets to ``datasets={...}`` and give
both an entry in ``DataConfig.execution_options`` so they're each
scoped to their own subcluster:
.. code-block:: python
from ray.data import ExecutionOptions
dataset_config = ray.train.DataConfig(
datasets_to_split=["train", "validation"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "training"}
),
"validation": ExecutionOptions(
label_selector={"ray-subcluster": "validation"}
),
},
)
Tuning asynchronous validation
------------------------------
Overlapping validation and training
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Asynchronous validation is most beneficial when training and validation fully overlap. If one
finishes before the other, some workers sit idle. :ref:`Autoscaling <vms-autoscaling>` lets you
spin up workers only for the duration of validation, which mitigates this but doesn't fully
eliminate the gap.
You can tune the following knobs to overlap validation and training as closely as possible:
* **Number of workers**: Tune the number of validation workers relative to training workers so that
the two phases overlap as closely as possible.
* **Batch size**: A larger batch size typically improves throughput, but it can negatively impact
training convergence and may lead to out-of-memory (OOM) errors.
* **Validation frequency**: Choose a validation cadence and dataset size that balance overlap with
training. Validating too frequently or over too many rows can create a long validation tail.
Also note that breaking early from a Ray Data iterator may lead to resource leaks - this will be
fixed in a future release.
Ray Data production vs consumption
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
See :ref:`balancing-data-production-consumption` for tips on balancing data production and consumption rates.
Checkpoint metrics lifecycle
-----------------------------
During the training loop the following happens to your checkpoints and metrics :
1. You report a checkpoint with some initial metrics, such as training loss, as well as a
:class:`~ray.train.v2.api.report_config.ValidationTaskConfig` object that contains the keyword
arguments to pass to the ``validation_fn``.
2. Ray Train asynchronously runs your ``validation_fn`` with that checkpoint and configuration.
3. When that validation task completes, Ray Train associates the metrics returned by your ``validation_fn``
with that checkpoint.
4. After training is done, you can access your checkpoints and their associated metrics with the
:class:`ray.train.Result` object. See :ref:`train-inspect-results` for more details.
.. figure:: ../images/checkpoint_metrics_lifecycle.png
How Ray Train populates checkpoint metrics during training and how you access them after training.
Experiment tracking
-------------------
In normal :ref:`experiment tracking with Ray Train <train-experiment-tracking-native>`,
you handle creating, logging to, and finishing the experiment tracking run from
the rank 0 training worker. However, asynchronous validation complicates this because
validation metrics are computed outside of the training worker, in a separate
Ray task.
Most modern experiment tracking configurations (for example,
`W&B distributed training <https://docs.wandb.ai/models/track/log/distributed-training#track-all-processes-to-a-single-run>`_)
support writing to the same run from different threads or processes. Other configurations,
such as the `MLflow fluent API <https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html>`_, may not.
Writing to the same run
~~~~~~~~~~~~~~~~~~~~~~~
If your experiment tracking library supports writing to the same run from different
processes, the rank 0 training worker can start the run and the validation task can
join it and log validation metrics directly.
.. tab-set::
.. tab-item:: W&B
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __exp_tracking_same_run_wandb_start__
:end-before: __exp_tracking_same_run_wandb_end__
.. tab-item:: MLflow (non-fluent)
.. literalinclude:: ../doc_code/asynchronous_validation.py
:language: python
:start-after: __exp_tracking_same_run_mlflow_start__
:end-before: __exp_tracking_same_run_mlflow_end__
Reliability
~~~~~~~~~~~
If experiment tracking logging fails (for example, due to a transient network error),
you have two options for retrying:
1. **Wrap your logging calls in a try/except block** within the ``validation_fn`` and
retry the logging manually with your experiment tracker's API.
2. **Use** :func:`ray.train.get_all_reported_checkpoints` **periodically during training** to
retrieve all reported checkpoints and their associated metrics, then re-log any missing
entries to your experiment tracker.
Writing to different runs
~~~~~~~~~~~~~~~~~~~~~~~~~
If your experiment tracking library does not support writing to the same run from different
processes, the validation task must start a new run each time it logs validation metrics.
Many tracking libraries provide ways to group related runs together so that training and
validation runs are still associated.
.. tab-set::
.. tab-item:: W&B
Use `W&B run grouping <https://docs.wandb.ai/models/runs/grouping>`_ to group
the training run and validation runs together.
.. tab-item:: MLflow
Use `MLflow parent and child runs <https://mlflow.org/docs/latest/ml/traditional-ml/tutorials/hyperparameter-tuning/part1-child-runs/#adapting-for-parent-and-child-runs>`_
to group the training run and validation runs together.
doc/source/train/user-guides/checkpoints.rst
+470
View File
@@ -0,0 +1,470 @@
.. _train-checkpointing:
Saving and Loading Checkpoints
==============================
Ray Train provides a way to snapshot training progress with :class:`Checkpoints <ray.train.Checkpoint>`.
This is useful for:
1. **Storing the best-performing model weights:** Save your model to persistent storage, and use it for downstream serving or inference.
2. **Fault tolerance:** Handle worker process and node failures in a long-running training job and leverage pre-emptible machines.
3. **Distributed checkpointing:** Ray Train checkpointing can be used to
:ref:`upload model shards from multiple workers in parallel. <train-distributed-checkpointing>`
.. _train-dl-saving-checkpoints:
Saving checkpoints during training
----------------------------------
The :class:`Checkpoint <ray.train.Checkpoint>` is a lightweight interface provided
by Ray Train that represents a *directory* that exists on local or remote storage.
For example, a checkpoint could point to a directory in cloud storage:
``s3://my-bucket/my-checkpoint-dir``.
A locally available checkpoint points to a location on the local filesystem:
``/tmp/my-checkpoint-dir``.
Here's how you save a checkpoint in the training loop:
1. Write your model checkpoint to a local directory.
- Since a :class:`Checkpoint <ray.train.Checkpoint>` just points to a directory, the contents are completely up to you.
- This means that you can use any serialization format you want.
- This makes it **easy to use familiar checkpoint utilities provided by training frameworks**, such as
``torch.save``, ``pl.Trainer.save_checkpoint``, Accelerate's ``accelerator.save_model``,
Transformers' ``save_pretrained``, ``tf.keras.Model.save``, etc.
2. Create a :class:`Checkpoint <ray.train.Checkpoint>` from the directory using :meth:`Checkpoint.from_directory <ray.train.Checkpoint.from_directory>`.
3. Report the checkpoint to Ray Train using :func:`ray.train.report(metrics, checkpoint=...) <ray.train.report>`.
- The metrics reported alongside the checkpoint are used to :ref:`keep track of the best-performing checkpoints <train-dl-configure-checkpoints>`.
- This will **upload the checkpoint to persistent storage** if configured. See :ref:`persistent-storage-guide`.
.. figure:: ../images/checkpoint_lifecycle.png
The lifecycle of a :class:`~ray.train.Checkpoint`, from being saved locally
to disk to being uploaded to persistent storage via ``train.report``.
As shown in the figure above, the best practice for saving checkpoints is to
first dump the checkpoint to a local temporary directory. Then, the call to ``train.report``
uploads the checkpoint to its final persistent storage location.
Then, the local temporary directory can be safely cleaned up to free up disk space
(e.g., from exiting the ``tempfile.TemporaryDirectory`` context).
.. tip::
In standard DDP training, where each worker has a copy of the full-model, you should
only save and report a checkpoint from a single worker to prevent redundant uploads.
This typically looks like:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_from_single_worker_start__
:end-before: __checkpoint_from_single_worker_end__
If using parallel training strategies such as DeepSpeed Zero and FSDP, where
each worker only has a shard of the full training state, you can save and report a checkpoint
from each worker. See :ref:`train-distributed-checkpointing` for an example.
Here are a few examples of saving checkpoints with different training frameworks:
.. tab-set::
.. tab-item:: Native PyTorch
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __pytorch_save_start__
:end-before: __pytorch_save_end__
.. tip::
You most likely want to unwrap the DDP model before saving it to a checkpoint.
``model.module.state_dict()`` is the state dict without each key having a ``"module."`` prefix.
.. tab-item:: PyTorch Lightning
Ray Train leverages PyTorch Lightning's ``Callback`` interface to report metrics
and checkpoints. We provide a simple callback implementation that reports
``on_train_epoch_end``.
Specifically, on each train epoch end, it
- collects all the logged metrics from ``trainer.callback_metrics``
- saves a checkpoint via ``trainer.save_checkpoint``
- reports to Ray Train via :func:`ray.train.report(metrics, checkpoint) <ray.train.report>`
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_save_example_start__
:end-before: __lightning_save_example_end__
You can always get the saved checkpoint path from :attr:`result.checkpoint <ray.train.Result>` and
:attr:`result.best_checkpoints <ray.train.Result>`.
For more advanced usage (e.g. reporting at different frequency, reporting
customized checkpoint files), you can implement your own customized callback.
Here is a simple example that reports a checkpoint every 3 epochs:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_custom_save_example_start__
:end-before: __lightning_custom_save_example_end__
.. tab-item:: Hugging Face Transformers
Ray Train leverages Hugging Face Transformers Trainer's ``Callback`` interface
to report metrics and checkpoints.
**Option 1: Use Ray Train's default report callback**
We provide a simple callback implementation :class:`~ray.train.huggingface.transformers.RayTrainReportCallback` that
reports on checkpoint save. You can change the checkpointing frequency by ``save_strategy`` and ``save_steps``.
It collects the latest logged metrics and report them together with the latest saved checkpoint.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __transformers_save_example_start__
:end-before: __transformers_save_example_end__
Note that :class:`~ray.train.huggingface.transformers.RayTrainReportCallback`
binds the latest metrics and checkpoints together,
so users can properly configure ``logging_strategy``, ``save_strategy`` and ``evaluation_strategy``
to ensure the monitoring metric is logged at the same step as checkpoint saving.
For example, the evaluation metrics (``eval_loss`` in this case) are logged during
evaluation. If users want to keep the best 3 checkpoints according to ``eval_loss``, they
should align the saving and evaluation frequency. Below are two examples of valid configurations:
.. testcode::
:skipif: True
args = TrainingArguments(
...,
evaluation_strategy="epoch",
save_strategy="epoch",
)
args = TrainingArguments(
...,
evaluation_strategy="steps",
save_strategy="steps",
eval_steps=50,
save_steps=100,
)
# And more ...
**Option 2: Implement your customized report callback**
If you feel that Ray Train's default :class:`~ray.train.huggingface.transformers.RayTrainReportCallback`
is not sufficient for your use case, you can also implement a callback yourself!
Below is a example implementation that collects latest metrics
and reports on checkpoint save.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __transformers_custom_save_example_start__
:end-before: __transformers_custom_save_example_end__
You can customize when (``on_save``, ``on_epoch_end``, ``on_evaluate``) and
what (customized metrics and checkpoint files) to report by implementing your own
Transformers Trainer callback.
.. _train-distributed-checkpointing:
Saving checkpoints from multiple workers (distributed checkpointing)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In model parallel training strategies where each worker only has a shard of the full-model,
you can save and report checkpoint shards in parallel from each worker.
.. figure:: ../images/persistent_storage_checkpoint.png
Distributed checkpointing in Ray Train. Each worker uploads its own checkpoint shard
to persistent storage independently.
Distributed checkpointing is the best practice for saving checkpoints
when doing model-parallel training (e.g., DeepSpeed, FSDP, Megatron-LM).
There are two major benefits:
1. **It is faster, resulting in less idle time.** Faster checkpointing incentivizes more frequent checkpointing!
Each worker can upload its checkpoint shard in parallel,
maximizing the network bandwidth of the cluster. Instead of a single node
uploading the full model of size ``M``, the cluster distributes the load across
``N`` nodes, each uploading a shard of size ``M / N``.
2. **Distributed checkpointing avoids needing to gather the full model onto a single worker's CPU memory.**
This gather operation puts a large CPU memory requirement on the worker that performs checkpointing
and is a common source of OOM errors.
Here is an example of distributed checkpointing with PyTorch:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __distributed_checkpointing_start__
:end-before: __distributed_checkpointing_end__
.. note::
Checkpoint files with the same name will collide between workers.
You can get around this by adding a rank-specific suffix to checkpoint files.
Note that having filename collisions does not error, but it will result in the last
uploaded version being the one that is persisted. This is fine if the file
contents are the same across all workers.
Model shard saving utilities provided by frameworks such as DeepSpeed will create
rank-specific filenames already, so you usually do not need to worry about this.
.. _train-checkpoint-upload-modes:
Checkpoint upload modes
-----------------------
By default, when you call :func:`~ray.train.report`, Ray Train synchronously pushes
your checkpoint from ``checkpoint.path`` on local disk to ``checkpoint_dir_name`` on
your ``storage_path``. This is equivalent to calling :func:`~ray.train.report` with
:class:`~ray.train.CheckpointUploadMode` set to ``ray.train.CheckpointUploadMode.SYNC``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_sync_start__
:end-before: __checkpoint_upload_mode_sync_end__
.. _train-checkpoint-upload-mode-async:
Asynchronous checkpoint uploading
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
You may want to upload your checkpoint asynchronously instead so that
the next training step can start in parallel. If so, you should use
``ray.train.CheckpointUploadMode.ASYNC``, which kicks off a new thread
to upload the checkpoint. This is helpful for larger
checkpoints that might take longer to upload, but might add unnecessary
complexity (see below) if you want to immediately upload only a small checkpoint.
Each ``report`` blocks until the previous ``report``\'s checkpoint
upload completes before starting a new checkpoint upload thread. Ray Train does this
to avoid accumulating too many upload threads and potentially running out of memory.
Because ``report`` returns without waiting for the checkpoint upload to complete,
you must ensure that the local checkpoint directory stays alive until the checkpoint
upload completes. This means you can't use a temporary directory that Ray Train may
delete before the upload finishes, for example from ``tempfile.TemporaryDirectory``.
``report`` also exposes the ``delete_local_checkpoint_after_upload`` parameter, which
defaults to ``True`` if ``checkpoint_upload_mode`` is ``ray.train.CheckpointUploadMode.ASYNC``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_async_start__
:end-before: __checkpoint_upload_mode_async_end__
.. figure:: ../images/sync_vs_async_checkpointing.png
This figure illustrates the difference between synchronous and asynchronous
checkpoint uploading.
Custom checkpoint uploading
~~~~~~~~~~~~~~~~~~~~~~~~~~~
:func:`~ray.train.report` defaults to uploading from disk to the remote ``storage_path``
with the PyArrow filesystem copying utilities before reporting the checkpoint to Ray Train.
If you would rather upload the checkpoint manually or with a third-party library
such as `Torch Distributed Checkpointing <https://docs.pytorch.org/docs/stable/distributed.checkpoint.html>`_,
you have the following options:
.. tab-set::
.. tab-item:: Synchronous
If you want to upload the checkpoint synchronously, you can first upload the checkpoint
to the ``storage_path`` and then report a reference to the uploaded checkpoint with
``ray.train.CheckpointUploadMode.NO_UPLOAD``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_mode_no_upload_start__
:end-before: __checkpoint_upload_mode_no_upload_end__
.. tab-item:: Asynchronous
If you want to upload the checkpoint asynchronously, you can set ``checkpoint_upload_mode``
to ``ray.train.CheckpointUploadMode.ASYNC`` and pass a ``checkpoint_upload_fn`` to
``ray.train.report``. This function takes the ``Checkpoint`` and ``checkpoint_dir_name``
passed to ``ray.train.report`` and returns the persisted ``Checkpoint``.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __checkpoint_upload_fn_start__
:end-before: __checkpoint_upload_fn_end__
.. warning::
In your ``checkpoint_upload_fn``, you should not call ``ray.train.report``, which may
lead to unexpected behavior. You should also avoid collective operations, such as
:func:`~ray.train.report` or ``model.state_dict()``, which can cause deadlocks. Finally,
the upload function should only the return a checkpoint object once all checkpoint data
has been saved.
.. note::
Do not pass a ``checkpoint_upload_fn`` with ``checkpoint_upload_mode=ray.train.CheckpointUploadMode.NO_UPLOAD``
because Ray Train will simply ignore ``checkpoint_upload_fn``. You can pass a ``checkpoint_upload_fn`` with
``checkpoint_upload_mode=ray.train.CheckpointUploadMode.SYNC``, but this is equivalent to uploading the
checkpoint yourself and reporting the checkpoint with ``ray.train.CheckpointUploadMode.NO_UPLOAD``.
.. _train-dl-configure-checkpoints:
Configure checkpointing
-----------------------
Ray Train provides some configuration options for checkpointing via :class:`~ray.train.CheckpointConfig`.
The primary configuration is keeping only the top ``K`` checkpoints with respect to a metric.
Lower-performing checkpoints are deleted to save storage space. By default, all checkpoints are kept.
.. literalinclude:: ../doc_code/key_concepts.py
:language: python
:start-after: __checkpoint_config_start__
:end-before: __checkpoint_config_end__
.. note::
If you want to save the top ``num_to_keep`` checkpoints with respect to a metric via
:py:class:`~ray.train.CheckpointConfig`,
please ensure that the metric is always reported together with the checkpoints.
Using checkpoints during training
----------------------------------
During training, you may want to access checkpoints you've reported and their associated metrics
from training workers for a variety of reasons, such as
reporting the best checkpoint so far to an experiment tracker. You can do this by calling
:func:`~ray.train.get_all_reported_checkpoints` from within your training function. This function returns
a list of :class:`~ray.train.ReportedCheckpoint` objects that represent all the
:class:`~ray.train.Checkpoint`\s and their associated metrics that you've reported so far
and have been kept based on the :ref:`checkpoint configuration <train-dl-configure-checkpoints>`.
This function supports two consistency modes:
- ``CheckpointConsistencyMode.COMMITTED``: Block until the checkpoint from the latest ``ray.train.report``
has been uploaded to persistent storage and committed.
- ``CheckpointConsistencyMode.VALIDATED``: Block until the checkpoint from the latest ``ray.train.report``
has been uploaded to persistent storage, committed, and validated (see :ref:`train-validating-checkpoints`).
This is the default consistency mode and has the same behavior as ``CheckpointConsistencyMode.COMMITTED``
if your report did not kick off validation.
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __get_all_reported_checkpoints_example_start__
:end-before: __get_all_reported_checkpoints_example_end__
Using checkpoints after training
--------------------------------
The latest saved checkpoint can be accessed with :attr:`Result.checkpoint <ray.train.Result>`.
The full list of persisted checkpoints can be accessed with :attr:`Result.best_checkpoints <ray.train.Result>`.
If :class:`CheckpointConfig(num_to_keep) <ray.train.CheckpointConfig>` is set, this list will contain the best ``num_to_keep`` checkpoints.
See :ref:`train-inspect-results` for a full guide on inspecting training results.
:meth:`Checkpoint.as_directory <ray.train.Checkpoint.as_directory>`
and :meth:`Checkpoint.to_directory <ray.train.Checkpoint.to_directory>`
are the two main APIs to interact with Train checkpoints:
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_checkpoint_example_start__
:end-before: __inspect_checkpoint_example_end__
For Lightning and Transformers, if you are using the default `RayTrainReportCallback` for checkpoint saving in your training function,
you can retrieve the original checkpoint files as below:
.. tab-set::
.. tab-item:: PyTorch Lightning
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_lightning_checkpoint_example_start__
:end-before: __inspect_lightning_checkpoint_example_end__
.. tab-item:: Transformers
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __inspect_transformers_checkpoint_example_start__
:end-before: __inspect_transformers_checkpoint_example_end__
.. _train-dl-loading-checkpoints:
Restore training state from a checkpoint
----------------------------------------
In order to enable fault tolerance, you should modify your training loop to restore
training state from a :class:`~ray.train.Checkpoint`.
The :class:`Checkpoint <ray.train.Checkpoint>` to restore from can be accessed in the
training function with :func:`ray.train.get_checkpoint <ray.train.get_checkpoint>`.
The checkpoint returned by :func:`ray.train.get_checkpoint <ray.train.get_checkpoint>` is populated
as the latest reported checkpoint during :ref:`automatic failure recovery <train-fault-tolerance>`.
See :ref:`train-fault-tolerance` for more details on restoration and fault tolerance.
.. tab-set::
.. tab-item:: Native PyTorch
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __pytorch_restore_start__
:end-before: __pytorch_restore_end__
.. tab-item:: PyTorch Lightning
.. literalinclude:: ../doc_code/checkpoints.py
:language: python
:start-after: __lightning_restore_example_start__
:end-before: __lightning_restore_example_end__
.. note::
In these examples, :meth:`Checkpoint.as_directory <ray.train.Checkpoint.as_directory>`
is used to view the checkpoint contents as a local directory.
*If the checkpoint points to a local directory*, this method just returns the
local directory path without making a copy.
*If the checkpoint points to a remote directory*, this method will download the
checkpoint to a local temporary directory and return the path to the temporary directory.
**If multiple processes on the same node call this method simultaneously,**
only a single process will perform the download, while the others
wait for the download to finish. Once the download finishes, all processes receive
the same local (temporary) directory to read from.
Once all processes have finished working with the checkpoint, the temporary directory
is cleaned up.
doc/source/train/user-guides/data-loading-preprocessing.rst
+755
View File
@@ -0,0 +1,755 @@
.. _data-ingest-torch:
Data Loading and Preprocessing
==============================
Ray Train integrates with :ref:`Ray Data <data>` to offer a performant and scalable streaming solution for loading and preprocessing large datasets.
Key advantages include:
- Streaming data loading and preprocessing, scalable to petabyte-scale data.
- Scaling out heavy data preprocessing to CPU nodes, to avoid bottlenecking GPU training.
- Automatic and fast failure recovery.
- Automatic on-the-fly data splitting across distributed training workers.
For more details about Ray Data, check out the :ref:`Ray Data documentation<data>`.
.. note::
In addition to Ray Data, you can continue to use framework-native data utilities with Ray Train, such as PyTorch Dataset, Hugging Face Dataset, and Lightning DataModule.
In this guide, we will cover how to incorporate Ray Data into your Ray Train script, and different ways to customize your data ingestion pipeline.
.. TODO: Replace this image with a better one.
.. figure:: ../images/train_ingest.png
:align: center
:width: 300px
Quickstart
----------
Install Ray Data and Ray Train:
.. code-block:: bash
pip install -U "ray[data,train]"
Data ingestion can be set up with four basic steps:
1. Create a Ray Dataset from your input data.
2. Apply preprocessing operations to your Ray Dataset.
3. Input the preprocessed Dataset into the Ray Train Trainer, which internally splits the dataset equally in a streaming way across the distributed training workers.
4. Consume the Ray Dataset in your training function.
.. tab-set::
.. tab-item:: PyTorch
.. code-block:: python
:emphasize-lines: 14,21,29,33-35,53
import torch
import ray
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer
# Set this to True to use GPU.
# If False, do CPU training instead of GPU training.
use_gpu = False
# Step 1: Create a Ray Dataset from in-memory Python lists.
# You can also create a Ray Dataset from many other sources and file
# formats.
train_dataset = ray.data.from_items([{"x": [x], "y": [2 * x]} for x in range(200)])
# Step 2: Preprocess your Ray Dataset.
def increment(batch):
batch["y"] = batch["y"] + 1
return batch
train_dataset = train_dataset.map_batches(increment, batch_size="auto")
def train_func():
batch_size = 16
# Step 4: Access the dataset shard for the training worker via
# ``get_dataset_shard``.
train_data_shard = train.get_dataset_shard("train")
# `iter_torch_batches` returns an iterable object that
# yield tensor batches. Ray Data automatically moves the Tensor batches
# to GPU if you enable GPU training.
train_dataloader = train_data_shard.iter_torch_batches(
batch_size=batch_size, dtypes=torch.float32
)
for epoch_idx in range(1):
for batch in train_dataloader:
inputs, labels = batch["x"], batch["y"]
assert type(inputs) == torch.Tensor
assert type(labels) == torch.Tensor
assert inputs.shape[0] == batch_size
assert labels.shape[0] == batch_size
# Only check one batch for demo purposes.
# Replace the above with your actual model training code.
break
# Step 3: Create a TorchTrainer. Specify the number of training workers and
# pass in your Ray Dataset.
# The Ray Dataset is automatically split across all training workers.
trainer = TorchTrainer(
train_func,
datasets={"train": train_dataset},
scaling_config=ScalingConfig(num_workers=2, use_gpu=use_gpu)
)
result = trainer.fit()
.. tab-item:: PyTorch Lightning
.. code-block:: python
:emphasize-lines: 4-5,10-11,14-15,26-27,33
from ray import train
# Create the train and validation datasets.
train_data = ray.data.read_csv("./train.csv")
val_data = ray.data.read_csv("./validation.csv")
def train_func_per_worker():
# Access Ray datsets in your train_func via ``get_dataset_shard``.
# Ray Data shards all datasets across workers by default.
train_ds = train.get_dataset_shard("train")
val_ds = train.get_dataset_shard("validation")
# Create Ray dataset iterables via ``iter_torch_batches``.
train_dataloader = train_ds.iter_torch_batches(batch_size=16)
val_dataloader = val_ds.iter_torch_batches(batch_size=16)
...
trainer = pl.Trainer(
# ...
)
# Feed the Ray dataset iterables to ``pl.Trainer.fit``.
trainer.fit(
model,
train_dataloaders=train_dataloader,
val_dataloaders=val_dataloader
)
trainer = TorchTrainer(
train_func,
# You can pass in multiple datasets to the Trainer.
datasets={"train": train_data, "validation": val_data},
scaling_config=ScalingConfig(num_workers=4),
)
trainer.fit()
.. tab-item:: HuggingFace Transformers
.. code-block:: python
:emphasize-lines: 7-9,14-15,18-19,25,31-32,42
import ray
import ray.train
from huggingface_hub import HfFileSystem
...
# Create the train and evaluation datasets using HfFileSystem.
fs = HfFileSystem()
train_data = ray.data.read_parquet("hf://datasets/your-dataset/train/", filesystem=fs)
eval_data = ray.data.read_parquet("hf://datasets/your-dataset/validation/", filesystem=fs)
def train_func():
# Access Ray datsets in your train_func via ``get_dataset_shard``.
# Ray Data shards all datasets across workers by default.
train_ds = ray.train.get_dataset_shard("train")
eval_ds = ray.train.get_dataset_shard("evaluation")
# Create Ray dataset iterables via ``iter_torch_batches``.
train_iterable_ds = train_ds.iter_torch_batches(batch_size=16)
eval_iterable_ds = eval_ds.iter_torch_batches(batch_size=16)
...
args = transformers.TrainingArguments(
...,
max_steps=max_steps # Required for iterable datasets
)
trainer = transformers.Trainer(
...,
model=model,
train_dataset=train_iterable_ds,
eval_dataset=eval_iterable_ds,
)
# Prepare your Transformers Trainer
trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
trainer.train()
trainer = TorchTrainer(
train_func,
# You can pass in multiple datasets to the Trainer.
datasets={"train": train_data, "evaluation": val_data},
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
trainer.fit()
.. _train-datasets-load:
Loading data
~~~~~~~~~~~~
Ray Datasets can be created from many different data sources and formats. For more details, see :ref:`Loading Data <loading_data>`.
.. _train-datasets-preprocess:
Preprocessing data
~~~~~~~~~~~~~~~~~~
Ray Data supports a wide range of preprocessing operations that you can use to transform data prior to training.
- For general preprocessing, see :ref:`Transforming Data <transforming_data>`.
- For tabular data, see :ref:`Preprocessing Structured Data <preprocessing_structured_data>`.
- For PyTorch tensors, see :ref:`Transformations with torch tensors <transform_pytorch>`.
- For optimizing expensive preprocessing operations, see :ref:`Caching the preprocessed dataset <dataset_cache_performance>`.
.. _train-datasets-input:
Inputting and splitting data
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Your preprocessed datasets can be passed into a Ray Train Trainer (e.g. :class:`~ray.train.torch.TorchTrainer`) through the ``datasets`` argument.
The datasets passed into the Trainer's ``datasets`` can be accessed inside of the ``train_loop_per_worker`` run on each distributed training worker by calling :meth:`ray.train.get_dataset_shard`.
Ray Data splits all datasets across the training workers by default. :meth:`~ray.train.get_dataset_shard` returns ``1/n`` of the dataset, where ``n`` is the number of training workers.
Ray Data does data splitting in a streaming fashion on the fly.
.. note::
Be aware that because Ray Data splits the evaluation dataset, you have to aggregate the evaluation results across workers.
You might consider using `TorchMetrics <https://torchmetrics.readthedocs.io/en/latest/>`_ (:doc:`example <../examples/deepspeed/deepspeed_example>`) or
utilities available in other frameworks that you can explore.
This behavior can be overwritten by passing in the ``dataset_config`` argument. For more information on configuring splitting logic, see :ref:`Splitting datasets <train-datasets-split>`.
.. _train-datasets-consume:
Consuming data
~~~~~~~~~~~~~~
Inside the ``train_loop_per_worker``, each worker can access its shard of the dataset via :meth:`ray.train.get_dataset_shard`.
This data can be consumed in a variety of ways:
- To create a generic Iterable of batches, you can call :meth:`~ray.data.DataIterator.iter_batches`.
- To create a replacement for a PyTorch DataLoader, you can call :meth:`~ray.data.DataIterator.iter_torch_batches`.
For more details on how to iterate over your data, see :ref:`Iterating over data <iterating-over-data>`.
.. _train-datasets-pytorch:
Starting with PyTorch data
--------------------------
Some frameworks provide their own dataset and data loading utilities. For example:
- **PyTorch:** `Dataset & DataLoader <https://pytorch.org/tutorials/beginner/basics/data_tutorial.html>`_
- **Hugging Face:** `Dataset <https://huggingface.co/docs/datasets/index>`_
- **PyTorch Lightning:** `LightningDataModule <https://lightning.ai/docs/pytorch/stable/data/datamodule.html>`_
You can still use these framework data utilities directly with Ray Train.
At a high level, you can compare these concepts as follows:
.. list-table::
:header-rows: 1
* - PyTorch API
- HuggingFace API
- Ray Data API
* - `torch.utils.data.Dataset <https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset>`_
- `datasets.Dataset <https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset>`_
- :class:`ray.data.Dataset`
* - `torch.utils.data.DataLoader <https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader>`_
- n/a
- :meth:`ray.data.Dataset.iter_torch_batches`
For more details, see the following sections for each framework:
.. tab-set::
.. tab-item:: PyTorch DataLoader
**Option 1 (with Ray Data):**
1. Convert your PyTorch Dataset to a Ray Dataset.
2. Pass the Ray Dataset into the TorchTrainer via ``datasets`` argument.
3. Inside your ``train_loop_per_worker``, you can access the dataset via :meth:`ray.train.get_dataset_shard`.
4. Create a dataset iterable via :meth:`ray.data.DataIterator.iter_torch_batches`.
For more details, see the :ref:`Migrating from PyTorch Datasets and DataLoaders <migrate_pytorch>`.
**Option 2 (without Ray Data):**
1. Instantiate the Torch Dataset and DataLoader directly in the ``train_loop_per_worker``.
2. Use the :meth:`ray.train.torch.prepare_data_loader` utility to set up the DataLoader for distributed training.
.. tab-item:: LightningDataModule
The ``LightningDataModule`` is created with PyTorch ``Dataset``\s and ``DataLoader``\s. You can apply the same logic here.
.. tab-item:: Hugging Face Dataset
**Option 1 (with Ray Data):**
1. Convert your Hugging Face Dataset to a Ray Dataset. For instructions, see :ref:`Ray Data for Hugging Face <loading_datasets_from_ml_libraries>`.
2. Pass the Ray Dataset into the TorchTrainer via the ``datasets`` argument.
3. Inside your ``train_loop_per_worker``, access the sharded dataset via :meth:`ray.train.get_dataset_shard`.
4. Create a iterable dataset via :meth:`ray.data.DataIterator.iter_torch_batches`.
5. Pass the iterable dataset while initializing ``transformers.Trainer``.
6. Wrap your transformers trainer with the :meth:`ray.train.huggingface.transformers.prepare_trainer` utility.
**Option 2 (without Ray Data):**
1. Instantiate the Hugging Face Dataset directly in the ``train_loop_per_worker``.
2. Pass the Hugging Face Dataset into ``transformers.Trainer`` during initialization.
.. tip::
When using Torch or Hugging Face Datasets directly without Ray Data, make sure to instantiate your Dataset *inside* the ``train_loop_per_worker``.
Instantiating the Dataset outside of the ``train_loop_per_worker`` and passing it in via global scope
can cause errors due to the Dataset not being serializable.
.. note::
When using PyTorch DataLoader with more than 1 worker, you should set the
process start method to be `forkserver` or `spawn`.
:ref:`Forking Ray Actors and Tasks is an anti-pattern <forking-ray-processes-antipattern>` that
can lead to unexpected issues such as deadlocks.
.. code-block:: python
data_loader = DataLoader(
dataset,
num_workers=2,
multiprocessing_context=multiprocessing.get_context("forkserver"),
...
)
.. _train-datasets-split:
Splitting datasets
------------------
By default, Ray Train splits all datasets across workers using :meth:`Dataset.streaming_split <ray.data.Dataset.streaming_split>`. Each worker sees a disjoint subset of the data, instead of iterating over the entire dataset.
If want to customize which datasets are split, pass in a :class:`DataConfig <ray.train.DataConfig>` to the Trainer constructor.
For example, to split only the training dataset, do the following:
.. testcode::
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
ds = ray.data.read_text(
"s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)
train_ds, val_ds = ds.train_test_split(0.3)
def train_loop_per_worker():
# Get the sharded training dataset
train_ds = train.get_dataset_shard("train")
for _ in range(2):
for batch in train_ds.iter_batches(batch_size=128):
print("Do some training on batch", batch)
# Get the unsharded full validation dataset
val_ds = train.get_dataset_shard("val")
for _ in range(2):
for batch in val_ds.iter_batches(batch_size=128):
print("Do some evaluation on batch", batch)
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": train_ds, "val": val_ds},
dataset_config=ray.train.DataConfig(
datasets_to_split=["train"],
),
)
my_trainer.fit()
Full customization (advanced)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For use cases not covered by the default config class, you can also fully customize exactly how your input datasets are split. Define a custom :class:`DataConfig <ray.train.DataConfig>` class (DeveloperAPI). The :class:`DataConfig <ray.train.DataConfig>` class is responsible for that shared setup and splitting of data across nodes.
.. testcode::
# Note that this example class is doing the same thing as the basic DataConfig
# implementation included with Ray Train.
from typing import Optional, Dict, List
import ray
from ray import train
from ray.train.torch import TorchTrainer
from ray.train import DataConfig, ScalingConfig
from ray.data import Dataset, DataIterator, NodeIdStr
from ray.actor import ActorHandle
ds = ray.data.read_text(
"s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)
def train_loop_per_worker():
# Get an iterator to the dataset we passed in below.
it = train.get_dataset_shard("train")
for _ in range(2):
for batch in it.iter_batches(batch_size=128):
print("Do some training on batch", batch)
class MyCustomDataConfig(DataConfig):
def configure(
self,
datasets: Dict[str, Dataset],
world_size: int,
worker_handles: Optional[List[ActorHandle]],
worker_node_ids: Optional[List[NodeIdStr]],
**kwargs,
) -> List[Dict[str, DataIterator]]:
assert len(datasets) == 1, "This example only handles the simple case"
# Configure Ray Data for ingest.
ctx = ray.data.DataContext.get_current()
ctx.execution_options = DataConfig.default_ingest_options()
# Split the stream into shards.
iterator_shards = datasets["train"].streaming_split(
world_size, equal=True, locality_hints=worker_node_ids
)
# Return the assigned iterators for each worker.
return [{"train": it} for it in iterator_shards]
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": ds},
dataset_config=MyCustomDataConfig(),
)
my_trainer.fit()
The subclass must be serializable, since Ray Train copies it from the driver script to the driving actor of the Trainer. Ray Train calls its :meth:`configure <ray.train.DataConfig.configure>` method on the main actor of the Trainer group to create the data iterators for each worker.
In general, you can use :class:`DataConfig <ray.train.DataConfig>` for any shared setup that has to occur ahead of time before the workers start iterating over data. The setup runs at the start of each Trainer run.
Random shuffling
----------------
Randomly shuffling data for each epoch can be important for model quality depending on what model you are training.
Ray Data provides multiple options for random shuffling, see :ref:`Shuffling Data <shuffling_data>` for more details.
Enabling reproducibility
------------------------
When developing or hyperparameter tuning models, reproducibility is important during data ingest so that data ingest does not affect model quality. Follow these three steps to enable reproducibility:
**Step 1:** Enable deterministic execution in Ray Datasets by setting the `preserve_order` flag in the :class:`DataContext <ray.data.context.DataContext>`.
.. testcode::
import ray
# Preserve ordering in Ray Datasets for reproducibility.
ctx = ray.data.DataContext.get_current()
ctx.execution_options.preserve_order = True
ds = ray.data.read_text(
"s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)
**Step 2:** Set a seed for any shuffling operations:
* `seed` argument to :meth:`random_shuffle <ray.data.Dataset.random_shuffle>`
* `seed` argument to :meth:`randomize_block_order <ray.data.Dataset.randomize_block_order>`
* `local_shuffle_seed` argument to :meth:`iter_batches <ray.data.DataIterator.iter_batches>`
**Step 3:** Follow the best practices for enabling reproducibility for your training framework of choice. For example, see the `Pytorch reproducibility guide <https://pytorch.org/docs/stable/notes/randomness.html>`_.
.. _preprocessing_structured_data:
Preprocessing structured data
-----------------------------
.. note::
This section is for tabular/structured data. The recommended way for preprocessing unstructured data is to use
Ray Data operations such as `map_batches`. See the :ref:`Ray Data Working with Pytorch guide <working_with_pytorch>` for more details.
For tabular data, use Ray Data :ref:`preprocessors <preprocessor-ref>`, which implement common data preprocessing operations.
You can use this with Ray Train Trainers by applying them on the dataset before passing the dataset into a Trainer. For example:
.. testcode::
import base64
import numpy as np
from tempfile import TemporaryDirectory
import ray
from ray import train
from ray.train import Checkpoint, ScalingConfig
from ray.train.torch import TorchTrainer
from ray.data.preprocessors import Concatenator, StandardScaler
dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
# Create preprocessors to scale some columns and concatenate the results.
scaler = StandardScaler(columns=["mean radius", "mean texture"])
columns_to_concatenate = dataset.columns()
columns_to_concatenate.remove("target")
concatenator = Concatenator(columns=columns_to_concatenate, dtype=np.float32)
# Compute dataset statistics and get transformed datasets. Note that the
# fit call is executed immediately, but the transformation is lazy.
dataset = scaler.fit_transform(dataset)
dataset = concatenator.fit_transform(dataset)
def train_loop_per_worker():
context = train.get_context()
print(context.get_metadata()) # prints {"preprocessor_pkl": ...}
# Get an iterator to the dataset we passed in below.
it = train.get_dataset_shard("train")
for _ in range(2):
# Prefetch 10 batches at a time.
for batch in it.iter_batches(batch_size=128, prefetch_batches=10):
print("Do some training on batch", batch)
# Save a checkpoint.
with TemporaryDirectory() as temp_dir:
train.report(
{"score": 2.0},
checkpoint=Checkpoint.from_directory(temp_dir),
)
# Serialize the preprocessor. Since serialize() returns bytes,
# convert to base64 string for JSON compatibility.
serialized_preprocessor = base64.b64encode(scaler.serialize()).decode("ascii")
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": dataset},
metadata={"preprocessor_pkl": serialized_preprocessor},
)
# Get the fitted preprocessor back from the result metadata.
metadata = my_trainer.fit().checkpoint.get_metadata()
# Decode from base64 before deserializing
serialized_data = base64.b64decode(metadata["preprocessor_pkl"])
print(StandardScaler.deserialize(serialized_data))
This example persists the fitted preprocessor using the ``Trainer(metadata={...})`` constructor argument. This arg specifies a dict that is available from ``TrainContext.get_metadata()`` and ``checkpoint.get_metadata()`` for checkpoints that the Trainer saves. This design enables the recreation of the fitted preprocessor for inference.
Performance tips
----------------
Prefetching batches
~~~~~~~~~~~~~~~~~~~
While iterating over a dataset for training, you can increase ``prefetch_batches`` in :meth:`iter_batches <ray.data.DataIterator.iter_batches>` or :meth:`iter_torch_batches <ray.data.DataIterator.iter_torch_batches>` to further increase performance. While training on the current batch, this approach launches background threads to fetch and process the next ``N`` batches.
This approach can help if training is bottlenecked on cross-node data transfer or on last-mile preprocessing such as converting batches to tensors or executing ``collate_fn``. However, increasing ``prefetch_batches`` leads to more data that needs to be held in heap memory. By default, ``prefetch_batches`` is set to 1.
For example, the following code prefetches 10 batches at a time for each training worker:
.. testcode::
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
ds = ray.data.read_text(
"s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
)
def train_loop_per_worker():
# Get an iterator to the dataset we passed in below.
it = train.get_dataset_shard("train")
for _ in range(2):
# Prefetch 10 batches at a time.
for batch in it.iter_batches(batch_size=128, prefetch_batches=10):
print("Do some training on batch", batch)
my_trainer = TorchTrainer(
train_loop_per_worker,
scaling_config=ScalingConfig(num_workers=2),
datasets={"train": ds},
)
my_trainer.fit()
Avoid heavy transformation in collate_fn
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The ``collate_fn`` parameter in :meth:`iter_batches <ray.data.DataIterator.iter_batches>` or :meth:`iter_torch_batches <ray.data.DataIterator.iter_torch_batches>` allows you to transform data before feeding it to the model. This operation happens locally in the training workers. Avoid adding a heavy transformation in this function as it may become the bottleneck. Instead, :ref:`apply the transformation with map or map_batches <transforming_data>` before passing the dataset to the Trainer. When your expensive transformation requires batch_size as input, such as text tokenization, you can :ref:`scale it out to Ray Data <scaling_collation_functions>` for better performance.
.. _dataset_cache_performance:
Caching the preprocessed dataset
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If your preprocessed Dataset is small enough to fit in Ray object store memory (by default this is 30% of total cluster RAM), *materialize* the preprocessed dataset in Ray's built-in object store, by calling :meth:`materialize() <ray.data.Dataset.materialize>` on the preprocessed dataset. This method tells Ray Data to compute the entire preprocessed and pin it in the Ray object store memory. As a result, when iterating over the dataset repeatedly, the preprocessing operations do not need to be re-run. However, if the preprocessed data is too large to fit into Ray object store memory, this approach will greatly decreases performance as data needs to be spilled to and read back from disk.
Transformations that you want to run per-epoch, such as randomization, should go after the materialize call.
.. testcode::
from typing import Dict
import numpy as np
import ray
# Load the data.
train_ds = ray.data.read_parquet("s3://anonymous@ray-example-data/iris.parquet")
# Define a preprocessing function.
def normalize_length(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
new_col = batch["sepal.length"] / np.max(batch["sepal.length"])
batch["normalized.sepal.length"] = new_col
del batch["sepal.length"]
return batch
# Preprocess the data. Transformations that are made before the materialize call
# below are only run once.
train_ds = train_ds.map_batches(normalize_length, batch_size="auto")
# Materialize the dataset in object store memory.
# Only do this if train_ds is small enough to fit in object store memory.
train_ds = train_ds.materialize()
# Dummy augmentation transform.
def augment_data(batch):
return batch
# Add per-epoch preprocessing. Transformations that you want to run per-epoch, such
# as data augmentation or randomization, should go after the materialize call.
train_ds = train_ds.map_batches(augment_data, batch_size="auto")
# Pass train_ds to the Trainer
Adding CPU-only nodes to your cluster
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If the GPU training is bottlenecked on expensive CPU preprocessing and the preprocessed Dataset is too large to fit in object store memory, then materializing the dataset doesn't work. In this case, Ray's native support for heterogeneous resources enables you to simply add more CPU-only nodes to your cluster, and Ray Data automatically scales out CPU-only preprocessing tasks to CPU-only nodes, making GPUs more saturated.
In general, adding CPU-only nodes can help in two ways:
* Adding more CPU cores helps further parallelize preprocessing. This approach is helpful when CPU compute time is the bottleneck.
* Increasing object store memory, which 1) allows Ray Data to buffer more data in between preprocessing and training stages, and 2) provides more memory to make it possible to :ref:`cache the preprocessed dataset <dataset_cache_performance>`. This approach is helpful when memory is the bottleneck.
Isolating Ray Data worker processes from training nodes
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
You may sometimes want to prevent Ray Data CPU tasks from running on training worker nodes
when training workers themselves run CPU or RAM-heavy operations
such as storing large local shuffle buffers or running expensive collate functions.
Launching more Ray Data processes would oversubscribe the training worker nodes.
Instead, the Ray Data tasks should run on a separate set of CPU nodes in your heterogeneous
cluster (for example, 4 GPU training nodes and 4 CPU-only nodes).
One workaround is to force full-node exclusion by reserving all CPUs per training worker via
``resources_per_worker={"CPU": node_cpus // num_gpus_per_node, "GPU": 1}`` in ``ScalingConfig``.
This method is fragile since it's tied to node shapes, and Ray Data also doesn't exclude other
resources such as object store memory properly, since the typical configuration is to only take up logical
CPUs and GPUs.
The recommended approach is to use :ref:`subclusters <data_concurrent_execution>` to pin the
training Dataset to CPU-only nodes. This correctly scopes the memory budget to only the nodes
where data tasks can actually run. It requires adding labels to your worker node configurations and setting the
``label_selector`` in two places:
.. code-block:: python
import ray
from ray.data import ExecutionOptions
from ray.train import DataConfig
from ray.train.torch import TorchTrainer
# (1) Pin construction-time tasks (schema inference, file listing).
ctx = ray.data.DataContext.get_current().copy()
ctx.execution_options.label_selector = {"ray-subcluster": "data"}
with ray.data.DataContext.current(ctx):
train_dataset = ray.data.read_parquet(...)
# (2) Pin per-worker ingest — Train replaces ds.context options
# wholesale, so the selector must be restated here.
trainer = TorchTrainer(
...,
datasets={"train": train_dataset},
dataset_config=DataConfig(
datasets_to_split=["train"],
execution_options={
"train": ExecutionOptions(
label_selector={"ray-subcluster": "data"}
),
},
),
)
.. tip::
Before resorting to isolating Ray Data tasks from training nodes, consider offloading that
heavy work from training workers to the data pipeline instead:
:ref:`scale out expensive collation <scaling_collation_functions>`
and use :ref:`map_batches-based shuffling <map_batches_shuffle>` in place of large local
shuffle buffers. These reduce CPU pressure on training workers and often eliminate the
need for node isolation entirely.
.. _balancing-data-production-consumption:
Balancing data production and consumption
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Ideally, data production (dataset processing) and data consumption (training ingestion) happen at
the same rate. When production outpaces consumption, excess data is written to the
:ref:`object store <object-spilling-internals>`, which can spill to disk, which in turn decreases
Ray Data throughput and leads to out of disk errors. Ray Data's backpressure system automatically
balances production and consumption, but if you are still running into issues, you can try
tuning the following:
* **Use fewer CPUs for data production**: If you are using :func:`~ray.data.Dataset.map_batches`,
you can set the number of workers with ``compute`` and the number of CPUs per worker with ``num_cpus``.
* **Limit object store usage per dataset**: Set a per-dataset object store memory limit using
each dataset's execution options. Ray Data's backpressure system will slow down production
once the object store memory limit is reached.
.. code-block:: python
train_ds = ray.data.read_parquet("s3://bucket/train")
val_ds = ray.data.read_parquet("s3://bucket/val")
train_ds.context.execution_options.resource_limits = ray.data.ExecutionResources(
object_store_memory=50 * 1024**3,
)
val_ds.context.execution_options.resource_limits = ray.data.ExecutionResources(
object_store_memory=50 * 1024**3,
)
See :ref:`data_performance_tips` for more info on how to tune Ray Data.
More data ingest guides
-----------------------
- :ref:`Weighted Dataset Mixing <mixing_data>` — combine multiple datasets with target row ratios for training.
- :ref:`Scaling Collation Functions <scaling_collation_functions>` — scale out expensive collation functions to Ray Data.

Some files were not shown because too many files have changed in this diff Show More