def _custom_param_init_fn(module: nn.Module):
            """Custom parameter initialization function for the model's modules.

            Args:
                module: The module to initialize.
            """
            if te is not None:
                # Initialize transformer engine modules
                if isinstance(
                    module,
                    (te.LayerNormMLP, te.LayerNormLinear, te.Linear),
                ):
                    if hasattr(module, 'reset_parameters') and callable(
                        module.reset_parameters,
                    ):
                        module.reset_parameters(defer_init=False)
                    return

            # Use the model's default initialization method
            model._init_weights(module)  # type: ignore

            # Initialize modules that are skipped in model._init_weights
            if hasattr(module, 'weight') and module.weight is not None:
                weight = module.weight
                if isinstance(weight, torch.Tensor):
                    if module.__class__.__name__ in EXPLICIT_INIT_ONES_NAMES:
                        torch.nn.init.ones_(weight)
                    elif torch.isnan(weight).any():
                        # Log a warning if a layer is left uninitialized and contains NaN values
                        log.warning(
                            f'{module.__class__.__name__} weight contains NaN values after model._init_weights call.',
                        )
Custom parameter initialization function for the model's modules.

            Args:
                module: The module to initialize.
            
llmfoundry/models/hf/hf_base.py