| """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 |
|
|
