本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。
张量并行性
张量并行性是模型并行性的一种,它在设备之间拆分特定的模型权重、梯度和优化器状态。与流水线并行度形成鲜明对比,后者保持单个权重不变,但将权重、梯度或优化器集分成多个设备,而张量并行度则对单个权重进行分片。这通常涉及对模型的特定运算、模块或层进行分布式计算。
在单个参数占用大部分 GPU 内存的情况下(例如词汇表很大的大型嵌入表,或者具有大量类的大型 softmax 层),则需要张量并行性。在这种情况下,将这种大张量或运算视为原子单元的效率会很低,并且会阻碍内存负载的平衡。
SMP v2 与 Transform er Engin
实际上,张量并行性在以下场景中特别有用。
-
当在较长的上下文长度下训练时,仅使用 FSDP 就会产生较高的激活记忆。
-
当使用非常大的集群进行训练时,其全局批次大小超过了所需限制。
Hugging Face Transformer 型号兼容 SMP 张量并行度
SMP v2 目前为以下 Hugging Face 变压器模型提供张量并行度支持。
-
GPT-NeoX
-
Llama 2
有关在这些模型上应用张量并行性的参考配置,请参阅。配置提示
配置张量并行度
对于tensor_parallel_degree
,您可以为张量并行度选择一个值。该值必须平均除以集群中 GPU 的数量。例如,要在使用具有 8 个 GPU 的实例时对模型进行分片,请选择 2、4 或 8。我们建议您从较小的数字开始,然后逐渐增加该数字,直到模型适合 GPU 内存。
以下代码片段展示了如何将 SMP 初始化模块torch.sagemaker.init()
添加到训练脚本中,以及如何为训练作业启动器设置 JSON 格式的 SMP 配置字典,同时遵循中介绍的两步流程。开始使用 SageMaker 模型并行度库 v2您无需对 PyTorch 模型或 PyTorch FSDPtensor_parallel_degree
和 random_seed
参数的更多信息,请参阅 SMP v2 核心功能配置参数。
SMP 配置
{ "tensor_parallel_degree": 8, "random_seed": 0 }
在你的训练脚本中
使用初始化torch.sagemaker.init()
以激活 SMP v2 并使用 API 封装您的模型。torch.sagemaker.transform
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 库提供了 API 来保存变换后的模型中的检查点,以 Hugging Face Transformer 表示形式,还提供了 torch.sagemaker.transform
API 来加载 Hugging Face Transformer 模型检查点进行微调。
有关在使用 SMP v2 的张量并行度功能时保存检查点的更多信息,请参阅。使用 SMP 时保存和加载检查点
有关应用 SMP v2 的张量并行度功能微调模型的更多信息,请参见。微调