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_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_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()
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: