register_optim_in_bwd_hooks

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

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

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

Parameters:

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

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