def config_moe_args(
    ffn_config: dict,
    d_model: int,
    expansion_ratio: Union[int, float],
    n_layers: int,
) -> dict:
    """Configures `ffn_config` for MoE.

    Args:
        ffn_config (dict): FFN configuration before the MoE is configured.
        d_model (int): Hidden size of the network.
        expansion_ratio (int, float): Expansion ratio in FFN.
        n_layers (int): Number of blocks used in the network.

    Returns:
        ffn_config (dict): FFN configuration with MoE configured.
    """
    if ffn_config['ffn_type'] in ffns_with_megablocks:
        return config_megablocks_moe_args(
            ffn_config=ffn_config,
            d_model=d_model,
            expansion_ratio=expansion_ratio,
            n_layers=n_layers,
            get_device_mesh=get_megablocks_device_mesh,
        )
    else:
        ffn_type = ffn_config['ffn_type']
        raise ValueError(f'Invalid ffn_type ({ffn_type}).')
Configures `ffn_config` for MoE.

    Args:
        ffn_config (dict): FFN configuration before the MoE is configured.
        d_model (int): Hidden size of the network.
        expansion_ratio (int, float): Expansion ratio in FFN.
        n_layers (int): Number of blocks used in the network.

    Returns:
        ffn_config (dict): FFN configuration with MoE configured.
    
llmfoundry/models/utils/config_moe_args.py