Remove FSDP wrapping from sub-models.

[eljandoubi]

What does this PR do?

Fixes #34113

Who can review?

Library:

• trainer: @muellerzr and @SunMarc

[SunMarc]

Thnaks for fixing the issue @eljandoubi ! Do you think there is a simpler way to handle this edge case @muellerzr ?

[SunMarc]

You can use unwarp_model function in transformers instead. Also, why do we need to set recursive to True ? Also, please leave a comment above as this specific path is only to make it functional with auto_find_batch_size .

[eljandoubi]

unwrap_model does not provide access to the recursive argument. Auto-wrap policies wrap submodules with FSDP, and unwrap_model is unable to remove them. You can test this on the toy example from the PyTorch FSDP tutorial for rank=0 and world_size=1, then experiment with the line I provided in a notebook.

my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=20000
    )
torch.cuda.set_device(rank)
model = Net().to(rank)
print(model)
fsdp_model = FSDP(model,
    auto_wrap_policy=my_auto_wrap_policy)
print(fsdp_model)
unwrap_model = unwarp_model(fsdp_model)
print(unwrap_model)

VS
You need to reinstantiates model and fsdp_model:

model = Net().to(rank)

fsdp_model = FSDP(model,
    auto_wrap_policy=my_auto_wrap_policy)

extract_model = extract_model_from_parallel(fsdp_model, recursive=True)
print(extract_model)

[SunMarc]

I'm talking about this function in transformers. It uses extract_model_from_parallel under the hood so it should be comparable.

[eljandoubi]

Euh I see.

[eljandoubi]

@SunMarc @muellerzr Did you get a different result than I did?

[muellerzr]

Thanks for the fix, can you add a test in tests/test_trainer.py for this?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks ! Left an suggestion for unwrap_model

[eljandoubi]

@SunMarc I migrated to unwrap_model.

[LysandreJik]

Let's merge it if you're both ok with it @SunMarc @muellerzr

[SunMarc]

Please rebase this PR on main in order to pass the CI @eljandoubi !