chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,20 @@
|
||||
# Minimal makefile for Sphinx documentation
|
||||
#
|
||||
|
||||
# You can set these variables from the command line, and also
|
||||
# from the environment for the first two.
|
||||
SPHINXOPTS ?=
|
||||
SPHINXBUILD ?= sphinx-build
|
||||
SOURCEDIR = source
|
||||
BUILDDIR = build
|
||||
|
||||
# Put it first so that "make" without argument is like "make help".
|
||||
help:
|
||||
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
|
||||
|
||||
.PHONY: help Makefile
|
||||
|
||||
# Catch-all target: route all unknown targets to Sphinx using the new
|
||||
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
|
||||
%: Makefile
|
||||
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
|
||||
Executable
+4
@@ -0,0 +1,4 @@
|
||||
#!/bin/bash
|
||||
|
||||
sphinx-apidoc -f -o source ../../deepspeed
|
||||
make html
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 677 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 235 KiB |
@@ -0,0 +1,38 @@
|
||||
Activation Checkpointing
|
||||
========================
|
||||
|
||||
The activation checkpointing API's in DeepSpeed can be used to enable a range
|
||||
of memory optimizations relating to activation checkpointing. These include
|
||||
activation partitioning across GPUs when using model parallelism, CPU
|
||||
checkpointing, contiguous memory optimizations, etc.
|
||||
|
||||
Please see the `DeepSpeed JSON config <https://www.deepspeed.ai/docs/config-json/>`_
|
||||
for the full set.
|
||||
|
||||
Here we present the activation checkpointing API. Please see the enabling
|
||||
DeepSpeed for `Megatron-LM tutorial <https://www.deepspeed.ai/tutorials/megatron/>`_
|
||||
for example usage.
|
||||
|
||||
Configuring Activation Checkpointing
|
||||
------------------------------------
|
||||
.. autofunction:: deepspeed.checkpointing.configure
|
||||
|
||||
.. autofunction:: deepspeed.checkpointing.is_configured
|
||||
|
||||
|
||||
Using Activation Checkpointing
|
||||
------------------------------
|
||||
.. autofunction:: deepspeed.checkpointing.checkpoint
|
||||
|
||||
.. autofunction:: deepspeed.checkpointing.reset
|
||||
|
||||
|
||||
Configuring and Checkpointing Random Seeds
|
||||
------------------------------------------
|
||||
.. autofunction:: deepspeed.checkpointing.get_cuda_rng_tracker
|
||||
|
||||
.. autofunction:: deepspeed.checkpointing.model_parallel_cuda_manual_seed
|
||||
|
||||
.. autoclass:: deepspeed.checkpointing.CudaRNGStatesTracker
|
||||
|
||||
.. autoclass:: deepspeed.checkpointing.CheckpointFunction
|
||||
@@ -0,0 +1,106 @@
|
||||
AutoEP (Automatic Expert Parallelism)
|
||||
=====================================
|
||||
|
||||
AutoEP automatically detects MoE layers in Hugging Face models and replaces them
|
||||
with EP-enabled versions, requiring zero model code changes. It follows the
|
||||
pattern of AutoTP (Automatic Tensor Parallelism).
|
||||
|
||||
This API is separate from the explicit ``deepspeed.moe.layer.MoE`` layer API.
|
||||
For the explicit DeepSpeed MoE layer API, see :doc:`moe`.
|
||||
|
||||
**Built-in AutoEP presets:** ``mixtral`` (Mixtral), ``qwen3_moe`` (Qwen3-MoE),
|
||||
``qwen3_5_moe`` (Qwen3.5-MoE), ``deepseek_v2`` (DeepSeek-V2), and
|
||||
``deepseek_v3`` (DeepSeek-V3).
|
||||
|
||||
The preset name means AutoEP knows the router, expert, and weight naming
|
||||
patterns for that model family. Running a Hugging Face model also requires a
|
||||
Transformers build that exposes the matching config/model classes,
|
||||
``model.config.model_type`` value, and fused expert layout.
|
||||
|
||||
.. list-table:: AutoEP preset compatibility by Transformers version
|
||||
:header-rows: 1
|
||||
|
||||
* - Preset
|
||||
- Minimum Transformers version
|
||||
- Notes
|
||||
* - ``mixtral``
|
||||
- ``5.0.0``
|
||||
-
|
||||
* - ``qwen3_moe``
|
||||
- ``5.0.0``
|
||||
- Also covers Qwen2-MoE when the installed Transformers build uses the
|
||||
validated fused expert layout. Qwen3-MoE classes appear in ``4.51.3``,
|
||||
but the tested ``4.x`` builds do not match the validated AutoEP layout.
|
||||
* - ``qwen3_5_moe``
|
||||
- ``5.2.0``
|
||||
- Requires the Qwen3.5 text-backbone ``qwen3_5_moe_text`` model type;
|
||||
for performance on Qwen3.5's Gated DeltaNet layers, install optimized
|
||||
kernels. See the `Hugging Face Transformers kernel loading docs
|
||||
<https://huggingface.co/docs/transformers/kernel_doc/loading_kernels>`__
|
||||
and the `Qwen FlashQLA blog <https://qwen.ai/blog?id=flashqla>`__.
|
||||
* - ``deepseek_v2``
|
||||
- ``5.0.0``
|
||||
- ``load_balance_coeff`` / expert-bias auxiliary-loss-free load balancing
|
||||
is not currently supported; non-null values are rejected.
|
||||
* - ``deepseek_v3``
|
||||
- ``5.0.0``
|
||||
- ``load_balance_coeff`` / expert-bias auxiliary-loss-free load balancing
|
||||
is not currently supported; non-null values are rejected.
|
||||
|
||||
**ZeRO compatibility:** Stages 0, 1, and 2, plus constrained Stage 3
|
||||
support. Stage 3 requires AutoEP-managed MoE layers and does not support native
|
||||
DeepSpeed MoE layers, AutoTP, tensor model parallelism from ``mpu``, sequence
|
||||
parallelism, MiCS, hpZeRO secondary tensor groups, non-1 expert tensor
|
||||
parallelism, or quantized gradients. Stage 3 AutoEP checkpoints are saved
|
||||
partition-natively in the ``zero_pp_rank_*`` shard files and support
|
||||
same-topology load, module-only loads (``load_module_only``),
|
||||
optimizer-state-free loads (``load_optimizer_states=False``), and Universal
|
||||
Checkpoint conversion. Optimizer-including Universal Checkpoint loads can
|
||||
resume with a different data-parallel world size, a different ``autoep_size``,
|
||||
or both, when the target ``autoep_size`` divides the model's expert count.
|
||||
Weights-only/module-only Universal Checkpoint loads use the converted
|
||||
``fp32.pt`` parameter files and support the same data-parallel and
|
||||
``autoep_size`` topology changes.
|
||||
|
||||
**Usage:**
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"expert_parallel": {
|
||||
"enabled": true,
|
||||
"autoep_size": 4,
|
||||
"preset_model": "mixtral"
|
||||
}
|
||||
}
|
||||
|
||||
**How it works:**
|
||||
|
||||
1. During ``deepspeed.initialize()``, AutoEP scans the model for MoE layers
|
||||
using preset-defined patterns (router name, expert name, weight shapes).
|
||||
2. Detected MoE blocks are replaced with ``AutoEPMoELayer``, which uses
|
||||
TorchTitan's grouped GEMM kernels and AllToAll token dispatch.
|
||||
3. EP/EDP process groups are created automatically based on ``autoep_size``.
|
||||
4. Expert parameters are marked for expert-data-parallel gradient reduction;
|
||||
router and shared-expert parameters use standard data-parallel reduction.
|
||||
|
||||
**Constraints:**
|
||||
|
||||
- ``autoep_size`` must divide ``num_experts`` for all detected MoE layers.
|
||||
- ``autoep_size=1`` is valid: all experts remain local (no AllToAll), useful
|
||||
for functional testing on a single GPU.
|
||||
- AutoEP currently cannot be combined with AutoTP
|
||||
(``tensor_parallel.autotp_size > 1``) or tensor model parallelism from
|
||||
``mpu``; support is planned as follow-up work.
|
||||
- AutoEP with ZeRO Stage 3 is supported only without sequence parallelism,
|
||||
MiCS, hpZeRO secondary tensor groups, non-1 expert tensor parallelism, or
|
||||
quantized gradients.
|
||||
- Regular checkpoint save/load requires matching ``autoep_size``. To change
|
||||
``autoep_size`` or data-parallel world size across runs for the same
|
||||
AutoEP-detected model topology, convert the checkpoint to Universal
|
||||
Checkpoint format and load it with ``checkpoint.load_universal``; see the
|
||||
`Universal Checkpointing tutorial </tutorials/universal-checkpointing/>`__
|
||||
for the detailed flow and constraints.
|
||||
- DeepSeek-V2 and DeepSeek-V3 AutoEP do not support load-balance expert bias
|
||||
yet. The built-in DeepSeek presets disable it by default; explicit non-null
|
||||
values fail.
|
||||
@@ -0,0 +1,14 @@
|
||||
Autotuning
|
||||
==============
|
||||
One pain point in model training is to figure out good performance-relevant configurations such as micro-batch size to fully utilize the hardware and achieve a high throughput number. This configuration exploring process is commonly done manually but is important since model training is repeated many times and benefits from using a good configuration. Not only is the hand-tuning process time-consuming, but the outcome is hardware-dependent. This means that a good configuration on one hardware might not be the best on another different hardware. The user thus has to hand tune the configuration again. With DeepSpeed, there are more configuration parameters that could potentially affect the training speed, thus making it more tedious to manually tune the configuration.
|
||||
|
||||
The DeepSpeed Autotuner mitigates this pain point and automatically discovers the optimal DeepSpeed configuration that delivers good training speed.
|
||||
The Autotuner uses model information, system information, and heuristics to efficiently tune system knobs that affect compute and memory efficiencies, such as ZeRO optimization stages, micro-batch sizes, and many other ZeRO optimization configurations.
|
||||
It not only reduces the time and resources users spend on tuning, but also can discover configurations better than hand-tuned methods.
|
||||
|
||||
Please see the `Autotuning tutorial <https://www.deepspeed.ai/tutorials/autotuning/>`_ for usage details.
|
||||
|
||||
Autotuner
|
||||
---------------------------------------------------
|
||||
.. autoclass:: deepspeed.autotuning.autotuner
|
||||
:members:
|
||||
@@ -0,0 +1,101 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# Configuration file for the Sphinx documentation builder.
|
||||
#
|
||||
# This file only contains a selection of the most common options. For a full
|
||||
# list see the documentation:
|
||||
# https://www.sphinx-doc.org/en/master/usage/configuration.html
|
||||
|
||||
# -- Path setup --------------------------------------------------------------
|
||||
|
||||
# If extensions (or modules to document with autodoc) are in another directory,
|
||||
# add these directories to sys.path here. If the directory is relative to the
|
||||
# documentation root, use os.path.abspath to make it absolute, like shown here.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# -- Project information -----------------------------------------------------
|
||||
|
||||
project = 'DeepSpeed'
|
||||
copyright = '2020, Microsoft'
|
||||
author = 'Microsoft'
|
||||
|
||||
# The full version, including alpha/beta/rc tags
|
||||
with open("../../../version.txt", "r") as f:
|
||||
release = f.readline().rstrip()
|
||||
|
||||
master_doc = 'index'
|
||||
|
||||
autodoc_member_order = 'bysource'
|
||||
|
||||
# -- General configuration ---------------------------------------------------
|
||||
|
||||
# Add any Sphinx extension module names here, as strings. They can be
|
||||
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
|
||||
# ones.
|
||||
extensions = [
|
||||
'sphinx.ext.autodoc',
|
||||
'sphinx.ext.napoleon',
|
||||
'sphinx.ext.viewcode',
|
||||
'recommonmark',
|
||||
'sphinx_rtd_theme',
|
||||
'sphinxcontrib.autodoc_pydantic',
|
||||
'sphinx.ext.autosectionlabel',
|
||||
]
|
||||
|
||||
pygments_style = 'sphinx'
|
||||
|
||||
# autodoc_pyandtic config
|
||||
autodoc_pydantic_model_show_field_summary = False
|
||||
autodoc_pydantic_field_signature_prefix = ' '
|
||||
autodoc_pydantic_model_signature_prefix = 'class'
|
||||
autodoc_pydantic_model_show_json = False
|
||||
autodoc_pydantic_model_show_config_summary = False
|
||||
autodoc_pydantic_model_show_config_member = False
|
||||
autodoc_pydantic_model_show_validator_summary = False
|
||||
autodoc_pydantic_model_show_validator_members = False
|
||||
autodoc_pydantic_model_summary_list_order = 'bysource'
|
||||
autodoc_pydantic_model_member_order = 'bysource'
|
||||
autodoc_pydantic_field_list_validators = False
|
||||
|
||||
# Add any paths that contain templates here, relative to this directory.
|
||||
templates_path = ['_templates']
|
||||
|
||||
# List of patterns, relative to source directory, that match files and
|
||||
# directories to ignore when looking for source files.
|
||||
# This pattern also affects html_static_path and html_extra_path.
|
||||
exclude_patterns = []
|
||||
|
||||
# -- Options for HTML output -------------------------------------------------
|
||||
|
||||
# The theme to use for HTML and HTML Help pages. See the documentation for
|
||||
# a list of builtin themes.
|
||||
#
|
||||
html_theme = 'sphinx_rtd_theme'
|
||||
|
||||
# Add any paths that contain custom static files (such as style sheets) here,
|
||||
# relative to this directory. They are copied after the builtin static files,
|
||||
# so a file named "default.css" will overwrite the builtin "default.css".
|
||||
html_static_path = ['_static']
|
||||
|
||||
# GitHub integration
|
||||
html_context = {
|
||||
"display_github": True,
|
||||
"github_user": "microsoft",
|
||||
"github_repo": "DeepSpeed",
|
||||
"github_version": "master",
|
||||
"conf_py_path": "/docs/code-docs/source/",
|
||||
}
|
||||
|
||||
sys.path.insert(0, os.path.abspath('../../../'))
|
||||
|
||||
# Prepend module names to class descriptions?
|
||||
add_module_names = True
|
||||
|
||||
autoclass_content = 'both'
|
||||
|
||||
autodoc_mock_imports = ["apex", "mpi4py", "tensorboardX", "numpy", "cupy", "einops", "transformers", "msgpack"]
|
||||
@@ -0,0 +1,16 @@
|
||||
Flops Profiler
|
||||
|
||||
==============
|
||||
|
||||
The flops profiler in DeepSpeed profiles the forward pass of a model and measures its parameters, latency, and floating point operations. The DeepSpeed flops profiler can be used with the DeepSpeed runtime or as a standalone package.
|
||||
|
||||
When using DeepSpeed for model training, the flops profiler can be configured in the deepspeed_config file without user code changes. To use the flops profiler outside of the DeepSpeed runtime, one can simply install DeepSpeed and import the flops_profiler package to use the APIs directly.
|
||||
|
||||
Please see the `Flops Profiler tutorial <https://www.deepspeed.ai/tutorials/flops-profiler/>`_ for usage details.
|
||||
|
||||
Flops Profiler
|
||||
---------------------------------------------------
|
||||
|
||||
.. automodule:: deepspeed.profiling.flops_profiler.profiler
|
||||
:members:
|
||||
:show-inheritance:
|
||||
@@ -0,0 +1,114 @@
|
||||
DeepSpeed
|
||||
=========
|
||||
|
||||
Model Setup
|
||||
-----------
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
initialize
|
||||
inference-init
|
||||
|
||||
Training API
|
||||
------------
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
training
|
||||
|
||||
Inference API
|
||||
-------------
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
inference-engine
|
||||
|
||||
Checkpointing API
|
||||
-----------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
model-checkpointing
|
||||
activation-checkpointing
|
||||
|
||||
|
||||
ZeRO API
|
||||
--------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
zero3
|
||||
|
||||
Mixture of Experts
|
||||
------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
autoep
|
||||
moe
|
||||
|
||||
Transformer Kernel API
|
||||
----------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
kernel
|
||||
|
||||
Pipeline Parallelism
|
||||
--------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
pipeline
|
||||
|
||||
Optimizers
|
||||
----------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
optimizers
|
||||
|
||||
Learning Rate Schedulers
|
||||
------------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
schedulers
|
||||
|
||||
Flops Profiler
|
||||
--------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
flops-profiler
|
||||
|
||||
Autotuning
|
||||
--------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
autotuning
|
||||
|
||||
Memory Usage
|
||||
------------------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
memory
|
||||
|
||||
Monitoring
|
||||
----------
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
monitor
|
||||
|
||||
Indices and tables
|
||||
------------------
|
||||
|
||||
* :ref:`genindex`
|
||||
* :ref:`modindex`
|
||||
* :ref:`search`
|
||||
@@ -0,0 +1,15 @@
|
||||
Inference API
|
||||
=============
|
||||
|
||||
:func:`deepspeed.init_inference` returns an *inference engine*
|
||||
of type :class:`InferenceEngine`.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for step, batch in enumerate(data_loader):
|
||||
#forward() method
|
||||
loss = engine(batch)
|
||||
|
||||
Forward Propagation
|
||||
-------------------
|
||||
.. autofunction:: deepspeed.InferenceEngine.forward
|
||||
@@ -0,0 +1,42 @@
|
||||
Inference Setup
|
||||
-----------------------
|
||||
The entrypoint for inference with DeepSpeed is ``deepspeed.init_inference()``.
|
||||
|
||||
Example usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
engine = deepspeed.init_inference(model=net, config=config)
|
||||
|
||||
The ``DeepSpeedInferenceConfig`` is used to control all aspects of initializing
|
||||
the ``InferenceEngine``. The config should be passed as a dictionary to
|
||||
``init_inference``, but parameters can also be passed as keyword arguments.
|
||||
|
||||
.. _DeepSpeedInferenceConfig:
|
||||
.. autopydantic_model:: deepspeed.inference.config.DeepSpeedInferenceConfig
|
||||
|
||||
.. _DeepSpeedTPConfig:
|
||||
.. autopydantic_model:: deepspeed.inference.config.DeepSpeedTPConfig
|
||||
|
||||
.. _DeepSpeedMoEConfig:
|
||||
.. autopydantic_model:: deepspeed.inference.config.DeepSpeedMoEConfig
|
||||
|
||||
.. _QuantizationConfig:
|
||||
.. autopydantic_model:: deepspeed.inference.config.QuantizationConfig
|
||||
|
||||
.. _InferenceCheckpointConfig:
|
||||
.. autopydantic_model:: deepspeed.inference.config.InferenceCheckpointConfig
|
||||
|
||||
|
||||
Example config:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
config = {
|
||||
"kernel_inject": True,
|
||||
"tensor_parallel": {"tp_size": 4},
|
||||
"dtype": "fp16",
|
||||
"enable_cuda_graph": False
|
||||
}
|
||||
|
||||
.. autofunction:: deepspeed.init_inference
|
||||
@@ -0,0 +1,44 @@
|
||||
Training Setup
|
||||
==============
|
||||
|
||||
.. _deepspeed-args:
|
||||
|
||||
Argument Parsing
|
||||
----------------
|
||||
DeepSpeed uses the `argparse <https://docs.python.org/3/library/argparse.html>`_ library to
|
||||
supply commandline configuration to the DeepSpeed runtime. Use ``deepspeed.add_config_arguments()``
|
||||
to add DeepSpeed's builtin arguments to your application's parser.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
parser = argparse.ArgumentParser(description='My training script.')
|
||||
parser.add_argument('--local_rank', type=int, default=-1,
|
||||
help='local rank passed from distributed launcher')
|
||||
# Include DeepSpeed configuration arguments
|
||||
parser = deepspeed.add_config_arguments(parser)
|
||||
cmd_args = parser.parse_args()
|
||||
|
||||
.. autofunction:: deepspeed.add_config_arguments
|
||||
|
||||
|
||||
.. _deepspeed-init:
|
||||
|
||||
Training Initialization
|
||||
-----------------------
|
||||
The entrypoint for all training with DeepSpeed is ``deepspeed.initialize()``. Will initialize distributed backend if it is not initialized already.
|
||||
|
||||
Example usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
|
||||
model=net,
|
||||
model_parameters=net.parameters())
|
||||
|
||||
.. autofunction:: deepspeed.initialize
|
||||
|
||||
Distributed Initialization
|
||||
--------------------------
|
||||
Optional distributed backend initialization separate from ``deepspeed.initialize()``. Useful in scenarios where the user wants to use torch distributed calls before calling ``deepspeed.initialize()``, such as when using model parallelism, pipeline parallelism, or certain data loader scenarios.
|
||||
|
||||
.. autofunction:: deepspeed.init_distributed
|
||||
@@ -0,0 +1,17 @@
|
||||
Transformer Kernels
|
||||
===================
|
||||
|
||||
The transformer kernel API in DeepSpeed can be used to create BERT transformer layer for
|
||||
more efficient pre-training and fine-tuning, it includes the transformer layer configurations and
|
||||
transformer layer module initialization.
|
||||
|
||||
Here we present the transformer kernel API.
|
||||
Please see the `BERT pre-training tutorial <https://www.deepspeed.ai/tutorials/bert-pretraining/>`_ for usage details.
|
||||
|
||||
DeepSpeed Transformer Config
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.DeepSpeedTransformerConfig
|
||||
|
||||
DeepSpeed Transformer Layer
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.DeepSpeedTransformerLayer
|
||||
@@ -0,0 +1,292 @@
|
||||
Memory Requirements
|
||||
-----------------------
|
||||
|
||||
|
||||
API To Estimate Memory Usage
|
||||
============================
|
||||
|
||||
ZeRO2:
|
||||
|
||||
.. autofunction:: deepspeed.runtime.zero.stage_1_and_2.estimate_zero2_model_states_mem_needs_all_live
|
||||
|
||||
.. autofunction:: deepspeed.runtime.zero.stage_1_and_2.estimate_zero2_model_states_mem_needs_all_cold
|
||||
|
||||
Examples:
|
||||
|
||||
Let's try a 3B model with just 1 node with 8 gpus, using live model:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python -c 'from transformers import AutoModel; \
|
||||
from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_live; \
|
||||
model = AutoModel.from_pretrained("t5-3b"); \
|
||||
estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
|
||||
|
||||
Estimated memory needed for params, optim states and gradients for a:
|
||||
HW: Setup with 1 node, 8 GPUs per node.
|
||||
SW: Model with 2851M total params.
|
||||
per CPU | per GPU | Options
|
||||
127.48GB | 5.31GB | offload_optimizer=cpu
|
||||
127.48GB | 15.93GB | offload_optimizer=none
|
||||
|
||||
Now, without the actual model, which requires us to know ``total_params`` and
|
||||
``largest_layer_params``, but we got those from the run above, so future estimators are now much
|
||||
faster as we don't need to load the model.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python -c 'from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_cold; \
|
||||
estimate_zero2_model_states_mem_needs_all_cold(total_params=2851e6, num_gpus_per_node=8, num_nodes=1)'
|
||||
|
||||
Estimated memory needed for params, optim states and gradients for a:
|
||||
HW: Setup with 1 node, 8 GPUs per node.
|
||||
SW: Model with 2851M total params.
|
||||
per CPU | per GPU | Options
|
||||
127.45GB | 5.31GB | offload_optimizer=cpu
|
||||
127.45GB | 15.93GB | offload_optimizer=none
|
||||
|
||||
There is a slight difference due to rounding - the actual live model has a few more params
|
||||
|
||||
|
||||
ZeRO3:
|
||||
|
||||
.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live
|
||||
|
||||
.. autofunction:: deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_cold
|
||||
|
||||
Examples:
|
||||
|
||||
Let's try a 3B model with just 1 node with 8 gpus, using live model:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python -c 'from transformers import AutoModel; \
|
||||
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
|
||||
model = AutoModel.from_pretrained("t5-3b"); \
|
||||
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'
|
||||
|
||||
Estimated memory needed for params, optim states and gradients for a:
|
||||
HW: Setup with 1 node, 8 GPUs per node.
|
||||
SW: Model with 2851M total params, 32M largest layer params.
|
||||
per CPU | per GPU | Options
|
||||
71.71GB | 0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
|
||||
127.48GB | 0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
|
||||
63.74GB | 0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=1
|
||||
127.48GB | 0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=0
|
||||
1.47GB | 6.10GB | offload_param=none, offload_optimizer=none, zero_init=1
|
||||
127.48GB | 6.10GB | offload_param=none, offload_optimizer=none, zero_init=0
|
||||
|
||||
Now, without the actual model, which requires us to know ``total_params`` and
|
||||
``largest_layer_params``, but we got those from the run above, so future estimators are now much
|
||||
faster as we don't need to load the model.
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
python -c 'from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_cold; \
|
||||
estimate_zero3_model_states_mem_needs_all_cold(total_params=2851e6, largest_layer_params=32e6, num_gpus_per_node=8, num_nodes=1)'
|
||||
|
||||
Estimated memory needed for params, optim states and gradients for a:
|
||||
HW: Setup with 1 node, 8 GPUs per node.
|
||||
SW: Model with 2851M total params, 32M largest layer params.
|
||||
per CPU | per GPU | Options
|
||||
71.69GB | 0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
|
||||
127.45GB | 0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
|
||||
63.72GB | 0.78GB | offload_param=none, offload_optimizer=cpu , zero_init=1
|
||||
127.45GB | 0.78GB | offload_param=none, offload_optimizer=cpu , zero_init=0
|
||||
1.43GB | 6.09GB | offload_param=none, offload_optimizer=none, zero_init=1
|
||||
127.45GB | 6.09GB | offload_param=none, offload_optimizer=none, zero_init=0
|
||||
|
||||
There is a slight difference due to rounding - the actual live model has a few more params
|
||||
|
||||
|
||||
|
||||
Discussion
|
||||
==========
|
||||
|
||||
Let's look in detail how the memory estimator API calculates these numbers and also discuss some additional numbers that aren't covered by the API.
|
||||
|
||||
In the following discussion:
|
||||
|
||||
- ``params`` - total number of model params, which can be calculated as:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
print(sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()))
|
||||
|
||||
Some models already include the number of params in the model name, e.g. t5-11b (11B params), gpt-neo-1.3B (1.3B params), etc.
|
||||
|
||||
Also if the model weights are stored in ``fp32`` the other quick way to calculate the size of the model is to simply divide the size of the ``state_dict`` file by 4 (fp32 == 4 bytes). For example, you can see that `t5-11b's pytorch_model.bin <https://huggingface.co/t5-11b/tree/main>`__ is 42.1GB in size, so if we divide it by 4, we can immediately tell it's an 11B model.
|
||||
|
||||
The following calculations show how much memory is required by model params, gradients and optimizer states. In addition to those you will need enough memory to fit activation calculations and any temporary memory for intermediate calculations, which for long sequences could be very significant (e.g. could take the same amount of memory as params+grads+optim_states combined).
|
||||
|
||||
The optimizer states assume that ``Adam`` is used, where 4 bytes per parameter are used by momentum and another 4 by variance (8 in total).
|
||||
|
||||
Gradients at ``fp32`` take 4 bytes, and parameters take 2 bytes at ``fp16`` and 4 bytes at ``fp32``.
|
||||
|
||||
**GPU RAM**
|
||||
|
||||
The big question is how big of a model you can fit on the hardware you have? Or rather what size of a GPU RAM do you need to fit the desired model.
|
||||
|
||||
|
||||
* ZeRO-2:
|
||||
|
||||
- ``"offload_optimizer": {"device": "cpu"}``: 2 * params
|
||||
|
||||
Example: a 40GB GPU can fit ~11B param model (regardless of how many GPUs are used). Here the model is loaded in ``fp16`` so just the model weights take about 22GB and the remaining 18GB are used by other components. You can barely fit a very small batch size in this scenario.
|
||||
|
||||
- ``"offload_optimizer": {"device": "none"}``: 4 * params + 16 * params/ (total number of gpus)
|
||||
|
||||
* ZeRO-3:
|
||||
|
||||
``largest_layer_memory = 4*largest_layer_params`` - GPU memory needed to gather the largest layer on a single GPU. 2 bytes fp16 params are gathered and 2 bytes fp16 grads are computed (total 4x). The optimizer states and fp32 parameters are updated in partitioned form and copied to fp16 params in partitioned form. This happens during the optimizer step. After that the fp16 params are sufficient.
|
||||
|
||||
- case 1: ``"offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"}`` - largest_layer_memory + 18 * params / total number of gpus across all nodes
|
||||
- case 2: ``"offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}``- largest_layer_memory. The main limit here is general RAM.
|
||||
- case 3: ``"offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}``- largest_layer_memory + 2 * params / total number of gpus across all nodes
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from transformers import AutoModel
|
||||
model = AutoModel.from_pretrained("t5-large")
|
||||
|
||||
# shared params calculated only ones
|
||||
total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
|
||||
|
||||
largest_layer_params = 0
|
||||
for m in model.modules():
|
||||
# assuming no shared params within a single layer
|
||||
layer_params = sum(p.numel() for p in m.parameters(recurse=False))
|
||||
largest_layer_params = max(largest_layer_params, layer_params)
|
||||
|
||||
largest_layer_memory = (4*largest_layer_params)
|
||||
|
||||
total_gpus = 4
|
||||
|
||||
case1 = largest_layer_memory + int(18*total_params/total_gpus)
|
||||
case2 = largest_layer_memory
|
||||
case3 = largest_layer_memory + int(2*total_params/total_gpus)
|
||||
|
||||
print(f"total params: {total_params/1e6:6.2f}M")
|
||||
print(f"largest layer params: {largest_layer_params/1e6:6.2f}M")
|
||||
print(f"largest layer memory: {largest_layer_memory>>20:6}MB")
|
||||
print(f"case1 gpu memory: {(case1)>>20:6}MB")
|
||||
print(f"case2 gpu memory: {(case2)>>20:6}MB")
|
||||
print(f"case3 gpu memory: {(case3)>>20:6}MB")
|
||||
|
||||
total params: 737.67M
|
||||
largest layer params: 32.90M
|
||||
largest layer memory: 125MB
|
||||
case1 gpu memory: 3291MB
|
||||
case2 gpu memory: 125MB
|
||||
case3 gpu memory: 477MB
|
||||
|
||||
|
||||
**General RAM**:
|
||||
|
||||
One of the key features of ZeRO is its CPU offload which can dramatically extend the total memory pool accessible to the project by using general RAM. One can easily expand their general RAM by 10x times, at a significantly lower cost than what it'd take to have the same GPU RAM. And often, it's not even possible to buy GPUs with a lot of RAM (112GB GPU anybody?) since they simply don't yet exist.
|
||||
|
||||
In the following calculations we will use:
|
||||
|
||||
- ``additional_buffer_factor=1.5`` as an additional buffer factor to be conservative
|
||||
- ``n_gpus`` the number of GPUs on a single node (machine)
|
||||
- ``total_gpus`` the total number of GPUs across all nodes
|
||||
- ``params`` - total number of model params (see above for how to get this number)
|
||||
|
||||
* ZeRO-2:
|
||||
|
||||
- ``"offload_optimizer": {"device": "none"}``:
|
||||
|
||||
params * 4 * n_gpus * additional_buffer_factor - this is the memory needed only at the beginning to initialize the model on CPU memory
|
||||
|
||||
- ``"offload_optimizer": {"device": "cpu"}``:
|
||||
|
||||
params * max(4 * n_gpus, 16) * additional_buffer_factor
|
||||
|
||||
Example: xxx
|
||||
|
||||
* ZeRO-3:
|
||||
|
||||
gpus_factor = n_gpus / total_gpus
|
||||
|
||||
- case 1: ``"offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"}``:
|
||||
|
||||
Without ``zero.Init``:
|
||||
|
||||
params * 4 * n_gpus * additional_buffer_factor
|
||||
|
||||
this is the memory needed only at the beginning to initialize the model on CPU memory. Once the model is transferred to GPUs this memory is freed.
|
||||
|
||||
With ``zero.Init``:
|
||||
|
||||
largest_layer_params * 4 * n_gpus * additional_buffer_factor
|
||||
|
||||
assuming Pytorch is deallocating the memory once the tensors are moved to the GPU by ZeRO.Init
|
||||
|
||||
- case 2: ``"offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}``:
|
||||
|
||||
Without ``zero.Init``:
|
||||
|
||||
params * max(4 * n_gpus, 18 * gpus_factor) * additional_buffer_factor
|
||||
|
||||
With ``zero.Init``:
|
||||
|
||||
params * 18 * gpus_factor * additional_buffer_factor
|
||||
|
||||
- case 3: ``"offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}``:
|
||||
|
||||
Without ``zero.Init``:
|
||||
|
||||
params * max(4 * n_gpus, 16 * gpus_factor) * additional_buffer_factor
|
||||
|
||||
With ``zero.Init``:
|
||||
|
||||
params * 16 * gpus_factor * additional_buffer_factor
|
||||
|
||||
|
||||
Here is a breakdown for the 16 and 18 multipliers (b = bytes):
|
||||
|
||||
4 (in ``4*n_gpus``):
|
||||
|
||||
- when pytorch creates a model it creates it in fp32 by default (4 bytes)
|
||||
|
||||
16:
|
||||
|
||||
- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
|
||||
|
||||
18:
|
||||
|
||||
- 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
|
||||
- +2b for fp16 params
|
||||
|
||||
Note about gradients: While gradients are stored in fp16 (2 bytes), during the weight update, all of them are converted into fp32 before doing the weight updates since the weight updates are done at almost the entire model granularity (param_group granularity) in FusedAdam Optimizer in DeepSpeed. So after that conversion we would need the 4 bytes per gradient for nearly the entire set of weights.
|
||||
|
||||
|
||||
**Pinned Memory**
|
||||
|
||||
Pinned general RAM is included in normal general RAM allocations (i.e. this is not extra memory allocations but simply shows how much of the general RAM is pinned)
|
||||
|
||||
* ZeRO-2: can't be controlled
|
||||
|
||||
* ZeRO-3
|
||||
|
||||
To enable add: ``"cpu_offload_use_pin_memory" : true``
|
||||
|
||||
Now there are 2 sub-cases:
|
||||
|
||||
1. ``"cpu_offload_params": true``:
|
||||
|
||||
- 6 * params (2b for fp16 params + 4b for fp32 gradients)
|
||||
- if ``gradient_accumulation_steps > 1`` an additional 2b for fp16 gradients are pinned
|
||||
|
||||
2. ``"cpu_offload_params": false``:
|
||||
|
||||
- 4b for fp32 gradients
|
||||
|
||||
|
||||
**Activation Memory**
|
||||
|
||||
XXX: For Transformers is probably around (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu
|
||||
|
||||
This needs to be completed.
|
||||
@@ -0,0 +1,56 @@
|
||||
Model Checkpointing
|
||||
===================
|
||||
|
||||
DeepSpeed provides routines for checkpointing model state during training.
|
||||
|
||||
Loading Training Checkpoints
|
||||
----------------------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.load_checkpoint
|
||||
|
||||
Saving Training Checkpoints
|
||||
---------------------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.save_checkpoint
|
||||
|
||||
|
||||
ZeRO Checkpoint fp32 Weights Recovery
|
||||
-------------------------------------
|
||||
|
||||
DeepSpeed provides routines for extracting fp32 weights from the saved ZeRO checkpoint's optimizer states.
|
||||
|
||||
.. autofunction:: deepspeed.utils.zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint
|
||||
|
||||
.. autofunction:: deepspeed.utils.zero_to_fp32.load_state_dict_from_zero_checkpoint
|
||||
|
||||
.. autofunction:: deepspeed.utils.zero_to_fp32.convert_zero_checkpoint_to_fp32_state_dict
|
||||
|
||||
|
||||
Avoiding ZeRO Checkpoint Bloat
|
||||
------------------------------
|
||||
ZeRO stage 1 and 2 checkpoints created using ``torch.save()`` can sometimes be larger than expected. This bloat
|
||||
is caused by the interaction of ZeRO's tensor flattening and torch's tensor `storage management <https://pytorch.org/docs/stable/notes/serialization.html#preserve-storage-sharing>`_ .
|
||||
You can avoid this problem by using the ``clone_tensors_for_torch_save`` utility of DeepSpeed as illustrated below.
|
||||
|
||||
.. autofunction:: deepspeed.checkpoint.utils.clone_tensors_for_torch_save
|
||||
|
||||
The following code snippet illustrates this functionality for creating a HuggingFace model checkpoint:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ds_config = {
|
||||
...
|
||||
}
|
||||
model = AutoModelForCausalLM.from_pretrained("facebook/opt-13b", torch_dtype=torch.float16)
|
||||
ds_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)
|
||||
lean_state_dict = deepspeed.checkpoint.utils.clone_tensors_for_torch_save(ds_engine.module.state_dict())
|
||||
ds_engine.module.save_pretrained("lean_after", state_dict=lean_state_dict)
|
||||
|
||||
|
||||
|
||||
Universal Checkpoints (under development)
|
||||
------------------------------------------
|
||||
Parallelism techniques such as ZeRO data parallelism (DP), Tensor parallelism (TP), Pipeline parallelism (TP), which shard model and/or
|
||||
optimizer states make it difficult to resume training with a checkpoint that was created on a different number of GPUs. DeepSpeed provides the
|
||||
Universal Checkpoint mechanism to address this problem. Universal Checkpoints give users the flexibility of changing the number of GPUs when training
|
||||
with 3D (TP, PP, and DP) parallelism, and enables more efficient use of elastic training hardware. The easiest way to get started with
|
||||
using Universal Checkpoints is to consult the `Megatron-DeepSpeed <https://github.com/deepspeedai/Megatron-DeepSpeed/blob/main/examples_deepspeed/universal_checkpointing/README.md>`_
|
||||
and `BLOOM <https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/README.md#checkpoint-reshaping>`_ examples.
|
||||
@@ -0,0 +1,14 @@
|
||||
Mixture of Experts (DeepSpeed MoE)
|
||||
==================================
|
||||
|
||||
DeepSpeed provides two forms of MoE support: DeepSpeed MoE and :doc:`AutoEP
|
||||
(Automatic Expert Parallelism) <autoep>`. DeepSpeed MoE is the explicit
|
||||
``deepspeed.moe.layer.MoE`` API for constructing MoE layers in model code. This
|
||||
page introduces the DeepSpeed MoE API.
|
||||
|
||||
See also the `Mixture of Experts (DeepSpeed MoE) tutorial
|
||||
<https://www.deepspeed.ai/tutorials/mixture-of-experts/>`__ for training
|
||||
examples and configuration details.
|
||||
|
||||
.. autoclass:: deepspeed.moe.layer.MoE
|
||||
:members:
|
||||
@@ -0,0 +1,40 @@
|
||||
Monitoring
|
||||
==========
|
||||
|
||||
Deepspeed’s Monitor module can log training details into a
|
||||
Tensorboard-compatible file, to WandB, or to simple CSV files. Below is an
|
||||
overview of what DeepSpeed will log automatically.
|
||||
|
||||
.. csv-table:: Automatically Logged Data
|
||||
:header: "Field", "Description", "Condition"
|
||||
:widths: 20, 20, 10
|
||||
|
||||
`Train/Samples/train_loss`,"The training loss.",None
|
||||
`Train/Samples/lr`,"The learning rate during training.",None
|
||||
`Train/Samples/loss_scale`,"The loss scale when training using `fp16`.",`fp16` must be enabled.
|
||||
`Train/Eigenvalues/ModelBlockParam_{i}`,"Eigen values per param block.",`eigenvalue` must be enabled.
|
||||
`Train/Samples/elapsed_time_ms_forward`,"The global duration of the forward pass.",`flops_profiler.enabled` or `wall_clock_breakdown`.
|
||||
`Train/Samples/elapsed_time_ms_backward`,"The global duration of the forward pass.",`flops_profiler.enabled` or `wall_clock_breakdown`.
|
||||
`Train/Samples/elapsed_time_ms_backward_inner`,"The backward time that does not include the gradient reduction time. Only in cases where the gradient reduction is not overlapped, if it is overlapped then the inner time should be about the same as the entire backward time.",`flops_profiler.enabled` or `wall_clock_breakdown`.
|
||||
`Train/Samples/elapsed_time_ms_backward_allreduce`,"The global duration of the allreduce operation.",`flops_profiler.enabled` or `wall_clock_breakdown`.
|
||||
`Train/Samples/elapsed_time_ms_step`,"The optimizer step time.",`flops_profiler.enabled` or `wall_clock_breakdown`.
|
||||
|
||||
TensorBoard
|
||||
-----------
|
||||
.. _TensorBoardConfig:
|
||||
.. autopydantic_model:: deepspeed.monitor.config.TensorBoardConfig
|
||||
|
||||
WandB
|
||||
-----
|
||||
.. _WandbConfig:
|
||||
.. autopydantic_model:: deepspeed.monitor.config.WandbConfig
|
||||
|
||||
Comet
|
||||
-----
|
||||
.. _CometConfig:
|
||||
.. autopydantic_model:: deepspeed.monitor.config.CometConfig
|
||||
|
||||
CSV Monitor
|
||||
-----------
|
||||
.. _CSVConfig:
|
||||
.. autopydantic_model:: deepspeed.monitor.config.CSVConfig
|
||||
Executable
+28
@@ -0,0 +1,28 @@
|
||||
Optimizers
|
||||
===================
|
||||
|
||||
DeepSpeed offers high-performance implementations of ``Adam`` optimizer on CPU; ``FusedAdam``, ``FusedLamb``, ``OnebitAdam``, ``OnebitLamb`` optimizers on GPU.
|
||||
|
||||
Adam (CPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.ops.adam.DeepSpeedCPUAdam
|
||||
|
||||
FusedAdam (GPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.ops.adam.FusedAdam
|
||||
|
||||
FusedLamb (GPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.ops.lamb.FusedLamb
|
||||
|
||||
OneBitAdam (GPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.runtime.fp16.onebit.adam.OnebitAdam
|
||||
|
||||
ZeroOneAdam (GPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.runtime.fp16.onebit.zoadam.ZeroOneAdam
|
||||
|
||||
OnebitLamb (GPU)
|
||||
----------------------------
|
||||
.. autoclass:: deepspeed.runtime.fp16.onebit.lamb.OnebitLamb
|
||||
@@ -0,0 +1,26 @@
|
||||
Pipeline Parallelism
|
||||
====================
|
||||
|
||||
Model Specification
|
||||
--------------------
|
||||
.. autoclass:: deepspeed.pipe.PipelineModule
|
||||
:members:
|
||||
|
||||
.. autoclass:: deepspeed.pipe.LayerSpec
|
||||
:members:
|
||||
|
||||
.. autoclass:: deepspeed.pipe.TiedLayerSpec
|
||||
:members:
|
||||
|
||||
.. autoclass:: deepspeed.runtime.pipe.ProcessTopology
|
||||
:members:
|
||||
|
||||
Training
|
||||
--------
|
||||
.. automodule:: deepspeed.runtime.pipe.engine
|
||||
:members:
|
||||
|
||||
Extending Pipeline Parallelism
|
||||
------------------------------
|
||||
.. automodule:: deepspeed.runtime.pipe.schedule
|
||||
:members:
|
||||
Executable
+30
@@ -0,0 +1,30 @@
|
||||
Learning Rate Schedulers
|
||||
=================================
|
||||
|
||||
DeepSpeed offers implementations of ``LRRangeTest``, ``OneCycle``, ``WarmupLR``, ``WarmupDecayLR``, ``WarmupCosineLR`` learning rate schedulers. When using a DeepSpeed's learning rate scheduler (specified in the `ds_config.json` file), DeepSpeed calls the `step()` method of the scheduler at every training step (when `model_engine.step()` is executed). When not using a DeepSpeed's learning rate scheduler:
|
||||
* if the schedule is supposed to execute at every training step, then the user can pass the scheduler to `deepspeed.initialize` when initializing the DeepSpeed engine and let DeepSpeed manage it for update or save/restore.
|
||||
* if the schedule is supposed to execute at any other interval (e.g., training epochs), then the user should NOT pass the scheduler to DeepSpeed during initialization and must manage it explicitly.
|
||||
|
||||
LRRangeTest
|
||||
---------------------------
|
||||
.. autoclass:: deepspeed.runtime.lr_schedules.LRRangeTest
|
||||
|
||||
|
||||
OneCycle
|
||||
---------------------------
|
||||
.. autoclass:: deepspeed.runtime.lr_schedules.OneCycle
|
||||
|
||||
|
||||
WarmupLR
|
||||
---------------------------
|
||||
.. autoclass:: deepspeed.runtime.lr_schedules.WarmupLR
|
||||
|
||||
|
||||
WarmupDecayLR
|
||||
---------------------------
|
||||
.. autoclass:: deepspeed.runtime.lr_schedules.WarmupDecayLR
|
||||
|
||||
|
||||
WarmupCosineLR
|
||||
---------------------------
|
||||
.. autoclass:: deepspeed.runtime.lr_schedules.WarmupCosineLR
|
||||
@@ -0,0 +1,683 @@
|
||||
Training API
|
||||
############
|
||||
|
||||
:func:`deepspeed.initialize` returns a *training engine* in its first argument
|
||||
of type :class:`DeepSpeedEngine`. This engine is used to progress training:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for step, batch in enumerate(data_loader):
|
||||
#forward() method
|
||||
loss = model_engine(batch)
|
||||
|
||||
#runs backpropagation
|
||||
model_engine.backward(loss)
|
||||
|
||||
#weight update
|
||||
model_engine.step()
|
||||
|
||||
Note that ``model_engine.backward()`` accepts only a scalar loss tensor produced by a forward pass.
|
||||
Starting from v0.18.3, DeepSpeed also supports direct calls to ``tensor.backward()``. You can now call
|
||||
``loss.backward()`` or ``tensor.backward(out_grad)`` when your PyTorch version supports the necessary APIs.
|
||||
If your PyTorch version does not support these APIs, a direct call to ``tensor.backward()`` will raise an error.
|
||||
|
||||
Forward Propagation
|
||||
-------------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.forward
|
||||
|
||||
Backward Propagation
|
||||
--------------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.backward
|
||||
|
||||
Loss Scaling for Manual Backward Passes
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.scale
|
||||
|
||||
When using mixed precision training (fp16, bf16, or torch.autocast), DeepSpeed applies loss scaling
|
||||
to prevent gradient underflow. If you prefer to call ``loss.backward()`` directly instead of
|
||||
``engine.backward(loss)``, you must use ``engine.scale(loss)`` to apply the appropriate loss scaler:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Option 1: Use engine.backward() (recommended)
|
||||
loss = model_engine(batch)
|
||||
model_engine.backward(loss)
|
||||
|
||||
# Option 2: Manual backward with scaling
|
||||
loss = model_engine(batch)
|
||||
scaled_loss = model_engine.scale(loss)
|
||||
scaled_loss.backward()
|
||||
|
||||
Both approaches produce identical gradients. The ``scale()`` method automatically applies the
|
||||
appropriate scaler based on your configuration (ZeRO optimizer scaler, torch.autocast GradScaler, etc.).
|
||||
|
||||
Optimizer Step
|
||||
--------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.step
|
||||
|
||||
Gradient Accumulation
|
||||
---------------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.is_gradient_accumulation_boundary
|
||||
|
||||
Coalesced Gradient Reduction
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.coalesce_grad_reduction
|
||||
|
||||
Use this when one optimizer step needs multiple ``engine.backward()`` calls
|
||||
and per-backward reduction is wasted work. Typical cases are GradCache-style
|
||||
cached contrastive losses that replay backward over chunked representations,
|
||||
and custom ``torch.autograd.Function`` subclasses that call
|
||||
``torch.autograd.backward`` from inside their ``forward``. Results are
|
||||
bit-exact against the per-backward baseline.
|
||||
|
||||
Under ZeRO-3, each backward inside the block leaves param-shaped gradients
|
||||
on the leaf modules instead of triggering the per-backward reduce-scatter.
|
||||
On exit, a single pass drives the reducer over the accumulated grads and
|
||||
restores the partitioned ``averaged_gradients`` for ``step()``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
for batch in data_loader:
|
||||
chunks = batch.split(chunk_size)
|
||||
with model_engine.coalesce_grad_reduction():
|
||||
for chunk in chunks:
|
||||
loss = model_engine(chunk)
|
||||
model_engine.backward(loss)
|
||||
model_engine.step()
|
||||
|
||||
Communication
|
||||
^^^^^^^^^^^^^
|
||||
|
||||
With ``N`` back-to-back ``backward()`` calls per step, ZeRO-2 and ZeRO-3
|
||||
normally issue ``N`` gradient collectives (one per backward). Inside
|
||||
``coalesce_grad_reduction()`` those collapse to one collective on exit.
|
||||
ZeRO-1 already reduces only at the accumulation boundary, so its collective
|
||||
count is unchanged; the context still removes the per-backward bucket setup
|
||||
cost.
|
||||
|
||||
Memory
|
||||
^^^^^^
|
||||
|
||||
Suppressing the per-backward reduction means each rank holds a full local
|
||||
gradient copy for the duration of the ``with`` block.
|
||||
|
||||
* ZeRO-2: window-resident memory equals ZeRO-1 with
|
||||
:meth:`deepspeed.DeepSpeedEngine.no_sync`, one full gradient per rank
|
||||
held until flush. On a 2-GPU, 134M-param bf16 rig with ``N=4``, peak
|
||||
window memory drops from 640 MiB (baseline) to 384 MiB.
|
||||
* ZeRO-3: window-resident is one full gradient per rank vs the
|
||||
``1/world_size`` partition the per-backward path holds throughout. Peak
|
||||
is roughly equal to baseline (the in-flight backward already needs
|
||||
full-grad room and the accumulator reuses it).
|
||||
|
||||
Constraints
|
||||
^^^^^^^^^^^
|
||||
|
||||
* ZeRO stage 0 and pipeline parallelism raise ``NotImplementedError``.
|
||||
* The BF16/FP16 optimizer wrappers (``BF16_Optimizer``, ``FP16_Optimizer``)
|
||||
route grads through their own ``backward_epilogue`` path and are not yet
|
||||
supported; the context raises ``NotImplementedError`` at entry. Use raw
|
||||
ZeRO-1/2/3 for now.
|
||||
* ``engine.step()`` inside the ``with`` block raises.
|
||||
* Cannot be nested inside :meth:`deepspeed.DeepSpeedEngine.no_sync`.
|
||||
* Do not split one ``gradient_accumulation_steps`` window across multiple
|
||||
``with`` blocks: the flush overwrites ``averaged_gradients`` on each exit.
|
||||
|
||||
:meth:`deepspeed.DeepSpeedEngine.no_sync` raises ``AssertionError`` for
|
||||
ZeRO-2 and ZeRO-3 (``zero_optimization_partition_gradients()`` is true for
|
||||
stage >= 2), so it cannot collapse collectives for those stages.
|
||||
``coalesce_grad_reduction()`` is the equivalent for ZeRO-2/3.
|
||||
|
||||
|
||||
Mixed Precision Training
|
||||
-------------------------
|
||||
DeepSpeed supports mixed precision training using either native or PyTorch mechanisms. The desired mixed precision mode can be selected through the configuration dict.
|
||||
Mixed precision training can used with ZeRO (i.e., stages > 0) and without ZeRO (i.e., stage=0).
|
||||
|
||||
|
||||
Native Mixed Precision
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
DeepSpeed provides native support for
|
||||
`fp16 <https://www.deepspeed.ai/docs/config-json/#fp16-training-options>`_ and `bf16 <https://www.deepspeed.ai/docs/config-json/#bfloat16-training-options>`_ mixed precsion training.
|
||||
|
||||
|
||||
PyTorch Automatic Mixed Precision (AMP)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
DeepSpeed provides torch-compatible automatic mixed precision (AMP) training via
|
||||
`torch.autocast <https://docs.pytorch.org/docs/stable/amp.html>`_ functionality. The following snippet illustrates how to enable Torch AMP.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
{
|
||||
"torch_autocast": {
|
||||
"enabled": true,
|
||||
"dtype": "bfloat16",
|
||||
"lower_precision_safe_modules": ["torch.nn.Linear", "torch.nn.Conv2d"]
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
Each configuration works as follows:
|
||||
|
||||
* ``enabled``: Enable ``torch.autocast`` when set to ``True``. You don't need to call ``torch.autocast`` in your code. The grad scaler is also applied in the DeepSpeed optimizer.
|
||||
* ``dtype``: Lower precision dtype passed to ``torch.autocast``. Gradients for all-reduce (reduce-scatter) and parameters for all-gather (only for ZeRO3) of ``lower_precision_safe_modules`` are also downcasted to this ``dtype``.
|
||||
* ``lower_precision_safe_modules``: The list of modules that will be downcasted for all-reduce (reduce-scatter) and all-gather (ZeRO3). The precision for PyTorch operators in forward/backward follows ``torch.autocast``'s policy, not this list. If you don't set this item, DeepSpeed uses the default list: ``[torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d]``.
|
||||
|
||||
Manual Backward with torch.autocast
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
When using ``torch.autocast`` with manual backward passes (``loss.backward()`` instead of ``engine.backward()``),
|
||||
you must use ``engine.scale(loss)`` to apply the gradient scaler:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Training loop with torch.autocast and manual backward
|
||||
for batch in data_loader:
|
||||
loss = model_engine(batch)
|
||||
|
||||
# Apply loss scaling before manual backward
|
||||
scaled_loss = model_engine.scale(loss)
|
||||
scaled_loss.backward()
|
||||
|
||||
model_engine.step()
|
||||
|
||||
The ``scale()`` method ensures that the ``torch.amp.GradScaler`` is properly applied when ``torch.autocast``
|
||||
is enabled with fp16. For bf16 or when no mixed precision is used, ``scale()`` returns the loss unchanged.
|
||||
|
||||
If you call ``loss.backward()`` directly without using ``engine.scale()`` or ``engine.backward()``, DeepSpeed
|
||||
will raise a ``RuntimeError`` to prevent training with unscaled gradients, which can lead to incorrect results
|
||||
or gradient underflow.
|
||||
|
||||
Using torch.autocast Outside the Engine
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
DeepSpeed applies ``torch.autocast`` internally during ``engine.forward()``.
|
||||
However, you may also want autocast to cover code that runs **outside** the engine,
|
||||
such as a loss function or post-processing logic. In that case, wrap the entire
|
||||
forward-plus-loss block in your own ``torch.autocast`` context:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Autocast covers both the engine forward AND the loss computation
|
||||
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
||||
logits = model_engine(input_ids)
|
||||
loss = loss_fn(logits.view(-1, vocab_size), labels.view(-1))
|
||||
|
||||
Without the outer ``torch.autocast``, only the model's forward pass benefits from
|
||||
autocast; the loss function would run in full precision.
|
||||
|
||||
When DeepSpeed detects a nested autocast context, it handles it as follows:
|
||||
|
||||
* If ``torch_autocast`` is **enabled** in the DeepSpeed config, the engine overrides the
|
||||
outer context with the dtype from the config. An info message is logged once.
|
||||
* If ``torch_autocast`` is **disabled** in the config (i.e., you are using DeepSpeed's
|
||||
built-in bf16/fp16 support instead), the engine disables autocast inside
|
||||
``engine.forward()`` and a warning is logged once.
|
||||
|
||||
In both cases, PyTorch's ``torch.autocast`` is idempotent when nested with the same
|
||||
dtype, so there is no performance or correctness penalty from the nesting.
|
||||
|
||||
.. autofunction:: deepspeed.runtime.torch_autocast.init_autocast_params
|
||||
.. autofunction:: deepspeed.runtime.torch_autocast.is_autocast_initialized
|
||||
.. autofunction:: deepspeed.runtime.torch_autocast.get_default_autocast_lower_precision_modules
|
||||
.. autofunction:: deepspeed.runtime.torch_autocast.has_autocast_dtype
|
||||
|
||||
|
||||
Configuring ZeRO Leaf Modules
|
||||
-----------------------------
|
||||
|
||||
ZeRO-3 relies on module execution order to gather partitioned parameters.
|
||||
When models select submodules dynamically (for example, MoE routers), different data-parallel ranks may gather different sets of parameters, which can cause the all-gather collective to deadlock.
|
||||
To avoid this problem, you can designate the parent of dynamically activated submodules (e.g., MoE experts) as a "leaf" module.
|
||||
When a module is marked as a leaf, ZeRO gathers all of its descendants immediately and stops inserting hooks beneath it.
|
||||
|
||||
Programmatic API
|
||||
~~~~~~~~~~~~~~~~
|
||||
|
||||
Use :func:`deepspeed.utils.set_z3_leaf_modules` to flag modules by class, class
|
||||
name, or both. Optionally combine with
|
||||
:func:`deepspeed.utils.set_z3_leaf_modules_by_name` to target specific entries
|
||||
from ``model.named_modules()`` or
|
||||
:func:`deepspeed.utils.set_z3_leaf_modules_by_suffix` to match suffixes of those
|
||||
names.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from deepspeed.utils import (
|
||||
set_z3_leaf_modules,
|
||||
set_z3_leaf_modules_by_name,
|
||||
set_z3_leaf_modules_by_suffix,
|
||||
)
|
||||
|
||||
# Match by class or subclass
|
||||
set_z3_leaf_modules(model, [CustomMoEBlock])
|
||||
|
||||
# Match by fully qualified class name
|
||||
set_z3_leaf_modules(model, ["my_package.layers.CustomMoEBlock"])
|
||||
|
||||
# Match by module name returned from model.named_modules()
|
||||
set_z3_leaf_modules_by_name(model, ["transformer.layers.0.experts"])
|
||||
|
||||
# Match by suffix of names returned from model.named_modules()
|
||||
set_z3_leaf_modules_by_suffix(model, ["experts"])
|
||||
|
||||
Configuration in DeepSpeed config
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The same behavior can be controlled from the DeepSpeed config. Add a
|
||||
``leaf_module`` block to ``zero_optimization`` specifying either classes,
|
||||
module names, or name suffixes (or any combination). While the example below shows three different ways (``classes``, ``names``, and ``name_suffixes``) to specify modules as leaf modules, typically you will use just one of these.
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"leaf_module": {
|
||||
"classes": ["my_package.layers.CustomMoEBlock"],
|
||||
"names": ["transformer.layers.0.experts"],
|
||||
"name_suffixes": ["experts"]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
``names`` must match exactly what ``model.named_modules()`` produces. The
|
||||
``name_suffixes`` field compares each suffix against the end of those same
|
||||
module paths, making it convenient to apply a rule across repeated structures.
|
||||
Entries in ``classes`` may be either bare class names (for example,
|
||||
``MixtralSparseMoeBlock``) or fully qualified dotted paths; both forms are
|
||||
accepted.
|
||||
|
||||
You can mix and match the API and configuration approaches; all referenced
|
||||
modules are flagged before ZeRO installs its hooks.
|
||||
|
||||
By default DeepSpeed marks several Hugging Face MoE blocks—including Mixtral and Qwen MoE sparse blocks so that they behave well with ZeRO3. The default class list currently contains:
|
||||
|
||||
* ``transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock``
|
||||
* ``transformers.models.qwen2_moe.modeling_qwen2_moe.Qwen2MoeSparseMoeBlock``
|
||||
* ``transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeSparseMoeBlock``
|
||||
|
||||
|
||||
Model Saving
|
||||
------------
|
||||
.. autofunction:: deepspeed.DeepSpeedEngine.save_16bit_model
|
||||
|
||||
|
||||
Additionally when a DeepSpeed checkpoint is created, a script ``zero_to_fp32.py`` is added there which can be used to reconstruct fp32 master weights into a single pytorch ``state_dict`` file.
|
||||
|
||||
|
||||
Training Multiple Models
|
||||
------------------------
|
||||
DeepSpeed supports training multiple models, which is a useful feature in `scenarios <https://huggingface.co/docs/accelerate/en/usage_guides/deepspeed_multiple_model>`_ such as knowledge distillation and post-training RLHF.
|
||||
The core approach is to create individual DeepSpeedEngines for each model.
|
||||
|
||||
|
||||
Training Independent Models
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The following code snippet illustrates independently training multiple models on the same dataset.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model_engines = [engine for engine, _, _, _ in [deepspeed.initialize(m, ...,) for m in models]]
|
||||
for batch in data_loader:
|
||||
losses = [engine(batch) for engine in model_engines]
|
||||
for engine, loss in zip(model_engines, losses):
|
||||
engine.backward(loss)
|
||||
|
||||
The above is similar to typical DeepSpeed usage except for the creation of multiple DeepSpeedEngines (one for each model).
|
||||
|
||||
|
||||
Jointly Training Models With Shared Loss
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
The following code snippet illustrates jointly training multiple models on a shared loss value.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model_engines = [engine for engine, _, _, _ in [deepspeed.initialize(m, ...,) for m in models]]
|
||||
for batch in data_loader:
|
||||
losses = [engine(batch[0], batch[1]) for engine in model_engines]
|
||||
loss = sum(l / (i + 1) for i, l in enumerate(losses))
|
||||
loss.backward()
|
||||
|
||||
for engine in model_engines:
|
||||
engine.step()
|
||||
|
||||
for engine in model_engines:
|
||||
engine.optimizer.zero_grad()
|
||||
|
||||
Besides the use of multiple DeepSpeedEngines, the above differs from typical usage in two key ways:
|
||||
|
||||
#. The **backward** call is made using the common loss value rather on individual model engines.
|
||||
|
||||
You can call ``loss.backward()`` once for the shared loss.
|
||||
|
||||
**Note:** Previously, you had to call ``_backward_epilogue`` on each model engine after ``loss.backward()``. However, starting from v0.18.3, DeepSpeed automatically handles this internally, so you no longer need to call ``_backward_epilogue`` manually.
|
||||
|
||||
|
||||
Automatic Tensor Parallel Training
|
||||
----------------------------------
|
||||
DeepSpeed supports **Automatic Tensor Parallel (AutoTP) training** for sharding
|
||||
model weights across GPUs while remaining compatible with ZeRO and standard
|
||||
training workflows. This training API is different from the inference-only
|
||||
tensor parallel API exposed by ``deepspeed.init_inference``.
|
||||
|
||||
Tensor parallelism (TP) splits the computations and parameters of large layers
|
||||
across multiple GPUs so each rank holds only a shard of the weight matrix. This
|
||||
is an efficient way to train large-scale transformer models by reducing per-GPU
|
||||
memory pressure while keeping the layer math distributed across the TP group.
|
||||
|
||||
AutoTP training is enabled by setting ``tensor_parallel`` in the DeepSpeed
|
||||
config and passing it to ``deepspeed.initialize``. DeepSpeed applies AutoTP
|
||||
sharding during engine initialization; calling ``deepspeed.tp_model_init``, which we previously used to initialize AutoTP, is now optional.
|
||||
See :ref:`autotp-training-init-details` for more details.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import deepspeed
|
||||
|
||||
ds_config = {
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"zero_optimization": {"stage": 2},
|
||||
"tensor_parallel": {"autotp_size": 4},
|
||||
}
|
||||
|
||||
engine, optimizer, _, _ = deepspeed.initialize(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
config=ds_config,
|
||||
mpu=mpu, # optional: TP/DP process groups
|
||||
)
|
||||
|
||||
.. note::
|
||||
AutoTP training supports ZeRO stages 0, 1, and 2. ZeRO Stage 3 is not supported.
|
||||
|
||||
.. _autotp-training-init-details:
|
||||
|
||||
Initialization behavior
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
AutoTP previously required calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` before ``deepspeed.initialize``. Now we can include all the necessary configurations in the DeepSpeed config.
|
||||
We still support the traditional initialization path for backward compatibility.
|
||||
When you use both (i.e. calling ``set_autotp_mode(training=True)`` and ``deepspeed.tp_model_init`` and passing the config to ``deepspeed.initialize``), we will merge the settings at initialization. When we have conflicting settings, we will error out.
|
||||
|
||||
Parameter partitioning
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
TP sharding needs to know which parameter tensors should be partitioned and
|
||||
along which dimensions. AutoTP provides three ways to balance ready-to-use
|
||||
defaults with customizability:
|
||||
|
||||
* **Heuristics**: automatic sharding based on parameter names and model rules.
|
||||
* **Preset**: choose a built-in model family via ``preset_model``.
|
||||
* **Custom specs**: define regex patterns and partition rules via ``partition_config``.
|
||||
* **HuggingFace tp_plan**: automatically detected from ``model.config.base_model_tp_plan`` or ``model._tp_plan``.
|
||||
|
||||
HuggingFace tp_plan
|
||||
^^^^^^^^^^^^^^^^^^^
|
||||
Many HuggingFace models (e.g. Llama, Qwen, Gemma2) define a
|
||||
``base_model_tp_plan`` in their model config. When present, DeepSpeed
|
||||
automatically extracts and converts this plan into internal partition rules.
|
||||
This means you do not need ``preset_model`` or ``partition_config`` for these
|
||||
models -- just set ``autotp_size``.
|
||||
|
||||
The resolution priority is:
|
||||
|
||||
1. ``partition_config`` (user-defined custom specs -- highest priority)
|
||||
2. HuggingFace ``tp_plan`` (from model config)
|
||||
3. AutoTP heuristics / ``preset_model`` (lowest priority)
|
||||
|
||||
Currently only ``colwise`` and ``rowwise`` partition types from the HuggingFace
|
||||
``tp_plan`` are supported. Other types (``colwise_rep``, ``local_colwise``,
|
||||
``local_rowwise``, ``local_packed_rowwise``, ``gather``, ``sequence_parallel``)
|
||||
are not yet handled and will raise an error.
|
||||
|
||||
Heuristic rules
|
||||
^^^^^^^^^^^^^^^
|
||||
Heuristics use parameter names and model-specific rules to decide how to shard
|
||||
layers. If you are training a supported model (see
|
||||
:ref:`autotp-supported-models`), the heuristic rules automatically shard the
|
||||
model, so you only need to add ``autotp_size``.
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
...
|
||||
"tensor_parallel": {
|
||||
"autotp_size": 4
|
||||
},
|
||||
"zero_optimization": {
|
||||
...
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
Preset-based partitioning
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
You can explicitly specify the model family with ``preset_model``:
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"tensor_parallel": {
|
||||
"autotp_size": 4,
|
||||
"preset_model": "llama"
|
||||
}
|
||||
}
|
||||
|
||||
See :ref:`autotp-supported-models` for the supported preset names and the
|
||||
implementation in `AutoTPPresets <https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/autotp_config.py>`_.
|
||||
If you add a new model family, you can easily add a new preset by defining
|
||||
patterns like the existing presets, and we welcome PRs for those additions.
|
||||
|
||||
Custom layer specs
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
If you are training a custom model, you can use ``partition_config`` to specify
|
||||
custom regex-based patterns and partition settings.
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"tensor_parallel": {
|
||||
"autotp_size": 4,
|
||||
"partition_config": {
|
||||
"use_default_specs": false,
|
||||
"layer_specs": [
|
||||
{
|
||||
"patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"],
|
||||
"partition_type": "row"
|
||||
},
|
||||
{
|
||||
"patterns": [".*\\.[qkv]_proj\\.weight$"],
|
||||
"partition_type": "column"
|
||||
},
|
||||
{
|
||||
"patterns": [".*\\.gate_up_proj\\.weight$"],
|
||||
"partition_type": "column",
|
||||
"shape": [2, -1],
|
||||
"partition_dim": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
You can also set ``use_default_specs`` to ``true`` to merge your custom
|
||||
patterns on top of the preset (when ``preset_model`` is provided).
|
||||
|
||||
For fused or packed weights (for example QKV or gate/up projections), the
|
||||
``shape`` and ``partition_dim`` options control sub-parameter partitioning.
|
||||
Sub-parameter partitioning lets AutoTP split a single weight tensor into
|
||||
logical chunks before applying tensor-parallel sharding. For example, the
|
||||
``gate_up_proj`` weight can be viewed as two packed matrices (gate and up) by
|
||||
setting ``shape`` to ``[2, -1]`` and ``partition_dim`` to ``0``; AutoTP then
|
||||
partitions each chunk consistently across tensor-parallel ranks.
|
||||
|
||||
.. image:: /_static/autotp-subparams-gate-up.png
|
||||
:alt: AutoTP sub-parameter partitioning
|
||||
|
||||
Another example is GQA-style fused QKV weights. The tensor can contain unequal
|
||||
Q/K/V segments stacked along the output dimension. For example, set ``shape``
|
||||
to the explicit sizes (for example ``[(q_size, kv_size, kv_size), -1]``) and
|
||||
``partition_dim`` to ``0`` so AutoTP splits the Q, K, and V regions first, then
|
||||
shards each region across tensor-parallel ranks.
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"patterns": [".*\\.qkv_proj\\.weight$"],
|
||||
"partition_type": "column",
|
||||
"shape": [[q_size, kv_size, kv_size], -1],
|
||||
"partition_dim": 0
|
||||
}
|
||||
|
||||
.. image:: /_static/autotp-subparams-gqa.png
|
||||
:alt: AutoTP sub-parameter partitioning
|
||||
|
||||
|
||||
Model-type filtering for shared configs
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
Use ``model_types`` when you want a single config to work across multiple model
|
||||
families but apply different specs. This is useful in shared training scripts
|
||||
or when patterns overlap across architectures.
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"tensor_parallel": {
|
||||
"autotp_size": 4,
|
||||
"partition_config": {
|
||||
"layer_specs": [
|
||||
{
|
||||
"patterns": [".*\\.qkv_proj\\.weight$"],
|
||||
"partition_type": "column",
|
||||
"shape": [[q_size, kv_size, kv_size], -1],
|
||||
"partition_dim": 0,
|
||||
"model_types": ["llama"]
|
||||
},
|
||||
{
|
||||
"patterns": [".*\\.qkv_proj\\.weight$"],
|
||||
"partition_type": "column",
|
||||
"shape": [3, -1],
|
||||
"partition_dim": 0,
|
||||
"model_types": ["qwen2"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
.. _autotp-supported-models:
|
||||
|
||||
Supported models
|
||||
~~~~~~~~~~~~~~~~
|
||||
The following model families are supported by built-in AutoTP presets:
|
||||
|
||||
- ``llama``
|
||||
- ``bloom``
|
||||
- ``chatglm``
|
||||
- ``mixtral``
|
||||
- ``deepseek_v2``
|
||||
- ``qwen2``
|
||||
- ``phi3``
|
||||
|
||||
Preset definitions live in `AutoTPPresets <https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/autotp_config.py>`_.
|
||||
If you add a new model family, you can easily add a new preset by defining
|
||||
patterns like the existing presets, and we welcome PRs for those additions.
|
||||
|
||||
These strings are the values accepted by ``preset_model`` and are matched
|
||||
against the model type in ``model.config.model_type`` (case-insensitive). When
|
||||
``preset_model`` is not set, AutoTP uses the legacy automatic sharding rules
|
||||
unless you provide a custom ``partition_config``.
|
||||
These presets are also useful when you want to extend the default patterns:
|
||||
set ``use_default_specs`` to ``true`` in ``partition_config`` to merge your custom
|
||||
specs on top of the selected preset.
|
||||
|
||||
|
||||
Automatic Sequence Parallel Training
|
||||
------------------------------------
|
||||
DeepSpeed supports **Automatic Sequence Parallel (AutoSP) training** for enabling
|
||||
compiler-based sequence parallelism to unlock long-context LLM training. AutoSP
|
||||
leverages defines custom passes to automatically shard inputs along the
|
||||
sequence dimension and enable Ulysses-styled sequence parallelism.
|
||||
|
||||
AutoSP training is enabled by setting ``compile`` and ``passes`` in the DeepSpeed
|
||||
config and calling ``prepare_autosp_inputs()`` to prepare inputs before each forward pass.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import deepspeed
|
||||
from deepspeed.compile.passes.sp_compile import prepare_autosp_inputs
|
||||
|
||||
ds_config = {
|
||||
"train_micro_batch_size_per_gpu": 1,
|
||||
"zero_optimization": {"stage": 0},
|
||||
"compile": {
|
||||
"deepcompile": True,
|
||||
"passes": ["autosp"],
|
||||
}
|
||||
}
|
||||
|
||||
engine, optimizer, _, _ = deepspeed.initialize(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
config=ds_config,
|
||||
)
|
||||
|
||||
# Compile the model before training
|
||||
engine.compile(backend='inductor')
|
||||
|
||||
for batch in dataloader:
|
||||
input_ids = prepare_autosp_inputs(
|
||||
input_id=batch["input_ids"],
|
||||
label_id=batch["labels"],
|
||||
position_id=batch.get("position_ids"),
|
||||
seq_dim=1
|
||||
)
|
||||
loss = engine(input_ids)
|
||||
engine.backward(loss)
|
||||
engine.step()
|
||||
|
||||
.. note::
|
||||
AutoSP requires ZeRO stage 0 (no ZeRO optimization). Using AutoSP with ZeRO stages 1, 2, or 3 is not currently supported.
|
||||
AutoSP also requires ``torch.nn.functional.scaled_dot_product_attention()`` as the attention backend.
|
||||
|
||||
Input Preparation
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
Before each forward pass, inputs must be prepared using ``prepare_autosp_inputs()`` to
|
||||
mark the sequence dimension as dynamic and annotate tensors for identification during
|
||||
automatic sharding:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from deepspeed.compile.passes.sp_compile import prepare_autosp_inputs
|
||||
|
||||
input_ids = prepare_autosp_inputs(
|
||||
input_id=input_ids,
|
||||
label_id=labels,
|
||||
position_id=position_ids, # optional
|
||||
attention_mask=attention_mask, # optional
|
||||
seq_dim=1
|
||||
)
|
||||
|
||||
This serves as a hint to the compiler to know which inputs should be sharded across which dimension.
|
||||
|
||||
Memory Optimization
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
AutoSP includes selective activation checkpointing that recomputes matmul operations
|
||||
during backpropagation while preserving attention activations. This is effective for
|
||||
long-context training because attention operations scale quadratically with sequence
|
||||
length and dominate computation latency, while matmul operations scale linearly and are relatively cheaper
|
||||
to recompute. This provides significant memory savings with minimal computational
|
||||
overhead
|
||||
|
||||
Limitations
|
||||
~~~~~~~~~~~
|
||||
|
||||
AutoSP currently supports only ``torch.nn.functional.scaled_dot_product_attention``. Other attention patterns require additional pattern matching logic.
|
||||
|
||||
AutoSP requires a fully connected computation graph without breaks. Graph breaks destroy the use-def chains across graphs and the compiler cannot propoaget sequence dimension sharding information.
|
||||
@@ -0,0 +1,564 @@
|
||||
ZeRO
|
||||
####
|
||||
|
||||
The Zero Redundancy Optimizer (ZeRO) removes the memory redundancies across
|
||||
data-parallel processes by partitioning the three model states (optimizer
|
||||
states, gradients, and parameters) across data-parallel processes instead of
|
||||
replicating them. By doing this, it boosts memory efficiency compared to
|
||||
classic data-parallelism while retaining its computational granularity and
|
||||
communication efficiency.
|
||||
|
||||
#. **ZeRO Stage 1**: The optimizer states (e.g., for `Adam optimizer <https://arxiv.org/abs/1412.6980>`_, 32-bit weights, and the first, and second moment estimates) are partitioned across the processes, so that each process updates only its partition.
|
||||
|
||||
#. **ZeRO Stage 2**: The reduced 16-bit gradients for updating the model weights are also partitioned such that each process retains only the gradients corresponding to its portion of the optimizer states.
|
||||
|
||||
#. **ZeRO Stage 3**: The 16-bit model parameters are partitioned across the processes. ZeRO-3 will automatically collect and partition them during the forward and backward passes.
|
||||
|
||||
In addition, ZeRO-3 includes the *infinity offload engine* to form
|
||||
ZeRO-Infinity ([paper](https://arxiv.org/abs/2104.07857)), which can offload
|
||||
all model states to both CPU and NVMe memory for huge memory savings.
|
||||
|
||||
|
||||
For a deep dive of our algorithms, please see our `papers <https://www.deepspeed.ai/#publications>`_ on `ZeRO
|
||||
<https://arxiv.org/abs/1910.02054>`_, `ZeRO-Offload
|
||||
<https://arxiv.org/abs/2101.06840>`_,
|
||||
and `ZeRO-Infinity <https://arxiv.org/abs/2104.07857>`_.
|
||||
|
||||
.. note::
|
||||
DeepSpeed first included offloading capabilities with **ZeRO-Offload**, a
|
||||
system for offloading optimizer and gradient states to CPU memory within
|
||||
ZeRO-2. **ZeRO-Infinity** is the next generation of offloading
|
||||
capabilities, accessible to ZeRO-3. ZeRO-Infinity has all of the savings
|
||||
of ZeRO-Offload, plus is able to offload more the model weights and has
|
||||
more effective bandwidth utilization and overlapping of computation and
|
||||
communication.
|
||||
|
||||
|
||||
|
||||
Getting Started
|
||||
---------------
|
||||
|
||||
If you are new to DeepSpeed, check out our `Getting Started <https://www.deepspeed.ai/getting-started/>`_ page.
|
||||
|
||||
Once you are training with DeepSpeed, enabling ZeRO-3 offload is as simple as enabling it
|
||||
in your DeepSpeed configuration! Below are a few examples of ZeRO-3 configurations. Please see
|
||||
our `config guide <https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training>`_
|
||||
for a complete list of options for configuration and performance tuning.
|
||||
|
||||
.. note::
|
||||
ZeRO-Infinity and ZeRO-Offload work best with our heavily optimized
|
||||
:class:`deepspeed.ops.adam.DeepSpeedCPUAdam` optimizer. We recommend using
|
||||
our `optimizer config <https://www.deepspeed.ai/docs/config-json/#optimizer-parameters>`_
|
||||
to instruct :meth:`deepspeed.initialize` to build the optimizer for you.
|
||||
|
||||
ZeRO Configurations
|
||||
===================
|
||||
|
||||
All the settings for DeepSpeed ZeRO are set with the `DeepSpeedZeroConfig`_.
|
||||
The dictionary provided under the ``zero_optimization`` entry of the main
|
||||
DeepSpeed configuration dict will be parsed and validated with this class.
|
||||
Sub-configurations for parameter offload and optimizer offload settings are
|
||||
parsed by `DeepSpeedZeroOffloadParamConfig`_ and
|
||||
`DeepSpeedZeroOffloadOptimizerConfig`_.
|
||||
|
||||
.. _DeepSpeedZeroConfig:
|
||||
.. autopydantic_model:: deepspeed.runtime.zero.config.DeepSpeedZeroConfig
|
||||
|
||||
.. _DeepSpeedZeroOffloadParamConfig:
|
||||
.. autopydantic_model:: deepspeed.runtime.zero.config.DeepSpeedZeroOffloadParamConfig
|
||||
|
||||
.. _DeepSpeedZeroOffloadOptimizerConfig:
|
||||
.. autopydantic_model:: deepspeed.runtime.zero.config.DeepSpeedZeroOffloadOptimizerConfig
|
||||
|
||||
|
||||
Example ZeRO-3 Configurations
|
||||
=============================
|
||||
|
||||
#. Use ZeRO to partition the optimizer states (stage 1), gradients (stage 2),
|
||||
and parameters (stage 3).
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 3
|
||||
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
},
|
||||
"fp16": {
|
||||
"enabled": true
|
||||
},
|
||||
"optimizer": {
|
||||
"type": "AdamW",
|
||||
"params": {
|
||||
"lr": 0.001,
|
||||
"betas": [
|
||||
0.8,
|
||||
0.999
|
||||
],
|
||||
"eps": 1e-8,
|
||||
"weight_decay": 3e-7
|
||||
}
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
|
||||
#. Additionally offload the optimizer states and computations to the CPU with ZeRO-Infinity.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {
|
||||
"device": "cpu"
|
||||
}
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
|
||||
#. Save even more memory by offloading parameters to the CPU memory.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {
|
||||
"device": "cpu"
|
||||
}
|
||||
"offload_param": {
|
||||
"device": "cpu"
|
||||
}
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
|
||||
#. Save even MORE memory by offloading to NVMe (if available on your system):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {
|
||||
"device": "nvme",
|
||||
"nvme_path": "/nvme_data"
|
||||
}
|
||||
"offload_param": {
|
||||
"device": "nvme",
|
||||
"nvme_path": "/nvme_data"
|
||||
}
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
MiCS Configurations
|
||||
===================
|
||||
|
||||
All MiCS configurations are set with `DeepSpeedZeroConfig`. MiCS assumes ZeRO
|
||||
stage 3 optimization is enabled. For now, there are two configuration fields of
|
||||
MiCS `mics_shard_size` and `mics_hierarchical_params_gather`. `mics_shard_size`
|
||||
controls how many devices are used for partitioning the model states.
|
||||
`mics_hierarchical_params_gather` controls whether we use a two-stage
|
||||
hierarchical way to gather parameters in the forward computation.
|
||||
`mics_hierarchical_params_gather` is useful when model states are partitioned
|
||||
across multiple nodes and the cross-node bandwidth is slow. By default this is
|
||||
turned off.
|
||||
|
||||
|
||||
Example MiCS Configurations
|
||||
===========================
|
||||
|
||||
#. Use MiCS to partition the model states (including optimizer states,
|
||||
gradients, and parameters). The following config example partitions the model
|
||||
states to eight devices, and assumes the eight devices are located within a
|
||||
single node (`mics_hierarchical_params_gather` is `False`).
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 3
|
||||
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"mics_shard_size": 8,
|
||||
"mics_hierarchical_params_gather": False,
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
Assumptions
|
||||
===========
|
||||
|
||||
DeepSpeed automatically coordinates the collection (*i.e.,* all-gather),
|
||||
partitioning (*i.e.,* scatter), and offloading of parameters at the
|
||||
granularity of (sub)module ``forward()`` methods. The backward pass is
|
||||
handled similarly. This strategy has two underlying assumptions:
|
||||
|
||||
#. The forward and backward passes of submodules must individually fit in device memory.
|
||||
If this not the case, :class:`deepspeed.zero.TiledLinear` implements
|
||||
**memory-centric tiling** and works with ZeRO-3 to break linear layers
|
||||
into a sequence of smaller submodules that can fit in memory.
|
||||
|
||||
#. A module's parameters are only accessed within its own ``__init__`` and ``forward()`` methods.
|
||||
Otherwise, DeepSpeed must be instructed to collect and re-partition the parameter.
|
||||
See :ref:`external-parameters` for manually coordinating parameters.
|
||||
|
||||
|
||||
Constructing Massive Models
|
||||
---------------------------
|
||||
|
||||
ZeRO-3 enables massive models whose parameters exceed the size of individual
|
||||
nodes in a system. For the typical case of training without model parallelism,
|
||||
you can simply allocate your model in our context:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with deepspeed.zero.Init():
|
||||
model = MyLargeModel()
|
||||
|
||||
|
||||
.. autoclass:: deepspeed.zero.Init
|
||||
:members:
|
||||
|
||||
|
||||
.. _external-parameters:
|
||||
|
||||
Manual Parameter Coordination
|
||||
-----------------------------
|
||||
|
||||
Most models require no modification to be trained with ZeRO-3. However, in
|
||||
some cases one may need to access model weights outside of the training loop,
|
||||
or to share weights across submodules during training. DeepSpeed has
|
||||
several mechanisms to coordinate partitioned weights for ZeRO-3.
|
||||
|
||||
|
||||
Gathering Parameters
|
||||
====================
|
||||
|
||||
DeepSpeed provides mechanisms for collecting (or *gathering*) a partitioned parameter.
|
||||
|
||||
Some models partitioned with :class:`deepspeed.zero.Init` may need to access
|
||||
a module’s weights outside of the class constructor or its ``forward()``
|
||||
method. We refer to these weights as **external parameters**, since these
|
||||
parameters are accessed outside of the module that created them. To do so, use
|
||||
:class:`deepspeed.zero.GatheredParameters` or :meth:`deepspeed.zero.register_external_parameter`.
|
||||
|
||||
.. autoclass:: deepspeed.zero.GatheredParameters
|
||||
:members:
|
||||
|
||||
|
||||
Registering External Parameters
|
||||
===============================
|
||||
|
||||
ZeRO-3 will automatically collect and partition the model parameters as they
|
||||
are needed during the forward and backward passes. However, in some cases a
|
||||
parameter may be used outside of its module's forward pass. We call these
|
||||
*external* parameters. ZeRO-3 can coordinate these parameters if they are
|
||||
registered either automatically or manually.
|
||||
|
||||
|
||||
.. note::
|
||||
DeepSpeed version ``0.3.15`` includes automatic external parameter
|
||||
discovery and registration to support the most common cases. Parameters
|
||||
can still be manually registered if they cannot be automatically
|
||||
detected.
|
||||
|
||||
|
||||
DeepSpeed can automatically detect the following external parameter scenarios:
|
||||
|
||||
|
||||
#. Parameter access: consider the following pattern common in language models such as GPT:
|
||||
|
||||
The tensor ``embeddings.weight`` is used in both ``embeddings.forward()`` and
|
||||
``compute_logits()``. We call ``embeddings.weight`` an *external* parameter
|
||||
because it is used in the training loop outside of its owning module's
|
||||
forward pass.
|
||||
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class LanguageModel(torch.nn.Module):
|
||||
...
|
||||
def forward(self, inputs):
|
||||
embeds = self.embeddings(inputs)
|
||||
...
|
||||
logits = compute_logits(output, self.embeddings.weight)
|
||||
...
|
||||
|
||||
|
||||
#. Returning a parameter:
|
||||
|
||||
``CustomLinear`` returns both an output and its own ``bias`` parameter. DeepSpeed
|
||||
will detect the external ``bias`` parameter and register it with submodules that
|
||||
use ``CustomLinear``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class CustomLinear(torch.nn.Linear):
|
||||
def forward(self, *input):
|
||||
output = super().forward(*input)
|
||||
return output, self.bias
|
||||
|
||||
|
||||
|
||||
.. autofunction:: deepspeed.zero.register_external_parameter
|
||||
|
||||
.. autofunction:: deepspeed.zero.unregister_external_parameter
|
||||
|
||||
|
||||
.. `Module.apply <https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=module+apply#torch.nn.Module.apply>`_
|
||||
|
||||
Overriding Module.apply
|
||||
===============================
|
||||
A convenient mechanism for customizing model initialization is `Module.apply <https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=module+apply#torch.nn.Module.apply>`_.
|
||||
With ZeRO stage 3, ``Module.apply`` implementations must account for parameter partitioning by ``zero.Init`` during model initialization. The default behavior of ZeRO stage 3 is to automatically
|
||||
handle this issue by overriding ``Module.apply`` to ensure that parameters are gathered before access by ``Module.apply``. The benefit of this approach is development convenience, since
|
||||
users are saved the burden of manual parameter coordination in ``Module.apply``. However, the downside is slow model initialization, since all the model parameters (e.g., billions) are gathered
|
||||
even though the common usage of ``Module.apply`` is to customize a few parameters. Developers can disable this default behavior by setting the ``override_module_apply`` configuration knob to ``False``,
|
||||
for faster model initialization at the cost of manually handling partitioned parameters in their ``Module.apply`` implementations.
|
||||
|
||||
|
||||
Memory-Centric Tiling
|
||||
---------------------
|
||||
|
||||
To reduce the working memory requirements of DL training for large models,
|
||||
ZeRO-Infinity includes technique called *memory-centric tiling* that exploits
|
||||
the data fetch and release pattern of ZeRO-3 to reduce the working memory
|
||||
requirements by breaking down a large operator into smaller tiles that can be
|
||||
executed sequentially. When combined with ZeRO-3, the parameter and gradients
|
||||
of each tile can be fetched and released one at a time, reducing the working
|
||||
memory proportional to the number of tiles. Therefore, ZeRO-Infinity can
|
||||
support operators of arbitrary sizes, without refactoring for model
|
||||
parallelism to fit them in limited GPU memory.
|
||||
|
||||
|
||||
.. autoclass:: deepspeed.zero.TiledLinear
|
||||
:members:
|
||||
|
||||
|
||||
Debugging
|
||||
---------
|
||||
|
||||
Debugging ZeRO training is complicated by the partitioning of parameters, gradients, and optimizer states. None of these 3 groups of tensors (model states) can be normally accessed because of that. To overcome that DeepSpeed provides the following routines for accessing individual model states in both their partitioned (local) and unpartitioned (full) forms.
|
||||
|
||||
Important notes:
|
||||
|
||||
# These APIs return tensors that are on accelerator device even if the corresponding model state is offloaded to CPU or NVMe.
|
||||
|
||||
# To access the unpartitioned (full) form, these utilities must be called by all processes participating in the training, even if you decide to do something with the result only in the main process. If all processes don't participate these utilities will hang waiting for all processes to send their contribution.
|
||||
|
||||
# You must be aware that these routines return correct data only in specific phases of the training. So for examples the gradients are valid after ``backward`` and before ``step``. The optimizer states are updated after ``step``. Same goes for fp32 master weights.
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_full_fp32_param
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_full_grad
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_full_optimizer_state
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_local_fp32_param
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_local_grad
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_get_local_optimizer_state
|
||||
|
||||
|
||||
These routines can be used in a training loop as shown in the following snippet.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
backward(loss)
|
||||
[...]
|
||||
from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
|
||||
for n, lp in model.named_parameters():
|
||||
# 1. Access the full states
|
||||
# 1.1) gradient lookup
|
||||
# For zero1 and zero2, gradient lookup must be called after `backward` and before `step`
|
||||
# For zero3, gradient lookup must be called after `backward`
|
||||
hp_grad = safe_get_full_grad(lp)
|
||||
|
||||
|
||||
# 1.2) fp32 and optim states can probably be called anywhere in the training loop, but will be updated after `step`
|
||||
hp = safe_get_full_fp32_param(lp)
|
||||
exp_avg = safe_get_full_optimizer_state(lp, "exp_avg")
|
||||
exp_avg_sq = safe_get_full_optimizer_state(lp, "exp_avg_sq")
|
||||
|
||||
# 2. Access the local states (zero3)
|
||||
# For zero3, all of the parameters, gradients, and optimizer states are partitioned,
|
||||
# and each process can access its corresponding local state.
|
||||
local_hp = safe_get_local_fp32_param(lp)
|
||||
local_hp_grad = safe_get_local_grad(lp)
|
||||
local_exp_avg = safe_get_local_optimizer_state(lp, "exp_avg")
|
||||
local_exp_avg_sq = safe_get_local_optimizer_state(lp, "exp_avg_sq")
|
||||
|
||||
[...]
|
||||
optimizer.step()
|
||||
|
||||
|
||||
|
||||
Modifying Partitioned States
|
||||
----------------------------
|
||||
|
||||
Sometimes, a user may want to modify parameters, gradients, or optimizer states outside of the regular training loop. This is currently difficult in ZeRO training because of partitioning. To overcome that, DeepSpeed provides the following routines for modifying the fp32 master parameters and the fp32 optimizer states.
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_full_fp32_param
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_full_optimizer_state
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_full_grad
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_local_fp32_param
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_local_grad
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_set_local_optimizer_state
|
||||
|
||||
.. autofunction:: deepspeed.utils.safe_update_full_grad_vectorized
|
||||
|
||||
The routines for modifying parameters and optimizer states can be used at any point after initialization of the DeepSpeed engine (i.e., ``deepspeed.initialize()``) as shown in the following snippet.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
[...]
|
||||
from deepspeed.runtime.zero.utils import is_zero_param
|
||||
from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_optimizer_state
|
||||
from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_optimizer_state
|
||||
# Here is an example to zero all the fp32 parameters and optimizer states.
|
||||
for n, lp in model.named_parameters():
|
||||
# 1. For zero stage 1, 2, or 3 set the full fp32 and their full optim states
|
||||
zero_tensor = torch.zeros(lp.ds_shape) if is_zero_param(lp) else torch.zeros(lp.shape)
|
||||
|
||||
safe_set_full_fp32_param(lp, zero_tensor)
|
||||
safe_get_full_optimizer_state(lp, zero_tensor, "exp_avg")
|
||||
safe_get_full_optimizer_state(lp, zero_tensor, "exp_avg_sq")
|
||||
|
||||
# 2. For zero stage 3, each process sets its local fp32 parameters and their local optimizer states individually
|
||||
zero_tensor_local = torch.zeros(lp.ds_tensor.shape)
|
||||
|
||||
safe_set_local_fp32_param(lp, zero_tensor_local)
|
||||
safe_set_local_optimizer_state(lp, zero_tensor_local, "exp_avg")
|
||||
safe_set_local_optimizer_state(lp, zero_tensor_local, "exp_avg_sq")
|
||||
|
||||
[...]
|
||||
|
||||
|
||||
The routines for modifying gradients can be used after ``backward`` but before ``step`` as shown in the following snippet.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
backward(loss)
|
||||
[...]
|
||||
from deepspeed.runtime.zero.utils import is_zero_param
|
||||
from deepspeed.utils import safe_set_full_grad, safe_set_local_grad
|
||||
# Here is an example of how to zero all the gradients.
|
||||
for n, lp in model.named_parameters():
|
||||
# 1. For zero stage 1, 2, or 3 set the full gradient.
|
||||
zero_tensor = torch.zeros(lp.ds_shape) if is_zero_param(lp) else torch.zeros(lp.shape)
|
||||
|
||||
safe_set_full_grad(lp, zero_tensor)
|
||||
|
||||
# 2. For zero stage 3, each process sets its local gradient partition.
|
||||
zero_tensor_local = torch.zeros_like(lp.ds_tensor.shape)
|
||||
|
||||
safe_set_local_grad(lp, zero_tensor_local)
|
||||
|
||||
[...]
|
||||
optimizer.step()
|
||||
|
||||
|
||||
|
||||
GPU Memory Management
|
||||
---------------------
|
||||
|
||||
By default at the end of training with ZeRO stage 3 some parameters could remain unpartitioned and use up some gpu memory.
|
||||
This is done on purpose as an optimization should you resume training again. If you'd like to clear out the cached
|
||||
parameters that use up gpu memory, you can call ``empty_partition_cache`` method of a DeepSpeed engine.
|
||||
|
||||
.. autofunction::deepspeed.DeepSpeedEngine.empty_partition_cache
|
||||
|
||||
The following code snippet illustrates this functionality.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with zero.Init():
|
||||
model = MyLargeModel()
|
||||
|
||||
ds_engine, _, _, _ = deepspeed.initialize(model, ...)
|
||||
for batch in ...:
|
||||
loss = ds_engine(batch)
|
||||
ds_engine.backward(batch)
|
||||
ds_engine.step()
|
||||
|
||||
# Free GPU memory consumed by model parameters
|
||||
ds_engine.empty_partition_cache()
|
||||
|
||||
|
||||
Offload States
|
||||
--------------
|
||||
|
||||
The DeepSpeed engine maintains a set of states in device memory (e.g., CUDA memory). The following API allows you to offload these states to a different device (currently, only CPU memory is supported), reducing the memory footprint on the device.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def offload_states(self,
|
||||
include: Container[OffloadStateTypeEnum] = None,
|
||||
device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
|
||||
pin_memory: bool = True,
|
||||
non_blocking: bool = False) -> None:
|
||||
"""Offload the engine's states to the specified device.
|
||||
|
||||
Arguments:
|
||||
include: Optional. The set of states to offload. If not provided, all states are offloaded.
|
||||
device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported.
|
||||
pin_memory: Optional. Whether to pin the memory of the offloaded states.
|
||||
non_blocking: Optional. Whether to offload the states asynchronously.
|
||||
"""
|
||||
|
||||
You can selectively offload specific states by specifying the ``OffloadStateTypeEnum`` in the include argument. ``OffloadStateTypeEnum`` is an enum that defines the states that can be offloaded. The following states are supported:
|
||||
|
||||
* ``OffloadStateTypeEnum.optim_states``: Optimizer states. Currently, only states of DeepSpeed's FusedAdam optimizer are supported.
|
||||
* ``OffloadStateTypeEnum.hp_params``: FP32 parameters.
|
||||
* ``OffloadStateTypeEnum.lp_params``: BF16/FP16 parameters.
|
||||
* ``OffloadStateTypeEnum.lp_grads``: BF16/FP16 gradients.
|
||||
* ``OffloadStateTypeEnum.contiguous_grad_buffer``: The contiguous gradient buffer for reduce operations.
|
||||
|
||||
Note that offloading states comes with a trade-off between memory savings and computational overhead. This API allows states to be reloaded back into device memory when needed.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def reload_states(self, non_blocking: bool = False) -> None:
|
||||
"""Reload the engine states to the original device.
|
||||
|
||||
Arguments:
|
||||
non_blocking: Optional. Whether to offload the states asynchronously.
|
||||
"""
|
||||
|
||||
Below is an example code snippet demonstrating how to offload FP32 parameters and optimizer states to CPU memory:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Offload after forward, backward, and step
|
||||
ds_engine.offload_states(include=[OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.optim_states])
|
||||
|
||||
# Do something requiring a lot of device memory
|
||||
...
|
||||
# Load states back to device memory
|
||||
ds_engine.reload_states()
|
||||
|
||||
``deepspeed.runtime.zero.offload_states.get_state_devices`` returns devices of the specified state.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]:
|
||||
"""Retrieve the devices of the specified state of the model.
|
||||
|
||||
Args:
|
||||
model (DeepSpeedEngine): The model whose device allocations are to be checked.
|
||||
state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved.
|
||||
|
||||
Returns:
|
||||
Set[torch.device]: A set of devices of the specified state.
|
||||
|
||||
"""
|
||||
Reference in New Issue
Block a user