fix: restore streaming with device_map=auto and max_memory

This commit is contained in:
Christian Medina
2026-07-03 01:41:51 -04:00
parent 5018be86ec
commit 0878b11e98

View File

@@ -39,9 +39,10 @@ def streaming_quantize(model_path: str, output_path: str):
bnb_4bit_use_double_quant=True,
)
print("Loading model with BnB 4-bit (shards streamed to GPUs)...")
print(" This will distribute across both GPUs\n")
print("Loading model with BnB 4-bit (shards streamed via accelerate)...")
print(" This will stream shards and quantize them\n")
# This streams shards and applies BnB correctly
model = AutoModelForCausalLM.from_pretrained(
model_path,
quantization_config=bnb_config,