Added inter document masking for manual and flash attention.#434
Added inter document masking for manual and flash attention.#434BlueCrescent wants to merge 7 commits intomainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Adds inter-document masking support to the GPT-2 attention stack so sequences containing multiple concatenated documents can prevent cross-document attention for both manual attention and DAO flash attention (via varlen).
Changes:
- Added
CausalSelfAttention.prepare_inter_document_masking()and threaded optional masking info through attention execution paths. - Implemented DAO flash varlen execution path to support document-wise masking/splitting without padding leakage.
- Added extensive unit tests for inter-document masking behaviors and edge cases.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 8 comments.
| File | Description |
|---|---|
tests/models/test_causal_self_attention.py |
Adds comprehensive tests for inter-document masking across manual and DAO flash attention implementations. |
src/modalities/models/gpt2/gpt2_model.py |
Implements inter-document masking preparation, DAO flash varlen execution, and integrates masking into GPT2LLM/GPT2Block/CausalSelfAttention. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| torch.manual_seed(0) # FIXME remove or do within tests? | ||
|
|
||
|
|
There was a problem hiding this comment.
Setting torch.manual_seed(0) at module import time mutates global RNG state for the entire test run, which can make unrelated tests order-dependent. Prefer seeding inside the specific tests that need determinism (or via a fixture) instead of at import scope.
| torch.manual_seed(0) # FIXME remove or do within tests? |
| device = self.c_proj.weight.device | ||
| if self.attention_impl == AttentionImplementation.MANUAL: | ||
| batch_size = len(in_batch_seq_lens) | ||
| attn_mask = torch.zeros((batch_size, max_seq_len, max_seq_len), dtype=torch.bool, device=device) | ||
| for i, doc_seq_lens in enumerate(in_batch_seq_lens): | ||
| doc_boundaries = torch.cumsum(torch.tensor([0] + doc_seq_lens, device=device), dim=0) | ||
| for j in range(len(doc_boundaries) - 1): | ||
| start_idx = doc_boundaries[j] | ||
| end_idx = doc_boundaries[j + 1] | ||
| attn_mask[i, start_idx:end_idx, start_idx:end_idx] = True | ||
| return attn_mask |
There was a problem hiding this comment.
prepare_inter_document_masking() validates in_batch_seq_lens for the dao_flash path but not for the manual path. If a batch item's subsequence lengths sum to more than max_seq_len (or contain invalid values), the manual path can produce incorrect masks or raise a low-level indexing error. Consider applying the same validation (e.g., reuse _build_concatenated_lengths_tensor checks) for the manual implementation too.
| attention_masking_information (torch.Tensor | tuple[torch.Tensor, torch.Tensor, int] | None): | ||
| Optional tensor containing masking information for inter-document attention. | ||
|
|
||
| Returns: | ||
| torch.Tensor: The output tensor. |
There was a problem hiding this comment.
execute_attention() now accepts attention_masking_information, but inter-document masking is only applied in the MANUAL / DAO_FLASH implementations. For PYTORCH_FLASH, the mask is currently silently ignored (it always passes attn_mask=None). To avoid surprising behavior, consider raising NotImplementedError (or asserting the mask is None) when attention_impl is PYTORCH_FLASH and a mask is provided.
…ce lengths required for inter document attention masking.
… manual attention. - Also applied some review comments.
…ed eod_token_id to eos_token_id
…doc string Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 participant
What does this PR do?
Adds inter document masking for manual and flash attention.
General Changes
prepare_inter_document_masking()to CausalSelfAttention which computes 3D attention masks for manual attention and cu_seqlens for DAO flash attention. The input are the sub sequence lengths for each sequence. Thus, padded sequences are also supported.forward()call, inter document masking is applied.Breaking Changes
Checklist before submitting final PR
python tests/tests.py)CHANGELOG_DEV.md)