monkeypatch.models.nemotron_h.modeling

monkeypatch.models.nemotron_h.modeling

Sample-packing and context-parallelism patch for NemotronH (Mamba2/Attention/MoE hybrid).

Threads seq_idx (derived from position_ids) into the Mamba2 SSM kernels so packed-sequence boundaries reset SSM state. Upstream never passes one, which leaks hidden state across boundaries. Attention and MoE blocks need no changes — transformers builds block-diagonal masks from position_ids for attention.

CP correction (ring-shift of SSM state + additive output fix) is handled by wrap_mamba_scan_for_cp from mamba_utils, which wraps the chunk-scan call at the module level.

Functions

Name	Description
guard_nemotron_h_fused_scan	Skip the fused Mamba2 training kernel where its assumptions break.
patch_nemotron_h_modeling_packing	Patch NemotronH for sample packing: seq_idx threading into Mamba2 SSM kernels.

guard_nemotron_h_fused_scan

monkeypatch.models.nemotron_h.modeling.guard_nemotron_h_fused_scan(mod=None)

Skip the fused Mamba2 training kernel where its assumptions break.

mamba_split_conv1d_scan_combined multiplies by out_proj.weight itself, which fails on a bitsandbytes-quantized weight, and it never surfaces the final SSM state that the CP correction needs. Returning None is upstream’s own “kernel unavailable” signal, so the mixer falls back to the separate conv / chunk-scan calls.

patch_nemotron_h_modeling_packing

monkeypatch.models.nemotron_h.modeling.patch_nemotron_h_modeling_packing(
    kernels_enabled=False,
)

Patch NemotronH for sample packing: seq_idx threading into Mamba2 SSM kernels.

_get_unpad_data is handled by SUPPORTED_MULTIPACK_MODEL_TYPES / patch_for_multipack(). This function only applies the seq_idx patches that are unique to nemotron_h.

kernels_enabled mirrors use_kernels: transformers then kernelizes the mixer with the Hub Mamba2 kernels, which take seq_idx like the pip ones.