""" Thin quantized fork of Qwen3.5-MoE (35B-A3B: hybrid DeltaNet + full-attn, MoE, MM). Same idea as model/qwen3_5_quantized.py (the 27B dense-MLP fork), but this model is a Mixture-of-Experts: every decoder layer's MLP is a Qwen3_5MoeSparseMoeBlock holding - shared_expert : a normal MLP (gate/up/down nn.Linear) -> quantized as BODY (K) - experts : 256 routed experts as 3D nn.Parameters -> quantized at EXPERT_BITS - gate / shared_expert_gate : router -> kept fp16, never quantized MIXED PRECISION: the body (attention + shared_expert) is quantized at quip_params['K'] (=4 bit); the routed experts at quip_params['expert_bits'] (=2 bit). Two things are swapped in __init__: 1. every text projection nn.Linear (attn q/k/v/o or DeltaNet in_proj_*/out_proj, and mlp.shared_expert.{gate,up,down}_proj) -> QuantizedLinear (K bit) 2. mlp.experts (Qwen3_5MoeExperts, 3D params) -> QuantizedExperts, which holds one QuantizedLinear per (expert, projection) at EXPERT_BITS and reproduces the exact routed-expert forward. Result: a self-loading COMPRESSED checkpoint. from_pretrained builds the QuantizedLinear tree (trellis buffers loaded by name from the saved state_dict) and keeps the vision encoder + embeddings + lm_head + norms + router gates in fp16. Everything else — the hybrid DeltaNet/full-attn logic, the vision encoder, the hybrid cache, rope — is inherited unchanged from the real Qwen3_5MoeForConditionalGeneration. Left fp16 (never quantized): router mlp.gate.weight, shared_expert_gate, conv1d, norms, q/k_norm, embed_tokens, lm_head, and the entire vision tower. """ import torch from torch import nn import torch.nn.functional as F from transformers.activations import ACT2FN from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( Qwen3_5MoeForConditionalGeneration, ) from .qtip_quantized_linear import QuantizedLinear from .qtip_utils import get_hadK # --- trust_remote_code copies only DIRECT relative imports of THIS file (not # transitive ones), so import every helper module here to force them all to be # fetched alongside this model. --- from .qtip_misc import clean as _clean # noqa: F401 from .qtip_matmul_had import get_hadK as _ghk # noqa: F401 from .qtip_kernel_check import has_kernel as _hk # noqa: F401 from .qtip_kernel_decompress import decode_compressed as _dc # noqa: F401 from .qtip_codebook import kdict as _kd # noqa: F401 from .qtip_bitshift import BitshiftLinear as _bl # noqa: F401 from .qtip_quantized_linear import QuantizedLinear as _ql # noqa: F401 def materialize_hadK(model): """Recompute the non-persistent Hadamard buffers of every QuantizedLinear. QuantizedLinear registers had_left/had_right as persistent=False buffers built in __init__. Under from_pretrained(low_cpu_mem_usage=True) modules are created on the meta device, so those buffers are meta/empty and are NOT restored from the checkpoint (they're not in the state_dict) -> NaN in the Hadamard transform. Call this after loading a compressed checkpoint to rebuild them on the right device. (Covers body AND expert QuantizedLinear — it iterates every module.) """ for mod in model.modules(): if isinstance(mod, QuantizedLinear): dev = mod.SU.device hl, kl = get_hadK(mod.in_features) hr, kr = get_hadK(mod.out_features) mod.had_left = None if hl is None else hl.to(dev) mod.had_right = None if hr is None else hr.to(dev) mod.K_left, mod.K_right = kl, kr return model # save-name for each projection -> used to check skip_list ("{idx}_{name}") _FULL_ATTN = {"q_proj": "q", "k_proj": "k", "v_proj": "v", "o_proj": "o"} _DELTA_ATTN = {"in_proj_qkv": "in_proj_qkv", "in_proj_z": "in_proj_z", "out_proj": "out_proj"} # MoE body MLP = the shared expert (router mlp.gate stays fp16). _SHARED_MLP = {"gate_proj": "shared_gate", "up_proj": "shared_up", "down_proj": "shared_down"} def _resolve_dtype(config): dtype = getattr(config, "torch_dtype", None) or torch.bfloat16 if isinstance(dtype, str): dtype = getattr(torch, dtype.replace("torch.", "")) return dtype def _make_qlin(in_f, out_f, qp, K, dtype): q = QuantizedLinear( in_f, out_f, qp["td_x"], qp["td_y"], qp["L"], K, qp["V"], qp["tlut_bits"], qp["decode_mode"], dtype=dtype, bias=False, ) # The trellis was packed with one layout (kernel vs plain) at quantize time; decode # MUST use the same. Honor qp['packed_for_kernel'] (default True — old checkpoints # were packed when kernels worked). The kernel-format decode is pure torch, so this # is correct even when the CUDA .so can't import. q.has_kernel = bool(qp.get("packed_for_kernel", True)) return q def _swap(parent, attr, qp, K, dtype): """Replace parent. (an nn.Linear) with a QuantizedLinear of the same shape.""" lin = getattr(parent, attr) setattr(parent, attr, _make_qlin(lin.in_features, lin.out_features, qp, K, dtype)) class QuantizedExperts(nn.Module): """Drop-in replacement for Qwen3_5MoeExperts with per-expert QuantizedLinear. The original stores gate_up_proj [E, 2*inter, hidden] and down_proj [E, hidden, inter] as 3D nn.Parameters and, per hit expert, does F.linear(state, W[e]). Here each expert's two matrices are a QuantizedLinear (K=expert_bits) so the checkpoint stays compressed on disk; the forward is otherwise identical. """ def __init__(self, text_config, qp, dtype, dead=None): super().__init__() E = text_config.num_experts hidden = text_config.hidden_size inter = text_config.moe_intermediate_size self.num_experts = E self.hidden_dim = hidden self.intermediate_dim = inter self.act_fn = ACT2FN[text_config.hidden_act] Ke = qp.get("expert_bits") or qp["K"] # dead = {(e, 'gate_up_proj'|'down_proj')} kept as fp16 nn.Linear so their # ORIGINAL bf16 weight loads (variant a: dead experts left unquantized). dead = dead or set() def _mk(e, proj, in_f, out_f): if (e, proj) in dead: return nn.Linear(in_f, out_f, bias=False, dtype=dtype) return _make_qlin(in_f, out_f, qp, Ke, dtype) # gate_up: hidden -> 2*inter ; down: inter -> hidden (out = weight.shape[0]) self.gate_up = nn.ModuleList( [_mk(e, 'gate_up_proj', hidden, 2 * inter) for e in range(E)]) self.down = nn.ModuleList( [_mk(e, 'down_proj', inter, hidden) for e in range(E)]) def forward(self, hidden_states, top_k_index, top_k_weights): final_hidden_states = torch.zeros_like(hidden_states) with torch.no_grad(): expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts) expert_mask = expert_mask.permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_idx = expert_idx[0] if expert_idx == self.num_experts: continue top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) current_state = hidden_states[token_idx] gate, up = self.gate_up[expert_idx](current_state).chunk(2, dim=-1) current_hidden_states = self.act_fn(gate) * up current_hidden_states = self.down[expert_idx](current_hidden_states) current_hidden_states = current_hidden_states * top_k_weights[ token_idx, top_k_pos, None] final_hidden_states.index_add_( 0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) return final_hidden_states class Qwen3_5MoeQuantizedForConditionalGeneration( Qwen3_5MoeForConditionalGeneration): """Qwen3.5-MoE MM model with text body + routed experts swapped to QuantizedLinear.""" def __init__(self, config): super().__init__(config) qp = getattr(config, "quip_params", None) if qp is None: return # behaves as the vanilla model qp = dict(qp) skip = set(qp.get("skip_list") or []) K = qp["K"] dtype = _resolve_dtype(config) tc = config.text_config if hasattr(config, "text_config") else config # Dead routed experts (no collected factor) kept in fp16 nn.Linear so their # ORIGINAL bf16 weight loads (variant a). qp['dead_experts'] = [[idx, e, proj], ...]. dead_by_layer = {} for entry in (qp.get("dead_experts") or []): li, e, proj = entry dead_by_layer.setdefault(int(li), set()).add((int(e), proj)) layers = self.model.language_model.layers for idx, layer in enumerate(layers): # attention: full vs DeltaNet, mutually exclusive if hasattr(layer, "self_attn"): for attr, name in _FULL_ATTN.items(): if hasattr(layer.self_attn, attr) and f"{idx}_{name}" not in skip: _swap(layer.self_attn, attr, qp, K, dtype) elif hasattr(layer, "linear_attn"): for attr, name in _DELTA_ATTN.items(): if hasattr(layer.linear_attn, attr) and f"{idx}_{name}" not in skip: _swap(layer.linear_attn, attr, qp, K, dtype) # MoE block: shared_expert MLP = body (K); routed experts = QuantizedExperts. mlp = layer.mlp if hasattr(mlp, "shared_expert"): for attr, name in _SHARED_MLP.items(): if hasattr(mlp.shared_expert, attr) and f"{idx}_{name}" not in skip: _swap(mlp.shared_expert, attr, qp, K, dtype) if hasattr(mlp, "experts"): mlp.experts = QuantizedExperts(tc, qp, dtype, dead=dead_by_layer.get(idx))