From 41d3adf379966c8935cc069ed30519315d44cd5b Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 16:58:55 -0500 Subject: [PATCH 1/2] [shardformer] fall back to native RMSNorm when apex is not installed Fixes #6260. `FusedRMSNorm.from_native_module` is meant to keep the native module and warn when apex is missing, but the fallback is unreachable: without apex `FusedRMSNormWithHook` stays `None`, so calling it raises `TypeError: 'NoneType' object is not callable`, while the code only catches `ImportError`. As a result `enable_fused_normalization=True` (and therefore `enable_all_optimization=True`) crashes for every RMSNorm model (LLaMA, Qwen2, Mistral, DeepSeek, ...) on installations without apex: RuntimeError: Failed to replace norm of type LlamaRMSNorm with FusedRMSNorm with the exception: 'NoneType' object is not callable. Check for `None` explicitly so the intended fallback is taken, matching the behaviour `FusedLayerNorm` already has. Add a unit test that simulates the missing dependency, so it is exercised on CI where apex is installed. --- colossalai/shardformer/layer/normalization.py | 20 ++++++++++--------- .../test_layer/test_layernorm.py | 18 ++++++++++++++++- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/colossalai/shardformer/layer/normalization.py b/colossalai/shardformer/layer/normalization.py index 2edca3b3e1c4..909532a27b1f 100644 --- a/colossalai/shardformer/layer/normalization.py +++ b/colossalai/shardformer/layer/normalization.py @@ -329,19 +329,21 @@ def from_native_module(module: nn.Module, sp_partial_derived: bool = False, *arg eps = module.variance_epsilon if hasattr(module, "variance_epsilon") else module.eps elementwise_affine = getattr(module, "elementwise_affine", True) - try: - rmsnorm = FusedRMSNormWithHook( - normalized_shape=normalized_shape, - eps=eps, - elementwise_affine=elementwise_affine, - ) - except ImportError: + # FusedRMSNormWithHook stays None when apex (or torch_npu) is not available + if FusedRMSNormWithHook is None: warnings.warn( - "Module replacement failed.\ - Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMS normalization kernel" + "Module replacement failed. " + "Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMS normalization kernel. " + "Using native RMSNorm instead." ) return module + rmsnorm = FusedRMSNormWithHook( + normalized_shape=normalized_shape, + eps=eps, + elementwise_affine=elementwise_affine, + ) + rmsnorm.weight = module.weight if sp_partial_derived: diff --git a/tests/test_shardformer/test_layer/test_layernorm.py b/tests/test_shardformer/test_layer/test_layernorm.py index b0deff6b8fa4..e7948639c82a 100644 --- a/tests/test_shardformer/test_layer/test_layernorm.py +++ b/tests/test_shardformer/test_layer/test_layernorm.py @@ -1,12 +1,13 @@ from contextlib import nullcontext +import pytest import torch import torch.nn as nn from torch.testing import assert_close import colossalai from colossalai.lazy import LazyInitContext -from colossalai.shardformer.layer import FusedLayerNorm +from colossalai.shardformer.layer import FusedLayerNorm, FusedRMSNorm, normalization from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn @@ -50,5 +51,20 @@ def test_layernorm(): spawn(run_dist, nprocs=2) +def test_fused_rmsnorm_fallback_without_apex(monkeypatch): + # FusedRMSNormWithHook is None when apex is not installed, the native module should be kept in that case + monkeypatch.setattr(normalization, "FusedRMSNormWithHook", None) + + class RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + norm = RMSNorm(128) + with pytest.warns(UserWarning, match="apex"): + assert FusedRMSNorm.from_native_module(norm) is norm + + if __name__ == "__main__": test_layernorm() From 2fe785df04e7e25fa1edb70edd7b4baf76099a70 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:58:24 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_checkpoint_io/test_gemini_checkpoint_io.py | 2 +- tests/test_zero/test_gemini/test_optim.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index ba60de0bd1c1..3649b18ec8bb 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -23,7 +23,7 @@ ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False] try: - import apex.amp + pass APEX_AVAILABLE = True except ImportError: diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index b89c591882d6..a0970e43db9e 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -15,7 +15,7 @@ from tests.kit.model_zoo import model_zoo, run_fwd_bwd try: - import apex.amp + pass APEX_AMP_AVAILABLE = True except ImportError: