fix: add bnb_4bit_quant_storage for FSDP sharding of 4-bit weights

2026-07-02 09:05:00 -04:00
parent 4c9522072e
commit 4eb06a625e
1 changed files with 1 additions and 0 deletions
--- a/train.py
+++ b/train.py
@@ -47,6 +47,7 @@ def train(config_path):
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_storage=torch.bfloat16,  # Enable FSDP sharding of 4-bit weights
        )
        model = AutoModelForCausalLM.from_pretrained(
            config["base_model"],