Not supporting closure.
(self, closure=None)
| 2542 | |
| 2543 | @instrument_w_nvtx |
| 2544 | def step(self, closure=None): |
| 2545 | """ |
| 2546 | Not supporting closure. |
| 2547 | """ |
| 2548 | self._pre_step() |
| 2549 | self._partition_all_parameters() |
| 2550 | |
| 2551 | #checks for overflow, adjust the loss scale accordingly |
| 2552 | if self._overflow_check_and_loss_scale_update(): |
| 2553 | if self.swap_optimizer: |
| 2554 | self.optimizer_swapper.log_timers() |
| 2555 | return |
| 2556 | |
| 2557 | norm_groups = self._get_norm_groups() |
| 2558 | scaled_global_grad_norm = torch.linalg.vector_norm(torch.stack(norm_groups)) |
| 2559 | |
| 2560 | # Stash unscaled gradient norm |
| 2561 | self._global_grad_norm = scaled_global_grad_norm / self.loss_scale |
| 2562 | |
| 2563 | timer_names = set() |
| 2564 | |
| 2565 | timer_names.add(OPTIMIZER_STEP_TIMER) |
| 2566 | self.timers(OPTIMIZER_STEP_TIMER).start() |
| 2567 | |
| 2568 | #update parameters one sub group at a time |
| 2569 | for sub_group_id, group in enumerate(self.fp16_groups): |
| 2570 | |
| 2571 | #prepare optimizer states, gradients and fp32 parameters for update |
| 2572 | self._prepare_sub_group(sub_group_id, timer_names) |
| 2573 | |
| 2574 | #scale the fp32 gradients |
| 2575 | self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) |
| 2576 | |
| 2577 | #apply the optimizer step on the sub group and copy fp32 parameters to fp16 |
| 2578 | self._optimizer_step(sub_group_id) |
| 2579 | |
| 2580 | #put fp16 parameters in appropriate location |
| 2581 | self._reassign_or_swap_out_partitioned_parameters(sub_group_id) |
| 2582 | |
| 2583 | #release memory or swap out optimizer states of fp32 parameters |
| 2584 | self._release_sub_group(sub_group_id, timer_names) |
| 2585 | |
| 2586 | self.timers(OPTIMIZER_STEP_TIMER).stop() |
| 2587 | |
| 2588 | self._post_step(timer_names) |
| 2589 | |
| 2590 | # warn user about caching allocator flushes |
| 2591 | memory_stats = get_accelerator().memory_stats() |
| 2592 | alloc_retries = memory_stats.get("num_alloc_retries") |
| 2593 | if alloc_retries is None: |
| 2594 | alloc_retries = 0 |
| 2595 | if alloc_retries > self.n_caching_allocator_flushes: |
| 2596 | if dist.get_rank() == 0: |
| 2597 | logger.warning( |
| 2598 | "%d pytorch allocator cache flushes since last step. this happens " |
| 2599 | "when there is high memory pressure and is detrimental to " |
| 2600 | "performance. if this is happening frequently consider adjusting " |
| 2601 | "settings to reduce memory consumption. If you are unable to " |
no test coverage detected