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/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) 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..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,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 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..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 @@ -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()