fix assertion for offloading states

microsoft · Dec 11, 2024 · ad235e7 · ad235e7
1 parent 074d5c6
commit ad235e7
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: