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

    Extends the Megablocks Sparse MLP case to an additional weight v1 for GLUs.
    This additional weight v1 has the same initialization procedure as w1 for MLPs.

    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.
    """
    # Init for w1 and w2 matrices
    _megablocks_sparse_mlp_generic_param_init_fn_(
        module=module,
        init_fn_=init_fn_,
        init_div_is_residual=init_div_is_residual,
        div_is_residual=div_is_residual,
    )

    # Init ported from _megablocks_sparse_mlp_generic_param_init_fn_ for v1
    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

    # Separately initialize v1
    v1 = module.v1
    if isinstance(v1, DTensor):
        v1 = v1._local_tensor
    v1_size = list(v1.shape)  # type: ignore
    v1_size[0] = v1_size[0] * expert_process_group_size

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

    _v1 = v1.new_empty(v1_size)  # type: ignore
    fused_param_init_helper(_v1, init_fn_, _fused)
    _v1_local = _v1.chunk(expert_process_group_size, dim=0)[rank]
    with torch.no_grad():
        v1.copy_(_v1_local)  # type: ignore
Initializes MegaBlocks Sparse GLU.

    Extends the Megablocks Sparse MLP case to an additional weight v1 for GLUs.
    This additional weight v1 has the same initialization procedure as w1 for MLPs.

    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