张量并行性 - Amazon SageMaker
Hugging Face Transformer 模型与张量并SMP行度兼容配置张量并行度保存和加载 Hugging Face Transformer 检查点

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

张量并行性

张量并行性是模型并行性的一种，它在设备之间拆分特定的模型权重、梯度和优化器状态。与流水线并行度形成鲜明对比，后者保持单个权重不变，但将权重、梯度或优化器集分成多个设备，而张量并行度则对单个权重进行分片。这通常涉及对模型的特定运算、模块或层进行分布式计算。

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

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

实际上，张量并行性在以下场景中特别有用。

  • 当在较长的上下文长度下训练时，FSDP单独使用会产生较高的激活记忆。

  • 当使用非常大的集群进行训练时，其全局批次大小超过了所需限制。

Hugging Face Transformer 模型与张量并SMP行度兼容

SMPv2 目前为以下 Hugging Face 变换器模型提供张量并行度支持。

  • GPT-neoX

  • Llama 2

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

配置张量并行度

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

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

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 Transformer 检查点

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

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

有关应用 v2 的张量并行度功能微调模型的SMP更多信息，请参见。微调