dnaihao/phi-3-tablellm

1"""2Orginally Taken verbatim from xformers library3https://github.com/facebookresearch/xformers/blob/bcb707576c6a80eaf850aa80e8643d3497ec2bc4/xformers/components/positional_embedding/rotary.py4 5The difference is that xformers seems to assume the inputs to be6(bs, head, seq_len, dim) while we assume (bs, seq_len, head, dim)7 8"""9# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.10#11# This source code is licensed under the BSD license found in the12# LICENSE file in the root directory of this source tree.13 14 15# CREDITS: This implementation is inspired by GPT-NeoX https://github.com/EleutherAI/gpt-neox16# NOTE: Almost the same right now, moving parts to Triton is the next step17 18import math19from typing import List, Optional, Tuple, Dict, Union20 21import torch22import dataclasses23from transformers.utils import logging24 25from transformers import PretrainedConfig26 27is_dacite_available = False28try:29    import dacite30    is_dacite_available = True31except ImportError:32    pass33 34logger = logging.get_logger(__name__)35 36@dataclasses.dataclass37class LongRopeConfig(object):38    short_factor: List[float]39    long_factor: List[float]40    original_max_position_embeddings: int41    type: str = "longrope"42    short_mscale: float = -143    long_mscale: float = -144 45 46    def __post_init__(self):47        assert self.type in ("longrope", "su"), f"Invalid type {self.type} for LongRopeConfig. Expected longrope / su"48 49 50    @classmethod51    def from_dict(cls, config_dict: Dict[str, Union[float, List[float], int]]) -> "LongRopeConfig":52        if is_dacite_available:53            # Preferred since we can also type check the input54            return dacite.from_dict(data_class=cls, data=config_dict)55        kwargs = {}56        for field in dataclasses.fields(cls):57            if field.name in config_dict:58                if field.init:59                    kwargs[field.name] = config_dict[field.name]60                else:61                    raise ValueError(f"Field {field.name} is not initiable")62            else:63                if field.default is dataclasses.MISSING:64                    raise ValueError(f"Field {field.name} is required")65        extra_keys = set(config_dict.keys()) - set(kwargs.keys())66        if len(extra_keys) > 0:67            for key in extra_keys:68                logger.error(f"Unrecognized key {key} in config_dict")69            raise ValueError(f"Unrecognized keys in config_dict")70        return cls(**kwargs)71 72def rotate_half(x):73    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]74    return torch.cat((-x2, x1), dim=x1.ndim - 1)75 76 77 78@torch.jit.script79def apply_rotary_pos_emb(x, cos, sin, seq_dimension: int):80    # NOTE: This could probably be moved to Triton81 82    if seq_dimension == 0:83        cos = cos[: x.shape[0], None, None, :]84        sin = sin[: x.shape[0], None, None, :]85    elif seq_dimension == 1:86        # Handle a possible sequence length mismatch in between q and k87        cos = cos[None, : x.shape[1], None, :]88        sin = sin[None, : x.shape[1], None, :]89    elif seq_dimension == 2:90        cos = cos[None, None, : x.shape[2], :]91        sin = sin[None, None, : x.shape[2], :]92 93    return (x * cos) + (rotate_half(x) * sin)94 95 96 97class RotaryEmbedding(torch.nn.Module):98    """99    Adapted from the xformers library100 101    The rotary position embeddings from RoFormer_ (Su et. al).102    A crucial insight from the method is that the query and keys are103    transformed by rotation matrices which depend on the relative positions.104    Other implementations are available in the Rotary Transformer repo_ and in105    GPT-NeoX_, GPT-NeoX was an inspiration106    .. _RoFormer: https://arxiv.org/abs/2104.09864107    .. _repo: https://github.com/ZhuiyiTechnology/roformer108    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox109    .. warning: Please note that this embedding is not registered on purpose, as it is transformative110        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis111 112    # Arguments113    :param dim_mode: head dimention114    :param max_seq_len:115    :param default_seq_dimension: which dim is the sequence length116    :param dtype: cos/sin dtype117    :param use_fused_kernel: if to use customized fused kernel.118        Note: if used, q, k will be modified inplace. Ok for both forward & backward.119    """120 121    def __init__(122        self,123        dim_model: int,124        *,125        max_seq_len: Optional[int] = None,126        dtype: Optional[torch.dtype] = None,127        base=10000,128        position_scale=1,129        device: Optional[torch.device] = None,130        longrope_config: Optional[LongRopeConfig] = None,131    ):132        super().__init__()133        self.base = base134        self.dim_model = dim_model135        self.max_seq_len = max_seq_len136        self.longrope_config = longrope_config137 138        if self.is_longrope:139            # Keep the maximum range vector, and slice from it as needed140            self.register_buffer(141                "range_vector",142                torch.arange(max_seq_len, device=device, dtype=torch.float32),143                persistent=False144            )145            self.register_buffer(146                "short_factors",147                torch.tensor(self.longrope_config.short_factor, dtype=torch.float32),148                persistent=False149            )150            self.register_buffer(151                "long_factors",152                torch.tensor(self.longrope_config.long_factor, dtype=torch.float32),153                persistent=False154            )155        else:156            # Generate and save the inverse frequency buffer (non trainable)157            inv_freq = 1.0 / (base ** (torch.arange(0, dim_model, 2).float().to(device) / self.dim_model))158            self.register_buffer("inv_freq", inv_freq)159 160        self.position_scale = position_scale161        162        if not self.is_longrope:163            dtype = dtype or torch.get_default_dtype()164            self._set_cos_sin_cache(165                seq_len=max_seq_len,166                device=self.inv_freq.device,167                dtype=dtype,168            )169    @property170    def is_longrope(self):171        return self.longrope_config is not None172 173    @property174    def original_max_seq_len(self):175        if self.longrope_config is not None:176            return self.longrope_config.original_max_position_embeddings177        logger.warning_once(178            (179                "``original_max_seq_len'' is being accessed, but longrope_config has not been set. "180                "Please only do this if you are sure about the context."181            )182        )183        return self.max_seq_len184 185    def get_range_vector(self, seq_len: int, device: torch.device):186        if self.is_longrope:187            assert seq_len < self.range_vector.shape[0], f"Found seq_len {seq_len} greater than max_seq_len {self.range_vector.shape[0]}"188            if self.range_vector.device != device:189                self.range_vector = self.range_vector.to(device)190            return self.range_vector[:seq_len]191        return torch.arange(seq_len, device=device, dtype=torch.float32)192 193 194    def _calc_mscale(self, scale: torch.Tensor) -> torch.Tensor:195        if scale <= 1.0:196            return 1.0197        return math.sqrt(1 + math.log(scale) / math.log(self.original_max_seq_len))198 199    def _set_cos_sin_cache(200        self,201        seq_len: int,202        device: Optional[torch.device] = None,203        dtype: Optional[torch.dtype] = None,204    ) -> None:205        dtype = dtype or torch.get_default_dtype()206        self.max_seq_len_cached = seq_len207        t = (torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32) * self.position_scale).type_as(self.inv_freq)208        device_type = device.type if device is not None else "cpu"209        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"210        with torch.autocast(device_type=device_type, enabled=False):211            # shape: (seq_len, dim_model // 2)212            freqs = torch.outer(t, self.inv_freq)213            # shape: (seq_len, dim_model)214            emb = torch.cat((freqs, freqs), dim=-1)215            cos = emb.cos()216            sin = emb.sin()217        self.register_buffer("cos_cached", cos.to(dtype), persistent=False)218        self.register_buffer("sin_cached", sin.to(dtype), persistent=False)219 220    def forward(221        self, q: torch.Tensor,222        k: torch.Tensor,223        seq_dimension: int = 1,224        seqlen_offset: int = 0,225    ) -> Tuple[torch.Tensor, torch.Tensor]:226        """q, k does not include `seqlen_offset`227        q: Either (bs, seq_len, num_heads, head_dim) or (seq_len, bs, num_heads, head_dim)228        k: Either (bs, seq_len, num_heads, head_dim) or (seq_len, bs, num_heads, head_dim)229        """230        if seq_dimension < 0:231            seq_dimension = k.ndim + seq_dimension232        assert seq_dimension in (0, 1, 2)233        seq_len = k.shape[seq_dimension] + seqlen_offset234 235        if self.is_longrope:236            if seq_len > self.original_max_seq_len:237                t = self.get_range_vector(seq_len, device=q.device)238                rescale_factors = self.long_factors.to(q.device)239                long_mscale = self.longrope_config.long_mscale240                mscale = long_mscale if long_mscale > 0 else self._calc_mscale(self.max_seq_len / self.original_max_seq_len)241            else:242                t = self.get_range_vector(self.original_max_seq_len, device=q.device)243                rescale_factors = self.short_factors.to(q.device)244                short_mscale = self.longrope_config.short_mscale245                mscale = short_mscale if short_mscale > 0 else 1.0246            assert rescale_factors.shape == (self.dim_model // 2, ), (247                f"misaligned shape for LongRoPE rescale factors:\n"248                f"\tExpected {(self.dim_model // 2, )}, got {rescale_factors.shape}."249            )250            inv_freq = 1.0 / (rescale_factors * (self.base ** (torch.arange(0, self.dim_model, 2).float().to(q.device) / self.dim_model)))251            device_type = q.device.type if q.device is not None else "cpu"252            device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"253            with torch.autocast(device_type=device_type, enabled=False):254                freqs = torch.outer(t, inv_freq)255                emb = torch.cat((freqs, freqs), dim=-1)256                cos = emb.cos() * mscale257                sin = emb.sin() * mscale258            cos_cached = cos.to(q.dtype)259            sin_cached = sin.to(q.dtype)260        else:261            if seq_len > self.max_seq_len_cached:262                self._set_cos_sin_cache(263                    seq_len=seq_len,264                    device=k.device,265                    dtype=k.dtype,266                )267            cos_cached = self.cos_cached268            sin_cached = self.sin_cached269        return (270            apply_rotary_pos_emb(271                q, cos_cached[seqlen_offset:seq_len], sin_cached[seqlen_offset:seq_len], seq_dimension=seq_dimension272            ).to(q.dtype),273            apply_rotary_pos_emb(274                k, cos_cached[seqlen_offset:seq_len], sin_cached[seqlen_offset:seq_len], seq_dimension=seq_dimension275            ).to(k.dtype),276        )277 278    @classmethod279    def from_config(cls, config: PretrainedConfig) -> "RotaryEmbedding":280        kwargs = dict(281            dim_model=config.hidden_size // config.num_attention_heads,282            max_seq_len=config.max_position_embeddings,283            base=config.rope_embedding_base,284            position_scale=config.rope_position_scale,285        )286        if config.rope_scaling is not None:287            kwargs["longrope_config"] = LongRopeConfig.from_dict(config.rope_scaling)288        return cls(**kwargs)289