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]) – 从参数到优化器的映射。