misc - #59
Conversation
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
e02fc86 to
51ec4bd
Compare
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| for expert_idx in active_experts.tolist(): | ||
| expert = getattr(self, str(expert_idx)) | ||
| gate_weights.append(_extract_expert_dense_weight(expert.gate_proj, hidden_dim, intermediate_dim, hidden_states.dtype, device)) | ||
| up_weights.append(_extract_expert_dense_weight(expert.up_proj, hidden_dim, intermediate_dim, hidden_states.dtype, device)) | ||
| down_weights.append(_extract_expert_dense_weight(expert.down_proj, intermediate_dim, hidden_dim, hidden_states.dtype, device)) | ||
| if gate_weights[-1] is None or up_weights[-1] is None or down_weights[-1] is None: | ||
| raise RuntimeError( | ||
| f"Unable to extract dense weights for expert {expert_idx} in {self.__class__.__name__}" | ||
| ) | ||
| gate_biases.append(_extract_expert_bias(expert.gate_proj, intermediate_dim, hidden_states.dtype, device)) | ||
| up_biases.append(_extract_expert_bias(expert.up_proj, intermediate_dim, hidden_states.dtype, device)) | ||
| down_biases.append(_extract_expert_bias(expert.down_proj, hidden_dim, hidden_states.dtype, device)) | ||
|
|
||
| gate_stack = torch.stack(gate_weights, dim=0) | ||
| up_stack = torch.stack(up_weights, dim=0) | ||
| down_stack = torch.stack(down_weights, dim=0) | ||
|
|
||
| grouped_mm = torch.nn.functional.grouped_mm | ||
| gate_out = grouped_mm(selected_hidden, gate_stack, offs=offsets) |
There was a problem hiding this comment.
🔍 Grouped fast path silently changes numerics for quantized backends that use integer GEMM
_grouped_mm_dequant_experts_forward always dequantizes expert weights to hidden_states.dtype and runs an fp16/bf16 grouped_mm (defuser/modeling/moe_experts_interface.py:238-256). For QuantLinear backends whose per-expert forward normally performs a true low-precision/integer GEMM, this dequantize-then-fp-matmul substitution can produce numerically different outputs than the original per-expert loop. The PR only validates equivalence for TritonV2Linear (max diff 0) and dense nn.Linear bf16. Worth confirming that other quant backends stay within acceptable tolerance since this path becomes the default for CUDA fp16/bf16 inference.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Addressed. The orientation probe also acts as a numeric guard: it compares the projection's native forward() against the dense dequantized GEMM (x @ weight + bias). If the two do not agree within tolerance, the class is marked unsupported and linear_loop_experts_forward falls back to the original per-expert loop. This automatically excludes backends whose integer/fused kernels cannot be reproduced by a dense matmul.
A fake-quantized regression fixture is added to exercise the [in, out] dequantize_weight convention in tests/test_moe_experts_interface.py::test_grouped_mm_experts_fake_quant_square_orientation.
| except Exception as exc: # noqa: BLE001 | ||
| if DEBUG_ON: | ||
| logger.debug( | ||
| f"grouped_mm experts forward failed for {self.__class__.__name__}: {exc}; " | ||
| f"falling back to per-expert loop" | ||
| ) |
There was a problem hiding this comment.
🟡 Repeated failing fast-path attempts make the optimized path slower than before
When the grouped fast path throws at runtime, the failure is swallowed (except Exception at defuser/modeling/moe_experts_interface.py:482) without disabling the fast path, so every subsequent call redoes the whole expensive grouped setup and then throws again before falling back.
Impact: On a model whose experts pass the eligibility check but whose grouped kernel actually fails, each layer does the costly dequantize-and-stack work, errors out, and then runs the full per-expert loop anyway on every token, making the "optimization" consistently slower than the original code.
Why the guard never gets disarmed after a runtime failure
_can_use_grouped_mm (defuser/modeling/moe_experts_interface.py:180-217) only validates that dense weights can be extracted for expert0; it never actually invokes grouped_mm. It caches success via self._grouped_mm_ok = True. Thus a module can pass the gate yet still fail inside _grouped_mm_dequant_experts_forward when torch.nn.functional.grouped_mm rejects the shapes/alignment (e.g. a dim that is a multiple of 8 but not the kernel's required multiple, since the check at :204 is % 8).
On failure, the handler at :482-487 logs (only if DEBUG_ON) and falls through to _per_expert_linear_loop_experts_forward, but leaves self._grouped_mm_ok set to True and never sets self._grouped_mm_failed. So the next forward re-enters the fast path (:479), rebuilds all active-expert weight stacks in _project (:319-346), calls grouped_mm again, raises again, and only then runs the per-expert loop — permanently doubling work on the hot MoE path.
Setting self._grouped_mm_failed = True (and clearing _grouped_mm_ok) inside the except handler would latch the fallback after the first runtime failure.
| except Exception as exc: # noqa: BLE001 | |
| if DEBUG_ON: | |
| logger.debug( | |
| f"grouped_mm experts forward failed for {self.__class__.__name__}: {exc}; " | |
| f"falling back to per-expert loop" | |
| ) | |
| except Exception as exc: # noqa: BLE001 | |
| self._grouped_mm_ok = False | |
| self._grouped_mm_failed = True | |
| if DEBUG_ON: | |
| logger.debug( | |
| f"grouped_mm experts forward failed for {self.__class__.__name__}: {exc}; " | |
| f"falling back to per-expert loop" | |
| ) |
Was this helpful? React with 👍 or 👎 to provide feedback.
devin-ai-integration
Bot
changed the title
6fa0fb4 to
ccee99c
Compare
Qubitium
changed the title
1 participant
No description provided.