DeepSpeed 加速器抽象接口

目录

简介

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.<interface>(...) 的形式出现。使用 DeepSpeed 加速器抽象接口，此类加速器运行时调用可以写成 get_accelerator().<interface>(...) 的形式，这将是加速器无关的。

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

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

–>

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

对于大多数 torch.cuda.<interface>(...) 调用，我们可以直接用 get_accelerator() 替换 torch.cuda。但是，有一些例外需要引起注意

  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 模型

加速器设置指南 提供了有关如何为 DeepSpeed 设置不同加速器的指南。它还附带了在不同加速器上运行 deepspeed 的简单示例。提供以下指南

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

实现新的加速器扩展

可以实现新的 DeepSpeed 加速器扩展来支持 DeepSpeed 中的新加速器。一个可供参考的示例是 英特尔 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。一种常见的做法是实现一个基础操作构建器（在英特尔 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 中添加隐式检测代码。

更新时间：