create _preprare_fsdp to pre- prepare fsdp model training

[eljandoubi]

Fixes #3209

Who can review?

• Fully-Sharded Data Parallism: @muellerzr

[eljandoubi]

@muellerzr have you any feedback?

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

Thanks for working on this important issue.

I wonder on the overall approach to solving this issue. Right now, IIUC, the idea is to re-initialize the optimizer(s) using the FSDP-wrapped models. This could potentially be error prone and ideally should be well tested before we merge.

Would it be possible instead to check if this has happened and raise an error, then direct users towards not initializing the optimizer prematurely, and handle the initialization with accelerate so that it's in the right order? Maybe that can't work, but I think it's worth considering different solutions.

Also, in case this is really the same issue as this one reported in PEFT, it means that it used to work correctly in previous versions of accelerate/transformers. In that case, I wonder what changed that resulted in the issue.

[BenjaminBossan]

Would it make sense to move these checks to the very start of the method?

[eljandoubi]

@BenjaminBossan The method! Do you mean .prepare? What are the benefits of doing so?

[BenjaminBossan]

I meant _prepare_fsdp. It is a common pattern to perform all checks as early as possible, so following this makes the code easier to understand for readers. This is especially so if there are early returns.

The reason why it's common is so that we don't do any unnecessary work if the checks fail anyway. In this case, there is no need to determine models or optimizers if we're going to raise an error later. By skipping the unnecessary work, we ensure faster execution and prevent possibly unwanted side-effects (this might not be relevant here right now but code will change in the future and then it could be true).

[BenjaminBossan]

Is it better to call group["params"].clear()? That would affect references.

[eljandoubi]

@BenjaminBossan

• You are right; the solution needs more thorough testing. I have conducted a small test in a multi-GPU environment, but it should be tested against other configurations and models.
• Raising an error when the optimizer is initialized prematurely could be a solution, but two cases emerge:
  1. If Accelerate initializes the optimizer internally, either it creates a default optimizer with limited flexibility, or the user provides optimizer specifications. This would require external special treatment for FSDP, causing Accelerate to lose its unified interface for distributed training.
  2. If the user initializes the optimizer post-FSDP setup, Accelerate also loses generalization in this case.
• For PEFT training, which involves mixing frozen and non-frozen parameters, use_orig_params=True must be used. This ensures non-flattened parameters, meaning no changes are needed in the optimizer (I assume).
• The Transformers Trainer internally uses delay_optimizer_creation and creates the optimizer after FSDP wrapping.

[BenjaminBossan]

Thanks for the updates and explaining your testing and reasoning further.

Raising an error when the optimizer is initialized prematurely could be a solution, but two cases emerge

These are valid concerns. At the end of the day, it's a trade off and we need to decide which cost we'd rather pay.

The Transformers Trainer internally uses delay_optimizer_creation and creates the optimizer after FSDP wrapping.

I wonder if this logic could be moved/copied to accelerate.

For PEFT training, which involves mixing frozen and non-frozen parameters, use_orig_params=True must be used

I don't think this is strictly necessary, we have examples in PEFT with use_orig_params=False. But I have to admit I don't know what exactly changes under the hood in FSDP when this parameter is set. Note also that in the linked PEFT issue, I tried setting use_orig_params=True and it didn't help.

[BenjaminBossan]

I meant _prepare_fsdp. It is a common pattern to perform all checks as early as possible, so following this makes the code easier to understand for readers. This is especially so if there are early returns.

The reason why it's common is so that we don't do any unnecessary work if the checks fail anyway. In this case, there is no need to determine models or optimizers if we're going to raise an error later. By skipping the unnecessary work, we ensure faster execution and prevent possibly unwanted side-effects (this might not be relevant here right now but code will change in the future and then it could be true).

[eljandoubi]

• The PyTorch FSDP documentation mentions that mixed frozen and non-frozen parameters are only supported when use_orig_params=True. I believe this is due to the flattening process — we cannot flatten parameters with gradients together with those without. My contribution does not handle the case where use_orig_params=True because, in that scenario, the parameters remain unchanged.
• Another solution would be to adapt an approach similar to the DeepSpeed integration. This involves using a single model and optimizer, and after wrapping the model with model = FSDP(model), we create the optimizer as optimizer = optimizer_class(model.parameters(), **defaults). See Select the DeepSpeedCPUOptimizer based on the original optimizer class. #3255

[BenjaminBossan]

The PyTorch FSDP documentation mentions that mixed frozen and non-frozen parameters are only supported when use_orig_params=True. I believe this is due to the flattening process — we cannot flatten parameters with gradients together with those without.

This is true but using the right FSDP auto wrap policy, trainable and frozen parameters should be prevented from being mixed.

Another solution would be to adapt an approach similar to the DeepSpeed integration.

If this is an option here, it sounds like the more robust solution to me. Let's wait for @muellerzr return to office and get his opinion on this.