fix opt step on cpu

PrimeIntellect-ai · Sep 11, 2024 · e926e6c · e926e6c
1 parent 04dd585
commit e926e6c
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/open_diloco/train_pure_fsdp.py b/open_diloco/train_pure_fsdp.py
@@ -110,7 +110,14 @@ def get_model(config: Config) -> LlamaForCausalLM:
 
 
 def get_offloaded_param(model: LlamaForCausalLM) -> list[torch.Tensor]:
-    return [param.data.detach().clone().to("cpu") for param in model.parameters()]
+    offloaded_params = []
+    for param in model.parameters():
+        if param.requires_grad:
+            offloaded_param = param.data.detach().clone().to("cpu")
+            offloaded_param.requires_grad = True
+            offloaded_params.append(offloaded_param)
+
+    return offloaded_params
 
 
 def train(config: Config):
@@ -241,6 +248,9 @@ def train(config: Config):
             else:
                 dist.all_reduce(param_offloaded.grad, op=dist.ReduceOp.AVG, group=global_pg)
 
+        for param in outer_optimizer.param_groups[0]["params"]:
+            print(param.requires_grad)
+
         outer_optimizer.step()
         outer_optimizer.zero_grad()