From d26e148b04de4d1c20c42d865a64f8f530f463b8 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 29 May 2026 17:23:52 +0000 Subject: [PATCH 1/3] skip the test if the env variable is not turned on for single grouped weight Signed-off-by: Varun Thumbe --- tests/pytorch/test_sanity.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 27eafbecdc..1055db0ea2 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -598,7 +598,11 @@ def test_sanity_grouped_linear( # Small batch size used to catch bug from https://github.com/NVIDIA/TransformerEngine/pull/1527. bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") == "0" and single_param: + pytest.skip( + "single parameter grouped linear requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) skip_unsupported_backward_override("grouped_linear", fp8_recipe, backward_override) if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) From 70b99a9a78b23dc1564bf5d9b9571fa021b9f03a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 17:28:04 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_sanity.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 1055db0ea2..9aab728ba6 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -599,10 +599,7 @@ def test_sanity_grouped_linear( bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") == "0" and single_param: - pytest.skip( - "single parameter grouped linear requires" - " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" - ) + pytest.skip("single parameter grouped linear requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") skip_unsupported_backward_override("grouped_linear", fp8_recipe, backward_override) if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) From e8c403775c9b4dc8020b73db1df8292036725944 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 29 May 2026 10:32:58 -0700 Subject: [PATCH 3/3] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 --- tests/pytorch/test_sanity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 9aab728ba6..404fae85fd 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -598,7 +598,7 @@ def test_sanity_grouped_linear( # Small batch size used to catch bug from https://github.com/NVIDIA/TransformerEngine/pull/1527. bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) - if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") == "0" and single_param: + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_param: pytest.skip("single parameter grouped linear requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") skip_unsupported_backward_override("grouped_linear", fp8_recipe, backward_override) if fp8_recipe is not None: