Skip to content

Fix GPT-2 FlashAttention mask preparation without pipeline - #6434

Draft
sdjasj wants to merge 1 commit into
hpcaitech:mainfrom
sdjasj:agent/fix-gpt2-flash-mask-6389
Draft

Fix GPT-2 FlashAttention mask preparation without pipeline#6434
sdjasj wants to merge 1 commit into
hpcaitech:mainfrom
sdjasj:agent/fix-gpt2-flash-mask-6389

Conversation

@sdjasj

@sdjasj sdjasj commented Jul 29, 2026

Copy link
Copy Markdown

What changed

  • use ColossalAI's GPT-2 model forward whenever FlashAttention is enabled without pipeline parallelism, not only when sequence parallelism is enabled
  • let GPT-2 attention-mask preparation handle calls without an explicit mask
  • add the reported tp_size=1, pp_size=1, enable_all_optimization=True configuration to the GPT-2 regression matrix

Why

For a non-pipeline, non-sequence-parallel GPT-2 model, the policy replaced GPT2Attention.forward with the ColoAttention implementation but left the Hugging Face GPT2Model.forward unchanged. Hugging Face passed its additive attention-mask Tensor directly to the replacement, which attempted **attention_mask and raised a TypeError because the value was not a mapping.

The existing ColossalAI GPT-2 model forward already converts masks into the keyword dictionary expected by ColoAttention and preserves causal-mask behavior.

Impact

GPT-2 works with HybridParallelPlugin(enable_all_optimization=True) when TP, PP, and SP sizes are all one.

Fixes #6389.

Validation

  • 2-GPU minimal GPT-2 training regression with the reported plugin configuration and a padded Tensor mask: forward, backward, and optimizer step passed
  • compileall and git diff --check: passed

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.

[BUG]: TypeError in ColoAttention.attention() when enable_all_optimization=True in HybridParallelPlugin

1 participant