chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,473 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Embedding Quantization Script
|
||||
This script converts ggml-model-f32.gguf to multiple quantized versions
|
||||
with different token embedding types.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import argparse
|
||||
import re
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class EmbeddingQuantizer:
|
||||
def __init__(self, input_model, output_dir, quantize_bin="../build/bin/llama-quantize",
|
||||
bench_bin="../build/bin/llama-bench", stats_dir="../stats", csv_output=None):
|
||||
self.input_model = Path(input_model)
|
||||
self.output_dir = Path(output_dir)
|
||||
self.quantize_bin = Path(quantize_bin)
|
||||
self.bench_bin = Path(bench_bin)
|
||||
self.stats_dir = Path(stats_dir)
|
||||
self.csv_output = Path(csv_output) if csv_output else None
|
||||
|
||||
# Verify input file exists
|
||||
if not self.input_model.exists():
|
||||
raise FileNotFoundError(f"Input model not found: {self.input_model}")
|
||||
|
||||
# Verify quantize tool exists
|
||||
if not self.quantize_bin.exists():
|
||||
raise FileNotFoundError(f"Quantize binary not found: {self.quantize_bin}")
|
||||
|
||||
# Verify bench tool exists
|
||||
if not self.bench_bin.exists():
|
||||
raise FileNotFoundError(f"Benchmark binary not found: {self.bench_bin}")
|
||||
|
||||
# Create output directories
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.stats_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.results = []
|
||||
self.newly_created_files = set() # Track newly created files
|
||||
|
||||
def quantize(self, embedding_type, output_suffix):
|
||||
"""
|
||||
Perform single quantization
|
||||
|
||||
Args:
|
||||
embedding_type: Token embedding type (uppercase format, e.g., Q6_K)
|
||||
output_suffix: Output file suffix (lowercase format, e.g., q6_k)
|
||||
|
||||
Returns:
|
||||
bool: Whether successful
|
||||
"""
|
||||
output_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
# Check if file already exists
|
||||
file_already_existed = output_file.exists()
|
||||
|
||||
if file_already_existed:
|
||||
print(f"ℹ️ File already exists: {output_file}")
|
||||
print(f" Skipping quantization, will use existing file for benchmark")
|
||||
return True
|
||||
|
||||
cmd = [
|
||||
str(self.quantize_bin),
|
||||
"--token-embedding-type", embedding_type,
|
||||
str(self.input_model),
|
||||
str(output_file),
|
||||
"I2_S",
|
||||
"1",
|
||||
"1"
|
||||
]
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🔄 Quantizing with embedding type: {embedding_type}")
|
||||
print(f"📥 Input: {self.input_model}")
|
||||
print(f"📤 Output: {output_file}")
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=600 # 10 minute timeout
|
||||
)
|
||||
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
if result.returncode == 0:
|
||||
# Get output file size
|
||||
file_size_mb = output_file.stat().st_size / (1024 * 1024)
|
||||
|
||||
print(f"✅ Success! Duration: {duration:.2f}s, Size: {file_size_mb:.2f} MB")
|
||||
|
||||
# Record newly created file
|
||||
if not file_already_existed:
|
||||
self.newly_created_files.add(output_file)
|
||||
|
||||
# Print part of output
|
||||
if result.stdout:
|
||||
print("\n📊 Quantization output:")
|
||||
print(result.stdout[-500:] if len(result.stdout) > 500 else result.stdout)
|
||||
|
||||
return True
|
||||
else:
|
||||
print(f"❌ Failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr}")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"❌ Timeout (exceeded 10 minutes)")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Exception: {e}")
|
||||
return False
|
||||
|
||||
def benchmark_model(self, output_suffix):
|
||||
"""
|
||||
Benchmark model
|
||||
|
||||
Args:
|
||||
output_suffix: Output file suffix (lowercase format, e.g., q6_k)
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with benchmark results, or None if failed
|
||||
"""
|
||||
model_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
if not model_file.exists():
|
||||
print(f"❌ Model file not found for benchmarking: {model_file}")
|
||||
return None
|
||||
|
||||
cmd = [
|
||||
str(self.bench_bin),
|
||||
"-m", str(model_file),
|
||||
"-p", "128",
|
||||
"-n", "0",
|
||||
"-t", "1,2,4,8",
|
||||
"-ngl", "0"
|
||||
]
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🏃 Running benchmark for: {output_suffix}")
|
||||
print(f"💻 Command: {' '.join(cmd)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.getcwd(),
|
||||
timeout=300 # 5 minute timeout
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
print("✅ Benchmark completed successfully")
|
||||
print("\n📊 Benchmark output:")
|
||||
print(result.stdout)
|
||||
|
||||
# 解析输出
|
||||
bench_results = self.parse_benchmark_output(result.stdout, output_suffix)
|
||||
return bench_results
|
||||
else:
|
||||
print(f"❌ Benchmark failed with return code {result.returncode}")
|
||||
print(f"Error: {result.stderr}")
|
||||
return None
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f"❌ Benchmark timeout (exceeded 5 minutes)")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Benchmark exception: {e}")
|
||||
return None
|
||||
|
||||
def parse_benchmark_output(self, output, output_suffix):
|
||||
"""
|
||||
Parse benchmark output to extract t/s data (mean±std)
|
||||
|
||||
Args:
|
||||
output: Benchmark command output
|
||||
output_suffix: Output file suffix
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with parsed results
|
||||
"""
|
||||
results = {
|
||||
'embedding_type': output_suffix,
|
||||
'threads_1': None,
|
||||
'threads_2': None,
|
||||
'threads_4': None,
|
||||
'threads_8': None,
|
||||
}
|
||||
|
||||
# Parse table data
|
||||
# Find lines containing pp128 and t/s
|
||||
lines = output.strip().split('\n')
|
||||
|
||||
for line in lines:
|
||||
# Skip header and separator lines
|
||||
if '|' not in line or 'model' in line or '---' in line:
|
||||
continue
|
||||
|
||||
# Try to extract data
|
||||
# Format similar to: | bitnet-25 2B I2_S - 2 bpw ternary | 1012.28 MiB | 2.74 B | CPU | 12 | pp128 | 405.73 ± 3.69 |
|
||||
parts = [p.strip() for p in line.split('|')]
|
||||
|
||||
if len(parts) >= 8 and 'pp128' in parts[6]:
|
||||
threads_str = parts[5].strip()
|
||||
throughput_str = parts[7].strip()
|
||||
|
||||
# Extract thread count
|
||||
try:
|
||||
threads = int(threads_str)
|
||||
except:
|
||||
continue
|
||||
|
||||
# Extract t/s data (format: "405.73 ± 3.69" or "405.73")
|
||||
# Try to match "mean ± std" format
|
||||
match_with_std = re.search(r'([\d.]+)\s*±\s*([\d.]+)', throughput_str)
|
||||
if match_with_std:
|
||||
mean = float(match_with_std.group(1))
|
||||
std = float(match_with_std.group(2))
|
||||
throughput = f"{mean:.2f}±{std:.2f}"
|
||||
else:
|
||||
# Only mean, no std
|
||||
match = re.search(r'([\d.]+)', throughput_str)
|
||||
if match:
|
||||
throughput = f"{float(match.group(1)):.2f}"
|
||||
else:
|
||||
continue
|
||||
|
||||
# Store result based on thread count
|
||||
if threads == 1:
|
||||
results['threads_1'] = throughput
|
||||
elif threads == 2:
|
||||
results['threads_2'] = throughput
|
||||
elif threads == 4:
|
||||
results['threads_4'] = throughput
|
||||
elif threads == 8:
|
||||
results['threads_8'] = throughput
|
||||
|
||||
return results
|
||||
|
||||
def cleanup_model(self, output_suffix):
|
||||
"""
|
||||
Cleanup model files (only delete newly created files)
|
||||
|
||||
Args:
|
||||
output_suffix: Output file suffix
|
||||
"""
|
||||
model_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf"
|
||||
|
||||
if model_file in self.newly_created_files:
|
||||
try:
|
||||
model_file.unlink()
|
||||
print(f"🗑️ Deleted newly created file: {model_file}")
|
||||
self.newly_created_files.remove(model_file)
|
||||
except Exception as e:
|
||||
print(f"⚠️ Failed to delete {model_file}: {e}")
|
||||
else:
|
||||
print(f"ℹ️ Keeping existing file: {model_file}")
|
||||
|
||||
def run_all_quantizations(self, types_to_quantize):
|
||||
"""
|
||||
Run all quantizations
|
||||
|
||||
Args:
|
||||
types_to_quantize: List of quantization types, tuples of (embedding_type, output_suffix)
|
||||
"""
|
||||
print(f"\n{'='*80}")
|
||||
print(f"🚀 Starting Embedding Quantization and Benchmarking")
|
||||
print(f"{'='*80}")
|
||||
print(f"📥 Input model: {self.input_model}")
|
||||
print(f"📤 Output directory: {self.output_dir}")
|
||||
print(f"📊 Stats directory: {self.stats_dir}")
|
||||
print(f"🔢 Total quantizations: {len(types_to_quantize)}")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
total_start = datetime.now()
|
||||
|
||||
for i, (embedding_type, output_suffix) in enumerate(types_to_quantize, 1):
|
||||
print(f"\n{'#'*80}")
|
||||
print(f"[{i}/{len(types_to_quantize)}] Processing {output_suffix} ({embedding_type})")
|
||||
print(f"{'#'*80}\n")
|
||||
|
||||
# Quantize model
|
||||
success = self.quantize(embedding_type, output_suffix)
|
||||
|
||||
if not success:
|
||||
print(f"⚠️ Skipping benchmark for {output_suffix} due to quantization failure")
|
||||
continue
|
||||
|
||||
# Run benchmark
|
||||
bench_results = self.benchmark_model(output_suffix)
|
||||
|
||||
if bench_results:
|
||||
self.results.append(bench_results)
|
||||
else:
|
||||
print(f"⚠️ Benchmark failed for {output_suffix}")
|
||||
|
||||
# Cleanup model files (only delete newly created files)
|
||||
self.cleanup_model(output_suffix)
|
||||
|
||||
print(f"\n{'#'*80}")
|
||||
print(f"✅ Completed {output_suffix}")
|
||||
print(f"{'#'*80}\n")
|
||||
|
||||
total_end = datetime.now()
|
||||
total_duration = (total_end - total_start).total_seconds()
|
||||
|
||||
# 保存结果到CSV
|
||||
self.save_results_to_csv()
|
||||
|
||||
# 打印总结
|
||||
self.print_summary(total_duration)
|
||||
|
||||
def save_results_to_csv(self):
|
||||
"""将benchmark结果保存到CSV文件"""
|
||||
if not self.results:
|
||||
print("⚠️ No results to save")
|
||||
return
|
||||
|
||||
# Use user-specified CSV path, otherwise use default path
|
||||
if self.csv_output:
|
||||
csv_file = self.csv_output
|
||||
# Ensure parent directory exists
|
||||
csv_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
csv_file = self.stats_dir / f"embedding_benchmark.csv"
|
||||
|
||||
print(f"\n💾 Saving results to: {csv_file}")
|
||||
|
||||
try:
|
||||
with open(csv_file, 'w', newline='') as f:
|
||||
fieldnames = ['embedding_type', 'threads_1', 'threads_2', 'threads_4', 'threads_8']
|
||||
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||
|
||||
writer.writeheader()
|
||||
for result in self.results:
|
||||
writer.writerow(result)
|
||||
|
||||
print(f"✅ Results saved successfully")
|
||||
|
||||
# Also print table
|
||||
print(f"\n📊 Benchmark Results:")
|
||||
print(f"{'Type':<15} {'1 thread':<18} {'2 threads':<18} {'4 threads':<18} {'8 threads':<18}")
|
||||
print("-" * 87)
|
||||
for result in self.results:
|
||||
t1 = result['threads_1'] if result['threads_1'] else "N/A"
|
||||
t2 = result['threads_2'] if result['threads_2'] else "N/A"
|
||||
t4 = result['threads_4'] if result['threads_4'] else "N/A"
|
||||
t8 = result['threads_8'] if result['threads_8'] else "N/A"
|
||||
print(f"{result['embedding_type']:<15} {t1:<18} {t2:<18} {t4:<18} {t8:<18}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to save results: {e}")
|
||||
|
||||
def print_summary(self, total_duration):
|
||||
"""Print quantization summary"""
|
||||
print(f"\n\n{'='*80}")
|
||||
print(f"📊 QUANTIZATION AND BENCHMARK SUMMARY")
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
successful = len(self.results)
|
||||
total = len(self.results)
|
||||
|
||||
print(f"✅ Completed: {successful} benchmarks")
|
||||
print(f"⏱️ Total duration: {total_duration/60:.2f} minutes\n")
|
||||
|
||||
if self.results:
|
||||
if self.csv_output and self.csv_output.exists():
|
||||
print(f"📁 Results saved to: {self.csv_output}")
|
||||
else:
|
||||
csv_files = list(self.stats_dir.glob("embedding_benchmark*.csv"))
|
||||
if csv_files:
|
||||
latest_csv = max(csv_files, key=lambda p: p.stat().st_mtime)
|
||||
print(f"📁 Results saved to: {latest_csv}")
|
||||
|
||||
print(f"\n{'='*80}\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Quantize model embeddings to multiple formats')
|
||||
parser.add_argument('--input', '-i',
|
||||
default='../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf',
|
||||
help='Input model path (default: ../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf)')
|
||||
parser.add_argument('--output-dir', '-o',
|
||||
default='../models/BitNet-b1.58-2B-4T',
|
||||
help='Output directory (default: ../models/BitNet-b1.58-2B-4T)')
|
||||
parser.add_argument('--quantize-bin', '-q',
|
||||
default='../build/bin/llama-quantize',
|
||||
help='Path to llama-quantize binary (default: ../build/bin/llama-quantize)')
|
||||
parser.add_argument('--bench-bin', '-b',
|
||||
default='../build/bin/llama-bench',
|
||||
help='Path to llama-bench binary (default: ../build/bin/llama-bench)')
|
||||
parser.add_argument('--stats-dir',
|
||||
default='../stats',
|
||||
help='Directory to save benchmark results (default: ../stats)')
|
||||
parser.add_argument('--csv-output', '-c',
|
||||
help='Custom path for CSV output file (e.g., stats/my_results.csv)')
|
||||
parser.add_argument('--types', '-t',
|
||||
nargs='+',
|
||||
help='Specific types to quantize (e.g., f32 q6_k q4_0)')
|
||||
parser.add_argument('--skip-existing', '-s',
|
||||
action='store_true',
|
||||
help='Skip quantization if output file already exists (will still benchmark existing files)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Define all supported quantization types
|
||||
# Format: (embedding_type for command line, output_suffix for filename)
|
||||
all_types = [
|
||||
('F32', 'f32'),
|
||||
('F16', 'f16'),
|
||||
('Q8_0', 'q8_0'),
|
||||
('Q6_K', 'q6_k'),
|
||||
('Q5_0', 'q5_0'),
|
||||
('Q4_0', 'q4_0'),
|
||||
('Q3_K', 'q3_k'),
|
||||
('TQ2_0', 'tq2_0'),
|
||||
]
|
||||
|
||||
# If specific types are specified, filter the list
|
||||
if args.types:
|
||||
types_lower = [t.lower() for t in args.types]
|
||||
types_to_quantize = [(et, os) for et, os in all_types if os.lower() in types_lower]
|
||||
if not types_to_quantize:
|
||||
print(f"❌ No valid types specified. Available types: {', '.join([os for _, os in all_types])}")
|
||||
return
|
||||
else:
|
||||
types_to_quantize = all_types
|
||||
|
||||
# If skip existing files is enabled, no need to filter
|
||||
# Because new logic will automatically detect and skip during quantization, but will still benchmark
|
||||
|
||||
# 创建量化器并运行
|
||||
try:
|
||||
quantizer = EmbeddingQuantizer(
|
||||
args.input,
|
||||
args.output_dir,
|
||||
args.quantize_bin,
|
||||
args.bench_bin,
|
||||
args.stats_dir,
|
||||
args.csv_output
|
||||
)
|
||||
quantizer.run_all_quantizations(types_to_quantize)
|
||||
except FileNotFoundError as e:
|
||||
print(f"❌ Error: {e}")
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n⚠️ Quantization interrupted by user")
|
||||
return 1
|
||||
except Exception as e:
|
||||
print(f"\n❌ Unexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main() or 0)
|
||||
Reference in New Issue
Block a user