model-conversion : add device option to run-org-model.py (#18318)

* model-conversion : add device option to run-org-model.py

This commit refactors the `run-org-model.py` script to include a
`--device` argument, to allow users to specify the device on which to
run the model (e.g., cpu, cuda, mps, auto).
It also extracts a few common functions to prepare for future changes
where some code duplication will be removed which there currently
exists in embedding scripts.

The Makefile is also been updated to pass the device argument, for
example:
```console
(venv) $ make causal-verify-logits DEVICE=cpu
```

* fix error handling and remove parser reference

This commit fixes the error handling which previously referenced an
undefined 'parser' variable.

examples/model-conversion/Makefile

@@ -25,6 +25,8 @@ define quantize_model
 	@echo "Export the quantized model path to $(2) variable in your environment"
 endef
 
+DEVICE ?= auto
+
 ###
 ### Casual Model targets/recipes
 ###
@@ -53,7 +55,7 @@ causal-convert-mm-model:
 
 causal-run-original-model:
 	$(call validate_model_path,causal-run-original-model)
-	@MODEL_PATH="$(MODEL_PATH)" ./scripts/causal/run-org-model.py
+	@MODEL_PATH="$(MODEL_PATH)" ./scripts/causal/run-org-model.py --device "$(DEVICE)"
 
 causal-run-converted-model:
 	@CONVERTED_MODEL="$(CONVERTED_MODEL)" ./scripts/causal/run-converted-model.sh

examples/model-conversion/scripts/causal/run-org-model.py

@@ -4,45 +4,45 @@ import argparse
 import os
 import sys
 import importlib
+import torch
+import numpy as np
+
 from pathlib import Path
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, AutoConfig
 
 # Add parent directory to path for imports
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, AutoConfig
-import torch
-import numpy as np
 from utils.common import debug_hook
 
+def parse_arguments():
     parser = argparse.ArgumentParser(description="Process model with specified path")
     parser.add_argument("--model-path", "-m", help="Path to the model")
     parser.add_argument("--prompt-file", "-f", help="Optional prompt file", required=False)
     parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose debug output")
-args = parser.parse_args()
+    parser.add_argument("--device", "-d", help="Device to use (cpu, cuda, mps, auto)", default="auto")
-
+    return parser.parse_args()
-model_path = os.environ.get("MODEL_PATH", args.model_path)
-if model_path is None:
-    parser.error(
-        "Model path must be specified either via --model-path argument or MODEL_PATH environment variable"
-    )
-
-### If you want to dump RoPE activations, uncomment the following lines:
-### === START ROPE DEBUG ===
-# from utils.common import setup_rope_debug
-# setup_rope_debug("transformers.models.apertus.modeling_apertus")
-### == END ROPE DEBUG ===
-
 
+def load_model_and_tokenizer(model_path, device="auto"):
     print("Loading model and tokenizer using AutoTokenizer:", model_path)
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
     multimodal = False
     full_config = config
 
+    # Determine device_map based on device argument
+    if device == "cpu":
+        device_map = {"": "cpu"}
+        print("Forcing CPU usage")
+    elif device == "auto":
+        device_map = "auto"
+    else:
+        device_map = {"": device}
+
     print("Model type:       ", config.model_type)
     if "vocab_size" not in config and "text_config" in config:
         config = config.text_config
         multimodal = True
+
     print("Vocab size:       ", config.vocab_size)
     print("Hidden size:      ", config.hidden_size)
     print("Number of layers: ", config.num_hidden_layers)
@@ -59,45 +59,72 @@ if unreleased_model_name:
         print(f"Importing unreleased model module: {unreleased_module_path}")
 
         try:
-        model_class = getattr(
+            model_class = getattr(importlib.import_module(unreleased_module_path), class_name)
-            importlib.import_module(unreleased_module_path), class_name
-        )
             model = model_class.from_pretrained(
-            model_path
+                    model_path,
-        )  # Note: from_pretrained, not fromPretrained
+                    device_map=device_map,
+                    offload_folder="offload",
+                    trust_remote_code=True,
+                    config=config
+            )
         except (ImportError, AttributeError) as e:
             print(f"Failed to import or load model: {e}")
             exit(1)
     else:
         if multimodal:
             model = AutoModelForImageTextToText.from_pretrained(
-            model_path, device_map="auto", offload_folder="offload", trust_remote_code=True, config=full_config
+                    model_path,
+                    device_map=device_map,
+                    offload_folder="offload",
+                    trust_remote_code=True,
+                    config=full_config
             )
         else:
             model = AutoModelForCausalLM.from_pretrained(
-            model_path, device_map="auto", offload_folder="offload", trust_remote_code=True, config=config
+                    model_path,
+                    device_map=device_map,
+                    offload_folder="offload",
+                    trust_remote_code=True,
+                    config=config
             )
 
-if args.verbose:
+    print(f"Model class: {model.__class__.__name__}")
+
+    return model, tokenizer, config
+
+def enable_torch_debugging(model):
         for name, module in model.named_modules():
             if len(list(module.children())) == 0:  # only leaf modules
                 module.register_forward_hook(debug_hook(name))
 
-model_name = os.path.basename(model_path)
+def get_prompt(args):
-# Printing the Model class to allow for easier debugging. This can be useful
-# when working with models that have not been publicly released yet and this
-# migth require that the concrete class is imported and used directly instead
-# of using AutoModelForCausalLM.
-print(f"Model class: {model.__class__.__name__}")
-
-device = next(model.parameters()).device
     if args.prompt_file:
         with open(args.prompt_file, encoding='utf-8') as f:
-        prompt = f.read()
+            return f.read()
     elif os.getenv("MODEL_TESTING_PROMPT"):
-    prompt = os.getenv("MODEL_TESTING_PROMPT")
+        return os.getenv("MODEL_TESTING_PROMPT")
     else:
-    prompt = "Hello, my name is"
+        return "Hello, my name is"
+
+def main():
+    args = parse_arguments()
+    model_path = os.environ.get("MODEL_PATH", args.model_path)
+    if model_path is None:
+        print("Error: Model path must be specified either via --model-path argument or MODEL_PATH environment variable")
+        sys.exit(1)
+
+
+    model, tokenizer, config = load_model_and_tokenizer(model_path, args.device)
+
+    if args.verbose:
+        enable_torch_debugging(model)
+
+    model_name = os.path.basename(model_path)
+
+    # Iterate over the model parameters (the tensors) and get the first one
+    # and use it to get the device the model is on.
+    device = next(model.parameters()).device
+    prompt = get_prompt(args)
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
     print(f"Input tokens: {input_ids}")
@@ -150,3 +177,6 @@ with torch.no_grad():
 
         print(f"Saved bin logits to: {bin_filename}")
         print(f"Saved txt logist to: {txt_filename}")
+
+if __name__ == "__main__":
+    main()