Skip to content

fix(engine): use model-native packed attention masks - #1558

Open
wm19999 wants to merge 2 commits into
areal-project:mainfrom
wm19999:fix/fsdp-packed-attention-mask
Open

fix(engine): use model-native packed attention masks#1558
wm19999 wants to merge 2 commits into
areal-project:mainfrom
wm19999:fix/fsdp-packed-attention-mask

Conversation

@wm19999

@wm19999 wm19999 commented Jul 22, 2026

Copy link
Copy Markdown

Description

Remove the model-family allowlist that injected a hard-coded per-layer attention-mask dictionary into FSDP packed microbatches. Passing None lets each Transformers model build the backend-specific mask it expects from AReaL's reset position IDs, avoiding the Llama/MiniCPM5 dict.ndim crash while keeping Qwen and multimodal preparation model-native.

This also adds CPU regressions for:

  • Llama/MiniCPM5-style, Qwen2, Qwen3, Qwen3-MoE, Qwen3.5, and Gemma3 configs
  • Qwen2.5-VL and Qwen3-VL mRoPE preparation
  • Packed-sequence isolation under eager attention and SDPA

Related Issue

Fixes #1557

Related to #1132, #1153, and #1442.

Type of Change

  • 🐛 Bug fix
  • ✨ New feature
  • 💥 Breaking change
  • 📝 Documentation update
  • ♻️ Refactoring
  • ⚡ Performance improvement
  • ✅ Test coverage improvement

Checklist

  • I have read the Contributing Guide
  • Pre-commit hooks pass (pre-commit run --all-files)
  • Relevant tests pass; new tests added for new functionality
  • Documentation updated (not applicable; no user-facing API or configuration changes)
  • Branch is up to date with main
  • Self-reviewed via /review-pr command
  • This PR was created by a coding agent via /create-pr
  • This PR is a breaking change

Breaking Change Details (if applicable):

None.

Additional Context

Validation:

  • PYTHONPATH=. python -m pytest -q tests/test_fsdp_packed_attention_mask.py — 10 passed
  • pre-commit run --all-files — passed
  • git diff --check — passed

Documentation is not affected because this only changes internal packed-mask plumbing.

wm19999 added 2 commits July 23, 2026 02:17
Let each Transformers model construct its backend-specific packed attention mask from reset position IDs.

Signed-off-by: amos_wang <86055438+wm19999@users.noreply.github.com>
Cover ordinary, multimodal, eager, and SDPA packed-mask paths.

Signed-off-by: amos_wang <86055438+wm19999@users.noreply.github.com>
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@github-actions

github-actions Bot commented Aug 6, 2026

Copy link
Copy Markdown

This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days.

Please add a comment or push new commits to keep it active.

Thank you for your contribution!

@github-actions github-actions Bot added the stale label Aug 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] FSDP packed SFT passes a dict attention mask to Llama models

1 participant