[https://nvbugs/6437410][fix] fix nemotron weight update test

jenkins/L0_Test.groovy

@@ -4431,6 +4431,12 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
             trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install opencv-python-headless")
             if (stageName.contains("-Ray-")) {
                 trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install ray[default]==2.55.1")
+                trtllm_utils.llmExecStepWithRetry(pipeline, script: """
+                    mambaArch=\$(uname -m)
+                    pip3 install --no-deps \
+                        "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.2/causal_conv1d-1.6.1%2Bcu13torch26.04cxx11abiTRUE-cp312-cp312-linux_\${mambaArch}.whl" \
+                        "https://github.com/state-spaces/mamba/releases/download/v2.3.0/mamba_ssm-2.3.0%2Bcu13torch26.01cxx11abiTRUE-cp312-cp312-linux_\${mambaArch}.whl"
+                """)
             }
             if (!skipInstallWheel) {
                 trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && pip3 install --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl")

jenkins/scripts/slurm_install.sh

@@ -59,6 +59,10 @@ slurm_install_setup() {
         nvidia-smi && nvidia-smi -q && nvidia-smi topo -m
         if [[ $pytestCommand == *--run-ray* ]]; then
             retry_command --timeout 2700 pip3 install --retries 10 "ray[default]==2.55.1"
+            mambaArch=$(uname -m)
+            retry_command --timeout 2700 pip3 install --retries 10 --no-deps \
+                "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.2/causal_conv1d-1.6.1%2Bcu13torch26.04cxx11abiTRUE-cp312-cp312-linux_${mambaArch}.whl" \
+                "https://github.com/state-spaces/mamba/releases/download/v2.3.0/mamba_ssm-2.3.0%2Bcu13torch26.01cxx11abiTRUE-cp312-cp312-linux_${mambaArch}.whl"
         fi
         retry_command --timeout 2700 bash -c "pip3 install --retries 10 opencv-python-headless"
         retry_command --timeout 2700 bash -c "cd $llmSrcNode && pip3 install --retries 10 -r requirements-dev.txt"

tests/integration/test_lists/waives.txt

@@ -324,7 +324,6 @@ unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestN
 unittest/_torch/modules/tests_lora_modules/test_qwen3_sanity.py::TestQwen3LoRA::test_qwen3_fp8_lora SKIP (https://nvbugs/6607487)
 unittest/_torch/multi_gpu/test_linear.py::test_row_linear[2-balanced] SKIP (https://nvbugs/6507113)
 unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404)
-unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part4" SKIP (https://nvbugs/6437410)
 unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" SKIP (https://nvbugs/6619882)
 unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819)
 unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TorchSampler-no_cuda_graph_and_overlap-stop_token_ids0-1-1-True-True-True] SKIP (https://nvbugs/6581048)

tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py

@@ -3,15 +3,10 @@
 
 import base64
 import gc
-import importlib.util
 import json
-import multiprocessing
 import os
 import pickle
 import re
-import subprocess
-import sys
-import traceback
 from typing import Callable, List, Optional, Tuple
 
 import pytest
@@ -1127,75 +1122,37 @@ def filter_fn(name: str) -> bool:
         compare_logits(llm_logits, ref_logits, threshold=0.8)
 
 
-@pytest.fixture
-def mamba_deps():
-    """Install mamba-ssm and causal-conv1d for the duration of the test, then
-    restore the full pip environment. Uses a pip-freeze diff so transitive
-    dependencies (e.g. quack-kernels pinning nvidia-cutlass-dsl==4.6.0.dev0,
-    which breaks tensorrt-llm's pin of 4.5.0) are also reverted."""
-
-    def _freeze():
-        out = subprocess.check_output(
-            [sys.executable, "-m", "pip", "freeze", "--disable-pip-version-check"],
-            text=True,
-        )
-        result = {}
-        for line in out.splitlines():
-            line = line.strip()
-            if not line or line.startswith("#") or " @ " in line:
-                continue
-            if "==" in line:
-                name, ver = line.split("==", 1)
-                result[name.lower()] = ver
-        return result
-
-    pkgs = ["mamba-ssm", "causal-conv1d"]
-    mod_names = {"mamba-ssm": "mamba_ssm", "causal-conv1d": "causal_conv1d"}
-    need_install = [p for p in pkgs if importlib.util.find_spec(mod_names[p]) is None]
-
-    before = _freeze() if need_install else None
+@pytest.mark.part4
+@skip_pre_hopper
+def test_llm_update_weights_nemotron_h():
+    """Weight update on Nemotron-H, a hybrid model mixing mamba, MoE and
+    attention layers.
+
+    Requires mamba-ssm and causal-conv1d to be importable: without them HF
+    falls back to the naive Python selective_scan path, which OOMs on
+    Nemotron-H and produces unmatched logits. The Ray CI stage installs both
+    next to ray -- see jenkins/scripts/slurm_install.sh."""
     try:
-        if need_install:
-            # --no-deps: avoid pulling in optional kernel deps (quack-kernels,
-            # tilelang) that upgrade nvidia-cutlass-dsl and break tensorrt-llm.
-            # The container already provides torch/einops/etc.
-            subprocess.check_call(
-                [
-                    sys.executable,
-                    "-m",
-                    "pip",
-                    "install",
-                    "--no-build-isolation",
-                    "--no-deps",
-                    *need_install,
-                ]
-            )
-            importlib.invalidate_caches()
-        yield
-    finally:
-        if before is None:
-            return
-        after = _freeze()
-        new_pkgs = [p for p in after if p not in before]
-        changed = [(p, before[p]) for p in after if p in before and after[p] != before[p]]
-        if new_pkgs:
-            subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", *new_pkgs])
-        if changed:
-            subprocess.check_call(
-                [sys.executable, "-m", "pip", "install", *[f"{p}=={v}" for p, v in changed]]
-            )
-
-
-def _nemotron_h_body():
-    """Body of test_llm_update_weights_nemotron_h. Executed in a fresh
-    subprocess via spawn so HF transformers re-imports cleanly and the
-    mamba-ssm / causal-conv1d fast path (installed by the mamba_deps
-    fixture) is picked up. Running this in-process would let the parent
-    pytest's already-resolved negative caches force the naive Python
-    selective_scan path, which OOMs on Nemotron-H and produces unmatched logits."""
+        import causal_conv1d  # noqa: F401
+        import mamba_ssm  # noqa: F401
+    except ImportError as e:
+        # Fail loudly here rather than let the naive fallback OOM further in,
+        # which is a much harder failure to read.
+        pytest.fail(
+            f"{e.name} is not installed, so the mamba fast path is unavailable. "
+            "The Ray CI stage installs mamba-ssm and causal-conv1d alongside ray; "
+            "see jenkins/scripts/slurm_install.sh."
+        )
     model_dir = str(llm_models_root() / "NVIDIA-Nemotron-3-Nano-30B-A3B-BF16")
     num_hidden_layers = 7
-    hf_model = RefHFModelWithIPCHandles(model_dir, num_hidden_layers=num_hidden_layers)
+    # NemotronHConfig derives num_hidden_layers from ``layers_block_type``
+    # and silently ignores direct assignment, so truncation must go through
+    # the layer-type list. The first 7 entries of the checkpoint's pattern
+    # ("MEMEM*E") keep all three layer types: mamba, MoE and attention.
+    layers_block_type = AutoConfig.from_pretrained(model_dir).layers_block_type[:num_hidden_layers]
+    hf_model = RefHFModelWithIPCHandles(
+        model_dir, num_hidden_layers=num_hidden_layers, layers_block_type=layers_block_type
+    )
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
     # Nemotron-H's Mamba state dominates the cache budget; 0.25 of free memory
     # leaves enough room for HF (resident on cuda:0 + replicas on cuda:1..3)
@@ -1216,7 +1173,7 @@ def _nemotron_h_body():
         kv_cache_config=kv_cache_config,
         moe_config=moe_config,
         max_batch_size=4,
-        model_kwargs={"num_hidden_layers": num_hidden_layers},
+        model_kwargs={"layers_block_type": layers_block_type},
     ) as llm:
         prompts_texts = [
             "Hello, my name is",
@@ -1261,30 +1218,13 @@ def filter_fn(name: str) -> bool:
         llm._collective_rpc("update_weights", (None,))
 
         llm_logits, ref_logits = run_generate(llm, hf_model, prompts, sampling_params)
-        compare_logits(llm_logits, ref_logits)
-
-
-def _nemotron_h_subprocess_entry(result_queue):
-    try:
-        _nemotron_h_body()
-        result_queue.put(None)
-    except BaseException:
-        result_queue.put(traceback.format_exc())
-
-
-@pytest.mark.part4
-@skip_pre_hopper
-def test_llm_update_weights_nemotron_h(mamba_deps):
-    """Runs _nemotron_h_body in a spawned subprocess so HF transformers
-    sees the mamba-ssm / causal-conv1d fast path installed by the
-    mamba_deps fixture. See _nemotron_h_body docstring for why."""
-    ctx = multiprocessing.get_context("spawn")
-    queue = ctx.Queue()
-    proc = ctx.Process(target=_nemotron_h_subprocess_entry, args=(queue,))
-    proc.start()
-    proc.join()
-    err = queue.get() if not queue.empty() else None
-    if proc.exitcode != 0:
-        pytest.fail(f"Subprocess exited with code {proc.exitcode}\n{err or ''}")
-    if err is not None:
-        pytest.fail(err)
+        # Looser threshold: Nemotron-H logits are compared against a BF16
+        # reference and the mamba SSM / selective-scan path introduces small
+        # numerical differences. Measured over 5 runs x 4 prompts (GB200, TP=4,
+        # BF16): mean top-20 overlap 0.891, overall range 0.867-0.928. The
+        # spread is dominated by which prompt it is, not by run-to-run noise --
+        # the weakest prompt stays in 0.867-0.875 across all 5 runs, so 0.85
+        # clears the worst observation by ~0.017 while 0.88 would already flake.
+        compare_logits(llm_logits, ref_logits, threshold=0.85)
+
+    del hf_model

tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py

@@ -7,7 +7,7 @@
 import pytest
 import torch
 from torch.multiprocessing.reductions import reduce_tensor
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from utils.llm_data import llm_models_root
 from utils.torch_ref import RefHFModel
 from utils.util import getSMVersion, skip_pre_hopper
@@ -41,13 +41,36 @@ def release_shared_cuda_memory():
 
 
 class RefHFModelWithIPCHandles(RefHFModel):
-    def __init__(self, model_dir: str, device_id: int = 0, num_hidden_layers: int = 4):
+    def __init__(
+        self,
+        model_dir: str,
+        device_id: int = 0,
+        *,
+        num_hidden_layers: Optional[int] = None,
+        layers_block_type: Optional[List[str]] = None,
+    ):
         self.device_id = device_id
-        config = AutoConfig.from_pretrained(model_dir)
-        config.num_hidden_layers = num_hidden_layers
+        model_kwargs = {}
+        if num_hidden_layers is not None:
+            model_kwargs["num_hidden_layers"] = num_hidden_layers
+        if layers_block_type is not None:
+            model_kwargs["layers_block_type"] = layers_block_type
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_dir, config=config, torch_dtype=torch.bfloat16, attn_implementation="eager"
+            model_dir,
+            torch_dtype=torch.bfloat16,
+            attn_implementation="eager",
+            **model_kwargs,
         ).to(f"cuda:{device_id}")
+        # Hybrid configs (e.g. NemotronH) derive num_hidden_layers from
+        # ``layers_block_type`` and silently ignore the num_hidden_layers
+        # override; callers must pass a truncated ``layers_block_type`` for
+        # such models. Catch a silently ignored override loudly here.
+        if num_hidden_layers is not None:
+            assert self.model.config.num_hidden_layers == num_hidden_layers, (
+                f"num_hidden_layers override silently ignored: "
+                f"HF loaded {self.model.config.num_hidden_layers}, "
+                f"expected {num_hidden_layers}"
+            )
         self.all_weights = {}
         self.device_uuid = [get_device_uuid(i) for i in range(torch.cuda.device_count())]
         self._replicate_weights()