| 2754 | |
| 2755 | @instrument_w_nvtx |
| 2756 | def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): |
| 2757 | # Skip gradient reduction when DeepCompile is enabled |
| 2758 | # DeepCompile handles its own gradient reduction through compiled graph operations |
| 2759 | if self.is_deepcompile_active() and not self.compile_autosp(): |
| 2760 | return |
| 2761 | |
| 2762 | # Pass (PP) gas boundary flag to optimizer (required for zero) |
| 2763 | self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary() |
| 2764 | if self.is_gradient_accumulation_boundary(): |
| 2765 | self._reduce_autoep_folding_tp_replicated_gradients() |
| 2766 | # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well |
| 2767 | if self.zero_optimization_partition_gradients(): |
| 2768 | self.optimizer.overlapping_partition_gradients_reduce_epilogue() |
| 2769 | |
| 2770 | # Communicate only at gradient accumulation boundaries |
| 2771 | elif self.is_gradient_accumulation_boundary(): |
| 2772 | if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr( |
| 2773 | self.optimizer, 'reduce_gradients'): |
| 2774 | self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) |
| 2775 | else: |
| 2776 | grads = None |
| 2777 | self.buffered_allreduce_fallback(grads=grads, elements_per_buffer=bucket_size) |
| 2778 | elif self.zenflow: |
| 2779 | self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism) |
| 2780 | |
| 2781 | def _reduce_autoep_folding_tp_replicated_gradients(self): |
| 2782 | folding_spec = getattr(self, "_autoep_folding_spec", None) |