Fix FSDP resume Initialization issue (#34032)

* Fix FSDP Initialization for resume training

* Added init_fsdp function to work with dummy values

* Fix FSDP initialization for resuming training

* Added CUDA decorator for tests

* Added torch_gpu decorator to FSDP tests

* Fixup for failing code quality tests

tests/trainer/test_trainer.py

@@ -4914,3 +4914,34 @@ class OptimizerAndModelInspectionTest(unittest.TestCase):
             param = next(model.parameters())
             group = trainer.get_optimizer_group(param)
             self.assertIn(param, group["params"])
+
+
+@require_torch_gpu
+@require_torch
+@require_accelerate
+class TestFSDPInitialization(unittest.TestCase):
+    def test_fsdp_initialization(self):
+        config = RegressionModelConfig(a=1, b=1, double_output=False)
+        model = RegressionPreTrainedModel(config)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = TrainingArguments(
+                output_dir=tmp_dir,
+                fsdp=True,
+                fsdp_config={"min_num_params": 1},
+                no_cuda=True,
+            )
+            trainer = Trainer(model=model, args=training_args)
+
+            # Check for FSDP enabled
+            self.assertTrue(trainer.is_fsdp_enabled)
+
+            # Check if model is wrapped with FSDP
+            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+            self.assertTrue(trainer.model, FSDP)
+
+            # Running a forward pass to ensure FSDP is initialized
+            dummy_input = torch.ones((1, 1), dtype=torch.float)
+            output = trainer.model(dummy_input)
+            self.assertTrue(output)