Revert "Fix FSDP resume Initialization issue" (#34193)
Revert "Fix FSDP resume Initialization issue (#34032)"
This reverts commit 4de1bdbf63.
This commit is contained in:
@@ -4914,34 +4914,3 @@ 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)
|
||||
|
||||
Reference in New Issue
Block a user