Refuse context parallelism for models with sliding-window or chunked attention layers

[qgallouedec]

_attach_context_parallel_hooks attaches this hook to every self_attn module:

def _self_attn_pre_forward_hook(_module, module_args, module_kwargs):
    if "attention_mask" in module_kwargs:
        module_kwargs["attention_mask"] = None
        module_kwargs["is_causal"] = True
    return module_args, module_kwargs

Its own docstring says it will "check if it is a causal mask, if yes, will add a kwarg is_causal=True, otherwise will raise an error". The implementation does neither check nor raise: it discards whatever mask the layer was given.

Replacing the mask with is_causal=True is only equivalent for a plain causal mask. For a model whose layers
use a stricter mask (sliding-window or chunked attention) the layer is silently switched to full
causal attention.

That is most of the current crop, not a legacy corner: gpt-oss makes every other layer sliding (window 128), Gemma 4 makes 5 of every 6 (window 512), Gemma 3 the same ratio (window 1024), Muse-Glimmer-30B 39 of its 52 layers (window 2048), and Mistral/Ministral every layer.

Training runs, loss looks plausible, and the model is trained with the wrong attention pattern. There is one unanswered user report of exactly this (PyTorch forums: "training speed improved significantly — but model performance dropped").

Note this hook is what creates the silence: with the hook removed, torch itself raises a shape error for these models (The expanded size of the tensor (512) must match the existing size (256) …). So the behavior being replaced is not "working", it is a wrong-but-quiet run where torch would have refused.

Repro

Weight-independent probe: perturb one token at position 0 and count how many output positions change. Under a correctly applied window W, only positions within reach may change (two layers reach 2W); under full causal, every later position changes.

torchrun --nproc_per_node=2 repro_cp_mask_drop.py

On main:

no CP: 24.8% of positions change (expected ~25% for 2 layers of window 64)
CP=2 : 100.0%  (100% means the window was dropped -> full causal)

With this PR the same command stops instead:

ValueError: Context parallelism does not support attention layers of type ['sliding_attention']
(model Qwen3ForCausalLM). 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.

The same probe with two packed documents (block-diagonal mask via restarting position_ids) gives 50.0% without CP and 100.0% with CP: packed documents attend across their boundaries.

repro_cp_mask_drop.py

"""Does context parallelism preserve a sliding-window mask? (It does not.)

Weight-independent probe: perturb the token at position 0 and count how many output positions
change. Under a correctly applied window `W`, only positions within reach of the window may change
(two layers of window W reach 2*W); under full causal attention, every later position changes.

Run: torchrun --nproc_per_node=2 repro_cp_mask_drop.py [--dump PATH]
"""

import argparse
import os

import torch
import torch.distributed as dist
from accelerate.big_modeling import _attach_context_parallel_hooks
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from transformers import AutoModelForCausalLM, Qwen3Config

SEQ, WINDOW = 512, 64


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dump", help="save the per-position change mask for plotting")
    args = parser.parse_args()

    rank = int(os.environ["RANK"])
    torch.cuda.set_device(rank)
    dist.init_process_group("nccl")
    device = torch.device("cuda", rank)
    mesh = init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("cp",))

    config = Qwen3Config(vocab_size=1024, hidden_size=256, intermediate_size=512, num_hidden_layers=2,
                         num_attention_heads=4, num_key_value_heads=2, head_dim=64,
                         use_sliding_window=True, sliding_window=WINDOW, max_window_layers=0,
                         attn_implementation="sdpa")
    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_config(config).to(device, torch.bfloat16).eval()

    ids_a = torch.randint(0, 1024, (1, SEQ), device=device)
    ids_b = ids_a.clone()
    ids_b[0, 0] = (ids_a[0, 0] + 1) % 1024  # perturb position 0 only
    positions = torch.arange(SEQ, device=device).unsqueeze(0)

    def changed(cp):
        def logits(ids):
            kwargs = dict(input_ids=ids, position_ids=positions, use_cache=False)
            if not cp:
                with torch.no_grad():
                    return model(**kwargs).logits.float()
            with torch.no_grad(), context_parallel(mesh, buffers=[ids, positions], buffer_seq_dims=[1, 1]):
                return model(**kwargs).logits.float()

        delta = (logits(ids_a) - logits(ids_b)).abs().amax(-1)[0] > 1e-3
        if not cp:
            return delta
        gathered = [torch.zeros_like(delta) for _ in range(dist.get_world_size())]
        dist.all_gather(gathered, delta)
        return torch.cat(gathered)

    reference = changed(cp=False)
    _attach_context_parallel_hooks(model)
    with_cp = changed(cp=True)

    if rank == 0:
        print(f"no CP: {reference.float().mean():.1%} of positions change "
              f"(expected ~{2 * WINDOW / SEQ:.0%} for 2 layers of window {WINDOW})")
        print(f"CP={dist.get_world_size()} : {with_cp.float().mean():.1%}  "
              f"(100% means the window was dropped -> full causal)")
        if args.dump:
            torch.save({"reference": reference.cpu(), "with_cp": with_cp.cpu(),
                        "seq": SEQ, "window": WINDOW}, args.dump)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

The fix

Reject these models when the hooks are attached, before any training happens:

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:
    raise ValueError(...)

layer_types alone is not enough: only 80 of the 491 model configs in Transformers define it, and Mistral deliberately does not (it warns and points you at Ministral instead) while still building a sliding-window mask for every layer whenever config.sliding_window is set. Checking layer_types first and falling back to sliding_window keeps the models that merely carry a stale sliding_window value (Qwen3 sets it while layer_types is all full_attention) from being rejected.

and correct the docstring to describe what the hook actually does.

Verified

case	before	after
sliding-window model (layer_types has sliding_attention)	trains silently with full causal attention	raises with an explanatory error
sliding-window model with no layer_types (Mistral-7B-v0.1)	trains silently with full causal attention	raises with an explanatory error
plain causal model (Qwen3-0.6B/8B/32B, Qwen3-30B-A3B)	works	works (no false positive)

Checked against real configs: Qwen3-8B, Qwen3-32B, Qwen3-0.6B, Qwen3-30B-A3B-Base and the VLM Qwen3-VL-2B-Instruct, Mixtral-8x7B-v0.1 and Mistral-7B-v0.3 (which turned its sliding window off) are allowed; google/gemma-2-2b, google/gemma-3-4b-it, Ministral-8B-Instruct-2410 and Mistral-7B-v0.1 raise. Edge cases are inert rather than fatal: a module with no config at all, or a config without layer_types, passes through untouched.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Extends the guard to the Ulysses path, which had none.
Linear-attention layers carry a recurrent state along the sequence and are computed inside the model's own layer code rather than through the attention interface Ulysses wraps, so every rank restarts that state from zero and never exchanges it. Measured against an unsharded fp32 reference, the first shard's forward output stays right while its gradients are wrong, which is not something a loss curve shows.
Scoped to linear_attention only: I have not verified whether sliding-window is affected under Ulysses, so it is left alone.

[SunMarc]

Thanks ! I don't mind having them here but if possible, it is better to have these in transformers.