def attach_ffn_mb_args(
    ffn: nn.Module,
    expert_parallel_group: ProcessGroup,
    args: 'megablocks.layers.arguments.Arguments',
):
    """Attach arguments used in parameter initialization to the FFN.

    Args:
        ffn (nn.Module): The FFN module.
        expert_parallel_group (ProcessGroup): The expert parallel process group.
        args (megablocks.layers.arguments.Arguments): The arguments for MegaBlocks.
    """
    ffn.experts.mlp.hidden_size = args.ffn_hidden_size  # type: ignore
    ffn.experts.mlp.expert_parallel_group = expert_parallel_group  # type: ignore
Attach arguments used in parameter initialization to the FFN.

    Args:
        ffn (nn.Module): The FFN module.
        expert_parallel_group (ProcessGroup): The expert parallel process group.
        args (megablocks.layers.arguments.Arguments): The arguments for MegaBlocks.
    
llmfoundry/models/layers/ffn.py