激活检查点

DeepSpeed 中的激活检查点 API 可用于启用与激活检查点相关的各种内存优化。这些优化包括在使用模型并行时跨 GPU 划分激活、CPU 检查点、连续内存优化等。

请参阅 DeepSpeed JSON 配置，了解完整设置。

在此，我们介绍激活检查点 API。请参阅为 Megatron-LM 教程启用 DeepSpeed，了解使用示例。

配置激活检查点

deepspeed.checkpointing.configure(mpu_, deepspeed_config=None, partition_activations=None, contiguous_checkpointing=None, num_checkpoints=None, checkpoint_in_cpu=None, synchronize=None, profile=None)

配置 DeepSpeed 激活检查点。

参数

mpu – 可选：实现以下方法的 Object：get_model_parallel_rank/group/world_size 以及 get_data_parallel_rank/group/world_size
deepspeed_config – 可选：提供时，DeepSpeed Config json 文件将用于配置 DeepSpeed 激活检查点
partition_activations – 可选：启用后，在模型并行 GPU 间分配激活检查点。默认情况下为 False。如果提供，将覆盖 deepspeed_config
contiguous_checkpointing – 可选：将激活检查点复制到连续内存缓冲区。仅在启用 partition_activations 时与同质检查点一起使用。必须提供 num_checkpoints。默认情况下为 False。如果提供，将覆盖 deepspeed_config
num_checkpoints – 可选：激活检查点的数量，在模型的前向传播期间存储。用于计算 contiguous_checkpointing 的缓冲区大小，如果提供，将覆盖 deepspeed_config
checkpoint_in_cpu – 可选：将激活检查点移动到 CPU。仅适用于 partition_activations。默认值为 false。如果提供，将覆盖 deepspeed_config
synchronize – 可选：在每次对 deepspeed.checkpointing.checkpoint 调用时，在前向和后向传递中执行 get_accelerator().synchronize()。默认情况下为 false。如果提供，将覆盖 deepspeed_config
profile – 可选：记录每次 deepspeed.checkpointing.checkpoint 调用的前向和后向时间。如果提供，将覆盖 deepspeed_config

返回

无

deepspeed.checkpointing.is_configured()

如果已配置 deepspeed 激活检查点，则为 True: 通过调用 deepspeed.checkpointing.configure，否则返回 false

参数: 无 –
返回: 如果配置则为 True，否则为 False

使用激活检查点

deepspeed.checkpointing.checkpoint(function, *args): 检查点模型或模型的部件。这已直接从 torch.utils.checkpoint 复制过来。

deepspeed.checkpointing.reset()

重置与连续内存优化相关的内存缓冲区。当计算多个前向传播且没有通常用来清除这些缓冲区的后向传播时，应在评估中调用此函数。：param 无

返回: 无

配置和检查点随机种子

deepspeed.checkpointing.get_cuda_rng_tracker(): 获取 cuda rng 跟踪器。

deepspeed.checkpointing.model_parallel_cuda_manual_seed(seed)

初始化模型并行 cuda 种子。

在模型并行初始化后应当调用此函数。此外，此函数后不得调用 get_accelerator().manual_seed。基本上，此函数取代了该函数。跟踪两组 RNG 状态

默认状态：此状态用于数据并行，且在一组
模型并行 GPU 中相同，但在不同模型并行组中不同。例如，用于非模型并行区域中的丢弃。

模型并行状态：此状态在一组模型
并行 GPU 中不同，但在数据并行组中相同。例如，用于模型并行区域中的丢弃。

class deepspeed.checkpointing.CudaRNGStatesTracker

cuda RNG 状态的跟踪器。

使用 add 方法，将根据输入 seed 初始化 cuda rng 状态，并分配给 name。稍后，通过派生 rng 状态，我们可以执行操作并返回到我们的初始 cuda 状态。

class deepspeed.checkpointing.CheckpointFunction(*args, **kwargs)

该函数改编自 torch.utils.checkpoint，有两个主要改动

torch.cuda.set_rng_state 已替换为 _set_cuda_rng_state #ignore-cuda

模型并行跟踪器中的状态也被正确跟踪/设置/重置。

性能激活分区，连续内存优化

CPU 检查点

对前向和后向函数进行概要分析