Skip to content

Add supports_context_parallel to PreTrainedModel - #48442

Open
qgallouedec wants to merge 2 commits into
mainfrom
supports-context-parallel
Open

Add supports_context_parallel to PreTrainedModel#48442
qgallouedec wants to merge 2 commits into
mainfrom
supports-context-parallel

Conversation

@qgallouedec

@qgallouedec qgallouedec commented Aug 31, 2026

Copy link
Copy Markdown
Member

CPU CI GPU run-slow

@SunMarc suggested in huggingface/accelerate#4177 that this check belongs here rather than in accelerate

It reads layer_types and sliding_window, which are transformers concepts, and there is already _supports_sdpa / _supports_flex_attn for this kind of capability flag.

CP can only express full causal attention. A layer with a stricter mask (sliding-window, chunked) gets that mask dropped and silently trains as full causal, and a layer carrying a recurrent state along the sequence (linear attention) never has that state exchanged between ranks.

huggingface/accelerate#4177 currently pattern-matches the config to detect this. With this property it becomes if not model.supports_context_parallel: raise.

_supports_context_parallel = False also lets a model rule it out whatever its config, which the config check cannot express (gpt-oss, where attention sinks need a custom kernel).

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 33431074964
Result: success | Grafana metrics are not available yet.

@ArthurZucker ArthurZucker left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, but this belongs in the DistributedMixin? (or not? TrainingMixing would be a better name?)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants