Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shuyixiong wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ee4bbdd
Fix nemotron weight update test
shuyixiong Jul 22, 2026
74ecf49
Waive test
shuyixiong Jul 22, 2026
efeaa76
[https://nvbugs/6437410][fix] pre-install mamba deps in Ray CI stage,…
shikicloud Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions jenkins/L0_Test.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" \

@brnguyen2 brnguyen2 Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cpver is a parameter of this function (defaults to cp312, set to cp310 for PY310 stages at line 5687), but the wheel URLs hardcode cp312-cp312. No PY310 Ray stage exists today, so this isn't broken now — but it will fail with a 404 the moment one is added, and the failure will look like a network flake rather than a version mismatch. Use ${cpver}-${cpver} for consistency with the rest of the file.

"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")
Expand Down
4 changes: 4 additions & 0 deletions jenkins/scripts/slurm_install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
zhaoyangwang-nvidia marked this conversation as resolved.
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"
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
138 changes: 39 additions & 99 deletions tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(

@brnguyen2 brnguyen2 Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

e.name is only the module that directly failed to import. Both wheels are installed with --no-deps, so if a transitive dependency is missing from the container (einops is the likely one for mamba_ssm), this reports einops is not installed, so the mamba fast path is unavailable followed by advice about installing mamba-ssm — which is already installed. That sends the next person down the wrong path.

Include the original exception so the real cause is visible:

pytest.fail(
    f"mamba fast path unavailable ({e!r}). The Ray CI stage installs "
    "mamba-ssm and causal-conv1d alongside ray; see jenkins/scripts/slurm_install.sh."
)

Same message also fires on an ABI mismatch (ImportError: undefined symbol) after a container torch bump, which is the more likely failure mode given the pinned torch26.01/torch26.04 wheels — repr(e) makes that case readable too.

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)
Expand All @@ -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",
Expand Down Expand Up @@ -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)

@brnguyen2 brnguyen2 Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good that the number is backed by 5 runs rather than tuned until green. One question though: every other BF16-vs-BF16 comparison in this file uses the default 0.9, and 0.8 is reserved for NVFP4-vs-BF16. A BF16 reference topping out at ~0.89 mean top-20 overlap is a bigger gap than "small numerical differences" in the SSM path would suggest.

Did you check whether the gap is dominated by the mamba path specifically (e.g. does the overlap improve if you compare only the attention/MoE layer outputs, or with mamba_ssm_cache_dtype left at default)? If it's genuinely inherent to chunked-scan vs HF's kernel, say so in the comment; if it's unexplained, 0.85 buys a green test at the cost of the check no longer catching a real regression in the mamba weight-update path.


del hf_model

@brnguyen2 brnguyen2 Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the last statement in the function, so the local dies with the frame either way — the del only matters if you intend it to run before the with block exits (it doesn't; it's after). The autouse release_shared_cuda_memory fixture already does gc.collect() + ipc_collect(). Drop it, or move it inside the with if the intent was to free the HF replicas before LLM teardown.

33 changes: 28 additions & 5 deletions tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
):
Comment on lines +44 to +51

@coderabbitai coderabbitai Bot Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Verify the repository Python target before applying Python 3.10+ syntax.
fd -HI '^(pyproject\.toml|setup\.py|\.python-version|\.tool-versions)$' . -0 |
  xargs -0 -r rg -n 'requires-python|python_requires|target-version|python'

Repository: NVIDIA/TensorRT-LLM

Length of output: 6821

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"

printf '%s\n' '--- repository guidance ---'
if [ -f CODING_GUIDELINES.md ]; then
  cat -n CODING_GUIDELINES.md
else
  printf '%s\n' 'CODING_GUIDELINES.md not found'
fi

printf '%s\n' '--- target file ---'
cat -n "$file"

printf '%s\n' '--- relevant imports and constructor usages ---'
rg -n -C 3 '^(from typing|import typing)|RefHFModelWithIPCHandles|num_hidden_layers|layers_block_type' "$file"

printf '%s\n' '--- changed-file summary ---'
git diff --stat -- "$file"

printf '%s\n' '--- test-list references ---'
rg -n 'test_llm_update_weights|update_weights' tests/integration/test_lists tests 2>/dev/null | head -200

Repository: NVIDIA/TensorRT-LLM

Length of output: 50377

🏁 Script executed:

#!/bin/bash
set -euo pipefail

target="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"

printf '%s\n' '--- target constructor and nearby code ---'
sed -n '43,95p' "$target"

printf '%s\n' '--- test functions in target ---'
rg -n '^def test_|^async def test_|^class ' "$target"

printf '%s\n' '--- base class definition ---'
rg -n -C 8 'class RefHFModel|def __init__' tests utils tensorrt_llm 2>/dev/null | head -240

printf '%s\n' '--- test-list entries for this module ---'
rg -n -C 3 'test_llm_update_weights|ray_orchestrator.*single_gpu|update_weights' tests/integration/test_lists 2>/dev/null || true

printf '%s\n' '--- typing imports and usages in target ---'
python3 - <<'PY'
from pathlib import Path
import ast

path = Path("tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py")
tree = ast.parse(path.read_text())
imports = {}
for node in tree.body:
    if isinstance(node, ast.ImportFrom) and node.module == "typing":
        for alias in node.names:
            imports[alias.asname or alias.name] = alias.name

used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}
for local_name, imported_name in imports.items():
    print(f"{imported_name}: {'used' if local_name in used else 'unused'}")
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 20284

🏁 Script executed:

#!/bin/bash
set -euo pipefail

target="tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py"

printf '%s\n' '--- target diff ---'
git diff --unified=3 -- "$target"

printf '%s\n' '--- base helper candidates ---'
fd -HI -t f 'torch_ref|ref.*model' tests | head -100

printf '%s\n' '--- base helper definitions ---'
rg -n -C 12 'class RefHFModel|def __init__' tests/unittest tests/utils 2>/dev/null | rg -B 12 -A 12 'RefHFModel|torch_ref' | head -240 || true

printf '%s\n' '--- test-list references ---'
if [ -d tests/integration/test_lists ]; then
  rg -n -C 3 'test_llm_update_weights|ray_orchestrator|update_weights' tests/integration/test_lists || true
else
  printf '%s\n' 'tests/integration/test_lists not found'
fi

printf '%s\n' '--- typing import usage in target ---'
python3 - <<'PY'
from pathlib import Path
import ast

path = Path("tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py")
tree = ast.parse(path.read_text())
typing_names = {}
for node in tree.body:
    if isinstance(node, ast.ImportFrom) and node.module == "typing":
        for alias in node.names:
            typing_names[alias.asname or alias.name] = alias.name
used = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}
for local_name, imported_name in typing_names.items():
    print(f"{imported_name}: {'used' if local_name in used else 'unused'}")
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 21409

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- target test decorators and calls ---'
sed -n '170,360p' tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py

printf '%s\n' '--- cross-file helper and Nemotron-H references ---'
rg -n -C 4 'RefHFModelWithIPCHandles|layers_block_type|num_hidden_layers' \
  tests/unittest/_torch/ray_orchestrator tests/unittest/utils/torch_ref.py

printf '%s\n' '--- exact CI list entries ---'
sed -n '238,255p' tests/integration/test_lists/test-db/l0_h100.yml

Repository: NVIDIA/TensorRT-LLM

Length of output: 40864

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- Nemotron-H test function ---'
sed -n '620,715p' tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py
rg -n -C 2 'test.*nemotron|test_llm_update_weights.*nemotron|multi_gpu' \
  tests/integration/test_lists/test-db tests/integration/test_lists/qa 2>/dev/null || true

Repository: NVIDIA/TensorRT-LLM

Length of output: 50377

Use Python 3.10 type syntax and annotate the constructor return.

RefHFModelWithIPCHandles.__init__ needs a -> None annotation. Replace Optional[int] and Optional[List[str]] with int | None and list[str] | None.

Test coverage summary

  • Changed test functions: none. Changed helper: RefHFModelWithIPCHandles.__init__.
  • Existing consumers include the single-GPU update-weight tests and test_llm_update_weights_nemotron_h.
  • Test lists: l0_h100.yml covers parts 0–2; l0_dgx_b200.yml covers part 4.
  • Verdict: sufficient.
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py`
around lines 44 - 51, Update RefHFModelWithIPCHandles.__init__ to annotate its
return as None and replace Optional[int] with int | None and Optional[List[str]]
with list[str] | None, preserving the existing constructor behavior and
parameters.

Sources: Coding guidelines, Learnings

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
Comment thread
tongyuantongyu marked this conversation as resolved.
# ``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()
Expand Down
Loading