def mb_setup_args(
    d_model: int,
    expansion_ratio: Union[int, float],
    ffn_hidden_size: Optional[int],
    ffn_act_fn: Optional[dict],
    device: Optional[str],
    bias: bool,
    kwargs: dict[str, Any],
) -> tuple['megablocks.layers.arguments.Arguments', int, ProcessGroup]:
    """Setup the MegaBlocks args.

    Args:
        d_model (int): The dimension of the input and output of the FFN.
        expansion_ratio (Union[int, float]): The expansion ratio of the FFN.
        ffn_hidden_size (Optional[int]): The hidden size of the FFN.
        ffn_act_fn (Optional[dict]): The activation function of the FFN.
        device (Optional[str]): The device to run the FFN on.
        bias (bool): Whether to include bias in the FFN.
        kwargs (dict[str, Any]): Additional kwargs.

    Returns:
        tuple['megablocks.layers.arguments.Arguments', int, ProcessGroup]:
            The MegaBlocks args, the MoE world size, and the expert parallel group.
    """
    if megablocks is None:
        raise RuntimeError(
            'Requirements for megablocks not installed; see install instructions in `README.md`.',
        )
    args = kwargs['args']
    args.bias = bias
    args.hidden_size = d_model
    args.device = device

    ffn_hidden_size = resolve_ffn_hidden_size(
        d_model,
        expansion_ratio,
        ffn_hidden_size,
    )
    args.ffn_hidden_size = ffn_hidden_size

    if ffn_act_fn is not None:
        args.activation_fn = resolve_ffn_act_fn(ffn_act_fn)

    moe_world_size = 1
    expert_parallel_group = args.expert_parallel_group
    if expert_parallel_group is not None:
        moe_world_size = expert_parallel_group.size()
    if kwargs.get('moe_world_size') != moe_world_size:
        raise RuntimeError(
            f'MoE expert_parallel_group configured with incorrect world size.',
        )

    return args, moe_world_size, expert_parallel_group
Setup the MegaBlocks args.

    Args:
        d_model (int): The dimension of the input and output of the FFN.
        expansion_ratio (Union[int, float]): The expansion ratio of the FFN.
        ffn_hidden_size (Optional[int]): The hidden size of the FFN.
        ffn_act_fn (Optional[dict]): The activation function of the FFN.
        device (Optional[str]): The device to run the FFN on.
        bias (bool): Whether to include bias in the FFN.
        kwargs (dict[str, Any]): Additional kwargs.

    Returns:
        tuple['megablocks.layers.arguments.Arguments', int, ProcessGroup]:
            The MegaBlocks args, the MoE world size, and the expert parallel group.
    
llmfoundry/models/layers/ffn.py