fix(model): reject CP>1 in Gemma4DenseProvider with a clear error (#5895)

Signed-off-by: dafu-wu <wuchengyi2006@163.com>
Co-authored-by: gautham-kollu <gkollu@nvidia.com>

Author: dafu-wu

src/megatron/bridge/models/gemma/gemma4_provider.py

@@ -200,6 +200,20 @@ def provide(
         if vp_stage is not None or getattr(self, "pipeline_model_parallel_size", 1) != 1:
             raise NotImplementedError("Gemma4DenseProvider currently supports PP=1 only.")
 
+        if getattr(self, "context_parallel_size", 1) != 1:
+            # get_gemma4_layer_spec() builds the dense layer from LocalSpecProvider,
+            # whose core_attention() is the non-TE DotProductAttention. That module
+            # asserts context_parallel_size == 1, so CP fails deep inside layer
+            # construction with "Context parallelism is only supported by
+            # TEDotProductAttention! when instantiating Gemma4DenseSelfAttention".
+            # Surface it here instead, next to the PP guard, so the message names the
+            # provider and the knob to change.
+            raise NotImplementedError(
+                "Gemma4DenseProvider currently supports CP=1 only: its layer spec uses "
+                "the local DotProductAttention, which does not implement context "
+                "parallelism."
+            )
+
         return self.build(
             pre_process=True if pre_process is None else pre_process,
             post_process=True if post_process is None else post_process,

tests/unit_tests/models/gemma/test_gemma4_provider.py

@@ -142,6 +142,11 @@ def test_provide_rejects_pipeline_parallel(self, provider):
         with pytest.raises(NotImplementedError, match="PP=1"):
             provider.provide()
 
+    def test_provide_rejects_context_parallel(self, provider):
+        provider.context_parallel_size = 2
+        with pytest.raises(NotImplementedError, match="CP=1"):
+            provider.provide()
+
     def test_provide_rejects_virtual_pipeline_stage(self, provider):
         with pytest.raises(NotImplementedError, match="PP=1"):
             provider.provide(vp_stage=0)