From ee4bbddd1c6480ac885f35bde09ce388397a1ef6 Mon Sep 17 00:00:00 2001 From: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:07:00 -0700 Subject: [PATCH 1/3] Fix nemotron weight update test Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com> --- .../test_llm_update_weights_multi_gpu.py | 111 +++++++++++++----- .../single_gpu/test_llm_update_weights.py | 18 ++- 2 files changed, 96 insertions(+), 33 deletions(-) diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py index ac5e6bfeee27..b905cd25b1e8 100644 --- a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py +++ b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py @@ -11,7 +11,6 @@ import re import subprocess import sys -import traceback from typing import Callable, List, Optional, Tuple import pytest @@ -1188,14 +1187,22 @@ def _freeze(): 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.""" + ``python -m pytest`` subprocess (via test_nemotron_h_body_impl) 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.""" 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 +1223,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 +1268,78 @@ 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) + # Looser threshold: Nemotron-H logits are compared against a BF16 + # reference and the mamba SSM / selective-scan path introduces small + # numerical differences (observed top-20 overlap ~0.89 vs the 0.9 + # default). + compare_logits(llm_logits, ref_logits, threshold=0.8) -def _nemotron_h_subprocess_entry(result_queue): - try: - _nemotron_h_body() - result_queue.put(None) - except BaseException: - result_queue.put(traceback.format_exc()) +# Guard so this inner test only runs inside the subprocess launched by +# ``test_llm_update_weights_nemotron_h`` (which sets the env var and targets it +# by node id). It carries no ``part*`` marker, so marker-filtered CI runs +# deselect it, and the env guard skips it in an unfiltered in-process run. +_NEMOTRON_H_BODY_ENV = "_TLLM_RUN_NEMOTRON_H_BODY" + + +@pytest.mark.skipif( + os.environ.get(_NEMOTRON_H_BODY_ENV) != "1", + reason="Inner body of test_llm_update_weights_nemotron_h; only run in the " + "subprocess spawned by that test.", +) +def test_nemotron_h_body_impl(): + _nemotron_h_body() @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) + """Runs the Nemotron-H body in a fresh ``python -m pytest`` subprocess so + HF transformers re-imports cleanly and picks up the mamba-ssm / + causal-conv1d fast path installed by the ``mamba_deps`` fixture (a plain + in-process run would keep the parent's negative import caches; see the + _nemotron_h_body docstring). Driving it as a subprocess — instead of a + hand-managed ``multiprocessing`` child — lets ``subprocess.run`` and the + inner pytest own the process lifecycle: a hang is bounded by ``timeout=``, + a crash surfaces as a non-zero return code, and the failure detail is the + inner pytest's own traceback.""" + # Must stay under the outer pytest ``--timeout`` so a genuine hang (e.g. a + # Ray/NCCL/CUDA deadlock) is reported here with useful output instead of + # the whole test being hard-killed at the pytest timeout. + subprocess_timeout_s = 1800.0 + + node_id = f"{os.path.abspath(__file__)}::test_nemotron_h_body_impl" + cmd = [ + sys.executable, + "-m", + "pytest", + node_id, + "--run-ray", + "-p", + "no:cacheprovider", + "-p", + "no:xdist", + "--tb=short", + "-s", + "-v", + ] + env = {**os.environ, _NEMOTRON_H_BODY_ENV: "1"} + try: + result = subprocess.run( + cmd, + env=env, + capture_output=True, + text=True, + timeout=subprocess_timeout_s, + ) + except subprocess.TimeoutExpired as e: + out = (e.stdout or "") + (e.stderr or "") + pytest.fail( + f"Nemotron-H subprocess did not complete within " + f"{subprocess_timeout_s:.0f}s (likely hung); terminated.\n{out}" + ) + if result.returncode != 0: + pytest.fail( + f"Nemotron-H subprocess failed (exit code {result.returncode}).\n" + f"{result.stdout}\n{result.stderr}" + ) diff --git a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py index df9322b10921..3290ecba0fa9 100644 --- a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py +++ b/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,21 @@ 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, **model_kwargs): self.device_id = device_id - config = AutoConfig.from_pretrained(model_dir) - config.num_hidden_layers = num_hidden_layers 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. + num_hidden_layers = model_kwargs.get("num_hidden_layers") + if num_hidden_layers is not None: + assert self.model.config.num_hidden_layers == num_hidden_layers self.all_weights = {} self.device_uuid = [get_device_uuid(i) for i in range(torch.cuda.device_count())] self._replicate_weights() From 74ecf492760883bfa5c4c96af3784dddfe761898 Mon Sep 17 00:00:00 2001 From: shuyixiong <219646547+shuyixiong@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:26:40 -0700 Subject: [PATCH 2/3] Waive test Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index fc8bae3f0f10..4a84e727b01c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -329,7 +329,6 @@ unittest/_torch/modules/test_w4a16_nvfp4_linear.py::test_nvfp4_attention_keeps_h unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) 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/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) unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) From efeaa76e881212ccbc73a11ff1874babb319ef0b Mon Sep 17 00:00:00 2001 From: shikicloud Date: Mon, 3 Aug 2026 19:42:48 -0700 Subject: [PATCH 3/3] [https://nvbugs/6437410][fix] pre-install mamba deps in Ray CI stage, drop nested pytest Signed-off-by: shikicloud --- jenkins/L0_Test.groovy | 6 + jenkins/scripts/slurm_install.sh | 4 + .../test_llm_update_weights_multi_gpu.py | 171 +++--------------- .../single_gpu/test_llm_update_weights.py | 21 ++- 4 files changed, 56 insertions(+), 146 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index d152129766c6..7e5254e53019 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4371,6 +4371,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") diff --git a/jenkins/scripts/slurm_install.sh b/jenkins/scripts/slurm_install.sh index 27c591760f5c..619a370caf63 100644 --- a/jenkins/scripts/slurm_install.sh +++ b/jenkins/scripts/slurm_install.sh @@ -28,6 +28,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" diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py index b905cd25b1e8..33f187cec84c 100644 --- a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py +++ b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py @@ -3,14 +3,10 @@ import base64 import gc -import importlib.util import json -import multiprocessing import os import pickle import re -import subprocess -import sys from typing import Callable, List, Optional, Tuple import pytest @@ -1126,73 +1122,27 @@ 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 - ``python -m pytest`` subprocess (via test_nemotron_h_body_impl) 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 # NemotronHConfig derives num_hidden_layers from ``layers_block_type`` @@ -1270,76 +1220,11 @@ def filter_fn(name: str) -> bool: llm_logits, ref_logits = run_generate(llm, hf_model, prompts, sampling_params) # Looser threshold: Nemotron-H logits are compared against a BF16 # reference and the mamba SSM / selective-scan path introduces small - # numerical differences (observed top-20 overlap ~0.89 vs the 0.9 - # default). - compare_logits(llm_logits, ref_logits, threshold=0.8) - - -# Guard so this inner test only runs inside the subprocess launched by -# ``test_llm_update_weights_nemotron_h`` (which sets the env var and targets it -# by node id). It carries no ``part*`` marker, so marker-filtered CI runs -# deselect it, and the env guard skips it in an unfiltered in-process run. -_NEMOTRON_H_BODY_ENV = "_TLLM_RUN_NEMOTRON_H_BODY" - - -@pytest.mark.skipif( - os.environ.get(_NEMOTRON_H_BODY_ENV) != "1", - reason="Inner body of test_llm_update_weights_nemotron_h; only run in the " - "subprocess spawned by that test.", -) -def test_nemotron_h_body_impl(): - _nemotron_h_body() - - -@pytest.mark.part4 -@skip_pre_hopper -def test_llm_update_weights_nemotron_h(mamba_deps): - """Runs the Nemotron-H body in a fresh ``python -m pytest`` subprocess so - HF transformers re-imports cleanly and picks up the mamba-ssm / - causal-conv1d fast path installed by the ``mamba_deps`` fixture (a plain - in-process run would keep the parent's negative import caches; see the - _nemotron_h_body docstring). Driving it as a subprocess — instead of a - hand-managed ``multiprocessing`` child — lets ``subprocess.run`` and the - inner pytest own the process lifecycle: a hang is bounded by ``timeout=``, - a crash surfaces as a non-zero return code, and the failure detail is the - inner pytest's own traceback.""" - # Must stay under the outer pytest ``--timeout`` so a genuine hang (e.g. a - # Ray/NCCL/CUDA deadlock) is reported here with useful output instead of - # the whole test being hard-killed at the pytest timeout. - subprocess_timeout_s = 1800.0 - - node_id = f"{os.path.abspath(__file__)}::test_nemotron_h_body_impl" - cmd = [ - sys.executable, - "-m", - "pytest", - node_id, - "--run-ray", - "-p", - "no:cacheprovider", - "-p", - "no:xdist", - "--tb=short", - "-s", - "-v", - ] - env = {**os.environ, _NEMOTRON_H_BODY_ENV: "1"} - try: - result = subprocess.run( - cmd, - env=env, - capture_output=True, - text=True, - timeout=subprocess_timeout_s, - ) - except subprocess.TimeoutExpired as e: - out = (e.stdout or "") + (e.stderr or "") - pytest.fail( - f"Nemotron-H subprocess did not complete within " - f"{subprocess_timeout_s:.0f}s (likely hung); terminated.\n{out}" - ) - if result.returncode != 0: - pytest.fail( - f"Nemotron-H subprocess failed (exit code {result.returncode}).\n" - f"{result.stdout}\n{result.stderr}" - ) + # 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 diff --git a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py index 3290ecba0fa9..959b4c5ebab6 100644 --- a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py +++ b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py @@ -41,8 +41,20 @@ def release_shared_cuda_memory(): class RefHFModelWithIPCHandles(RefHFModel): - def __init__(self, model_dir: str, device_id: int = 0, **model_kwargs): + 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 + 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, torch_dtype=torch.bfloat16, @@ -53,9 +65,12 @@ def __init__(self, model_dir: str, device_id: int = 0, **model_kwargs): # ``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. - num_hidden_layers = model_kwargs.get("num_hidden_layers") if num_hidden_layers is not None: - assert self.model.config.num_hidden_layers == num_hidden_layers + 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()