Spaces:
Running on Zero
Running on Zero
File size: 6,675 Bytes
bd97ee9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """
Frox AI Morph 1.1 — Quantization
Morph 1.0's `_load_4bit()` / `_load_8bit()` in the inference engine were
stubs: they printed a message and then returned a full-precision model
regardless. This module makes quantized loading actually work.
Two paths:
- bitsandbytes NF4 (QLoRA-style) — best for training/fine-tuning
- torchao / weight-only int8 — best for pure inference, no bnb dependency
Both operate on an already-constructed MorphForCausalLM by replacing
nn.Linear layers in-place, since Morph is a custom architecture (not a
HuggingFace AutoModel) and can't go through `from_pretrained(...,
quantization_config=...)` directly.
"""
from __future__ import annotations
from typing import List, Optional
import torch
import torch.nn as nn
# ── bitsandbytes NF4 (4-bit) ──────────────────────────────────────
def quantize_4bit(
model: nn.Module,
compute_dtype: torch.dtype = torch.float16,
skip_modules: Optional[List[str]] = None,
) -> nn.Module:
"""
Replace nn.Linear layers with bitsandbytes Linear4bit (NF4).
Embedding and lm_head are skipped by default (quantizing the
vocab projection tanks quality for a tiny VRAM saving).
"""
try:
import bitsandbytes as bnb
except ImportError:
raise ImportError(
"bitsandbytes is required for 4-bit quantization. "
"Install with: pip install bitsandbytes"
)
skip = set(skip_modules or ["lm_head", "embed_tokens"])
replaced = 0
def _replace(module: nn.Module, prefix: str = ""):
nonlocal replaced
for name, child in module.named_children():
full_name = f"{prefix}.{name}" if prefix else name
if any(s in full_name for s in skip):
continue
if isinstance(child, nn.Linear):
new_layer = bnb.nn.Linear4bit(
child.in_features,
child.out_features,
bias=child.bias is not None,
compute_dtype=compute_dtype,
quant_type="nf4",
)
new_layer.weight = bnb.nn.Params4bit(
child.weight.data.clone(),
requires_grad=False,
quant_type="nf4",
)
if child.bias is not None:
new_layer.bias = nn.Parameter(child.bias.data.clone())
setattr(module, name, new_layer)
replaced += 1
else:
_replace(child, full_name)
_replace(model)
print(f"✓ 4-bit NF4 quantization applied to {replaced} linear layers")
return model
# ── Weight-only int8 (inference-only, no bnb dependency) ─────────
class Int8Linear(nn.Module):
"""
Weight-only int8 linear layer. Weights stored as int8 + per-channel
scale; activations stay in the compute dtype. ~4x smaller weights
than fp16 with a small quality cost — good for inference on cards
without bitsandbytes support (e.g. some ARM / edge deployments).
"""
def __init__(self, in_features: int, out_features: int, bias: bool = False):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.register_buffer("weight_int8", torch.zeros(out_features, in_features, dtype=torch.int8))
self.register_buffer("scale", torch.ones(out_features, dtype=torch.float32))
self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None
@classmethod
def from_linear(cls, linear: nn.Linear) -> "Int8Linear":
layer = cls(linear.in_features, linear.out_features, bias=linear.bias is not None)
w = linear.weight.data.float()
scale = w.abs().max(dim=1).values / 127.0
scale = scale.clamp(min=1e-8)
w_int8 = (w / scale.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
layer.weight_int8.copy_(w_int8)
layer.scale.copy_(scale)
if linear.bias is not None:
layer.bias.data.copy_(linear.bias.data)
return layer
def forward(self, x: torch.Tensor) -> torch.Tensor:
w = self.weight_int8.to(x.dtype) * self.scale.unsqueeze(1).to(x.dtype)
out = torch.nn.functional.linear(x, w, self.bias)
return out
def quantize_8bit(
model: nn.Module,
skip_modules: Optional[List[str]] = None,
) -> nn.Module:
"""Replace nn.Linear layers with weight-only Int8Linear (no bnb needed)."""
skip = set(skip_modules or ["lm_head", "embed_tokens"])
replaced = 0
def _replace(module: nn.Module, prefix: str = ""):
nonlocal replaced
for name, child in module.named_children():
full_name = f"{prefix}.{name}" if prefix else name
if any(s in full_name for s in skip):
continue
if isinstance(child, nn.Linear):
setattr(module, name, Int8Linear.from_linear(child))
replaced += 1
else:
_replace(child, full_name)
_replace(model)
print(f"✓ Weight-only int8 quantization applied to {replaced} linear layers")
return model
# ── Size estimation ────────────────────────────────────────────────
def estimate_memory_footprint(
num_params: int,
dtype: str = "float16",
) -> dict:
"""Estimate model weight memory footprint at various precisions."""
bytes_per_param = {
"float32": 4, "float16": 2, "bfloat16": 2,
"int8": 1, "nf4": 0.5, "int4": 0.5,
}
b = bytes_per_param.get(dtype, 2)
total_bytes = num_params * b
return {
"dtype": dtype,
"params_billions": round(num_params / 1e9, 3),
"weights_gb": round(total_bytes / (1024 ** 3), 2),
# Rule of thumb: inference needs weights + ~20% for activations/KV cache
"estimated_inference_vram_gb": round(total_bytes * 1.2 / (1024 ** 3), 2),
}
def print_quantization_report(model: nn.Module, dtype_label: str = "float16"):
total_params = sum(p.numel() for p in model.parameters())
footprint = estimate_memory_footprint(total_params, dtype_label)
print("\n📦 Model Memory Footprint")
print(f" Parameters: {footprint['params_billions']}B")
print(f" Precision: {footprint['dtype']}")
print(f" Weight size: {footprint['weights_gb']} GB")
print(f" Est. inference: {footprint['estimated_inference_vram_gb']} GB "
f"(weights + activations/KV headroom)\n")
|