File size: 2,857 Bytes
6980f6d | 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 | import os
import torch
import logging
import gc
import json
from pathlib import Path
from unsloth import FastVisionModel
from safetensors.torch import save_file
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)
INPUT_MODEL = "punjabi_gemma/ankahi"
OUTPUT_DIR = "artifacts/deploy/ankahi-gemma4-e4b-int8"
def main():
log.info(f"Loading merged BF16 model from {INPUT_MODEL} for quantization...")
# Reload with INT8
model_int8, processor = FastVisionModel.from_pretrained(
INPUT_MODEL,
load_in_4bit=False, # we want 8bit
load_in_8bit=True, # unsloth handles this via bnb
device_map="auto",
dtype=torch.bfloat16,
)
log.info(f"Saving INT8 model to {OUTPUT_DIR}...")
Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
log.info(" Performing manual Safetensors save to avoid serialization issues...")
# 1. Save weights manually using safetensors
state_dict = model_int8.state_dict()
# Handle shared tensors: model.language_model.embed_tokens.weight and lm_head.weight
shared_keys = [
("model.language_model.embed_tokens.weight", "lm_head.weight"),
("model.embed_tokens.weight", "lm_head.weight")
]
for k1, k2 in shared_keys:
if k1 in state_dict and k2 in state_dict:
if state_dict[k1].data_ptr() == state_dict[k2].data_ptr():
log.info(f" Detected shared weights ({k1} and {k2}). Removing {k2} for safetensors save.")
del state_dict[k2]
break
save_file(state_dict, os.path.join(OUTPUT_DIR, "model.safetensors"))
# 2. Manually fix config.json by copying from INPUT_MODEL and adding quantization info
with open(os.path.join(INPUT_MODEL, "config.json"), "r") as f:
clean_config = json.load(f)
# Scrub non-serializable or problematic fields if any
# (Just in case, though the file on disk should be clean)
clean_config["quantization_config"] = {
"bits": 8,
"quant_method": "bitsandbytes",
"load_in_8bit": True
}
clean_config["_name_or_path"] = "ankahi-gemma4-e4b-int8"
with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
json.dump(clean_config, f, indent=2)
# 3. Save processor
processor.save_pretrained(OUTPUT_DIR)
# 4. Copy other necessary files
for filename in ["tokenizer.json", "tokenizer_config.json", "generation_config.json", "chat_template.jinja", "processor_config.json"]:
src = os.path.join(INPUT_MODEL, filename)
if os.path.exists(src):
import shutil
shutil.copy(src, os.path.join(OUTPUT_DIR, filename))
log.info(f"Manual INT8 save complete in {OUTPUT_DIR}")
if __name__ == "__main__":
main()
|