llama-stack-mirror/llama_stack
Charlie Doern 6494658a10 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>
2025-06-12 13:33:33 -04:00
..
apis fix: set appropriate defaults for params (#2434) 2025-06-11 17:30:34 -07:00
cli fix: resolve template name to config path in llama stack run (#2361) 2025-06-03 14:39:12 -07:00
distribution ci: fix external provider test (#2438) 2025-06-12 16:14:32 +02:00
models chore: remove usage of load_tiktoken_bpe (#2276) 2025-06-02 07:33:37 -07:00
providers feat: add finetune_multi_device recipe with fsdp support 2025-06-12 13:33:33 -04:00
strong_typing chore: enable pyupgrade fixes (#1806) 2025-05-01 14:23:50 -07:00
templates feat: add finetune_multi_device recipe with fsdp support 2025-06-12 13:33:33 -04:00
ui build: Bump version to 0.2.10 2025-06-05 22:56:39 +00:00
__init__.py export LibraryClient 2024-12-13 12:08:00 -08:00
env.py refactor(test): move tools, evals, datasetio, scoring and post training tests (#1401) 2025-03-04 14:53:47 -08:00
log.py ci: fix external provider test (#2438) 2025-06-12 16:14:32 +02:00
schema_utils.py chore: enable pyupgrade fixes (#1806) 2025-05-01 14:23:50 -07:00