def _megablocks_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

    _init_fn_ = _flip_fan_mode(init_fn_)

    # Initialize w1
    w1_size = list(module.w1.shape)  # type: ignore
    w1_size[0] = w1_size[0] * expert_process_group_size
    _w1 = module.w1.new_empty(w1_size)  # type: ignore
    stacked_param_init_helper(_w1, _init_fn_, module._stack_dim)  # type: ignore
    _w1_local = _w1.chunk(expert_process_group_size, dim=0)[rank]
    with torch.no_grad():
        module.w1.copy_(_w1_local)  # type: ignore

    # Initialize w2
    w2_size = list(module.w2.shape)  # type: ignore
    w2_size[0] = w2_size[0] * expert_process_group_size
    _w2 = module.w2.new_empty(w2_size)  # type: ignore
    stacked_param_init_helper(_w2, _init_fn_, module._stack_dim)  # type: ignore
    _w2_local = _w2.chunk(expert_process_group_size, dim=0)[rank]
    with torch.no_grad():
        module.w2.copy_(_w2_local)  # type: ignore
    if init_div_is_residual is not False:
        with torch.no_grad():
            module.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