What causes "Got {config.axes_dim} but expected positional dim {pe_dim}" in hpcaitech/Open-Sora?

Compute axes_dim as a partition of hidden_size // num_heads (e.g. pe_dim=64 → [16, 24, 24])

hpcaitech/Open-Sora · error · ValueError

Got {config.axes_dim} but expected positional dim {pe_dim}

Error message

Got {config.axes_dim} but expected positional dim {pe_dim}

What it means

Immediately after the divisibility check, MMDiT verifies that sum(config.axes_dim) equals pe_dim = hidden_size // num_heads, because the N-D rotary embedding (EmbedND/LigerEmbedND) allocates rope dims per axis and they must exactly fill the head dimension. A mismatch means rotary embeddings would be truncated or oversized.

Source

Thrown at opensora/models/mmdit/model.py:87

class MMDiTModel(nn.Module):
    config_class = MMDiTConfig

    def __init__(self, config: MMDiTConfig):
        super().__init__()

        self.config = config
        self.in_channels = config.in_channels
        self.out_channels = self.in_channels
        self.patch_size = config.patch_size

        if config.hidden_size % config.num_heads != 0:
            raise ValueError(
                f"Hidden size {config.hidden_size} must be divisible by num_heads {config.num_heads}"
            )

        pe_dim = config.hidden_size // config.num_heads
        if sum(config.axes_dim) != pe_dim:
            raise ValueError(
                f"Got {config.axes_dim} but expected positional dim {pe_dim}"
            )

        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        pe_embedder_cls = LigerEmbedND if config.use_liger_rope else EmbedND
        self.pe_embedder = pe_embedder_cls(
            dim=pe_dim, theta=config.theta, axes_dim=config.axes_dim
        )

        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
        self.vector_in = MLPEmbedder(config.vec_in_dim, self.hidden_size)
        self.guidance_in = (
            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
            if config.guidance_embed
            else nn.Identity()
        )

View on GitHub (pinned to 7ad6a96a13)

Solutions

Compute axes_dim as a partition of hidden_size // num_heads (e.g. pe_dim=64 → [16, 24, 24])
If you changed num_heads/hidden_size, rescale axes_dim entries to sum to the new pe_dim
Use symmetric spatial dims (h/w equal) unless the task requires otherwise; only the SUM is validated

Example fix

# before
# hidden_size=3072, num_heads=24 → pe_dim=128, axes_dim=[16, 56, 56] (sum=128) ok;
config.axes_dim = [16, 32, 32]  # sum=80 ≠ 128 → error
# after
config.axes_dim = [16, 56, 56]  # sum == 3072 // 24
model = MMDit(config)

Defensive patterns

Strategy: validation

Validate before calling

pe_dim = config.hidden_size // config.num_heads
assert sum(config.axes_dim) == pe_dim, f"axes_dim sums to {sum(config.axes_dim)}, need {pe_dim}"

Type guard

def is_valid_axes_dim(axes_dim, hidden_size, num_heads) -> bool:
    return sum(axes_dim) == hidden_size // num_heads

Prevention

Derive axes_dim from hidden_size//num_heads instead of hardcoding
Re-validate rope config after any head-count change
Keep h/w rope dims equal for square video patches

When it happens

Trigger: Configuring rope axes_dim (e.g. [16, 24, 24] for t/h/w) whose sum does not equal hidden_size // num_heads; changing num_heads or hidden_size without updating axes_dim (or vice versa).

Common situations: Porting a config between MMDiT sizes (e.g. from a 2-axis image model to 3-axis video), editing head counts, or hand-writing rope configs.

Related errors

AI-assisted analysis of hpcaitech/Open-Sora@7ad6a96a13 (2026-08-28). Data as JSON: /api/errors/c7b99d02b150395f. Report an issue: GitHub.