优化器

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。

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

  • Adam 更新的逐元素操作的融合

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

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

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() 方法时是否将 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 在 `大型批次深度学习优化：76分钟训练 BERT` 中被提出。 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)

ZeroOneAdam (GPU)

OnebitLamb (GPU)