diff --git a/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/src/transformers/models/nemotron_h/modeling_nemotron_h.py index 6af7fd477564..2f746be11bdc 100644 --- a/src/transformers/models/nemotron_h/modeling_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -973,22 +973,27 @@ def _init_weights(self, module): """Initialize the weights.""" super()._init_weights(module) if isinstance(module, NemotronHMamba2Mixer): - # Initialize A_log and D parameters - A = torch.arange(1, self.config.mamba_num_heads + 1) - init.copy_(module.A_log, torch.log(A)) - init.ones_(module.D) - - dt = torch.exp( - torch.rand(self.config.mamba_num_heads) - * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) - + math.log(self.config.time_step_min) - ).clamp(min=self.config.time_step_floor) - - # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 - inv_dt = dt + torch.log(-torch.expm1(-dt)) - with torch.no_grad(): - init.copy_(module.dt_bias, inv_dt) - module.dt_bias._no_reinit = True + # Only re-initialise params that were NOT loaded from a checkpoint. + # `_is_hf_initialized` is set by `from_pretrained` on each loaded + # parameter; without this guard a post-load safety pass of + # `_init_weights` would overwrite checkpoint values of + # A_log / D / dt_bias with fresh random draws. + if not getattr(module.A_log, "_is_hf_initialized", False): + A = torch.arange(1, self.config.mamba_num_heads + 1) + init.copy_(module.A_log, torch.log(A)) + if not getattr(module.D, "_is_hf_initialized", False): + init.ones_(module.D) + if not getattr(module.dt_bias, "_is_hf_initialized", False): + dt = torch.exp( + torch.rand(self.config.mamba_num_heads) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min) + ).clamp(min=self.config.time_step_floor) + + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + init.copy_(module.dt_bias, inv_dt) elif isinstance(module, NemotronHTopkRouter): init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) init.zeros_(module.e_score_correction_bias) @@ -999,7 +1004,7 @@ def _init_weights(self, module): if isinstance(module, nn.Linear): if module.bias is not None: - if not getattr(module.bias, "_no_reinit", False): + if not getattr(module.bias, "_is_hf_initialized", False): init.zeros_(module.bias) elif isinstance(module, nn.Embedding): init.normal_(module.weight, std=self.config.initializer_range) @@ -1013,10 +1018,12 @@ def _init_weights(self, module): # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py for name, p in module.named_parameters(): if name == "out_proj.weight": + # Skip checkpoint-loaded weights so a post-load safety + # pass of `_init_weights` doesn't silently overwrite them. + if getattr(p, "_is_hf_initialized", False): + continue # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) - # We need to reinit p since this code could be called multiple times - # Having just p *= scale would repeatedly scale it down init.kaiming_uniform_(p, a=math.sqrt(5)) with torch.no_grad(): p_new = p / math.sqrt(self.config.num_hidden_layers) diff --git a/src/transformers/models/nemotron_h/modular_nemotron_h.py b/src/transformers/models/nemotron_h/modular_nemotron_h.py index f49597f43140..c0a6f596049e 100644 --- a/src/transformers/models/nemotron_h/modular_nemotron_h.py +++ b/src/transformers/models/nemotron_h/modular_nemotron_h.py @@ -326,22 +326,27 @@ def _init_weights(self, module): """Initialize the weights.""" super()._init_weights(module) if isinstance(module, NemotronHMamba2Mixer): - # Initialize A_log and D parameters - A = torch.arange(1, self.config.mamba_num_heads + 1) - init.copy_(module.A_log, torch.log(A)) - init.ones_(module.D) - - dt = torch.exp( - torch.rand(self.config.mamba_num_heads) - * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) - + math.log(self.config.time_step_min) - ).clamp(min=self.config.time_step_floor) - - # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 - inv_dt = dt + torch.log(-torch.expm1(-dt)) - with torch.no_grad(): - init.copy_(module.dt_bias, inv_dt) - module.dt_bias._no_reinit = True + # Only re-initialise params that were NOT loaded from a checkpoint. + # `_is_hf_initialized` is set by `from_pretrained` on each loaded + # parameter; without this guard a post-load safety pass of + # `_init_weights` would overwrite checkpoint values of + # A_log / D / dt_bias with fresh random draws. + if not getattr(module.A_log, "_is_hf_initialized", False): + A = torch.arange(1, self.config.mamba_num_heads + 1) + init.copy_(module.A_log, torch.log(A)) + if not getattr(module.D, "_is_hf_initialized", False): + init.ones_(module.D) + if not getattr(module.dt_bias, "_is_hf_initialized", False): + dt = torch.exp( + torch.rand(self.config.mamba_num_heads) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min) + ).clamp(min=self.config.time_step_floor) + + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + init.copy_(module.dt_bias, inv_dt) elif isinstance(module, NemotronHTopkRouter): init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) init.zeros_(module.e_score_correction_bias) @@ -352,7 +357,7 @@ def _init_weights(self, module): if isinstance(module, nn.Linear): if module.bias is not None: - if not getattr(module.bias, "_no_reinit", False): + if not getattr(module.bias, "_is_hf_initialized", False): init.zeros_(module.bias) elif isinstance(module, nn.Embedding): init.normal_(module.weight, std=self.config.initializer_range) @@ -366,10 +371,12 @@ def _init_weights(self, module): # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py for name, p in module.named_parameters(): if name == "out_proj.weight": + # Skip checkpoint-loaded weights so a post-load safety + # pass of `_init_weights` doesn't silently overwrite them. + if getattr(p, "_is_hf_initialized", False): + continue # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) - # We need to reinit p since this code could be called multiple times - # Having just p *= scale would repeatedly scale it down init.kaiming_uniform_(p, a=math.sqrt(5)) with torch.no_grad(): p_new = p / math.sqrt(self.config.num_hidden_layers)