121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
"""
|
|
Quantization detection and quantization-aware utilities for ROCM Ninodes.
|
|
|
|
Provides functions for detecting model quantization types and performing
|
|
quantization-aware memory calculations.
|
|
"""
|
|
|
|
from typing import Tuple, Optional, Any
|
|
|
|
import torch
|
|
|
|
from ..utils.memory import get_gpu_memory_info
|
|
|
|
|
|
def detect_model_quantization(model: Any) -> str:
|
|
"""
|
|
Detect if a model is quantized and return quantization type
|
|
|
|
Args:
|
|
model: Model object with model_dtype() method
|
|
|
|
Returns:
|
|
Quantization type string: "fp8", "int8", "int4", "bf16", "fp32", or "unknown"
|
|
"""
|
|
try:
|
|
# Check model dtype
|
|
model_dtype = model.model_dtype()
|
|
dtype_name = str(model_dtype).lower()
|
|
|
|
if 'int8' in dtype_name:
|
|
return "int8"
|
|
elif 'int4' in dtype_name:
|
|
return "int4"
|
|
elif 'float8' in dtype_name or 'fp8' in dtype_name:
|
|
return "fp8"
|
|
elif 'bfloat16' in dtype_name or 'bf16' in dtype_name:
|
|
return "bf16"
|
|
else:
|
|
return "fp32"
|
|
except Exception as e:
|
|
print(f"⚠️ Could not detect model quantization: {e}")
|
|
return "unknown"
|
|
|
|
|
|
def check_quantized_memory_safety(model_size: float, quantization_type: str = "fp8") -> Tuple[bool, str]:
|
|
"""
|
|
Check memory safety with quantization-aware estimation
|
|
|
|
Args:
|
|
model_size: Model size in bytes (FP32 equivalent)
|
|
quantization_type: Type of quantization ("fp8", "int8", "int4", "fp32")
|
|
|
|
Returns:
|
|
Tuple of (is_safe, message)
|
|
"""
|
|
if not torch.cuda.is_available():
|
|
return True, "CUDA not available"
|
|
|
|
total_memory, allocated_memory, reserved_memory, free_memory = get_gpu_memory_info()
|
|
|
|
if total_memory is None:
|
|
return True, "Memory info unavailable"
|
|
|
|
# Calculate memory requirements based on quantization
|
|
if quantization_type == "fp8":
|
|
required_memory = model_size * 0.5 # 50% of FP32
|
|
safety_margin = 1.2 # Less aggressive for quantized
|
|
elif quantization_type in ["int8", "int4"]:
|
|
required_memory = model_size * 0.25 # 25% of FP32
|
|
safety_margin = 1.1 # Even less aggressive
|
|
else:
|
|
required_memory = model_size # Full FP32
|
|
safety_margin = 1.5 # Conservative for non-quantized
|
|
|
|
required_gb = (required_memory * safety_margin) / (1024**3)
|
|
free_gb = free_memory / (1024**3)
|
|
|
|
if free_gb < required_gb:
|
|
return False, f"Only {free_gb:.2f}GB free, need {required_gb:.2f}GB for {quantization_type}"
|
|
|
|
return True, f"Memory OK for {quantization_type}: {free_gb:.2f}GB free"
|
|
|
|
|
|
def quantized_memory_cleanup(quantization_type: str) -> bool:
|
|
"""
|
|
Quantization-aware memory cleanup - less aggressive for quantized models
|
|
|
|
Args:
|
|
quantization_type: Type of quantization ("fp8", "int8", "int4", "fp32")
|
|
|
|
Returns:
|
|
True if cleanup succeeded, False otherwise
|
|
"""
|
|
try:
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
|
|
# Different cleanup strategies based on quantization type
|
|
if quantization_type in ["fp8", "int8", "int4"]:
|
|
# Quantized models are already memory-efficient, minimal cleanup
|
|
torch.cuda.empty_cache()
|
|
print(f"🧹 Light memory cleanup for {quantization_type} model")
|
|
else:
|
|
# Standard cleanup for non-quantized models
|
|
from ..utils.memory import gentle_memory_cleanup
|
|
gentle_memory_cleanup()
|
|
print(f"🧹 Standard memory cleanup for {quantization_type} model")
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f" Quantized memory cleanup error: {e}")
|
|
return False
|
|
|
|
|
|
__all__ = [
|
|
'detect_model_quantization',
|
|
'check_quantized_memory_safety',
|
|
'quantized_memory_cleanup',
|
|
]
|
|
|