Update on "[DCP][DSD] Add AdamW to distributed state dict unit tests"

wz337 · wz337 · commit 0e57f157ca24 · 2024-03-14T16:52:15.000-07:00
cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj tianyu-l wconstab yf225 chauhang LucasLLC Thanks fegin for removing the fsdp root module check in DCP to unblock test updates. #121544 This PR adds "optimzer_class" as a kwarg for the subtests of the following tests to add AdamW as an option. - test_fsdp - test_compiled_fsdp - test_fsdp2 - test_ddp - test_fsdp_ddp - test_cpu_offload_full_state_dict In addition, we temporarily remove the two _verify_osd_by_load in _test_save_load, as state dict loading seems affect parameters. Creating an issue #121186 to keep track. [ghstack-poisoned]
diff --git a/test/distributed/checkpoint/test_state_dict.py b/test/distributed/checkpoint/test_state_dict.py
@@ -236,8 +236,12 @@ def _test_fsdp2(
     ):
         def init_model_optim():
             orig_model = CompositeParamModel(device=torch.device("cuda"))
-            orig_optim = optimizer_class(orig_model.parameters(), lr=1e-3, foreach=foreach)
-            copy_optim = optimizer_class(orig_model.parameters(), lr=1e-3, foreach=foreach)
+            orig_optim = optimizer_class(
+                orig_model.parameters(), lr=1e-3, foreach=foreach
+            )
+            copy_optim = optimizer_class(
+                orig_model.parameters(), lr=1e-3, foreach=foreach
+            )
 
             dist_model = FSDP2(
                 copy.deepcopy(orig_model),
@@ -246,7 +250,9 @@ def init_model_optim():
 
             if compile_model:
                 dist_model = torch.compile(dist_model)
-            dist_optim = optimizer_class(dist_model.parameters(), lr=1e-3, foreach=foreach)
+            dist_optim = optimizer_class(
+                dist_model.parameters(), lr=1e-3, foreach=foreach
+            )
 
             return orig_model, orig_optim, copy_optim, dist_model, dist_optim
 
