Comfy-Org/ComfyUI · error · ValueError

dim_x={dim_x} should be divisible by num_heads={num_heads}

Error message

dim_x={dim_x} should be divisible by num_heads={num_heads}

What it means

The Mochi joint attention module splits the visual stream dim_x across num_heads (head_dim = dim_x // num_heads) and requires exact divisibility. A non-divisible pair makes per-head RMSNorm and attention shapes ill-defined, so __init__ raises ValueError. dim_x/dim_y/num_heads come from the Mochi model config.

Source

Thrown at comfy/ldm/genmo/joint_model/asymm_models_joint.py:80

        update_y: bool = True,
        out_bias: bool = True,
        attend_to_padding: bool = False,
        softmax_scale: Optional[float] = None,
        device: Optional[torch.device] = None,
        dtype=None,
        operations=None,
    ):
        super().__init__()
        self.dim_x = dim_x
        self.dim_y = dim_y
        self.num_heads = num_heads
        self.head_dim = dim_x // num_heads
        self.attn_drop = attn_drop
        self.update_y = update_y
        self.attend_to_padding = attend_to_padding
        self.softmax_scale = softmax_scale
        if dim_x % num_heads != 0:
            raise ValueError(
                f"dim_x={dim_x} should be divisible by num_heads={num_heads}"
            )

        # Input layers.
        self.qkv_bias = qkv_bias
        self.qkv_x = operations.Linear(dim_x, 3 * dim_x, bias=qkv_bias, device=device, dtype=dtype)
        # Project text features to match visual features (dim_y -> dim_x)
        self.qkv_y = operations.Linear(dim_y, 3 * dim_x, bias=qkv_bias, device=device, dtype=dtype)

        # Query and key normalization for stability.
        assert qk_norm
        self.q_norm_x = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
        self.k_norm_x = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
        self.q_norm_y = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)
        self.k_norm_y = operations.RMSNorm(self.head_dim, eps=1e-5, device=device, dtype=dtype)

        # Output layers. y features go back down from dim_x -> dim_y.
        self.proj_x = operations.Linear(dim_x, dim_x, bias=out_bias, device=device, dtype=dtype)

View on GitHub (pinned to 1c6d8d45b3)

Solutions

  1. Pick num_heads that divides dim_x exactly (stock Mochi: dim_x 3072 with 24 heads).
  2. When scaling dim_x, scale num_heads to keep head_dim integral (e.g. 48*64 -> 24/32/48 heads).
  3. Validate dim_x % num_heads == 0 in your loader before building the model.

Example fix

# before
AsymmetricJointBlock(dim_x=3072, dim_y=1536, num_heads=28, ...)

# after
AsymmetricJointBlock(dim_x=3072, dim_y=1536, num_heads=24, ...)
Defensive patterns

Strategy: validation

Validate before calling

assert dim_x % num_heads == 0, f"dim_x={dim_x} not divisible by num_heads={num_heads}"

Type guard

def valid_mochi_head_config(dim_x: int, num_heads: int) -> bool:
    return num_heads > 0 and dim_x % num_heads == 0

Prevention

When it happens

Trigger: Constructing AsymmetricJointBlock with dim_x not a multiple of num_heads (e.g. 3072 with 28 heads); a partial config override that changes dim_x (widened model) but keeps the original head count.

Common situations: Experimenting with Mochi architecture variants; community checkpoints with modified widths; config typos when hand-writing the Mochi params dict.

Related errors


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