diff --git a/deploy/functional/__init__.py b/deploy/functional/__init__.py index 29951ec..2d74dcc 100644 --- a/deploy/functional/__init__.py +++ b/deploy/functional/__init__.py @@ -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 diff --git a/deploy/nn/linear.py b/deploy/nn/linear.py index 026cb5d..cb7a16b 100644 --- a/deploy/nn/linear.py +++ b/deploy/nn/linear.py @@ -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 diff --git a/deploy/nn/quantization.py b/deploy/nn/quantization.py index 31e6a49..91dedd5 100644 --- a/deploy/nn/quantization.py +++ b/deploy/nn/quantization.py @@ -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) diff --git a/flatquant/args_utils.py b/flatquant/args_utils.py index 1cda369..ee49361 100644 --- a/flatquant/args_utils.py +++ b/flatquant/args_utils.py @@ -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.') diff --git a/flatquant/data_utils.py b/flatquant/data_utils.py index 5ddc680..ffc53dc 100644 --- a/flatquant/data_utils.py +++ b/flatquant/data_utils.py @@ -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') diff --git a/flatquant/flat_linear.py b/flatquant/flat_linear.py index 6e142de..45172d9 100644 --- a/flatquant/flat_linear.py +++ b/flatquant/flat_linear.py @@ -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) diff --git a/flatquant/model_tools/phi3_utils.py b/flatquant/model_tools/phi3_utils.py new file mode 100644 index 0000000..cc62529 --- /dev/null +++ b/flatquant/model_tools/phi3_utils.py @@ -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 diff --git a/flatquant/model_utils.py b/flatquant/model_utils.py index 2d64c98..b181a02 100644 --- a/flatquant/model_utils.py +++ b/flatquant/model_utils.py @@ -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}') diff --git a/flatquant/pack_utils.py b/flatquant/pack_utils.py new file mode 100644 index 0000000..ff8be0a --- /dev/null +++ b/flatquant/pack_utils.py @@ -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 diff --git a/flatquant/utils.py b/flatquant/utils.py index 6de9dbe..b23faac 100644 --- a/flatquant/utils.py +++ b/flatquant/utils.py @@ -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) diff --git a/gptq_utils.py b/gptq_utils.py index 714849f..c28f78d 100644 --- a/gptq_utils.py +++ b/gptq_utils.py @@ -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'], diff --git a/quant_agent_inference_adapter.py b/quant_agent_inference_adapter.py new file mode 100644 index 0000000..7326b82 --- /dev/null +++ b/quant_agent_inference_adapter.py @@ -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