Clip gradients (if needed) and step the base optimizer. Always return successful since there is no overflow.
(self)
| 472 | |
| 473 | @torch.no_grad() |
| 474 | def step(self): |
| 475 | """Clip gradients (if needed) and step the base optimizer. |
| 476 | Always return successful since there is no overflow.""" |
| 477 | |
| 478 | # Copy main_grads to grads. |
| 479 | if self.params_have_main_grad: |
| 480 | for param_group in self.optimizer.param_groups: |
| 481 | for param in param_group["params"]: |
| 482 | param.grad = param.main_grad |
| 483 | |
| 484 | # Clip gradients. |
| 485 | grad_norm = None |
| 486 | if self.clip_grad > 0.0: |
| 487 | grad_norm = self.clip_grad_norm(self.clip_grad) |
| 488 | |
| 489 | # count the zeros in the grads |
| 490 | num_zeros_in_grad = self.count_zeros() if self.log_num_zeros_in_grad else None |
| 491 | |
| 492 | # Update parameters. |
| 493 | self.optimizer.step() |
| 494 | |
| 495 | # No overflow for FP32 optimizer. |
| 496 | return True, grad_norm, num_zeros_in_grad |
| 497 | |
| 498 | def reload_model_params(self): |
| 499 | pass |
no test coverage detected