优化器

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 实现了两种融合。

  • 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 正则化或权重衰减 True 表示解耦权重衰减（也称为 AdamW） (默认值：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 在 `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)

ZeroOneAdam (GPU)

OnebitLamb (GPU)