merge fsdp2 main and solve conflicts

verl/utils/checkpoint/fsdp_checkpoint_manager.py

@@ -84,7 +84,7 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte
         # every rank download its own checkpoint
         state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True) if self.save_model else None
         optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True) if self.save_optimizer else None
-        with FSDP.state_dict_type(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_config=state_dict_cfg, optim_state_dict_config=optim_cfg):
+        with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
             if self.save_model:
                 remote_model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
                 local_model_path = copy_to_local(remote_model_path)
@@ -146,7 +146,7 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
         optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True) if self.save_optimizer else None
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            with FSDP.state_dict_type(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_config=state_dict_cfg, optim_state_dict_config=optim_cfg):
+            with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
                 model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
                 optim_path = os.path.join(local_path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt")
                 extra_path = os.path.join(local_path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt")