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"],