| 428 | amp._clear_cache() |
| 429 | |
| 430 | def _step_with_closure(self, closure): |
| 431 | def wrapped_closure(): |
| 432 | # helpful for debugging |
| 433 | # print("Calling wrapped_closure, first_closure_call_this_step = {}" |
| 434 | # .format(self.first_closure_call_this_step)) |
| 435 | if self.first_closure_call_this_step: |
| 436 | # We expect that the fp16 params are initially fresh on entering |
| 437 | # self.step(), so _master_params_to_model_params() is unnecessary |
| 438 | # the first time wrapped_closure() is called within |
| 439 | # self.optimizer.step(). |
| 440 | self.first_closure_call_this_step = False |
| 441 | else: |
| 442 | # If self.optimizer.step() internally calls wrapped_closure more |
| 443 | # than once, it may update the fp32 params after each call. |
| 444 | # However, self.optimizer doesn't know about the fp16 params at |
| 445 | # all. If the fp32 params get updated, we can't rely on |
| 446 | # self.optimizer to refresh the fp16 params. We need to handle |
| 447 | # that manually: |
| 448 | self._master_params_to_model_params() |
| 449 | # Our API expects the user to give us ownership of the backward() |
| 450 | # call by replacing all calls to loss.backward() with |
| 451 | # optimizer.backward(loss). |
| 452 | # This requirement holds whether or not the call to backward() is |
| 453 | # made within a closure. If the user is properly calling |
| 454 | # optimizer.backward(loss) within "closure," calling closure() here |
| 455 | # will give the fp32 master params fresh gradients for the optimizer |
| 456 | # to play with, so all wrapped_closure needs to do is call |
| 457 | # closure() and return the loss. |
| 458 | temp_loss = closure() |
| 459 | if self.overflow: |
| 460 | self.overflow = False |
| 461 | raise ValueError("OVERFLOW within closure! Skipping step. " |
| 462 | "Reducing loss scale to : {}".format(self.loss_scale)) |
| 463 | return temp_loss |
| 464 | |
| 465 | retval = self.optimizer.step(wrapped_closure) |
| 466 | |
| 467 | self.first_closure_call_this_step = True |
| 468 | |
| 469 | return retval |
| 470 | |
| 471 | def backward(self, loss, update_master_grads=True): |
| 472 | """ |