ENH: Allow FSDP ignored modules to be regex

[BenjaminBossan]

What does this PR do?

For FSDP, there is an option to indicate ignored_modules, which should be a list of modules are ignored by FSDP. Even though this argument was supported in accelerate, it was not very usable:

1. Listing all modules can tricky, especially with something like PEFT, where the whole model is wrapped and thus the module structure changes.
2. When configuring this argument, accelerate takes a detour via environment variables. These can only be strings. Therefore, passing a list of modules is not feasible.

Moreover, I noticed that the environment variable for ignored_modules was not even set, so configuring this argument didn't even work.

Status

Don't merge yet This PR is lacking tests. I would be happy for pointers on how to add those.

Context

When using PEFT with LoRA and the new target_parameters feature, I ran into an issue training such a model with FSDP. The only working fix I found was to ignore the layers targeted by LoRA. However, I could not configure accelerate to do that. With this PR, it is possible: I could successfully train such a PEFT model that targets q_proj and v_proj by setting fsdp_ignored_modules: '.*\.(q_proj$|v_proj$)'.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@SunMarc

[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.

[BenjaminBossan]

The failing test is unrelated and caused by an error in transformers that makes it not work with torch < 2.4. It should be patched soon.

AttributeError: module 'torch.nn' has no attribute 'RMSNorm'

[SunMarc]

Thanks for fixing this ! You can have a look at the tests/fsdp/test_fsdp.py file cc @S1ro1 for more info

[BenjaminBossan]

You can have a look at the tests/fsdp/test_fsdp.py

@SunMarc Thanks for the pointer but the issue I encountered is the following: I want to initialize an Accelerator instance with a given config file that specifies fsdp_ignored_modules. Then I want to:

1. Check that ignored_modules is correctly set on accelerator.state.fsdp_plugin.
2. Possibly also prepare a model and directly check the FSDP wrapped modules.

I couldn't find any test that would do something along these lines. There are some tests that use configs, like in test_performance, but those are shelling out, which means I can't inspect the Python objects as described above. Is there a simple way to achieve what I described above?

[S1ro1]

Check that ignored_modules is correctly set on accelerator.state.fsdp_plugin.
We usually test stuff like this with tests as these:

accelerate/tests/fsdp/test_fsdp.py

Line 90 in 2f075c7

def test_sharding_strategy(self):

[BenjaminBossan]

We usually test stuff like this with tests as these:

Thanks for the pointer @S1ro1 and sorry for my basic questions, but I'm still struggling to see how I can test the change. The two problems I have when using said function as a starting point:

1. In test_sharding_strategy, the FullyShardedDataParallelPlugin is initialized from env vars, but I want to initialize "end to end" from a config file, as the route is: config file -> env var -> fsdp plugin instance -> FSDP wrapper.
2. In that test, attributes on fsdp_plugin are checked. However, this doesn't cover the changes in Accelerator.prepare.

I guess I could try to add 3 tests: 1) Check that the config correctly sets the corresponding env var. 2) A test similar to test_sharding_strategy that checks the env var -> fsdp plugin part. 3) A test that starts with an initialized fsdp plugin and then calls accelerator.prepare, maybe mocking the FSDP wrapper to check that the correct ignored_modules are passed. Would that be the idea? I was hoping to test this all in one go.

[S1ro1]

I think reasonable is to test only the path after config file, the config_file -> env is tested quite okay + it's usually a place that is touched once and never again. IMO testing from env/fsdp_plugin -> model wrapper.

[SunMarc]

Merging for the release, we can add the tests later

[BenjaminBossan]

I created a PR, #3719, that only checks the env -> fsdp plugin -> model wrapper part.

I think reasonable is to test only the path after config file, the config_file -> env is tested quite okay

Well, that part was missing the ignored_modules arguments, which is part of the fix in this PR, so I can't quite agree to that :)