Fix assertion for offloading states (#6855)

This PR fixes the assertions in `offload_states` method mentioned in #6833. Co-authored-by: Logan Adams <[email protected]>
microsoft · Dec 16, 2024 · db98cc3 · db98cc3
1 parent fc7c070
commit db98cc3
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -3738,6 +3738,11 @@ def offload_states(self,
         assert self.zero_optimization_stage(
         ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3."
 
+        opt_offload_config = self.zero_offload_optimizer()
+        assert opt_offload_config is None or opt_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded optimizer states."
+        param_offload_config = self.zero_offload_param()
+        assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters."
+
         assert not self.zero_offload_param(), "Moving states across devices is not supported for offloaded parameters."
 
         if device == OffloadDeviceEnum.none: