Lomo

class Lomo(model, lr=0.001, clip_grad_norm=None, clip_grad_value=None, weight_decay=0.0, loss_scale_args={})

一个自定义的优化器类Lomo，用于在分布式训练中的梯度更新。

该类实现两个梯度更新函数 fuse_update() 和 fuse_update_zero3()，分别用于非ZeRO和ZeRO模式下的梯度更新。

参数:

• model -- 待优化的模型

• lr (default: 0.001) -- 学习率，默认值为1e-3

• clip_grad_norm (default: None) --

  梯度裁剪的范数阈值

  备注

  clip_grad_norm须为正数

• zero3_enabled -- 是否开启了 zero3

• clip_grad_value (default: None) -- 梯度裁剪的值域阈值

• loss_scale_args (default: {}) -- 用于初始化 DynamicLossScaler 的参数

fuse_update()

在非ZeRO模式下更新模型参数的梯度。

返回:

func，一个闭包函数，用于更新模型参数的梯度

fuse_update_zero3()

在ZeRO模式下更新模型参数的梯度。

返回:

func，一个闭包函数，用于更新模型参数的梯度。

fused_backward(loss, lr)

执行一步反向传播并更新模型的梯度。

参数:

• loss -- 模型的loss值

• lr -- 学习率

grad_norm(loss)

计算梯度的范数。

参数:

loss -- 模型的loss值