134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
CONFIG_FILE = Path(".QuantTensor_config.json")
|
|
__version__ = "0.0.1"
|
|
|
|
# === Chargement ou création de la config ===
|
|
if CONFIG_FILE.exists():
|
|
with open(CONFIG_FILE, "r") as f:
|
|
config = json.load(f)
|
|
else:
|
|
config = None
|
|
|
|
# Vérification si on doit relancer les tests
|
|
if config is None or config.get("version") != __version__:
|
|
print("[QuantTensor] Première exécution ou mise à jour détectée : test des capacités Triton...")
|
|
|
|
# Test réel de Triton (import conditionnel)
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
import torch
|
|
print(f"[QuantTensor] Triton est fonctionnel.")
|
|
|
|
try:
|
|
pass # Tester le premier kernel de Quant
|
|
|
|
print(f" [X] Kernels de quantifications Q4_1")
|
|
Q4_1__triton_quant = True
|
|
except Exception as e :
|
|
|
|
print(f" [ ] Kernels de quantifications Q4_1 :\n{e}")
|
|
Q4_1__triton_quant = False
|
|
|
|
try:
|
|
pass # Tester le premier kernel de DeQuant
|
|
|
|
print(f" [X] Kernels de déquantifications Q4_1")
|
|
Q4_1__triton_dequant = True
|
|
except Exception as e :
|
|
|
|
print(f" [ ] Kernels de déquantifications Q4_1 :\n{e}")
|
|
Q4_1__triton_dequant = False
|
|
|
|
try:
|
|
pass # Tester kernel de GEMM même si les précédents passent pas
|
|
|
|
print(f" [X] Kernels de fused-GEMM Q4_1")
|
|
Q4_1__triton_fused_gemm = True
|
|
except Exception as e :
|
|
|
|
print(f" [ ] Kernels de fused-GEMM Q4_1 :\n{e}")
|
|
Q4_1__triton_fused_gemm = False
|
|
|
|
# Et Etc....
|
|
|
|
except Exception as e :
|
|
print(f"[QuantTensor] La bibliothèque Triton ne fonctionne pas :\n{e}")
|
|
Q4_1__triton_dequant = False
|
|
Q4_1__triton_quant = False
|
|
Q4_1__triton_fused_gemm = False
|
|
|
|
|
|
config = {
|
|
"version": __version__,
|
|
"Q4_1__triton_dequant": Q4_1__triton_dequant,
|
|
"Q4_1__triton_quant": Q4_1__triton_quant,
|
|
"Q4_1__fused_gemm": Q4_1__triton_fused_gemm
|
|
}
|
|
|
|
with open(CONFIG_FILE, "w") as f:
|
|
json.dump(config, f, indent=2)
|
|
|
|
|
|
|
|
# === Binding des fonctions au moment de l'import ===
|
|
if config["triton_dequant"]:
|
|
from ops import dequantize_q4_1_triton as dequantize_q4_1
|
|
else:
|
|
from ops import dequantize_q4_1_pytorch as dequantize_q4_1
|
|
|
|
if config["triton_quant"]:
|
|
from ops import quantize_q4_1_triton as dequantize_q4_1
|
|
else:
|
|
from ops import quantize_q4_1_pytorch as dequantize_q4_1
|
|
|
|
if config["q4_1_fused_gemm"]:
|
|
from ops import q4_1_fused_gemm_triton as q4_1_fused_gemm
|
|
else:
|
|
from ops import q4_1_fused_gemm_pytorch as q4_1_fused_gemm
|
|
|
|
# =================
|
|
|
|
class QuantTensor:
|
|
def __init__(self, format: str, qs: torch.Tensor, d: torch.Tensor,
|
|
m: torch.Tensor, original_shape: torch.Size):
|
|
self.format = format
|
|
self.qs = qs
|
|
self.d = d
|
|
self.m = m
|
|
self.original_shape = original_shape
|
|
|
|
self.dtype = d.dtype
|
|
self.device = d.device
|
|
|
|
def numel(self) -> int:
|
|
return self.original_shape.numel()
|
|
|
|
@property
|
|
def nbytes(self) -> int:
|
|
total = 0
|
|
for attr in self.__dict__.values():
|
|
if isinstance(attr, torch.Tensor):
|
|
total += attr.numel() * attr.element_size()
|
|
return total
|
|
|
|
def dequantize(self):
|
|
if self.format == "Q4_1":
|
|
w = dequantize_q4_1(self.qs, self.d, self.m)
|
|
return w.view(self.original_shape)
|
|
else:
|
|
raise NotImplementedError(f"Format {self.format} non supporté")
|
|
|
|
def to(self, device=None, dtype=None):
|
|
if device is not None:
|
|
self.device = device
|
|
self.qs = self.qs.to(device=device)
|
|
self.d = self.d.to(device=device)
|
|
self.m = self.m.to(device=device)
|
|
|
|
if dtype is not None:
|
|
self.dtype = dtype
|
|
|
|
return self |