File size: 1,281 Bytes
c857c8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from importlib import import_module
from typing import Tuple
import torch
import transformers
from torch import nn
from torch.nn import functional as F
__all__ = ["patch"]
def _get_unpad_data(attention_mask: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, int]:
if hasattr(_get_unpad_data, "seqlens_in_batch"):
seqlens_in_batch = _get_unpad_data.seqlens_in_batch
else:
seqlens_in_batch = torch.sum(attention_mask, dim=1)
indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
max_seqlen_in_batch = seqlens_in_batch.max().item()
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
return indices, cu_seqlens, max_seqlen_in_batch
def set_seqlens_in_batch(seqlens_in_batch: torch.Tensor) -> None:
_get_unpad_data.seqlens_in_batch = seqlens_in_batch
def patch(model: nn.Module) -> None:
if transformers.__version__ < "4.43.0":
m = import_module(model.__module__)
if not hasattr(m, "_get_unpad_data"):
raise ValueError(f"Module {m} does not have function '_get_unpad_data' for packing")
m._get_unpad_data = _get_unpad_data
else:
transformers.modeling_flash_attention_utils._get_unpad_data = _get_unpad_data
|