Merge branch 'cp-5-docs' into cp-6-tp-x-cp

Author: 3outeille

src/transformers/distributed/context_parallel.py

@@ -242,6 +242,13 @@ def _check_context_parallelism(model: nn.Module) -> None:
             "disable context parallelism."
         )
 
+    if model._is_stateful:
+        raise ValueError(
+            f"Context parallelism does not support {model.__class__.__name__}: its recurrent layers carry a state "
+            "along the sequence, so they would run on this rank's slice only and be silently cut at the slice "
+            "boundaries. Use another model, or disable context parallelism."
+        )
+
     text_config = model.config.get_text_config()
     layer_types = set(getattr(text_config, "layer_types", None) or ())
     if unsupported := layer_types - set(SOFTMAX_ATTENTION_LAYER_TYPES):

tests/context_parallel/test_context_parallel.py

@@ -419,6 +419,13 @@ class LlamaWithOwnAttention(LlamaForCausalLM):
         with self.assertRaisesRegex(ValueError, "`_supports_attention_backend`"):
             self._check(LlamaWithOwnAttention(_tiny_config()))
 
+    def test_refuses_stateful_models(self):
+        class StatefulLlama(LlamaForCausalLM):
+            _is_stateful = True
+
+        with self.assertRaisesRegex(ValueError, "carry a state along the sequence"):
+            self._check(StatefulLlama(_tiny_config()))
+
     def test_refuses_layers_that_mix_tokens_outside_attention(self):
         config = _tiny_config()
         config.layer_types = ["full_attention", "linear_attention"]