Add missing quantize_base to llama 3.1

[felipemello1]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

Llama 3.1 component_builders were missing quantize_base. I noticed that when finetuning all layers, QLoRA size was ~6.2GB, but when i finetuned only KV, it would take much more memory. Thats because the layers are not quantized when we don't train them.

With this PR, the base model stays 6.2GB regardless of the layers you are finetuning.

Test plan

Compare LoRA vs QLoRA when finetuning all layers and only QV. Expected behavior is for memory: LoRA_all > LoRA_QV > QLoRA_all > QLoRA_KV

#llama_QLoRA_QV_only
tune run lora_finetune_single_device --config llama3_1/8B_qlora_single_device \
max_steps_per_epoch=25 \
epochs=1 \
model.lora_attn_modules=[q_proj,v_proj] \
model.apply_lora_to_mlp=False \
model.apply_lora_to_output=False \
metric_logger=torchtune.training.metric_logging.WandBLogger  \
metric_logger.name=llama_QLoRA_QV_only \
log_peak_memory_stats=True \
batch_size=4 \
gradient_accumulation_steps=1

#llama_LoRA_QV_only
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device \
max_steps_per_epoch=25 \
epochs=1  \
model.lora_attn_modules=[q_proj,v_proj] \
model.apply_lora_to_mlp=False \
model.apply_lora_to_output=False \
metric_logger=torchtune.training.metric_logging.WandBLogger  \
metric_logger.name=llama_LoRA_QV_only \
log_peak_memory_stats=True \
batch_size=4 \
gradient_accumulation_steps=1

#llama_QLoRA_all_layers
tune run lora_finetune_single_device --config llama3_1/8B_qlora_single_device \
max_steps_per_epoch=25 \
epochs=1  \
model.lora_attn_modules=[q_proj,v_proj,k_proj,output_proj] \
model.apply_lora_to_mlp=True \
model.apply_lora_to_output=False \
metric_logger=torchtune.training.metric_logging.WandBLogger  \
metric_logger.name=llama_QLoRA_all_layers \
log_peak_memory_stats=True \
batch_size=4 \
gradient_accumulation_steps=1

#llama_LoRA_all_layers
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device \
max_steps_per_epoch=25 \
epochs=1  \
model.lora_attn_modules=[q_proj,v_proj,k_proj,output_proj] \
model.apply_lora_to_mlp=True \
model.apply_lora_to_output=False \
metric_logger=torchtune.training.metric_logging.WandBLogger  \
metric_logger.name=llama_LoRA_all_layers \
log_peak_memory_stats=True \
batch_size=4 \
gradient_accumulation_steps=1

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1485

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 20796b8 with merge base 4dc1cd0 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.