feat: add finetune_multi_device recipe with fsdp support

the HF SFTTrainer supports distributed training using FSDP.

Add a new recipe, `finetune_multi_device` which supports multi-GPU (cuda) training
using FSDP and optionally LoRA.

transformers hides _alot_ of their usage of FSDP behind the training args:
a6b51e7341/src/transformers/training_args.py (L1535)

you need to pass both `fsdp` and `fsdp_config` to get it to work properly. However,
it seems many of the `fsdp_config` entries are silently ignored. The key things to get this working were:
full_shard
offload (cpu offload)
transformer_layer_cls_to_wrap (model specific wrapping)
cpu_ram_efficient_loading
sharding_strategy
limit_all_gathers
sync_module_states
backward_prefetch
use_orig_params

these can be seen both in `fsdp=` and `fsdp_config=` int he `SFTConfig` call.

I have tested this with different model architectures with and without LoRA with success.

the user can now toggle `recipe` in their provider config between `single` and `multi` to access the two different recipes.

for debugging purposes NCCL logging settings can now be accessed via the provider config as well

Signed-off-by: Charlie Doern <cdoern@redhat.com>
This commit is contained in:
Charlie Doern 2025-05-19 13:21:35 -04:00
parent 35c2817d0a
commit 6494658a10
5 changed files with 1146 additions and 2 deletions

View file

@ -91,6 +91,7 @@ providers:
checkpoint_format: huggingface
distributed_backend: null
device: cpu
recipe: single
tool_runtime:
- provider_id: brave-search
provider_type: remote::brave-search

View file

@ -89,6 +89,7 @@ providers:
checkpoint_format: huggingface
distributed_backend: null
device: cpu
recipe: single
tool_runtime:
- provider_id: brave-search
provider_type: remote::brave-search