|
13 | 13 | # See the License for the specific language governing permissions and |
14 | 14 | # limitations under the License. |
15 | 15 |
|
| 16 | +import inspect |
16 | 17 | import json |
17 | 18 | import os |
| 19 | +from collections.abc import Callable |
18 | 20 | from pathlib import Path |
| 21 | +from typing import TYPE_CHECKING |
| 22 | + |
| 23 | +if TYPE_CHECKING: |
| 24 | + from types import FrameType |
19 | 25 | from typing import Any |
20 | 26 |
|
21 | 27 | import numpy as np |
|
24 | 30 | from datasets import load_dataset |
25 | 31 | from PIL import Image |
26 | 32 | from scripts.ar_validate import validate_ar |
| 33 | +from torch.distributed.tensor.experimental._attention import _SDPAMerger |
27 | 34 | from torch.utils.data import Dataset |
28 | 35 | from transformers import AutoProcessor, Trainer, TrainerCallback |
29 | 36 | from transformers.trainer_pt_utils import LabelSmoother |
30 | 37 |
|
| 38 | +import modelopt.torch.speculative.plugins.transformers |
| 39 | +from modelopt.torch.speculative.utils import get_ttt_msk_func |
31 | 40 | from modelopt.torch.utils import print_rank_0 |
32 | 41 | from modelopt.torch.utils.distributed import is_master |
33 | 42 |
|
@@ -566,3 +575,153 @@ def on_step_end(self, args, state, control, **kwargs): |
566 | 575 | except Exception: |
567 | 576 | print_rank_0("AR validation not available.") |
568 | 577 | return control |
| 578 | + |
| 579 | + |
| 580 | +def _compute_ttt_attention_mask(batch_size, seq_length, ttt_step, dtype) -> torch.Tensor: |
| 581 | + """Return TTT attention_mask tensor of type BlockMask or Tensor depends on eagle attn impl.""" |
| 582 | + |
| 583 | + msk_func = get_ttt_msk_func(seq_length, ttt_step) |
| 584 | + |
| 585 | + dtypemin = torch.finfo(dtype).min |
| 586 | + q_len = seq_length |
| 587 | + kv_len = seq_length * (1 + ttt_step) |
| 588 | + # Return tensor mask for non-flex attention |
| 589 | + tensor_mask = msk_func( |
| 590 | + None, |
| 591 | + None, |
| 592 | + torch.arange(q_len).view(1, 1, q_len, 1), |
| 593 | + torch.arange(kv_len).view(1, 1, 1, kv_len), |
| 594 | + ).to(torch.cuda.current_device()) |
| 595 | + tensor_mask = torch.full_like( |
| 596 | + tensor_mask, 0, dtype=dtype, device=torch.cuda.current_device() |
| 597 | + ).masked_fill(~tensor_mask, dtypemin) |
| 598 | + return tensor_mask |
| 599 | + |
| 600 | + |
| 601 | +def get_patched_templated_ring_attn(orig_templated_attn: Callable): |
| 602 | + """ |
| 603 | + Return patched version of |
| 604 | + torch.distributed.tensor.experimental._attention._templated_ring_attention |
| 605 | + to support TTT. |
| 606 | + """ |
| 607 | + |
| 608 | + def _get_sharded_ttt_msk(i, rank, size, q_len, ttt_step, dtype): |
| 609 | + """Get chunk-interleaved TTT mask for current rank. |
| 610 | + e.g.: |
| 611 | + 2 ranks, ttt_step=1; |
| 612 | + full_ttt_mask = [[0, 0, 0, 0, x, 0, 0, 0], |
| 613 | + [x, 0, 0, 0, 0, x, 0, 0], |
| 614 | + [x, x, 0, 0, 0, 0, x, 0], |
| 615 | + [x, x, x, 0, 0, 0, 0, x], |
| 616 | +
|
| 617 | + rank 0, step0: [[0, 0, x, 0], |
| 618 | + [x, 0, 0, x]] |
| 619 | +
|
| 620 | + rank 1, step0: [[0, 0, x, 0], |
| 621 | + [x, 0, 0, x]] |
| 622 | +
|
| 623 | + rank 0, step1: [[0, 0, 0, 0], |
| 624 | + [0, 0, 0, 0]] |
| 625 | +
|
| 626 | + rank 1, step1: [[x, x, 0, 0], |
| 627 | + [x, x, 0, 0]] |
| 628 | +
|
| 629 | + """ |
| 630 | + device = torch.cuda.current_device() |
| 631 | + q_indices = torch.arange(q_len * rank, q_len * (rank + 1), device=device) |
| 632 | + kv_indices = ( |
| 633 | + torch.arange(q_len * size * (ttt_step + 1), device=device) |
| 634 | + .view(ttt_step + 1, size, q_len)[:, (rank - i) % size, :] |
| 635 | + .reshape(-1) |
| 636 | + ) |
| 637 | + msk_func = get_ttt_msk_func(q_len * size, ttt_step) |
| 638 | + attn_mask = msk_func( |
| 639 | + None, |
| 640 | + None, |
| 641 | + q_indices.view(1, 1, -1, 1), |
| 642 | + kv_indices.view(1, 1, 1, -1), |
| 643 | + ) |
| 644 | + attn_bias = torch.where( |
| 645 | + attn_mask, |
| 646 | + torch.zeros((), dtype=dtype, device=attn_mask.device), |
| 647 | + torch.full((), torch.finfo(dtype).min, dtype=dtype, device=attn_mask.device), |
| 648 | + ) |
| 649 | + |
| 650 | + return attn_bias |
| 651 | + |
| 652 | + def patched_templated_attn(*args, **kwargs): |
| 653 | + """Patched version of torch.distributed.tensor.experimental._attention._templated_ring_attention.""" |
| 654 | + # Get original attention op |
| 655 | + # Sensitive to impl of _templated_ring_attention |
| 656 | + original_op = args[2] |
| 657 | + |
| 658 | + # This patch is only enabled for eagle model by context manager, not base model. |
| 659 | + patch_enbabled = modelopt.torch.speculative.plugins.transformers.ENABLE_CP_TTT_PATCH |
| 660 | + |
| 661 | + if patch_enbabled and original_op != torch.ops.aten._scaled_dot_product_cudnn_attention: |
| 662 | + raise ValueError(f"CP TTT only supports cuddn attention now. Got: {original_op}") |
| 663 | + |
| 664 | + # Unset is_causal to use custom attn mask |
| 665 | + if patch_enbabled: |
| 666 | + kwargs["is_causal"] = False |
| 667 | + |
| 668 | + def patched_op(*args, **kwargs): |
| 669 | + # Inpect the parent frame to get current shard info |
| 670 | + # This is sensitive to torch _templated_ring_attention impl |
| 671 | + try: |
| 672 | + frame: FrameType = inspect.currentframe() |
| 673 | + f_back: FrameType = frame.f_back |
| 674 | + rank = f_back.f_locals["rank"] |
| 675 | + size = f_back.f_locals["size"] |
| 676 | + query = f_back.f_locals["query"] |
| 677 | + key = f_back.f_locals["key"] |
| 678 | + i = f_back.f_locals["i"] |
| 679 | + ttt_step = (key.shape[2] // query.shape[2]) - 1 |
| 680 | + except Exception as e: |
| 681 | + print(f"Failed to capture loop variables in patched _templated_ring_attention: {e}") |
| 682 | + # Set attn mask to permuted TTT mask |
| 683 | + if "attn_bias" in kwargs: |
| 684 | + kwargs["attn_bias"] = _get_sharded_ttt_msk( |
| 685 | + i, rank, size, query.shape[2], ttt_step, query.dtype |
| 686 | + ) |
| 687 | + # Perform shard attention |
| 688 | + return original_op(*args, **kwargs) |
| 689 | + |
| 690 | + return orig_templated_attn(args[0], args[1], patched_op, *args[3:], **kwargs) |
| 691 | + |
| 692 | + return patched_templated_attn |
| 693 | + |
| 694 | + |
| 695 | +def patch_ring_attention_for_ttt(): |
| 696 | + """Patch torch ring attention to support context parallelism for TTT.""" |
| 697 | + # Torch Ring Attention only supports no mask or causal mask. We apply the following patches to enable TTT mask. |
| 698 | + |
| 699 | + # 1. Disable load balance, which is designed for causal mask. |
| 700 | + # This affect how buffers are sharded. So need to be done permenantly before accelerate/hf trainer init. |
| 701 | + torch.distributed.tensor.experimental._attention._cp_options.enable_load_balance = False |
| 702 | + |
| 703 | + # 2. Patch templated ring attention for TTT mask. |
| 704 | + original_templated_ring_attention = ( |
| 705 | + torch.distributed.tensor.experimental._attention._templated_ring_attention |
| 706 | + ) |
| 707 | + original_templated_ring_attention_backward = ( |
| 708 | + torch.distributed.tensor.experimental._attention._templated_ring_attention_backward |
| 709 | + ) |
| 710 | + torch.distributed.tensor.experimental._attention._templated_ring_attention = ( |
| 711 | + get_patched_templated_ring_attn(original_templated_ring_attention) |
| 712 | + ) |
| 713 | + torch.distributed.tensor.experimental._attention._templated_ring_attention_backward = ( |
| 714 | + get_patched_templated_ring_attn(original_templated_ring_attention_backward) |
| 715 | + ) |
| 716 | + |
| 717 | + # 3. Patch merger to skip the blank shard to avoid difference in output. |
| 718 | + original_sdpa_merger_step = _SDPAMerger.step |
| 719 | + |
| 720 | + def patched_sdpa_merger_step( |
| 721 | + self, out: torch.Tensor, lse: torch.Tensor, partial: bool |
| 722 | + ) -> torch.Tensor: |
| 723 | + if lse.sum() <= 0: |
| 724 | + return |
| 725 | + return original_sdpa_merger_step(self, out, lse, partial) |
| 726 | + |
| 727 | + _SDPAMerger.step = patched_sdpa_merger_step |
0 commit comments