From 10dd4ecf21f7ba31e8b42f7d6e23a784dc630803 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 31 Aug 2026 19:39:55 +0000
Subject: [PATCH 1/2] Ask transformers whether a model supports context
 parallelism

---
 src/accelerate/big_modeling.py | 31 +++++++++++--------------------
 1 file changed, 11 insertions(+), 20 deletions(-)

diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 98de9074de5..8cb2e550200 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -802,27 +802,18 @@ def _attach_context_parallel_hooks(
 
     """
 
-    # The hook below discards the attention mask and forces `is_causal=True`. That is only
-    # equivalent to the model's own masking for plain causal attention. Models whose layers use
-    # a *stricter* mask (sliding-window or chunked attention) would otherwise be trained with
-    # full causal attention silently, so refuse them up front. Without this hook torch raises a
-    # shape error for such models, so nothing that works today starts failing here.
-    config = getattr(model, "config", None)
-    config = config.get_text_config() if hasattr(config, "get_text_config") else config
-    layer_types = getattr(config, "layer_types", None)
-    if layer_types is not None:
-        non_full = {layer_type for layer_type in layer_types if layer_type != "full_attention"}
-    else:
-        # Models that predate `layer_types` (Mistral, for one) apply a sliding window to every layer
-        # whenever `sliding_window` is set.
-        non_full = {"sliding_attention"} if getattr(config, "sliding_window", None) else set()
-    if non_full:
+    # The hook below discards the attention mask and forces `is_causal=True`, which only matches the
+    # model's own masking for plain causal attention. `supports_context_parallel` is transformers'
+    # answer for whether that substitution is safe; it also lets a model rule context parallelism out
+    # whatever its config, which inspecting `layer_types` here could not express. Defaults to allowed,
+    # so a transformers without the property behaves as before.
+    if not getattr(model, "supports_context_parallel", True):
         raise ValueError(
-            f"Context parallelism does not support attention layers of type {sorted(non_full)} "
-            f"(model {model.__class__.__name__}). Context parallelism can only express full causal "
-            "attention: the per-layer mask has to be dropped, so those layers would silently be "
-            "trained with full causal attention instead. Use a full-attention model, or disable "
-            "context parallelism."
+            f"{model.__class__.__name__} does not support context parallelism. Context parallelism can "
+            "only express full causal attention: the per-layer mask has to be dropped, so a layer using "
+            "a stricter mask (sliding-window or chunked) would silently be trained with full causal "
+            "attention instead, and a layer carrying a recurrent state along the sequence never has that "
+            "state exchanged between ranks. Use a full-attention model, or disable context parallelism."
         )
 
     def _self_attn_pre_forward_hook(_module, module_args, module_kwargs):

From c19955a4009ccc1e0b5e37ee80aabee633f3222c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 31 Aug 2026 19:47:33 +0000
Subject: [PATCH 2/2] Trim the context-parallel guard comment

---
 src/accelerate/big_modeling.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 8cb2e550200..a26624784a1 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -802,11 +802,8 @@ def _attach_context_parallel_hooks(
 
     """
 
-    # The hook below discards the attention mask and forces `is_causal=True`, which only matches the
-    # model's own masking for plain causal attention. `supports_context_parallel` is transformers'
-    # answer for whether that substitution is safe; it also lets a model rule context parallelism out
-    # whatever its config, which inspecting `layer_types` here could not express. Defaults to allowed,
-    # so a transformers without the property behaves as before.
+    # The hook below drops the attention mask and forces `is_causal=True`; transformers says whether
+    # that substitution is safe. Defaults to allowed, so an older transformers behaves as before.
     if not getattr(model, "supports_context_parallel", True):
         raise ValueError(
             f"{model.__class__.__name__} does not support context parallelism. Context parallelism can "