From 34ae0edd1fb963c3033f6b4cd219c75ce861740b Mon Sep 17 00:00:00 2001 From: Jscaldwell55 Date: Thu, 17 Jul 2025 15:41:29 -0500 Subject: [PATCH 1/2] feat: Enable QLoRA and FSDP compatibility --- src/llama_cookbook/finetuning.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/llama_cookbook/finetuning.py b/src/llama_cookbook/finetuning.py index 93929dde3..d4f0979c4 100644 --- a/src/llama_cookbook/finetuning.py +++ b/src/llama_cookbook/finetuning.py @@ -116,7 +116,7 @@ def main(**kwargs): if not train_config.enable_fsdp or rank == 0: wandb_run = setup_wandb(train_config, fsdp_config, **kwargs) - # setting quantization configs + # setting quantization configs bnb_config = None if train_config.quantization: if type(train_config.quantization) == type(True): @@ -135,6 +135,17 @@ def main(**kwargs): update_config(quant_config, **kwargs) bnb_config = quant_config.create_bnb_config(train_config.quantization) + + if train_config.enable_fsdp: + if train_config.quantization == "4bit": + bnb_config.bnb_4bit_quant_storage = bnb_config.bnb_4bit_compute_dtype + from logging import getLogger + logger = getLogger() + logger.warning( + "FSDP and 4-bit QLoRA enabled. Setting `bnb_4bit_quant_storage` " + f"to {bnb_config.bnb_4bit_compute_dtype} for compatibility." + ) + # Load the pre-trained model and setup its configuration use_cache = False if train_config.enable_fsdp else None config = AutoConfig.from_pretrained(train_config.model_name) @@ -264,10 +275,9 @@ def main(**kwargs): elif torch.cuda.is_available(): device_id = torch.cuda.current_device() - if train_config.freeze_LLM_only: - use_orig_params = True - else: - use_orig_params = False + # This single line robustly sets use_orig_params for both freeze_LLM_only and PEFT (LoRA/QLoRA) + use_orig_params = train_config.freeze_LLM_only or train_config.use_peft + model = FSDP( model, auto_wrap_policy=( From 7e17d897fc3848f34ddd614d8c8744cb183dbc0f Mon Sep 17 00:00:00 2001 From: Jay Caldwell <111952840+jscaldwell55@users.noreply.github.com> Date: Thu, 17 Jul 2025 16:13:31 -0500 Subject: [PATCH 2/2] Update finetuning.py Deleting comment --- src/llama_cookbook/finetuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama_cookbook/finetuning.py b/src/llama_cookbook/finetuning.py index d4f0979c4..05d5babc9 100644 --- a/src/llama_cookbook/finetuning.py +++ b/src/llama_cookbook/finetuning.py @@ -275,7 +275,7 @@ def main(**kwargs): elif torch.cuda.is_available(): device_id = torch.cuda.current_device() - # This single line robustly sets use_orig_params for both freeze_LLM_only and PEFT (LoRA/QLoRA) + use_orig_params = train_config.freeze_LLM_only or train_config.use_peft model = FSDP(