def _megablocks_sparse_mlp_generic_param_init_fn_(
    module: nn.Module,
    init_fn_: Callable,
    init_div_is_residual: bool = False,
    div_is_residual: float = 1.0,
):
    """Initializes MegaBlocks MLP.

    To enable elastic deterministic initialization, this method creates the entire
    weight matrix then slice into the weight tensors such that the sampled weights
    should not vary between moe world size for the same random seed.

    Args:
        module (nn.Module): The module to initialize.
        init_fn_ (Callable): Initialization method.
        init_div_is_residual (bool): Flag enabling parameters tagged with _is_residual
            flag to be divided by div_is_residual.
        div_is_residual (float): The value by which parameter initialization is divided
            if init_div_is_residual flag is enabled.
    """
    expert_process_group_size, rank = 1, 0
    if module.expert_parallel_group is not None:
        expert_process_group_size = int(
            module.expert_parallel_group.size(),  # type: ignore
        )
        rank = int(module.expert_parallel_group.rank())  # type: ignore

    hidden_size = int(module.hidden_size)  # type: ignore

    # Initialize w1
    w1 = module.w1
    if isinstance(w1, DTensor):
        w1 = w1._local_tensor
    w1_size = list(w1.shape)  # type: ignore
    w1_size[0] = w1_size[0] * expert_process_group_size

    n_exp = w1_size[0] // hidden_size
    _fused = (0, [(n + 1) * hidden_size for n in range(n_exp - 1)])

    _w1 = w1.new_empty(w1_size)  # type: ignore
    fused_param_init_helper(_w1, init_fn_, _fused)
    _w1_local = _w1.chunk(expert_process_group_size, dim=0)[rank]
    with torch.no_grad():
        w1.copy_(_w1_local)  # type: ignore

    # Initialize w2
    w2 = module.w2
    if isinstance(w2, DTensor):
        w2 = w2._local_tensor
    w2_size = list(w2.shape)  # type: ignore
    w2_size[0] = w2_size[0] * expert_process_group_size
    _w2 = w2.new_empty(w2_size)  # type: ignore
    # MegaBlocks operates on w2 as x @ w2, so needs flipped fan mode
    fused_param_init_helper(_w2, _flip_fan_mode(init_fn_), _fused)
    _w2_local = _w2.chunk(expert_process_group_size, dim=0)[rank]
    with torch.no_grad():
        w2.copy_(_w2_local)  # type: ignore
    if init_div_is_residual is not False:
        with torch.no_grad():
            w2.div_(div_is_residual)  # type: ignore
Initializes MegaBlocks MLP.

    To enable elastic deterministic initialization, this method creates the entire
    weight matrix then slice into the weight tensors such that the sampled weights
    should not vary between moe world size for the same random seed.

    Args:
        module (nn.Module): The module to initialize.
        init_fn_ (Callable): Initialization method.
        init_div_is_residual (bool): Flag enabling parameters tagged with _is_residual
            flag to be divided by div_is_residual.
        div_is_residual (float): The value by which parameter initialization is divided
            if init_div_is_residual flag is enabled.
    
llmfoundry/models/utils/param_init_fns.py