File size: 5,743 Bytes
525f312 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Teach SGLang-Omni's Drama 3 loader the full int4 checkpoint (quant_method "drama3-int4").
Routed FusedMoE experts run SGLang's moe_wna16 kernel (W4A16, GPTQ packing, symmetric, group 128);
the dense target linears (full-attention q/k/v/o, gated-deltanet in_proj_qkv/in_proj_z/out_proj,
shared-expert gate/up/down) run GPTQ-Marlin (standard GPTQ sym checkpoint fields, repacked by
GPTQMarlinLinearMethod.process_weights_after_loading); every other layer stays unquantized bf16.
Idempotent. Usage: python apply_patch.py /path/to/sglang-omni
"""
import sys
from pathlib import Path
root = Path(sys.argv[1]) / "sglang_omni/models/fishaudio_drama3"
q = root / "quantization.py"
s = q.read_text()
ADD = '''
# --- added by drama-3-int4 (int4 experts AND dense linears) ----------------------------------
import re as _re # noqa: E402
from sglang.srt.layers.quantization.gptq.gptq import ( # noqa: E402
GPTQMarlinConfig,
GPTQMarlinLinearMethod,
)
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config, MoeWNA16Method # noqa: E402
DRAMA3_INT4_QUANTIZATION_METHOD = "drama3-int4"
# Dense checkpoint-side prefixes that hold int4 GPTQ fields. Matches by checkpoint name, which
# is what sglang's weight_loader prefix strings carry (model.language_model.layers.N....).
# The checkpoint's quantization_config.dense.targets says exactly which dense linears were
# quantized offline; the default below matches the full drama3-int4 release. Never matches
# mlp.gate (router), mlp.shared_expert_gate, linear_attn.in_proj_a/b, conv1d, norms, embeddings.
_ALL_DENSE_TARGETS = [
"self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
"linear_attn.in_proj_qkv", "linear_attn.in_proj_z", "linear_attn.out_proj",
"mlp.shared_expert.gate_proj", "mlp.shared_expert.up_proj", "mlp.shared_expert.down_proj",
]
class Drama3Int4Config(QuantizationConfig):
"""int4 W4A16 everywhere it pays: moe_wna16 for FusedMoE, GPTQ-Marlin for dense targets."""
# SGLang fuses modules at load time; the loader prefix is the FUSED name while the
# checkpoint (and our targets list) uses per-projection names. Map back before matching.
MERGED = {"qkv_proj": ("q_proj", "k_proj", "v_proj"),
"gate_up_proj": ("gate_proj", "up_proj"),
"in_proj_qkvz": ("in_proj_qkv", "in_proj_z")}
def _checkpoint_names(self, prefix: str) -> list:
p = prefix.replace("model.layers.", "model.language_model.layers.", 1)
head, _, last = p.rpartition(".")
return [f"{head}.{n}" for n in self.MERGED.get(last, (last,))]
def __init__(self, group_size: int = 128, targets: list | None = None) -> None:
super().__init__()
full_moe = {"quant_method": "gptq", "bits": 4, "group_size": group_size,
"sym": True, "desc_act": False}
self.wna16 = MoeWNA16Config("gptq", 4, group_size, False, False, None, full_moe)
full_dense = {"quant_method": "gptq_marlin", "bits": 4, "group_size": group_size,
"sym": True, "desc_act": False}
self.marlin = GPTQMarlinConfig(4, group_size, False, True, False, {}, full_dense)
targets = _ALL_DENSE_TARGETS if targets is None else targets
self.dense_re = (_re.compile("|".join(_re.escape(t) for t in targets))
if targets else None)
def get_name(self) -> str:
return DRAMA3_INT4_QUANTIZATION_METHOD
def get_supported_act_dtypes(self) -> list[torch.dtype]:
return [torch.bfloat16, torch.half]
@classmethod
def get_min_capability(cls) -> int:
return 80
@staticmethod
def get_config_filenames() -> list[str]:
return []
@classmethod
def from_config(cls, config: Mapping[str, object]) -> "Drama3Int4Config":
if (config.get("quant_method") != DRAMA3_INT4_QUANTIZATION_METHOD
or config.get("bits") != 4 or not config.get("sym")):
raise ValueError(
f"expected quant_method={DRAMA3_INT4_QUANTIZATION_METHOD!r}, bits=4, sym=True; "
f"got {dict(config)}"
)
gs = int(config.get("group_size", 128))
if gs <= 0 or gs % 32:
raise ValueError(f"group_size must be a positive multiple of 32; got {gs}")
dense_cfg = config.get("dense") or {}
targets = dense_cfg.get("targets") if isinstance(dense_cfg, Mapping) else None
return cls(gs, targets)
def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
if isinstance(layer, FusedMoE):
return MoeWNA16Method(self.wna16)
elif isinstance(layer, LinearBase):
if (prefix and self.dense_re
and any(self.dense_re.search(n) for n in self._checkpoint_names(prefix))):
return GPTQMarlinLinearMethod(self.marlin)
return UnquantizedLinearMethod()
else:
return None
def get_scaled_act_names(self) -> list[str]:
return []
'''
if "Drama3Int4Config" not in s:
q.write_text(s + ADD)
e = root / "engine_builder.py"
s = e.read_text()
if "Drama3Int4Config" not in s:
s = s.replace(" ExpertFp8Config,\n", " ExpertFp8Config,\n DRAMA3_INT4_QUANTIZATION_METHOD,\n Drama3Int4Config,\n", 1)
s = s.replace(" QUANTIZATION_METHODS[QUANTIZATION_METHOD] = ExpertFp8Config\n",
" QUANTIZATION_METHODS[QUANTIZATION_METHOD] = ExpertFp8Config\n"
" QUANTIZATION_METHODS[DRAMA3_INT4_QUANTIZATION_METHOD] = Drama3Int4Config\n", 1)
assert "QUANTIZATION_METHODS[DRAMA3_INT4_QUANTIZATION_METHOD]" in s, "engine_builder layout changed"
e.write_text(s)
print("patched", q, e)
|