| |
| |
| |
| |
| @@ -1,3 +1,18 @@ |
| from .quantization import pack_i4, unpack_i4, asym_quant_dequant, sym_quant_dequant |
| -from .online_trans import ( |
| - matmul_hadU_cuda, get_decompose_dim) |
| + |
| +try: |
| + from .online_trans import matmul_hadU_cuda, get_decompose_dim |
| +except ImportError: # Hadamard extra is optional; packed W4A4 uses CUDA matmul |
| + matmul_hadU_cuda = None |
| + |
| + def get_decompose_dim(n): |
| + a = int(n ** 0.5) |
| + if a * a < n: |
| + a += 1 |
| + while True: |
| + tmp = a * a - n |
| + b = int(tmp ** 0.5) |
| + if b * b == tmp: |
| + break |
| + a += 1 |
| + return a - b, a + b |
| |
| |
| |
| |
| @@ -27,8 +27,9 @@ class Linear4bit(torch.nn.Module): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| + # CUDA sym_dequant takes half* scale_col; default zeros() is float32. |
| self.register_buffer('weight_scales', |
| - torch.zeros((self.out_features, 1), requires_grad=False)) |
| + torch.zeros((self.out_features, 1), dtype=torch.float16, requires_grad=False)) |
| self.register_buffer('weight', (torch.randint(1, 7, (self.out_features, self.in_features // 2), |
| # SubByte weight |
| dtype=torch.uint8, requires_grad=False))) |
| @@ -46,12 +47,12 @@ class Linear4bit(torch.nn.Module): |
| #shape_handler = ShapeHandler(quantized_x) |
| #quantized_x = shape_handler.flatten(quantized_x) |
| x = deploy.matmul(x, self.weight) |
| - #out = shape_handler.unflatten( |
| - # deploy.sym_dequant(int_result, scales_x, self.weight_scales)) |
| + scales_x = scales_x.to(dtype=torch.float16) |
| + weight_scales = self.weight_scales.to(dtype=torch.float16) |
| if self.bias is not None: |
| - return deploy.sym_dequant(x, scales_x, self.weight_scales) + self.bias |
| + return deploy.sym_dequant(x, scales_x, weight_scales) + self.bias |
| else: |
| - return deploy.sym_dequant(x, scales_x, self.weight_scales) |
| + return deploy.sym_dequant(x, scales_x, weight_scales) |
| |
| @staticmethod |
| def from_float(module: torch.nn.Linear, weight_scales=None,): |
| @@ -63,15 +64,16 @@ class Linear4bit(torch.nn.Module): |
| weight_matrix = module.weight.data |
| |
| |
| - int_module = Linear4bit(module.in_features, module.out_features, bias=module.bias is not None, dtype=weight_matrix.dtype).to(weight_matrix.dtype) |
| + int_module = Linear4bit(module.in_features, module.out_features, bias=module.bias is not None, dtype=torch.float16) |
| if weight_scales is not None: |
| assert weight_scales.shape == (module.out_features, 1), 'weight_scales should have shape (out_features, 1)' |
| weight_matrix = weight_matrix.cuda() |
| - int_module.weight_scales.copy_(weight_scales.to(weight_matrix.dtype)) |
| - int_rounded_weight = (weight_matrix/weight_scales.cuda()).round() |
| + scales = weight_scales.detach().to(device=weight_matrix.device, dtype=torch.float16) |
| + int_module.weight_scales.copy_(scales.cpu() if int_module.weight_scales.device.type == "cpu" else scales) |
| + int_rounded_weight = (weight_matrix / scales).round() |
| int_module.weight.copy_(deploy.functional.pack_i4(int_rounded_weight.to(torch.int8)).cpu()) |
| |
| if module.bias is not None: |
| - int_module.bias.copy_(module.bias) |
| + int_module.bias.copy_(module.bias.to(dtype=torch.float16)) |
| |
| return int_module |
| |
| |
| |
| |
| @@ -12,6 +12,7 @@ class Quantizer(torch.nn.Module): |
| |
| def forward(self, x): |
| if not isinstance(x, deploy.PackedQuantizedTensor): |
| + x = x.to(dtype=torch.float16) |
| if self.lac: |
| reshaped_x = x.reshape((-1, x.shape[-1])) |
| xmax, xmin = reshaped_x.amax(1, keepdim=True), reshaped_x.amin(1, keepdim=True) |
| |
| |
| |
| |
| @@ -21,6 +21,8 @@ supported_models = [ |
| './modelzoo/meta-llama/Llama-3.3-70B-Instruct', |
| './modelzoo/Qwen/Qwen2.5-7B-Instruct', |
| './modelzoo/Qwen/Qwen2.5-32B-Instruct', |
| + './modelzoo/microsoft/Phi-3-mini-4k-instruct', |
| + 'microsoft/Phi-3-mini-4k-instruct', |
| ] |
| supported_datasets = ['wikitext2', 'c4', 'pile'] |
| |
| @@ -30,7 +32,7 @@ def parser_gen(): |
| |
| # General Arguments |
| parser.add_argument('--model', type=str, default='meta-llama/Llama-2-7b-hf', |
| - help='Model to load.', choices=supported_models) |
| + help='Model to load.') |
| parser.add_argument('--seed', type=int, default=0, help='Random seed for HuggingFace and PyTorch.') |
| parser.add_argument('--hf_token', type=str, default=None, help='HuggingFace token for model access.') |
| |
| |
| |
| |
| |
| @@ -9,13 +9,20 @@ class TokenizerWrapper: |
| self.input_ids = input_ids |
| |
| |
| +def _load_wikitext(split): |
| + local_dir = './datasets/wikitext' |
| + if os.path.isdir(local_dir): |
| + return datasets.load_dataset(local_dir, 'wikitext-2-raw-v1', split=split) |
| + return datasets.load_dataset('wikitext', 'wikitext-2-raw-v1', split=split) |
| + |
| + |
| def get_wikitext2(nsamples, seqlen, tokenizer, eval_mode=False): |
| if eval_mode: |
| - testdata = datasets.load_dataset('./datasets/wikitext', 'wikitext-2-raw-v1', split='test') |
| + testdata = _load_wikitext('test') |
| testenc = tokenizer("\n\n".join(testdata['text']), return_tensors='pt') |
| return testenc |
| else: |
| - traindata = datasets.load_dataset('./datasets/wikitext', 'wikitext-2-raw-v1', split='train') |
| + traindata = _load_wikitext('train') |
| traindata = traindata.filter(lambda x: len(x) > 0) |
| traindata = traindata.map(lambda x : {'text': x['text'].strip()}) |
| trainenc = tokenizer("\n\n".join(traindata['text']), return_tensors='pt') |
| |
| |
| |
| |
| @@ -24,6 +24,7 @@ class FlatQuantizedLinear(nn.Module): |
| self.sigmoid = nn.Sigmoid() |
| |
| self._eval_mode = False |
| + self._packed = False |
| |
| def apply_wclip(self, weight): |
| wmin, wmax = weight.min(1, keepdim=True)[0], weight.max(1, keepdim=True)[0] |
| @@ -73,6 +74,8 @@ class FlatQuantizedLinear(nn.Module): |
| return self._eval_forward(hidden_states) |
| |
| def _eval_forward(self, hidden_states): |
| + if getattr(self, "_packed", False): |
| + return self.linear(self.a4(hidden_states)) |
| x_dtype = hidden_states.dtype |
| hidden_states = self.act_quantizer(hidden_states).to(x_dtype) |
| |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,423 @@ |
| +import math |
| + |
| +import torch |
| +import torch.nn as nn |
| +import torch.nn.functional as F |
| + |
| +from flatquant.quant_utils import ActivationQuantizer |
| +from flatquant.utils import skip_initialization |
| +from flatquant.function_utils import get_init_scale, get_decompose_dim |
| +from flatquant.trans_utils import SVDSingleTransMatrix, SVDDecomposeTransMatrix |
| +from flatquant.trans_utils import InvSingleTransMatrix, InvDecomposeTransMatrix |
| +from flatquant.flat_linear import FlatQuantizedLinear |
| + |
| +from transformers.models.phi3.modeling_phi3 import apply_rotary_pos_emb, repeat_kv |
| + |
| + |
| +def _clip_scalar(value): |
| + """Triton kron_matmul needs a Python float; 0-d tensors compile as pointers.""" |
| + if value is None: |
| + return 1.0 |
| + if torch.is_tensor(value): |
| + return float(value.detach().reshape(-1)[0].cpu()) |
| + return float(value) |
| + |
| + |
| +def _eval_proj(module, hidden_states, trans=None): |
| + """Eval projection: fuse decompose T + act quant into repo kron_matmul when packed. |
| + |
| + Import deploy.kernels.kron_matmul, not deploy.functional.online_trans (that |
| + module imports optional fast_hadamard_transform at load). |
| + """ |
| + packed = getattr(module, "_packed", False) |
| + if ( |
| + packed |
| + and trans is not None |
| + and getattr(trans, "_eval_mode", False) |
| + and hasattr(trans, "matrix_left") |
| + and hasattr(trans, "matrix_right") |
| + and hidden_states.dim() == 3 |
| + ): |
| + from deploy.kernels.kron_matmul import kron_matmul |
| + |
| + a4 = module.a4 |
| + bsz, seq_len, _ = hidden_states.shape |
| + left = trans.matrix_left.detach().to(dtype=hidden_states.dtype).T.contiguous() |
| + right = trans.matrix_right.detach().to(dtype=hidden_states.dtype).contiguous() |
| + x = hidden_states.contiguous().reshape(-1, left.shape[0], right.shape[0]) |
| + packed_x = kron_matmul( |
| + left, |
| + x, |
| + right, |
| + seq_len, |
| + _clip_scalar(getattr(a4, "clip_factor_a_max", 1.0)), |
| + _clip_scalar(getattr(a4, "clip_factor_a_min", 1.0)), |
| + ) |
| + packed_x.quantized_x = packed_x.quantized_x.reshape(bsz, seq_len, -1) |
| + packed_x.scales_x = packed_x.scales_x.reshape(bsz, 1, seq_len) |
| + return module.linear(packed_x) |
| + if trans is not None: |
| + hidden_states = trans(hidden_states) |
| + if packed: |
| + return module.linear(module.a4(hidden_states)) |
| + return module(hidden_states) |
| + |
| + |
| +class FlatQuantPhi3MLP(nn.Module): |
| + """FlatQuant wrapper for Phi-3 fused gate_up_proj + down_proj.""" |
| + |
| + def __init__(self, args, module): |
| + super().__init__() |
| + self.args = args |
| + self.config = getattr(module, "config", None) |
| + self.hidden_size = getattr(module, "hidden_size", module.gate_up_proj.in_features) |
| + self.intermediate_size = module.down_proj.in_features |
| + self.act_fn = module.activation_fn |
| + self.gate_up_proj = FlatQuantizedLinear(args, module.gate_up_proj) |
| + self.down_proj = FlatQuantizedLinear(args, module.down_proj) |
| + self.add_fq_trans() |
| + |
| + self._ori_mode = False |
| + self.diag_init = args.diag_init |
| + if self.diag_init == "sq_style": |
| + self.up_smax = torch.ones_like(self.gate_up_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5 |
| + self.down_smax = torch.ones_like(self.down_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5 |
| + |
| + def add_fq_trans(self): |
| + if self.args.direct_inv: |
| + DecomposeTransMatrix = InvDecomposeTransMatrix |
| + else: |
| + DecomposeTransMatrix = SVDDecomposeTransMatrix |
| + if self.args.w_bits < 16 or self.args.a_bits < 16: |
| + up_dim_left, up_dim_right = get_decompose_dim(self.gate_up_proj.linear.weight.shape[1]) |
| + self.up_gate_trans = DecomposeTransMatrix(up_dim_left, up_dim_right, add_diag=self.args.add_diag) |
| + down_dim_left, down_dim_right = get_decompose_dim(self.down_proj.linear.weight.shape[1]) |
| + self.down_trans = DecomposeTransMatrix(down_dim_left, down_dim_right, add_diag=self.args.add_diag) |
| + else: |
| + self.up_gate_trans, self.down_trans = None, None |
| + |
| + def _trans_forward(self, x): |
| + up_states = _eval_proj(self.gate_up_proj, x, self.up_gate_trans) |
| + gate, up_states = up_states.chunk(2, dim=-1) |
| + x_act_fn = self.act_fn(gate) * up_states |
| + return _eval_proj(self.down_proj, x_act_fn, self.down_trans) |
| + |
| + def _ori_forward(self, x): |
| + if self.diag_init == "sq_style": |
| + self.up_smax = torch.maximum( |
| + self.up_smax, x.reshape(-1, x.shape[-1]).abs().max(0)[0].clone().detach() |
| + ) |
| + up_states = self.gate_up_proj._ori_forward(x) |
| + gate, up_states = up_states.chunk(2, dim=-1) |
| + x = self.act_fn(gate) * up_states |
| + if self.diag_init == "sq_style": |
| + self.down_smax = torch.maximum( |
| + self.down_smax, x.reshape(-1, x.shape[-1]).abs().max(0)[0].clone().detach() |
| + ) |
| + return self.down_proj._ori_forward(x) |
| + |
| + def forward(self, x): |
| + if self._ori_mode: |
| + return self._ori_forward(x) |
| + return self._trans_forward(x) |
| + |
| + def reparameterize(self): |
| + if self.up_gate_trans is not None: |
| + self.up_gate_trans.to_eval_mode() |
| + self.down_trans.to_eval_mode() |
| + self.gate_up_proj.reparameterize(qa_trans=self.up_gate_trans) |
| + self.down_proj.reparameterize(qa_trans=self.down_trans) |
| + if self.up_gate_trans is not None: |
| + self.up_gate_trans.use_diag = False |
| + if self.down_trans is not None and self.down_trans.add_diag: |
| + up_weight = self.gate_up_proj.linear.weight |
| + ori_dtype = up_weight.dtype |
| + # Fuse down-channel scales into the up half of the fused gate_up matrix. |
| + up_half = up_weight[self.intermediate_size :].to(torch.float64) |
| + up_half = up_half.T.mul(self.down_trans.diag_scale.to(torch.float64)).T |
| + up_weight = up_weight.clone() |
| + up_weight[self.intermediate_size :] = up_half.to(ori_dtype) |
| + self.gate_up_proj.linear.weight.data = up_weight |
| + self.down_trans.use_diag = False |
| + |
| + def init_diag_scale(self, alpha=0.5): |
| + assert hasattr(self, "up_smax") and hasattr(self, "down_smax") |
| + upw_smax = self.gate_up_proj.linear.weight.abs().max(dim=0)[0] |
| + downw_smax = self.down_proj.linear.weight.abs().max(dim=0)[0] |
| + if self.up_gate_trans is not None: |
| + self.up_gate_trans.diag_scale.data = get_init_scale(upw_smax, self.up_smax, alpha) |
| + if self.down_trans is not None: |
| + self.down_trans.diag_scale.data = get_init_scale(downw_smax, self.down_smax, alpha) |
| + del self.up_smax, self.down_smax |
| + self.diag_init = None |
| + |
| + def rep_matrix_only(self): |
| + if self.up_gate_trans is not None: |
| + self.up_gate_trans.to_eval_mode() |
| + self.down_trans.to_eval_mode() |
| + |
| + |
| +class FlatQuantPhi3Attention(nn.Module): |
| + """FlatQuant wrapper for Phi-3 fused qkv_proj + o_proj.""" |
| + |
| + def __init__(self, args, module): |
| + super().__init__() |
| + self.args = args |
| + self.config = module.config |
| + self.layer_idx = module.layer_idx |
| + self.hidden_size = module.hidden_size |
| + self.num_heads = module.num_heads |
| + self.head_dim = module.head_dim |
| + self.num_key_value_heads = module.num_key_value_heads |
| + self.num_key_value_groups = module.num_key_value_groups |
| + self.attention_dropout = module.attention_dropout |
| + self.rotary_emb = module.rotary_emb |
| + self.query_pos = self.num_heads * self.head_dim |
| + self.kv_dim = self.num_key_value_heads * self.head_dim |
| + |
| + self.qkv_proj = FlatQuantizedLinear(args, module.qkv_proj) |
| + self.o_proj = FlatQuantizedLinear(args, module.o_proj) |
| + self.add_fq_trans() |
| + |
| + if args.q_bits < 16: |
| + self.q_cache_quantizer = ActivationQuantizer( |
| + bits=args.q_bits, sym=not (args.q_asym), lac=args.lac, groupsize=-1 |
| + ) |
| + if args.k_bits < 16: |
| + self.k_cache_quantizer = ActivationQuantizer( |
| + bits=args.k_bits, sym=not (args.k_asym), lac=args.lac, groupsize=-1 |
| + ) |
| + if args.v_bits < 16: |
| + self.v_cache_quantizer = ActivationQuantizer( |
| + bits=args.v_bits, sym=not (args.v_asym), lac=args.lac, groupsize=-1 |
| + ) |
| + |
| + self._ori_mode = False |
| + self._eval_mode = False |
| + self.diag_init = args.diag_init |
| + if self.diag_init == "sq_style": |
| + self.ln_smax = torch.ones_like(self.qkv_proj.linear.weight.abs().max(dim=0)[0]).cuda() * 1e-5 |
| + |
| + def add_fq_trans(self): |
| + if self.args.direct_inv: |
| + SingleTransMatrix, DecomposeTransMatrix = InvSingleTransMatrix, InvDecomposeTransMatrix |
| + else: |
| + SingleTransMatrix, DecomposeTransMatrix = SVDSingleTransMatrix, SVDDecomposeTransMatrix |
| + if self.args.w_bits < 16 or self.args.a_bits < 16: |
| + ln_dim_left, ln_dim_right = get_decompose_dim(self.qkv_proj.linear.weight.shape[1]) |
| + self.ln_trans = DecomposeTransMatrix(ln_dim_left, ln_dim_right, add_diag=self.args.add_diag) |
| + self.o_trans = SingleTransMatrix(self.config.num_attention_heads) |
| + else: |
| + self.ln_trans, self.o_trans = None, None |
| + |
| + if self.args.k_bits < 16 or self.args.q_bits < 16: |
| + self.kcache_trans = SingleTransMatrix(self.head_dim) |
| + else: |
| + self.kcache_trans = None |
| + if self.args.v_bits < 16 or self.args.w_bits < 16 or self.args.a_bits < 16: |
| + self.vcache_trans = SingleTransMatrix(self.head_dim) |
| + else: |
| + self.vcache_trans = None |
| + |
| + def _split_qkv(self, qkv): |
| + query_states = qkv[..., : self.query_pos] |
| + key_states = qkv[..., self.query_pos : self.query_pos + self.kv_dim] |
| + value_states = qkv[..., self.query_pos + self.kv_dim :] |
| + return query_states, key_states, value_states |
| + |
| + def _trans_forward_after_ln(self, hidden_states): |
| + qkv = _eval_proj(self.qkv_proj, hidden_states, self.ln_trans) |
| + query_states, key_states, value_states = self._split_qkv(qkv) |
| + # Official llama applies T_v as out_trans on V weights, then fuses at |
| + # reparameterize. After _eval_mode, do not transform V activations again. |
| + if ( |
| + self.vcache_trans is not None |
| + and not self.args.separate_vtrans |
| + and not self._eval_mode |
| + ): |
| + value_states = self.vcache_trans(value_states) |
| + return query_states, key_states, value_states |
| + |
| + def _ori_forward_after_ln(self, hidden_states): |
| + if self.diag_init == "sq_style" and hasattr(self, "ln_smax"): |
| + self.ln_smax = torch.maximum( |
| + self.ln_smax, |
| + hidden_states.reshape(-1, hidden_states.shape[-1]).abs().max(0)[0].clone().detach(), |
| + ) |
| + qkv = self.qkv_proj._ori_forward(hidden_states) |
| + return self._split_qkv(qkv) |
| + |
| + def quant_vcache(self, value_states): |
| + if self.args.separate_vtrans and self.vcache_trans is not None: |
| + value_states = self.vcache_trans(value_states) |
| + if self.args.v_bits < 16: |
| + value_states = self.v_cache_quantizer(value_states) |
| + return value_states |
| + |
| + def quant_kcache(self, q, k): |
| + if not (self.args.k_bits < 16 or self.args.q_bits < 16): |
| + return q, k |
| + if self.kcache_trans is not None: |
| + q = self.kcache_trans(q, inv_t=True) |
| + k = self.kcache_trans(k) |
| + if self.args.q_bits < 16: |
| + q = self.q_cache_quantizer(q).to(q) |
| + if self.args.k_bits < 16: |
| + k = self.k_cache_quantizer(k).to(q) |
| + return q, k |
| + |
| + def forward( |
| + self, |
| + hidden_states, |
| + attention_mask=None, |
| + position_ids=None, |
| + past_key_value=None, |
| + output_attentions=False, |
| + use_cache=False, |
| + cache_position=None, |
| + **kwargs, |
| + ): |
| + bsz, q_len, _ = hidden_states.size() |
| + if self._ori_mode: |
| + query_states, key_states, value_states = self._ori_forward_after_ln(hidden_states) |
| + else: |
| + query_states, key_states, value_states = self._trans_forward_after_ln(hidden_states) |
| + |
| + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| + |
| + kv_seq_len = key_states.shape[-2] |
| + if past_key_value is not None: |
| + if self.layer_idx is None: |
| + raise ValueError( |
| + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " |
| + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " |
| + "with a layer index." |
| + ) |
| + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) |
| + cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len) |
| + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) |
| + |
| + if not self._ori_mode: |
| + query_states, key_states = self.quant_kcache(query_states, key_states) |
| + value_states = self.quant_vcache(value_states) |
| + |
| + if past_key_value is not None: |
| + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| + key_states, value_states = past_key_value.update( |
| + key_states, value_states, self.layer_idx, cache_kwargs |
| + ) |
| + |
| + key_states = repeat_kv(key_states, self.num_key_value_groups) |
| + value_states = repeat_kv(value_states, self.num_key_value_groups) |
| + dropout_p = self.attention_dropout if self.training else 0.0 |
| + kv_len = key_states.shape[-2] |
| + if attention_mask is not None: |
| + causal_mask = attention_mask[:, :, :, :kv_len] |
| + attn_output = F.scaled_dot_product_attention( |
| + query_states, |
| + key_states, |
| + value_states, |
| + attn_mask=causal_mask, |
| + dropout_p=dropout_p, |
| + is_causal=False, |
| + ) |
| + else: |
| + # is_causal requires q_len == kv_len. Decode with KV cache is q_len=1. |
| + attn_output = F.scaled_dot_product_attention( |
| + query_states, |
| + key_states, |
| + value_states, |
| + dropout_p=dropout_p, |
| + is_causal=(q_len == kv_len), |
| + ) |
| + attn_output = attn_output.transpose(1, 2).contiguous() |
| + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) |
| + |
| + if self._ori_mode: |
| + attn_output = self.o_proj._ori_forward(attn_output) |
| + else: |
| + if self.o_trans is None and self.vcache_trans is not None: |
| + init_shape = attn_output.shape |
| + attn_output = attn_output.reshape( |
| + -1, self.config.num_attention_heads, self.config.hidden_size // self.config.num_attention_heads |
| + ) |
| + attn_output = torch.matmul( |
| + attn_output, self.vcache_trans.get_matrix(inv_t=True).T.to(attn_output) |
| + ).reshape(init_shape) |
| + attn_output = self.o_proj(attn_output) |
| + else: |
| + init_shape = attn_output.shape |
| + attn_output = attn_output.reshape( |
| + -1, self.config.num_attention_heads, self.config.hidden_size // self.config.num_attention_heads |
| + ) |
| + attn_output = torch.matmul( |
| + self.o_trans.get_matrix().T.to(attn_output), attn_output |
| + ).reshape(init_shape) |
| + if not self._eval_mode: |
| + attn_o_og_it = self.o_trans.get_matrix(inv_t=True) |
| + attn_v_og_it = self.vcache_trans.get_matrix(inv_t=True) |
| + attn_output = self.o_proj(attn_output, qa_trans=[attn_o_og_it, attn_v_og_it]) |
| + else: |
| + attn_output = self.o_proj(attn_output) |
| + |
| + if not output_attentions: |
| + attn_weights = None |
| + else: |
| + attn_weights = None |
| + return attn_output, attn_weights, past_key_value |
| + |
| + def reparameterize(self): |
| + if self.ln_trans is not None: |
| + self.ln_trans.to_eval_mode() |
| + if self.kcache_trans is not None: |
| + self.kcache_trans.to_eval_mode() |
| + if self.vcache_trans is not None: |
| + self.vcache_trans.to_eval_mode() |
| + if self.o_trans is not None: |
| + self.o_trans.to_eval_mode() |
| + self.qkv_proj.reparameterize(qa_trans=self.ln_trans) |
| + if self.vcache_trans is not None and not self.args.separate_vtrans: |
| + weight = self.qkv_proj.linear.weight.data |
| + ori_dtype = weight.dtype |
| + v_weight = weight[self.query_pos + self.kv_dim :].to(torch.float64) |
| + v_weight = self.vcache_trans(v_weight.T).T |
| + weight = weight.clone() |
| + weight[self.query_pos + self.kv_dim :] = v_weight.to(ori_dtype) |
| + self.qkv_proj.linear.weight.data = weight |
| + if self.o_trans is not None and self.vcache_trans is not None: |
| + attn_o_og_it = self.o_trans.get_matrix(inv_t=True) |
| + attn_v_og_it = self.vcache_trans.get_matrix(inv_t=True) |
| + self.o_proj.reparameterize(qa_trans=[attn_o_og_it, attn_v_og_it]) |
| + else: |
| + self.o_proj.reparameterize() |
| + self._eval_mode = True |
| + |
| + def init_diag_scale(self, alpha=0.5): |
| + assert hasattr(self, "ln_smax") |
| + qkvw_smax = self.qkv_proj.linear.weight.abs().max(dim=0)[0] |
| + if self.ln_trans is not None: |
| + self.ln_trans.diag_scale.data = get_init_scale(qkvw_smax, self.ln_smax, alpha) |
| + del self.ln_smax |
| + self.diag_init = None |
| + |
| + def rep_matrix_only(self): |
| + if self.ln_trans is not None: |
| + self.ln_trans.to_eval_mode() |
| + if self.kcache_trans is not None: |
| + self.kcache_trans.to_eval_mode() |
| + if self.vcache_trans is not None: |
| + self.vcache_trans.to_eval_mode() |
| + if self.o_trans is not None: |
| + self.o_trans.to_eval_mode() |
| + |
| + |
| +def apply_flatquant_to_phi3(args, model): |
| + skip_initialization() |
| + for layer in range(model.config.num_hidden_layers): |
| + model.model.layers[layer].self_attn = FlatQuantPhi3Attention( |
| + args, model.model.layers[layer].self_attn |
| + ) |
| + model.model.layers[layer].mlp = FlatQuantPhi3MLP(args, model.model.layers[layer].mlp) |
| + return model |
| |
| |
| |
| |
| @@ -1,3 +1,4 @@ |
| +import os |
| import torch |
| import transformers |
| import logging |
| @@ -62,6 +63,31 @@ def get_qwen2(model_name, hf_token): |
| return model, apply_flatquant_to_qwen |
| |
| |
| +def get_phi3(model_name, hf_token): |
| + skip_initialization() |
| + from transformers import Phi3Config, Phi3ForCausalLM |
| + from flatquant.model_tools.phi3_utils import apply_flatquant_to_phi3 |
| + |
| + is_local = os.path.isdir(model_name) |
| + config = Phi3Config.from_pretrained( |
| + model_name, use_auth_token=hf_token, local_files_only=is_local |
| + ) |
| + config._attn_implementation = "eager" |
| + config._attn_implementation_internal = "eager" |
| + model = Phi3ForCausalLM.from_pretrained( |
| + model_name, |
| + torch_dtype="auto", |
| + config=config, |
| + use_auth_token=hf_token, |
| + low_cpu_mem_usage=True, |
| + attn_implementation="eager", |
| + local_files_only=is_local, |
| + ) |
| + model.seqlen = 2048 |
| + logging.info(f"---> Loading {model_name} Model with seq_len: {model.seqlen}") |
| + return model, apply_flatquant_to_phi3 |
| + |
| + |
| def get_opt(model_name): |
| skip_initialization() |
| model = transformers.OPTForCausalLM.from_pretrained(model_name, |
| @@ -80,6 +106,8 @@ def get_model(model_name, hf_token=None): |
| return get_llama(model_name, hf_token) |
| elif 'qwen-2.5' in model_name: |
| return get_qwen2(model_name, hf_token) |
| + elif 'phi-3' in model_name.lower() or 'phi3' in model_name.lower(): |
| + return get_phi3(model_name, hf_token) |
| else: |
| raise ValueError(f'Unknown model {model_name}') |
| |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,127 @@ |
| +"""Convert a reparameterized FlatQuant model onto official packed W4A4 Linear4bit.""" |
| +from __future__ import annotations |
| + |
| + |
| +def _kernel_scale_dtype(): |
| + import torch |
| + |
| + return torch.float16 |
| + |
| + |
| +def cast_linear4bit_kernel_dtypes(model) -> int: |
| + """Match CUDA sym_quant/sym_dequant: scale_row/scale_col must be float16. |
| + |
| + Empty Linear4bit buffers default to float32, so load_state_dict keeps |
| + float32 even when the packed file stored fp16 scales. |
| + """ |
| + import torch |
| + from deploy.nn.linear import Linear4bit |
| + |
| + dtype = _kernel_scale_dtype() |
| + converted = 0 |
| + for module in model.modules(): |
| + if not isinstance(module, Linear4bit): |
| + continue |
| + scales = module.weight_scales.detach().to(dtype=dtype) |
| + module.register_buffer("weight_scales", scales) |
| + if module.bias is not None: |
| + module.register_buffer("bias", module.bias.detach().to(dtype=dtype)) |
| + converted += 1 |
| + return converted |
| + |
| + |
| +def prepare_packed_slots(model) -> None: |
| + """Replace inner nn.Linear with empty Linear4bit so a packed state_dict can load.""" |
| + import torch |
| + from deploy.nn.linear import Linear4bit |
| + from deploy.nn.quantization import Quantizer |
| + from flatquant.flat_linear import FlatQuantizedLinear |
| + |
| + dtype = _kernel_scale_dtype() |
| + for module in model.modules(): |
| + if not isinstance(module, FlatQuantizedLinear): |
| + continue |
| + if getattr(module, "_packed", False) and isinstance(module.linear, Linear4bit): |
| + continue |
| + linear = module.linear |
| + weight_dtype = linear.weight.dtype if linear.weight is not None else dtype |
| + packed = Linear4bit( |
| + linear.in_features, |
| + linear.out_features, |
| + bias=linear.bias is not None, |
| + dtype=weight_dtype, |
| + ) |
| + module.linear = packed |
| + if not hasattr(module, "a4"): |
| + module.add_module("a4", Quantizer(lac=True)) |
| + module._packed = True |
| + module._eval_mode = True |
| + cast_linear4bit_kernel_dtypes(model) |
| + |
| + |
| +def restore_packed_eval_runtime(model) -> None: |
| + """Restore eval-only Python flags that are not stored in state_dict. |
| + |
| + Packed / reparameterized weights already fused T and diag(T). The adapter |
| + must call transform.to_eval_mode() *before* load so matrix buffers exist, |
| + then set these flags *after* load. Do not run reparameterize_model again |
| + on Linear4bit weights. |
| + """ |
| + from flatquant.flat_linear import FlatQuantizedLinear |
| + |
| + for layer in model.model.layers: |
| + attn = layer.self_attn |
| + mlp = layer.mlp |
| + attn._eval_mode = True |
| + for trans in ( |
| + getattr(attn, "ln_trans", None), |
| + getattr(attn, "o_trans", None), |
| + getattr(attn, "kcache_trans", None), |
| + getattr(attn, "vcache_trans", None), |
| + getattr(mlp, "up_gate_trans", None), |
| + getattr(mlp, "down_trans", None), |
| + ): |
| + if trans is None: |
| + continue |
| + if hasattr(trans, "use_diag"): |
| + trans.use_diag = False |
| + trans._eval_mode = True |
| + for module in layer.modules(): |
| + if isinstance(module, FlatQuantizedLinear): |
| + module._eval_mode = True |
| + |
| + |
| +def convert_model_to_w4a4(model, quantizers: dict) -> int: |
| + """Pack RTN-quantized inner linears with official pack_i4 / Linear4bit.""" |
| + import torch |
| + from deploy.nn.linear import Linear4bit |
| + from deploy.nn.quantization import Quantizer |
| + from flatquant.flat_linear import FlatQuantizedLinear |
| + |
| + dtype = _kernel_scale_dtype() |
| + converted = 0 |
| + named = dict(model.named_modules()) |
| + for name, module in named.items(): |
| + if not isinstance(module, FlatQuantizedLinear): |
| + continue |
| + key = f"{name}.linear" |
| + quantizer = quantizers.get(key) |
| + if quantizer is None: |
| + raise RuntimeError(f"missing RTN quantizer for {key}") |
| + scale = quantizer.scale.detach().reshape(-1, 1).to(dtype=dtype) |
| + packed = Linear4bit.from_float(module.linear, weight_scales=scale) |
| + module.linear = packed |
| + a4 = Quantizer(lac=True) |
| + act = module.act_quantizer |
| + if hasattr(act, "clip_factor_a_max"): |
| + with torch.no_grad(): |
| + a4.clip_factor_a_max.fill_(float(act.clip_factor_a_max.detach().reshape(-1)[0])) |
| + a4.clip_factor_a_min.fill_(float(act.clip_factor_a_min.detach().reshape(-1)[0])) |
| + module.add_module("a4", a4) |
| + module._packed = True |
| + module._eval_mode = True |
| + converted += 1 |
| + if converted == 0: |
| + raise RuntimeError("no FlatQuantizedLinear modules to pack") |
| + cast_linear4bit_kernel_dtypes(model) |
| + return converted |
| |
| |
| |
| |
| @@ -52,7 +52,7 @@ def cleanup_memory(verbose=True) -> None: |
| |
| def distribute_model(model) -> None: |
| """Distribute the model across available GPUs. NB: only implemented for Llama-2/3/Qwen-2.""" |
| - no_split_module_classes = ['LlamaDecoderLayer', 'Qwen2DecoderLayer'] |
| + no_split_module_classes = ['LlamaDecoderLayer', 'Qwen2DecoderLayer', 'Phi3DecoderLayer'] |
| max_memory = get_balanced_memory(model, no_split_module_classes=no_split_module_classes) |
| |
| device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=no_split_module_classes) |
| |
| |
| |
| |
| @@ -200,12 +200,21 @@ def gptq_fwrd(model, dataloader, dev, args): |
| position_ids = cache['position_ids'] |
| |
| quantizers = {} |
| - sequential = [ |
| - ['self_attn.k_proj.linear', 'self_attn.v_proj.linear', 'self_attn.q_proj.linear'], |
| - ['self_attn.o_proj.linear'], |
| - ['mlp.up_proj.linear', 'mlp.gate_proj.linear'], |
| - ['mlp.down_proj.linear'] |
| - ] |
| + first_attn = layers[0].self_attn |
| + if hasattr(first_attn, "qkv_proj"): |
| + sequential = [ |
| + ['self_attn.qkv_proj.linear'], |
| + ['self_attn.o_proj.linear'], |
| + ['mlp.gate_up_proj.linear'], |
| + ['mlp.down_proj.linear'] |
| + ] |
| + else: |
| + sequential = [ |
| + ['self_attn.k_proj.linear', 'self_attn.v_proj.linear', 'self_attn.q_proj.linear'], |
| + ['self_attn.o_proj.linear'], |
| + ['mlp.up_proj.linear', 'mlp.gate_proj.linear'], |
| + ['mlp.down_proj.linear'] |
| + ] |
| # sequential = [ |
| # ['self_attn.k_proj', 'self_attn.v_proj', 'self_attn.q_proj'], |
| # ['self_attn.o_proj'], |
| |
| new file mode 100644 |
| |
| |
| |
| @@ -0,0 +1,66 @@ |
| +"""Reload packed FlatQuant W4A4 weights (Linear4bit + transforms) from the artifact.""" |
| +from __future__ import annotations |
| + |
| +import json |
| +import os |
| +import sys |
| +from argparse import Namespace |
| + |
| +import torch |
| +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer |
| + |
| +QUANT_AGENT_ADAPTER_API = 1 |
| + |
| + |
| +def load_model_and_tokenizer(model_path, model_id, dtype, device, trust_remote_code): |
| + method_repo = os.environ.get("QUANT_AGENT_METHOD_REPO") |
| + if method_repo: |
| + sys.path.insert(0, method_repo) |
| + import flatquant.utils as utils |
| + from flatquant.model_tools.phi3_utils import apply_flatquant_to_phi3 |
| + import flatquant.pack_utils as pack_utils |
| + |
| + packed_path = os.path.join(model_path, "packed_w4a4.pt") |
| + if not os.path.isfile(packed_path): |
| + raise FileNotFoundError("packed W4A4 artifact missing: %s" % packed_path) |
| + dtypes = {"float16": torch.float16, "bfloat16": torch.bfloat16, "auto": "auto"} |
| + torch_dtype = dtypes.get(dtype, torch.float16) |
| + tokenizer = AutoTokenizer.from_pretrained( |
| + model_path, trust_remote_code=trust_remote_code, local_files_only=True |
| + ) |
| + if tokenizer.pad_token is None: |
| + tokenizer.pad_token = tokenizer.eos_token |
| + config = AutoConfig.from_pretrained( |
| + model_path, trust_remote_code=trust_remote_code, local_files_only=True |
| + ) |
| + config._attn_implementation = "eager" |
| + config._attn_implementation_internal = "eager" |
| + utils.skip_initialization() |
| + model = AutoModelForCausalLM.from_config( |
| + config, |
| + torch_dtype=torch_dtype if device == "cuda" else torch.float32, |
| + trust_remote_code=trust_remote_code, |
| + attn_implementation="eager", |
| + ) |
| + model.seqlen = 2048 |
| + with open(os.path.join(model_path, "flatquant_args.json"), "r", encoding="utf-8") as handle: |
| + args = Namespace(**json.load(handle)) |
| + model = apply_flatquant_to_phi3(args, model) |
| + for layer in model.model.layers: |
| + layer.self_attn.rep_matrix_only() |
| + layer.mlp.rep_matrix_only() |
| + pack_utils.prepare_packed_slots(model) |
| + payload = torch.load(packed_path, map_location="cpu", weights_only=False) |
| + missing, unexpected = model.load_state_dict(payload, strict=False) |
| + packed_missing = [key for key in missing if "linear.weight" in key or key.endswith("a4.clip_factor_a_max")] |
| + if packed_missing: |
| + raise RuntimeError( |
| + "packed_w4a4.pt did not restore Linear4bit weights: missing=%s unexpected=%s" |
| + % (packed_missing[:12], list(unexpected)[:12]) |
| + ) |
| + pack_utils.cast_linear4bit_kernel_dtypes(model) |
| + pack_utils.restore_packed_eval_runtime(model) |
| + model.to(device) |
| + pack_utils.cast_linear4bit_kernel_dtypes(model) |
| + model.eval() |
| + return model, tokenizer |
|
|
