From c6ce7585cdb7e0beaaea52a2270010ff354c53d6 Mon Sep 17 00:00:00 2001 From: Bowen Fu <5812640+BowenFu@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:29:40 +0000 Subject: [PATCH] [https://nvbugs/6566765][fix] Release Qwen MoE CUDA memory between tests Signed-off-by: Bowen Fu <5812640+BowenFu@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- .../_torch/modeling/test_modeling_qwen_moe.py | 22 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c709cc7d3984..1f5d5196ef97 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -355,9 +355,6 @@ unittest/_torch/attention/test_attention_backends.py::test_attention_backend[dee unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-cpp_scheduler-TorchSampler] SKIP (https://nvbugs/6561559) unittest/_torch/misc/test_autotuner.py::test_cutedsl_nvfp4_heuristic_matches_full_sweep SKIP (https://nvbugs/6490028) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) -unittest/_torch/modeling/test_modeling_qwen_moe.py::TestQwenMoe::test_qwen_moe_allclose_to_hf[backend:trtllm-use_cuda_graph:False] SKIP (https://nvbugs/6566765) -unittest/_torch/modeling/test_modeling_qwen_moe.py::TestQwenMoe::test_qwen_moe_allclose_to_hf[backend:trtllm-use_cuda_graph:True] SKIP (https://nvbugs/6575012) -unittest/_torch/modeling/test_modeling_qwen_moe.py::TestQwenMoe::test_qwen_moe_allclose_to_hf[backend:vanilla-use_cuda_graph:False] SKIP (https://nvbugs/6566765) 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) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen_moe.py index 39cbf33b823d..a3ada975b53c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen_moe.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen_moe.py @@ -1,3 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc import unittest from copy import deepcopy from dataclasses import dataclass @@ -67,6 +71,24 @@ def __repr__(self) -> str: class TestQwenMoe(unittest.TestCase): + @staticmethod + def _release_cuda_memory(): + gc.collect() + torch.cuda.empty_cache() + + @classmethod + def setUpClass(cls): + super().setUpClass() + # The A30 suite runs every Qwen modeling test in one process. Clear + # cyclic model references left by earlier classes before allocating + # both the HF and TensorRT-LLM models used by these comparisons. + cls._release_cuda_memory() + + def tearDown(self): + # Do not leave either model alive for the next parameterized case. + self._release_cuda_memory() + super().tearDown() + @parameterized.expand([None, "FP8", "NVFP4"]) def test_qwen_moe_sanity(self, quant_algo): config_dict = deepcopy(QWEN2_57B_A14B_CONFIG)