NeoChen1024 commited on
Commit
44215bf
·
verified ·
1 Parent(s): 94c4ba6

Upload folder using huggingface_hub

Browse files
Files changed (1) hide show
  1. quantize.py +4 -10
quantize.py CHANGED
@@ -1,7 +1,8 @@
1
  from transformers import AutoProcessor, LlavaForConditionalGeneration
2
 
3
  from llmcompressor.modifiers.quantization import QuantizationModifier
4
- from llmcompressor.transformers import oneshot
 
5
 
6
  MODEL_ID = "llama-joycaption-beta-one-hf-llava"
7
 
@@ -17,17 +18,10 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
17
  recipe = QuantizationModifier(
18
  targets="Linear",
19
  scheme="FP8_DYNAMIC",
20
- ignore=["re:.*lm_head", "re:multi_modal_projector.*", "re:vision_tower.*"],
21
  )
22
 
23
  # Apply quantization and save to disk in compressed-tensors format.
24
  SAVE_DIR = MODEL_ID + "-FP8-Dynamic"
25
- oneshot(model=model, recipe=recipe, output_dir=SAVE_DIR)
26
  processor.save_pretrained(SAVE_DIR)
27
-
28
- # Confirm generations of the quantized model look sane.
29
- print("========== SAMPLE GENERATION ==============")
30
- input_ids = processor(text="Hello my name is", return_tensors="pt").input_ids.to("cuda")
31
- output = model.generate(input_ids, max_new_tokens=20)
32
- print(processor.decode(output[0]))
33
- print("==========================================")
 
1
  from transformers import AutoProcessor, LlavaForConditionalGeneration
2
 
3
  from llmcompressor.modifiers.quantization import QuantizationModifier
4
+ from llmcompressor import oneshot
5
+ from llmcompressor.utils import dispatch_for_generation
6
 
7
  MODEL_ID = "llama-joycaption-beta-one-hf-llava"
8
 
 
18
  recipe = QuantizationModifier(
19
  targets="Linear",
20
  scheme="FP8_DYNAMIC",
21
+ ignore=["re:.*lm_head", "re:.*multi_modal_projector.*", "re:.*vision_tower.*"],
22
  )
23
 
24
  # Apply quantization and save to disk in compressed-tensors format.
25
  SAVE_DIR = MODEL_ID + "-FP8-Dynamic"
26
+ oneshot(model=model, recipe=recipe, output_dir=SAVE_DIR, save_compressed=True)
27
  processor.save_pretrained(SAVE_DIR)