Update FSDP hardcoded params + dynamic copy per review

Author: amanzoni1

src/transformers/training_args.py

@@ -2676,7 +2676,6 @@ def _process_fsdp_args(self):
             with open(self.fsdp_config, encoding="utf-8") as f:
                 self.fsdp_config = json.load(f)
 
-        fsdp_version = self.fsdp_config.get("fsdp_version", 1)
         if self.fsdp_config is not None and isinstance(self.fsdp_config, dict):
             for k in list(self.fsdp_config.keys()):
                 if k.startswith("fsdp_") and k != "fsdp_version":
@@ -2749,14 +2748,28 @@ def _process_fsdp_args(self):
                             self.fsdp_config["transformer_layer_cls_to_wrap"]
                         )
 
-            # Pull allowed parameters from fsdp_config
-            ALLOWED_FSDP_PARAMS = {f.name for f in fields(FullyShardedDataParallelPlugin)}
-            for key in ALLOWED_FSDP_PARAMS:
-                if key in self.fsdp_config and key not in fsdp_plugin_args:
-                    fsdp_plugin_args[key] = self.fsdp_config[key]
-            fsdp_plugin_args["fsdp_version"] = fsdp_version
+            fsdp_plugin_args["fsdp_version"] = self.fsdp_config.get("fsdp_version", 1)
+            prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH")
+            fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper()
+            if fsdp_plugin_args["fsdp_version"] == 1:
+                # forward_prefetch is not yet implemented in FSDP2
+                fsdp_plugin_args["forward_prefetch"] = str(self.fsdp_config.get("forward_prefetch", "false")).lower()
+
+            sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower()
+            cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower()
+            if sync_module_states == "false" and cpu_ram_efficient_loading == "true":
+                # In this case, all the processes except the main process would have random weights leading
+                # to unexpected behaviour during training, thus throwing error here to prevent it.
+                raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`')
+
+            # we need to set the env here as otherwise we get a warning in accelerate + we need to set it for transformers
+            fsdp_plugin_args["cpu_ram_efficient_loading"] = cpu_ram_efficient_loading
+            os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading
 
-            # Special HF-to-plugin map: transformer_layer_cls_to_wrap → joined cls_names
+            fsdp_plugin_args["sync_module_states"] = sync_module_states
+            fsdp_plugin_args["use_orig_params"] = str(self.fsdp_config.get("use_orig_params", "true")).lower()
+
+            # HF-to-plugin map
             if (
                 "transformer_layer_cls_to_wrap" in self.fsdp_config
                 and "transformer_cls_names_to_wrap" not in fsdp_plugin_args
@@ -2765,26 +2778,11 @@ def _process_fsdp_args(self):
                     self.fsdp_config["transformer_layer_cls_to_wrap"]
                 )
 
-            # Validation: sync_module_states vs cpu_ram_efficient_loading
-            sync_states = fsdp_plugin_args.get("sync_module_states", "true")
-            cpu_loading = fsdp_plugin_args.get("cpu_ram_efficient_loading", "false")
-
-            if isinstance(sync_states, str):
-                sync_states = sync_states.lower()
-            if isinstance(cpu_loading, str):
-                cpu_loading = cpu_loading.lower()
-            if sync_states == "false" and cpu_loading == "true":
-                raise ValueError('`sync_module_states` must be `"true"` if `cpu_ram_efficient_loading` is `"true"`.')
-
-            # CRITICAL: Set environment variable for cpu_ram_efficient_loading
-            if "cpu_ram_efficient_loading" in fsdp_plugin_args:
-                cpu_ram_value = fsdp_plugin_args["cpu_ram_efficient_loading"]
-                # Handle both bool and string values
-                if isinstance(cpu_ram_value, bool):
-                    cpu_ram_value = str(cpu_ram_value).lower()
-                elif isinstance(cpu_ram_value, str):
-                    cpu_ram_value = cpu_ram_value.lower()
-                os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_value
+            # Pull allowed parameters from fsdp_config
+            ALLOWED_FSDP_PARAMS = {f.name for f in fields(FullyShardedDataParallelPlugin)}
+            for key in ALLOWED_FSDP_PARAMS:
+                if key in self.fsdp_config and key not in fsdp_plugin_args:
+                    fsdp_plugin_args[key] = self.fsdp_config[key]
 
         return fsdp_plugin_args