Comfy-Org/ComfyUI · error · NotImplementedError

insupportable scoring function for MoE gating: {self.scoring

Error message

insupportable scoring function for MoE gating: {self.scoring_func}

What it means

HiDream's MoE router computes gating scores and only implements scoring_func == 'softmax'; anything else raises NotImplementedError. The attribute comes from the checkpoint config of HiDream's MoE LLM layers (based on Qwen/LLaVA-style MoE routing).

Source

Thrown at comfy/ldm/hidream/model.py:272

        self.gating_dim = embed_dim
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim), dtype=dtype, device=device))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        pass
        # import torch.nn.init  as init
        # init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape

        ### compute gating score
        hidden_states = hidden_states.view(-1, h)
        logits = F.linear(hidden_states, comfy.model_management.cast_to(self.weight, dtype=hidden_states.dtype, device=hidden_states.device), None)
        if self.scoring_func == 'softmax':
            scores = logits.softmax(dim=-1)
        else:
            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')

        ### select top-k experts
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        ### norm gate to sum 1
        if self.top_k > 1 and self.norm_topk_prob:
            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
            topk_weight = topk_weight / denominator

        aux_loss = None
        return topk_idx, topk_weight, aux_loss


# Modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/model.py
class MOEFeedForwardSwiGLU(nn.Module):
    def __init__(
        self,
        dim: int,

View on GitHub (pinned to 1c6d8d45b3)

Solutions

  1. Use a HiDream checkpoint whose MoE config specifies scoring_func='softmax' (the release models).
  2. If porting a sigmoid-router variant, add a branch: scores = logits.sigmoid() when scoring_func == 'sigmoid', matching the source model's normalization.
  3. Verify the config key is being read from the right section (per-layer MoE config, not global model config).

Example fix

# before (config)
"scoring_func": "sigmoid"

# after (config)
"scoring_func": "softmax"

# or, if the checkpoint truly uses sigmoid routing, extend forward():
if self.scoring_func == 'softmax':
    scores = logits.softmax(dim=-1)
elif self.scoring_func == 'sigmoid':
    scores = logits.sigmoid()
Defensive patterns

Strategy: validation

Validate before calling

assert layer.scoring_func == "softmax", f"MoE gate scoring_func {layer.scoring_func!r} is not implemented"

Type guard

def is_softmax_router(scoring_func: str) -> bool:
    return scoring_func == "softmax"

Prevention

When it happens

Trigger: Loading a HiDream checkpoint whose MoE config sets a different scoring function (e.g. 'sigmoid' as used in some Qwen3/Llama-4 style routers), or constructing the gate module manually with a non-softmax string.

Common situations: HiDream checkpoints built on newer MoE architectures that switched gating to sigmoid scoring; community merges of config files from different model generations.

Related errors


AI-assisted analysis of Comfy-Org/ComfyUI@1c6d8d45b3 (2026-08-14). Data as JSON: /api/errors/7cb493c699fed211. Report an issue: GitHub.