张量并行性 - 亚马逊 SageMaker AI

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

张量并行性

张量并行性是模型并行性的一种,它在设备之间拆分特定的模型权重、梯度和优化器状态。管道并行性保持单个权重不变,但将权重、梯度或优化器集分割到不同设备上,而张量并行则将单个权重分割开来。这通常涉及对模型的特定运算、模块或层进行分布式计算。

在单个参数占用大部分 GPU 内存的情况下(例如词汇表很大的大型嵌入表,或者具有大量类的大型 softmax 层),则需要张量并行性。在这种情况下,将这种大张量或运算视为原子单元的效率会很低,并且会阻碍内存负载的平衡。

SMP v2 与 Transform er Engin e 集成以实现张量并行性,并在 FSDP 之上运行。 PyTorch APIs您可以同时启用 PyTorch FSDP 和 SMP 张量并行度,并确定最佳模型并行度以获得最佳性能。

在实践中,张量并行在以下场景下尤其有用。

  • 当在较长的上下文长度下训练时,仅使用 FSDP 就会导致较高的激活内存。

  • 当使用真正的大集群进行训练时,其全局批次大小会超过所需的限制。

Hugging Face 转换器模型兼容 SMP 张量并行

SMP v2 目前可为以下 Hugging Face 转换器模型提供张量并行支持。

有关在这些模型上应用张量并行的参考配置,请参阅 配置提示

配置张量并行

对于 tensor_parallel_degree,您可以为张量并行选择一个值。该值必须平均除以集群 GPUs 中的数量。例如,要在使用带有 8 的实例时对模型进行分片 GPUs,请选择 2、4 或 8。我们建议您从一个较小的数字开始,然后逐渐增加,直到模型适合 GPU 内存。

以下代码片段显示了如何在训练脚本中添加 SMP 初始化模块 torch.sagemaker.init(),并按照 使用 SageMaker 模型并行度库 v2 中介绍的两步流程,为训练作业启动器设置 JSON 格式的 SMP 配置字典。您无需对 PyTorch 模型或 PyTorch FSDP 配置进行任何更改。有关 tensor_parallel_degreerandom_seed 参数的更多信息,请参阅 SMP v2 核心功能配置参数

SMP 配置

{ "tensor_parallel_degree": 8, "random_seed": 0 }

在您的训练脚本中

使用 torch.sagemaker.init() 初始化以激活 SMP v2,然后使用 torch.sagemaker.transform API 封装模型。

import torch.sagemaker as tsm tsm.init() from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_config(..) model = tsm.transform(model)

保存和加载 Hugging Face 转换器检查点

SMP 库转换模型后,会更改模型的状态字典 (state_dict)。这意味着此模型与最初的 Hugging Face 转换器检查点功能不兼容。为了处理这个问题,SMP 库提供了 APIs 以 Hugging Face Transformer 表示形式保存变换模型中的检查点,以及用于加载 Hugging Face Transformer 模型检查点以进行微调的 torch.sagemaker.transform API。

有关在使用 SMP v2 张量并行功能时保存检查点的更多信息,请参阅 使用 SMP 的检查点

有关应用 SMP v2 的张量并行功能对模型进行微调的更多信息,请参阅 微调