From 228b5169b9780d0678ece1a21c2b93ce0099a41e Mon Sep 17 00:00:00 2001 From: OMEGA Agent Date: Wed, 1 Jul 2026 09:52:21 +0200 Subject: [PATCH] [OpenVINO] Add NemotronH (Nemotron-3 Mamba2) support Adds support for nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 (NemotronH) hybrid Mamba2 + Attention + MoE model export to OpenVINO IR. Changes: - _ov_ops.py: Add convert_recurrent_mamba2_cell() for Mamba2 SSM recurrence (arXiv:2405.21060) via ov::Loop with ModuleExtension - input_generators.py: Add NemotronHDummyPastKeyValuesGenerator for conv_state + recurrent_state (mamba layers) + KV cache (attn layers) - model_configs.py: Add NemotronHOpenVINOConfig registered for 'nemotron_h' - model_patcher.py: Add NemotronHModelPatcher with Mamba2RecurrentCell, nemotron_h_mamba_mixer_forward, patched_nemotron_h_moe_forward - utils.py: Add 'nemotron_h' to SSM_MODELS list - modeling_decoder.py: Add NemotronH cache init (conv_state shape uses mamba_num_heads * mamba_head_dim, not mamba_expand * hidden_size) - tests: Add nemotron_h to SUPPORTED_SSM_ARCHITECTURES The exported IR uses ov::Loop for the Mamba2 recurrence. The FuseMamba2Loop pass in OV (PR #36412) will fuse the loop into ov::op::internal::Mamba2 at inference time (after that PR merges). Resolves: omega issue #44 References: openvinotoolkit/openvino#36412, openvinotoolkit/openvino#36372 Based on: PR #1789 by openvino-agent AI Assistance: yes - OMEGA agent (REQUIRES HUMAN REVIEW) --- docs/source/openvino/models.mdx | 1 + optimum/exporters/openvino/_ov_ops.py | 104 +++++ .../exporters/openvino/input_generators.py | 71 ++++ optimum/exporters/openvino/model_configs.py | 78 ++++ optimum/exporters/openvino/model_patcher.py | 375 +++++++++++++++++- optimum/exporters/openvino/utils.py | 1 + optimum/intel/openvino/modeling_decoder.py | 37 +- tests/openvino/test_decoder.py | 3 +- tests/openvino/test_export.py | 1 + tests/openvino/utils_tests.py | 2 + 10 files changed, 669 insertions(+), 4 deletions(-) diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index 6220d3630f..386683db2b 100644 --- a/docs/source/openvino/models.mdx +++ b/docs/source/openvino/models.mdx @@ -115,6 +115,7 @@ Here is the list of the supported architectures : - MobileNet v1 - MobileNet v2 - MobileViT +- Nemotron-H (NVIDIA-Nemotron-3-Nano) - Nystromformer - OLMo - OLMo 2 diff --git a/optimum/exporters/openvino/_ov_ops.py b/optimum/exporters/openvino/_ov_ops.py index 78e5b6d23b..74e97fbf06 100644 --- a/optimum/exporters/openvino/_ov_ops.py +++ b/optimum/exporters/openvino/_ov_ops.py @@ -111,3 +111,107 @@ def convert_recurrent_attention_cell(context): final_output = ops.concat([core_attn_out_new, last_recurrent_state_new], 0) return [final_output.output(0)] + + + + +# Conversion rule for the `Mamba2RecurrentCellOp` operation in a Torch graph. +# +# This generalizes the recurrent-cell-to-`ov::Loop` approach (originally introduced for the +# GatedDeltaNet block, see `convert_recurrent_attention_cell` above) to the Mamba2 selective +# state-space recurrence used by hybrid Mamba2 models such as NemotronH. +# +# The Mamba2 single-step recurrence over the SSM state is a linear recurrence: +# state_t = state_{t-1} * dA_t + dBx_t # [B, H, P, N] +# y_t = reduce_sum(state_t * C_t, axis=N) # [B, H, P] +# where `dA` (discretized A), `dBx` (discretized B * x) and `C` are precomputed and vectorized +# over the sequence in the patched mixer forward. The skip connection `x_t * D` does not depend +# on the recurrent state and is therefore added outside the loop. +# +# The `Mamba2RecurrentCellOp` appears in the Torch graph as a result of replacing the +# `Mamba2RecurrentCell` `torch.nn.Module` via a registered `ModuleExtension` in the OpenVINO +# PyTorch frontend; OpenVINO then applies this conversion rule to the resulting operation. +def convert_recurrent_mamba2_cell(context): + # Inputs match the forward signature of `Mamba2RecurrentCell`. + # `dA` is broadcastable ([B, H, T, 1, 1]) while `dBx` carries the full state shape + # ([B, H, T, P, N]); shapes for the accumulator and the trip count are therefore + # derived from `dBx`. + dA = context.get_input(0) # [B, H, T, 1, 1] + dBx = context.get_input(1) # [B, H, T, P, N] + C = context.get_input(2) # [B, H, T, N] + last_state_old = context.get_input(3) # [B, H, P, N] + + const_zero_axis = ops.constant(0, dtype=np.int32) + const_two = ops.constant(2, dtype=np.int32) + const_minus_one = ops.constant(-1, dtype=np.int32) + + # Build the zero-initialized output accumulator with shape [B, H, T, P]. + dBx_shape = ops.shape_of(dBx) + core_shape = ops.gather(dBx_shape, ops.constant([0, 1, 2, 3], dtype=np.int32), const_zero_axis) + const_zero_f32 = ops.constant(0, dtype=np.float32) + core_out = ops.broadcast(const_zero_f32, core_shape) + + # Trip count for the loop equals the sequence length (dim 2). + seq_len = ops.gather(dBx_shape, const_two, const_zero_axis) + seq_len = ops.convert(seq_len, "i32") + + # Body parameters (one timestep slice each). + timestep_param = ops.parameter([], np.int32, "timestep") + dA_t_param = ops.parameter([-1, -1, 1, -1, -1], np.float32, "dA_t") + dBx_t_param = ops.parameter([-1, -1, 1, -1, -1], np.float32, "dBx_t") + C_t_param = ops.parameter([-1, -1, 1, -1], np.float32, "C_t") + last_state_t = ops.parameter([-1, -1, -1, -1], np.float32, "last_state_t") + core_out_t = ops.parameter([-1, -1, -1, -1], np.float32, "core_out_t") + + # Drop the singleton sequence dimension introduced by slicing. + dA_t = ops.squeeze(dA_t_param, const_two) # [B, H, 1, 1] (broadcastable to [B, H, P, N]) + dBx_t = ops.squeeze(dBx_t_param, const_two) # [B, H, P, N] + C_t = ops.squeeze(C_t_param, const_two) # [B, H, N] + + # state_t = state_{t-1} * dA_t + dBx_t + last_state_new = ops.add(ops.multiply(last_state_t, dA_t), dBx_t) # [B, H, P, N] + + # y_t = reduce_sum(state_t * C_t, axis=N) -> [B, H, P] + const_minus_two = ops.constant(-2, dtype=np.int32) + y_t = ops.multiply(last_state_new, ops.unsqueeze(C_t, const_minus_two)) # [B, H, P, N] + y_t = ops.reduce_sum(y_t, const_minus_one, False) # [B, H, P] + y_t = ops.unsqueeze(y_t, const_two) # [B, H, 1, P] + + timestep = ops.unsqueeze(timestep_param, const_zero_axis) + core_out_res = ops.scatter_update(core_out_t, timestep, y_t, const_two) + last_state_res = last_state_new + + body_cond = ops.constant([True], dtype=bool) + body_model = ov.Model( + [body_cond, last_state_res, core_out_res], + [ + timestep_param, + dA_t_param, + dBx_t_param, + C_t_param, + last_state_t, + core_out_t, + ], + "mamba2_body_model", + ) + + loop = ops.loop(seq_len, ops.constant(True, dtype="bool")) + loop.set_function(body_model) + + loop.set_sliced_input(dA_t_param, dA, 0, 1, 1, -1, 2) + loop.set_sliced_input(dBx_t_param, dBx, 0, 1, 1, -1, 2) + loop.set_sliced_input(C_t_param, C, 0, 1, 1, -1, 2) + loop.set_merged_input(last_state_t, last_state_old, last_state_res.output(0)) + loop.set_merged_input(core_out_t, core_out.output(0), core_out_res.output(0)) + loop.set_special_body_ports([0, 0]) + + core_out_new = loop.get_iter_value(core_out_res.output(0), -1) + last_state_new = loop.get_iter_value(last_state_res.output(0), -1) + + flatten_shape = ops.constant([-1], dtype=np.int32) + core_out_new = ops.reshape(core_out_new, flatten_shape, False) + last_state_new = ops.reshape(last_state_new, flatten_shape, False) + + final_output = ops.concat([core_out_new, last_state_new], 0) + + return [final_output.output(0)] \ No newline at end of file diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 65af52b286..7f38bc92bc 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -1703,6 +1703,77 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int return cache_params + +class NemotronHDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for NemotronH (hybrid Mamba2 + Attention + MoE) architectures. + + The cache is a flat list of tensors ordered as: + * for each mamba layer: conv_state, recurrent_state + * for each attention layer: key cache, value cache + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + + config = normalized_config.config + self.layers_block_type = config.layers_block_type + self.num_mamba_layers = config.layers_block_type.count("mamba") + self.num_attn_layers = config.layers_block_type.count("attention") + + self.conv_kernel_size = config.conv_kernel + self.ssm_state_size = config.ssm_state_size + self.n_groups = config.n_groups + self.mamba_num_heads = config.mamba_num_heads + self.mamba_head_dim = config.mamba_head_dim + self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim + self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + cache_params = [] + + for _ in range(self.num_mamba_layers): + conv_state_shape = (self.batch_size, self.conv_dim, self.conv_kernel_size) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(conv_state) + recurrent_state_shape = ( + self.batch_size, + self.mamba_num_heads, + self.mamba_head_dim, + self.ssm_state_size, + ) + recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(recurrent_state) + + for _ in range(self.num_attn_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + cache_params.append(k) + cache_params.append(v) + + return cache_params + + + class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): """ Generates dummy cache_params inputs for Qwen3.5 architectures. diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index f4a492c723..9d946250bf 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -75,6 +75,7 @@ LTXTransformerDummyInputGenerator, LTXVaeDummyInputGenerator, MambaCacheDummyInputGenerator, + NemotronHDummyPastKeyValuesGenerator, OVFalconDummyPastKeyValuesGenerator, OVMiniCPM3DummyPastKeyValuesGenerator, PooledProjectionsDummyInputGenerator, @@ -141,6 +142,7 @@ MixtralModelPatcher, ModelPatcher, MPTModelPatcher, + NemotronHModelPatcher, OVDecoderModelPatcher, OVSeq2SeqModelPatcher, PegasusModelPatcher, @@ -6101,3 +6103,79 @@ class TrOCROpenVINOConfig(TextSeq2SeqOpenVINOConfig): decoder_num_attention_heads="decoder_attention_heads", hidden_size="hidden_size", ) + + +@register_in_tasks_manager( + "nemotron_h", + *["text-generation", "text-generation-with-past"], + library_name="transformers", +) +class NemotronHOpenVINOConfig(TextDecoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "5.0.0" + + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + cache_name_prefix = "cache_params.past" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + cache_name_prefix = "cache_params.present" + + layers_block_type = self._normalized_config.layers_block_type + self.num_mamba_layers = layers_block_type.count("mamba") + self.num_attn_layers = layers_block_type.count("attention") + + for i in range(self.num_mamba_layers): + # [batch_size, conv_dim, conv_kernel_size] + inputs_or_outputs[f"{cache_name_prefix}.conv.{i}"] = {0: "batch_size"} + # [batch_size, mamba_num_heads, mamba_head_dim, ssm_state_size] + inputs_or_outputs[f"{cache_name_prefix}.ssm.{i}"] = {0: "batch_size"} + + for i in range(self.num_attn_layers): + inputs_or_outputs[f"{cache_name_prefix}.key.{i}"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{cache_name_prefix}.value.{i}"] = {0: "batch_size", 2: decoder_sequence_name} + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + return common_inputs + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + # need to override `generate_dummy_inputs` since the hybrid Mamba2 model exposes + # conv/recurrent states (for mamba layers) and key/value caches (for attention layers) + # bundled together under a single flat `cache_params` input. + dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) + + dummy_inputs = {} + input_names = [key for key in self.inputs.keys() if not key.startswith("cache_params")] + if self.use_past_in_inputs: + input_names.extend(["cache_params"]) + + for input_name in input_names: + input_was_inserted = False + for dummy_input_gen in dummy_inputs_generators: + if dummy_input_gen.supports_input(input_name): + dummy_inputs[input_name] = self.overwrite_shape_and_generate_input( + dummy_input_gen, + input_name, + framework, + input_shapes=kwargs, + ) + input_was_inserted = True + break + if not input_was_inserted: + raise RuntimeError( + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' + ) + + return dummy_inputs + diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index ed51e3d2f6..5bda6cb111 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -47,7 +47,7 @@ from transformers.processing_utils import Unpack from transformers.utils import ModelOutput -from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell +from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell, convert_recurrent_mamba2_cell from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.openvino.patching_utils import ( ModelPatcher, @@ -10317,3 +10317,376 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) self._model.forward = self._model._orig_forward + + +# This torch.nn.Module represents the Mamba2 selective-scan recurrence in its recurrent form. +# It is required for converting the Mamba2 mixer with OpenVINO using the ModuleExtension mechanism. +# +# The discretized quantities `dA` (decay), `dBx` (input contribution) and `C` (output projection) +# are precomputed and vectorized over the sequence dimension in the patched mixer forward +# (`nemotron_h_mamba_mixer_forward`). The recurrence over the SSM state is: +# state_t = state_{t-1} * dA_t + dBx_t +# y_t = reduce_sum(state_t * C_t, axis=N) +# This loop has no known vectorized form that can be correctly traced by torch.jit.trace, +# so it is replaced with an `ov::Loop` operation via `convert_recurrent_mamba2_cell`. +class Mamba2RecurrentCell(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward( + self, + dA, # (B, H, T, 1, 1) + dBx, # (B, H, T, P, N) + C, # (B, H, T, N) + last_state, # (B, H, P, N) + ): + sequence_length = dBx.shape[2] + core_out = torch.zeros(dBx.shape[:4], dtype=dBx.dtype) # (B, H, T, P) + + for i in range(sequence_length): + last_state = last_state * dA[:, :, i] + dBx[:, :, i] + core_out[:, :, i] = (last_state * C[:, :, i].unsqueeze(-2)).sum(dim=-1) + + # This is a workaround to ensure a single output from the torch.nn.Module. + # The OpenVINO ModuleExtension mechanism has a limitation and expects + # the module to produce only one output. + output_cell = torch.cat([core_out.flatten(), last_state.flatten()], dim=0) + return output_cell + + +# Recurrent form of the NemotronH Mamba2 mixer (`NemotronHMamba2Mixer`). +# Adapted from `torch_forward` of: +# https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/blob/main/modeling_nemotron_h.py +# +# Compared with the original implementation this patch: +# * replaces the `CausalConv1d` with the generic, cache-aware `ov_causal_conv1d()` helper; +# * expresses the selective scan in recurrent form (precompute dA/dBx/C, then run +# `Mamba2RecurrentCell`), which is replaced by an `ov::Loop` during conversion. +# It runs a single unified code path for both the prefill and the decoding stages. +def nemotron_h_mamba_mixer_forward( + self, + hidden_states: torch.Tensor, + cache_params=None, + conv_state: Optional[torch.Tensor] = None, + recurrent_state: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, +): + batch_size, seq_len, _ = hidden_states.shape + + def apply_mask_to_padding_states(hidden_states, attention_mask): + # Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66. + # The mask is sliced to the current `seq_len` so that it broadcasts correctly both during + # the prefill stage (seq_len > 1) and the decoding stage (seq_len == 1, the slice selects the + # current — always valid — token), keeping a single traceable code path. + if attention_mask is not None and attention_mask.shape[0] > 1: + dtype = hidden_states.dtype + attention_mask = attention_mask[:, -seq_len:] + hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) + return hidden_states + + # 1. Gated MLP's linear projection + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + projected_states = self.in_proj(hidden_states) + gate, hidden_states_B_C, dt = projected_states.split( + [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 + ) + + # 2. Convolution sequence transformation with cached state + hidden_states_B_C = hidden_states_B_C.transpose(1, 2) + if conv_state is None: + conv_state = torch.zeros(batch_size, self.conv_dim, self.conv_kernel_size, dtype=hidden_states_B_C.dtype) + new_hidden_states_B_C, new_conv_state = ov_causal_conv1d( + conv_state, hidden_states_B_C, self.conv1d.weight, self.conv1d.bias + ) + hidden_states_B_C = self.act(new_hidden_states_B_C).transpose(1, 2) + hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask) + + hidden_states_ssm, B, C = torch.split( + hidden_states_B_C, + [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], + dim=-1, + ) + + # 3. State Space Model transformation in recurrent form + A = -torch.exp(self.A_log.float()) # (num_heads,) + dt = torch.nn.functional.softplus(dt + self.dt_bias) # (B, seq, num_heads) + dt = torch.clamp(dt, self.time_step_min) + + x = hidden_states_ssm.reshape(batch_size, seq_len, self.num_heads, self.head_dim).float() + B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size).float() + C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size).float() + B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2) + C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2) + + # rearrange to head-major layout: (B, H, T, ...) + dt = dt.transpose(1, 2).float() # (B, H, T) + x = x.permute(0, 2, 1, 3) # (B, H, T, P) + B = B.permute(0, 2, 1, 3) # (B, H, T, N) + C = C.permute(0, 2, 1, 3) # (B, H, T, N) + + # discretize: dA (decay), dBx (input contribution) + dA = torch.exp(dt * A.view(1, -1, 1)) # (B, H, T) + dA = dA[..., None, None] # (B, H, T, 1, 1) + dBx = dt[..., None, None] * B[..., None, :] * x[..., :, None] # (B, H, T, P, N) + + if recurrent_state is None: + recurrent_state = torch.zeros( + batch_size, self.num_heads, self.head_dim, self.ssm_state_size, dtype=torch.float32 + ) + recurrent_state = recurrent_state.float() + + output_cell = self.mamba2_recurrent_cell(dA, dBx, C, recurrent_state) + + num_elems = batch_size * self.num_heads * seq_len * self.head_dim + y = output_cell[:num_elems].reshape(batch_size, self.num_heads, seq_len, self.head_dim) + new_recurrent_state = output_cell[num_elems:].reshape(recurrent_state.shape) + + # D skip connection (independent of the recurrent state) + y = y + x * self.D.view(1, -1, 1, 1) + + # (B, H, T, P) -> (B, T, intermediate_size) + y = y.permute(0, 2, 1, 3).reshape(batch_size, seq_len, -1) + + scan_output = self.norm(y, gate) + + # 4. Final linear projection + contextualized_states = self.out_proj(scan_output.to(hidden_states.dtype)) + + return contextualized_states, new_conv_state, new_recurrent_state + + +# Vectorized implementation of the NemotronH MoE block (`NemotronHMoE`). +# The vectorized form is required to ensure correct torch.jit tracing. +# The token routing (sigmoid scoring + grouped top-k) is kept identical to the +# original `route_tokens_to_experts`; only the per-expert MLP loop is replaced +# with batched matrix multiplications over the 3D expert weight tensors. +def patched_nemotron_h_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residuals = hidden_states + orig_shape = hidden_states.shape + num_experts = self.n_routed_experts + + router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) + topk_weights = topk_weights.to(hidden_states.dtype) + + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + num_tokens = hidden_states.shape[0] + + # scatter the top-k weights into a dense (num_tokens, num_experts) tensor + routing_weights = torch.zeros(num_tokens, num_experts, dtype=topk_weights.dtype) + routing_weights.scatter_(1, topk_indices, topk_weights) + + # NemotronH-specific latent projection (Identity when moe_latent_size is None) + h = self.fc1_latent_proj(hidden_states) + + # compute all experts in a vectorized form: experts are non-gated MLPs + h = h.repeat(num_experts, 1).view(num_experts, num_tokens, -1) + up = torch.bmm(h, self.up_projs.transpose(1, 2)) + up = self.experts.act_fn(up) + next_states = torch.bmm(up, self.down_projs.transpose(1, 2)) # (num_experts, num_tokens, in_dim) + next_states = next_states * routing_weights.transpose(0, 1)[..., None] + next_states = next_states.sum(dim=0) + + next_states = self.fc2_latent_proj(next_states) + next_states = next_states.view(*orig_shape) + + output = next_states + self.shared_experts(residuals) + return output + + +class NemotronHModelPatcher(OVDecoderModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: "PreTrainedModel", + model_kwargs: Optional[Dict[str, Any]] = None, + ): + from openvino.frontend.pytorch import ConversionExtension, ModuleExtension + + super().__init__(config, model, model_kwargs) + + layers_block_type = self.real_config._config.layers_block_type + + # Lightweight cache wrapper. Unlike the upstream hybrid cache, it stores the + # convolution/recurrent states for mamba layers and the key/value caches for + # attention layers in plain lists, mapping the global layer index to the + # corresponding per-type cache index. The patched mixer reads conv/recurrent + # states directly, while the (unmodified) attention block uses `update()`. + class NemotronHCacheWrap: + def __init__(self, config, conv_states, recurrent_states, key_cache, value_cache, attention_mask=None): + self.config = config + self.conv_states = conv_states + self.recurrent_states = recurrent_states + self.key_cache = key_cache + self.value_cache = value_cache + # The raw 2D attention mask is stored here so that the patched mamba mixer can + # always zero out padding positions. We cannot rely on the model's internal + # `_update_mamba_mask`, since it returns `None` whenever the mask is all ones — + # which is the case for the all-ones dummy input used during tracing, baking + # "no padding masking" into the exported graph. + self.raw_attention_mask = attention_mask + self.mamba_mapping = {} + self.attn_mapping = {} + mamba_idx = 0 + attn_idx = 0 + for i, block_type in enumerate(config.layers_block_type): + if block_type == "mamba": + self.mamba_mapping[i] = mamba_idx + mamba_idx += 1 + elif block_type == "attention": + self.attn_mapping[i] = attn_idx + attn_idx += 1 + self.num_attn_layers = attn_idx + + def update(self, key_states, value_states, layer_idx, cache_kwargs=None): + idx = self.attn_mapping[layer_idx] + if self.key_cache[idx] is None or self.key_cache[idx].shape[-2] == 0: + self.key_cache[idx] = key_states + self.value_cache[idx] = value_states + else: + self.key_cache[idx] = torch.cat([self.key_cache[idx], key_states], dim=2) + self.value_cache[idx] = torch.cat([self.value_cache[idx], value_states], dim=2) + return self.key_cache[idx], self.value_cache[idx] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + if self.num_attn_layers == 0 or self.key_cache[0] is None: + return 0 + return self.key_cache[0].shape[-2] + + def get_mask_sizes(self, query_length: int, layer_idx: int): + kv_length = self.get_seq_length() + query_length + return kv_length, 0 + + @property + def is_sliding(self): + return [False] + + is_compileable = False + + def has_previous_state(self, layer_idx: Optional[int] = None) -> bool: + # The recurrent forward always runs with valid conv/recurrent states. + return True + + # The patch is needed to include conv, recurrent, key and value states in the + # model inputs and outputs as a flat list of tensors (`cache_params`). + def patched_forward( + input_ids, + attention_mask=None, + cache_params=None, + ): + num_mamba_layers = layers_block_type.count("mamba") + num_attn_layers = layers_block_type.count("attention") + + use_cache = False + wrapped_cache_params = None + if cache_params is not None: + use_cache = True + conv_states = [] + recurrent_states = [] + key_cache = [] + value_cache = [] + + for idx in range(num_mamba_layers): + conv_states.append(cache_params[2 * idx]) + recurrent_states.append(cache_params[2 * idx + 1]) + + for idx in range(num_attn_layers): + key_cache.append(cache_params[2 * num_mamba_layers + 2 * idx]) + value_cache.append(cache_params[2 * num_mamba_layers + 2 * idx + 1]) + + wrapped_cache_params = NemotronHCacheWrap( + self.real_config._config, conv_states, recurrent_states, key_cache, value_cache, attention_mask + ) + + causal_lm_output = self.model_orig_forward( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=wrapped_cache_params, + use_cache=use_cache, + ) + outputs = { + "logits": causal_lm_output.logits, + } + + if use_cache: + present = wrapped_cache_params + present_key_values = [] + for idx in range(num_mamba_layers): + present_key_values.append(present.conv_states[idx]) + present_key_values.append(present.recurrent_states[idx]) + + for idx in range(num_attn_layers): + present_key_values.append(present.key_cache[idx]) + present_key_values.append(present.value_cache[idx]) + + outputs["present_key_values"] = present_key_values + + return outputs + + self.patched_forward = patched_forward + self.model_orig_forward = self.orig_forward + self.orig_forward = patched_forward + + self.module_extensions = { + Mamba2RecurrentCell: ModuleExtension(Mamba2RecurrentCell, "Mamba2RecurrentCellOp"), + } + self.conversion_extensions = [ + ConversionExtension("Mamba2RecurrentCellOp", convert_recurrent_mamba2_cell), + ] + + def __enter__(self): + from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHMamba2Mixer, NemotronHMoE + + super().__enter__() + setattr(self._model, self.orig_forward_name, self.patched_forward) + + # Patch the mamba mixer to read/write conv and recurrent states from the wrapped cache. + def make_mamba_forward(mixer): + def _forward(hidden_states, cache_params=None, attention_mask=None, **kwargs): + conv_state = None + recurrent_state = None + raw_attention_mask = None + if cache_params is not None: + mamba_idx = cache_params.mamba_mapping[mixer.layer_idx] + conv_state = cache_params.conv_states[mamba_idx] + recurrent_state = cache_params.recurrent_states[mamba_idx] + # Use the raw 2D mask rather than the (possibly `None`) mamba mask the model + # computes internally, to keep padding handling correct in the traced graph. + raw_attention_mask = cache_params.raw_attention_mask + out, new_conv_state, new_recurrent_state = nemotron_h_mamba_mixer_forward( + mixer, hidden_states, cache_params, conv_state, recurrent_state, raw_attention_mask + ) + if cache_params is not None: + mamba_idx = cache_params.mamba_mapping[mixer.layer_idx] + cache_params.conv_states[mamba_idx] = new_conv_state + cache_params.recurrent_states[mamba_idx] = new_recurrent_state + return out + + return _forward + + for decoder_layer in self._model.model.layers: + mixer = decoder_layer.mixer + if isinstance(mixer, NemotronHMamba2Mixer): + mixer._orig_forward = mixer.forward + mixer.mamba2_recurrent_cell = Mamba2RecurrentCell() + mixer.forward = make_mamba_forward(mixer) + elif isinstance(mixer, NemotronHMoE): + mixer._orig_forward = mixer.forward + mixer.forward = types.MethodType(patched_nemotron_h_moe_forward, mixer) + # TODO: remove `float()` casting when CVS-181449 is fixed; + # it is currently needed for MoE optimizations to be applied. + mixer.up_projs = mixer.experts.up_proj.float() + mixer.down_projs = mixer.experts.down_proj.float() + + def __exit__(self, exc_type, exc_value, traceback): + from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHMamba2Mixer, NemotronHMoE + + super().__exit__(exc_type, exc_value, traceback) + setattr(self._model, self.orig_forward_name, self.model_orig_forward) + for decoder_layer in self._model.model.layers: + mixer = decoder_layer.mixer + if isinstance(mixer, (NemotronHMamba2Mixer, NemotronHMoE)): + mixer.forward = mixer._orig_forward + if isinstance(mixer, NemotronHMoE): + del mixer.up_projs, mixer.down_projs + diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index de5fd4b355..9d27c66b54 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -349,6 +349,7 @@ def _get_kokoro_submodels(model): "qwen3_next", "qwen3_5_text", "qwen3_5_moe_text", + "nemotron_h", ] # All transformers, diffusers, timm and sentence transformers models that were supported via optimum-onnx OnnxConfigs for which support is now removed diff --git a/optimum/intel/openvino/modeling_decoder.py b/optimum/intel/openvino/modeling_decoder.py index eadbbf130f..bcad926cc8 100644 --- a/optimum/intel/openvino/modeling_decoder.py +++ b/optimum/intel/openvino/modeling_decoder.py @@ -1162,6 +1162,20 @@ def __init__( self.mamba_headdim = getattr(config, "mamba_d_head", None) self.num_mamba_layers = layer_types.count("mamba") self.num_attn_layers = layer_types.count("attention") + elif config.model_type == "nemotron_h": + # NemotronH is a hybrid Mamba2 + Attention + MoE model. `config.intermediate_size` + # refers to the (MoE) MLP size, while the mamba intermediate size is derived from + # the mamba head configuration. + layers_block_type = getattr(config, "layers_block_type", None) + self.num_key_value_heads = getattr(config, "num_key_value_heads", None) + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_attention_heads) + self.mamba_ngroups = getattr(config, "n_groups", None) + self.n_mamba_heads = getattr(config, "mamba_num_heads", None) + self.ssm_state_size = getattr(config, "ssm_state_size", None) + self.mamba_headdim = getattr(config, "mamba_head_dim", None) + self.mamba_intermediate_size = self.n_mamba_heads * self.mamba_headdim + self.num_mamba_layers = layers_block_type.count("mamba") + self.num_attn_layers = layers_block_type.count("attention") else: # Mamba 2 specific parameters hybrid_layer_ids = getattr(config, "hybrid_layer_ids", None) @@ -1182,7 +1196,14 @@ def __init__( self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.conv_states = conv_states - if self.conv_states is None: + if self.conv_states is None and config.model_type == "nemotron_h": + self.conv_states = [] + conv_dim = self.mamba_intermediate_size + 2 * self.mamba_ngroups * self.ssm_state_size + for _ in range(self.num_mamba_layers): + conv_state_shape = (self.max_batch_size, conv_dim, self.conv_kernel_size) + conv_state: torch.Tensor = torch.zeros(conv_state_shape, device=self.device, dtype=dtype) + self.conv_states.append(conv_state) + elif self.conv_states is None: self.conv_states = [] for _ in range(self.num_mamba_layers): if ( @@ -1206,7 +1227,18 @@ def __init__( self.conv_states.append(conv_state) self.ssm_states = ssm_states - if self.ssm_states is None: + if self.ssm_states is None and config.model_type == "nemotron_h": + self.ssm_states = [] + for _ in range(self.num_mamba_layers): + ssm_state_shape = ( + self.max_batch_size, + self.n_mamba_heads, + self.mamba_headdim, + self.ssm_state_size, + ) + ssm_state: torch.Tensor = torch.zeros(ssm_state_shape, device=self.device, dtype=dtype) + self.ssm_states.append(ssm_state) + elif self.ssm_states is None: self.ssm_states: List[torch.Tensor] = [] for _ in range(self.num_mamba_layers): if self.n_mamba_heads and self.mamba_headdim: @@ -1518,6 +1550,7 @@ def prepare_inputs_for_generation( "qwen3_next", "qwen3_5_text", "qwen3_5_moe_text", + "nemotron_h", ]: # LFM2, GraniteMoeHybrid (Granite-4.0), and Qwen3-Next require the attention mask # to be the length of the full context, so default mask from OVModelForCausalLM needs to be used. diff --git a/tests/openvino/test_decoder.py b/tests/openvino/test_decoder.py index 079c78c875..4627266df0 100644 --- a/tests/openvino/test_decoder.py +++ b/tests/openvino/test_decoder.py @@ -104,7 +104,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase): SUPPORTED_SSM_ARCHITECTURES += ("qwen3_next",) if is_transformers_version(">=", "5.0"): - SUPPORTED_SSM_ARCHITECTURES += ("lfm2_moe",) + SUPPORTED_SSM_ARCHITECTURES += ("lfm2_moe", "nemotron_h") SUPPORTED_ARCHITECTURES += SUPPORTED_SSM_ARCHITECTURES @@ -265,6 +265,7 @@ class OVModelForCausalLMIntegrationTest(unittest.TestCase): "bitnet": 6, "hunyuan_v1_dense": 2, "qwen3_next": 1, + "nemotron_h": 2, } TASK = "text-generation" diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index 77a84a3b99..404a438dc0 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -140,6 +140,7 @@ class ExportModelTest(unittest.TestCase): if is_transformers_version(">=", "5.0"): SUPPORTED_ARCHITECTURES.update({"lfm2_moe": OVModelForCausalLM}) + SUPPORTED_ARCHITECTURES.update({"nemotron_h": OVModelForCausalLM}) EXPECTED_DIFFUSERS_SCALE_FACTORS = { "stable-diffusion-xl": {"vae_encoder": "128.0", "vae_decoder": "128.0"}, diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index c67cd383db..620c8465b6 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -269,6 +269,7 @@ def _create_tiny_kokoro_model(): "mobilenet_v2": "optimum-intel-internal-testing/tiny-random-MobileNetV2Model", "mobilevit": "optimum-intel-internal-testing/tiny-random-mobilevit", "mpt": "optimum-intel-internal-testing/tiny-random-MptForCausalLM", + "nemotron_h": "optimum-intel-internal-testing/tiny-random-nemotron-h", "mpnet": "optimum-intel-internal-testing/tiny-random-MPNetModel", "mt5": "optimum-intel-internal-testing/mt5-tiny-random", "llava-qwen2": "optimum-intel-internal-testing/tiny-random-nanollava", @@ -568,6 +569,7 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "qwen3_eagle3": {"model": 20}, "qwen3_vl_eagle3": {"model": 18}, "qwen3_next": {"model": 100}, + "nemotron_h": {"model": 46}, "gemma4": { "lm_model": 54, "text_embeddings_model": 1,