add dispatch utility

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

examples/quantization_w4a16/llama3_example.py

@@ -3,6 +3,7 @@
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.transformers import oneshot
+from llmcompressor.utils.dev import dispatch_for_generation
 
 # Select model and load it.
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
@@ -61,18 +62,17 @@ def tokenize(sample):
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
 )
 
-# Save to disk compressed.
-SAVE_DIR = model_id.split("/")[1] + "-W4A16-G128"
-model.save_pretrained(SAVE_DIR, save_compressed=True)
-tokenizer.save_pretrained(SAVE_DIR)
-
-# Load model after saving
-model = AutoModelForCausalLM.from_pretrained(SAVE_DIR, device_map="auto")
-
 # Confirm generations of the quantized model look sane.
 print("\n\n")
 print("========== SAMPLE GENERATION ==============")
-input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
-output = model.generate(input_ids, max_new_tokens=100)
+dispatch_for_generation(model)
+sample = tokenizer("Hello my name is", return_tensors="pt")
+sample = {key: value.to("cuda") for key, value in sample.items()}
+output = model.generate(**sample, max_new_tokens=100)
 print(tokenizer.decode(output[0]))
 print("==========================================\n\n")
+
+# Save to disk compressed.
+SAVE_DIR = model_id.split("/")[-1] + "-W4A16-G128"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)

src/llmcompressor/utils/dev.py

@@ -2,18 +2,20 @@
 import logging
 import os
 import tempfile
-from typing import Type
+from typing import Type, Dict, Any, Union
 
 import torch
 from huggingface_hub import snapshot_download
 from safetensors.torch import save_file
 from transformers import AutoModelForCausalLM, PreTrainedModel
 from transformers.modeling_utils import TORCH_INIT_FUNCTIONS
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, WEIGHTS_INDEX_NAME
+from accelerate import dispatch_model, infer_auto_device_map
+from accelerate.utils import get_balanced_memory
 
 from llmcompressor.utils.helpers import patch_attr
 
-__all__ = ["skip_weights_download", "patch_transformers_logger_level"]
+__all__ = ["skip_weights_download", "patch_transformers_logger_level", "dispatch_for_generation"]
 
 
 @contextlib.contextmanager
@@ -106,3 +108,14 @@ def patch_transformers_logger_level(level: int = logging.ERROR):
     transformers_logger.setLevel(level=level)
     yield
     transformers_logger.setLevel(level=restore_log_level)
+
+
+def dispatch_for_generation(model: PreTrainedModel) -> PreTrainedModel:
+    max_memory = get_balanced_memory(
+        model,
+        dtype=model.dtype,
+        no_split_module_classes=model._get_no_split_modules("auto")
+    )
+    device_map = infer_auto_device_map(model, dtype=model.dtype, max_memory=max_memory)
+
+    return dispatch_model(model, device_map=device_map)