ankahi / scripts /quantize_main_brain.py

Add files using upload-large-folder tool

6980f6d verified 6 months ago

2.86 kB

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()

	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()