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