优化器

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)[来源]

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)[来源]

实现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 的直接替代品使用， 或者与 torch.optim.Adam 一起使用，当 adam_w_mode=False 时：

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

apex.optimizers.FusedAdam 可以在使用或不使用 Amp 的情况下使用。如果您希望在使用 Amp 时使用 FusedAdam，您可以选择任何 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`_ 中提出的。

Parameters

  • 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 (布尔值, 可选) – 是否使用来自论文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()方法时，是否将梯度设置为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)[来源]

实现了LAMB算法。目前仅支持GPU。

LAMB 是在 `大规模深度学习优化：在76分钟内训练BERT 中提出的。 https://arxiv.org/abs/1904.00962

Parameters

  • 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, 可选) – lamb系数的最小值（默认值：0.01）

  • amsgrad (布尔值, 可选) – 在FusedLamb中不支持！

OneBitAdam (GPU)

ZeroOneAdam (GPU)

OnebitLamb (GPU)