b3e9ee96fa
- Fix ambiguous fullwidth characters (commas, parentheses) in strings and comments - Replace Chinese comments with English equivalents - Fix unused imports with proper noqa annotations for intentional imports - Fix bare except clauses with specific exception types - Fix redefined variables and undefined names - Add ruff noqa annotations for generated protobuf files - Add lint and format check to GitHub Actions CI pipeline
651 lines
23 KiB
Python
651 lines
23 KiB
Python
# python embedd_micro.py --use_int8 Fastest
|
|
|
|
import argparse
|
|
import time
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from tqdm import tqdm
|
|
from transformers import AutoModel, BitsAndBytesConfig
|
|
|
|
|
|
@dataclass
|
|
class BenchmarkConfig:
|
|
model_path: str
|
|
batch_sizes: list[int]
|
|
seq_length: int
|
|
num_runs: int
|
|
use_fp16: bool = True
|
|
use_int4: bool = False
|
|
use_int8: bool = False # Add this parameter
|
|
use_cuda_graphs: bool = False
|
|
use_flash_attention: bool = False
|
|
use_linear8bitlt: bool = False
|
|
|
|
|
|
class GraphContainer:
|
|
"""Container for managing graphs for different batch sizes (CUDA graphs on NVIDIA, regular on others)."""
|
|
|
|
def __init__(self, model: nn.Module, seq_length: int):
|
|
self.model = model
|
|
self.seq_length = seq_length
|
|
self.graphs: dict[int, GraphWrapper] = {}
|
|
|
|
def get_or_create(self, batch_size: int) -> "GraphWrapper":
|
|
if batch_size not in self.graphs:
|
|
self.graphs[batch_size] = GraphWrapper(self.model, batch_size, self.seq_length)
|
|
return self.graphs[batch_size]
|
|
|
|
|
|
class GraphWrapper:
|
|
"""Wrapper for graph capture and replay (CUDA graphs on NVIDIA, regular on others)."""
|
|
|
|
def __init__(self, model: nn.Module, batch_size: int, seq_length: int):
|
|
self.model = model
|
|
self.device = self._get_device()
|
|
self.static_input = self._create_random_batch(batch_size, seq_length)
|
|
self.static_attention_mask = torch.ones_like(self.static_input)
|
|
|
|
# Warm up
|
|
self._warmup()
|
|
|
|
# Only use CUDA graphs on NVIDIA GPUs
|
|
if torch.cuda.is_available() and hasattr(torch.cuda, "CUDAGraph"):
|
|
# Capture graph
|
|
self.graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(self.graph):
|
|
self.static_output = self.model(
|
|
input_ids=self.static_input, attention_mask=self.static_attention_mask
|
|
)
|
|
self.use_cuda_graph = True
|
|
else:
|
|
# For MPS or CPU, just store the model
|
|
self.use_cuda_graph = False
|
|
self.static_output = None
|
|
|
|
def _get_device(self) -> str:
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
elif torch.backends.mps.is_available():
|
|
return "mps"
|
|
else:
|
|
return "cpu"
|
|
|
|
def _create_random_batch(self, batch_size: int, seq_length: int) -> torch.Tensor:
|
|
return torch.randint(
|
|
0, 1000, (batch_size, seq_length), device=self.device, dtype=torch.long
|
|
)
|
|
|
|
def _warmup(self, num_warmup: int = 3):
|
|
with torch.no_grad():
|
|
for _ in range(num_warmup):
|
|
self.model(input_ids=self.static_input, attention_mask=self.static_attention_mask)
|
|
|
|
def __call__(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
|
|
if self.use_cuda_graph:
|
|
self.static_input.copy_(input_ids)
|
|
self.static_attention_mask.copy_(attention_mask)
|
|
self.graph.replay()
|
|
return self.static_output
|
|
else:
|
|
# For MPS/CPU, just run normally
|
|
return self.model(input_ids=input_ids, attention_mask=attention_mask)
|
|
|
|
|
|
class ModelOptimizer:
|
|
"""Applies various optimizations to the model."""
|
|
|
|
@staticmethod
|
|
def optimize(model: nn.Module, config: BenchmarkConfig) -> nn.Module:
|
|
print("\nApplying model optimizations:")
|
|
|
|
if model is None:
|
|
raise ValueError("Cannot optimize None model")
|
|
|
|
# Move to GPU
|
|
if torch.cuda.is_available():
|
|
model = model.cuda()
|
|
device = "cuda"
|
|
elif torch.backends.mps.is_available():
|
|
model = model.to("mps")
|
|
device = "mps"
|
|
else:
|
|
model = model.cpu()
|
|
device = "cpu"
|
|
print(f"- Model moved to {device}")
|
|
|
|
# FP16
|
|
if config.use_fp16 and not config.use_int4:
|
|
model = model.half()
|
|
# use torch compile
|
|
model = torch.compile(model)
|
|
print("- Using FP16 precision")
|
|
|
|
# Check if using SDPA (only on CUDA)
|
|
if (
|
|
torch.cuda.is_available()
|
|
and torch.version.cuda
|
|
and float(torch.version.cuda[:3]) >= 11.6
|
|
):
|
|
if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
|
|
print("- Using PyTorch SDPA (scaled_dot_product_attention)")
|
|
else:
|
|
print("- PyTorch SDPA not available")
|
|
|
|
# Flash Attention (only on CUDA)
|
|
if config.use_flash_attention and torch.cuda.is_available():
|
|
try:
|
|
from flash_attn.flash_attention import FlashAttention # noqa: F401
|
|
|
|
print("- Flash Attention 2 available")
|
|
if hasattr(model.config, "attention_mode"):
|
|
model.config.attention_mode = "flash_attention_2"
|
|
print(" - Enabled Flash Attention 2 mode")
|
|
except ImportError:
|
|
print("- Flash Attention not available")
|
|
|
|
# Memory efficient attention (only on CUDA)
|
|
if torch.cuda.is_available():
|
|
try:
|
|
from xformers.ops import memory_efficient_attention # noqa: F401
|
|
|
|
if hasattr(model, "enable_xformers_memory_efficient_attention"):
|
|
model.enable_xformers_memory_efficient_attention()
|
|
print("- Enabled xformers memory efficient attention")
|
|
else:
|
|
print("- Model doesn't support xformers")
|
|
except (ImportError, AttributeError):
|
|
print("- Xformers not available")
|
|
|
|
model.eval()
|
|
print("- Model set to eval mode")
|
|
|
|
return model
|
|
|
|
|
|
class Timer:
|
|
"""Handles accurate GPU timing using GPU events or CPU timing."""
|
|
|
|
def __init__(self):
|
|
if torch.cuda.is_available():
|
|
self.start_event = torch.cuda.Event(enable_timing=True)
|
|
self.end_event = torch.cuda.Event(enable_timing=True)
|
|
self.use_gpu_timing = True
|
|
elif torch.backends.mps.is_available():
|
|
# MPS doesn't have events, use CPU timing
|
|
self.use_gpu_timing = False
|
|
else:
|
|
# CPU timing
|
|
self.use_gpu_timing = False
|
|
|
|
@contextmanager
|
|
def timing(self):
|
|
if self.use_gpu_timing:
|
|
self.start_event.record()
|
|
yield
|
|
self.end_event.record()
|
|
self.end_event.synchronize()
|
|
else:
|
|
# Use CPU timing for MPS/CPU
|
|
start_time = time.time()
|
|
yield
|
|
self.cpu_elapsed = time.time() - start_time
|
|
|
|
def elapsed_time(self) -> float:
|
|
if self.use_gpu_timing:
|
|
return self.start_event.elapsed_time(self.end_event) / 1000 # ms to seconds
|
|
else:
|
|
return self.cpu_elapsed
|
|
|
|
|
|
class Benchmark:
|
|
"""Main benchmark runner."""
|
|
|
|
def __init__(self, config: BenchmarkConfig):
|
|
self.config = config
|
|
try:
|
|
self.model = self._load_model()
|
|
if self.model is None:
|
|
raise ValueError("Model initialization failed - model is None")
|
|
|
|
# Only use CUDA graphs on NVIDIA GPUs
|
|
if config.use_cuda_graphs and torch.cuda.is_available():
|
|
self.graphs = GraphContainer(self.model, config.seq_length)
|
|
else:
|
|
self.graphs = None
|
|
self.timer = Timer()
|
|
except Exception as e:
|
|
print(f"ERROR in benchmark initialization: {e!s}")
|
|
raise
|
|
|
|
def _load_model(self) -> nn.Module:
|
|
print(f"Loading model from {self.config.model_path}...")
|
|
|
|
try:
|
|
# Int4 quantization using HuggingFace integration
|
|
if self.config.use_int4:
|
|
import bitsandbytes as bnb
|
|
|
|
print(f"- bitsandbytes version: {bnb.__version__}")
|
|
|
|
# Check if using custom 8bit quantization
|
|
if hasattr(self.config, "use_linear8bitlt") and self.config.use_linear8bitlt:
|
|
print("- Using custom Linear8bitLt replacement for all linear layers")
|
|
|
|
# Load original model (without quantization config)
|
|
import bitsandbytes as bnb
|
|
import torch
|
|
|
|
# set default to half
|
|
torch.set_default_dtype(torch.float16)
|
|
compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
|
|
model = AutoModel.from_pretrained(
|
|
self.config.model_path,
|
|
torch_dtype=compute_dtype,
|
|
)
|
|
|
|
# Define replacement function
|
|
def replace_linear_with_linear8bitlt(model):
|
|
"""Recursively replace all nn.Linear layers with Linear8bitLt"""
|
|
for name, module in list(model.named_children()):
|
|
if isinstance(module, nn.Linear):
|
|
# Get original linear layer parameters
|
|
in_features = module.in_features
|
|
out_features = module.out_features
|
|
bias = module.bias is not None
|
|
|
|
# Create 8bit linear layer
|
|
# print size
|
|
print(f"in_features: {in_features}, out_features: {out_features}")
|
|
new_module = bnb.nn.Linear8bitLt(
|
|
in_features, out_features, bias=bias, has_fp16_weights=False
|
|
)
|
|
|
|
# Copy weights and bias
|
|
new_module.weight.data = module.weight.data
|
|
if bias:
|
|
new_module.bias.data = module.bias.data
|
|
|
|
# Replace module
|
|
setattr(model, name, new_module)
|
|
else:
|
|
# Process child modules recursively
|
|
replace_linear_with_linear8bitlt(module)
|
|
|
|
return model
|
|
|
|
# Replace all linear layers
|
|
model = replace_linear_with_linear8bitlt(model)
|
|
# add torch compile
|
|
model = torch.compile(model)
|
|
|
|
# Move model to GPU (quantization happens here)
|
|
device = (
|
|
"cuda"
|
|
if torch.cuda.is_available()
|
|
else "mps"
|
|
if torch.backends.mps.is_available()
|
|
else "cpu"
|
|
)
|
|
model = model.to(device)
|
|
|
|
print("- All linear layers replaced with Linear8bitLt")
|
|
|
|
else:
|
|
# Use original Int4 quantization method
|
|
print("- Using bitsandbytes for Int4 quantization")
|
|
|
|
# Create quantization config
|
|
|
|
compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
|
|
quantization_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=compute_dtype,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
)
|
|
|
|
print("- Quantization config:", quantization_config)
|
|
|
|
# Load model directly with quantization config
|
|
model = AutoModel.from_pretrained(
|
|
self.config.model_path,
|
|
quantization_config=quantization_config,
|
|
torch_dtype=compute_dtype,
|
|
device_map="auto", # Let HF decide on device mapping
|
|
)
|
|
|
|
# Check if model loaded successfully
|
|
if model is None:
|
|
raise ValueError("Model loading returned None")
|
|
|
|
print(f"- Model type: {type(model)}")
|
|
|
|
# Apply optimizations directly here
|
|
print("\nApplying model optimizations:")
|
|
|
|
if hasattr(self.config, "use_linear8bitlt") and self.config.use_linear8bitlt:
|
|
print("- Model moved to GPU with Linear8bitLt quantization")
|
|
else:
|
|
# Skip moving to GPU since device_map="auto" already did that
|
|
print("- Model already on GPU due to device_map='auto'")
|
|
|
|
# Skip FP16 conversion since we specified compute_dtype
|
|
print(f"- Using {compute_dtype} for compute dtype")
|
|
|
|
# Check CUDA and SDPA
|
|
if (
|
|
torch.cuda.is_available()
|
|
and torch.version.cuda
|
|
and float(torch.version.cuda[:3]) >= 11.6
|
|
):
|
|
if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
|
|
print("- Using PyTorch SDPA (scaled_dot_product_attention)")
|
|
else:
|
|
print("- PyTorch SDPA not available")
|
|
|
|
# Try xformers if available (only on CUDA)
|
|
if torch.cuda.is_available():
|
|
try:
|
|
from xformers.ops import memory_efficient_attention # noqa: F401
|
|
|
|
if hasattr(model, "enable_xformers_memory_efficient_attention"):
|
|
model.enable_xformers_memory_efficient_attention()
|
|
print("- Enabled xformers memory efficient attention")
|
|
else:
|
|
print("- Model doesn't support xformers")
|
|
except (ImportError, AttributeError):
|
|
print("- Xformers not available")
|
|
|
|
# Set to eval mode
|
|
model.eval()
|
|
print("- Model set to eval mode")
|
|
# Int8 quantization using HuggingFace integration
|
|
elif self.config.use_int8:
|
|
print("- Using INT8 quantization")
|
|
# For now, just use standard loading with INT8 config
|
|
compute_dtype = torch.float16 if self.config.use_fp16 else torch.float32
|
|
quantization_config = BitsAndBytesConfig(
|
|
load_in_8bit=True,
|
|
llm_int8_threshold=6.0,
|
|
llm_int8_has_fp16_weight=False,
|
|
)
|
|
|
|
model = AutoModel.from_pretrained(
|
|
self.config.model_path,
|
|
quantization_config=quantization_config,
|
|
torch_dtype=compute_dtype,
|
|
device_map="auto",
|
|
)
|
|
|
|
if model is None:
|
|
raise ValueError("Model loading returned None")
|
|
|
|
print(f"- Model type: {type(model)}")
|
|
model.eval()
|
|
print("- Model set to eval mode")
|
|
|
|
else:
|
|
# Standard loading for FP16/FP32
|
|
model = AutoModel.from_pretrained(self.config.model_path)
|
|
print("- Model loaded in standard precision")
|
|
print(f"- Model type: {type(model)}")
|
|
|
|
# Apply standard optimizations
|
|
# set default to half
|
|
import torch
|
|
|
|
torch.set_default_dtype(torch.bfloat16)
|
|
model = ModelOptimizer.optimize(model, self.config)
|
|
model = model.half()
|
|
# add torch compile
|
|
model = torch.compile(model)
|
|
|
|
# Final check to ensure model is not None
|
|
if model is None:
|
|
raise ValueError("Model is None after optimization")
|
|
|
|
print(f"- Final model type: {type(model)}")
|
|
return model
|
|
|
|
except Exception as e:
|
|
print(f"ERROR loading model: {e!s}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
def _create_random_batch(self, batch_size: int) -> torch.Tensor:
|
|
device = (
|
|
"cuda"
|
|
if torch.cuda.is_available()
|
|
else "mps"
|
|
if torch.backends.mps.is_available()
|
|
else "cpu"
|
|
)
|
|
return torch.randint(
|
|
0, 1000, (batch_size, self.config.seq_length), device=device, dtype=torch.long
|
|
)
|
|
|
|
def _run_inference(
|
|
self, input_ids: torch.Tensor, graph_wrapper: GraphWrapper | None = None
|
|
) -> tuple[float, torch.Tensor]:
|
|
attention_mask = torch.ones_like(input_ids)
|
|
|
|
with torch.no_grad(), self.timer.timing():
|
|
if graph_wrapper is not None:
|
|
output = graph_wrapper(input_ids, attention_mask)
|
|
else:
|
|
output = self.model(input_ids=input_ids, attention_mask=attention_mask)
|
|
|
|
return self.timer.elapsed_time(), output
|
|
|
|
def run(self) -> dict[int, dict[str, float]]:
|
|
results = {}
|
|
|
|
# Reset peak memory stats
|
|
if torch.cuda.is_available():
|
|
torch.cuda.reset_peak_memory_stats()
|
|
elif torch.backends.mps.is_available():
|
|
# MPS doesn't have reset_peak_memory_stats, skip it
|
|
pass
|
|
else:
|
|
print("- No GPU memory stats available")
|
|
|
|
for batch_size in self.config.batch_sizes:
|
|
print(f"\nTesting batch size: {batch_size}")
|
|
times = []
|
|
|
|
# Get or create graph for this batch size
|
|
graph_wrapper = (
|
|
self.graphs.get_or_create(batch_size) if self.graphs is not None else None
|
|
)
|
|
|
|
# Pre-allocate input tensor
|
|
input_ids = self._create_random_batch(batch_size)
|
|
print(f"Input shape: {input_ids.shape}")
|
|
|
|
# Run benchmark
|
|
for i in tqdm(range(self.config.num_runs), desc=f"Batch size {batch_size}"):
|
|
try:
|
|
elapsed_time, output = self._run_inference(input_ids, graph_wrapper)
|
|
if i == 0: # Only print on first run
|
|
print(f"Output shape: {output.last_hidden_state.shape}")
|
|
times.append(elapsed_time)
|
|
except Exception as e:
|
|
print(f"Error during inference: {e}")
|
|
break
|
|
|
|
if not times:
|
|
print(f"No successful runs for batch size {batch_size}, skipping")
|
|
continue
|
|
|
|
# Calculate statistics
|
|
avg_time = np.mean(times)
|
|
std_time = np.std(times)
|
|
throughput = batch_size / avg_time
|
|
|
|
results[batch_size] = {
|
|
"avg_time": avg_time,
|
|
"std_time": std_time,
|
|
"throughput": throughput,
|
|
}
|
|
|
|
print(f"Avg Time: {avg_time:.4f}s ± {std_time:.4f}s")
|
|
print(f"Throughput: {throughput:.2f} sequences/second")
|
|
|
|
# Log memory usage
|
|
if torch.cuda.is_available():
|
|
peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3)
|
|
elif torch.backends.mps.is_available():
|
|
# MPS doesn't have max_memory_allocated, use 0
|
|
peak_memory_gb = 0.0
|
|
else:
|
|
peak_memory_gb = 0.0
|
|
print("- No GPU memory usage available")
|
|
|
|
if peak_memory_gb > 0:
|
|
print(f"\nPeak GPU memory usage: {peak_memory_gb:.2f} GB")
|
|
else:
|
|
print("\n- GPU memory usage not available")
|
|
|
|
# Add memory info to results
|
|
for batch_size in results:
|
|
results[batch_size]["peak_memory_gb"] = peak_memory_gb
|
|
|
|
return results
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Model Inference Benchmark")
|
|
parser.add_argument(
|
|
"--model_path",
|
|
type=str,
|
|
default="facebook/contriever",
|
|
help="Path to the model",
|
|
)
|
|
parser.add_argument(
|
|
"--batch_sizes",
|
|
type=str,
|
|
default="1,2,4,8,16,32",
|
|
help="Comma-separated list of batch sizes",
|
|
)
|
|
parser.add_argument(
|
|
"--seq_length",
|
|
type=int,
|
|
default=256,
|
|
help="Sequence length for input",
|
|
)
|
|
parser.add_argument(
|
|
"--num_runs",
|
|
type=int,
|
|
default=5,
|
|
help="Number of runs for each batch size",
|
|
)
|
|
parser.add_argument(
|
|
"--use_fp16",
|
|
action="store_true",
|
|
help="Enable FP16 inference",
|
|
)
|
|
parser.add_argument(
|
|
"--use_int4",
|
|
action="store_true",
|
|
help="Enable INT4 quantization using bitsandbytes",
|
|
)
|
|
parser.add_argument(
|
|
"--use_int8",
|
|
action="store_true",
|
|
help="Enable INT8 quantization for both activations and weights using bitsandbytes",
|
|
)
|
|
parser.add_argument(
|
|
"--use_cuda_graphs",
|
|
action="store_true",
|
|
help="Enable CUDA Graphs optimization (only on NVIDIA GPUs)",
|
|
)
|
|
parser.add_argument(
|
|
"--use_flash_attention",
|
|
action="store_true",
|
|
help="Enable Flash Attention 2 if available (only on NVIDIA GPUs)",
|
|
)
|
|
parser.add_argument(
|
|
"--use_linear8bitlt",
|
|
action="store_true",
|
|
help="Enable Linear8bitLt quantization for all linear layers",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Print arguments for debugging
|
|
print("\nCommand line arguments:")
|
|
for arg, value in vars(args).items():
|
|
print(f"- {arg}: {value}")
|
|
|
|
config = BenchmarkConfig(
|
|
model_path=args.model_path,
|
|
batch_sizes=[int(bs) for bs in args.batch_sizes.split(",")],
|
|
seq_length=args.seq_length,
|
|
num_runs=args.num_runs,
|
|
use_fp16=args.use_fp16,
|
|
use_int4=args.use_int4,
|
|
use_int8=args.use_int8, # Add this line
|
|
use_cuda_graphs=args.use_cuda_graphs,
|
|
use_flash_attention=args.use_flash_attention,
|
|
use_linear8bitlt=args.use_linear8bitlt,
|
|
)
|
|
|
|
# Print configuration for debugging
|
|
print("\nBenchmark configuration:")
|
|
for field, value in vars(config).items():
|
|
print(f"- {field}: {value}")
|
|
|
|
try:
|
|
benchmark = Benchmark(config)
|
|
results = benchmark.run()
|
|
|
|
# Save results to file
|
|
import json
|
|
import os
|
|
|
|
# Create results directory if it doesn't exist
|
|
os.makedirs("results", exist_ok=True)
|
|
|
|
# Generate filename based on configuration
|
|
precision_type = (
|
|
"int4"
|
|
if config.use_int4
|
|
else "int8"
|
|
if config.use_int8
|
|
else "fp16"
|
|
if config.use_fp16
|
|
else "fp32"
|
|
)
|
|
model_name = os.path.basename(config.model_path)
|
|
output_file = f"results/benchmark_{model_name}_{precision_type}.json"
|
|
|
|
# Save results
|
|
with open(output_file, "w") as f:
|
|
json.dump(
|
|
{
|
|
"config": {
|
|
k: str(v) if isinstance(v, list) else v for k, v in vars(config).items()
|
|
},
|
|
"results": {str(k): v for k, v in results.items()},
|
|
},
|
|
f,
|
|
indent=2,
|
|
)
|
|
print(f"Results saved to {output_file}")
|
|
|
|
except Exception as e:
|
|
print(f"Benchmark failed: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|