Qwen3-VL-4B-Instruct-per-grp-quant / run_qwen3_vl_4b_quant_model.py

Adding the quantized model files

ba0ba68 8 months ago

9.77 kB

import sys
import os
import torch
import copy
import argparse

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from transformers import AutoConfig, AutoTokenizer, AutoProcessor, PreTrainedTokenizer

from quark.torch import (
    LLMTemplate,
    ModelQuantizer,
    export_safetensors,
    import_model_from_safetensors,
)

from quark.torch.utils.llm import (
    get_calib_dataloader,
)

# Import the correct model class
try:
    from transformers import Qwen3VLForConditionalGeneration
    model_class = Qwen3VLForConditionalGeneration
    print("Using Qwen3VLForConditionalGeneration")
except ImportError:
    print("Failed to load the model using Qwen3_VLForConditionalGeneration")


def parse_args():
    parser = argparse.ArgumentParser(description='Quantize Qwen3-VL model with various configurations')
    
    # Model paths
    parser.add_argument('--model_dir', type=str, default='/scratch/dwchenna/github/hf-models/Qwen3-VL-4B-Instruct',
                        help='Path to the input model directory')
    parser.add_argument('--output_dir', type=str, default='quantized_models/Qwen3-VL-4B-Instruct-per-grp-quant',
                        help='Output directory for quantized model')
    
    # Evaluation options
    parser.add_argument('--eval_float', action='store_true', default=False,
                        help='Evaluate original float model')
    parser.add_argument('--eval_quantized', action='store_true', default=False,
                        help='Evaluate quantized model')
    parser.add_argument('--skip_quantization', action='store_true',
                        help='Skip quantization and only evaluate existing models')
    
    # Device
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                        help='Device to use for computation')
    
    return parser.parse_args()


args = parse_args()

# Setup
model_dir = args.model_dir
model_out_dir = args.output_dir
device = args.device

print(f"Device: {device}")
print("Loading model...")

# Load config, tokenizer, processor
config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)

# Load model
model = model_class.from_pretrained(
    model_dir, 
    trust_remote_code=True, 
    torch_dtype=torch.bfloat16,
    device_map="cpu"
).to(device)

print(f"Model loaded: {model.__class__.__name__}")
model.eval()

# Create a deep copy before quantization to preserve the original float model
float_model = copy.deepcopy(model)
# Filter out visual part
# del model.visual
print(model)

# Create output directory
os.makedirs(model_out_dir, exist_ok=True)
tokenizer.save_pretrained(model_out_dir)

# Get model type
model_config_type = (
    model.config.model_type if hasattr(model.config, "model_type") else model.config.architectures[0]
)
print(f"Model config type: {model_config_type}")

# Register template for qwen3_vl if needed
if model_config_type == "qwen3_vl":
    qwen3_vl_template = LLMTemplate(
        model_type="qwen3_vl",
        kv_layers_name=["model.layers.*.self_attn.q_proj", 
                       "model.layers.*.self_attn.k_proj", 
                       "model.layers.*.self_attn.v_proj"],
        q_layer_name="model.layers.*.self_attn.q_proj",
        exclude_layers_name=["visual*", "*vision*"],
    )
    LLMTemplate.register_template(qwen3_vl_template)
    print(f"Registered template for '{qwen3_vl_template.model_type}'")

# Check if model type is supported
if model_config_type not in LLMTemplate.list_available():
    print(f"Available templates: {LLMTemplate.list_available()}")
    raise ValueError(f"Model type '{model_config_type}' is not supported.")

template = LLMTemplate.get(model_config_type)
print(f"Using template: {model_config_type}")

# Quantization configuration
quant_scheme = "uint4_wo_128"
quant_algo = None # "awq"
exclude_layers = ["visual*", "*vision*"]

quant_config = template.get_config(
    scheme=quant_scheme,
    algorithm=quant_algo,
    exclude_layers=exclude_layers,
)

print(f"Quantization config: {quant_config}")

# Quantization section
if not args.skip_quantization:
    # Create calibration dataloader
    print("Loading calibration dataset...")
    main_device = model.device
    dataset = "pileval_for_awq_benchmark"
    batch_size = 1
    num_calib_data = 128
    seq_len = 512

    calib_dataloader = get_calib_dataloader(
        dataset_name=dataset,
        processor=None,  # Set to None to avoid multimodal issues
        tokenizer=tokenizer,
        batch_size=batch_size,
        num_calib_data=num_calib_data,
        seqlen=seq_len,
        device=main_device,
    )

    # Quantize model
    print("Starting quantization...")
    try:
        quantizer = ModelQuantizer(quant_config)
        model = quantizer.quantize_model(model, calib_dataloader)
        
        # Freeze model
        model = quantizer.freeze(model)
        print("✓ Quantization completed successfully!")
        
        # Export quantized model
        print("Exporting quantized model...")
        with torch.no_grad():
            export_safetensors(
                model=model,
                output_dir=model_out_dir,
                custom_mode="quark",
                weight_format="real_quantized",
                pack_method="reorder",
            )
        
        print(f"✓ Model exported to: {model_out_dir}")
        
    except Exception as e:
        print(f"✗ Quantization failed: {e}")
        print("This is likely due to AWQ not being compatible with multimodal models.")
        print("Try using a simpler quantization scheme without AWQ.")

    print("Quantization script completed!")
else:
    print("[INFO]: Skipping quantization as requested")

# Evaluation section
print("\n" + "="*60)
print("MODEL EVALUATION")
print("="*60)

if args.eval_float or args.eval_quantized:
    # Load quantized model if needed for evaluation
    model_quant = None
    if args.eval_quantized:
        print("Loading quantized model...")
        try:
            # Load a fresh model instance first
            fresh_model = model_class.from_pretrained(
                model_dir,  # Use original model dir for architecture
                trust_remote_code=True, 
                torch_dtype=torch.bfloat16,
                device_map="cpu"
            )
            # Load quantized weights using Quark's import function
            model_quant = import_model_from_safetensors(
                model=fresh_model,  
                model_dir=model_out_dir  # Directory with quantized safetensors
            )
            model_quant = model_quant.to(device)
            print("Successfully loaded quantized model from safetensors")
            print(f"Model type: {type(model_quant)}")
        except Exception as e:
            print(f"✗ Failed to load from {model_out_dir}: {e}")
            if not args.skip_quantization:
                print("Using the quantized model directly from memory...")
                model_quant = model
            else:
                print("[WARNING]: No quantized model available for evaluation")
                args.eval_quantized = False

    # Load the evaluation dataset
    from datasets import Dataset, load_dataset
    testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
    
    # Use tokenizer from model directory or output directory
    tokenizer_path = model_out_dir if os.path.exists(model_out_dir) else model_dir
    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(  # type: ignore
        tokenizer_path,
        trust_remote_code=True,
    )
    
    # Load the test dataset
    testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")
    
    # Results storage
    results = {}
    
    from quark.contrib.llm_eval import ppl_eval
    
    # Evaluate original float model
    if args.eval_float:
        print("\n[INFO]: Evaluating original float model...")
        print("-" * 40)
        try:
            main_device = float_model.device
            ppl_float = ppl_eval(float_model, testenc, main_device)
            ppl_float_value = ppl_float.item()
            results['float_ppl'] = ppl_float_value
            print(f"[RESULT] Float Model Perplexity: {ppl_float_value:.4f}")
        except Exception as e:
            print(f"[ERROR] Float model evaluation failed: {e}")
            results['float_ppl'] = None
    
    # Evaluate quantized model
    if args.eval_quantized and model_quant is not None:
        print("\n[INFO]: Evaluating quantized model...")
        print("-" * 40)
        try:
            main_device = model_quant.device
            ppl_quant = ppl_eval(model_quant, testenc, main_device)
            ppl_quant_value = ppl_quant.item()
            results['quantized_ppl'] = ppl_quant_value
            print(f"[RESULT] Quantized Model Perplexity: {ppl_quant_value:.4f}")
        except Exception as e:
            print(f"[ERROR] Quantized model evaluation failed: {e}")
            results['quantized_ppl'] = None
    
    # Summary
    print("\n" + "="*60)
    print("EVALUATION SUMMARY")
    print("="*60)
    print(f"Model: {model_dir}")
    print(f"Output: {model_out_dir}")
    
    if 'float_ppl' in results and results['float_ppl'] is not None:
        print(f"Float Model Perplexity: {results['float_ppl']:.4f}")
    
    if 'quantized_ppl' in results and results['quantized_ppl'] is not None:
        print(f"Quantized Model Perplexity: {results['quantized_ppl']:.4f}")
    
    print("="*60)
else:
    print("\n[INFO]: Skipping evaluation as requested")

print("\n[INFO]: Script completed successfully!")

	import sys
	import os
	import torch
	import copy
	import argparse

	sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
	from transformers import AutoConfig, AutoTokenizer, AutoProcessor, PreTrainedTokenizer

	from quark.torch import (
	LLMTemplate,
	ModelQuantizer,
	export_safetensors,
	import_model_from_safetensors,
	)

	from quark.torch.utils.llm import (
	get_calib_dataloader,
	)

	# Import the correct model class
	try:
	from transformers import Qwen3VLForConditionalGeneration
	model_class = Qwen3VLForConditionalGeneration
	print("Using Qwen3VLForConditionalGeneration")
	except ImportError:
	print("Failed to load the model using Qwen3_VLForConditionalGeneration")


	def parse_args():
	parser = argparse.ArgumentParser(description='Quantize Qwen3-VL model with various configurations')

	# Model paths
	parser.add_argument('--model_dir', type=str, default='/scratch/dwchenna/github/hf-models/Qwen3-VL-4B-Instruct',
	help='Path to the input model directory')
	parser.add_argument('--output_dir', type=str, default='quantized_models/Qwen3-VL-4B-Instruct-per-grp-quant',
	help='Output directory for quantized model')

	# Evaluation options
	parser.add_argument('--eval_float', action='store_true', default=False,
	help='Evaluate original float model')
	parser.add_argument('--eval_quantized', action='store_true', default=False,
	help='Evaluate quantized model')
	parser.add_argument('--skip_quantization', action='store_true',
	help='Skip quantization and only evaluate existing models')

	# Device
	parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
	help='Device to use for computation')

	return parser.parse_args()


	args = parse_args()

	# Setup
	model_dir = args.model_dir
	model_out_dir = args.output_dir
	device = args.device

	print(f"Device: {device}")
	print("Loading model...")

	# Load config, tokenizer, processor
	config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
	tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
	processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)

	# Load model
	model = model_class.from_pretrained(
	model_dir,
	trust_remote_code=True,
	torch_dtype=torch.bfloat16,
	device_map="cpu"
	).to(device)

	print(f"Model loaded: {model.__class__.__name__}")
	model.eval()

	# Create a deep copy before quantization to preserve the original float model
	float_model = copy.deepcopy(model)
	# Filter out visual part
	# del model.visual
	print(model)

	# Create output directory
	os.makedirs(model_out_dir, exist_ok=True)
	tokenizer.save_pretrained(model_out_dir)

	# Get model type
	model_config_type = (
	model.config.model_type if hasattr(model.config, "model_type") else model.config.architectures[0]
	)
	print(f"Model config type: {model_config_type}")

	# Register template for qwen3_vl if needed
	if model_config_type == "qwen3_vl":
	qwen3_vl_template = LLMTemplate(
	model_type="qwen3_vl",
	kv_layers_name=["model.layers.*.self_attn.q_proj",
	"model.layers.*.self_attn.k_proj",
	"model.layers.*.self_attn.v_proj"],
	q_layer_name="model.layers.*.self_attn.q_proj",
	exclude_layers_name=["visual", "vision*"],
	)
	LLMTemplate.register_template(qwen3_vl_template)
	print(f"Registered template for '{qwen3_vl_template.model_type}'")

	# Check if model type is supported
	if model_config_type not in LLMTemplate.list_available():
	print(f"Available templates: {LLMTemplate.list_available()}")
	raise ValueError(f"Model type '{model_config_type}' is not supported.")

	template = LLMTemplate.get(model_config_type)
	print(f"Using template: {model_config_type}")

	# Quantization configuration
	quant_scheme = "uint4_wo_128"
	quant_algo = None # "awq"
	exclude_layers = ["visual", "vision*"]

	quant_config = template.get_config(
	scheme=quant_scheme,
	algorithm=quant_algo,
	exclude_layers=exclude_layers,
	)

	print(f"Quantization config: {quant_config}")

	# Quantization section
	if not args.skip_quantization:
	# Create calibration dataloader
	print("Loading calibration dataset...")
	main_device = model.device
	dataset = "pileval_for_awq_benchmark"
	batch_size = 1
	num_calib_data = 128
	seq_len = 512

	calib_dataloader = get_calib_dataloader(
	dataset_name=dataset,
	processor=None, # Set to None to avoid multimodal issues
	tokenizer=tokenizer,
	batch_size=batch_size,
	num_calib_data=num_calib_data,
	seqlen=seq_len,
	device=main_device,
	)

	# Quantize model
	print("Starting quantization...")
	try:
	quantizer = ModelQuantizer(quant_config)
	model = quantizer.quantize_model(model, calib_dataloader)

	# Freeze model
	model = quantizer.freeze(model)
	print("✓ Quantization completed successfully!")

	# Export quantized model
	print("Exporting quantized model...")
	with torch.no_grad():
	export_safetensors(
	model=model,
	output_dir=model_out_dir,
	custom_mode="quark",
	weight_format="real_quantized",
	pack_method="reorder",
	)

	print(f"✓ Model exported to: {model_out_dir}")

	except Exception as e:
	print(f"✗ Quantization failed: {e}")
	print("This is likely due to AWQ not being compatible with multimodal models.")
	print("Try using a simpler quantization scheme without AWQ.")

	print("Quantization script completed!")
	else:
	print("[INFO]: Skipping quantization as requested")

	# Evaluation section
	print("\n" + "="*60)
	print("MODEL EVALUATION")
	print("="*60)

	if args.eval_float or args.eval_quantized:
	# Load quantized model if needed for evaluation
	model_quant = None
	if args.eval_quantized:
	print("Loading quantized model...")
	try:
	# Load a fresh model instance first
	fresh_model = model_class.from_pretrained(
	model_dir, # Use original model dir for architecture
	trust_remote_code=True,
	torch_dtype=torch.bfloat16,
	device_map="cpu"
	)
	# Load quantized weights using Quark's import function
	model_quant = import_model_from_safetensors(
	model=fresh_model,
	model_dir=model_out_dir # Directory with quantized safetensors
	)
	model_quant = model_quant.to(device)
	print("Successfully loaded quantized model from safetensors")
	print(f"Model type: {type(model_quant)}")
	except Exception as e:
	print(f"✗ Failed to load from {model_out_dir}: {e}")
	if not args.skip_quantization:
	print("Using the quantized model directly from memory...")
	model_quant = model
	else:
	print("[WARNING]: No quantized model available for evaluation")
	args.eval_quantized = False

	# Load the evaluation dataset
	from datasets import Dataset, load_dataset
	testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")

	# Use tokenizer from model directory or output directory
	tokenizer_path = model_out_dir if os.path.exists(model_out_dir) else model_dir
	tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained( # type: ignore
	tokenizer_path,
	trust_remote_code=True,
	)

	# Load the test dataset
	testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")

	# Results storage
	results = {}

	from quark.contrib.llm_eval import ppl_eval

	# Evaluate original float model
	if args.eval_float:
	print("\n[INFO]: Evaluating original float model...")
	print("-" * 40)
	try:
	main_device = float_model.device
	ppl_float = ppl_eval(float_model, testenc, main_device)
	ppl_float_value = ppl_float.item()
	results['float_ppl'] = ppl_float_value
	print(f"[RESULT] Float Model Perplexity: {ppl_float_value:.4f}")
	except Exception as e:
	print(f"[ERROR] Float model evaluation failed: {e}")
	results['float_ppl'] = None

	# Evaluate quantized model
	if args.eval_quantized and model_quant is not None:
	print("\n[INFO]: Evaluating quantized model...")
	print("-" * 40)
	try:
	main_device = model_quant.device
	ppl_quant = ppl_eval(model_quant, testenc, main_device)
	ppl_quant_value = ppl_quant.item()
	results['quantized_ppl'] = ppl_quant_value
	print(f"[RESULT] Quantized Model Perplexity: {ppl_quant_value:.4f}")
	except Exception as e:
	print(f"[ERROR] Quantized model evaluation failed: {e}")
	results['quantized_ppl'] = None

	# Summary
	print("\n" + "="*60)
	print("EVALUATION SUMMARY")
	print("="*60)
	print(f"Model: {model_dir}")
	print(f"Output: {model_out_dir}")

	if 'float_ppl' in results and results['float_ppl'] is not None:
	print(f"Float Model Perplexity: {results['float_ppl']:.4f}")

	if 'quantized_ppl' in results and results['quantized_ppl'] is not None:
	print(f"Quantized Model Perplexity: {results['quantized_ppl']:.4f}")

	print("="*60)
	else:
	print("\n[INFO]: Skipping evaluation as requested")

	print("\n[INFO]: Script completed successfully!")