register_optim_in_bwd_hooks

torchtune.training.register_optim_in_bwd_hooks(model: Module, optim_dict: Dict[Parameter, Optimizer]) → None

注册在反向传播中运行的优化器步骤挂钩。

当将优化器步骤融合到反向传播中时，我们需要在优化器的给定参数的梯度准备就绪后立即调用 .step()。此实用程序在模型中的所有参数上注册后累积梯度挂钩以实现此目的。

参数：

• model (torch.nn.Module) – 将优化其参数的模型。请注意，目前将在模型中的所有参数上注册挂钩。

• optim_dict (Dict[torch.nn.Parameter, torch.optim.Optimizer]) – 参数到优化器的映射。