cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: restore streaming with device_map=auto and max_memory

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 01:41:51 -04:00
parent 5018be86ec
commit 0878b11e98
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
quantize_streaming.py
View File

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