137 lines
5.4 KiB
ReStructuredText
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/>`_.
|