From 4eb06a625e6729c2ff25177646a55f0aa8cfb8f2 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 09:05:00 -0400
Subject: [PATCH] fix: add bnb_4bit_quant_storage for FSDP sharding of 4-bit
 weights

---
 train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train.py b/train.py
index c32c05f..0cee768 100644
--- 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"],