DeepSpeed 加速器抽象接口

目录

介绍

DeepSpeed加速器抽象允许用户在各种深度学习加速硬件上无缝运行大型语言模型。它提供了一组加速器运行时和加速器操作构建器接口，这些接口可以为不同的硬件实现。这意味着用户可以编写大型语言模型代码，而无需编写特定硬件的代码。通过DeepSpeed加速器抽象，同一个大型语言模型可以在不同的硬件平台上运行，而无需重写模型代码。这使得在不同硬件上运行大型语言模型变得更加容易。

本文档涵盖了与DeepSpeed加速器抽象接口相关的三个主题：

  1. 使用DeepSpeed加速器抽象接口编写与加速器无关的模型。
  2. 在不同的加速器上运行DeepSpeed模型。
  3. 为DeepSpeed加速器抽象接口实现新的加速器扩展。

编写与加速器无关的模型

在这一部分，你将学习如何编写一个不包含硬件特定代码的模型，或者如何将一个只能在特定硬件上运行的模型移植为与加速器无关的模型。为此，我们首先从deepspeed.accelerator导入get_accelerator

from deepspeed.accelerator import get_accelerator

注意：get_accelerator() 是 DeepSpeed 加速器抽象接口的入口

端口加速器运行时调用

首先，我们需要移植加速器运行时调用。在CUDA设备上，加速器运行时调用以torch.cuda.(...)的形式出现。通过DeepSpeed加速器抽象接口，这样的加速器运行时调用可以写成get_accelerator().(...)的形式，这将与加速器无关。

一个典型的转换如下例所示：

if torch.cuda.is_available():
    ...

–>

if get_accelerator().is_available():
    ...

对于大多数torch.cuda.(...)调用，我们可以直接将torch.cuda替换为get_accelerator()。然而，有一些例外需要注意：

  1. 对于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()
  2. 对于 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模型

Accelerator Setup Guide 提供了如何为DeepSpeed设置不同加速器的指南。它还附带了一个简单的示例，展示了如何为不同的加速器运行deepspeed。提供了以下指南：

  1. 在CPU上运行DeepSpeed模型
  2. 在XPU上运行DeepSpeed模型
  3. 在华为昇腾NPU上运行DeepSpeed模型

实现新的加速器扩展

可以实现一个新的DeepSpeed加速器扩展来支持DeepSpeed中的新加速器。一个可以参考的例子是Intel Extension For DeepSpeed。一个加速器扩展包含以下组件：

  1. XYZ_Accelerator(DeepSpeedAccelerator) 类定义，其中‘XYZ’是加速器名称，例如‘XPU’或‘CPU’。 该类实现了 class DeepSpeedAccelerator 并将由 DeepSpeed 中的 get_accelerator() 返回。
  2. 遵循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），并继承这个基础操作构建器。
  3. 操作内核如下 链接。

请注意，扩展不需要一次性实现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 中添加隐式检测代码。

更新: