aditi-ps18-state / code /quantize_w4a16.py

Push 18-19 Aug drift: v10-v13_b adapters, corpora, all 22 gate reports, corrected qat_train.py, STATE_2026-08-21

6c8db72 verified 2 months ago

2.65 kB

#!/usr/bin/env python
"""W4A16 compressed-tensors export - runs in venv_quant. int4 / group 32 / symmetric,
matching google/gemma-4-12B-it-qat-w4a16-ct. Embeddings + lm_head stay high precision
(Google's ignore list) unless --quant-embeddings."""
import os, sys, json, argparse
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
ap=argparse.ArgumentParser()
ap.add_argument("--src",required=True); ap.add_argument("--out",required=True)
ap.add_argument("--base",default="/workspace/models/gemma4-qat-unquantized")
A=ap.parse_args()
from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
IGNORE=["lm_head","re:.*embed_vision.*","re:.*vision_embedder.*","re:.*embed_audio.*","re:.*audio_tower.*","re:.*vision_tower.*","re:.*multimodal_embedder.*","re:.*patch_dense.*","re:.*patch_ln.*","re:.*pos_norm.*","re:.*embedding_projection.*"]
# Match google/gemma-4-12B-it-qat-w4a16-ct EXACTLY: int4, group 32, symmetric.
# The QAT weights are calibrated for a group-32 grid; the llm-compressor default of
# group 128 is a coarser grid and forfeits part of that calibration.
from compressed_tensors.quantization import (QuantizationArgs, QuantizationScheme,
                                             QuantizationStrategy, QuantizationType)
_w=QuantizationArgs(num_bits=4, type=QuantizationType.INT, symmetric=True,
                    strategy=QuantizationStrategy.GROUP, group_size=32)
recipe=QuantizationModifier(
    config_groups={"group_0":QuantizationScheme(targets=["Linear"], weights=_w)},
    ignore=IGNORE)
tok=AutoTokenizer.from_pretrained(A.src)
m=AutoModelForCausalLM.from_pretrained(A.src, dtype="bfloat16", device_map="cpu")
oneshot(model=m, recipe=recipe, output_dir=A.out, tokenizer=tok)
# re-assert the three keys AFTER save_pretrained (handover 6.7)
cfg=os.path.join(A.out,"config.json"); c=json.load(open(cfg))
b=json.load(open(os.path.join(A.base,"config.json")))
for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
              ("audio_config","output_proj_dims")):
    c.setdefault(sec,{})
    if c[sec].get(k) is None: c[sec][k]=b[sec][k]
json.dump(c,open(cfg,"w"),indent=2)
c=json.load(open(cfg))
for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
              ("audio_config","output_proj_dims")):
    assert c[sec].get(k) is not None, f"MISSING {sec}.{k}"
sz=sum(os.path.getsize(os.path.join(A.out,f)) for f in os.listdir(A.out) if f.endswith('.safetensors'))
print(f"[quant] W4A16 written -> {A.out}  {sz/1e9:.2f} GB  (config keys asserted)")

	#!/usr/bin/env python
	"""W4A16 compressed-tensors export - runs in venv_quant. int4 / group 32 / symmetric,
	matching google/gemma-4-12B-it-qat-w4a16-ct. Embeddings + lm_head stay high precision
	(Google's ignore list) unless --quant-embeddings."""
	import os, sys, json, argparse
	os.environ.setdefault("HF_HOME","/workspace/.hf_home")
	ap=argparse.ArgumentParser()
	ap.add_argument("--src",required=True); ap.add_argument("--out",required=True)
	ap.add_argument("--base",default="/workspace/models/gemma4-qat-unquantized")
	A=ap.parse_args()
	from transformers import AutoTokenizer, AutoModelForCausalLM
	from llmcompressor import oneshot
	from llmcompressor.modifiers.quantization import QuantizationModifier
	IGNORE=["lm_head","re:.embed_vision.","re:.vision_embedder.","re:.embed_audio.","re:.audio_tower.","re:.vision_tower.","re:.multimodal_embedder.","re:.patch_dense.","re:.patch_ln.","re:.pos_norm.","re:.embedding_projection."]
	# Match google/gemma-4-12B-it-qat-w4a16-ct EXACTLY: int4, group 32, symmetric.
	# The QAT weights are calibrated for a group-32 grid; the llm-compressor default of
	# group 128 is a coarser grid and forfeits part of that calibration.
	from compressed_tensors.quantization import (QuantizationArgs, QuantizationScheme,
	QuantizationStrategy, QuantizationType)
	_w=QuantizationArgs(num_bits=4, type=QuantizationType.INT, symmetric=True,
	strategy=QuantizationStrategy.GROUP, group_size=32)
	recipe=QuantizationModifier(
	config_groups={"group_0":QuantizationScheme(targets=["Linear"], weights=_w)},
	ignore=IGNORE)
	tok=AutoTokenizer.from_pretrained(A.src)
	m=AutoModelForCausalLM.from_pretrained(A.src, dtype="bfloat16", device_map="cpu")
	oneshot(model=m, recipe=recipe, output_dir=A.out, tokenizer=tok)
	# re-assert the three keys AFTER save_pretrained (handover 6.7)
	cfg=os.path.join(A.out,"config.json"); c=json.load(open(cfg))
	b=json.load(open(os.path.join(A.base,"config.json")))
	for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
	("audio_config","output_proj_dims")):
	c.setdefault(sec,{})
	if c[sec].get(k) is None: c[sec][k]=b[sec][k]
	json.dump(c,open(cfg,"w"),indent=2)
	c=json.load(open(cfg))
	for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
	("audio_config","output_proj_dims")):
	assert c[sec].get(k) is not None, f"MISSING {sec}.{k}"
	sz=sum(os.path.getsize(os.path.join(A.out,f)) for f in os.listdir(A.out) if f.endswith('.safetensors'))
	print(f"[quant] W4A16 written -> {A.out} {sz/1e9:.2f} GB (config keys asserted)")