From 89ab2a9e481af119fc4ea30c4497c1efe232d27f Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:46:02 -0700 Subject: [PATCH 1/2] [https://nvbugs/6572835][fix] Use HashStore for single-rank MegaMoE tests Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - .../unittest/_torch/modules/moe/test_moe_backend.py | 13 ++++++------- .../unittest/_torch/modules/moe/test_moe_module.py | 9 +++++++++ 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1dd2fdce974d..0021dcc9640b 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -317,7 +317,6 @@ unittest/_torch/modeling/test_modeling_qwen_moe.py::TestQwenMoe::test_qwen_moe_a unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6602176) unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_CUTEDSL" SKIP (https://nvbugs/6601578) -unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_CUTEDSL" SKIP (https://nvbugs/6572835) unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" SKIP (https://nvbugs/6464169) unittest/_torch/modules/test_w4a16_nvfp4_linear.py::test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin SKIP (https://nvbugs/6581071) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 2021a0f6b999..97c44b666f72 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -17,7 +17,6 @@ import importlib import itertools import logging -import os from types import SimpleNamespace from typing import List, Optional, Tuple from unittest.mock import MagicMock @@ -138,13 +137,13 @@ def _ensure_single_proc_dist_for_megamoe(backend_type: MoeBackendType, rank: int pytest.skip("CUDA required for MegaMoE tests") if dist.is_initialized(): return - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29561") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", str(rank)) torch.cuda.set_device(rank) - dist.init_process_group(backend="nccl", rank=0, world_size=1) + dist.init_process_group( + backend="nccl", + store=dist.HashStore(), + rank=0, + world_size=1, + ) def should_skip_gptoss( diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index a7fae44529b5..c418fb2f1a29 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -147,6 +147,15 @@ def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> No pytest.skip("CUDA required for MegaMoE tests") if dist.is_initialized(): return + if world_size == 1: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + store=dist.HashStore(), + rank=rank, + world_size=world_size, + ) + return os.environ.setdefault("MASTER_ADDR", "127.0.0.1") os.environ.setdefault("MASTER_PORT", "29561") os.environ["RANK"] = str(rank) From 41512e6bc7d42a2cc27f25d31f02c5534c89e91e Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:00:13 -0700 Subject: [PATCH 2/2] [https://nvbugs/6572835][fix] Use HashStore in single-rank MoE benchmark Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- tests/microbenchmarks/bench_moe/utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/microbenchmarks/bench_moe/utils.py b/tests/microbenchmarks/bench_moe/utils.py index 40e7b0d75852..c5fbb4e6414b 100644 --- a/tests/microbenchmarks/bench_moe/utils.py +++ b/tests/microbenchmarks/bench_moe/utils.py @@ -74,6 +74,14 @@ def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> No raise RuntimeError("CUDA required for MegaMoE backend") if dist.is_initialized(): return + if world_size == 1: + dist.init_process_group( + backend="nccl", + store=dist.HashStore(), + rank=rank, + world_size=world_size, + ) + return rendezvous = None if rank == 0: master_addr = os.environ.get("MASTER_ADDR") or "127.0.0.1"