modeling_nemotron_h.py
| 1 | # coding=utf-8 |
| 2 | # Copyright 2024 HuggingFace Inc. team. |
| 3 | # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | # you may not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
| 16 | """PyTorch NemotronH model.""" |
| 17 | |
| 18 | import math |
| 19 | from dataclasses import dataclass |
| 20 | from typing import Any, Dict, Optional, Tuple, Union |
| 21 | |
| 22 | import torch |
| 23 | import torch.utils.checkpoint |
| 24 | from torch import nn |
| 25 | from torch.nn import CrossEntropyLoss |
| 26 | |
| 27 | from transformers.activations import ACT2FN |
| 28 | from transformers.cache_utils import DynamicCache # we need __iter__ and __len__ of pkv |
| 29 | from transformers.generation import GenerationMixin |
| 30 | from transformers.modeling_attn_mask_utils import ( |
| 31 | AttentionMaskConverter, |
| 32 | ) |
| 33 | from transformers.modeling_utils import PreTrainedModel |
| 34 | from transformers.utils import ( |
| 35 | ModelOutput, |
| 36 | add_code_sample_docstrings, |
| 37 | add_start_docstrings, |
| 38 | add_start_docstrings_to_model_forward, |
| 39 | logging, |
| 40 | ) |
| 41 | from transformers.utils.import_utils import ( |
| 42 | is_causal_conv1d_available, |
| 43 | is_flash_attn_2_available, |
| 44 | is_flash_attn_greater_or_equal_2_10, |
| 45 | is_mamba_2_ssm_available, |
| 46 | ) |
| 47 | from .configuration_nemotron_h import NemotronHConfig |
| 48 | |
| 49 | |
| 50 | logger = logging.get_logger(__name__) |
| 51 | |
| 52 | |
| 53 | # Copied from transformers.models.mamba.modeling_mamba2.modeling_mamba2.py with MAMBA2->NEMOTRONH,Mamba2->NemotronH |
| 54 | # For Mamba2 components Mamba2->NemotronHMamba2 |
| 55 | if is_mamba_2_ssm_available(): |
| 56 | from mamba_ssm.ops.triton.selective_state_update import selective_state_update |
| 57 | from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined |
| 58 | else: |
| 59 | mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined, selective_state_update = None, None, None |
| 60 | |
| 61 | try: |
| 62 | #from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated |
| 63 | from mamba_ssm.ops.triton.layernorm_gated import rmsnorm_fn |
| 64 | except ImportError: |
| 65 | raise ImportError("mamba-ssm is required by the Mamba model but cannot be imported") |
| 66 | |
| 67 | if is_causal_conv1d_available(): |
| 68 | from causal_conv1d import causal_conv1d_fn, causal_conv1d_update |
| 69 | else: |
| 70 | causal_conv1d_update, causal_conv1d_fn = None, None |
| 71 | |
| 72 | if is_flash_attn_2_available(): |
| 73 | from transformers.modeling_flash_attention_utils import _flash_attention_forward |
| 74 | |
| 75 | is_fast_path_available = all( |
| 76 | ( |
| 77 | selective_state_update, |
| 78 | mamba_chunk_scan_combined, |
| 79 | mamba_split_conv1d_scan_combined, |
| 80 | causal_conv1d_fn, |
| 81 | causal_conv1d_update, |
| 82 | ) |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | _CHECKPOINT_FOR_DOC = "nvidia/Nemotron-H-56B-Base-8K" |
| 87 | _CONFIG_FOR_DOC = "NemotronHConfig" |
| 88 | |
| 89 | |
| 90 | # Helper methods for segment sum computation |
| 91 | |
| 92 | |
| 93 | def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int): |
| 94 | """ |
| 95 | Padding x tensor with `pad_size` on the seq_len dim (dim=1) |
| 96 | |
| 97 | Assumes that we only have tensors of either size 4 or 3 |
| 98 | """ |
| 99 | pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0) |
| 100 | |
| 101 | return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0) |
| 102 | |
| 103 | |
| 104 | def reshape_into_chunks(input_tensor, pad_size, chunk_size): |
| 105 | """ |
| 106 | Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and |
| 107 | simultaneously splitting it into chunk sequences. |
| 108 | |
| 109 | Assumes that we only have tensors of either size 4 or 3 |
| 110 | """ |
| 111 | # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...] |
| 112 | input_tensor = pad_tensor_by_size(input_tensor, pad_size) |
| 113 | |
| 114 | if len(input_tensor.shape) == 3: |
| 115 | # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads] |
| 116 | return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2]) |
| 117 | else: |
| 118 | # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size] |
| 119 | return input_tensor.reshape( |
| 120 | input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3] |
| 121 | ) |
| 122 | |
| 123 | |
| 124 | def segment_sum(input_tensor): |
| 125 | """ |
| 126 | More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions. |
| 127 | """ |
| 128 | chunk_size = input_tensor.size(-1) |
| 129 | # 1. expand input tensor to have an additional dimension and repeat along that dimension |
| 130 | # [..., chunk_size] -> [..., chunk_size, chunk_size] |
| 131 | input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size) |
| 132 | # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag |
| 133 | mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1) |
| 134 | input_tensor = input_tensor.masked_fill(~mask, 0) |
| 135 | # 3. compute actual cumsum |
| 136 | tensor_segsum = torch.cumsum(input_tensor, dim=-2) |
| 137 | |
| 138 | # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time) |
| 139 | mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0) |
| 140 | tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf) |
| 141 | return tensor_segsum |
| 142 | |
| 143 | |
| 144 | def apply_mask_to_padding_states(hidden_states, attention_mask): |
| 145 | """ |
| 146 | Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 |
| 147 | """ |
| 148 | if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: |
| 149 | dtype = hidden_states.dtype |
| 150 | hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) |
| 151 | |
| 152 | return hidden_states |
| 153 | |
| 154 | # Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py |
| 155 | class HybridMambaAttentionDynamicCache(DynamicCache): |
| 156 | """ |
| 157 | A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache |
| 158 | (which has a constant shape regardless of seq_len). |
| 159 | |
| 160 | This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states` |
| 161 | and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor |
| 162 | For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`, |
| 163 | while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors). |
| 164 | For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors), |
| 165 | while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`, |
| 166 | and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`. |
| 167 | """ |
| 168 | |
| 169 | def __init__(self, config, batch_size, dtype=torch.float16, device=None): |
| 170 | super().__init__() |
| 171 | self.dtype = dtype |
| 172 | self.hybrid_override_pattern = config.hybrid_override_pattern |
| 173 | self.has_previous_state = False # only used by mamba |
| 174 | intermediate_size = config.expand * config.hidden_size |
| 175 | ssm_state_size = config.ssm_state_size |
| 176 | conv_kernel_size = config.conv_kernel |
| 177 | self.conv_states = [] |
| 178 | self.ssm_states = [] |
| 179 | self.transformer_layers = [] |
| 180 | for i in range(config.num_hidden_layers): |
| 181 | if self.hybrid_override_pattern[i] == "M": |
| 182 | # Mamba layer |
| 183 | self.conv_states += [ |
| 184 | torch.zeros(batch_size, intermediate_size, conv_kernel_size, device=device, dtype=dtype) |
| 185 | ] |
| 186 | self.ssm_states += [ |
| 187 | torch.zeros(batch_size, intermediate_size, ssm_state_size, device=device, dtype=dtype) |
| 188 | ] |
| 189 | else: |
| 190 | # Attention or MLP layer |
| 191 | self.conv_states += [torch.tensor([[]] * batch_size, device=device)] |
| 192 | self.ssm_states += [torch.tensor([[]] * batch_size, device=device)] |
| 193 | self.transformer_layers.append(i) |
| 194 | |
| 195 | self.key_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)] |
| 196 | self.value_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)] |
| 197 | |
| 198 | def update( |
| 199 | self, |
| 200 | key_states: torch.Tensor, |
| 201 | value_states: torch.Tensor, |
| 202 | layer_idx: int, |
| 203 | cache_kwargs: Optional[Dict[str, Any]] = None, |
| 204 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 205 | # Update the cache |
| 206 | if self.key_cache[layer_idx].shape[-1] == 0: |
| 207 | self.key_cache[layer_idx] = key_states |
| 208 | self.value_cache[layer_idx] = value_states |
| 209 | else: |
| 210 | self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2) |
| 211 | self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2) |
| 212 | |
| 213 | return self.key_cache[layer_idx], self.value_cache[layer_idx] |
| 214 | |
| 215 | def reorder_cache(self, beam_idx: torch.LongTensor): |
| 216 | """Reorders the cache for beam search, given the selected beam indices.""" |
| 217 | for layer_idx in range(len(self.key_cache)): |
| 218 | device = self.key_cache[layer_idx].device |
| 219 | self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) |
| 220 | device = self.value_cache[layer_idx].device |
| 221 | self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) |
| 222 | |
| 223 | device = self.conv_states[layer_idx].device |
| 224 | self.conv_states[layer_idx] = self.conv_states[layer_idx].index_select(0, beam_idx.to(device)) |
| 225 | device = self.ssm_states[layer_idx].device |
| 226 | self.ssm_states[layer_idx] = self.ssm_states[layer_idx].index_select(0, beam_idx.to(device)) |
| 227 | |
| 228 | def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: |
| 229 | """Returns the sequence length of the cached states. A layer index can be optionally passed.""" |
| 230 | # take any layer that contains cache and not empty tensor |
| 231 | layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx |
| 232 | if len(self.key_cache) <= layer_idx: |
| 233 | return 0 |
| 234 | return self.key_cache[layer_idx].shape[-2] |
| 235 | |
| 236 | def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]: |
| 237 | raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.") |
| 238 | |
| 239 | @classmethod |
| 240 | def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache": |
| 241 | raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.") |
| 242 | |
| 243 | # Copied from modeling_mamba2.py |
| 244 | def update_conv_state( |
| 245 | self, layer_idx: int, new_conv_state: torch.Tensor, cache_init: bool = False |
| 246 | ) -> torch.Tensor: |
| 247 | if cache_init: |
| 248 | self.conv_states[layer_idx] = new_conv_state.to(self.conv_states.device) |
| 249 | else: |
| 250 | self.conv_states[layer_idx] = self.conv_states[layer_idx].roll(shifts=-1, dims=-1) |
| 251 | self.conv_states[layer_idx][:, :, -1] = new_conv_state[:, 0, :].to(self.conv_states.device) |
| 252 | return self.conv_states[layer_idx] |
| 253 | |
| 254 | def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor): |
| 255 | self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states.device) |
| 256 | return self.ssm_states[layer_idx] |
| 257 | |
| 258 | def reset(self): |
| 259 | self.conv_states.zero_() |
| 260 | self.ssm_states.zero_() |
| 261 | |
| 262 | class MambaRMSNormGated(torch.nn.Module): |
| 263 | def __init__(self, hidden_size, group_size, eps=1e-5): |
| 264 | super().__init__() |
| 265 | self.weight = nn.Parameter(torch.ones(hidden_size)) |
| 266 | self.variance_epsilon = eps |
| 267 | self.group_size = group_size |
| 268 | |
| 269 | # jan28b version |
| 270 | def forward(self, hidden_states, gate=None): |
| 271 | return rmsnorm_fn(x=hidden_states, |
| 272 | weight=self.weight, |
| 273 | bias=None, # No bias |
| 274 | z=gate, |
| 275 | eps=self.variance_epsilon, |
| 276 | group_size=self.group_size, |
| 277 | norm_before_gate=False |
| 278 | ) |
| 279 | |
| 280 | class NemotronHMamba2Mixer(nn.Module): |
| 281 | """ |
| 282 | Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. |
| 283 | A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) |
| 284 | ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, |
| 285 | and is why Mamba is called **selective** state spaces) |
| 286 | """ |
| 287 | |
| 288 | def __init__(self, config: NemotronHConfig, layer_idx: int): |
| 289 | super().__init__() |
| 290 | self.num_heads = config.mamba_num_heads |
| 291 | self.hidden_size = config.hidden_size |
| 292 | self.ssm_state_size = config.ssm_state_size |
| 293 | self.conv_kernel_size = config.conv_kernel |
| 294 | self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim |
| 295 | self.layer_idx = layer_idx |
| 296 | self.use_conv_bias = config.use_conv_bias |
| 297 | self.activation = config.mamba_hidden_act |
| 298 | self.act = ACT2FN[config.mamba_hidden_act] |
| 299 | |
| 300 | self.layer_norm_epsilon = config.layer_norm_epsilon |
| 301 | |
| 302 | self.n_groups = config.n_groups |
| 303 | self.head_dim = config.mamba_head_dim |
| 304 | self.chunk_size = config.chunk_size |
| 305 | |
| 306 | self.time_step_limit = config.time_step_limit |
| 307 | self.time_step_min = config.time_step_min |
| 308 | self.time_step_max = config.time_step_max |
| 309 | |
| 310 | self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size |
| 311 | self.conv1d = nn.Conv1d( |
| 312 | in_channels=self.conv_dim, |
| 313 | out_channels=self.conv_dim, |
| 314 | bias=config.use_conv_bias, |
| 315 | kernel_size=config.conv_kernel, |
| 316 | groups=self.conv_dim, |
| 317 | padding=config.conv_kernel - 1, |
| 318 | ) |
| 319 | |
| 320 | # projection of the input hidden states |
| 321 | projection_size = self.intermediate_size + self.conv_dim + self.num_heads |
| 322 | self.in_proj = nn.Linear( |
| 323 | self.hidden_size, |
| 324 | projection_size, |
| 325 | bias=config.use_bias, |
| 326 | ) |
| 327 | # selective projection used to make dt, B and C input dependant |
| 328 | |
| 329 | # time step projection (discretization) |
| 330 | # instantiate once and copy inv_dt in init_weights of PretrainedModel |
| 331 | self.dt_bias = nn.Parameter(torch.ones(self.num_heads)) |
| 332 | |
| 333 | # S4D real initialization. These are not discretized! |
| 334 | # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded |
| 335 | A = torch.arange(1, self.num_heads + 1) |
| 336 | self.A_log = nn.Parameter(torch.log(A)) |
| 337 | self.A_log._no_weight_decay = True |
| 338 | self.norm = MambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon, group_size=self.intermediate_size // self.n_groups) |
| 339 | self.D = nn.Parameter(torch.ones(self.num_heads)) |
| 340 | self.D._no_weight_decay = True |
| 341 | |
| 342 | self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) |
| 343 | self.use_bias = config.use_bias |
| 344 | |
| 345 | if not is_fast_path_available: |
| 346 | logger.warning_once( |
| 347 | "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`" |
| 348 | " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and" |
| 349 | " https://github.com/Dao-AILab/causal-conv1d" |
| 350 | ) |
| 351 | |
| 352 | def cuda_kernels_forward( |
| 353 | self, |
| 354 | hidden_states: torch.Tensor, |
| 355 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None, |
| 356 | cache_position: Optional[torch.LongTensor] = None, |
| 357 | attention_mask: Optional[torch.Tensor] = None, |
| 358 | ): |
| 359 | # 1. Gated MLP's linear projection |
| 360 | hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) |
| 361 | projected_states = self.in_proj(hidden_states) |
| 362 | |
| 363 | # Set up dimensions for reshapes later |
| 364 | batch_size, seq_len, _ = hidden_states.shape |
| 365 | groups_time_state_size = self.n_groups * self.ssm_state_size |
| 366 | d_mlp = ( |
| 367 | projected_states.shape[-1] |
| 368 | - 2 * self.intermediate_size |
| 369 | - 2 * self.n_groups * self.ssm_state_size |
| 370 | - self.num_heads |
| 371 | ) // 2 |
| 372 | |
| 373 | # Single step calculations via cache |
| 374 | if cache_params is not None and cache_position is not None and cache_position[0] > 0: |
| 375 | _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split( |
| 376 | [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 |
| 377 | ) |
| 378 | |
| 379 | # 2. Convolution sequence transformation |
| 380 | hidden_states_B_C = causal_conv1d_update( |
| 381 | hidden_states_B_C, |
| 382 | cache_params.conv_states[self.layer_idx], |
| 383 | self.conv1d.weight.squeeze(1), |
| 384 | self.conv1d.bias, |
| 385 | self.activation, |
| 386 | ) |
| 387 | |
| 388 | hidden_states, B, C = torch.split( |
| 389 | hidden_states_B_C, |
| 390 | [self.intermediate_size, groups_time_state_size, groups_time_state_size], |
| 391 | dim=-1, |
| 392 | ) |
| 393 | |
| 394 | # 3. SSM transformation |
| 395 | A = -torch.exp(self.A_log.float()) # (nheads,) |
| 396 | A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) |
| 397 | dt = dt[:, :, None].expand(-1, -1, self.head_dim) |
| 398 | dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) |
| 399 | D = self.D[:, None, ...].expand(-1, self.head_dim) |
| 400 | B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups) |
| 401 | C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups) |
| 402 | hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim) |
| 403 | hidden_states = selective_state_update( |
| 404 | cache_params.ssm_states[self.layer_idx], |
| 405 | hidden_states_reshaped, |
| 406 | dt, |
| 407 | A, |
| 408 | B, |
| 409 | C, |
| 410 | D, |
| 411 | z=None, |
| 412 | dt_bias=dt_bias, |
| 413 | dt_softplus=True, |
| 414 | ) |
| 415 | hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim) |
| 416 | hidden_states = self.norm(hidden_states, gate) |
| 417 | |
| 418 | # 4. Final linear projection |
| 419 | out = self.out_proj(hidden_states)[:, None, ...] |
| 420 | |
| 421 | # Fused calculations or step by step if no initialized cache is found |
| 422 | else: |
| 423 | A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) |
| 424 | dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit} |
| 425 | |
| 426 | # 2-4. Fused kernel for conv1d, SSM, and the final projection |
| 427 | if self.training and cache_params is None: |
| 428 | out = mamba_split_conv1d_scan_combined( |
| 429 | projected_states, |
| 430 | self.conv1d.weight.squeeze(1), |
| 431 | self.conv1d.bias, |
| 432 | self.dt_bias, |
| 433 | A, |
| 434 | D=self.D, |
| 435 | chunk_size=self.chunk_size, |
| 436 | seq_idx=None, # was seq_idx |
| 437 | activation=self.activation, |
| 438 | rmsnorm_weight=self.norm.weight, |
| 439 | rmsnorm_eps=self.norm.variance_epsilon, |
| 440 | outproj_weight=self.out_proj.weight, |
| 441 | outproj_bias=self.out_proj.bias, |
| 442 | headdim=self.head_dim, |
| 443 | ngroups=self.n_groups, |
| 444 | norm_before_gate=False, |
| 445 | return_final_states=False, |
| 446 | **dt_limit_kwargs, |
| 447 | ) |
| 448 | |
| 449 | else: |
| 450 | _, _, gate, hidden_states_B_C, dt = projected_states.split( |
| 451 | [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 |
| 452 | ) |
| 453 | |
| 454 | # 2. Convolution sequence transformation |
| 455 | # Init cache |
| 456 | if cache_params is not None: |
| 457 | hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2) |
| 458 | conv_states = nn.functional.pad( |
| 459 | hidden_states_B_C_transposed, |
| 460 | (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0), |
| 461 | ) |
| 462 | cache_params.update_conv_state( |
| 463 | layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True |
| 464 | ) |
| 465 | |
| 466 | if self.activation not in ["silu", "swish"]: |
| 467 | hidden_states_B_C = self.act( |
| 468 | self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2) |
| 469 | ) |
| 470 | else: |
| 471 | hidden_states_B_C = causal_conv1d_fn( |
| 472 | x=hidden_states_B_C.transpose(1, 2), |
| 473 | weight=self.conv1d.weight.squeeze(1), |
| 474 | bias=self.conv1d.bias, |
| 475 | activation=self.activation, |
| 476 | ).transpose(1, 2) |
| 477 | hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) |
| 478 | hidden_states, B, C = torch.split( |
| 479 | hidden_states_B_C, |
| 480 | [self.intermediate_size, groups_time_state_size, groups_time_state_size], |
| 481 | dim=-1, |
| 482 | ) |
| 483 | |
| 484 | # 3. SSM transformation |
| 485 | scan_output, ssm_state = mamba_chunk_scan_combined( |
| 486 | hidden_states.view(batch_size, seq_len, -1, self.head_dim), |
| 487 | dt, |
| 488 | A, |
| 489 | B.view(batch_size, seq_len, self.n_groups, -1), |
| 490 | C.view(batch_size, seq_len, self.n_groups, -1), |
| 491 | chunk_size=self.chunk_size, |
| 492 | D=self.D, |
| 493 | z=None, |
| 494 | seq_idx=None, |
| 495 | return_final_states=True, |
| 496 | dt_bias=self.dt_bias, |
| 497 | dt_softplus=True, |
| 498 | **dt_limit_kwargs, |
| 499 | ) |
| 500 | |
| 501 | # Init cache |
| 502 | if ssm_state is not None and cache_params is not None: |
| 503 | cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state) |
| 504 | |
| 505 | scan_output = scan_output.view(batch_size, seq_len, -1) |
| 506 | |
| 507 | # Multiply "gate" branch and apply extra normalization layer |
| 508 | scan_output = self.norm(scan_output, gate) |
| 509 | |
| 510 | # 4. Final linear projection |
| 511 | out = self.out_proj(scan_output) |
| 512 | return out |
| 513 | |
| 514 | # fmt: off |
| 515 | def torch_forward(self, input_states, cache_params: Optional[HybridMambaAttentionDynamicCache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None): |
| 516 | batch_size, seq_len, _ = input_states.shape |
| 517 | dtype = input_states.dtype |
| 518 | |
| 519 | # 1. Gated MLP's linear projection |
| 520 | input_states = apply_mask_to_padding_states(input_states, attention_mask) |
| 521 | projected_states = self.in_proj(input_states) |
| 522 | d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size-self.num_heads) // 2 |
| 523 | _, _, gate, hidden_states_B_C, dt = projected_states.split( |
| 524 | [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 |
| 525 | ) |
| 526 | |
| 527 | # 2. Convolution sequence transformation |
| 528 | if cache_params is not None and cache_position is not None and cache_position[0] > 0: |
| 529 | cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=hidden_states_B_C, cache_init=False) |
| 530 | |
| 531 | # We need to guarantee that anything regarding the cache is on the same device |
| 532 | conv_states = cache_params.conv_states[self.layer_idx].to(device=self.conv1d.weight.device) |
| 533 | |
| 534 | hidden_states_B_C = torch.sum( |
| 535 | conv_states * self.conv1d.weight.squeeze(1), dim=-1 |
| 536 | ) |
| 537 | if self.use_conv_bias: |
| 538 | hidden_states_B_C = hidden_states_B_C + self.conv1d.bias |
| 539 | hidden_states_B_C = self.act(hidden_states_B_C) |
| 540 | else: |
| 541 | # Init cache |
| 542 | if cache_params is not None: |
| 543 | hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2) |
| 544 | conv_states = nn.functional.pad( |
| 545 | hidden_states_B_C_transposed, (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0) |
| 546 | ) |
| 547 | cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True) |
| 548 | |
| 549 | hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)) |
| 550 | |
| 551 | hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) |
| 552 | hidden_states, B, C = torch.split( |
| 553 | hidden_states_B_C, |
| 554 | [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], |
| 555 | dim=-1 |
| 556 | ) |
| 557 | |
| 558 | # 3. SSM transformation |
| 559 | A = -torch.exp(self.A_log.float()) # [num_heads] |
| 560 | if cache_params is not None and cache_position is not None and cache_position[0] > 0: |
| 561 | # We need to guarantee that anything regarding the cache is on the same device |
| 562 | cache_device = cache_params.ssm_states.device |
| 563 | |
| 564 | # Note: there is no need to pad parameter matrices here, as there is just one new token |
| 565 | # for batched generation |
| 566 | dt = dt[:, 0, :][:, None, ...] |
| 567 | dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim) |
| 568 | # [num_heads] -> [num_heads, head_dim] |
| 569 | dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim) |
| 570 | |
| 571 | dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype)) |
| 572 | dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) |
| 573 | A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32) |
| 574 | # [bsz, num_heads, head_dim, state_size] |
| 575 | dA = (torch.exp(dt[..., None] * A)).to(device=cache_device) |
| 576 | |
| 577 | # Discretize B |
| 578 | # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] -> |
| 579 | # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size] |
| 580 | B = B.reshape(batch_size, self.n_groups, -1)[..., None, :] |
| 581 | B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous() |
| 582 | B = B.reshape(batch_size, -1, B.shape[-1]) |
| 583 | # [bsz, num_heads, head_dim, state_size] |
| 584 | dB = dt[..., None] * B[..., None, :] |
| 585 | |
| 586 | # Discretize x into dB |
| 587 | # [bsz, intermediate_size] -> [bsz, num_heads, head_dim] |
| 588 | hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim) |
| 589 | dBx = (dB * hidden_states[..., None]).to(device=cache_device) |
| 590 | |
| 591 | # State calculation |
| 592 | cache_params.update_ssm_state( |
| 593 | layer_idx=self.layer_idx, |
| 594 | new_ssm_state=cache_params.ssm_states[self.layer_idx] * dA + dBx |
| 595 | ) |
| 596 | |
| 597 | # Subsequent output |
| 598 | # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size] |
| 599 | C = C.reshape(batch_size, self.n_groups, -1)[..., None, :] |
| 600 | C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous() |
| 601 | C = C.reshape(batch_size, -1, C.shape[-1]) |
| 602 | # [bsz, num_heads, head_dim] |
| 603 | |
| 604 | ssm_states = cache_params.ssm_states[self.layer_idx].to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n] |
| 605 | # Reshape ssm_states to merge the first two dimensions |
| 606 | ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) # Shape: [b*h, d, n] |
| 607 | C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1] |
| 608 | y = torch.bmm(ssm_states_reshaped, C_reshaped) |
| 609 | y = y.view(batch_size, self.num_heads, self.head_dim) |
| 610 | |
| 611 | # D skip connection |
| 612 | # [num_heads] -> [num_heads, head_dim] |
| 613 | D = self.D[..., None].expand(self.D.shape[0], self.head_dim) |
| 614 | y = (y + hidden_states * D).to(y.dtype) |
| 615 | |
| 616 | # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size] |
| 617 | y = y.reshape(batch_size, -1)[:, None, ...] |
| 618 | else: |
| 619 | # begin ssd naive implementation without einsums |
| 620 | dt = nn.functional.softplus(dt + self.dt_bias) |
| 621 | dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1]) |
| 622 | hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float() |
| 623 | B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() |
| 624 | C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float() |
| 625 | B = B.repeat(1, 1, self.num_heads // self.n_groups, 1) |
| 626 | C = C.repeat(1, 1, self.num_heads // self.n_groups, 1) |
| 627 | pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size |
| 628 | |
| 629 | D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size) |
| 630 | |
| 631 | # Discretize x and A |
| 632 | hidden_states = hidden_states * dt[..., None] |
| 633 | A = A.to(hidden_states.dtype) * dt |
| 634 | |
| 635 | # Rearrange into blocks/chunks |
| 636 | hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)] |
| 637 | |
| 638 | # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size] |
| 639 | A = A.permute(0, 3, 1, 2) |
| 640 | A_cumsum = torch.cumsum(A, dim=-1) |
| 641 | |
| 642 | # 1. Compute the output for each intra-chunk (diagonal blocks) |
| 643 | # This is the analog of a causal mask |
| 644 | L = torch.exp(segment_sum(A)) |
| 645 | |
| 646 | # Contraction of C and B to get G (attention-weights like) |
| 647 | G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n) |
| 648 | G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h) |
| 649 | |
| 650 | # Compute M, equivalent to applying attention mask to weights |
| 651 | M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None] |
| 652 | M = M_intermediate.sum(dim=-1) |
| 653 | |
| 654 | # Compute Y_diag (apply to values) |
| 655 | Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) |
| 656 | |
| 657 | # 2. Compute the state for each intra-chunk |
| 658 | # (right term of low-rank factorization of off-diagonal blocks; B terms) |
| 659 | decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum)) |
| 660 | B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] |
| 661 | states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) |
| 662 | |
| 663 | # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries |
| 664 | # (middle term of factorization of off-diag blocks; A terms) |
| 665 | if cache_params is not None and cache_position is not None and cache_position[0] > 0: |
| 666 | previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...].to(device=states.device) |
| 667 | else: |
| 668 | previous_states = torch.zeros_like(states[:, :1]) |
| 669 | states = torch.cat([previous_states, states], dim=1) |
| 670 | decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0)))) |
| 671 | decay_chunk = decay_chunk.transpose(1, 3) |
| 672 | new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) |
| 673 | states, ssm_state = new_states[:, :-1], new_states[:, -1] |
| 674 | |
| 675 | # 4. Compute state -> output conversion per chunk |
| 676 | # (left term of low-rank factorization of off-diagonal blocks; C terms) |
| 677 | state_decay_out = torch.exp(A_cumsum) |
| 678 | C_times_states = (C[..., None, :] * states[:, :, None, ...]) |
| 679 | state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1) |
| 680 | Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None]) |
| 681 | |
| 682 | # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) |
| 683 | y = Y_diag + Y_off |
| 684 | # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim] |
| 685 | y = y.reshape(batch_size, -1, self.num_heads, self.head_dim) |
| 686 | |
| 687 | y = y + D_residual |
| 688 | # Cutting off padded chunks |
| 689 | if pad_size > 0: |
| 690 | y = y[:, :seq_len, :, :] |
| 691 | y = y.reshape(batch_size, seq_len, -1) |
| 692 | |
| 693 | # Init cache |
| 694 | if ssm_state is not None and cache_params is not None: |
| 695 | cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state) |
| 696 | |
| 697 | scan_output = self.norm(y, gate) |
| 698 | |
| 699 | # end ssd naive |
| 700 | |
| 701 | # 4. Final linear projection |
| 702 | contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size] |
| 703 | return contextualized_states |
| 704 | # fmt: on |
| 705 | |
| 706 | def forward( |
| 707 | self, |
| 708 | hidden_states, |
| 709 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None, |
| 710 | cache_position: Optional[torch.LongTensor] = None, |
| 711 | attention_mask: Optional[torch.Tensor] = None, |
| 712 | ): |
| 713 | if is_fast_path_available and "cuda" in self.in_proj.weight.device.type: |
| 714 | return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask) |
| 715 | dtype = hidden_states.dtype |
| 716 | if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1: |
| 717 | # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66 |
| 718 | hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) |
| 719 | |
| 720 | return self.torch_forward(hidden_states, cache_params, cache_position, attention_mask) |
| 721 | |
| 722 | |
| 723 | class NemotronHRMSNorm(nn.Module): |
| 724 | def __init__(self, hidden_size, eps=1e-6): |
| 725 | """ |
| 726 | NemotronHRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm |
| 727 | """ |
| 728 | super().__init__() |
| 729 | self.weight = nn.Parameter(torch.ones(hidden_size)) |
| 730 | self.variance_epsilon = eps |
| 731 | |
| 732 | def forward(self, hidden_states): |
| 733 | input_dtype = hidden_states.dtype |
| 734 | hidden_states = hidden_states.to(torch.float32) |
| 735 | variance = hidden_states.pow(2).mean(-1, keepdim=True) |
| 736 | hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
| 737 | # Weights are in float32 |
| 738 | return (self.weight.to(torch.float32) * hidden_states).to(input_dtype) |
| 739 | |
| 740 | class NemotronHBlock(nn.Module): |
| 741 | def __init__(self, config, layer_idx): |
| 742 | super().__init__() |
| 743 | self.config = config |
| 744 | self.layer_idx = layer_idx |
| 745 | self.residual_in_fp32 = config.residual_in_fp32 |
| 746 | self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) |
| 747 | |
| 748 | # M: Mamba2, *: Attention, -: MLP |
| 749 | self.block_type = config.layers_block_type[layer_idx] |
| 750 | if self.block_type == "mamba": |
| 751 | self.mixer = NemotronHMamba2Mixer(config, layer_idx=layer_idx) |
| 752 | elif self.block_type == "attention": |
| 753 | self.mixer = NEMOTRONH_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx) |
| 754 | elif self.block_type == "mlp": |
| 755 | self.mixer = NemotronHMLP(config, layer_idx=layer_idx) |
| 756 | else: |
| 757 | raise ValueError(f"Invalid layer pattern {config.hybrid_override_pattern[layer_idx]}") |
| 758 | |
| 759 | def forward( |
| 760 | self, |
| 761 | hidden_states, |
| 762 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None, |
| 763 | cache_position: Optional[torch.LongTensor] = None, |
| 764 | attention_mask: Optional[torch.Tensor] = None, |
| 765 | ): |
| 766 | with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)): |
| 767 | # * Use torch.cuda.stream() to avoid NaN issues when using multiple GPUs |
| 768 | residual = hidden_states |
| 769 | hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) |
| 770 | if self.residual_in_fp32: |
| 771 | residual = residual.to(torch.float32) |
| 772 | |
| 773 | if self.block_type == "mamba": |
| 774 | hidden_states = self.mixer( |
| 775 | hidden_states, cache_params=cache_params, cache_position=cache_position |
| 776 | ) |
| 777 | elif self.block_type == "attention": |
| 778 | hidden_states = self.mixer( |
| 779 | hidden_states, cache_position=cache_position |
| 780 | ) |
| 781 | hidden_states = hidden_states[0] |
| 782 | elif self.block_type == "mlp": |
| 783 | hidden_states = self.mixer( |
| 784 | hidden_states |
| 785 | ) |
| 786 | else: |
| 787 | raise ValueError(f"Invalid block_type: {self.block_type}") |
| 788 | |
| 789 | hidden_states = residual + hidden_states |
| 790 | return hidden_states |
| 791 | |
| 792 | |
| 793 | # Copied from transformers.models.nemotron.modeling_nemotron Nemotron->NemotronH |
| 794 | class NemotronHMLP(nn.Module): |
| 795 | def __init__(self, config, layer_idx: Optional[int] = None): |
| 796 | super().__init__() |
| 797 | self.config = config |
| 798 | self.layer_idx = layer_idx |
| 799 | if layer_idx is None: |
| 800 | logger.warning_once( |
| 801 | f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " |
| 802 | "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " |
| 803 | "when creating this class." |
| 804 | ) |
| 805 | self.hidden_size = config.hidden_size |
| 806 | self.mlp_idx = config.hybrid_override_pattern[:layer_idx+1].count("-")-1 |
| 807 | if isinstance(config.intermediate_size, list): |
| 808 | if len(config.intermediate_size) == 1: |
| 809 | self.intermediate_size = config.intermediate_size[0] |
| 810 | else: |
| 811 | self.intermediate_size = config.intermediate_size[self.mlp_idx] |
| 812 | else: |
| 813 | self.intermediate_size = config.intermediate_size |
| 814 | self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) |
| 815 | self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) |
| 816 | self.act_fn = ACT2FN[config.mlp_hidden_act] |
| 817 | |
| 818 | def forward(self, x): |
| 819 | return self.down_proj(self.act_fn(self.up_proj(x))) |
| 820 | |
| 821 | |
| 822 | # Copied from transformers.models.llama.modeling_llama.repeat_kv |
| 823 | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: |
| 824 | """ |
| 825 | This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, |
| 826 | num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) |
| 827 | """ |
| 828 | batch, num_key_value_heads, slen, head_dim = hidden_states.shape |
| 829 | if n_rep == 1: |
| 830 | return hidden_states |
| 831 | hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) |
| 832 | return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) |
| 833 | |
| 834 | |
| 835 | class NemotronHAttention(nn.Module): |
| 836 | """Multi-headed attention from 'Attention Is All You Need' paper""" |
| 837 | |
| 838 | def __init__(self, config: NemotronHConfig, layer_idx: Optional[int] = None): |
| 839 | super().__init__() |
| 840 | self.config = config |
| 841 | self.layer_idx = layer_idx |
| 842 | if layer_idx is None: |
| 843 | logger.warning_once( |
| 844 | f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " |
| 845 | "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " |
| 846 | "when creating this class." |
| 847 | ) |
| 848 | |
| 849 | self.attention_dropout = config.attention_dropout |
| 850 | self.hidden_size = config.hidden_size |
| 851 | self.num_heads = config.num_attention_heads |
| 852 | if config.attention_head_dim is not None: |
| 853 | self.head_dim = config.attention_head_dim |
| 854 | else: |
| 855 | self.head_dim = config.hidden_size // config.num_attention_heads |
| 856 | self.num_key_value_heads = config.num_key_value_heads |
| 857 | self.num_key_value_groups = self.num_heads // self.num_key_value_heads |
| 858 | self.max_position_embeddings = config.max_position_embeddings |
| 859 | self.is_causal = True |
| 860 | |
| 861 | self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) |
| 862 | self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) |
| 863 | self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) |
| 864 | self.o_proj = nn.Linear(self.head_dim * self.num_heads, self.hidden_size, bias=config.attention_bias) |
| 865 | |
| 866 | def forward( |
| 867 | self, |
| 868 | hidden_states: torch.Tensor, |
| 869 | # position_embeddings: Tuple[torch.Tensor, torch.Tensor], #TODO |
| 870 | attention_mask: Optional[torch.Tensor] = None, |
| 871 | position_ids: Optional[torch.LongTensor] = None, |
| 872 | past_key_value: Optional[HybridMambaAttentionDynamicCache] = None, |
| 873 | output_attentions: bool = False, |
| 874 | use_cache: bool = False, |
| 875 | cache_position: Optional[torch.LongTensor] = None, |
| 876 | ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| 877 | bsz, q_len, _ = hidden_states.size() |
| 878 | |
| 879 | query_states = self.q_proj(hidden_states) |
| 880 | key_states = self.k_proj(hidden_states) |
| 881 | value_states = self.v_proj(hidden_states) |
| 882 | |
| 883 | query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| 884 | key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 885 | value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 886 | |
| 887 | if past_key_value is not None: |
| 888 | key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) |
| 889 | |
| 890 | key_states = repeat_kv(key_states, self.num_key_value_groups) |
| 891 | value_states = repeat_kv(value_states, self.num_key_value_groups) |
| 892 | |
| 893 | causal_mask = attention_mask |
| 894 | if attention_mask is not None: # no matter the length, we just slice it |
| 895 | causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] |
| 896 | |
| 897 | if query_states.device.type == "cuda" and attention_mask is not None: |
| 898 | query_states = query_states.contiguous() |
| 899 | key_states = key_states.contiguous() |
| 900 | value_states = value_states.contiguous() |
| 901 | |
| 902 | is_causal = True if causal_mask is None and q_len > 1 else False |
| 903 | |
| 904 | attn_output = torch.nn.functional.scaled_dot_product_attention( |
| 905 | query_states, |
| 906 | key_states, |
| 907 | value_states, |
| 908 | attn_mask=causal_mask, |
| 909 | dropout_p=self.attention_dropout if self.training else 0.0, |
| 910 | is_causal=is_causal, |
| 911 | ) |
| 912 | attn_output = attn_output.transpose(1, 2).contiguous() |
| 913 | #attn_output = attn_output.view(bsz, q_len, self.hidden_size) |
| 914 | attn_output = attn_output.view(bsz, q_len, self.num_heads * self.head_dim) |
| 915 | |
| 916 | attn_output = self.o_proj(attn_output) |
| 917 | |
| 918 | return attn_output, None, past_key_value |
| 919 | |
| 920 | |
| 921 | # Adapted from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->Jamba |
| 922 | #class JambaFlashAttention2(JambaAttention): |
| 923 | class NemotronHFlashAttention2(NemotronHAttention): |
| 924 | """ |
| 925 | Jamba flash attention module. This module inherits from `JambaAttention` as the weights of the module stays |
| 926 | untouched. The only required change would be on the forward pass where it needs to correctly call the public API of |
| 927 | flash attention and deal with padding tokens in case the input contains any of them. |
| 928 | """ |
| 929 | def __init__(self, *args, **kwargs): |
| 930 | super().__init__(*args, **kwargs) |
| 931 | |
| 932 | # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. |
| 933 | # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. |
| 934 | # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). |
| 935 | self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() |
| 936 | |
| 937 | def forward( |
| 938 | self, |
| 939 | hidden_states: torch.Tensor, |
| 940 | attention_mask: Optional[torch.Tensor] = None, |
| 941 | position_ids: Optional[torch.LongTensor] = None, |
| 942 | past_key_value: Optional[HybridMambaAttentionDynamicCache] = None, |
| 943 | output_attentions: bool = False, |
| 944 | use_cache: bool = False, |
| 945 | cache_position: Optional[torch.LongTensor] = None, |
| 946 | **kwargs, |
| 947 | ): |
| 948 | bsz, q_len, _ = hidden_states.size() |
| 949 | |
| 950 | query_states = self.q_proj(hidden_states) |
| 951 | key_states = self.k_proj(hidden_states) |
| 952 | value_states = self.v_proj(hidden_states) |
| 953 | |
| 954 | # Flash attention requires the input to have the shape |
| 955 | # batch_size x seq_length x head_dim x hidden_dim |
| 956 | # therefore we just need to keep the original shape |
| 957 | query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) |
| 958 | key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 959 | value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 960 | |
| 961 | if past_key_value is not None: |
| 962 | key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) |
| 963 | |
| 964 | # repeat k/v heads if n_kv_heads < n_heads |
| 965 | key_states = repeat_kv(key_states, self.num_key_value_groups) |
| 966 | value_states = repeat_kv(value_states, self.num_key_value_groups) |
| 967 | dropout_rate = 0.0 if not self.training else self.attention_dropout |
| 968 | |
| 969 | # In PEFT, usually we cast the layer norms in float32 for training stability reasons |
| 970 | # therefore the input hidden states gets silently casted in float32. Hence, we need |
| 971 | # cast them back in float16 just to be sure everything works as expected. |
| 972 | input_dtype = query_states.dtype |
| 973 | if input_dtype == torch.float32: |
| 974 | if torch.is_autocast_enabled(): |
| 975 | target_dtype = torch.get_autocast_gpu_dtype() |
| 976 | # Handle the case where the model is quantized |
| 977 | elif hasattr(self.config, "_pre_quantization_dtype"): |
| 978 | target_dtype = self.config._pre_quantization_dtype |
| 979 | else: |
| 980 | target_dtype = self.q_proj.weight.dtype |
| 981 | |
| 982 | logger.warning_once( |
| 983 | f"The input hidden states seems to be silently casted in float32, this might be related to" |
| 984 | f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" |
| 985 | f" {target_dtype}." |
| 986 | ) |
| 987 | |
| 988 | query_states = query_states.to(target_dtype) |
| 989 | key_states = key_states.to(target_dtype) |
| 990 | value_states = value_states.to(target_dtype) |
| 991 | |
| 992 | # Reashape to the expected shape for Flash Attention |
| 993 | key_states = key_states.transpose(1, 2) |
| 994 | value_states = value_states.transpose(1, 2) |
| 995 | |
| 996 | attn_output = _flash_attention_forward( |
| 997 | query_states, |
| 998 | key_states, |
| 999 | value_states, |
| 1000 | attention_mask, |
| 1001 | q_len, |
| 1002 | dropout=dropout_rate, |
| 1003 | sliding_window=getattr(self.config, "sliding_window", None), |
| 1004 | is_causal=self.is_causal, |
| 1005 | use_top_left_mask=self._flash_attn_uses_top_left_mask, |
| 1006 | ) |
| 1007 | |
| 1008 | #attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() |
| 1009 | attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous() |
| 1010 | attn_output = self.o_proj(attn_output) |
| 1011 | |
| 1012 | if not output_attentions: |
| 1013 | attn_weights = None |
| 1014 | |
| 1015 | return attn_output, attn_weights, past_key_value |
| 1016 | |
| 1017 | |
| 1018 | # Adapted from transformers.models.mistral.modeling_mistral.MistralSdpaAttention with Mistral->Jamba |
| 1019 | #class JambaSdpaAttention(JambaAttention): |
| 1020 | class NemotronHSdpaAttention(NemotronHAttention): |
| 1021 | """ |
| 1022 | Jamba attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from |
| 1023 | `JambaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to |
| 1024 | SDPA API. |
| 1025 | """ |
| 1026 | |
| 1027 | # Adapted from NemotronHAttention.forward |
| 1028 | def forward( |
| 1029 | self, |
| 1030 | hidden_states: torch.Tensor, |
| 1031 | attention_mask: Optional[torch.Tensor] = None, |
| 1032 | position_ids: Optional[torch.LongTensor] = None, |
| 1033 | past_key_value: Optional[HybridMambaAttentionDynamicCache] = None, |
| 1034 | output_attentions: bool = False, |
| 1035 | use_cache: bool = False, |
| 1036 | cache_position: Optional[torch.LongTensor] = None, |
| 1037 | ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
| 1038 | if output_attentions: |
| 1039 | # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. |
| 1040 | logger.warning_once( |
| 1041 | "NemotronHModel is using NemotronHSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " |
| 1042 | 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' |
| 1043 | ) |
| 1044 | return super().forward( |
| 1045 | hidden_states=hidden_states, |
| 1046 | attention_mask=attention_mask, |
| 1047 | position_ids=position_ids, |
| 1048 | past_key_value=past_key_value, |
| 1049 | output_attentions=output_attentions, |
| 1050 | use_cache=use_cache, |
| 1051 | ) |
| 1052 | |
| 1053 | bsz, q_len, _ = hidden_states.size() |
| 1054 | |
| 1055 | query_states = self.q_proj(hidden_states) |
| 1056 | key_states = self.k_proj(hidden_states) |
| 1057 | value_states = self.v_proj(hidden_states) |
| 1058 | |
| 1059 | query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) |
| 1060 | key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 1061 | value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) |
| 1062 | |
| 1063 | if past_key_value is not None: |
| 1064 | key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) |
| 1065 | |
| 1066 | key_states = repeat_kv(key_states, self.num_key_value_groups) |
| 1067 | value_states = repeat_kv(value_states, self.num_key_value_groups) |
| 1068 | |
| 1069 | causal_mask = attention_mask |
| 1070 | if attention_mask is not None: |
| 1071 | causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] |
| 1072 | |
| 1073 | # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, |
| 1074 | # Reference: https://github.com/pytorch/pytorch/issues/112577. |
| 1075 | if query_states.device.type == "cuda" and attention_mask is not None: |
| 1076 | query_states = query_states.contiguous() |
| 1077 | key_states = key_states.contiguous() |
| 1078 | value_states = value_states.contiguous() |
| 1079 | |
| 1080 | # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment |
| 1081 | # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. |
| 1082 | # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. |
| 1083 | is_causal = True if self.is_causal and causal_mask is None and q_len > 1 else False |
| 1084 | |
| 1085 | attn_output = torch.nn.functional.scaled_dot_product_attention( |
| 1086 | query_states, |
| 1087 | key_states, |
| 1088 | value_states, |
| 1089 | attn_mask=causal_mask, |
| 1090 | dropout_p=self.attention_dropout if self.training else 0.0, |
| 1091 | is_causal=is_causal, |
| 1092 | ) |
| 1093 | |
| 1094 | attn_output = attn_output.transpose(1, 2).contiguous() |
| 1095 | attn_output = attn_output.view(bsz, q_len, self.hidden_size) |
| 1096 | |
| 1097 | attn_output = self.o_proj(attn_output) |
| 1098 | |
| 1099 | return attn_output, None, past_key_value |
| 1100 | |
| 1101 | |
| 1102 | NEMOTRONH_ATTENTION_CLASSES = { |
| 1103 | "eager": NemotronHAttention, |
| 1104 | "flash_attention_2": NemotronHFlashAttention2, |
| 1105 | "sdpa": NemotronHSdpaAttention, |
| 1106 | } |
| 1107 | |
| 1108 | # Copied from transformers.models.mamba.modeling_mamba2.Mamba2PreTrainedModel |
| 1109 | class NemotronHPreTrainedModel(PreTrainedModel): |
| 1110 | """ |
| 1111 | An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained |
| 1112 | models. |
| 1113 | """ |
| 1114 | |
| 1115 | config_class = NemotronHConfig |
| 1116 | base_model_prefix = "backbone" |
| 1117 | _no_split_modules = ["NemotronHBlock"] |
| 1118 | supports_gradient_checkpointing = True |
| 1119 | _is_stateful = True |
| 1120 | |
| 1121 | def _init_weights(self, module): |
| 1122 | """Initialize the weights.""" |
| 1123 | if isinstance(module, NemotronHMamba2Mixer): |
| 1124 | module.A_log._no_weight_decay = True |
| 1125 | module.D._no_weight_decay = True |
| 1126 | |
| 1127 | dt = torch.exp( |
| 1128 | torch.rand(self.config.mamba_num_heads) |
| 1129 | * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) |
| 1130 | + math.log(self.config.time_step_min) |
| 1131 | ).clamp(min=self.config.time_step_floor) |
| 1132 | |
| 1133 | # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 |
| 1134 | inv_dt = dt + torch.log(-torch.expm1(-dt)) |
| 1135 | with torch.no_grad(): |
| 1136 | module.dt_bias.copy_(inv_dt) |
| 1137 | module.dt_bias._no_reinit = True |
| 1138 | |
| 1139 | if isinstance(module, nn.Linear): |
| 1140 | if module.bias is not None: |
| 1141 | if not getattr(module.bias, "_no_reinit", False): |
| 1142 | nn.init.zeros_(module.bias) |
| 1143 | elif isinstance(module, nn.Embedding): |
| 1144 | nn.init.normal_(module.weight, std=self.config.initializer_range) |
| 1145 | |
| 1146 | # TODO: Check |
| 1147 | if self.config.rescale_prenorm_residual: |
| 1148 | # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: |
| 1149 | # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale |
| 1150 | # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. |
| 1151 | # > -- GPT-2 :: https://openai.com/blog/better-language-models/ |
| 1152 | # |
| 1153 | # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py |
| 1154 | for name, p in module.named_parameters(): |
| 1155 | if name in ["out_proj.weight"]: |
| 1156 | # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block |
| 1157 | # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) |
| 1158 | # We need to reinit p since this code could be called multiple times |
| 1159 | # Having just p *= scale would repeatedly scale it down |
| 1160 | nn.init.kaiming_uniform_(p, a=math.sqrt(5)) |
| 1161 | with torch.no_grad(): |
| 1162 | p /= math.sqrt(self.config.num_hidden_layers) |
| 1163 | |
| 1164 | |
| 1165 | @dataclass |
| 1166 | # Copied from transformers.models.mamba.modeling_mamba2.Mamba2Output with MAMBA2->NemotronH,Mamba2->NemotronH |
| 1167 | class NemotronHOutput(ModelOutput): |
| 1168 | """ |
| 1169 | Class for the NemotronH model outputs. |
| 1170 | |
| 1171 | Args: |
| 1172 | last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): |
| 1173 | Sequence of hidden-states at the output of the last layer of the model. |
| 1174 | cache_params (`HybridMambaAttentionDynamicCache`): |
| 1175 | The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to |
| 1176 | avoid providing the old `input_ids`. |
| 1177 | |
| 1178 | Includes both the State space model state matrices after the selective scan, and the Convolutional states |
| 1179 | hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): |
| 1180 | Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + |
| 1181 | one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. |
| 1182 | |
| 1183 | Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. |
| 1184 | """ |
| 1185 | |
| 1186 | last_hidden_state: Optional[torch.FloatTensor] = None |
| 1187 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None |
| 1188 | hidden_states: Optional[Tuple[torch.FloatTensor]] = None |
| 1189 | attentions: Optional[Tuple[torch.FloatTensor]] = None |
| 1190 | |
| 1191 | |
| 1192 | @dataclass |
| 1193 | # Copied from transformers.models.mamba2.modeling_mamba2.MambaCausalLMOutput with Mamba2->NemotronH |
| 1194 | class NemotronHCausalLMOutput(ModelOutput): |
| 1195 | """ |
| 1196 | Base class for causal language model (or autoregressive) outputs. |
| 1197 | |
| 1198 | Args: |
| 1199 | loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): |
| 1200 | Language modeling loss (for next-token prediction). |
| 1201 | logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): |
| 1202 | Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). |
| 1203 | cache_params (`HybridMambaAttentionDynamicCache`): |
| 1204 | The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to |
| 1205 | avoid providing the old `input_ids`. |
| 1206 | |
| 1207 | Includes both the State space model state matrices after the selective scan, and the Convolutional states |
| 1208 | hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): |
| 1209 | Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + |
| 1210 | one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. |
| 1211 | |
| 1212 | Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. |
| 1213 | """ |
| 1214 | |
| 1215 | loss: Optional[torch.FloatTensor] = None |
| 1216 | logits: Optional[torch.FloatTensor] = None |
| 1217 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None |
| 1218 | hidden_states: Optional[Tuple[torch.FloatTensor]] = None |
| 1219 | attentions: Optional[Tuple[torch.FloatTensor]] = None |
| 1220 | |
| 1221 | |
| 1222 | NEMOTRONH_START_DOCSTRING = r""" |
| 1223 | |
| 1224 | This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the |
| 1225 | library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads |
| 1226 | etc.) |
| 1227 | |
| 1228 | This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. |
| 1229 | Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage |
| 1230 | and behavior. |
| 1231 | |
| 1232 | Parameters: |
| 1233 | config ([`NemotronHConfig`]): Model configuration class with all the parameters of the model. |
| 1234 | Initializing with a config file does not load the weights associated with the model, only the |
| 1235 | configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. |
| 1236 | """ |
| 1237 | |
| 1238 | NEMOTRONH_INPUTS_DOCSTRING = r""" |
| 1239 | Args: |
| 1240 | input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*): |
| 1241 | Indices of input sequence tokens in the vocabulary. |
| 1242 | |
| 1243 | If `cache_params.seqlen_offset>0`, only `input_ids` that do not have their past calculated should be passed as |
| 1244 | `input_ids`. |
| 1245 | |
| 1246 | Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and |
| 1247 | [`PreTrainedTokenizer.__call__`] for details. |
| 1248 | |
| 1249 | [What are input IDs?](../glossary#input-ids) |
| 1250 | inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): |
| 1251 | Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This |
| 1252 | is useful if you want more control over how to convert `input_ids` indices into associated vectors than the |
| 1253 | model's internal embedding lookup matrix. |
| 1254 | position_ids (`torch.LongTensor` of shape `(batch_size,)`, *optional*): |
| 1255 | Indices of positions of each input sequence tokens in the position embeddings. |
| 1256 | cache_params (`HybridMambaAttentionDynamicCache`, *optional*): |
| 1257 | If passed along, the model uses the previous state in all the blocks (which will give the output for the |
| 1258 | `input_ids` provided as if the model add `state_input_ids + input_ids` as context). |
| 1259 | use_cache (`bool`, *optional*): |
| 1260 | If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits. |
| 1261 | output_attentions (`bool`, *optional*): |
| 1262 | Whether or not to return the attentions tensors of all attention layers. |
| 1263 | output_hidden_states (`bool`, *optional*): |
| 1264 | Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for |
| 1265 | more detail. |
| 1266 | return_dict (`bool`, *optional*): |
| 1267 | Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. |
| 1268 | cache_position (`torch.LongTensor` of shape `(batch_size,)`, *optional*): |
| 1269 | The position of the current input in the cache. This is used to ensure that the cache is correctly updated. |
| 1270 | If `cache_params` is passed, `cache_position` should also be passed. |
| 1271 | attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| 1272 | Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: |
| 1273 | |
| 1274 | - 1 for tokens that are **not masked**, |
| 1275 | - 0 for tokens that are **masked**. |
| 1276 | |
| 1277 | [What are attention masks?](../glossary#attention-mask) |
| 1278 | """ |
| 1279 | |
| 1280 | |
| 1281 | @add_start_docstrings( |
| 1282 | "The bare NemotronH Model transformer outputting raw hidden-states without any specific head on top.", |
| 1283 | NEMOTRONH_START_DOCSTRING, |
| 1284 | ) |
| 1285 | class NemotronHModel(NemotronHPreTrainedModel): |
| 1286 | def __init__(self, config): |
| 1287 | super().__init__(config) |
| 1288 | |
| 1289 | self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) |
| 1290 | self.layers = nn.ModuleList([NemotronHBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) |
| 1291 | |
| 1292 | self.gradient_checkpointing = False |
| 1293 | self.norm_f = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) |
| 1294 | # Initialize weights and apply final processing |
| 1295 | self._register_load_state_dict_pre_hook(self.load_hook) |
| 1296 | self.post_init() |
| 1297 | |
| 1298 | def load_hook(self, state_dict, prefix, *args): |
| 1299 | for k in state_dict: |
| 1300 | if "embedding." in k: |
| 1301 | state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) |
| 1302 | break |
| 1303 | |
| 1304 | def get_input_embeddings(self): |
| 1305 | return self.embeddings |
| 1306 | |
| 1307 | def set_input_embeddings(self, new_embeddings): |
| 1308 | self.embeddings = new_embeddings |
| 1309 | |
| 1310 | @add_start_docstrings_to_model_forward(NEMOTRONH_INPUTS_DOCSTRING) |
| 1311 | @add_code_sample_docstrings( |
| 1312 | checkpoint=_CHECKPOINT_FOR_DOC, |
| 1313 | output_type=NemotronHOutput, |
| 1314 | config_class=_CONFIG_FOR_DOC, |
| 1315 | ) |
| 1316 | def forward( |
| 1317 | self, |
| 1318 | input_ids: Optional[torch.LongTensor] = None, |
| 1319 | inputs_embeds: Optional[torch.LongTensor] = None, |
| 1320 | position_ids: Optional[torch.LongTensor] = None, |
| 1321 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None, |
| 1322 | use_cache: Optional[bool] = None, |
| 1323 | output_attentions: Optional[bool] = None, |
| 1324 | output_hidden_states: Optional[bool] = None, |
| 1325 | return_dict: Optional[bool] = None, |
| 1326 | cache_position: Optional[torch.LongTensor] = None, |
| 1327 | attention_mask: Optional[torch.Tensor] = None, |
| 1328 | **kwargs, |
| 1329 | ) -> Union[Tuple, NemotronHOutput]: |
| 1330 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 1331 | output_hidden_states = ( |
| 1332 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 1333 | ) |
| 1334 | # use_cache = use_cache if use_cache is not None else self.config.use_cache |
| 1335 | use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) |
| 1336 | |
| 1337 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 1338 | |
| 1339 | if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor |
| 1340 | raise ValueError("You must specify exactly one of input_ids or inputs_embeds") |
| 1341 | |
| 1342 | if inputs_embeds is None: |
| 1343 | inputs_embeds = self.embeddings(input_ids) |
| 1344 | |
| 1345 | if self.gradient_checkpointing and self.training and use_cache: |
| 1346 | logger.warning_once( |
| 1347 | "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." |
| 1348 | ) |
| 1349 | use_cache = False |
| 1350 | |
| 1351 | # From zamba_modeling.py |
| 1352 | if use_cache and cache_params is None: |
| 1353 | logger.warning_once( |
| 1354 | "NemotronH requires an initialized `NemotronHHybridDynamicCache` to return a cache. None was " |
| 1355 | "provided, so no cache will be returned." |
| 1356 | ) |
| 1357 | |
| 1358 | hidden_states = inputs_embeds |
| 1359 | |
| 1360 | if cache_position is None: |
| 1361 | cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device) |
| 1362 | if position_ids is None: |
| 1363 | position_ids = cache_position.unsqueeze(0) |
| 1364 | |
| 1365 | causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position) |
| 1366 | mamba_mask = self._update_mamba_mask(attention_mask, cache_position) |
| 1367 | |
| 1368 | all_hidden_states = () if output_hidden_states else None |
| 1369 | all_self_attns = () if output_attentions else None |
| 1370 | # Until HERE |
| 1371 | |
| 1372 | for layer_idx, mixer_block in enumerate(self.layers): |
| 1373 | # Depending on the layer type we opt for 2D base attention mask (Mamba) or 4D causal mask (Attention) |
| 1374 | if mixer_block.block_type == "mamba": |
| 1375 | layer_mask = mamba_mask |
| 1376 | elif mixer_block.block_type == "attention": |
| 1377 | layer_mask = causal_mask |
| 1378 | elif mixer_block.block_type == "mlp": |
| 1379 | layer_mask = None |
| 1380 | else: |
| 1381 | raise ValueError(f"Invalid block_type: {self.block_type}") |
| 1382 | |
| 1383 | if output_hidden_states: |
| 1384 | all_hidden_states += (hidden_states,) |
| 1385 | |
| 1386 | if self.gradient_checkpointing and self.training: |
| 1387 | hidden_states = self._gradient_checkpointing_func( |
| 1388 | mixer_block.__call__, hidden_states, cache_params, cache_position, layer_mask |
| 1389 | ) |
| 1390 | else: |
| 1391 | hidden_states = mixer_block( |
| 1392 | hidden_states, |
| 1393 | cache_params=cache_params, |
| 1394 | cache_position=cache_position, |
| 1395 | attention_mask=layer_mask, |
| 1396 | ) |
| 1397 | |
| 1398 | # TODO: Store attentions |
| 1399 | # if output_attentions: |
| 1400 | # if layer_outputs[1] is not None: |
| 1401 | # # append attentions only of attention layers. Mamba layers return `None` as the attention weights |
| 1402 | # all_self_attns += (layer_outputs[1],) |
| 1403 | |
| 1404 | # TODO (Check): should it happen before the forward pass? |
| 1405 | # if output_hidden_states: |
| 1406 | # all_hidden_states = all_hidden_states + (hidden_states,) |
| 1407 | |
| 1408 | hidden_states = self.norm_f(hidden_states) |
| 1409 | |
| 1410 | if output_hidden_states: |
| 1411 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 1412 | |
| 1413 | if not return_dict: |
| 1414 | return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None) |
| 1415 | |
| 1416 | return NemotronHOutput( |
| 1417 | last_hidden_state=hidden_states, |
| 1418 | cache_params=cache_params if use_cache else None, |
| 1419 | hidden_states=all_hidden_states, |
| 1420 | attentions=all_self_attns, |
| 1421 | ) |
| 1422 | |
| 1423 | # Copied from transformers.models.jamba.modeling_jamba.JambaModel._update_causal_mask |
| 1424 | def _update_causal_mask(self, attention_mask, input_tensor, cache_position): |
| 1425 | if self.config._attn_implementation == "flash_attention_2": |
| 1426 | if attention_mask is not None and 0.0 in attention_mask: |
| 1427 | return attention_mask |
| 1428 | return None |
| 1429 | |
| 1430 | dtype, device = input_tensor.dtype, input_tensor.device |
| 1431 | min_dtype = torch.finfo(dtype).min |
| 1432 | sequence_length = input_tensor.shape[1] |
| 1433 | target_length = cache_position[-1] + 1 |
| 1434 | |
| 1435 | causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) |
| 1436 | if sequence_length != 1: |
| 1437 | causal_mask = torch.triu(causal_mask, diagonal=1) |
| 1438 | causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) |
| 1439 | causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) |
| 1440 | if attention_mask is not None: |
| 1441 | causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit |
| 1442 | if attention_mask.dim() == 2: |
| 1443 | mask_length = attention_mask.shape[-1] |
| 1444 | padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0) |
| 1445 | causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype) |
| 1446 | |
| 1447 | if ( |
| 1448 | self.config._attn_implementation == "sdpa" |
| 1449 | and attention_mask is not None |
| 1450 | and attention_mask.device.type == "cuda" |
| 1451 | ): |
| 1452 | # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when |
| 1453 | # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. |
| 1454 | # Details: https://github.com/pytorch/pytorch/issues/110213 |
| 1455 | causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) |
| 1456 | |
| 1457 | return causal_mask |
| 1458 | |
| 1459 | def _update_mamba_mask(self, attention_mask, cache_position): |
| 1460 | """ |
| 1461 | No need for zeroing states when |
| 1462 | 1. Cached forward |
| 1463 | 2. Attending to all inputs |
| 1464 | """ |
| 1465 | mamba_mask = attention_mask |
| 1466 | if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)): |
| 1467 | mamba_mask = None |
| 1468 | return mamba_mask |
| 1469 | |
| 1470 | |
| 1471 | @add_start_docstrings( |
| 1472 | """ |
| 1473 | The NEMOTRONH Model transformer with a language modeling head on top (linear layer with weights not tied to the input |
| 1474 | embeddings). |
| 1475 | """, |
| 1476 | NEMOTRONH_START_DOCSTRING, |
| 1477 | ) |
| 1478 | class NemotronHForCausalLM(NemotronHPreTrainedModel, GenerationMixin): |
| 1479 | _tied_weights_keys = ["lm_head.weight"] |
| 1480 | |
| 1481 | def __init__(self, config): |
| 1482 | super().__init__(config) |
| 1483 | self.backbone = NemotronHModel(config) |
| 1484 | self.vocab_size = config.vocab_size |
| 1485 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| 1486 | |
| 1487 | # Initialize weights and apply final processing |
| 1488 | self.post_init() |
| 1489 | |
| 1490 | def get_input_embeddings(self): |
| 1491 | return self.backbone.get_input_embeddings() |
| 1492 | |
| 1493 | def set_input_embeddings(self, new_embeddings): |
| 1494 | return self.backbone.set_input_embeddings(new_embeddings) |
| 1495 | |
| 1496 | def get_output_embeddings(self): |
| 1497 | return self.lm_head |
| 1498 | |
| 1499 | def set_output_embeddings(self, new_embeddings): |
| 1500 | self.lm_head = new_embeddings |
| 1501 | |
| 1502 | def get_decoder(self): |
| 1503 | return self.model |
| 1504 | |
| 1505 | def set_decoder(self, decoder): |
| 1506 | self.model = decoder |
| 1507 | |
| 1508 | def prepare_inputs_for_generation( |
| 1509 | self, |
| 1510 | input_ids, |
| 1511 | past_key_values=None, |
| 1512 | attention_mask=None, |
| 1513 | inputs_embeds=None, |
| 1514 | cache_position=None, |
| 1515 | position_ids=None, |
| 1516 | use_cache=True, |
| 1517 | **kwargs, |
| 1518 | ): |
| 1519 | # Copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py |
| 1520 | # Overwitten -- uses `cache_params` as opposed to `past_key_values` |
| 1521 | empty_past_kv = past_key_values is None |
| 1522 | |
| 1523 | # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens |
| 1524 | # Exception 1: when passing input_embeds, input_ids may be missing entries |
| 1525 | # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here |
| 1526 | # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. |
| 1527 | # (we can't check exception 3 while compiling) |
| 1528 | if not empty_past_kv: |
| 1529 | if ( |
| 1530 | inputs_embeds is not None # Exception 1 |
| 1531 | or cache_position[-1] >= input_ids.shape[1] # Exception 3 |
| 1532 | ): |
| 1533 | input_ids = input_ids[:, -cache_position.shape[0] :] |
| 1534 | elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) |
| 1535 | input_ids = input_ids[:, cache_position] |
| 1536 | else: |
| 1537 | past_key_values = HybridMambaAttentionDynamicCache( |
| 1538 | self.config, input_ids.shape[0], self.dtype, device=self.device |
| 1539 | ) |
| 1540 | |
| 1541 | if attention_mask is not None and position_ids is None: |
| 1542 | # create position_ids on the fly for batch generation |
| 1543 | position_ids = attention_mask.long().cumsum(-1) - 1 |
| 1544 | position_ids.masked_fill_(attention_mask == 0, 1) |
| 1545 | if not empty_past_kv: |
| 1546 | position_ids = position_ids[:, -input_ids.shape[1] :] |
| 1547 | |
| 1548 | # if `inputs_embeds` are passed, we only want to use them in the 1st generation step |
| 1549 | if inputs_embeds is not None and empty_past_kv: |
| 1550 | model_inputs = {"inputs_embeds": inputs_embeds} |
| 1551 | else: |
| 1552 | model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases |
| 1553 | |
| 1554 | model_inputs.update( |
| 1555 | { |
| 1556 | "position_ids": position_ids, |
| 1557 | "past_key_values": past_key_values, |
| 1558 | "use_cache": use_cache, |
| 1559 | "attention_mask": attention_mask, |
| 1560 | "logits_to_keep": self.config.num_logits_to_keep, |
| 1561 | "cache_position": cache_position, |
| 1562 | } |
| 1563 | ) |
| 1564 | return model_inputs |
| 1565 | |
| 1566 | @add_start_docstrings_to_model_forward(NEMOTRONH_INPUTS_DOCSTRING) |
| 1567 | @add_code_sample_docstrings( |
| 1568 | checkpoint=_CHECKPOINT_FOR_DOC, |
| 1569 | output_type=NemotronHCausalLMOutput, |
| 1570 | config_class=_CONFIG_FOR_DOC, |
| 1571 | ) |
| 1572 | def forward( |
| 1573 | self, |
| 1574 | input_ids: Optional[torch.LongTensor] = None, |
| 1575 | inputs_embeds: Optional[torch.FloatTensor] = None, |
| 1576 | position_ids: Optional[torch.LongTensor] = None, |
| 1577 | cache_params: Optional[HybridMambaAttentionDynamicCache] = None, |
| 1578 | labels: Optional[torch.LongTensor] = None, |
| 1579 | output_attentions: Optional[bool] = None, |
| 1580 | output_hidden_states: Optional[bool] = None, |
| 1581 | return_dict: Optional[bool] = None, |
| 1582 | use_cache: Optional[bool] = None, |
| 1583 | cache_position: Optional[torch.Tensor] = None, |
| 1584 | attention_mask: Optional[torch.Tensor] = None, |
| 1585 | **kwargs, # for now we need this for generation |
| 1586 | ) -> Union[Tuple, NemotronHCausalLMOutput]: |
| 1587 | r""" |
| 1588 | labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| 1589 | Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set |
| 1590 | `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` |
| 1591 | are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` |
| 1592 | """ |
| 1593 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 1594 | |
| 1595 | output_hidden_states = ( |
| 1596 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 1597 | ) |
| 1598 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 1599 | |
| 1600 | nemotron_h_outputs = self.backbone( |
| 1601 | input_ids, |
| 1602 | cache_params=cache_params, |
| 1603 | inputs_embeds=inputs_embeds, |
| 1604 | output_attentions=output_attentions, |
| 1605 | output_hidden_states=output_hidden_states, |
| 1606 | return_dict=return_dict, |
| 1607 | use_cache=use_cache, |
| 1608 | cache_position=cache_position, |
| 1609 | attention_mask=attention_mask, |
| 1610 | ) |
| 1611 | hidden_states = nemotron_h_outputs[0] |
| 1612 | |
| 1613 | # TODO: Check zamba_modeling.py: https://github.com/huggingface/transformers/blob/d7188ba600e36d3fd191b12e19f1b3bb81a8404f/src/transformers/models/zamba/modeling_zamba.py#L1284C1-L1286C2 |
| 1614 | #logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() |
| 1615 | logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() |
| 1616 | |
| 1617 | loss = None |
| 1618 | if labels is not None: |
| 1619 | # move labels to correct device to enable model parallelism |
| 1620 | labels = labels.to(logits.device) |
| 1621 | # Shift so that tokens < n predict n |
| 1622 | shift_logits = logits[..., :-1, :].contiguous() |
| 1623 | shift_labels = labels[..., 1:].contiguous() |
| 1624 | # Flatten the tokens |
| 1625 | loss_fct = CrossEntropyLoss() |
| 1626 | loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
| 1627 | |
| 1628 | if not return_dict: |
| 1629 | output = (logits,) + nemotron_h_outputs[1:] |
| 1630 | return ((loss,) + output) if loss is not None else output |
| 1631 | |
| 1632 | return NemotronHCausalLMOutput( |
| 1633 | loss=loss, |
| 1634 | logits=logits, |
| 1635 | cache_params=nemotron_h_outputs.cache_params, |
| 1636 | hidden_states=nemotron_h_outputs.hidden_states, |
| 1637 | attentions=nemotron_h_outputs.attentions, |
| 1638 | ) |
| 1639 | |