diff --git a/.github/workflows/nv-a6000.yml b/.github/workflows/nv-a6000.yml index 3ce406948432..484948b28e34 100644 --- a/.github/workflows/nv-a6000.yml +++ b/.github/workflows/nv-a6000.yml @@ -47,7 +47,8 @@ jobs: - name: Install deepspeed run: | python -m pip install docutils==0.18.1 jinja2==3.0 urllib3==1.26.11 ninja - python -m pip install pydantic==1.10.11 + # Update packages included in the container that do not support pydantic 2+ to versions that do + python -m pip install thinc spacy confection --upgrade python -m pip install .[dev,1bit,autotuning,inf] ds_report - name: Python environment diff --git a/.github/workflows/xpu-max1100.yml b/.github/workflows/xpu-max1100.yml index 1042db100a21..adeeb0acade2 100644 --- a/.github/workflows/xpu-max1100.yml +++ b/.github/workflows/xpu-max1100.yml @@ -21,7 +21,7 @@ on: - "deepspeed/runtime/zero/parameter_offload.py" - "deepspeed/runtime/pipe/engine.py" - "deepspeed/runtime/utils.py" - - "opbuilder/xpu/**" + - "op_builder/xpu/**" concurrency: group: ${{ github.workflow }}-${{ github.ref }} diff --git a/README.md b/README.md index 304169b56777..2f6661ef5860 100755 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ DeepSpeed empowers ChatGPT-like model training with a single click, offering 15x speedup over SOTA RLHF systems with unprecedented cost reduction at all scales; [learn how](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat). - +* [2024/08] [DeepSpeed on Windows](https://github.com/microsoft/DeepSpeed/tree/master/blogs/windows/08-2024/README.md) [[日本語](https://github.com/microsoft/DeepSpeed/tree/master/blogs/windows/08-2024/japanese/README.md)] * [2024/08] [DeepNVMe: Improving DL Applications through I/O Optimizations](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-gds/README.md) [[日本語](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-gds/japanese/README.md)] * [2024/07] [DeepSpeed Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-ucp/README.md) [[中文](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-ucp/chinese/README.md)] [[日本語](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-ucp/japanese/README.md)] * [2024/03] [DeepSpeed-FP6:The power of FP6-Centric Serving for Large Language Models](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024) [[English](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024/README.md)] [[中文](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024/README-Chinese.md)] diff --git a/blogs/windows/08-2024/README.md b/blogs/windows/08-2024/README.md new file mode 100644 index 000000000000..34e11bd47792 --- /dev/null +++ b/blogs/windows/08-2024/README.md @@ -0,0 +1,101 @@ +
+ +# DeepSpeed on Windows + +
+ +# Introduction + +DeepSpeed is a popular open-source deep learning optimization library that makes distributed training and inference easy, efficient, and effective. DeepSpeed has been widely used to train a variety of state-of-the-art models, including Phi-3, Megatron-Turing-530B, BLOOM-176B, and Arctic because of its rich suite of sophisticated optimizations (e.g., ZeRO, 3D parallelism, MoE, etc.). However, the lack of native support for Microsoft Windows, the most popular operating system, means that DeepSpeed innovations are inaccessible to many AI developers and users. To address this problem, we started an effort to make DeepSpeed run natively with full features on Windows, while ensuring the same ease-of-use enjoyed on Linux. + +In this blog, we are pleased to announce some early achievements on this journey: DeepSpeed can now be installed in Windows and run natively for single-GPU training, finetuning, and inferencing. Importantly, both the installation and usage experiences are identical to those on Linux. Furthermore, the finetuning and inferencing workloads demonstrate the functioning of three critical DeepSpeed features, HuggingFace Transformers integration, LoRA support, and CPU Offloading. DeepSpeed on Windows is available in DeepSpeed versions 0.14.5 and above. In the rest of this blog, we present examples to demonstrate these achievements. + +# Evaluation Environment +We conducted the experiments on a Surface Laptop Studio 2 running Windows 11 Version 23H2 and Build 22631.3880. The laptop is equipped with a single NVIDIA RTX A2000 GPU with 4GB VRAM. We used Pytorch version 2.3.0 and HuggingFace Transformers version 4.41.2. The example scripts used are from the [DeepSpeedExamples repo](https://github.com/microsoft/DeepSpeedExamples), therefore you need to clone the repo before running any of the following examples. + +# Installation +DeepSpeed can be installed on Windows in one of two ways. The easier way is to use the pip package manager, while the other is to build from source. The prerequisites for in both cases are Python 3.x and Pytorch with CUDA support. + +## Installing via pip +To install DeepSpeed, simply run: `pip install deepspeed`. This will install the latest version of DeepSpeed (0.14.5 at this time). Unlike the Linux counterpart, the Windows version comes with all the operators already prebuilt, so there is no need to have a CUDA SDK or C++ compiler installed. + +
+ +
+ +
+ pip installation of DeepSpeed on Windows. +
+ + +## Building from Source +To build DeepSpeed from source, you need to clone the DeepSpeed repository and run the `build_win.bat` compilation script. + + +## Validating Installation +Regardless of the installation choice, you can check that the installation was successful by running ds_report. The output should look like this: + + +
+ +
+ +
+ ds_report output confirming Windows installation of DeepSpeed. +
+ +# Pretraining Examples +We use an image classification model, CIFAR10, and a language model, BERT, to demonstrate pretraining on Windows with DeepSpeed. + +## Pretraining CIFAR10 +The scripts and codes required for CIFAR10 pretraining example are available in the following path: DeepSpeedExamples\training\cifar. You can launch the CIFAR10 pretraining experiment using the following command: `deepspeed cifar10_deepspeed.py –deepspeed`. The final output should look something like this: +
+ +
+ +
+ Pretraining CIFAR10 model on Windows using DeepSpeed. +
+ +## Pretraining BERT +The scripts and codes for the BERT pretraining example are available in the following path: DeepSpeedExamples\training\HelloDeepSpeed. You can launch the BERT pretraining experiment using the following command: `deepspeed train_bert_ds.py --checkpoint_dir experiment_deepspeed`. The final output should look like this: + +
+ +
+ +
+ Pretraining BERT model on Windows using DeepSpeed. +
+ +# Fine Tuning Example +We demonstrate fine tuning capability by using the supervised fine tuning (SFT) step of DeepSpeed-Chat application. We conduct SFT of the HuggingFace facebook/opt-125m model while enabling LoRA and CPU offloading memory optimizations. The command line for running this example is as follows: +deepspeed training\step1_supervised_finetuning\main.py --model_name_or_path facebook/opt-125m --gradient_accumulation_steps 8 --lora_dim 128 --only_optimize_lora --print_loss --zero_stage 2 --deepspeed --dtype bf16 --offload --output_dir output +The output should look like this: + +
+ +
+ +
+ Supervised Finetuning of facebook/opt-125m model on Windows using DeepSpeed. +
+ +# Inference Example +We demonstrate inference capability by using ZeRO-Inference for token generation. ZeRO-Inference reduces hardware cost of inferencing by offloading to CPU or NVMe memories. We use the example scripts here to run token generation using Llama-2-7B model from HuggingFace. We offload the model weights to CPU memory since the 4GB VRAM is insufficient to host both the model and the generation working set. We use the following command line to generate 32 tokens from a prompt of 8 tokens: +deepspeed run_model.py --model meta-llama/Llama-2-7b-hf --batch-size 64 --prompt-len 8 --gen-len 32 --cpu-offload +The output will look something like this: + +
+ +
+ +
+ LLAMA2-7B token generation on Windows using ZeRO-Inference. +
+ +# Summary +Enabling DeepSpeed, a popular deep learning framework, to run natively on Windows, the most popular operating system, is a crucial step towards empowering every person and every organization to benefit from the ongoing AI revolution. In this blog, we have shared early results of our work towards this goal. Although Windows support of DeepSpeed is a work-in-progress, we hope that the above updates are encouraging and already useful to users. The next items on our roadmap include running on multiple GPUs, weight quantization, and performance studies. + +# Acknowledgements +This work is a result of significant contributions from current and former DeepSpeed members including Costin Eseanu, Logan Adams, Elton Zheng, Reza Yazdani Aminabadi, Martin Cai, and Olatunji Ruwase. We also acknowledge the valuable contributions of DeepSpeed users who righteously demanded this feature, provided critical workarounds, partial solutions, and constructive feedback, and most importantly, stuck with us. diff --git a/blogs/windows/08-2024/japanese/README.md b/blogs/windows/08-2024/japanese/README.md new file mode 100644 index 000000000000..7e437f737f58 --- /dev/null +++ b/blogs/windows/08-2024/japanese/README.md @@ -0,0 +1,123 @@ +
+ +# DeepSpeedのWindowsサポート + +
+ +# はじめに + +DeepSpeedは、分散学習と推論を簡単かつ効率的に行うための人気のあるオープンソースの深層学習最適化ライブラリです。DeepSpeedは、その豊富かつ高度な最適化機能（例：ZeRO、3D parallelism, MoEなど）のおかげで、Phi-3、Megatron-Turing-530B、BLOOM-176B、Arcticなどの最先端モデルの学習に広く利用されています。しかし、最も普及しているオペレーティングシステムであるMicrosoft Windowsをネイティブにサポートしていなかったため、多くのAI開発者やユーザーが、DeepSpeedの革新的な機能を利用できない状態でした。この問題を解決するため、DeepSpeedの完全な機能をWindows上でネイティブに実行し、Linux上と同じ使いやすさを実現するための取り組みを開始しました。 + +このブログでは、この取り組みの最初の成果をお知らせします。現在、DeepSpeedはWindowsにインストールし、単一GPUでの学習、ファインチューニング、および推論をネイティブに実行できるようになりました。ここで重要なこととして、インストールと利用は、Linuxとまったく同じように行えます。ファインチューニングと推論のワークロードを通じて、HuggingFace Transformers との統合、LoRAのサポート、CPUオフロードの3つの重要なDeepSpeedの機能が、正しく動作していることが確認できました。このWindowsサポートは、バージョン0.14.5以降で利用可能です。このブログの残りの部分では、これらの成果を示す例を紹介します。 + +# テスト環境 + +Windows 11 Version 23H2 および Build 22631.3880 を実行している Surface Laptop Studio 2 でテストを行いました。このハードウェアには、4GBのVRAMを搭載した NVIDIA RTX A2000 GPU が1つ搭載されています。また、PyTorchバージョン 2.3.0 および HuggingFace Transformersバージョン 4.41.2 を使用しました。使用したサンプルスクリプトは[DeepSpeedExamplesリポジトリ](https://github.com/microsoft/DeepSpeedExamples)から取得できます。以下の例を実行する前にリポジトリをクローンしてください。 + +# インストール + +DeepSpeedは、2つの方法でWindowsにインストールできます。より簡単な方法は、pipパッケージマネージャーを使用することで、もう一方はソースからビルドする方法です。どちらの場合も、Python 3.xとCUDAサポート付きのPyTorchが必要です。 + +## pipを使用したインストール + +DeepSpeedをインストールするには、単に次のコマンドを実行します: `pip install deepspeed`。 +これにより、最新バージョンのDeepSpeed（現時点では0.14.5）がインストールされます。Linux版とは異なり、Windows版ではすべてのオペレーターがすでにビルド済みであるため、CUDA SDKやC++コンパイラをインストールする必要はありません。 + +
+ +
+ +
+ pipによるWindowsへのDeepSpeedのインストール +
+ + +## ソースからのビルド + +ソースからDeepSpeedをビルドするには、DeepSpeedリポジトリをクローンし、コンパイルスクリプトである `build_win.bat` を実行する必要があります。 + +## インストールの検証 + +インストール方法にかかわらず、`ds_report`を実行してインストールが成功したかどうかを確認できます。出力は次のようになります： + +
+ +
+ +
+ DeepSpeedのWindowsインストールを確認するds_reportの出力 +
+ +# 事前学習の例 + +Windows上でDeepSpeedを使用した事前学習の例として、画像分類モデルCIFAR10と言語モデルBERTの実行例を示します。 + +## CIFAR10の事前学習 + +CIFAR10の事前学習に必要なスクリプトとコードは、次のパスにあります: `DeepSpeedExamples\training\cifar` + +以下のコマンドを使用してCIFAR10の事前学習を開始できます: `deepspeed cifar10_deepspeed.py –deepspeed` + +出力は次のようになります。 + +
+ +
+ +
+ DeepSpeedによるWindowsでのCIFAR10モデルの事前学習 +
+ +## BERTの事前学習 + +BERTの事前学習に必要なスクリプトとコードは、次のパスにあります: `DeepSpeedExamples\training\HelloDeepSpeed` + +以下のコマンドを使用してBERTの事前学習を開始できます: `deepspeed train_bert_ds.py --checkpoint_dir experiment_deepspeed` + +出力は次のようになります。 + +
+ +
+ +
+ DeepSpeedによるWindowsでのBERTモデルの事前学習 +
+ +# ファインチューニングの例 + +DeepSpeed-Chatアプリケーションの教師ありファインチューニング（supervised fine tuning; SFT）を使用して、ファインチューニングの機能を示します。LoRAおよびCPUオフロードメモリ最適化を有効にして、 HuggingFace の `facebook/opt-125m` モデルのSFTを実施します。この例を実行するためのコマンドラインは次のとおりです: `deepspeed training\step1_supervised_finetuning\main.py --model_name_or_path facebook/opt-125m --gradient_accumulation_steps 8 --lora_dim 128 --only_optimize_lora --print_loss --zero_stage 2 --deepspeed --dtype bf16 --offload --output_dir output` + +出力は次のようになります。 + +
+ +
+ +
+ DeepSpeedを使用したWindowsでの facebook/opt-125m モデルのファインチューニング +
+ +# 推論の例 + +推論の機能を示すために、トークン生成のためのZeRO-Inferenceを使用します。ZeRO-Inferenceは、CPUまたはNVMeメモリにオフロードすることで推論のハードウェアコストを削減します。ここでは、サンプルスクリプトを使用して、HuggingFaceのLlama-2-7Bモデルを使用したトークン生成を実行します。4GBのVRAMではモデルと生成処理の両方を実効するのに十分ではないため、モデルパラメータをCPUメモリにオフロードします。 + +次のコマンドラインを使用して、8トークンのプロンプトから32トークンを生成します: `deepspeed run_model.py --model meta-llama/Llama-2-7b-hf --batch-size 64 --prompt-len 8 --gen-len 32 --cpu-offload` + +出力は次のようになります。 + +
+ +
+ +
+ DeepSpeedのZeRO-InferenceによるWindowsでのLLAMA2-7Bのトークン生成 +
+ +# まとめ + +最も広く使われているオペレーティングシステムであるWindowsで、深層学習フレームワークであるDeepSpeedをネイティブに実行できるようにすることは、多くの人と組織が、今まさに進行中のAI革命の恩恵を受けるための重要な一歩です。このブログでは、この目標に向けたプロジェクトの、最初の成果を共有しました。Windowsのサポートは現在進行中のプロジェクトですが、今回の成果が多くのユーザにとって活用され、またさらに発展していけることを願っています。次のロードマップには、複数のGPUでの実行、モデルパラメータの量子化、パフォーマンスの詳細な分析が含まれます。 + +# 謝辞 + +このプロジェクトは、Costin Eseanu、Logan Adams、Elton Zheng、Reza Yazdani Aminabadi、Martin Cai、Olatunji Ruwaseを含むDeepSpeedメンバーによる大きな貢献の結果です。また、この機能を必要とし、様々な問題の解決策や、建設的なフィードバックを提供し、私たちと共に歩んでくれたDeepSpeedユーザーの重要な貢献に感謝します。 diff --git a/blogs/windows/08-2024/media/bert_training.png b/blogs/windows/08-2024/media/bert_training.png new file mode 100644 index 000000000000..c5935e47747e Binary files /dev/null and b/blogs/windows/08-2024/media/bert_training.png differ diff --git a/blogs/windows/08-2024/media/cifar10_training.png b/blogs/windows/08-2024/media/cifar10_training.png new file mode 100644 index 000000000000..99f3fa25bc70 Binary files /dev/null and b/blogs/windows/08-2024/media/cifar10_training.png differ diff --git a/blogs/windows/08-2024/media/ds_report.png b/blogs/windows/08-2024/media/ds_report.png new file mode 100644 index 000000000000..43d82d724ed2 Binary files /dev/null and b/blogs/windows/08-2024/media/ds_report.png differ diff --git a/blogs/windows/08-2024/media/llama2-7b_inference.png b/blogs/windows/08-2024/media/llama2-7b_inference.png new file mode 100644 index 000000000000..f5874468a854 Binary files /dev/null and b/blogs/windows/08-2024/media/llama2-7b_inference.png differ diff --git a/blogs/windows/08-2024/media/opt125m_finetuning.png b/blogs/windows/08-2024/media/opt125m_finetuning.png new file mode 100644 index 000000000000..ed6d1522e3b3 Binary files /dev/null and b/blogs/windows/08-2024/media/opt125m_finetuning.png differ diff --git a/blogs/windows/08-2024/media/win_pip_install_deepspeed.png b/blogs/windows/08-2024/media/win_pip_install_deepspeed.png new file mode 100644 index 000000000000..3b87c95ef144 Binary files /dev/null and b/blogs/windows/08-2024/media/win_pip_install_deepspeed.png differ diff --git a/csrc/aio/py_lib/deepspeed_aio_op_desc.h b/csrc/aio/py_lib/deepspeed_aio_op_desc.h index 7305f6920c91..350d28d29d58 100644 --- a/csrc/aio/py_lib/deepspeed_aio_op_desc.h +++ b/csrc/aio/py_lib/deepspeed_aio_op_desc.h @@ -16,7 +16,7 @@ struct io_op_desc_t { const std::string _filename; const long long int _file_num_bytes; const int _num_threads; - const int _num_bytes_per_thread; + const long long int _num_bytes_per_thread; torch::Tensor _contiguous_buffer; const bool _validate; diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index fe0043547860..d8655299282f 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -77,6 +77,7 @@ def initialize(args=None, dist_init_required: Optional[bool] = None, collate_fn=None, config=None, + mesh_param=None, config_params=None): """Initialize the DeepSpeed Engine. @@ -144,10 +145,22 @@ def initialize(args=None, distributed_port=distributed_port, dist_init_required=dist_init_required) + ##TODO: combine reuse mpu as mesh device and vice versa # Set config using config_params for backwards compat if config is None and config_params is not None: config = config_params + mesh_device = None + if mesh_param: + logger.info(f"mesh_param to Initialize mesh device: {mesh_param}") + mesh_device = dist.initialize_mesh_device(mesh_param, ("data_parallel", "sequence_parallel")) + #if config file has sequence parallelize and data parallelize, then use them to initialize mesh device + elif config is not None: + if "sequence_parallel_size" in config and "data_parallel_size" in config: + logger.info(f"config to Initialize mesh device: {config}") + mesh_device = dist.initialize_mesh_device((config["data_parallel_size"], config["sequence_parallel_size"]), \ + ("data_parallel", "sequence_parallel")) + # Check for deepscale_config for backwards compat if hasattr(args, "deepscale_config") and args.deepscale_config is not None: logger.warning("************ --deepscale_config is deprecated, please use --deepspeed_config ************") @@ -162,9 +175,8 @@ def initialize(args=None, assert config is None, "Not sure how to proceed, we were given deepspeed configs in the deepspeed arguments and deepspeed.initialize() function call" config = args.deepspeed_config assert config is not None, "DeepSpeed requires --deepspeed_config to specify configuration file" - if not isinstance(model, PipelineModule): - config_class = DeepSpeedConfig(config, mpu) + config_class = DeepSpeedConfig(config, mpu, mesh_device=mesh_device) if config_class.hybrid_engine.enabled: engine = DeepSpeedHybridEngine(args=args, model=model, @@ -188,6 +200,7 @@ def initialize(args=None, dist_init_required=dist_init_required, collate_fn=collate_fn, config=config, + mesh_device=mesh_device, config_class=config_class) else: assert mpu is None, "mpu must be None with pipeline parallelism" @@ -208,7 +221,12 @@ def initialize(args=None, # Restore zero.Init context if necessary zero.partition_parameters.restore_init_context() - return_items = [engine, engine.optimizer, engine.training_dataloader, engine.lr_scheduler] + return_items = [ + engine, + engine.optimizer, + engine.training_dataloader, + engine.lr_scheduler, + ] return tuple(return_items) diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py old mode 100644 new mode 100755 index 85b7fab2c548..2895e0f2e011 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -600,6 +600,21 @@ def get_all_ranks_from_group(group=None): return group_ranks +def initialize_mesh_device(mesh_shape, mesh_dim_names): + global cdb + assert cdb is not None and cdb.is_initialized( + ), 'DeepSpeed backend not set, please initialize it using init_process_group()' + mesh_device = None + if hasattr(cdb, 'init_device_mesh'): + utils.logger.info(f"Initializing mesh device with backend {cdb.name} \ + with shape {mesh_shape} and dim names {mesh_dim_names}") + mesh_device = cdb.init_device_mesh(mesh_shape, mesh_dim_names) + else: + if get_rank() == 0: + utils.logger.warning_once(f"Backend {cdb.name} does not support mesh device initialization") + return mesh_device + + # Main DeepSpeed Comms. public API. def init_distributed(dist_backend=None, auto_mpi_discovery=True, diff --git a/deepspeed/comm/config.py b/deepspeed/comm/config.py index 1c441bb6bfe9..57501c9dd237 100644 --- a/deepspeed/comm/config.py +++ b/deepspeed/comm/config.py @@ -3,20 +3,12 @@ # DeepSpeed Team -from .constants import * -from ..pydantic_v1 import BaseModel - +from deepspeed.runtime.config_utils import DeepSpeedConfigModel -class CommsConfig(BaseModel): - - class Config: - validate_all = True - validate_assignment = True - use_enum_values = True - extra = 'forbid' +from .constants import * -class CommsLoggerConfig(CommsConfig): +class CommsLoggerConfig(DeepSpeedConfigModel): enabled: bool = COMMS_LOGGER_ENABLED_DEFAULT prof_all: bool = COMMS_LOGGER_PROF_ALL_DEFAULT prof_ops: list = COMMS_LOGGER_PROF_OPS_DEFAULT diff --git a/deepspeed/comm/torch.py b/deepspeed/comm/torch.py old mode 100644 new mode 100755 index 83754e98f033..ed2645d415c4 --- a/deepspeed/comm/torch.py +++ b/deepspeed/comm/torch.py @@ -386,6 +386,14 @@ def _reduce_op(self, op): op = torch.distributed.ReduceOp.BXOR return op + def init_device_mesh(self, mesh_shape, mesh_dim_names): + if not required_torch_version(min_version=2.2): + raise RuntimeError(f"Current torch version does not have device mesh" + f"api (torch.__version__: {torch.__version__})") + return torch.distributed.device_mesh.init_device_mesh(get_accelerator().current_device_name(), + mesh_shape, + mesh_dim_names=mesh_dim_names) + # This will become a light-weight wrapper around torch.distributed functions # TODO: create some example to show how this wrapper can help profile communication diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 1d5018aaa75b..c7c7684fff79 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -5,38 +5,25 @@ import torch import deepspeed -from deepspeed.pydantic_v1 import Field, validator +from pydantic import Field, field_validator from deepspeed.runtime.config_utils import DeepSpeedConfigModel from deepspeed.runtime.zero.config import DeepSpeedZeroConfig -from typing import Dict, Union +from typing import Dict, Union, Optional from enum import Enum class DtypeEnum(Enum): - # The torch dtype must always be the first value (so we return torch.dtype) - fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" - fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" - bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat" - int8 = torch.int8, "torch.int8", "int8" - - # Copied from https://stackoverflow.com/a/43210118 - # Allows us to use multiple values for each Enum index and returns first - # listed value when Enum is called - def __new__(cls, *values): - obj = object.__new__(cls) - # first value is canonical value - obj._value_ = values[0] - for other_value in values[1:]: - cls._value2member_map_[other_value] = obj - obj._all_values = values - return obj - - def __repr__(self): - return "<%s.%s: %s>" % ( - self.__class__.__name__, - self._name_, - ", ".join([repr(v) for v in self._all_values]), - ) + fp16 = (torch.float16, "torch.float16", "fp16", "float16", "half") + fp32 = (torch.float32, "torch.float32", "fp32", "float32", "float") + bf16 = (torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat") + int8 = (torch.int8, "torch.int8", "int8") + + @classmethod + def from_str(cls, value: str): + for dtype in cls: + if value in dtype.value: + return dtype + raise ValueError(f"'{value}' is not a valid DtypeEnum") class MoETypeEnum(str, Enum): @@ -91,24 +78,24 @@ class QuantTypeEnum(str, Enum): class BaseQuantConfig(DeepSpeedConfigModel): - enabled = True - num_bits = 8 + enabled: bool = True + num_bits: int = 8 q_type: QuantTypeEnum = QuantTypeEnum.sym q_groups: int = 1 class WeightQuantConfig(BaseQuantConfig): - enabled = True + enabled: bool = True quantized_initialization: Dict = {} post_init_quant: Dict = {} class ActivationQuantConfig(BaseQuantConfig): - enabled = True + enabled: bool = True class QKVQuantConfig(DeepSpeedConfigModel): - enabled = True + enabled: bool = True class QuantizationConfig(DeepSpeedConfigModel): @@ -120,9 +107,9 @@ class QuantizationConfig(DeepSpeedConfigModel): # todo: brainstorm on how to do ckpt loading for DS inference class InferenceCheckpointConfig(DeepSpeedConfigModel): - checkpoint_dir: str = None - save_mp_checkpoint_path: str = None - base_dir: str = None + checkpoint_dir: Optional[str] = None + save_mp_checkpoint_path: Optional[str] = None + base_dir: Optional[str] = None class DeepSpeedInferenceConfig(DeepSpeedConfigModel): @@ -136,7 +123,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): `(attention_output projection, transformer output projection)` """ - dtype: DtypeEnum = torch.float16 + dtype: torch.dtype = torch.float16 """ Desired model data type, will convert model to this type. Supported target types: `torch.half`, `torch.int8`, `torch.float` @@ -198,7 +185,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): """ #todo: refactor the following 3 into the new checkpoint_config - checkpoint: Union[str, Dict] = None + checkpoint: Optional[Union[str, Dict]] = None """ Path to deepspeed compatible checkpoint or path to JSON with load policy. """ @@ -214,7 +201,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): specifying whether the inference-module is created with empty or real Tensor """ - save_mp_checkpoint_path: str = None + save_mp_checkpoint_path: Optional[str] = None """ The path for which we want to save the loaded model with a checkpoint. This feature is used for adjusting the parallelism degree to help alleviate the @@ -243,19 +230,21 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): replace_method: str = Field( "auto", - deprecated=True, - deprecated_msg="This parameter is no longer needed, please remove from your call to DeepSpeed-inference") + json_schema_extra={ + "deprecated": True, + "deprecated_msg": "This parameter is no longer needed, please remove from your call to DeepSpeed-inference" + }) - injection_policy: Dict = Field(None, alias="injection_dict") + injection_policy: Optional[Dict] = Field(None, alias="injection_dict") """ Dictionary mapping a client nn.Module to its corresponding injection policy. e.g., `{BertLayer : deepspeed.inference.HFBertLayerPolicy}` """ - injection_policy_tuple: tuple = None + injection_policy_tuple: Optional[tuple] = None """ TODO: Add docs """ - config: Dict = Field(None, alias="args") # todo: really no need for this field if we can refactor + config: Optional[Dict] = Field(None, alias="args") # todo: really no need for this field if we can refactor max_out_tokens: int = Field(1024, alias="max_tokens") """ @@ -274,31 +263,49 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel): transposed_mode: bool = Field(False, alias="transposed_mode") - mp_size: int = Field(1, deprecated=True, new_param="tensor_parallel.tp_size") + mp_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.tp_size"}) """ Desired model parallel size, default is 1 meaning no model parallelism. Deprecated, please use the ``tensor_parallel` config to control model parallelism. """ - mpu: object = Field(None, deprecated=True, new_param="tensor_parallel.mpu") - ep_size: int = Field(1, deprecated=True, new_param="moe.ep_size") - ep_group: object = Field(None, alias="expert_group", deprecated=True, new_param="moe.ep_group") - ep_mp_group: object = Field(None, alias="expert_mp_group", deprecated=True, new_param="moe.ep_mp_group") - moe_experts: list = Field([1], deprecated=True, new_param="moe.moe_experts") - moe_type: MoETypeEnum = Field(MoETypeEnum.standard, deprecated=True, new_param="moe.type") - - @validator("moe") + mpu: object = Field(None, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.mpu"}) + ep_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "moe.ep_size"}) + ep_group: object = Field(None, + alias="expert_group", + json_schema_extra={ + "deprecated": True, + "new_param": "moe.ep_group" + }) + ep_mp_group: object = Field(None, + alias="expert_mp_group", + json_schema_extra={ + "deprecated": True, + "new_param": "moe.ep_mp_group" + }) + moe_experts: list = Field([1], json_schema_extra={"deprecated": True, "new_param": "moe.moe_experts"}) + moe_type: MoETypeEnum = Field(MoETypeEnum.standard, + json_schema_extra={ + "deprecated": True, + "new_param": "moe.type" + }) + + @field_validator("dtype", mode="before") + def validate_dtype(cls, field_value, values): + if isinstance(field_value, str): + return DtypeEnum.from_str(field_value).value[0] + if isinstance(field_value, torch.dtype): + return field_value + raise TypeError(f"Invalid type for dtype: {type(field_value)}") + + @field_validator("moe") def moe_backward_compat(cls, field_value, values): if isinstance(field_value, bool): return DeepSpeedMoEConfig(moe=field_value) return field_value - @validator("use_triton") + @field_validator("use_triton") def has_triton(cls, field_value, values): if field_value and not deepspeed.HAS_TRITON: raise ValueError('Triton needs to be installed to use deepspeed with triton kernels') return field_value - - class Config: - # Get the str representation of the datatype for serialization - json_encoders = {torch.dtype: lambda x: str(x)} diff --git a/deepspeed/inference/v2/checkpoint/huggingface_engine.py b/deepspeed/inference/v2/checkpoint/huggingface_engine.py index 46a84c61f884..d88d99ebebfd 100644 --- a/deepspeed/inference/v2/checkpoint/huggingface_engine.py +++ b/deepspeed/inference/v2/checkpoint/huggingface_engine.py @@ -15,13 +15,13 @@ class HuggingFaceCheckpointEngine(CheckpointEngineBase): - def __init__(self, model_name_or_path: str, auth_token: str = None) -> None: + def __init__(self, model_name_or_path: str, auth_token: str = None, **hf_kwargs) -> None: super().__init__() from transformers import AutoConfig, GenerationConfig self.model_name_or_path = model_name_or_path self.auth_token = auth_token - self.model_config = AutoConfig.from_pretrained(self.model_name_or_path) + self.model_config = AutoConfig.from_pretrained(self.model_name_or_path, **hf_kwargs) # Define this property here so we can use it in the model implementation if not hasattr(self.model_config, "max_seq_length"): if hasattr(self.model_config, "max_position_embeddings"): @@ -108,6 +108,12 @@ def parameters(self) -> Iterable[Tuple[str, torch.Tensor]]: for checkpoint in self._all_ckpt_paths: inference_logger().info(f"Loading checkpoint: {checkpoint}") checkpoint_sd = self._checkpoint_load_fn(checkpoint) + + # If the model has tied embeddings, we need to make sure the lm_head weights are tied to the embeddings weights + if hasattr(self.model_config, "tie_word_embeddings") and self.model_config.tie_word_embeddings: + if self.model_config.model_type == "qwen2": + checkpoint_sd["lm_head.weight"] = checkpoint_sd["model.embed_tokens.weight"] + param_keys = list(checkpoint_sd.keys()) for param_name in param_keys: param = checkpoint_sd[param_name] diff --git a/deepspeed/inference/v2/config_v2.py b/deepspeed/inference/v2/config_v2.py index 85e4b7a0e0a0..325b57d8f56a 100644 --- a/deepspeed/inference/v2/config_v2.py +++ b/deepspeed/inference/v2/config_v2.py @@ -3,8 +3,9 @@ # DeepSpeed Team +from pydantic import Field from typing import Optional -from deepspeed.pydantic_v1 import Field + from deepspeed.runtime.config_utils import DeepSpeedConfigModel from .ragged import DSStateManagerConfig diff --git a/deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h b/deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h index 2cc430ccfe34..f5104f899d9c 100644 --- a/deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h +++ b/deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h @@ -14,5 +14,8 @@ } else if (4 == N_TOP_K) { \ constexpr int CONST_TOP_K = 4; \ __VA_ARGS__(); \ + } else if (8 == N_TOP_K) { \ + constexpr int CONST_TOP_K = 8; \ + __VA_ARGS__(); \ } \ }() diff --git a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py b/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py index 7e1ec1a13cb9..aacbec0bd3ae 100644 --- a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py +++ b/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.py @@ -19,7 +19,7 @@ class BlockedRotaryEmbeddings(DSKernelBase): supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16] supported_head_sizes = [64, 80, 96, 128] - supported_q_ratios = [1, 2, 4, 5, 8, 16, 29, 35, 36, 71] + supported_q_ratios = [1, 2, 4, 5, 6, 7, 8, 16, 29, 35, 36, 71] def __init__(self, head_size: int, n_q_heads: int, n_kv_heads: int, dtype: torch.dtype, rotary_dim: int, theta_base: float) -> None: diff --git a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu b/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu index fbafece5ccf2..f7bc693eefee 100644 --- a/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu +++ b/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu @@ -265,6 +265,8 @@ void launch_kv_rotary_kernel(T* kv_cache, LAUNCH_KV_ROTARY_FOR_Q_RATIO(2) LAUNCH_KV_ROTARY_FOR_Q_RATIO(4) LAUNCH_KV_ROTARY_FOR_Q_RATIO(5) + LAUNCH_KV_ROTARY_FOR_Q_RATIO(6) + LAUNCH_KV_ROTARY_FOR_Q_RATIO(7) LAUNCH_KV_ROTARY_FOR_Q_RATIO(8) LAUNCH_KV_ROTARY_FOR_Q_RATIO(16) LAUNCH_KV_ROTARY_FOR_Q_RATIO(29) diff --git a/deepspeed/inference/v2/model_implementations/flat_model_helpers.py b/deepspeed/inference/v2/model_implementations/flat_model_helpers.py index ebdb59bca920..c5e02adaffc4 100644 --- a/deepspeed/inference/v2/model_implementations/flat_model_helpers.py +++ b/deepspeed/inference/v2/model_implementations/flat_model_helpers.py @@ -27,9 +27,9 @@ class TensorMetadata(DeepSpeedConfigModel): """ A class to represent a tensor specification. """ - dtype: Optional[str] - shape: Optional[Tuple[int, ...]] - strides: Optional[Tuple[int, ...]] + dtype: Optional[str] = None + shape: Optional[Tuple[int, ...]] = None + strides: Optional[Tuple[int, ...]] = None offset: int @@ -37,7 +37,7 @@ class ParameterMetadata(DeepSpeedConfigModel): """ A class to represent a parameter specification. """ - core_param: TensorMetadata = None + core_param: Optional[TensorMetadata] = None aux_params: Dict[str, TensorMetadata] = {} diff --git a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py index b4621257ff82..e499379da7e3 100644 --- a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py +++ b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/container.py @@ -8,45 +8,45 @@ from ..common_parameters import * from ..layer_container_base import LayerContainer ''' - # HF Qwen1.5-MoE-A2.7B model looks like this: + # HF Qwen2-57B-A14B model looks like this: Qwen2MoeForCausalLM( (model): Qwen2MoeModel( - (embed_tokens): Embedding(151936, 2048) + (embed_tokens): Embedding(151936, 3584) (layers): ModuleList( - (0-23): 24 x Qwen2MoeDecoderLayer( + (0-27): 28 x Qwen2MoeDecoderLayer( (self_attn): Qwen2MoeSdpaAttention( - (q_proj): Linear(in_features=2048, out_features=2048, bias=True) - (k_proj): Linear(in_features=2048, out_features=2048, bias=True) - (v_proj): Linear(in_features=2048, out_features=2048, bias=True) - (o_proj): Linear(in_features=2048, out_features=2048, bias=False) + (q_proj): Linear(in_features=3584, out_features=3584, bias=True) + (k_proj): Linear(in_features=3584, out_features=512, bias=True) + (v_proj): Linear(in_features=3584, out_features=512, bias=True) + (o_proj): Linear(in_features=3584, out_features=3584, bias=False) (rotary_emb): Qwen2MoeRotaryEmbedding() ) (mlp): Qwen2MoeSparseMoeBlock( - (gate): Linear(in_features=2048, out_features=60, bias=False) + (gate): Linear(in_features=3584, out_features=64, bias=False) (experts): ModuleList( - (0-59): 60 x Qwen2MoeMLP( - (gate_proj): Linear(in_features=2048, out_features=1408, bias=False) - (up_proj): Linear(in_features=2048, out_features=1408, bias=False) - (down_proj): Linear(in_features=1408, out_features=2048, bias=False) + (0-63): 64 x Qwen2MoeMLP( + (gate_proj): Linear(in_features=3584, out_features=2560, bias=False) + (up_proj): Linear(in_features=3584, out_features=2560, bias=False) + (down_proj): Linear(in_features=2560, out_features=3584, bias=False) (act_fn): SiLU() ) ) (shared_expert): Qwen2MoeMLP( - (gate_proj): Linear(in_features=2048, out_features=5632, bias=False) - (up_proj): Linear(in_features=2048, out_features=5632, bias=False) - (down_proj): Linear(in_features=5632, out_features=2048, bias=False) + (gate_proj): Linear(in_features=3584, out_features=20480, bias=False) + (up_proj): Linear(in_features=3584, out_features=20480, bias=False) + (down_proj): Linear(in_features=20480, out_features=3584, bias=False) (act_fn): SiLU() ) - (shared_expert_gate): Linear(in_features=2048, out_features=1, bias=False) + (shared_expert_gate): Linear(in_features=3584, out_features=1, bias=False) ) - (input_layernorm): Qwen2MoeRMSNorm() - (post_attention_layernorm): Qwen2MoeRMSNorm() + (input_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06) + (post_attention_layernorm): Qwen2MoeRMSNorm((3584,), eps=1e-06) ) ) - (norm): Qwen2MoeRMSNorm() + (norm): Qwen2MoeRMSNorm((3584,), eps=1e-06) ) - (lm_head): Linear(in_features=2048, out_features=151936, bias=False) + (lm_head): Linear(in_features=3584, out_features=151936, bias=False) ) ''' diff --git a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py index 7cddbf978369..c7841b24e5fc 100644 --- a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py +++ b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py @@ -73,7 +73,7 @@ def n_heads(self) -> int: @property def intermediate_dim(self) -> int: - return self._config.intermediate_size + return self._config.shared_expert_intermediate_size @property def n_heads_kv(self) -> int: diff --git a/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py b/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py index bd90cbd5d697..a9b01d1233cd 100644 --- a/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py +++ b/deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py @@ -42,7 +42,7 @@ def supports_config(config: DSMoEConfig) -> bool: if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16: return False - if config.top_k != 1 and config.top_k != 2 and config.top_k != 4: + if config.top_k != 1 and config.top_k != 2 and config.top_k != 4 and config.top_k != 8: return False return True diff --git a/deepspeed/inference/v2/ragged/manager_configs.py b/deepspeed/inference/v2/ragged/manager_configs.py index a5e98e5bcef1..17283b8bc0c4 100644 --- a/deepspeed/inference/v2/ragged/manager_configs.py +++ b/deepspeed/inference/v2/ragged/manager_configs.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Tuple -from deepspeed.pydantic_v1 import PositiveInt, validator +from pydantic import PositiveInt, model_validator from deepspeed.runtime.config_utils import DeepSpeedConfigModel from ..inference_utils import DtypeEnum @@ -173,11 +173,9 @@ class DSStateManagerConfig(DeepSpeedConfigModel): Enable tracking for offloading KV-cache to host memory. Currently unsupported. """ - @validator("max_ragged_sequence_count") - def max_ragged_sequence_count_validator(cls, v: int, values: dict): + @model_validator(mode="after") + def max_ragged_sequence_count_validator(self): # If the attributes below failed their validation they won't appear in the values dict. - if "max_tracked_sequences" in values and v > values["max_tracked_sequences"]: - raise ValueError("max_ragged_sequence_count must be less than max_tracked_sequences") - if "max_ragged_batch_size" in values and v > values["max_ragged_batch_size"]: - raise ValueError("max_ragged_sequence_count must be less than max_ragged_batch_size") - return v + assert self.max_ragged_sequence_count <= self.max_tracked_sequences, "max_ragged_sequence_count must be less than max_tracked_sequences" + assert self.max_ragged_sequence_count <= self.max_ragged_batch_size, "max_ragged_sequence_count must be less than max_ragged_batch_size" + return self diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py index c09a11e213db..340bc82de508 100644 --- a/deepspeed/moe/sharded_moe.py +++ b/deepspeed/moe/sharded_moe.py @@ -208,7 +208,7 @@ def top1gating(logits: Tensor, mask1 = einsum("s,se->se", used_token, mask1) # gating decisions - exp_counts = torch.sum(mask1, dim=0).detach().to('cpu') + exp_counts = torch.sum(mask1, dim=0).detach().to(logits.device) # if we don't want to drop any tokens if not drop_tokens: @@ -324,7 +324,7 @@ def top2gating(logits: Tensor, l_aux = torch.mean(me * ce) * num_experts * num_experts # gating decisions - exp_counts = torch.sum(mask1 + mask2, dim=0) + exp_counts = torch.sum(mask1 + mask2, dim=0).detach().to(logits.device) if drop_tokens: # Calculate configured capacity and remove locations outside capacity from mask @@ -368,7 +368,7 @@ def top2gating(logits: Tensor, combine_weights = combine1_sec + combine2_sec dispatch_mask = combine_weights.bool() - return l_aux, combine_weights, dispatch_mask, exp_counts.detach().to('cpu') + return l_aux, combine_weights, dispatch_mask, exp_counts def topkgating( diff --git a/deepspeed/monitor/config.py b/deepspeed/monitor/config.py index c4200877089a..960ce1ba997a 100644 --- a/deepspeed/monitor/config.py +++ b/deepspeed/monitor/config.py @@ -5,7 +5,7 @@ from typing import Optional -from deepspeed.pydantic_v1 import root_validator +from pydantic import model_validator from deepspeed.runtime.config_utils import DeepSpeedConfigModel @@ -36,10 +36,10 @@ class WandbConfig(DeepSpeedConfigModel): enabled: bool = False """ Whether logging to WandB is enabled. Requires `wandb` package is installed. """ - group: str = None + group: Optional[str] = None """ Name for the WandB group. This can be used to group together runs. """ - team: str = None + team: Optional[str] = None """ Name for the WandB team. """ project: str = "deepspeed" @@ -137,8 +137,8 @@ class DeepSpeedMonitorConfig(DeepSpeedConfigModel): csv_monitor: CSVConfig = {} """ Local CSV output of monitoring data. """ - @root_validator - def check_enabled(cls, values): - values["enabled"] = values.get("tensorboard").enabled or values.get("wandb").enabled or values.get( - "csv_monitor").enabled or values.get("comet").enabled - return values + @model_validator(mode="after") + def check_enabled(self): + enabled = self.tensorboard.enabled or self.wandb.enabled or self.csv_monitor.enabled or self.comet.enabled + self.__dict__["enabled"] = enabled + return self diff --git a/deepspeed/pydantic_v1.py b/deepspeed/pydantic_v1.py deleted file mode 100644 index 6aba072ad929..000000000000 --- a/deepspeed/pydantic_v1.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""Pydantic v1 compatibility module. - -Pydantic v2 introduced breaking changes that hinder its adoption: -https://docs.pydantic.dev/latest/migration/. To provide deepspeed users the option to -migrate to pydantic v2 on their own timeline, deepspeed uses this compatibility module -as a pydantic-version-agnostic alias for pydantic's v1 API. -""" - -try: - from pydantic.v1 import * # noqa: F401 -except ImportError: - from pydantic import * # noqa: F401 diff --git a/deepspeed/runtime/config.py b/deepspeed/runtime/config.py index b49b4a8b6086..8be2f7ac4055 100755 --- a/deepspeed/runtime/config.py +++ b/deepspeed/runtime/config.py @@ -705,7 +705,7 @@ def write_config(self, filename): class DeepSpeedConfig(object): - def __init__(self, config: Union[str, dict], mpu=None): + def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None): super(DeepSpeedConfig, self).__init__() if isinstance(config, dict): self._param_dict = config @@ -721,14 +721,16 @@ def __init__(self, config: Union[str, dict], mpu=None): ) try: self.global_rank = dist.get_rank() - if mpu is None: - self.world_size = dist.get_world_size() - else: + if mpu is not None: self.world_size = mpu.get_data_parallel_world_size() + elif mesh_device is not None: + self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) + else: + self.world_size = dist.get_world_size() except: self.global_rank = 0 self.world_size = 1 - + logger.info(f"Config mesh_device {mesh_device} world_size = {self.world_size}") # If elastic-mode enabled, update compute + update _param_dict self.elasticity_enabled = elasticity_enabled(self._param_dict) if self.elasticity_enabled: diff --git a/deepspeed/runtime/config_utils.py b/deepspeed/runtime/config_utils.py index 5522a8e79d69..d5c3a1548360 100755 --- a/deepspeed/runtime/config_utils.py +++ b/deepspeed/runtime/config_utils.py @@ -5,11 +5,12 @@ """ Collection of DeepSpeed configuration utilities """ -import json import collections -import collections.abc +import json +import torch from functools import reduce -from deepspeed.pydantic_v1 import BaseModel +from pydantic import BaseModel, ConfigDict, field_serializer + from deepspeed.utils import logger @@ -54,67 +55,73 @@ def __init__(self, strict=False, **data): if (not strict): # This is temporary until we refactor all DS configs, allows HF to load models data = {k: v for k, v in data.items() if (v != "auto" or k == "replace_method")} super().__init__(**data) - self._deprecated_fields_check(self) + self._deprecated_fields_check() - def _process_deprecated_field(self, pydantic_config, field): + def _process_deprecated_field(self, dep_field): # Get information about the deprecated field - fields_set = pydantic_config.__fields_set__ - dep_param = field.name - kwargs = field.field_info.extra + pydantic_config = self + fields_set = pydantic_config.model_fields_set + kwargs = pydantic_config.model_fields[dep_field].json_schema_extra new_param_fn = kwargs.get("new_param_fn", lambda x: x) - param_value = new_param_fn(getattr(pydantic_config, dep_param)) - new_param = kwargs.get("new_param", "") + param_value = new_param_fn(getattr(pydantic_config, dep_field)) + new_field = kwargs.get("new_param", "") dep_msg = kwargs.get("deprecated_msg", "") - if dep_param in fields_set: - logger.warning(f"Config parameter {dep_param} is deprecated" + - (f" use {new_param} instead" if new_param else "") + (f". {dep_msg}" if dep_msg else "")) + if dep_field in fields_set: + logger.warning(f"Config parameter {dep_field} is deprecated" + + (f" use {new_field} instead" if new_field else "") + (f". {dep_msg}" if dep_msg else "")) # Check if there is a new param and if it should be set with a value - if new_param and kwargs.get("set_new_param", True): + if new_field and kwargs.get("set_new_param", True): # Remove the deprecate field if there is a replacing field try: - delattr(pydantic_config, dep_param) + delattr(pydantic_config, dep_field) except Exception as e: - logger.error(f"Tried removing deprecated '{dep_param}' from config") + logger.error(f"Tried removing deprecated '{dep_field}' from config") raise e # Set new param value - new_param_nested = new_param.split(".") + new_param_nested = new_field.split(".") if len(new_param_nested) > 1: # If the new param exists in a subconfig, we need to get # the fields set for that subconfig pydantic_config = reduce(getattr, new_param_nested[:-1], pydantic_config) - fields_set = pydantic_config.__fields_set__ + fields_set = pydantic_config.model_fields_set new_param_name = new_param_nested[-1] assert ( new_param_name not in fields_set - ), f"Cannot provide deprecated parameter '{dep_param}' and replacing parameter '{new_param}' together" + ), f"Cannot provide deprecated parameter '{dep_field}' and replacing parameter '{new_field}' together" # A custom function for converting the old param value to new param value can be provided try: setattr(pydantic_config, new_param_name, param_value) except Exception as e: - logger.error(f"Tried setting value for '{new_param}' with value from deprecated '{dep_param}'") + logger.error(f"Tried setting value for '{new_field}' with value from deprecated '{dep_field}'") raise e - def _deprecated_fields_check(self, pydantic_config): - fields = pydantic_config.__fields__ - for field in fields.values(): - if field.field_info.extra.get("deprecated", False): - self._process_deprecated_field(pydantic_config, field) + def _deprecated_fields_check(self): + fields = self.model_fields + for field_name, field_info in fields.items(): + if field_info.json_schema_extra and field_info.json_schema_extra.get("deprecated", False): + self._process_deprecated_field(field_name) + + model_config = ConfigDict( + validate_default=True, + validate_assignment=True, + use_enum_values=True, + populate_by_name=True, + extra="forbid", + arbitrary_types_allowed=True, + protected_namespaces=(), + ) - class Config: - validate_all = True - validate_assignment = True - use_enum_values = True - allow_population_by_field_name = True - extra = "forbid" - arbitrary_types_allowed = True + @field_serializer("dtype", check_fields=False) + def serialize_torch_dtype(dtype: torch.dtype) -> str: + return str(dtype) def get_config_default(config, field_name): - assert field_name in config.__fields__, f"'{field_name}' is not a field in {config}" - assert not config.__fields__.get( - field_name).required, f"'{field_name}' is a required field and does not have a default value" - return config.__fields__.get(field_name).default + assert field_name in config.model_fields, f"'{field_name}' is not a field in {config}" + assert not config.model_fields.get( + field_name).is_required(), f"'{field_name}' is a required field and does not have a default value" + return config.model_fields.get(field_name).get_default() class pp_int(int): diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py old mode 100644 new mode 100755 index 27d294b3ae01..61e6da2663cf --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -194,6 +194,7 @@ def __init__(self, collate_fn=None, config=None, config_class=None, + mesh_device=None, dont_change_device=False): super(DeepSpeedEngine, self).__init__() self.dont_change_device = dont_change_device @@ -233,10 +234,14 @@ def __init__(self, self._is_gradient_accumulation_boundary = None self.scale_wrt_gas = None self.losses = None + self.mesh_device = mesh_device # for debug purposes - can then debug print: debug_get_module_name(module) debug_extract_module_and_param_names(model) + if self.mesh_device: + groups.mesh_device = self.mesh_device + self._do_args_sanity_check(args) self._configure_with_arguments(args, mpu) self._do_sanity_check() @@ -615,6 +620,9 @@ def random_ltd_initialize(self): raise ValueError(f'not yet support') #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler) + def get_sequence_parallel_group(self): + return self.seq_parallel_group + def wall_clock_breakdown(self): return self._config.wall_clock_breakdown @@ -1187,6 +1195,7 @@ def _configure_distributed_model(self, model): self.sequence_parallel_size = groups._get_sequence_parallel_world_size() if self.sequence_parallel_size > 1: self.communication_data_type = self._config.seq_parallel_communication_data_type + self.seq_parallel_group = groups._get_sequence_parallel_group() if not (self.amp_enabled() or is_zero_init_model): self._broadcast_model() diff --git a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py index 120723fae5ab..26fbf6164d54 100644 --- a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py +++ b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py @@ -121,7 +121,7 @@ def _configure_aio(self, ds_config): dtype=self.dtype, device=get_accelerator().device_name(), requires_grad=False) - self.aio_read_handle.new_device_locked_tensor(self.buffers) + self.aio_read_handle.pin_device_tensor(self.buffers) else: self.buffers = get_accelerator().pin_memory(torch.empty(int(self.aligned_elements_per_buffer * self.param_buffer_count), diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py index 2089d59dbce4..1cfcd784e2ce 100644 --- a/deepspeed/runtime/zero/config.py +++ b/deepspeed/runtime/zero/config.py @@ -6,7 +6,7 @@ import sys from typing import Optional from enum import Enum -from deepspeed.pydantic_v1 import Field, validator, root_validator +from pydantic import Field, model_validator from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel from deepspeed.utils import logger from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum @@ -30,7 +30,7 @@ "reduce_bucket_size": 500000000, "load_from_fp32_weights": [true|false], "cpu_offload": [true|false] (deprecated), - "cpu_offload_params" : [true|false] (deprecated), + "cpu_offload_param" : [true|false] (deprecated), "cpu_offload_use_pin_memory": [true|false] (deprecated), "sub_group_size" : 1000000000000, "offload_param": {...}, @@ -128,7 +128,7 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): the allgather for large model sizes """ - overlap_comm: bool = None # None for dynamic default value (see validator `overlap_comm_valid` below) + overlap_comm: Optional[bool] = None # None for dynamic default value (see validator `overlap_comm_valid` below) """ Attempts to overlap the reduction of the gradients with backward computation """ @@ -168,27 +168,37 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): parameters). Used by ZeRO3-Offload and ZeRO-Infinity """ - cpu_offload_param: bool = Field( + cpu_offload_param: Optional[bool] = Field( None, - deprecated=True, - new_param="offload_param", - new_param_fn=(lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) if val else None), + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, ) """ Deprecated, please use ``offload_param`` """ - cpu_offload_use_pin_memory: bool = Field( + cpu_offload_use_pin_memory: Optional[bool] = Field( None, - deprecated=True, - new_param="offload_param or offload_optimizer", - set_new_param=False, + json_schema_extra={ + "deprecated": True, + "new_param": "offload_param or offload_optimizer", + "set_new_param": False + }, ) """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """ - cpu_offload: bool = Field( + cpu_offload: Optional[bool] = Field( None, - deprecated=True, - new_param="offload_optimizer", - new_param_fn=(lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu) if val else None), + json_schema_extra={ + "deprecated": + True, + "new_param": + "offload_optimizer", + "new_param_fn": (lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu) + if val else None) + }, ) """ Deprecated, please use ``offload_optimizer`` """ @@ -242,8 +252,10 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): """ stage3_gather_fp16_weights_on_model_save: bool = Field(False, - deprecated=True, - new_param="gather_16bit_weights_on_model_save") + json_schema_extra={ + "deprecated": True, + "new_param": "gather_16bit_weights_on_model_save" + }) """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """ ignore_unused_parameters: bool = True @@ -309,16 +321,15 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel): """ # Validators - @validator("overlap_comm") - def overlap_comm_valid(cls, field_value, values): - if field_value is None: - assert ("stage" in values), "DeepSpeedZeroConfig: 'stage' must be defined before 'overlap_comm'" - field_value = values["stage"] == ZeroStageEnum.weights - return field_value - - @root_validator - def offload_ratio_check(cls, values): - offload_config = getattr(values, "offload_optimizer", {}) + @model_validator(mode="after") + def overlap_comm_valid(self): + if self.overlap_comm is None: + self.overlap_comm = self.stage == ZeroStageEnum.weights + return self + + @model_validator(mode="after") + def offload_ratio_check(self): + offload_config = self.offload_optimizer if offload_config and offload_config.ratio < 1.0: - assert values.get("stage") == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3." - return values + assert self.stage == ZeroStageEnum.weights, "Partial offloading only supported for ZeRO Stage 3." + return self diff --git a/deepspeed/runtime/zero/linear.py b/deepspeed/runtime/zero/linear.py index e9dd78864cde..8c8db60768eb 100644 --- a/deepspeed/runtime/zero/linear.py +++ b/deepspeed/runtime/zero/linear.py @@ -16,6 +16,7 @@ #when implemented outside of torch.autograd.Function import math +import functools import torch from torch import Tensor @@ -33,8 +34,14 @@ def print_rank_0(message, debug=False, force=False): try: - autocast_custom_fwd = get_accelerator().amp().custom_fwd - autocast_custom_bwd = get_accelerator().amp().custom_bwd + # Fix `torch.[device].amp.custom_fwd/bwd` FutureWarning in torch 2.4 + if hasattr(torch, 'amp') and hasattr(torch.amp, 'custom_fwd') and hasattr(torch.amp, 'custom_bwd'): + autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=get_accelerator().device_name()) + autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=get_accelerator().device_name()) + else: + # original implementation + autocast_custom_fwd = get_accelerator().amp().custom_fwd + autocast_custom_bwd = get_accelerator().amp().custom_bwd except (ImportError, AttributeError) as exp: autocast_custom_fwd = noop_decorator autocast_custom_bwd = noop_decorator diff --git a/deepspeed/runtime/zero/offload_config.py b/deepspeed/runtime/zero/offload_config.py index b7adc13a0ea2..74a5673bc1bc 100644 --- a/deepspeed/runtime/zero/offload_config.py +++ b/deepspeed/runtime/zero/offload_config.py @@ -5,7 +5,9 @@ from enum import Enum from pathlib import Path -from deepspeed.pydantic_v1 import Field, validator +from pydantic import Field, model_validator +from typing import Optional + from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int @@ -25,7 +27,7 @@ class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel): `nvme`. """ - nvme_path: Path = None + nvme_path: Optional[Path] = None """ Filesystem path for NVMe device for parameter offloading. """ buffer_count: int = Field(5, ge=0) @@ -56,7 +58,7 @@ class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): `nvme`. Optimizer computation is offload to CPU regardless of device option. """ - nvme_path: Path = None + nvme_path: Optional[Path] = None """ Filesystem path for NVMe device for optimizer state offloading. """ buffer_count: int = Field(4, ge=0) @@ -88,10 +90,11 @@ class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): fast_init: bool = False """ Enable fast optimizer initialization when offloading to NVMe. """ - @validator("pipeline_read", "pipeline_write", always=True) - def set_pipeline(cls, field_value, values): - values["pipeline"] = field_value or values.get("pipeline", False) - return field_value - ratio: float = Field(1.0, ge=0.0, le=1.0) """ Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3.""" + + @model_validator(mode="after") + def set_pipeline(self): + pipeline = self.pipeline_read or self.pipeline_write + self.__dict__["pipeline"] = pipeline + return self diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index 57e80911d645..83cf996ca019 100755 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -725,8 +725,9 @@ def reduce_gradients(self, pipeline_parallel=False): def get_first_param_index(self, group_id, param_group, partition_id): for index, param in enumerate(param_group): param_id = self.get_param_id(param) - if partition_id in self.param_to_partition_ids[group_id][param_id]: - return index + if group_id in self.param_to_partition_ids and param_id in self.param_to_partition_ids[group_id]: + if partition_id in self.param_to_partition_ids[group_id][param_id]: + return index return None def initialize_gradient_partitioning_data_structures(self): diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py new file mode 100644 index 000000000000..baa7bc1ea7a8 --- /dev/null +++ b/deepspeed/sequence/cross_entropy.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +import deepspeed.comm as dist + + +class _VocabSequenceParallelCrossEntropy(torch.autograd.Function): + + @staticmethod + def forward(ctx, vocab_seq_parallel_logits, target, sp_group): + # vocab_seq_parallel_logits: [S/P, B, V] + # target: [S/P, B] + # return: [S, B] + + # Need softmax for backward + softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1) + ctx.vocab_size = vocab_seq_parallel_logits.size(2) + loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none') + + sp_world_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + ctx.sp_world_size = sp_world_size + ctx.sp_rank = sp_rank + ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size + batch_size = vocab_seq_parallel_logits.size(1) + + loss_all = torch.empty(ctx.seqlen, + batch_size, + dtype=vocab_seq_parallel_logits.dtype, + device=vocab_seq_parallel_logits.device) + dist.all_gather_into_tensor(loss_all, loss, group=sp_group) + + ctx.save_for_backward(softmax, target) + + return loss_all + + @staticmethod + def backward(ctx, grad_output): + softmax, target = ctx.saved_tensors + + step_seqlen = ctx.seqlen // ctx.sp_world_size + sp_rank = ctx.sp_rank + grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :] + + grad_input = softmax + grad_2d = grad_input.view(-1, ctx.vocab_size) + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + + grad_2d[arange_1d, target.view(-1)] -= 1 + grad_input.mul_(grad_output_part.unsqueeze(dim=-1)) + + return grad_input, None, None, None + + +def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group): + return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group) diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py old mode 100644 new mode 100755 index c49f4520e16e..9dd288ef46db --- a/deepspeed/utils/groups.py +++ b/deepspeed/utils/groups.py @@ -30,6 +30,7 @@ from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size from deepspeed.utils.exceptions import DeprecatedException from deepspeed.accelerator import get_accelerator + # Expert parallel group that the current rank belongs to. _EXPERT_PARALLEL_GROUP = {} # Expert data parallel group that the current rank belongs to. @@ -47,6 +48,8 @@ _DATA_PARALLEL_GROUP = None +mesh_device = None + # Deprecated groups initialize function. def initialize(ep_size=1, mpu=None): @@ -398,8 +401,11 @@ def _get_data_parallel_group(): """Get the data parallel group the caller rank belongs to.""" assert dist.is_initialized(), 'dist is not initialized' global mpu + if mesh_device is not None: + return mesh_device.get_group(mesh_dim="data_parallel") if mpu is not None: return mpu.get_data_parallel_group() + # Return the clone of dist world group return _clone_world_group() @@ -442,6 +448,8 @@ def _get_expert_data_parallel_rank(group_name): def _get_data_parallel_world_size(): """Return world size for the data parallel group.""" + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) global mpu if mpu is not None: return mpu.get_data_parallel_world_size() @@ -464,6 +472,8 @@ def _get_data_parallel_rank(): def _get_sequence_parallel_world_size(): """Return world size for the model parallel group.""" global mpu + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel")) if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'): return mpu.get_sequence_parallel_world_size() return 1 @@ -479,9 +489,11 @@ def _get_sequence_parallel_rank(): def _get_sequence_parallel_group(): global mpu - if mpu is not None and hasattr(mpu, 'get_sequence_parallel_group'): - return mpu.get_sequence_parallel_group() - return None + if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'): + if mesh_device is None: + raise KeyError("No sequence parallel group found") + return mesh_device.get_group(mesh_dim="sequence_parallel") + return mpu.get_sequence_parallel_group() def _get_sequence_data_parallel_world_size(): diff --git a/docs/index.md b/docs/index.md index 127c7226e6d4..1efdcea132d2 100755 --- a/docs/index.md +++ b/docs/index.md @@ -7,11 +7,12 @@ title: "Latest News" --- DeepSpeed empowers ChatGPT-like model training with a single click, offering 15x speedup over SOTA RLHF systems with unprecedented cost reduction at all scales; [learn how](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat). +* [2024/08] [DeepSpeed on Windows](https://github.com/microsoft/DeepSpeed/blob/master/blogs/windows/08-2024/README.md)[[日本語](https://github.com/microsoft/DeepSpeed/blob/master/blogs/windows/08-2024/japanese/README.md)] [[中文](https://github.com/microsoft/DeepSpeed/blob/master/blogs/windows/08-2024/chinese/README.md)] + * [2024/08] [DeepNVMe: Improving DL Applications through I/O Optimizations](https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-gds/README.md)[[日本語](https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-gds/japanese/README.md)] [[中文](https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-gds/chinese/README.md)] * [2024/07] [DeepSpeed Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-ucp/README.md)[[日本語](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-ucp/japanese/README.md)] * [2024/03] [DeepSpeed-FP6: The Power of FP6-Centric Serving for Large Language Models](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024/README.md) [[English](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024/README.md)] [[中文](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fp6/03-05-2024/README-Chinese.md)] * [2024/01] [DeepSpeed-FastGen: Introducting Mixtral, Phi-2, and Falcon support with major performance and feature enhancements.](https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fastgen/2024-01-19) -* [2023/11] [Llama 2 Inference on 4th Gen Intel® Xeon® Scalable Processor with DeepSpeed](https://github.com/microsoft/DeepSpeed/tree/master/blogs/intel-inference) [[Intel version]](https://www.intel.com/content/www/us/en/developer/articles/technical/xllama-2-on-xeon-scalable-processor-with-deepspeed.html) @@ -19,12 +20,12 @@ title: "Latest News" More news diff --git a/op_builder/evoformer_attn.py b/op_builder/evoformer_attn.py index af3aa7429775..7f68ccf87290 100644 --- a/op_builder/evoformer_attn.py +++ b/op_builder/evoformer_attn.py @@ -52,11 +52,27 @@ def is_compatible(self, verbose=False): if verbose: self.warning("Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH") return False - with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f: - if '3.1.0' not in f.read(): + if os.path.exists(f'{self.cutlass_path}/CHANGELOG.md'): + with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f: + if '3.1.0' not in f.read(): + if verbose: + self.warning("Please use CUTLASS version >= 3.1.0") + return False + else: + # pip install nvidia-cutlass package + try: + import cutlass + except ImportError: + if verbose: + self.warning("Please pip install nvidia-cutlass if trying to pre-compile kernels") + return False + cutlass_major, cutlass_minor = cutlass.__version__.split('.')[:2] + cutlass_compatible = (int(cutlass_major) >= 3 and int(cutlass_minor) >= 1) + if not cutlass_compatible: if verbose: self.warning("Please use CUTLASS version >= 3.1.0") return False + cuda_okay = True if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda sys_cuda_major, _ = installed_cuda_version() diff --git a/op_builder/gds.py b/op_builder/gds.py index e024674e01d8..01c2d5a245d1 100644 --- a/op_builder/gds.py +++ b/op_builder/gds.py @@ -36,7 +36,13 @@ def extra_ldflags(self): return super().extra_ldflags() + ['-lcufile'] def is_compatible(self, verbose=False): - import torch.utils.cpp_extension + try: + import torch.utils.cpp_extension + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile GDS") + return False + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") gds_compatible = self.has_function(funcname="cuFileDriverOpen", diff --git a/op_builder/xpu/inference.py b/op_builder/xpu/inference.py index 9114dcc2c315..a9ac4f84c2ca 100644 --- a/op_builder/xpu/inference.py +++ b/op_builder/xpu/inference.py @@ -30,7 +30,10 @@ def cxx_args(self): def load(self): try: - import intel_extension_for_pytorch.deepspeed - return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + import intel_extension_for_pytorch + if hasattr(intel_extension_for_pytorch, "deepspeed"): + return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + else: + return intel_extension_for_pytorch.xpu.deepspeed except ImportError: raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.") diff --git a/requirements/requirements-readthedocs.txt b/requirements/requirements-readthedocs.txt index 1a2ad18611e7..a48a47e4428d 100644 --- a/requirements/requirements-readthedocs.txt +++ b/requirements/requirements-readthedocs.txt @@ -1,10 +1,10 @@ -autodoc_pydantic +autodoc_pydantic>=2.0.0 docutils<0.18 hjson packaging psutil py-cpuinfo -pydantic<2.0.0 +pydantic>=2.0.0 recommonmark sphinx_rtd_theme torch diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 6840d6dbcc98..70c94a745435 100755 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -4,6 +4,6 @@ numpy packaging>=20.0 psutil py-cpuinfo -pydantic +pydantic>=2.0.0 torch tqdm diff --git a/tests/unit/inference/v2/ragged/test_manager_configs.py b/tests/unit/inference/v2/ragged/test_manager_configs.py index a5f270cced8c..bdd513445ddb 100644 --- a/tests/unit/inference/v2/ragged/test_manager_configs.py +++ b/tests/unit/inference/v2/ragged/test_manager_configs.py @@ -5,7 +5,7 @@ import pytest -from deepspeed.pydantic_v1 import ValidationError +from pydantic import ValidationError from deepspeed.inference.v2.ragged import DSStateManagerConfig diff --git a/tests/unit/runtime/test_ds_config_dict.py b/tests/unit/runtime/test_ds_config_dict.py index c11c63d04867..d06b35e208fe 100644 --- a/tests/unit/runtime/test_ds_config_dict.py +++ b/tests/unit/runtime/test_ds_config_dict.py @@ -67,13 +67,11 @@ def _batch_assert(status, ds_config, batch, micro_batch, gas, success): if not success: assert not status - print("Failed but All is well") return assert ds_config.train_batch_size == batch assert ds_config.train_micro_batch_size_per_gpu == micro_batch assert ds_config.gradient_accumulation_steps == gas - print("All is well") #Tests different batch config provided in deepspeed json file diff --git a/tests/unit/runtime/test_ds_config_model.py b/tests/unit/runtime/test_ds_config_model.py index 87ea747cf423..4d184b2858a8 100644 --- a/tests/unit/runtime/test_ds_config_model.py +++ b/tests/unit/runtime/test_ds_config_model.py @@ -4,18 +4,25 @@ # DeepSpeed Team import pytest -import os import json -from typing import List -from deepspeed.pydantic_v1 import Field, ValidationError +import os +from typing import List, Optional + +from pydantic import Field, ValidationError + from deepspeed.runtime import config as ds_config from deepspeed.runtime.config_utils import DeepSpeedConfigModel class SimpleConf(DeepSpeedConfigModel): param_1: int = 0 - param_2_old: str = Field(None, deprecated=True, new_param="param_2", new_param_fn=(lambda x: [x])) - param_2: List[str] = None + param_2_old: Optional[str] = Field(None, + json_schema_extra={ + "deprecated": True, + "new_param": "param_2", + "new_param_fn": (lambda x: [x]) + }) + param_2: Optional[List[str]] = None param_3: int = Field(0, alias="param_3_alias") diff --git a/tests/unit/sequence_parallelism/test_ulysses.py b/tests/unit/sequence_parallelism/test_ulysses.py new file mode 100644 index 000000000000..915c89e0b00a --- /dev/null +++ b/tests/unit/sequence_parallelism/test_ulysses.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import pytest +import torch +import deepspeed.comm as dist +from deepspeed import initialize +from transformers import AutoModel +from unit.common import DistributedTest +from deepspeed.sequence.layer import _SeqAllToAll +from unit.util import skip_on_arch + + +#Use mesh device to create data and sequence parallel group +class TestUlyssesUtils(DistributedTest): + world_size = 4 + + def test_mesh_device_creation(self) -> None: + skip_on_arch(min_arch=8) + model = AutoModel.from_pretrained('bert-base-uncased') + sp_size = 2 + dp_size = 2 + ds_engine, _, _, _ = initialize( + model=model, + config_params={ + "train_batch_size": 8, + "data_parallel_size": dp_size, + "sequence_parallel_size": sp_size + }, + ) + assert ds_engine.seq_parallel_group is not None + assert ds_engine.data_parallel_group is not None + assert dist.get_world_size(group=ds_engine.seq_parallel_group) == sp_size + assert dist.get_world_size(group=ds_engine.data_parallel_group) == dp_size + assert dist.get_world_size() == sp_size * dp_size + + +#Sweep b,s,h,d to test all2all consistency +@pytest.mark.parametrize("d0", [2, 4]) #batch or sequence dimension +@pytest.mark.parametrize("d1", [4, 8]) #batch or sequence dimension +@pytest.mark.parametrize("num_heads", [4, 8]) +@pytest.mark.parametrize("head_dim", [16, 32]) +class TestUlyssesAll2All(DistributedTest): + world_size = 4 + + def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_heads: int) -> None: + skip_on_arch(min_arch=8) + model = AutoModel.from_pretrained('bert-base-uncased') + ds_engine, _, _, _ = initialize(model=model, config_params={"train_batch_size": 8}, mesh_param=(2, 2)) + #4D tensor : b,s,h,d or s,b,h,d + input_tensor = torch.randn(d0, d1, num_heads, head_dim, device=ds_engine.device) + scatter_idx = 2 + batch_dim_idx = 0 + outputs = [] + seq_dims = [0] #seq first API + #TODO: Add support for batch first (that seq_dims=[0,1]) after PR for bs>1 issue with batch first is fixed + ## See discussion in : https://github.com/microsoft/DeepSpeed/issues/5808 + for seq_dim in seq_dims: + gather_idx = seq_dim + #first all2all: sequence parallel to head parallel + s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx, + batch_dim_idx) + + #No op + # second all2all: head parallel to sequence parallel + h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx, + batch_dim_idx) + print( + f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}' + ) + outputs.append(h2s_tensor) + + # Check outputs are the same as input + for i in range(1, len(outputs)): + assert torch.allclose(input_tensor, outputs[i]), f"Outputs differ for sequence dim {seq_dims[i]}" diff --git a/version.txt b/version.txt index 226468ee5b2e..e815b861f023 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.14.6 +0.15.1