DeepSpeed 加速器抽象接口
目录
简介
DeepSpeed 加速器抽象允许用户使用 DeepSpeed 在各种深度学习加速硬件上无缝运行大型语言模型。它提供了一组加速器运行时和加速器运算构建器接口,可以为不同的硬件实现。这意味着用户可以编写大型语言模型代码,而无需硬件特定代码。使用 DeepSpeed 加速器抽象,相同的大型语言模型可以在不同的硬件平台上运行,而无需重写模型代码。这使得在不同的硬件上运行大型语言模型变得更加容易。
本文档涵盖了与 DeepSpeed 加速器抽象接口相关的三个主题
- 使用 DeepSpeed 加速器抽象接口编写加速器无关模型。
- 在不同的加速器上运行 DeepSpeed 模型。
- 为 DeepSpeed 加速器抽象接口实现新的加速器扩展。
编写加速器无关模型
在本部分中,您将学习如何编写不包含硬件特定代码的模型,或者如何移植仅在特定硬件上运行的模型,使其成为加速器无关的模型。为此,我们首先从 deepspeed.accelerator
中导入 get_accelerator
from deepspeed.accelerator import get_accelerator
注意:get_accelerator()
是 DeepSpeed 加速器抽象接口的入口
移植加速器运行时调用
首先,我们需要移植加速器运行时调用。在 CUDA 设备上,加速器运行时调用以 torch.cuda.<interface>(...)
的形式出现。使用 DeepSpeed 加速器抽象接口,此类加速器运行时调用可以以 get_accelerator().<interface>(...)
的形式编写,它将成为加速器无关的。
典型的转换如下例所示
if torch.cuda.is_available():
...
–>
if get_accelerator().is_available():
...
对于大多数 torch.cuda.<interface>(...)
调用,我们实际上可以将 torch.cuda
替换为 get_accelerator()
。但是,有一些例外情况需要注意
- 对于
torch.cuda.current_device()
,我们需要知道是否调用此接口是为了获取设备索引,还是将返回值作为设备提供。如果我们想将返回值用作设备字符串,我们需要调用get_accelerator().current_device_name()
。例如torch.empty(weight_shape, dtype=dtype, device=get_accelerator().current_device_name())
但是,如果我们希望获取设备索引作为数字,则应该调用
get_accelerator().current_device()
local_rank = get_accelerator().current_device()
- 对于
torch.cuda.default_generators[index]
,转换为get_accelerator().default_generator(index)
移植加速器设备名称
对于 CUDA 特定的设备名称,例如 'cuda'
或 'cuda:0'
,或 'cuda:1'
,我们将它们转换为 get_accelerator().device_name()
,get_accelerator().device_name(0)
和 get_accelerator().device_name(1)
。
如果没有索引的设备名称,则可以在模型需要对特定加速器执行特定操作时使用。我们建议尽量减少此类用法,仅用于无法以其他方式解决的情况。
张量运算
CUDA 特定的张量运算需要根据以下规则进行转换
-
当我们将一个 torch 张量转换为加速器设备,例如
my_tensor.cuda()
,我们将使用my_tensor.to(get_accelerator().device_name())
-
当我们检查一个 torch 张量是否在加速器设备上,例如
my_tensor.is_cuda
,我们将使用get_accelerator().on_accelerator(my_tensor)
-
当将张量固定到 GPU 内存中,例如
my_tensor.pin_memory()
,我们将使用get_accelerator().pin_memory(my_tensor)
通信后端
当使用通信后端字符串时,接口 get_accelerator().communication_backend_name()
用于获取通信后端名称。所以,代替
torch.distributed.init_process_group('nccl')
,我们将使用
torch.distributed.init_process_group(get_accelerator().communication_backend_name())
在不同的加速器上运行 DeepSpeed 模型
加速器设置指南 提供了有关如何为 DeepSpeed 设置不同加速器的指南。它还附带了在不同加速器上运行 deepspeed 的简单示例。提供了以下指南
- 在 CPU 上运行 DeepSpeed 模型
- 在 XPU 上运行 DeepSpeed 模型
- 在华为昇腾 NPU 上运行 DeepSpeed 模型
实现新的加速器扩展
可以实现新的 DeepSpeed 加速器扩展来支持 DeepSpeed 中的新加速器。一个示例是 Intel Extension For DeepSpeed。加速器扩展包含以下组件
- XYZ_Accelerator(DeepSpeedAccelerator) 类定义,其中“XYZ”是加速器名称,例如“XPU”或“CPU”。此类实现
class DeepSpeedAccelerator
,并将由 DeepSpeed 中的get_accelerator()
返回。 - 运算构建器,遵循 https://github.com/intel/intel-extension-for-deepspeed/tree/main/intel_extension_for_deepspeed/op_builder。所有运算构建器都需要直接或间接继承
deepspeed.ops.op_builder.builder.OpBuilder
。一种常见的做法是实现一个基本运算构建器(在 Intel Extension for DeepSpeed 的情况下为 SYCLOpBuilder),并继承此基本运算构建器。 - 运算内核,如以下 链接 所示。
请注意,扩展不必一次性实现 https://github.com/microsoft/DeepSpeed/tree/master/op_builder 下的所有运算构建器。缺少运算构建器通常意味着某些 DeepSpeed 功能无法用于该加速器,但未使用该功能的模型仍然可以运行。
在为加速器扩展实现运算构建器时,需要注意的是,运算构建器本机代码是由 DeepSpeed jit 加载机制构建的。这意味着要构建的本机源文件需要在 DeepSpeed 安装目录中。但是,这些文件是在加速器扩展安装目录中定义的,DeepSpeed 无法直接构建这些文件。为了解决这个问题,请参考 https://github.com/intel/intel-extension-for-deepspeed/blob/main/intel_extension_for_deepspeed/op_builder/cpu_adam.py 中的示例使用“sycl_kernel_path”和“sycl_kernel_include”(用户可以在自己的加速器扩展中将“sycl”更改为其他前缀)来允许本机代码在 DeepSpeed jit 加载期间构建。
当加速器扩展安装在环境中时,可以通过以下两种方式使用它:显式调用 deepspeed.accelerator.set_accelerator(XYZ_Accelerator())
,参考 https://github.com/microsoft/DeepSpeed/blob/master/accelerator/real_accelerator.py 中的示例,或者在上面同一文件中的 get_accelerator
中添加隐式检测代码。