本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。
激活分载
重要
在 SMP v2.2.0 中,SMP 库的激活卸载功能不起作用。改用本机 PyTorch 激活卸载。
通常,向前传递计算每层的激活次数,并将它们保存在GPU内存中,直到相应层的向后传递完成。在向前传递之后,将这些张量卸载到 CPU 内存中,并在需要向后传递层时将其提取回 GPU,可以节省大量的 GPU 内存使用量。 PyTorch支持卸载激活,但是该实现会导致 GPU 处于空闲状态,而激活是在向后传递期间从 CPU 中取回的。使用激活卸载时,这会导致性能严重降低。
SMP v2 改进了这种激活卸载,它会在需要激活之前提前获取激活,GPU 才能开始向后传递这些激活。这种预取功能有助于在没有闲置的 GPU 的情况下更有效地运行训练进度,从而在不降低性能的情况下降低内存使用量。
您可以在训练脚本中保留用于卸载激活的本机 PyTorch 模块。以下是在脚本中应用 SMP 激活卸载功能的示例结构。请注意,激活卸载仅在与激活检查点一起使用时才适用。要详细了解用于激活卸载的原生 PyTorch 检查点工具,另请参阅PyTorch GitHub存储库中的 checkpoint_wrapper.py
要在激活检查点上应用 SMP 激PyTorch 活卸载功能,请在此期间sm_activation_offloading
和activation_loading_horizon
参数添加到 SMP 配置字典中。步骤 2:启动训练作业
以下代码片段展示了如何将 SMP 初始化模块torch.sagemaker.init()
添加到训练脚本中,以及如何为训练作业启动器设置 JSON 格式的 SMP 配置字典,同时遵循中介绍的两步流程。开始使用 SageMaker 模型并行度库 v2您无需对 PyTorch 模型或 PyTorch FSDPsm_activation_offloading
和 activation_loading_horizon
参数的更多信息,请参阅 SMP v2 核心功能配置参数。
SMP 配置
{ "activation_loading_horizon": 2, "sm_activation_offloading": True }
在训练脚本中
注意
在激活 SMP 激活卸载功能时,请确保您也使用该 PyTorch offload_wrapper
功能并将其应用于根模块。SMP 激活卸载功能使用根模块来确定何时完成正向传递以开始预获取。
import torch.sagemaker as tsm tsm.init() # Native PyTorch module for activation offloading from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( apply_activation_checkpointing, offload_wrapper, ) model = FSDP(...) # Activation offloading requires activation checkpointing. apply_activation_checkpointing( model, check_fn=
checkpoint_transformer_layers_policy
, ) model = offload_wrapper(model)