[agents docs] update models.md with class attributes and attention mask#13665
Merged
Conversation
…sk guidance - Add "Model class attributes" section documenting _no_split_modules, _repeated_blocks, _skip_layerwise_casting_patterns, _keep_in_fp32_modules, _cp_plan, and _supports_gradient_checkpointing with their corresponding user-facing APIs and how they work - Improve attention mask guidance: recommend passing None when no real padding exists, document backend compatibility - Move _no_split_modules from gotchas to its own section with first-principles explanation of why it's needed (accelerate device hooks) Co-Authored-By: Claude Opus 4.6 <[email protected]>
sayakpaul
reviewed
May 1, 2026
sayakpaul
left a comment
Member
There was a problem hiding this comment.
I like the dead code path addition. It's quite nice!
| - **Padding mask → bool `(B, L)` or `(B, 1, 1, L)`.** Stays compatible with the `*_varlen` kernels via `_normalize_attn_mask` (`attention_dispatch.py:639`), which reduces bool masks to `cu_seqlens`. Dense additive-float masks *cannot* be reduced this way and so lose the varlen path. This is the Qwen pattern (`transformer_qwenimage.py:951`). | ||
| - **Structural mask (causal, sliding-window, band-diagonal) → dense `(1, 1, L, L)` is unavoidable.** Row-varying patterns can't be expressed as `(B, L)`. Expect SDPA/Flex-only for these layers; consider Flex's `sliding_window_mask_mod` or FA3's native `window_size=` kwarg if backend flexibility matters. Consult `src/models/transformers/transformer_kandinsky.py` as a reference. | ||
| - **Padding mask → bool `(B, L)` or `(B, 1, 1, L)`.** Only pass when the batch actually contains different-length sequences (i.e. there is real padding). If all sequences are the same length, set the mask to `None` — many backends (flash, sage, aiter) raise `ValueError` on any non-None mask, and even SDPA-based backends pay unnecessary overhead processing a no-op mask. See `pipeline_qwenimage.py` `encode_prompt` for the pattern: `if mask.all(): mask = None`. When a mask is needed, use bool format — it stays compatible with the `*_varlen` kernels via `_normalize_attn_mask` (`attention_dispatch.py:639`), which reduces bool masks to `cu_seqlens`. Dense additive-float masks *cannot* be reduced this way and so lose the varlen path. | ||
| - **Other mask types (structural, BlockMask, etc.)** — if the model requires a different mask pattern, figure out how to support as many backends as possible (e.g. use `window_size` kwarg for sliding window on flash, `BlockMask` for Flex) and document which backends are supported for that model. |
Member
There was a problem hiding this comment.
I think we might want to keep the kandinsky reference in case it ends up discovering that Flex would be the way to go.
Collaborator
Author
There was a problem hiding this comment.
let's see if it can figure it out by itself next time when it reviews a model that only works with flex
I think it's pretty straightforward, but we can always add a reference if it struggles
| - **`[]`** — split anywhere you like. | ||
| - **`["MyBlock"]`** — keep all `MyBlock` instances intact on one device. | ||
|
|
||
| **Why it's needed.** When `accelerate` splits a model across devices, it installs hooks on leaf modules that move inputs to the module's device before `forward` runs. Any inline operation (`+`, `*`, `torch.cat`) that combines tensors from different submodules has no hook — if those submodules landed on different devices, it crashes with "tensors on different devices". The fix is either: (a) list the parent module in `_no_split_modules` so all its children stay co-located, or (b) pack the operation into its own `nn.Module`. Inline ops on outputs from the **same** submodule call are fine since they're already on the same device. |
Member
There was a problem hiding this comment.
A good rule of thumb is any block where the forward has a residual connection.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.