r""" Context manager to disable gradient reduction during backward pass. This context manager has the following effects on other DeepSpeed features: 1. Incompatible with ZeRO stage 2/3 which rely on reduction for gradient partitioning. 2. It is illegal
(self)
| 2977 | |
| 2978 | @contextmanager |
| 2979 | def no_sync(self): |
| 2980 | r""" |
| 2981 | Context manager to disable gradient reduction during backward pass. |
| 2982 | This context manager has the following effects on other DeepSpeed features: |
| 2983 | 1. Incompatible with ZeRO stage 2/3 which rely on reduction for gradient partitioning. |
| 2984 | 2. It is illegal to call engine.step() within the context manager. |
| 2985 | 3. Tracking of gradient accumulation steps is disabled. |
| 2986 | """ |
| 2987 | assert not self.zero_optimization_partition_gradients(), \ |
| 2988 | f"no_sync context manager is incompatible with gradient partitioning logic of ZeRO stage {self.zero_optimization_stage()}" |
| 2989 | |
| 2990 | assert not self.inside_no_sync_ctxt, "no_sync context manager reentry is unsupported" |
| 2991 | |
| 2992 | self.inside_no_sync_ctxt = True |
| 2993 | try: |
| 2994 | yield |
| 2995 | finally: |
| 2996 | self.inside_no_sync_ctxt = False |
| 2997 | |
| 2998 | @contextmanager |
| 2999 | def coalesce_grad_reduction(self): |