nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 22 Packages Projects Releases Wiki Activity

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

Browse Source 
* 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.
This commit is contained in:
Daniel Bevenius
2025-12-23 14:07:25 +01:00
committed by GitHub
parent 12ee1763a6
commit 8e3ead6e4d
2 changed files with 154 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/model-conversion/Makefile
+3 -1
View File
@@ -25,6 +25,8 @@ define quantize_model
@echo "Export the quantized model path to $(2) variable in your environment" @echo "Export the quantized model path to $(2) variable in your environment"
endef endef
DEVICE ?= auto
### ###
### Casual Model targets/recipes ### Casual Model targets/recipes
### ###
@@ -53,7 +55,7 @@ causal-convert-mm-model:
causal-run-original-model: causal-run-original-model:
$(call validate_model_path,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: causal-run-converted-model:
@CONVERTED_MODEL="$(CONVERTED_MODEL)" ./scripts/causal/run-converted-model.sh @CONVERTED_MODEL="$(CONVERTED_MODEL)" ./scripts/causal/run-converted-model.sh
examples/model-conversion/scripts/causal/run-org-model.py
+94 -64
View File
@@ -4,53 +4,53 @@ import argparse
import os import os
import sys import sys
import importlib import importlib
import torch
import numpy as np
from pathlib import Path from pathlib import Path
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, AutoConfig
# Add parent directory to path for imports # Add parent directory to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) 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 from utils.common import debug_hook
parser = argparse.ArgumentParser(description="Process model with specified path") def parse_arguments():
parser.add_argument("--model-path", "-m", help="Path to the model") parser = argparse.ArgumentParser(description="Process model with specified path")
parser.add_argument("--prompt-file", "-f", help="Optional prompt file", required=False) parser.add_argument("--model-path", "-m", help="Path to the model")
parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose debug output") parser.add_argument("--prompt-file", "-f", help="Optional prompt file", required=False)
args = parser.parse_args() parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose debug output")
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) def load_model_and_tokenizer(model_path, device="auto"):
if model_path is None: print("Loading model and tokenizer using AutoTokenizer:", model_path)
parser.error( tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
"Model path must be specified either via --model-path argument or MODEL_PATH environment variable" config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
) multimodal = False
full_config = config
### If you want to dump RoPE activations, uncomment the following lines: # Determine device_map based on device argument
### === START ROPE DEBUG === if device == "cpu":
# from utils.common import setup_rope_debug device_map = {"": "cpu"}
# setup_rope_debug("transformers.models.apertus.modeling_apertus") print("Forcing CPU usage")
### == END ROPE DEBUG === elif device == "auto":
device_map = "auto"
else:
device_map = {"": device}
print("Model type: ", config.model_type)
print("Loading model and tokenizer using AutoTokenizer:", model_path) if "vocab_size" not in config and "text_config" in config:
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
print("Model type: ", config.model_type)
if "vocab_size" not in config and "text_config" in config:
config = config.text_config config = config.text_config
multimodal = True multimodal = True
print("Vocab size: ", config.vocab_size)
print("Hidden size: ", config.hidden_size)
print("Number of layers: ", config.num_hidden_layers)
print("BOS token id: ", config.bos_token_id)
print("EOS token id: ", config.eos_token_id)
unreleased_model_name = os.getenv("UNRELEASED_MODEL_NAME") print("Vocab size: ", config.vocab_size)
if unreleased_model_name: print("Hidden size: ", config.hidden_size)
print("Number of layers: ", config.num_hidden_layers)
print("BOS token id: ", config.bos_token_id)
print("EOS token id: ", config.eos_token_id)
unreleased_model_name = os.getenv("UNRELEASED_MODEL_NAME")
if unreleased_model_name:
model_name_lower = unreleased_model_name.lower() model_name_lower = unreleased_model_name.lower()
unreleased_module_path = ( unreleased_module_path = (
f"transformers.models.{model_name_lower}.modular_{model_name_lower}" f"transformers.models.{model_name_lower}.modular_{model_name_lower}"
@@ -59,54 +59,81 @@ if unreleased_model_name:
print(f"Importing unreleased model module: {unreleased_module_path}") print(f"Importing unreleased model module: {unreleased_module_path}")
try: 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 = 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: except (ImportError, AttributeError) as e:
print(f"Failed to import or load model: {e}") print(f"Failed to import or load model: {e}")
exit(1) exit(1)
else: else:
if multimodal: if multimodal:
model = AutoModelForImageTextToText.from_pretrained( 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: else:
model = AutoModelForCausalLM.from_pretrained( 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(): for name, module in model.named_modules():
if len(list(module.children())) == 0: # only leaf modules if len(list(module.children())) == 0: # only leaf modules
module.register_forward_hook(debug_hook(name)) 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 if args.prompt_file:
# 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: with open(args.prompt_file, encoding='utf-8') as f:
prompt = f.read() return f.read()
elif os.getenv("MODEL_TESTING_PROMPT"): elif os.getenv("MODEL_TESTING_PROMPT"):
prompt = os.getenv("MODEL_TESTING_PROMPT") return os.getenv("MODEL_TESTING_PROMPT")
else: else:
prompt = "Hello, my name is" return "Hello, my name is"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
print(f"Input tokens: {input_ids}") def main():
print(f"Input text: {repr(prompt)}") args = parse_arguments()
print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}") 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)
batch_size = 512
with torch.no_grad(): 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}")
print(f"Input text: {repr(prompt)}")
print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}")
batch_size = 512
with torch.no_grad():
past = None past = None
outputs = None outputs = None
for i in range(0, input_ids.size(1), batch_size): for i in range(0, input_ids.size(1), batch_size):
@@ -150,3 +177,6 @@ with torch.no_grad():
print(f"Saved bin logits to: {bin_filename}") print(f"Saved bin logits to: {bin_filename}")
print(f"Saved txt logist to: {txt_filename}") print(f"Saved txt logist to: {txt_filename}")
if __name__ == "__main__":
main()