fixing fsdp autowrap functionality (#475)

* fixing fsdp autowrap functionality

* updating version requirements

* update version to latest torch stable version

* quality

Author: pacman100

examples/by_feature/fsdp_with_peak_mem_tracking.py

@@ -112,11 +112,8 @@ def training_function(config, args):
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
         if accelerator.is_main_process:
-            run = os.path.split(__file__)[-1].split(".")[0]
-            if args.logging_dir:
-                run = os.path.join(args.logging_dir, run)
-                accelerator.print(run)
-            accelerator.init_trackers(run, config)
+            experiment_config = vars(args)
+            accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
     datasets = load_dataset("glue", "mrpc")

src/accelerate/accelerator.py

@@ -188,8 +188,8 @@ def __init__(
             deepspeed_plugin.set_deepspeed_weakref()
 
         if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
-            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
-                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+            if is_torch_version("<", "1.12.0"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0")
 
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None

src/accelerate/utils/dataclasses.py

@@ -442,7 +442,7 @@ class FullyShardedDataParallelPlugin:
 
     def __post_init__(self):
         from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy
-        from torch.distributed.fsdp.wrap import default_auto_wrap_policy
+        from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
 
         if self.sharding_strategy is None:
             self.sharding_strategy = ShardingStrategy(int(os.environ.get("FSDP_SHARDING_STRATEGY", 1)))
@@ -458,4 +458,6 @@ def __post_init__(self):
 
         if self.auto_wrap_policy is None:
             if self.min_num_params > 0:
-                self.auto_wrap_policy = functools.partial(default_auto_wrap_policy, min_num_params=self.min_num_params)
+                self.auto_wrap_policy = functools.partial(
+                    size_based_auto_wrap_policy, min_num_params=self.min_num_params
+                )