diff --git a/single_shot_inference.py b/single_shot_inference.py index 1c7b973d..9f18dc9f 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -621,39 +621,37 @@ def moe_forward(x, w, li, cfg, token_id, device): routed_scaling = cfg.get("routed_scaling_factor", 2.5) swiglu_limit = cfg.get("swiglu_limit", 10.0) mlp_inter = cfg["moe_intermediate_size"] - is_hash = li < 3 - # ---- Routing ---- + # Layers 0-2: hash routing (tid2eid lookup) + # Layers 3+: noaux_tc (sqrt(softplus) scoring + e_score_correction_bias for selection only) + # Config: topk_method='noaux_tc', scoring_func='sqrtsoftplus' expert_ids = None expert_weights = None - if is_hash: - tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid" - if tid2eid_key in w: - tid2eid = w[tid2eid_key] - tid = token_id.item() if token_id.numel() == 1 else token_id[0].item() - expert_ids = tid2eid[tid] # (top_k,) int64 - expert_weights = torch.ones(top_k, dtype=torch.float32, device=x.device) / top_k - else: - # Fallback: use dense routing even for hash layers - is_hash = False + tid2eid_key = f"model.layers.{li}.mlp.gate.tid2eid" + e_bias_key = f"model.layers.{li}.mlp.gate.e_score_correction_bias" + is_hash = tid2eid_key in w and e_bias_key not in w - if not is_hash: - # Dense routing: sqrt(softplus(X @ W_gate)) + e_bias for selection + if is_hash: + # Hash routing: deterministic per-token lookup, uniform weights + tid2eid = w[tid2eid_key] + tid = token_id.item() if token_id.numel() == 1 else token_id[0].item() + expert_ids = tid2eid[tid] # (top_k,) int64 + expert_weights = torch.ones(top_k, dtype=torch.float32, device=x.device) / top_k + else: + # Dense routing: sqrt(softplus(logits)) scoring gate_w = w[f"model.layers.{li}.mlp.gate.weight"] # (H, n_experts) BF16 logits = torch.nn.functional.linear(x, gate_w.bfloat16()) # (T, n_experts) - # Activation: sqrt(softplus(logits)) - activated = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) - # e_bias: learned per-expert bias for SELECTION ONLY (not in weights) - e_bias_key = f"model.layers.{li}.mlp.gate.e_bias" + # Scoring: sqrt(softplus(logits)) + scores = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) + # e_score_correction_bias: per-expert bias for SELECTION ONLY + selection_logits = scores.clone() if e_bias_key in w: - activated = activated + w[e_bias_key].float().unsqueeze(0) - # Top-k - scores, indices = activated.topk(top_k, dim=-1) # (T, top_k) - # Renormalize on UNBIASED activation (no e_bias in weights) - unbiased = torch.sqrt(torch.nn.functional.softplus(logits.float()) + 1e-6) - unbiased_scores = torch.gather(unbiased, -1, indices) - expert_weights = unbiased_scores / unbiased_scores.sum(dim=-1, keepdim=True) + selection_logits = selection_logits + w[e_bias_key].float().unsqueeze(0) + _, indices = selection_logits.topk(top_k, dim=-1) # (T, top_k) + # Weights from UNBIASED scores (no e_bias) + expert_weights = torch.gather(scores, -1, indices) + expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True) # For T=1 decode, squeeze if x.shape[0] == 1: expert_ids = indices[0]