Misplaced global variable warned (microsoft#5725)

Move the global variable `warned` from `deepspeed.runtime.zero.parameter_offload.py` to `deepspeed.runtime.zero.utils.py` to avoid `NameError: name 'warned' is not defined` when calling `ap ply_to_tensors_only()` (defined in `deepspeed.runtime.zero.utils.py`). Co-authored-by: Logan Adams <[email protected]>
samadejacobs · Jul 20, 2024 · 879c6cd · 879c6cd
1 parent 6a163e0
commit 879c6cd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py
@@ -17,9 +17,6 @@
 
 FWD_MODULE_STACK = list()
 
-# ensure we only warn once, otherwise every iteration will trigger a warning
-warned = False
-
 
 #for each tensor in outputs run the forward_function and register backward_function as hook
 def _apply_forward_and_backward_to_tensors_only(module, forward_function, backward_function, outputs):

diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py
@@ -16,6 +16,9 @@
 from deepspeed.utils.nvtx import instrument_w_nvtx
 from deepspeed.accelerator import get_accelerator
 
+# ensure we only warn once, otherwise every iteration will trigger a warning
+warned = False
+
 
 def _initialize_parameter_parallel_groups(parameter_parallel_size=None):
     data_parallel_size = int(dist.get_world_size())