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.