Skip to content

Commit 482b3fd

Browse files
committed
Enable local mamba kernel pre-population for NemotronH models
NemotronH has Mamba layers just like GraniteMoeHybrid and needs the same _use_local_mamba_kernels() call to avoid causal_conv1d_cuda import failures in torchrun subprocesses.
1 parent 3556af4 commit 482b3fd

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

src/instructlab/training/model.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,13 @@ def __init__(
7575

7676
# check model type & set on the mclasss
7777
self.is_granitemoehybrid = is_known_model(model_path, "granitemoehybrid")
78+
self.is_nemotronh = is_known_model(model_path, "nemotronh")
7879
self.is_gpt_oss = is_gpt_oss(model_path)
7980

8081
# Pre-populate the Hub kernel cache with locally installed mamba_ssm
8182
# and causal_conv1d to avoid PyTorch/CUDA ABI mismatches with the
8283
# Hub-provided kernel builds.
83-
if self.is_granitemoehybrid:
84+
if self.is_granitemoehybrid or self.is_nemotronh:
8485
self._use_local_mamba_kernels()
8586

8687
if self.is_gpt_oss:

0 commit comments

Comments
 (0)