Files
2026-07-13 13:17:40 +08:00

137 lines
5.4 KiB
ReStructuredText

.. _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/>`_.