Disable FSDP use original parameters

pomonam · Mar 19, 2024 · 2ee08aa · 2ee08aa
1 parent 834c645
commit 2ee08aa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/gpu_tests/fsdp_test.py b/tests/gpu_tests/fsdp_test.py
@@ -55,7 +55,8 @@ def setUpClass(cls) -> None:
         cls.model = cls.model.to(device=device)
         cls.model = DistributedDataParallel(cls.model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
         my_auto_wrap_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=100)
-        cls.model = FSDP(cls.model, use_orig_params=True, auto_wrap_policy=my_auto_wrap_policy)
+        cls.model = FSDP(cls.model, use_orig_params=False, auto_wrap_policy=my_auto_wrap_policy)
+        print(cls.model)
 
         cls.analyzer = Analyzer(
             analysis_name="gpu_test",