Merge pull request #327 from aws-samples/fsdp_sync_module_states_true

FSDP with meta device requires sync_module_states=True

Author: perifaws

3.test_cases/10.FSDP/train.py

@@ -161,6 +161,9 @@ def main(args):
         model = AutoModelForCausalLM.from_config(model_config)
     else:
         with torch.device("meta"):
+            # Instantiating model on `meta` device doesn't consume CPU memory,
+            # but requires specifing `param_init_fn=...`
+            # and `sync_module_states=True` in FSDP c-tor.
             model = AutoModelForCausalLM.from_config(model_config)
 
     num_params = compute_num_params(model)
@@ -197,6 +200,7 @@ def main(args):
         device_id=torch.cuda.current_device(),
         use_orig_params=False,
         sharding_strategy=sharding_strategy,
+        sync_module_states=True,
         param_init_fn=(lambda module: module.to_empty(device=torch.device("cuda"), recurse=False))
         if global_rank != 0 else None,
     )