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

    We prepare all necessary arguments for `megablocks.layers.arguments.Arguments` so that process
    groups can be initialized and shared across all blocks in the network.

    Args:
        ffn_config (dict): FFN configuration before the MegaBlocks MoE is configured.
        d_model (int): Hidden size of the network.
        expansion_ratio (Union[int, float]): Expansion ratio in FFN.
        n_layers (int): Number of blocks used in the network.
        get_device_mesh (Callable): Function to get the device mesh. Takes in the device mesh config and the MoE world size.

    Returns:
        ffn_config (dict): FFN configuration with MegaBlocks MoE configured.
    """
    try:
        import megablocks
    except:
        raise RuntimeError(
            'Requirements for MegaBlocks not installed; see install instructions in `README.md`.',
        )

    ffn_config.setdefault('fp16', False)
    ffn_config.setdefault('bf16', False)
    ffn_config['num_layers'] = n_layers

    ffn_type = ffn_config.pop('ffn_type')
    fc_type = ffn_config.pop('fc_type')
    ffn_act_fn = ffn_config.pop('ffn_act_fn', None)

    # Config for MegaBlocks MoE world size and device mesh
    world_size = 1  # default
    moe_world_size = ffn_config.pop('moe_world_size')
    device_mesh = None
    device_mesh_cfg = ffn_config.pop('device_mesh', None)
    if moe_world_size > 1:
        if version.parse(
            torch.__version__.split('.dev')[0],
        ) < version.parse('2.2.0'):  # type: ignore
            raise RuntimeError(
                'MoE world size > 1 is not supported in torch version {torch.__version__}<2.2.',
            )

        world_size = distributed.get_world_size()
        if world_size < moe_world_size or world_size % moe_world_size:
            raise ValueError(
                f'Invalid world size configuration: {world_size=} and {moe_world_size=}',
            )

        device_mesh = get_device_mesh(
            device_mesh_cfg=device_mesh_cfg,
            moe_world_size=moe_world_size,
            world_size=world_size,
        )

        ffn_config['moe_expert_model_parallelism'] = True
        ffn_config['expert_parallel_group'] = device_mesh[
            'expert_parallel'].get_group(0)  # type: ignore

    lbl_process_group = ffn_config.get('lbl_process_group', None)
    if lbl_process_group is not None:
        if lbl_process_group == 'expert_group':
            lbl_process_group = ffn_config['expert_parallel_group']
        elif lbl_process_group == 'global_group':
            lbl_process_group = distributed.group.WORLD
        elif isinstance(lbl_process_group, int):
            if lbl_process_group > 1:
                lbl_process_group = create_set_process_group(lbl_process_group)
            else:
                lbl_process_group = None
        elif not isinstance(lbl_process_group, distributed.ProcessGroup):
            raise ValueError(
                f'Unknown {lbl_process_group=}. Options are: none | a process group | ``expert_group`` | ``global_group`` | <GROUP_SIZE>.',
            )
        ffn_config['lbl_process_group'] = lbl_process_group

    ffn_hidden_size = resolve_ffn_hidden_size(d_model, expansion_ratio)
    ffn_config.setdefault('ffn_hidden_size', ffn_hidden_size)

    args_to_keep_in_ffn_config = inspect.signature(
        megablocks.layers.arguments.Arguments,
    ).parameters

    ffn_config = {
        k: v for k, v in ffn_config.items() if k in args_to_keep_in_ffn_config
    }

    args = megablocks.layers.arguments.Arguments(
        hidden_size=d_model,
        **ffn_config,
    )
    ffn_config['args'] = args
    ffn_config['device_mesh'] = device_mesh
    ffn_config['moe_world_size'] = moe_world_size
    ffn_config['ffn_type'] = ffn_type
    ffn_config['fc_type'] = fc_type
    ffn_config['ffn_act_fn'] = ffn_act_fn

    return ffn_config
Configures `ffn_config` for MegaBlocks MoE.

    We prepare all necessary arguments for `megablocks.layers.arguments.Arguments` so that process
    groups can be initialized and shared across all blocks in the network.

    Args:
        ffn_config (dict): FFN configuration before the MegaBlocks MoE is configured.
        d_model (int): Hidden size of the network.
        expansion_ratio (Union[int, float]): Expansion ratio in FFN.
        n_layers (int): Number of blocks used in the network.
        get_device_mesh (Callable): Function to get the device mesh. Takes in the device mesh config and the MoE world size.

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