From 0b6e185ac8a19e99efabceb3b1ca202641f3cae4 Mon Sep 17 00:00:00 2001 From: Matthew Hoffman Date: Thu, 29 Feb 2024 17:41:05 -0800 Subject: [PATCH] Remove unnecessary assert on sub_module.training --- deepspeed/runtime/zero/parameter_offload.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index e9e79c2647fb..3fed9aa1d0cc 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/deepspeed/runtime/zero/parameter_offload.py @@ -466,7 +466,6 @@ def post_sub_module_forward_function(self, sub_module): @torch.no_grad() def pre_sub_module_backward_function(self, sub_module): - assert sub_module.training, "backward pass is invalid for module in evaluation mode" param_coordinator = self.get_param_coordinator(training=True) param_coordinator.trace_prologue(sub_module) if param_coordinator.is_record_trace(): @@ -475,7 +474,6 @@ def pre_sub_module_backward_function(self, sub_module): @torch.no_grad() def post_sub_module_backward_function(self, sub_module): - assert sub_module.training, "backward pass is invalid for module in evaluation mode" see_memory_usage( f"After sub module backward function {sub_module.__class__.__name__} {sub_module.id} before release", force=False)