diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index c911d6874d0c..7e593d4857f0 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -23,6 +23,7 @@ def __init__( self, module, init_optimizer, + param_names, timers, ds_config, **kwargs, @@ -34,7 +35,7 @@ def __init__( self.async_cpuadam_num = 0 self.max_grad_numel = 0 - super().__init__(module, init_optimizer, timers, ds_config, **kwargs) + super().__init__(module, init_optimizer, param_names, timers, ds_config, **kwargs) optimizer_config = { "lr": self.optimizer.param_groups[0]["lr"],