def dmoe_fused_init_setup(
    ffn: nn.Module,
    args: 'megablocks.layers.arguments.Arguments',
    moe_world_size: int,
):
    """Attach the _fused attribute to the dMoE model.

    This is used for parameter initialization.

    Args:
        ffn (nn.Module): The FFN module.
        args (megablocks.layers.arguments.Arguments): The arguments for MegaBlocks.
        moe_world_size (int): The MoE world size.
    """
    n_exp = min(1, args.moe_num_experts // moe_world_size)
    ffn.experts.mlp._fused = (  # type: ignore
        0,
        [(n + 1) * args.ffn_hidden_size for n in range(n_exp - 1)],  # type: ignore
    )
Attach the _fused attribute to the dMoE model.

    This is used for parameter initialization.

    Args:
        ffn (nn.Module): The FFN module.
        args (megablocks.layers.arguments.Arguments): The arguments for MegaBlocks.
        moe_world_size (int): The MoE world size.
    
llmfoundry/models/layers/ffn.py