优化器

DeepSpeed 提供了在 CPU 上的 Adam 优化器的 高性能实现；在 GPU 上的 FusedAdam、FusedLamb、OnebitAdam、OnebitLamb 优化器。

Adam (CPU)

class deepspeed.ops.adam.DeepSpeedCPUAdam(model_params, lr=0.001, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False, adamw_mode=True, fp32_optimizer_states=True)[source]

FusedAdam (GPU)

class deepspeed.ops.adam.FusedAdam(params, lr=0.001, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, adam_w_mode=True, weight_decay=0.0, amsgrad=False, set_grad_none=True)[source]

实现 Adam 算法。

目前仅限 GPU。需要通过 pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ 安装 Apex。

此版本的 fused Adam 实现 2 种融合。

  • 将 Adam 更新的逐元素操作融合在一起

  • 一个多张量应用启动，将应用于模型所有参数的逐元素更新批处理到一个或几个内核启动中。

apex.optimizers.FusedAdam 可以用作 torch.optim.AdamW 的直接替代，或者 torch.optim.Adam，使用 adam_w_mode=False

opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....)
...
opt.step()

apex.optimizers.FusedAdam 可以与 Amp 一起使用或不使用。如果你希望使用 FusedAdam 和 Amp，你可以选择任何 opt_level

opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....)
model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
...
opt.step()

通常，建议使用 opt_level="O1"。

警告

以前版本的 FusedAdam 允许对 step 使用一些额外的参数。这些额外的参数现在已弃用且不再需要。

Adam 在 Adam: A Method for Stochastic Optimization 中被提出。

参数

  • params (iterable) – 要优化的参数的迭代器或定义参数组的字典。

  • lr (float, optional) – 学习率。 (默认: 1e-3)

  • betas (Tuple[float, float], optional) – 用于计算梯度及其平方运行平均值的系数。 (默认: (0.9, 0.999))

  • eps (float, optional) – 添加到分母的项，以提高数值稳定性。 (默认: 1e-8)

  • weight_decay (float, optional) – 权重衰减 (L2 惩罚) (默认: 0)

  • amsgrad (boolean, optional) – 是否使用来自论文 On the Convergence of Adam and Beyond 的该算法的 AMSGrad 变体 (默认: False) 在 FusedAdam 中不受支持！

  • adam_w_mode (boolean, optional) – 应用 L2 正则化或权重衰减，对解耦权重衰减（也称为 AdamW）为 True (默认: True)

  • set_grad_none (bool, optional) – 调用 zero_grad() 方法时，是否将梯度设置为 None。 (默认: True)

FusedLamb (GPU)

class deepspeed.ops.lamb.FusedLamb(params, lr=0.001, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, eps_inside_sqrt=False, weight_decay=0.0, max_grad_norm=0.0, max_coeff=10.0, min_coeff=0.01, amsgrad=False)[source]

实现 LAMB 算法。目前仅限 GPU。

LAMB 在 `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. https://arxiv.org/abs/1904.00962 中被提出。

参数

  • params (iterable) – 要优化的参数的迭代器或定义参数组的字典。

  • lr (float, optional) – 学习率。 (默认: 1e-3)

  • bias_correction (bool, optional) – 偏差校正 (默认: True)

  • betas (Tuple[float, float], optional) – 用于计算梯度及其平方运行平均值的系数。 (默认: (0.9, 0.999))

  • eps (float, optional) – 添加到分母的项，以提高数值稳定性。 (默认: 1e-8)

  • eps_inside_sqrt (boolean, optional) – 在“更新参数”步骤中，在评估平方根之前将 eps 添加到偏差校正后的二阶矩估计，而不是将其添加到二阶矩估计的平方根，如原始论文中所述。 (默认: False)

  • weight_decay (float, optional) – 权重衰减 (L2 惩罚) (默认: 0)

  • max_grad_norm (float, optional) – 用于剪裁全局梯度范数的值 (默认: 0.0)

  • max_coeff (float, optional) – lamb 系数的最大值 (默认: 10.0)

  • min_coeff (float, optional) – lamb 系数的最小值 (默认: 0.01)

  • amsgrad (boolean, optional) – 在 FusedLamb 中不受支持！

OneBitAdam (GPU)

class deepspeed.runtime.fp16.onebit.adam.OnebitAdam(params, deepspeed=None, lr=0.001, freeze_step=100000, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, eps_inside_sqrt=False, weight_decay=0.0, max_grad_norm=0.0, amsgrad=False, cuda_aware=False, comm_backend_name='nccl')[source]

实现 1 位 Adam 算法。目前仅支持 GPU。有关使用示例，请参阅 https://www.deepspeed.org.cn/tutorials/onebit-adam/。有关技术细节，请阅读 https://arxiv.org/abs/2102.02888

参数

  • params (iterable) – 要优化的参数的迭代器或定义参数组的字典。

  • lr (float, optional) – 学习率。 (默认: 1e-3)

  • freeze_step (int, optional) – 在开始使用压缩通信之前，用于预热（未压缩）阶段的步数。 (默认 100000)

  • betas (Tuple[float, float], optional) – 用于计算梯度及其平方运行平均值的系数。 (默认: (0.9, 0.999))

  • eps (float, optional) – 添加到分母的项，以提高数值稳定性。 (默认: 1e-8)

  • weight_decay (float, optional) – 权重衰减 (L2 惩罚) (默认: 0)

  • amsgrad (boolean, optional) – 是否使用 AMSGrad 变体，该变体来自论文 On the Convergence of Adam and Beyond (默认：False) 1 位 Adam 中不支持！

  • eps_inside_sqrt (boolean, optional) – 在“更新参数”步骤中，在评估平方根之前将 eps 添加到偏差校正后的二阶矩估计，而不是将其添加到二阶矩估计的平方根，如原始论文中所述。 (默认: False)

  • cuda_aware (boolean, required) – 如果底层 MPI 实现支持 CUDA 感知通信，则设置为 True。 (默认：False)

  • comm_backend_name (string, optional) – 如果需要，设置为 'mpi'。 (默认：'nccl'）

ZeroOneAdam (GPU)

class deepspeed.runtime.fp16.onebit.zoadam.ZeroOneAdam(params, deepspeed=None, lr=0.001, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, eps_inside_sqrt=False, weight_decay=0.0, max_grad_norm=0.0, var_freeze_step=100000, var_update_scaler=16, local_step_scaler=32678, local_step_clipper=16, amsgrad=False, cuda_aware=False, comm_backend_name='nccl')[source]

实现 0/1 Adam 算法。目前仅支持 GPU。有关使用示例，请参阅 https://www.deepspeed.org.cn/tutorials/zero-one-adam/。有关技术细节，请阅读 https://arxiv.org/abs/2202.06009 :param params: 要优化的参数的迭代器，或者定义参数组的字典。

参数

  • lr (float, optional) – 学习率。 (默认: 1e-3)

  • betas (Tuple[float, float], optional) – 用于计算梯度及其平方运行平均值的系数。 (默认: (0.9, 0.999))

  • eps (float, optional) – 添加到分母的项，以提高数值稳定性。 (默认: 1e-8)

  • weight_decay (float, optional) – 权重衰减 (L2 惩罚) (默认: 0)

  • var_freeze_step (int, optional) – 更新方差的最新步数，使用 https://arxiv.org/abs/2202.06009 中的符号，它表示 max{i|i in T_v}。请注意，这与 1 位 Adam 中的冻结步数不同。var_freeze_step 通常是学习率预热结束时的步数，因此不需要调整。 (默认：100000)

  • var_update_scaler (int, optional) – 更新方差的间隔。请注意，方差的更新策略遵循指数规则，其中 var_update_scaler 表示 0/1 Adam 论文中的 kappa。 (默认：16)

  • local_step_scaler (int, optional) – 根据学习率策略缩放局部步长间隔的间隔。 (默认：32678)

  • local_step_clipper (int, optional) – 使用学习率策略的局部步长的最大间隔。这对应于 0/1 Adam 论文中的变量 H。 (默认：16)

  • amsgrad (boolean, optional) – 是否使用 AMSGrad 变体，该变体来自论文 On the Convergence of Adam and Beyond (默认：False) 0/1 Adam 中不支持！

  • eps_inside_sqrt (boolean, optional) – 在“更新参数”步骤中，在评估平方根之前将 eps 添加到偏差校正后的二阶矩估计，而不是将其添加到二阶矩估计的平方根，如原始论文中所述。 (默认: False)

  • cuda_aware (boolean, required) – 如果底层 MPI 实现支持 CUDA 感知通信，则设置为 True。 (默认：False)

  • comm_backend_name (string, optional) – 如果需要，设置为 'mpi'。 (默认：'nccl'）

OnebitLamb (GPU)

class deepspeed.runtime.fp16.onebit.lamb.OnebitLamb(params, deepspeed=None, lr=0.001, freeze_step=100000, bias_correction=True, betas=(0.9, 0.999), eps=1e-08, eps_inside_sqrt=False, weight_decay=0.0, max_grad_norm=0.0, max_coeff=10.0, min_coeff=0.01, amsgrad=False, cuda_aware=False, comm_backend_name='nccl', coeff_beta=0.9, factor_max=4.0, factor_min=0.5, factor_threshold=0.1)[source]

实现 1 位 Lamb 算法。目前仅支持 GPU。有关使用示例，请参阅 https://www.deepspeed.org.cn/tutorials/onebit-lamb/。有关技术细节，请参阅我们的论文 https://arxiv.org/abs/2104.06069。

参数

  • params (iterable) – 要优化的参数的迭代器或定义参数组的字典。

  • lr (float, optional) – 学习率。 (默认: 1e-3)

  • freeze_step (int, optional) – 在开始使用压缩通信之前，用于预热（未压缩）阶段的步数。 (默认 100000)

  • betas (Tuple[float, float], optional) – 用于计算梯度及其平方运行平均值的系数。 (默认: (0.9, 0.999))

  • eps (float, optional) – 添加到分母的项，以提高数值稳定性。 (默认: 1e-8)

  • weight_decay (float, optional) – 权重衰减 (L2 惩罚) (默认: 0)

  • max_coeff (float, optional) – lamb 系数的最大值 (默认: 10.0)

  • min_coeff (float, optional) – lamb 系数的最小值 (默认: 0.01)

  • amsgrad (boolean, optional) – 是否使用 AMSGrad 变体，该变体来自论文 On the Convergence of Adam and Beyond (默认：False) 1 位 Lamb 中不支持！

  • eps_inside_sqrt (boolean, optional) – 在“更新参数”步骤中，在评估平方根之前将 eps 添加到偏差校正后的二阶矩估计，而不是将其添加到二阶矩估计的平方根，如原始论文中所述。 (默认: False)

  • cuda_aware (boolean, required) – 如果底层 MPI 实现支持 CUDA 感知通信，则设置为 True。 (默认：False)

  • comm_backend_name (string, optional) – 如果需要，设置为 'mpi'。 (默认：'nccl'）

  • coeff_beta (float, optional) – 用于计算 Lamb 系数的运行平均值的系数 (默认：0.9) 请注意，您可能需要根据选择的 freeze_step 增大或减小此 beta，因为 1/(1 - coeff_beta) 应小于或等于 freeze_step

  • factor_max (float, optional) – 在压缩阶段对冻结的 Lamb 系数进行缩放的因子最大值 (默认：4.0)

  • factor_min (float, optional) – 在压缩阶段对冻结的 Lamb 系数进行缩放的因子最小值 (默认：0.5)

  • factor_threshold (float, optional) – 缩放因子在步数之间可以波动的阈值 (默认：0.1)