diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..7d7885e34632cdfd5a155ab710a36602c96b517d 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +tb/20260114-0839/wandb/run-20260114_083920--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834/run--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834.wandb filter=lfs diff=lfs merge=lfs -text diff --git a/fla/__pycache__/__init__.cpython-312.pyc b/fla/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5aec7a8c1439875c1037fbbabf1de8a9df33ec6f Binary files /dev/null and b/fla/__pycache__/__init__.cpython-312.pyc differ diff --git a/fla/__pycache__/utils.cpython-312.pyc b/fla/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ae9fcd5ee831183e8d52e8b1d7f5930ba928969 Binary files /dev/null and b/fla/__pycache__/utils.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/__init__.cpython-312.pyc b/fla/layers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a085a0cc11e6c0d5bf8f694c1983f8821188347d Binary files /dev/null and b/fla/layers/__pycache__/__init__.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/abc.cpython-312.pyc b/fla/layers/__pycache__/abc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37be67749e7fab100470ecb80bafb3a65c1f6052 Binary files /dev/null and b/fla/layers/__pycache__/abc.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/attn.cpython-312.pyc b/fla/layers/__pycache__/attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15d43a9a803954a5aa87df709d4d298e80081e33 Binary files /dev/null and b/fla/layers/__pycache__/attn.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/based.cpython-312.pyc b/fla/layers/__pycache__/based.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fba56fbd1cf5c5778abaceee9c0dfe4474cd1c5 Binary files /dev/null and b/fla/layers/__pycache__/based.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/bitattn.cpython-312.pyc b/fla/layers/__pycache__/bitattn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb588ef9407464e5dd8ea4520abdc778c001fa26 Binary files /dev/null and b/fla/layers/__pycache__/bitattn.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/delta_net.cpython-312.pyc b/fla/layers/__pycache__/delta_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1f38ce871aed2c8aa8b9dcd76c0a3339fd7cd89 Binary files /dev/null and b/fla/layers/__pycache__/delta_net.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/forgetting_attn.cpython-312.pyc b/fla/layers/__pycache__/forgetting_attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6e458ff4e4d484a544f7b35a03da502ea81b721 Binary files /dev/null and b/fla/layers/__pycache__/forgetting_attn.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/gated_deltanet.cpython-312.pyc b/fla/layers/__pycache__/gated_deltanet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e00371cbea06eaf7860869e918fa0e29aff1a494 Binary files /dev/null and b/fla/layers/__pycache__/gated_deltanet.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/gated_deltaproduct.cpython-312.pyc b/fla/layers/__pycache__/gated_deltaproduct.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..def3475b90d93f5a5dc5ed1d5db6cbcfee5d1693 Binary files /dev/null and b/fla/layers/__pycache__/gated_deltaproduct.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/gsa.cpython-312.pyc b/fla/layers/__pycache__/gsa.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..945caaddde554dd7bb8a4d1e8f4f8e279d48f53f Binary files /dev/null and b/fla/layers/__pycache__/gsa.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/hgrn2.cpython-312.pyc b/fla/layers/__pycache__/hgrn2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dc9b498e646c8eaccb839f334fea5ba588f3e34 Binary files /dev/null and b/fla/layers/__pycache__/hgrn2.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/lightnet.cpython-312.pyc b/fla/layers/__pycache__/lightnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36ade4acf1682bcce87cd4161b35652abd087431 Binary files /dev/null and b/fla/layers/__pycache__/lightnet.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/linear_attn.cpython-312.pyc b/fla/layers/__pycache__/linear_attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..678825e48bbd65a0413b64ba8bcdb37ba72cb6d2 Binary files /dev/null and b/fla/layers/__pycache__/linear_attn.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/nsa.cpython-312.pyc b/fla/layers/__pycache__/nsa.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7f2e50afdad2c0aebb67a651335019c49423eae Binary files /dev/null and b/fla/layers/__pycache__/nsa.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/rebased.cpython-312.pyc b/fla/layers/__pycache__/rebased.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59b9c9acb3fb3fef9a7bd8b90a832eef333352de Binary files /dev/null and b/fla/layers/__pycache__/rebased.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/rwkv6.cpython-312.pyc b/fla/layers/__pycache__/rwkv6.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e3430b4f7966e78d20d00fe5b8aea2c8c321ec3 Binary files /dev/null and b/fla/layers/__pycache__/rwkv6.cpython-312.pyc differ diff --git a/fla/layers/__pycache__/rwkv7.cpython-312.pyc b/fla/layers/__pycache__/rwkv7.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c31ff317f5eda688e8055a5d515571c54a5fef53 Binary files /dev/null and b/fla/layers/__pycache__/rwkv7.cpython-312.pyc differ diff --git a/fla/layers/abc.py b/fla/layers/abc.py new file mode 100644 index 0000000000000000000000000000000000000000..3afc3ebd015f4414edf26474f3f0a98b0d18e0e4 --- /dev/null +++ b/fla/layers/abc.py @@ -0,0 +1,218 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.modules import FusedRMSNormGated, RMSNorm, RotaryEmbedding, ShortConvolution +from fla.modules.activations import swiglu, swish +from fla.ops.abc.chunk import chunk_abc + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class ABCAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 1024, + expand_k: float = 0.5, + expand_v: float = 1.0, + num_heads: int = 4, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_slots: Optional[int] = None, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + gate_low_rank_dim: int = 16, + gate_logit_normalizer: int = 16, + use_rope: bool = True, + use_input_gate: bool = False, + use_output_gate: bool = True, + use_norm: bool = True, + clamp_min: Optional[float] = -32, + clamp_max: Optional[float] = 32, + layer_idx: Optional[int] = None, + **kwargs + ) -> ABCAttention: + super().__init__() + + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.key_dim = int(self.hidden_size * self.expand_k) + self.value_dim = int(self.hidden_size * self.expand_v) + self.head_k_dim = self.key_dim // self.num_heads + self.head_v_dim = self.value_dim // self.num_heads + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.gate_low_rank_dim = gate_low_rank_dim + self.gate_logit_normalizer = gate_logit_normalizer + + self.use_rope = use_rope + self.use_input_gate = use_input_gate + self.use_output_gate = use_output_gate + self.use_norm = use_norm + + if num_slots is None: + num_slots = self.head_k_dim + self.num_slots = num_slots + + self.norm_eps = norm_eps + + self.clamp_min = clamp_min + self.clamp_max = clamp_max + self.layer_idx = layer_idx + + if layer_idx is None: + warnings.warn( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False) + + if use_output_gate: + self.g_proj = nn.Linear(self.hidden_size, self.value_dim, bias=False) + self.s_proj = nn.Linear(self.hidden_size, self.num_heads * self.num_slots, bias=False) + self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.k_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.v_conv1d = ShortConvolution(self.value_dim, conv_size, activation='silu') + + if self.use_norm: + if self.use_output_gate: + self.g_norm = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + else: + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + + if self.use_rope: + self.rotary = RotaryEmbedding(self.head_k_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if cu_seqlens is not None: + raise NotImplementedError("Training with cu_seqlens is not supported yet for ABCAttention") + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + if self.use_input_gate: + q, k, v = map(lambda x: swish(x), (q, k, v)) + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + if self.use_rope: + seqlen_offset = 0 + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset) + + s = rearrange(self.s_proj(hidden_states), '... (h m) -> ... h m', m=self.num_slots) + s = s.clamp_(self.clamp_min, self.clamp_max) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + o, recurrent_state = chunk_abc( + q=q, + k=k, + v=v, + s=s, + initial_state=recurrent_state, + output_final_state=use_cache, + head_first=False + ) + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + if self.use_norm and not self.use_output_gate: + o = self.g_norm(o) + elif self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.g_norm(o, g) if self.use_norm else swiglu(g, o) + o = rearrange(o, '... h d -> ... (h d)') + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, seq_len: int = 2048): + return 2 * self.num_slots * self.hidden_size diff --git a/fla/layers/attn.py b/fla/layers/attn.py new file mode 100644 index 0000000000000000000000000000000000000000..2d30743ee5f39bccfdd5d62366a718bbf44da0d3 --- /dev/null +++ b/fla/layers/attn.py @@ -0,0 +1,240 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange +from transformers.utils import logging + +from fla.modules import RMSNorm, RotaryEmbedding + +if TYPE_CHECKING: + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`." + " Falling back to use SDPA's attention implementation.", + category=ImportWarning + ) + flash_attn_func = None + +import os +if os.getenv("FLASH_ATTENTION_DISABLE", "0") == "1": + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +class Attention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: Optional[int] = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: Optional[int] = None, + rope_theta: Optional[float] = 10000., + max_position_embeddings: Optional[int] = None, + layer_idx: int = None + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + # equivalent to cu_seqlens in `flash_attn` + cu_seqlens = kwargs.get('cu_seqlens', None) + + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size) + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + # Contains at least one padding token in the sequence + if attention_mask is not None: + if flash_attn_func is not None: + q, k, v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(q, k, v, attention_mask, q_len) + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ) # B S H D + o = pad_input(o, indices_q, batch_size, q_len) # B S H D + else: + attention_mask = attention_mask.bool() + q = rearrange(q, 'b s h d -> b h s d') # B H S D + k = rearrange(k, 'b s h d -> b h s d') # B H S D + v = rearrange(v, 'b s h d -> b h s d') # B H S D + o = F.scaled_dot_product_attention( + q, k, v, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ) # B, H, S, D + o = rearrange(o, 'b h s d -> b s h d') + elif cu_seqlens is not None: + if flash_attn_func is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ).unsqueeze(0) + else: + # o = F.scaled_dot_product_attention( + # q.squeeze(0), k.squeeze(0), v.squeeze(0), + # dropout_p=0.0, + # is_causal=True + # ).unsqueeze(0) + raise NotImplementedError( + "SDPA does not support variable length inputs with cu_seqlens. " + "Please use flash_attn_func for variable length inputs." + ) + else: + if flash_attn_func is not None: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ) + else: + q = rearrange(q, 'b s h d -> b h s d') # B H S D + k = rearrange(k, 'b s h d -> b h s d') # B H S D + v = rearrange(v, 'b s h d -> b h s d') # B H S D + o = F.scaled_dot_product_attention( + q, k, v, + dropout_p=0.0, + is_causal=True + ) + o = rearrange(o, 'b h s d -> b s h d') + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values + + def _upad_input(self, q, k, v, attention_mask, q_len): + batch_size, seq_len, num_key_value_heads, head_dim = k.shape + cache_mask = attention_mask[:, -seq_len:] + seqlens = cache_mask.sum(-1, dtype=torch.int32) + indices_k = torch.nonzero(cache_mask.flatten(), as_tuple=False).flatten() + max_seqlen_k = seqlens.max().item() + cu_seqlens_k = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + + k = index_first_axis(k.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k) + v = index_first_axis(v.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k) + if q_len == seq_len: + q = index_first_axis(q.reshape(batch_size * seq_len, self.num_heads, head_dim), indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_q = max_seqlen_k + indices_q = indices_k + elif q_len == 1: + max_seqlen_q = 1 + # There is a memcpy here, that is very bad. + cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) + indices_q = cu_seqlens_q[:-1] + q = q.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -q_len:] + q, indices_q, cu_seqlens_q, max_seqlen_q, *_ = unpad_input(q, attention_mask) + + return q, k, v, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) diff --git a/fla/layers/based.py b/fla/layers/based.py new file mode 100644 index 0000000000000000000000000000000000000000..64a12364d246d7eeedcf40ac39ea21a6a4312971 --- /dev/null +++ b/fla/layers/based.py @@ -0,0 +1,96 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +Linear attention in Based. +https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py +""" + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.modules.feature_map import TaylorFeatureMap +from fla.ops.based import parallel_based +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn + + +class BasedLinearAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + feature_dim: int = 16, + num_key_value_heads: int = 12, + num_heads: int = 12, + feature_name: str = "taylor_exp", + eps: float = 1e-12, + causal: bool = True, + mode: str = "parallel", + ): + super().__init__() + + self.hidden_size = hidden_size + self.mode = mode + self.feature_name = feature_name + self.feature_dim = feature_dim + self.num_key_value_heads = num_key_value_heads + self.num_heads = num_heads + self.head_dim = self.hidden_size // self.num_key_value_heads + assert self.hidden_size % self.head_dim == 0 + self.causal = causal + + self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.dropout = nn.Identity() + self.feature_map = TaylorFeatureMap(feature_dim) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor, **kwargs): + mode = self.mode + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v]) + if mode == "fused_chunk": + q, k = self.feature_map(q), self.feature_map(k) + o, _ = fused_chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False) + elif mode == 'chunk': + q, k = self.feature_map(q), self.feature_map(k) + o, _ = chunk_linear_attn(q, k, v, normalize=True, scale=1, head_first=False) + elif mode == 'parallel': + assert q.shape[-1] <= 128 + o = parallel_based(q, k, v, scale=1, use_norm=True, head_first=False) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + o = self.dropout(o) + return o + + # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119 + + def forward_reference(self, hidden_states: torch.Tensor, filters: torch.Tensor = None, *args, **kwargs): + """ + x (torch.Tensor): tensor of shape (b, d, t) + y (torch.Tensor): tensor of shape (b, d, t) + """ + # hidden_states = hidden_states.transpose(1, 2) + b, t, _ = hidden_states.size() + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + + q = q.view(b, t, self.num_heads, self.feature_dim).transpose(1, 2) + k = k.view(b, t, self.num_key_value_heads, self.feature_dim).transpose(1, 2) + v = v.view(b, t, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + # Linear attention + q, k = self.feature_map(q), self.feature_map(k) + q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1) + + # Compute attention + if self.causal: + y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps)) + else: + y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps)) + y = rearrange(y, 'b h t d -> b t (h d)') + y = self.o_proj(y.to(hidden_states.dtype)) + y = self.dropout(y) + return y.to(hidden_states.dtype) diff --git a/fla/layers/bitattn.py b/fla/layers/bitattn.py new file mode 100644 index 0000000000000000000000000000000000000000..f797b164963161f7b292b844e06e16527ea0817a --- /dev/null +++ b/fla/layers/bitattn.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange +from transformers.utils import logging + +from fla.modules import RotaryEmbedding +from fla.modules.fused_bitlinear import FusedBitLinear + +if TYPE_CHECKING: + from fla.models.utils import Cache + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning + ) + flash_attn_func = None + +logger = logging.get_logger(__name__) + + +class BitAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: Optional[int] = None, + window_size: Optional[int] = None, + rope_theta: Optional[float] = 10000., + max_position_embeddings: Optional[int] = None, + norm_eps: float = 1e-5, + layer_idx: int = None + ): + super().__init__() + + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.hidden_size = hidden_size + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False) + self.k_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False) + self.v_proj = FusedBitLinear(self.hidden_size, self.kv_dim, bias=False) + self.o_proj = FusedBitLinear(self.hidden_size, self.hidden_size, bias=False) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + + # equivalent to cu_seqlens in `flash_attn` + cu_seqlens = kwargs.get('cu_seqlens', None) + + seqlen_offset, max_seqlen = 0, q_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=q_len, + cache_kwargs=dict(window_size=self.window_size) + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if flash_attn_func is None: + raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") + + # Contains at least one padding token in the sequence + if attention_mask is not None: + q, k, v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(q, k, v, attention_mask, q_len) + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_q, max_seqlen_k = max_seq_lens + o = flash_attn_varlen_func( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ) + o = pad_input(o, indices_q, batch_size, q_len) + elif cu_seqlens is not None: + o = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ).unsqueeze(0) + else: + o = flash_attn_func( + q, k, v, + causal=True, + window_size=(-1, -1) if self.window_size is None else (self.window_size-1, 0) + ) + o = o.reshape(batch_size, q_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values + + def _upad_input(self, q, k, v, attention_mask, q_len): + batch_size, seq_len, num_key_value_heads, head_dim = k.shape + cache_mask = attention_mask[:, -seq_len:] + seqlens = cache_mask.sum(-1, dtype=torch.int32) + indices_k = torch.nonzero(cache_mask.flatten(), as_tuple=False).flatten() + max_seqlen_k = seqlens.max().item() + cu_seqlens_k = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + + k = index_first_axis(k.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k) + v = index_first_axis(v.reshape(batch_size * seq_len, num_key_value_heads, head_dim), indices_k) + if q_len == seq_len: + q = index_first_axis(q.reshape(batch_size * seq_len, self.num_heads, head_dim), indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_q = max_seqlen_k + indices_q = indices_k + elif q_len == 1: + max_seqlen_q = 1 + # There is a memcpy here, that is very bad. + cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) + indices_q = cu_seqlens_q[:-1] + q = q.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -q_len:] + q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask) + + return q, k, v, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) diff --git a/fla/layers/delta_net.py b/fla/layers/delta_net.py new file mode 100644 index 0000000000000000000000000000000000000000..5190902eec586f0bb296995163eee0a98fe2cbfe --- /dev/null +++ b/fla/layers/delta_net.py @@ -0,0 +1,291 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +def elu_p1(x): + return (F.elu(x, 1., False) + 1.).to(x) + + +def sum_norm(x): + return (x / x.sum(-1, keepdim=True)).to(x) + + +class DeltaNet(nn.Module): + r""" + The layer implementaion for [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://arxiv.org/abs/2406.06484). # noqa: + DeltaNet was originally proposed in [Linear Transformers Are Secretly Fast Weight Programmers](https://arxiv.org/abs/2102.11174). # noqa + + Args: + mode (str, Optional): + Which DeltaNet kernel to use. + Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + use_beta (bool, Optional): + Whether to use beta. Default: `True`. + use_gate (bool, Optional): + Whether to use output gate. Default: `False`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + qk_activation (str, Optional): + The activation function for the query and key. Default: `silu`. + qk_norm (str, Optional): + The normalization method for the query and key. Default: `l2`. + """ + + def __init__( + self, + mode: str = 'chunk', + d_model: int = None, + hidden_size: int = 1024, + expand_k: float = 1.0, + expand_v: float = 1.0, + num_heads: int = 4, + use_beta: bool = True, + use_gate: bool = False, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + allow_neg_eigval: bool = False, + layer_idx: int = None, + qk_activation: str = 'silu', + qk_norm: str = 'l2', + norm_eps: float = 1e-5, + **kwargs + ) -> DeltaNet: + super().__init__() + + self.mode = mode + self.qk_activation = qk_activation + self.qk_norm = qk_norm + + assert self.qk_activation in ['silu', 'relu', 'elu', 'identity'] + assert self.qk_norm in ['l2', 'sum'] + + if d_model is not None: + hidden_size = d_model + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.allow_neg_eigval = allow_neg_eigval + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + self.layer_idx = layer_idx + + self.silu = nn.SiLU() + if mode == 'fused_chunk': + raise NotImplementedError("fused_chunk_delta_rule is now deprecated. Please use `chunk_delta_rule` instead.") + assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + self.use_beta = use_beta + if self.use_beta: + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation='silu' if qk_activation == 'silu' else None + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation='silu' if qk_activation == 'silu' else None + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + activation='silu' + ) + else: + raise UserWarning( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing." + ) + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps) + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # change to inference mode. + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + if self.qk_activation == 'silu': + q, k = self.silu(q), self.silu(k) + v = self.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + if self.qk_activation != 'silu': + if self.qk_activation == 'relu': + q, k = q.relu(), k.relu() + elif self.qk_activation == 'elu': + q, k = elu_p1(q), elu_p1(k) + elif self.qk_activation == 'identity': + pass + else: + raise NotImplementedError + + if self.qk_norm == 'sum': + q = sum_norm(q).to(q) + k = sum_norm(k).to(k) + + if self.use_beta: + beta = self.b_proj(hidden_states).sigmoid() + else: + beta = q.new_ones(q.shape[0], q.shape[1], q.shape[2]) + + if self.allow_neg_eigval: + beta = beta * 2. + + # dealing with padding + if attention_mask is not None: + beta = beta.mul(attention_mask[:, -beta.shape[-2]:, None]) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False + ) + elif mode == 'chunk': + o, recurrent_state = chunk_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True if self.qk_norm == 'l2' else False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + if self.use_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + + return o, None, past_key_values diff --git a/fla/layers/forgetting_attn.py b/fla/layers/forgetting_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..fc2ddc30ef531beedefd0f65836f4148771bcb7d --- /dev/null +++ b/fla/layers/forgetting_attn.py @@ -0,0 +1,109 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange +from transformers.utils import logging + +from fla.modules import GroupNorm +from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +logger = logging.get_logger(__name__) + + +class ForgettingAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 32, + num_kv_heads: Optional[int] = None, + qkv_bias: bool = False, + qk_norm: bool = False, + window_size: Optional[int] = None, + use_output_gate: bool = False, + layer_idx: int = None + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = self.hidden_size // self.num_heads + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + self.qk_norm = qk_norm + + self.window_size = window_size + self.use_output_gate = use_output_gate + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.f_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True) + + if use_output_gate: + self.g_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + if qk_norm: + self.q_norm = GroupNorm( + num_groups=self.num_heads, + hidden_size=self.hidden_size, + is_rms_norm=True, + ) + self.k_norm = GroupNorm( + num_groups=self.num_kv_heads, + hidden_size=self.kv_dim, + is_rms_norm=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + cu_seqlens = kwargs.get('cu_seqlens', None) + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + f = F.logsigmoid(self.f_proj(hidden_states).float()) + if self.qk_norm: + q, k = self.q_norm(q), self.k_norm(k) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + o = parallel_forgetting_attn(q, k, v, f, cu_seqlens=cu_seqlens) + o = rearrange(o, '... h d -> ... (h d)') + if self.use_output_gate: + o = self.g_proj(hidden_states).sigmoid() * o + o = self.o_proj(o) + + return o, None, past_key_values diff --git a/fla/layers/gated_deltanet.py b/fla/layers/gated_deltanet.py new file mode 100644 index 0000000000000000000000000000000000000000..01bc3385cb7d40e35aee2fc693387d1bf1496c22 --- /dev/null +++ b/fla/layers/gated_deltanet.py @@ -0,0 +1,293 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +@torch.compile +def elu_p1(x): + return (F.elu(x, 1., False) + 1.).to(x) + + +@torch.compile +def sum_norm(x): + return (x / x.sum(-1, keepdim=True)).to(x) + + +class GatedDeltaNet(nn.Module): + """ + The layer implementaion for [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://arxiv.org/abs/2412.06464). # noqa + + Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters. + + Parameter alloation when use_gate=True: + - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each + - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each + - Others are ignorably small. + - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size + NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim. + + Parameter allocation when use_gate=False: + - 1 * hidden_size * hidden_size for the q_proj and k_proj each + - 2 * hidden_size * hidden_size for the v_proj and o_proj each + - Others are ignorably small. + - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 4. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_beta (bool, Optional): + Whether to use beta. Default: `True`. + use_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + mode: str = 'chunk', + use_gate: bool = True, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs + ) -> GatedDeltaNet: + super().__init__() + + self.mode = mode + + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + + self.key_dim = int(self.num_heads * self.head_dim) + self.value_dim = int(self.key_dim * self.expand_v) + self.head_k_dim = head_dim + self.head_v_dim = int(head_dim * self.expand_v) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.key_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.key_dim * expand_v}, which is invalid for nn.Linear." + ) + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated." + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + + A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation='silu' + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation='silu' + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + activation='silu' + ) + else: + raise UserWarning( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing." + ) + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + beta = self.b_proj(hidden_states).sigmoid() + g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) + + # dealing with padding + if attention_mask is not None: + beta = beta.mul(attention_mask[:, -beta.shape[-2]:, None]) + g = g.mul(attention_mask[:, -g.shape[-2]:, None]) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True + ) + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + if self.use_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + + return o, None, past_key_values diff --git a/fla/layers/gated_deltaproduct.py b/fla/layers/gated_deltaproduct.py new file mode 100644 index 0000000000000000000000000000000000000000..77cad9a1143614ebb55233030f63f4c8b80c5a53 --- /dev/null +++ b/fla/layers/gated_deltaproduct.py @@ -0,0 +1,351 @@ +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution +from fla.ops.delta_rule import chunk_delta_rule +from fla.ops.gated_delta_rule import chunk_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +def elu_p1(x): + return (F.elu(x, 1.0, False) + 1.0).to(x) + + +def sum_norm(x): + return (x / x.sum(-1, keepdim=True)).to(x) + + +def interleave_multiple_sequences(*sequences): + """ + Interleave multiple sequences together. + For example, with sequences [A1, A2], [B1, B2], [C1, C2], + returns [A1, B1, C1, A2, B2, C2] + """ + if isinstance(sequences[0], (list, tuple)): + sequences = sequences[0] + + if len(sequences) == 1: + return sequences[0] + + # All sequences should have the same shape + assert all(s.shape == sequences[0].shape for s in sequences) + + # Get the original shape + batch_size, seq_len, *rest = sequences[0].shape + + # Stack sequences along a new dimension + stacked = torch.stack(sequences, dim=2) + + # Reshape to interleave + reshaped = stacked.view(batch_size, seq_len * len(sequences), *rest) + + return reshaped + + +class GatedDeltaProduct(nn.Module): + """ + Generalized version of GatedDoubleDeltaNet that supports arbitrary number of householder transformations. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_householder: int = 2, # New parameter for number of householder transformations + mode: str = "chunk", + use_gate: bool = True, + use_forget_gate: bool = True, # when true Gated DeltaProduct, when false DeltaProduct + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int | None = None, + norm_eps: float = 1e-5, + allow_neg_eigval: bool = False, # when true (Gated) DeltaProduct [-1, 1], when false (Gated) DeltaProduct [0, 1] + **kwargs, + ) -> None: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.head_dim = head_dim + self.num_heads = num_heads + self.num_householder = num_householder + self.allow_neg_eigval = allow_neg_eigval + self.use_forget_gate = use_forget_gate + self.key_dim = self.num_heads * self.head_dim + self.value_dim = int(self.key_dim * self.expand_v) + self.head_qk_dim = head_dim + self.head_v_dim = int(head_dim * self.expand_v) + self.layer_idx = layer_idx + self.silu = nn.SiLU() + assert mode in ["chunk", "fused_recurrent"], f"Not supported mode `{mode}`." + # Create multiple projection layers for each householder transformation + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + + self.k_projs = nn.ModuleList( + [ + nn.Linear(hidden_size, self.key_dim, bias=False) + for _ in range(num_householder) + ] + ) + self.v_projs = nn.ModuleList( + [ + nn.Linear(hidden_size, self.value_dim, bias=False) + for _ in range(num_householder) + ] + ) + self.b_projs = nn.ModuleList( + [ + nn.Linear(hidden_size, self.num_heads, bias=False) + for _ in range(num_householder) + ] + ) + if use_short_conv: + self.q_conv1ds = nn.ModuleList( + [ + ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation="silu", + ) + for _ in range(num_householder) + ] + ) + self.k_conv1ds = nn.ModuleList( + [ + ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + activation="silu", + ) + for _ in range(num_householder) + ] + ) + self.v_conv1ds = nn.ModuleList( + [ + ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + activation="silu", + ) + for _ in range(num_householder) + ] + ) + + if self.use_forget_gate: + self.a_proj = nn.Linear(hidden_size, self.num_heads, bias=False) + A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(0, 16) + A_log = torch.log(A) + self.A_log = nn.Parameter(A_log) + self.A_log._no_weight_decay = True + + # Initialize dt parameters + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min) + ) + dt = torch.clamp(dt, min=dt_init_floor) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + self.dt_bias._no_weight_decay = True + + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormSwishGate(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + self.k_id = torch.nn.Identity() + self.apply(self._initialize_weights) + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5) + if module.bias is not None: + nn.init.zeros_(module.bias) + module._is_hf_initialized = True + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding)." + ) + + mode = ( + "chunk" # 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + ) + if self.training: + assert mode == "chunk", "Only chunk mode is supported in training." + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + # Process each householder transformation + ks, vs, betas = [], [], [] + conv_states = [] + + for i in range(self.num_householder): + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"][ + i + ] + conv_mask = ( + attention_mask[:, -hidden_states.shape[1]:] + if attention_mask is not None + else None + ) + + k, conv_state_k = self.k_conv1ds[i]( + x=self.k_projs[i](hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + ) + v, conv_state_v = self.v_conv1ds[i]( + x=self.v_projs[i](hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + ) + conv_states.append((conv_state_q, conv_state_k, conv_state_v)) + else: + k = self.silu(self.k_projs[i](hidden_states)) + v = self.silu(self.v_projs[i](hidden_states)) + + ks.append(k) + vs.append(v) + + beta = self.b_projs[i]( + hidden_states + ).sigmoid() # bs, sequence_length, num_heads + if attention_mask is not None: + beta = beta.mul(attention_mask[:, -hidden_states.shape[1]:, None]) + if self.allow_neg_eigval: + beta = beta * 2 + betas.append(beta) + + if self.use_short_conv: + q, conv_state_q = self.q_conv1ds[0]( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + ) + else: + q = self.silu(self.q_proj(hidden_states)) + q = interleave_multiple_sequences( + [torch.zeros_like(q)] * (self.num_householder - 1) + [q] + ) + # Interleave all sequences + k = interleave_multiple_sequences(ks) + v = interleave_multiple_sequences(vs) + beta = interleave_multiple_sequences(betas) + + q, k, v = ( + rearrange(x, "b t (h d) -> b t h d", h=self.num_heads) for x in (q, k, v) + ) + + recurrent_state = ( + last_state["recurrent_state"] if last_state is not None else None + ) + offsets = kwargs.get("offsets") + + if mode == "chunk": + if self.use_forget_gate: + g = -self.A_log.float().exp() * F.softplus( + self.a_proj(hidden_states).float() + self.dt_bias + ) + if attention_mask is not None: + g = g.mul(attention_mask[:, -g.shape[-2]:, None]) + + # Interleave g with zeros for non-first transformations + g = interleave_multiple_sequences( + [g] + [torch.zeros_like(g)] * (self.num_householder - 1) + ) + + o, recurrent_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=offsets, + head_first=False, + use_qk_l2norm_in_kernel=True + ) + else: + o, recurrent_state = chunk_delta_rule( + q=q, + k=k, + v=v, + beta=beta, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=offsets, + head_first=False, + use_qk_l2norm_in_kernel=True + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + # Take every nth element for n householder transformations + o = o[:, self.num_householder - 1:: self.num_householder, :] + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=conv_states if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[2], + ) + + if self.use_gate: + g = rearrange( + self.g_proj(hidden_states), + "... (h d) -> ... h d", + h=self.num_heads, + ) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + + return o, None, past_key_values diff --git a/fla/layers/gla.py b/fla/layers/gla.py new file mode 100644 index 0000000000000000000000000000000000000000..54789157d6be22eae0a57cf7bc236fc08aec5e92 --- /dev/null +++ b/fla/layers/gla.py @@ -0,0 +1,294 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + + +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.activations import ACT2FN +from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedLinearAttention(nn.Module): + r""" + The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635). # noqa + + Args: + mode (str, Optional): + Which GLA kernel to use. + Currently available: `chunk`, `fused_recurrent`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 0.5. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + gate_logit_normalizer (int, Optional): + The normalizer for the gate logits, appied after `logsigmoid`. Default: 16. + gate_low_rank_dim (int, Optional): + The low rank dim for the gate projection. Default: 16. + clamp_min (float, Optional): + The minimum value for the gate logits. Default: None. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 0.5, + expand_v: float = 1.0, + num_heads: int = 4, + num_kv_heads: Optional[int] = None, + feature_map: Optional[str] = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + use_output_gate: bool = True, + gate_fn: str = 'swish', + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 16, + gate_low_rank_dim: int = 16, + clamp_min: Optional[float] = None, + fuse_norm: bool = True, + layer_idx: int = None, + ) -> GatedLinearAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.use_output_gate = use_output_gate + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.clamp_min = clamp_min + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + if self.use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu') + self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu') + + self.gk_proj = nn.Sequential(nn.Linear(hidden_size, gate_low_rank_dim, bias=False), + nn.Linear(gate_low_rank_dim, self.key_dim_per_group, bias=True)) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + if gate_fn == 'swish' and fuse_norm and use_output_gate: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.gate_fn = ACT2FN[gate_fn] + + self.gate_logit_normalizer = gate_logit_normalizer + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + gk = self.gk_proj(hidden_states) + + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + q = rearrange(q, 'b t (h d) -> b t h d', d=self.head_k_dim) + if self.num_kv_groups > 1: + k, gk = (repeat(x, 'b t (h d) -> b t (h g) d', g=self.num_kv_groups, d=self.head_k_dim) for x in (k, gk)) + v = repeat(v, 'b t (h d) -> b t (h g) d', g=self.num_kv_groups, d=self.head_v_dim) + else: + k, gk = (rearrange(x, 'b t (h d) -> b t h d', d=self.head_k_dim) for x in (k, gk)) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + gk = F.logsigmoid(gk) / self.gate_logit_normalizer + + if self.clamp_min is not None: + gk = torch.clamp_min(gk, self.clamp_min) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + head_first=False + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + if self.use_output_gate: + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, 'b t (h d) -> b t h d', d=self.head_v_dim) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, 'b t h d -> b t (h d)') + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = o * self.gate_fn(g) + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_v_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/gsa.py b/fla/layers/gsa.py new file mode 100644 index 0000000000000000000000000000000000000000..dc22b1d392d9dc117ae91b92167c3db32e746732 --- /dev/null +++ b/fla/layers/gsa.py @@ -0,0 +1,227 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.modules import RMSNorm, ShortConvolution +from fla.modules.feature_map import ReLUFeatureMap, SwishFeatureMap, T2RFeatureMap +from fla.modules.layernorm import rms_norm_linear +from fla.ops.gsa import chunk_gsa, fused_recurrent_gsa + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedSlotAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1., + expand_v: float = 1., + num_heads: int = 4, + num_kv_heads: Optional[int] = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_slots: Optional[int] = None, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 8, + feature_map: str = 'swish', + use_output_gate: bool = False, + use_norm: bool = True, + layer_idx: Optional[int] = None, + scale: Optional[float] = 1., + **kwargs + ) -> GatedSlotAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.head_k_dim = self.key_dim // self.num_heads + self.head_v_dim = self.value_dim // self.num_heads + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.gate_logit_normalizer = gate_logit_normalizer + + self.use_output_gate = use_output_gate + self.use_norm = use_norm + self.scale = scale + + if num_slots is None: + num_slots = self.head_k_dim + self.num_slots = num_slots + + self.layer_idx = layer_idx + + if layer_idx is None: + warnings.warn( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.register_module('feature_map', None) + if feature_map == 'swish': + self.feature_map = SwishFeatureMap() + elif feature_map == 'relu': + self.feature_map = ReLUFeatureMap() + elif feature_map == 't2r': + self.feature_map = T2RFeatureMap(self.head_k_dim, self.head_k_dim) + else: + raise NotImplementedError(f"Feature map `{feature_map}` is not supported now.") + + self.q_proj = nn.Linear(self.hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.value_dim_per_group, bias=False) + self.f_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.num_slots, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu') + self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu') + + self.g_norm = RMSNorm(self.hidden_size, elementwise_affine, eps=norm_eps) + self.o_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + f = self.f_proj(hidden_states) + + q = rearrange(q, 'b t (h d) -> b t h d', d=self.head_k_dim) + k = rearrange(k, 'b t (h d) -> b t h d', d=self.head_k_dim) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + f = rearrange(f, 'b t (h m) -> b t h m', m=self.num_slots) + + if self.feature_map is not None: + q, k = map(lambda x: self.feature_map(x), (q, k)) + v = F.silu(v) + + f = F.logsigmoid(f) / self.gate_logit_normalizer + s = (1 - f.exp()).to(f.dtype) + # dealing with left-padding + if attention_mask is not None: + s = s.mul_(attention_mask[:, -s.shape[1]:, None, None]) + v = v.mul_(attention_mask[:, -v.shape[1]:, None, None]) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + scale=self.scale, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gsa( + q=q, + k=k, + v=v, + s=s, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + scale=self.scale, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + o = rearrange(o, 'b t h d -> b t (h d)') + o = rms_norm_linear(F.silu(o), self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias) + return o, None, past_key_values + + def state_size(self, *args, **kwargs) -> int: + return 2 * self.num_slots * self.hidden_size diff --git a/fla/layers/hgrn.py b/fla/layers/hgrn.py new file mode 100644 index 0000000000000000000000000000000000000000..6dfa0f946f9b47ec2315605554e2362acef6cd56 --- /dev/null +++ b/fla/layers/hgrn.py @@ -0,0 +1,168 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# "Hierarchically Gated Recurrent Neural Network for Sequence Modeling" [https://arxiv.org/abs/2311.04823] + +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.modules.activations import swiglu +from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class HGRNAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_ratio: Optional[int] = 1, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + layer_idx: int = None + ) -> HGRNAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_ratio = expand_ratio + self.input_dim = int(hidden_size * expand_ratio) + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + + self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + self.f_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + self.g_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.input_dim, conv_size, activation=None) + self.f_conv1d = ShortConvolution(self.input_dim, conv_size, activation=None) + self.i_conv1d = ShortConvolution(self.input_dim, conv_size, activation=None) + + self.g_norm = FusedRMSNormGated( + hidden_size=self.input_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + lower_bound: Optional[torch.Tensor] = None, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if not self.training and hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_i, conv_state_f = None, None + if last_state is not None: + conv_state_i, conv_state_f = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + i, conv_state_i = self.i_conv1d( + x=self.i_proj(hidden_states), + mask=conv_mask, + cache=conv_state_i, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + f, conv_state_f = self.f_conv1d( + x=self.f_proj(hidden_states), + mask=conv_mask, + cache=conv_state_f, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + i = self.i_proj(hidden_states) + f = self.f_proj(hidden_states) + + # the lower bound for the first layer is zero + if lower_bound is None or self.layer_idx == 0: + i, f = swiglu(i, 1 - f.sigmoid()), F.logsigmoid(f) + else: + g = lower_bound + (1 - lower_bound) * f.sigmoid() + i, f = swiglu(i, 1 - g), g.log() + + # dealing with left-padding + if attention_mask is not None: + i = i.mul_(attention_mask[:, -i.shape[-2]:, None]) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + if cu_seqlens is not None: + raise NotImplementedError("Chunk mode does not support variable-length sequences.") + o, recurrent_state = chunk_hgrn( + x=i, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_hgrn( + x=i, + g=f, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_i, conv_state_f) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=i.shape[2] + ) + + o = self.g_norm(o, self.g_proj(hidden_states)) + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.hidden_size + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/hgrn2.py b/fla/layers/hgrn2.py new file mode 100644 index 0000000000000000000000000000000000000000..13fb04253790113c770faa93101b890eb63e4a19 --- /dev/null +++ b/fla/layers/hgrn2.py @@ -0,0 +1,211 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# "HGRN2: Gated Linear RNNs with State Expansion"[https://arxiv.org/abs/2404.07904] + +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.modules import RMSNorm, ShortConvolution +from fla.modules.activations import swish +from fla.modules.layernorm import rms_norm_linear +from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class HGRN2Attention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + num_heads: Optional[int] = None, + expand_ratio: Optional[int] = 128, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + layer_idx: int = None + ) -> HGRN2Attention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + + if expand_ratio is None and num_heads is not None: + expand_ratio = hidden_size // num_heads + elif expand_ratio is not None and num_heads is None: + num_heads = hidden_size // expand_ratio + elif expand_ratio is None and num_heads is None: + raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.") + self.num_heads = num_heads + self.expand_ratio = expand_ratio + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.forget_dim = int(self.num_heads * self.expand_ratio) + self.input_dim = hidden_size + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_recurrent', 'fused_chunk'], f"Not suppoerted mode `{mode}`." + assert self.forget_dim % num_heads == 0, f"forget dim must be divisible by num_heads of {num_heads}" + assert self.input_dim % num_heads == 0, f"input dim must be divisible by num_heads of {num_heads}" + + self.head_f_dim = self.expand_ratio + self.head_i_dim = self.hidden_size // num_heads + + self.q_proj = nn.Linear(hidden_size, self.forget_dim, bias=False) + self.f_proj = nn.Linear(hidden_size, self.forget_dim, bias=False) + self.i_proj = nn.Linear(hidden_size, self.input_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.forget_dim, conv_size, activation=None) + self.f_conv1d = ShortConvolution(self.forget_dim, conv_size, activation=None) + self.i_conv1d = ShortConvolution(self.input_dim, conv_size, activation=None) + + self.g_norm = RMSNorm(hidden_size=self.hidden_size, elementwise_affine=elementwise_affine, eps=norm_eps) + self.o_proj = nn.Linear(self.input_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + lower_bound: Optional[torch.Tensor] = None, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_f, conv_state_i = None, None, None + if last_state is not None: + conv_state_q, conv_state_f, conv_state_i = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + f, conv_state_f = self.f_conv1d( + x=self.f_proj(hidden_states), + mask=conv_mask, + cache=conv_state_f, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + i, conv_state_i = self.i_conv1d( + x=self.i_proj(hidden_states), + mask=conv_mask, + cache=conv_state_i, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + f = self.f_proj(hidden_states) + i = self.i_proj(hidden_states) + + # dealing with left-padding + if attention_mask is not None: + i = i.mul_(attention_mask[:, -i.shape[-2]:, None]) + + q = swish(q) + + # improve precision + f = f.float() + + # the lower bound for the first layer is zero + if lower_bound is None or self.layer_idx == 0: + k, g = 1 - f.sigmoid(), F.logsigmoid(f) + else: + g = lower_bound + (1 - lower_bound) * f.sigmoid() + k, g = 1 - g, g.log() + + q, k, g = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k.to(i), g)) + i = rearrange(i, '... (h d) -> ... h d', d=self.head_i_dim) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=i, + gk=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_gla( + q=q, + k=k, + v=i, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + head_first=False + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=i, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_f, conv_state_i) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + o = rearrange(o, '... h d -> ... (h d)') + o = rms_norm_linear(o, self.g_norm.weight, self.g_norm.bias, self.o_proj.weight, self.o_proj.bias) + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.forget_dim * self.head_i_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/lightnet.py b/fla/layers/lightnet.py new file mode 100644 index 0000000000000000000000000000000000000000..c7f3c01a88938a302502f12a54abeb159cf4db58 --- /dev/null +++ b/fla/layers/lightnet.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# ["You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet"](https://arxiv.org/abs/2405.21022) + +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.modules.fused_norm_gate import rms_norm_swish_gate_linear +from fla.ops.gla import chunk_gla, fused_recurrent_gla + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class LightNetAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + num_heads: Optional[int] = None, + expand_ratio: Optional[int] = 128, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + gate_low_rank_dim: int = 128, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + layer_idx: int = None + ) -> LightNetAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + + if expand_ratio is None and num_heads is not None: + expand_ratio = hidden_size // num_heads + elif expand_ratio is not None and num_heads is None: + num_heads = hidden_size // expand_ratio + elif expand_ratio is None and num_heads is None: + raise RuntimeError("One of `expand_ratio` or `num_heads` should be provided.") + self.num_heads = num_heads + self.expand_ratio = expand_ratio + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.key_dim = int(self.num_heads * self.expand_ratio) + self.value_dim = hidden_size + self.gate_low_rank_dim = gate_low_rank_dim + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_chunk'], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_f_dim = self.expand_ratio + self.head_i_dim = self.hidden_size // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation=None) + self.k_conv1d = ShortConvolution(self.key_dim, conv_size, activation=None) + self.v_conv1d = ShortConvolution(self.value_dim, conv_size, activation=None) + + self.g_proj = nn.Sequential( + nn.Linear(hidden_size, gate_low_rank_dim, bias=False), + nn.Linear(gate_low_rank_dim, hidden_size, bias=False) + ) + self.g_norm = FusedRMSNormGated( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[Dict] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + + q = F.silu(q) + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_f_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_i_dim) + # TODO: this 2 steps took huge amount of time, which should be optimized + z = k.float().logcumsumexp(1) + + if cu_seqlens is not None: + raise NotImplementedError("LightNet does not support variable-length sequences for now.") + k, g = torch.exp(k - z).to(k.dtype), (torch.cat((z[:, :1], z[:, :-1]), 1) - z).to(k.dtype) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + o = rms_norm_swish_gate_linear( + rearrange(o, 'b t h d -> b t (h d)'), + self.g_proj(hidden_states), + self.g_norm.weight, + self.g_norm.bias, + self.o_proj.weight, + self.o_proj.bias + ) + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_i_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/linear_attn.py b/fla/layers/linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..54c570b7b5759f0a5d7e9653e87b1ab07b3546da --- /dev/null +++ b/fla/layers/linear_attn.py @@ -0,0 +1,166 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.modules import RMSNorm +from fla.modules.feature_map import DPFPFeatureMap, HadamardFeatureMap, HedgehogFeatureMap, T2RFeatureMap +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn + + +class LinearAttention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: str = 1024, + expand_k: int = 1.0, + expand_v: int = 1.0, + num_heads: int = 8, + num_kv_heads: Optional[int] = None, + feature_map: str = 'elementwise_product', + tie_feature_map_qk: bool = False, + output_norm: str = 'rmsnorm', + norm_q: bool = False, + norm_k: bool = False, + do_feature_map_norm: bool = False, + elementwise_affine: bool = True, + norm_eps: float = 1e-5, + **kwargs + ): + super().__init__() + + self.hidden_size = hidden_size + self.mode = mode + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + + assert mode in ['chunk', 'fused_chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + self.do_feature_map_norm = do_feature_map_norm + + if feature_map == 'hedgehog': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = HedgehogFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = HedgehogFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 't2r': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = T2RFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = T2RFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'elementwise_product': + if tie_feature_map_qk: + self.feature_map_q = self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim) + else: + self.feature_map_q = HadamardFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = HadamardFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'dpfp': + self.feature_map_q = DPFPFeatureMap(head_dim=self.head_k_dim) + self.feature_map_k = DPFPFeatureMap(head_dim=self.head_k_dim) + + elif feature_map == 'elu': + def elu(x): + return F.elu(x) + 1 + self.feature_map_q = elu + self.feature_map_k = elu + + elif feature_map == 'relu': + self.feature_map_q = nn.ReLU() + self.feature_map_k = nn.ReLU() + + elif feature_map == 'identity': + self.feature_map_q = nn.Identity() + self.feature_map_k = nn.Identity() + else: + raise NotImplementedError(f"Not supported feature map `{feature_map}`.") + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + + if output_norm == 'rmsnorm': + self.norm = RMSNorm(hidden_size=self.head_v_dim, elementwise_affine=elementwise_affine, eps=norm_eps) + elif output_norm == 'identity': + self.norm = nn.Identity() + else: + raise NotImplementedError(f"Not supported output norm `{output_norm}`.") + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.norm_q = norm_q + self.norm_k = norm_k + + def forward( + self, + hidden_states: torch.Tensor, + **kwargs + ) -> torch.Tensor: + mode = self.mode + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + if self.num_kv_groups > 1: + k = repeat(k, '... (h d) -> ... (h g) d', d=self.head_k_dim, g=self.num_kv_groups) + v = repeat(v, '... (h d) -> ... (h g) d', d=self.head_v_dim, g=self.num_kv_groups) + else: + k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + q = self.feature_map_q(q) + k = self.feature_map_k(k) + + if self.norm_q: + q = q / (q.sum(-1, True) + 1e-4) + if self.norm_k: + k = k / (k.sum(-1, True) + 1e-4) + + if mode == 'chunk': + o, final_state = chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=self.do_feature_map_norm, + head_first=False + ) + elif mode == 'fused_chunk': + o, final_state = fused_chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=self.do_feature_map_norm, + ) + elif mode == 'fused_recurrent': + o, final_state = fused_recurrent_linear_attn( + q=q, + k=k, + v=v, + normalize=self.do_feature_map_norm, + ) + else: + raise NotImplementedError + o = self.norm(o) + o = rearrange(o, '... h d -> ... (h d)') + o = self.o_proj(o) + return o diff --git a/fla/layers/multiscale_retention.py b/fla/layers/multiscale_retention.py new file mode 100644 index 0000000000000000000000000000000000000000..e5364b356b245221b78c32c41eba1b4ca1cc5622 --- /dev/null +++ b/fla/layers/multiscale_retention.py @@ -0,0 +1,298 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from transformers.activations import ACT2FN + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.rotary import RotaryEmbedding +from fla.ops.retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class MultiScaleRetention(nn.Module): + r""" + The layer implementaion for [Retentive Network: A Successor to Transformer for Large Language Models](https://arxiv.org/pdf/2307.08621.pdf). # noqa + + Args: + mode (str, Optional): + Which Retention kernel to use. + Currently available: `chunk`, `fused_recurrent`, `parallel`, and `fused_chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + num_heads (int, Optional): + The number of heads. Default: 8. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1.0, + expand_v: float = 2.0, + num_heads: int = 8, + num_kv_heads: Optional[int] = None, + feature_map: Optional[str] = None, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + use_output_gate: bool = True, + gate_fn: str = 'swish', + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + fuse_norm: bool = True, + layer_idx: int = None, + **kwargs + ) -> MultiScaleRetention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.use_output_gate = use_output_gate + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.layer_idx = layer_idx + + assert mode in ['chunk', 'fused_chunk', 'parallel', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + if self.use_output_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu') + self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu') + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + if gate_fn == 'swish' and fuse_norm and use_output_gate: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.gate_fn = ACT2FN[gate_fn] + + # TODO: fix this issue + # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py#L180 + # Ideally, we would want to support arbitrary d_head_qk + assert self.head_k_dim <= 256, "head_k_dim must be less than or equal to 256" + self.rotary = RotaryEmbedding(dim=self.head_k_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + q = rearrange(q, '... (h d) -> ... h d', d=self.head_k_dim) + k = rearrange(k, '... (h d) -> ... h d', d=self.head_k_dim) + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + + seqlen_offset, max_seqlen = 0, q.shape[1] + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if self.num_kv_groups > 1: + k = repeat(k, 'b t h d -> b t (h g) d', g=self.num_kv_groups) + v = repeat(v, 'b t (h d) -> b t (h g) d', d=self.head_v_dim, g=self.num_kv_groups) + else: + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'fused_chunk': + o, recurrent_state = fused_chunk_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'parallel': + o, recurrent_state = parallel_retention( + q=q, + k=k, + v=v, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_retention( + q=q, + k=k, + v=v, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + if self.use_output_gate: + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, 'b t (h d) -> b t h d', d=self.head_v_dim) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, 'b t h d -> b t (h d)') + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = o * self.gate_fn(g) + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_v_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/layers/nsa.py b/fla/layers/nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..0a65b458d6c000bd52c56bd4c93c4357b5c3844f --- /dev/null +++ b/fla/layers/nsa.py @@ -0,0 +1,138 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple, Union + +import torch +import torch.nn as nn +from einops import rearrange +from transformers.utils import logging + +from fla.modules import RotaryEmbedding +from fla.ops.nsa.parallel import parallel_nsa + +if TYPE_CHECKING: + from fla.models.utils import Cache + +logger = logging.get_logger(__name__) + + +class NativeSparseAttention(nn.Module): + + def __init__( + self, + hidden_size: int = 2048, + num_heads: int = 64, + num_kv_heads: Optional[int] = 4, + head_dim: int = 64, + qkv_bias: bool = False, + block_size: Optional[int] = 64, + block_counts: Optional[Union[torch.LongTensor, int]] = 16, + window_size: Optional[int] = 512, + rope_theta: Optional[float] = 10000., + max_position_embeddings: Optional[int] = None, + layer_idx: int = None + ): + super().__init__() + + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None: + self.num_kv_heads = self.num_heads + else: + self.num_kv_heads = num_kv_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.head_dim = head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + self.qkv_bias = qkv_bias + + self.block_size = block_size + self.block_counts = block_counts + self.window_size = window_size + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.layer_idx = layer_idx + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.qkv_bias) + self.k_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.v_proj = nn.Linear(self.hidden_size, self.kv_dim, bias=self.qkv_bias) + self.g_proj = nn.Linear(self.hidden_size, self.num_heads * 3, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, seq_len, _ = hidden_states.size() + + q = rearrange(self.q_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + k = rearrange(self.k_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(self.v_proj(hidden_states), '... (h d) -> ... h d', d=self.head_dim) + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=3) + g_cmp, g_slc, g_swa = g.sigmoid().unbind(-1) + + cu_seqlens = kwargs.get('cu_seqlens', None) + + seqlen_offset, max_seqlen = 0, seq_len + if past_key_values is not None: + seqlen_offset = past_key_values.get_seq_length(self.layer_idx) + max_seqlen = q.shape[1] + seqlen_offset + + if attention_mask is not None: + # to deliminate the offsets of padding tokens + seqlen_offset = seqlen_offset + attention_mask.sum(-1) - attention_mask.shape[-1] + max_seqlen = q.shape[1] + max(seqlen_offset) + + if self.max_position_embeddings is not None: + max_seqlen = max(max_seqlen, self.max_position_embeddings) + q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen, cu_seqlens=cu_seqlens) + + if past_key_values is not None: + cache_has_content = past_key_values.get_seq_length(self.layer_idx) > 0 + k_cached, v_cached = past_key_values.update( + attn_state=(k.flatten(-2, -1), v.flatten(-2, -1)), + layer_idx=self.layer_idx, + offset=seq_len, + cache_kwargs=dict(window_size=self.window_size) + )['attn_state'] + if cache_has_content: + k, v = k_cached, v_cached + k = rearrange(k, '... (h d) -> ... h d', d=self.head_dim) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + o = parallel_nsa( + q=q, + k=k, + v=v, + g_cmp=g_cmp, + g_slc=g_slc, + g_swa=g_swa, + block_size=self.block_size, + block_counts=self.block_counts, + window_size=self.window_size, + cu_seqlens=cu_seqlens, + head_first=False + ) + o = o.reshape(batch_size, seq_len, -1) + o = self.o_proj(o) + + if not output_attentions: + attentions = None + + return o, attentions, past_key_values diff --git a/fla/layers/rebased.py b/fla/layers/rebased.py new file mode 100644 index 0000000000000000000000000000000000000000..47c2e17dfe088cc3fda45c3d3aa3605b8fdb1e73 --- /dev/null +++ b/fla/layers/rebased.py @@ -0,0 +1,133 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +""" +https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/layers/rebased_fast.py +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from fla.modules.feature_map import RebasedFeatureMap +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn +from fla.ops.rebased import parallel_rebased + + +class ReBasedLinearAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + l_max: int = 2048, + feature_dim: int = 16, + num_key_value_heads: int = 16, + num_heads: int = 16, + use_gamma: Optional[bool] = True, + use_beta: Optional[bool] = True, + normalize: Optional[bool] = True, + causal: bool = True, + eps: float = 1e-5, + mode: str = "parallel", + layer_idx: Optional[int] = None, + **kwargs + ) -> ReBasedLinearAttention: + super().__init__() + self.hidden_size = hidden_size + self.l_max = l_max + self.mode = mode + assert self.mode in ["fused_chunk", "parallel", 'chunk'] + + self.feature_dim = feature_dim + self.num_key_value_heads = num_key_value_heads + self.num_heads = num_heads + self.head_dim = self.hidden_size // self.num_key_value_heads + self.use_gamma = use_gamma + self.use_beta = use_beta + self.normalize = normalize + self.causal = causal + self.eps = eps + self.mode = mode + self.layer_idx = layer_idx + + self.feature_map = RebasedFeatureMap(self.feature_dim, use_gamma, use_beta, normalize) + self.q_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.k_proj = nn.Linear(self.hidden_size, self.feature_dim * self.num_heads, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.dropout = nn.Identity() + + def forward(self, hidden_states: torch.Tensor, **kwargs): + mode = self.mode + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + q, k, v = map(lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_dim), [q, k, v]) + q, k = self.feature_map(q, flatten=(mode != 'parallel')), self.feature_map(k, flatten=(mode != 'parallel')) + if mode == "fused_chunk": + o = fused_chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=True, + scale=1, + head_first=False + ) + elif mode == 'chunk': + o = chunk_linear_attn( + q=q, + k=k, + v=v, + normalize=True, + scale=1, + head_first=False + ) + elif mode == 'parallel': + assert q.shape[-1] <= 128 + o = parallel_rebased( + q=q, + k=k, + v=v, + eps=self.eps, + use_scale=True, + use_normalize=True, + head_first=False + ) + o = self.o_proj(o) + o = self.dropout(o) + return o + + # https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/based.py#L119 + def forward_reference( + self, + hidden_states: torch.Tensor, + filters: torch.Tensor = None, + *args, + **kwargs + ): + """ + x (torch.Tensor): tensor of shape (b, d, t) + y (torch.Tensor): tensor of shape (b, d, t) + """ + b, t, _ = hidden_states.size() + q, k, v = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) + + q = q.view(b, t, -1, self.feature_dim).transpose(1, 2) + k = k.view(b, t, -1, self.feature_dim).transpose(1, 2) + v = v.view(b, t, -1, self.head_dim).transpose(1, 2) + + # Linear attention + q, k = self.feature_map(q), self.feature_map(k) + q, k, v = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1) + + # Compute attention + if self.causal: + y = ((q * (k * v).cumsum(2)).sum(-1) / ((q * k.cumsum(2)).sum(-1) + self.eps)) + else: + y = ((q * (k * v).sum(2, True)).sum(-1) / ((q * k.sum(2, True)).sum(-1) + self.eps)) + y = rearrange(y, 'b h t d -> b t (h d)') + y = self.o_proj(y.to(hidden_states.dtype)) + y = self.dropout(y) + return y.to(hidden_states.dtype) diff --git a/fla/layers/rwkv7.py b/fla/layers/rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..44c6bf0b868965132710e2ad43edf663611ca618 --- /dev/null +++ b/fla/layers/rwkv7.py @@ -0,0 +1,221 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.layers.rwkv6 import LoRA +from fla.modules import GroupNorm +from fla.modules.l2norm import l2_norm +from fla.ops.rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class RWKV7Attention(nn.Module): + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + head_dim: Optional[int] = 64, + num_heads: Optional[int] = None, + decay_low_rank_dim: int = 64, + gate_low_rank_dim: int = 128, + a_low_rank_dim: int = 64, + v_low_rank_dim: int = 16, + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + layer_idx: int = None, + fuse_norm: bool = False, + value_dim: int = None, + **kwargs + ) -> RWKV7Attention: + super().__init__() + + self.mode = mode + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + self.hidden_size = hidden_size + + self.key_dim = hidden_size + self.value_dim = value_dim if value_dim is not None else hidden_size + if head_dim is None and num_heads is None: + raise ValueError("Either `head_dim` or `num_heads` must be specified.") + elif head_dim is not None: + self.head_dim = head_dim + self.num_heads = int(hidden_size // head_dim) + elif num_heads is not None: + self.head_dim = int(hidden_size // num_heads) + self.num_heads = num_heads + self.head_v_dim = int(self.value_dim // self.num_heads) + + self.decay_low_rank_dim = decay_low_rank_dim + self.gate_low_rank_dim = gate_low_rank_dim + self.a_low_rank_dim = a_low_rank_dim + self.v_low_rank_dim = v_low_rank_dim + self.layer_idx = layer_idx + self.fuse_norm = fuse_norm + + self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) + + self.x_x = nn.Parameter(torch.zeros(6, hidden_size)) + + self.k_k = nn.Parameter(torch.zeros(self.key_dim)) + self.k_a = nn.Parameter(torch.zeros(self.key_dim)) + self.r_k = nn.Parameter(torch.zeros(self.num_heads, self.head_dim)) + + self.r_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.w_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=decay_low_rank_dim, activation='tanh') + if self.layer_idx != 0: + self.v_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=v_low_rank_dim, activation=None) + self.a_lora = LoRA(hidden_size, self.key_dim, low_rank_dim=a_low_rank_dim, activation=None) + self.g_lora = LoRA(hidden_size, self.value_dim, low_rank_dim=gate_low_rank_dim, activation='sigmoid', bias=False) + + if self.fuse_norm: + self.g_norm = GroupNorm( + num_groups=self.num_heads, + hidden_size=self.value_dim, + elementwise_affine=elementwise_affine, + eps=self.head_dim*norm_eps, + bias=True, + ) + else: + self.g_norm = nn.GroupNorm( + num_groups=self.num_heads, + num_channels=self.value_dim, + eps=self.head_dim*norm_eps, + affine=elementwise_affine + ) + + self.apply(self._initialize_weights) + + def _initialize_weights(self, module: nn.Module): + if getattr(module, "_is_hf_initialized", False): + return + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight, gain=2 ** -2.5) + if module.bias is not None: + nn.init.zeros_(module.bias) + if isinstance(module, nn.Parameter): + nn.init.xavier_uniform_(module, gain=2 ** -2.5) + module._is_hf_initialized = True + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + v_first: torch.Tensor = None, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, seq_len, _ = hidden_states.shape + + if self.training: + # if training, use chunk mode no matter how short the sequence is + mode = 'chunk' + else: + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + if attention_mask is not None: + hidden_states = hidden_states.mul(attention_mask[:, -hidden_states.shape[-2]:, None]) + if hidden_states.shape[1] == 1 and last_state is not None: + shifted = last_state['conv_state'].unsqueeze(1) + else: + shifted = self.time_shift(hidden_states) + if last_state is not None: + shifted[:, 0] = last_state['conv_state'] + + # [batch_size, seq_len, hidden_size] + delta = shifted - hidden_states + xr, xw, xk, xv, xa, xg = hidden_states.addcmul(delta, self.x_x.view(6, 1, 1, -1)).unbind(0) + + r = self.r_proj(xr) + # -math.exp(-0.5) = -0.6065306597126334 + # I think .to(torch.float) is unnecessary here, since we calculate lora in bloat16 + # when we apply sigmoid, bf16 input will not have numerical issue + # FIXME: check if we can remove .to(torch.float) + w = -0.6065306597126334 * self.w_lora(xw).to(torch.float).sigmoid() + + k = self.k_proj(xk) + v = self.v_proj(xv) + + if self.layer_idx == 0: + v_first = v + else: + v = torch.lerp(v, v_first, self.v_lora(xv).sigmoid()) + a = self.a_lora(xa).sigmoid() + g = self.g_lora(xg) + + if self.fuse_norm: + kk = l2_norm(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim)) + else: + kk = F.normalize(rearrange(k * self.k_k, 'b t (h d) -> b t h d', d=self.head_dim), dim=-1, p=2.0) + + k = k.addcmul(k * (a - 1), self.k_a) + + # dealing with left-padding + if attention_mask is not None: + v = v * attention_mask[:, -v.shape[-2]:, None] + r, w, k, a = map(lambda x: rearrange(x, 'b t (h d) -> b t h d', d=self.head_dim), (r, w, k, a)) + v = rearrange(v, 'b t (h d) -> b t h d', d=self.head_v_dim) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + + rwkv7_fn = chunk_rwkv7 if mode == 'chunk' else fused_recurrent_rwkv7 + cu_seqlens = kwargs.get('cu_seqlens', None) + o, recurrent_state = rwkv7_fn( + r=r, + w=w, + k=k, + v=v, + a=-kk, + b=kk * a, + scale=1., + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=hidden_states[:, -1], + layer_idx=self.layer_idx, + offset=r.shape[1] + ) + + if self.fuse_norm: + o = self.g_norm(rearrange(o, '... h d -> ... (h d)')) + else: + o = self.g_norm(rearrange(o, 'b t h d -> (b t) (h d)')).view(batch_size, seq_len, -1) + + o = o + ((r * k * self.r_k).sum(-1, keepdim=True) * v).view(batch_size, seq_len, -1) + o = self.o_proj(o * g) + + return o, None, past_key_values, v_first diff --git a/fla/layers/simple_gla.py b/fla/layers/simple_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..0771a90ce2e15b1a43f4ffa081c63a37ce3ea124 --- /dev/null +++ b/fla/layers/simple_gla.py @@ -0,0 +1,261 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.activations import ACT2FN +from fla.ops.simple_gla import chunk_simple_gla, fused_recurrent_simple_gla + +if TYPE_CHECKING: + from fla.models.utils import Cache + + +class SimpleGatedLinearAttention(nn.Module): + r""" + The layer implementaion for [Gated Linear Attention Transformers with Hardware-Efficient Training](https://arxiv.org/abs/2312.06635). # noqa + This layer calls the simplified GLA kernel in which the gating is head-wise instead of elementwise. + + Args: + mode (str, Optional): + Which GLA kernel to use. + Currently available: `chunk`. + Default: `chunk`. + hidden_size (int, Optional): + The hidden size of the input. Default: 1024. + expand_k (float, Optional): + The expansion ratio for the key dim. Default: 1.0. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 1.0. + num_heads (int, Optional): + The number of heads. Default: 4. + num_kv_heads (int, Optional): + The number of key/value heads, used for MQA. Default: None. + feature_map (str, Optional): + Feature map function applied to queries/keys. Default: None. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + gate_fn (str, Optional): + The activation function for the output gate. Default: `swish`. + elementwise_affine (bool, Optional): + If `True`, applies elementwise affine to LayerNorm with learnable parameters. Default: `True`. + norm_eps (float, Optional): + The epsilon value for the layernorm/rmsnorm layer. Default: 1e-5. + gate_logit_normalizer (int, Optional): + The normalizer for the gate logits, appied after `logsigmoid`. Default: 16. + fuse_norm (bool, Optional): + Whether to fuse the norm and the output gate for better memory footprint. Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + """ + + def __init__( + self, + mode: str = 'chunk', + hidden_size: int = 1024, + expand_k: float = 1., + expand_v: float = 1., + num_heads: int = 4, + num_kv_heads: Optional[int] = None, + feature_map: Optional[str] = None, + use_short_conv: bool = True, + conv_size: int = 4, + conv_bias: bool = False, + gate_fn: str = 'swish', + elementwise_affine: Optional[bool] = True, + norm_eps: float = 1e-5, + gate_logit_normalizer: int = 16, + fuse_norm: bool = True, + layer_idx: int = None, + ) -> SimpleGatedLinearAttention: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_k = expand_k + self.expand_v = expand_v + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.feature_map_fn = ACT2FN[feature_map] if feature_map is not None else None + + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.key_dim = int(hidden_size * expand_k) + self.value_dim = int(hidden_size * expand_v) + self.key_dim_per_group = self.key_dim // self.num_kv_groups + self.value_dim_per_group = self.value_dim // self.num_kv_groups + self.layer_idx = layer_idx + + assert mode in ['chunk', "fused_recurrent"], f"Not suppoerted mode `{mode}`." + assert self.key_dim % num_heads == 0, f"key dim must be divisible by num_heads of {num_heads}" + assert self.value_dim % num_heads == 0, f"value dim must be divisible by num_heads of {num_heads}" + + self.head_k_dim = self.key_dim // num_heads + self.head_v_dim = self.value_dim // num_heads + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim_per_group, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim_per_group, bias=False) + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution(self.key_dim, conv_size, activation='silu') + self.k_conv1d = ShortConvolution(self.key_dim_per_group, conv_size, activation='silu') + self.v_conv1d = ShortConvolution(self.value_dim_per_group, conv_size, activation='silu') + + self.gk_proj = nn.Linear(hidden_size, self.num_heads) + + if gate_fn == 'swish' and fuse_norm: + self.g_norm_swish_gate = FusedRMSNormGated( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.fuse_norm_and_gate = True + else: + self.fuse_norm_and_gate = False + self.g_norm = RMSNorm( + hidden_size=self.head_v_dim, + elementwise_affine=elementwise_affine, + eps=norm_eps + ) + self.gate_fn = ACT2FN[gate_fn] + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + self.gate_logit_normalizer = gate_logit_normalizer + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + **kwargs + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + # launching the triton kernel for just one token will actually be slower + mode = 'fused_recurrent' if hidden_states.shape[1] <= 64 else self.mode + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens', None) + if self.use_short_conv: + conv_state_q, conv_state_k, conv_state_v = None, None, None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + conv_mask = attention_mask[:, -hidden_states.shape[1]:] if attention_mask is not None else None + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + mask=conv_mask, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + mask=conv_mask, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + mask=conv_mask, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens + ) + else: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + gk = self.gk_proj(hidden_states) + + if self.feature_map_fn is not None: + q, k = map(self.feature_map_fn, (q, k)) + # dealing with left-padding + if attention_mask is not None: + v = v.mul_(attention_mask[:, -v.shape[-2]:, None]) + q = rearrange(q, '... (h d) -> ... h d', h=self.num_heads) + if self.num_kv_groups > 1: + k, v = (repeat(x, '... (h d) -> ... (h g) d', h=self.num_kv_heads, g=self.num_kv_groups) for x in (k, v)) + else: + k, v = (rearrange(x, '... (h d) -> ... h d', h=self.num_kv_heads) for x in (k, v)) + gk = F.logsigmoid(gk) / self.gate_logit_normalizer + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_simple_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + head_first=False + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q.shape[1] + ) + + g = self.g_proj(hidden_states) + if self.fuse_norm_and_gate: + g = rearrange(g, 'b t (h d) -> b t h d', h=self.num_heads) + o = self.g_norm_swish_gate(o, g) + o = rearrange(o, 'b t h d -> b t (h d)') + else: + o = rearrange(self.g_norm(o), 'b t h d -> b t (h d)') + o = o * self.gate_fn(g) + o = self.o_proj(o) + + return o, None, past_key_values + + def state_size(self, **kwargs) -> int: + state_size = self.key_dim * self.head_v_dim + for module in self.children(): + if isinstance(module, ShortConvolution): + state_size += module.state_size + return state_size diff --git a/fla/models/__init__.py b/fla/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..211ab48bbf76ef22075c58b8143163ed6d0e8d2f --- /dev/null +++ b/fla/models/__init__.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- + +from fla.models.abc import ABCConfig, ABCForCausalLM, ABCModel +from fla.models.bitnet import BitNetConfig, BitNetForCausalLM, BitNetModel +from fla.models.delta_net import DeltaNetConfig, DeltaNetForCausalLM, DeltaNetModel +from fla.models.forgetting_transformer import ( + ForgettingTransformerConfig, + ForgettingTransformerForCausalLM, + ForgettingTransformerModel +) +from fla.models.gated_deltanet import GatedDeltaNetConfig, GatedDeltaNetForCausalLM, GatedDeltaNetModel +from fla.models.gated_deltaproduct import GatedDeltaProductConfig, GatedDeltaProductForCausalLM, GatedDeltaProductModel +from fla.models.gla import GLAConfig, GLAForCausalLM, GLAModel +from fla.models.gsa import GSAConfig, GSAForCausalLM, GSAModel +from fla.models.hgrn import HGRNConfig, HGRNForCausalLM, HGRNModel +from fla.models.hgrn2 import HGRN2Config, HGRN2ForCausalLM, HGRN2Model +from fla.models.lightnet import LightNetConfig, LightNetForCausalLM, LightNetModel +from fla.models.linear_attn import LinearAttentionConfig, LinearAttentionForCausalLM, LinearAttentionModel +from fla.models.mamba import MambaConfig, MambaForCausalLM, MambaModel +from fla.models.mamba2 import Mamba2Config, Mamba2ForCausalLM, Mamba2Model +from fla.models.nsa import NSAConfig, NSAForCausalLM, NSAModel +from fla.models.retnet import RetNetConfig, RetNetForCausalLM, RetNetModel +from fla.models.rwkv6 import RWKV6Config, RWKV6ForCausalLM, RWKV6Model +from fla.models.rwkv7 import RWKV7Config, RWKV7ForCausalLM, RWKV7Model +from fla.models.samba import SambaConfig, SambaForCausalLM, SambaModel +from fla.models.transformer import TransformerConfig, TransformerForCausalLM, TransformerModel +from fla.models.transformer_top import TOPTransformerConfig, TOPTransformerForCausalLM, TOPTransformerModel +from fla.models.transformer_mtp import MTPTransformerConfig, MTPTransformerForCausalLM, MTPTransformerModel +from fla.models.transformer_dsmtp import DSMTPTransformerConfig, DSMTPTransformerForCausalLM, DSMTPTransformerModel + +__all__ = [ + 'ABCConfig', 'ABCForCausalLM', 'ABCModel', + 'BitNetConfig', 'BitNetForCausalLM', 'BitNetModel', + 'DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel', + 'ForgettingTransformerConfig', 'ForgettingTransformerForCausalLM', 'ForgettingTransformerModel', + 'GatedDeltaNetConfig', 'GatedDeltaNetForCausalLM', 'GatedDeltaNetModel', + 'GLAConfig', 'GLAForCausalLM', 'GLAModel', + 'GSAConfig', 'GSAForCausalLM', 'GSAModel', + 'HGRNConfig', 'HGRNForCausalLM', 'HGRNModel', + 'HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model', + 'LightNetConfig', 'LightNetForCausalLM', 'LightNetModel', + 'LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel', + 'MambaConfig', 'MambaForCausalLM', 'MambaModel', + 'Mamba2Config', 'Mamba2ForCausalLM', 'Mamba2Model', + 'NSAConfig', 'NSAForCausalLM', 'NSAModel', + 'RetNetConfig', 'RetNetForCausalLM', 'RetNetModel', + 'RWKV6Config', 'RWKV6ForCausalLM', 'RWKV6Model', + 'RWKV7Config', 'RWKV7ForCausalLM', 'RWKV7Model', + 'SambaConfig', 'SambaForCausalLM', 'SambaModel', + 'TransformerConfig', 'TransformerForCausalLM', 'TransformerModel', + 'TOPTransformerConfig', 'TOPTransformerForCausalLM', 'TOPTransformerModel', + 'MTPTransformerConfig', 'MTPTransformerForCausalLM', 'MTPTransformerModel', + 'DSMTPTransformerConfig', 'DSMTPTransformerForCausalLM', 'DSMTPTransformerModel', + 'GatedDeltaProductConfig', 'GatedDeltaProductForCausalLM', 'GatedDeltaProductModel', +] diff --git a/fla/models/utils.py b/fla/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..da9d070fb95e2ebb7d90620d05ed2db0f8738c57 --- /dev/null +++ b/fla/models/utils.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +import torch +import transformers + + +class Cache(transformers.cache_utils.Cache): + """ + A cache used for storing hidden states produced by flash linear attention models. + + It stores the states of each layer as the tensor of shape `[batch_size, key_dim, value_dim]`. + """ + + is_compileable = True + + def __init__( + self, + seen_tokens: int = 0 + ) -> Cache: + super().__init__() + + self.states: List[Dict[str, Any]] = [] + + self._seen_tokens = seen_tokens # Used in `generate` to keep tally of how many tokens the cache has seen + + def __getitem__(self, layer_idx: int) -> Dict[str, Any]: + if layer_idx < len(self): + return self.states[layer_idx] + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __iter__(self): + for state in self.states: + yield state + + def __len__(self): + return len(self.states) + + def update( + self, + recurrent_state: torch.Tensor = None, + attn_state: Tuple[torch.Tensor, torch.Tensor] = None, + conv_state: Tuple[torch.Tensor] = None, + ffn_state: torch.Tensor = None, + layer_idx: int = 0, + offset: Optional[int] = 1, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """ + Updates the cache with the new `recurrent_state`/`attn_state`/`conv_state` for the layer `layer_idx`. + + Args: + recurrent_state (`torch.Tensor`, `optional`): + The new recurrent state to cache. + attn_state (`Tuple[torch.Tensor, torch.Tensor]`, `optional`): + The new attention key/value states to cache. + conv_state (`Tuple[torch.Tensor]`, `optional`): + The new convolution state to cache. + layer_idx (`int`, defaults to 0): + The index of the layer to cache the states for. + offset (`int`, `optional`, defaults to 1): + The number of new tokens being processed. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. + + Return: + Dictionary of the updated state. + """ + + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += offset + + if attn_state is not None: + input_size = attn_state[0].shape[-2] + window_size = cache_kwargs.get('window_size', None) + if not isinstance(attn_state, Tuple) or len(attn_state) != 2: + raise ValueError("`attn_state` must be a tuple of two tensors for key/value states") + if len(self.states) <= layer_idx: + if attn_state is not None: + if window_size is not None and input_size > window_size: + attn_state = (attn_state[0][..., -window_size:, :].contiguous(), + attn_state[1][..., -window_size:, :].contiguous()) + state = dict( + recurrent_state=recurrent_state, + attn_state=attn_state, + conv_state=conv_state, + ffn_state=ffn_state + ) + self.states.append(state) + else: + state = self.states[layer_idx] + if recurrent_state is not None: + state['recurrent_state'] = recurrent_state + if attn_state is not None: + key_state, value_state = state['attn_state'] + if window_size is not None and key_state.shape[-2] == window_size: + # DO NOT allocate new memory if the cache is full + # roll the key/value states to the left by `input_size` + key_state = key_state.roll(-input_size, -2) + value_state = value_state.roll(-input_size, -2) + # replace the last `input_size` tokens with the new key/value states + key_state[..., -input_size:, :] = attn_state[0] + value_state[..., -input_size:, :] = attn_state[1] + attn_state = (key_state, value_state) + else: + attn_state = (torch.cat([key_state, attn_state[0]], -2), + torch.cat([value_state, attn_state[1]], -2),) + state['attn_state'] = attn_state + if conv_state is not None: + state['conv_state'] = conv_state + if ffn_state is not None: + state['ffn_state'] = ffn_state + + return state + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + if len(self.states) <= layer_idx: + return 0 + return self._seen_tokens + + def get_max_length(self) -> Optional[int]: + """Returns the maximum sequence length of the cached states. Cache does not have a maximum length.""" + return None + + def to_legacy_cache(self) -> Tuple: + return tuple(self.states) + + @classmethod + @torch.compiler.disable + def from_legacy_cache( + cls, + past_key_values: Optional[Tuple] = None, + seen_tokens: int = 0 + ) -> Cache: + """Converts a cache in the legacy cache format into an equivalent `Cache`.""" + + cache = cls(seen_tokens) + if isinstance(past_key_values, list): + for layer_idx in range(len(past_key_values)): + cache.states.append(past_key_values[layer_idx]) + return cache diff --git a/fla/modules/__init__.py b/fla/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..28dd0861bb8e2a12a296f620f90a62366144e575 --- /dev/null +++ b/fla/modules/__init__.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- + +from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution +from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear +from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss +from fla.modules.fused_kl_div import FusedKLDivLoss +from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss +from fla.modules.fused_linear_listnet_loss import FusedLinearListNetLoss +from fla.modules.fused_norm_gate import ( + FusedLayerNormGated, + FusedLayerNormSwishGate, + FusedLayerNormSwishGateLinear, + FusedRMSNormGated, + FusedRMSNormSwishGate, + FusedRMSNormSwishGateLinear +) +from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear +from fla.modules.mlp import GatedMLP +from fla.modules.rotary import RotaryEmbedding + +__all__ = [ + 'ImplicitLongConvolution', 'LongConvolution', 'ShortConvolution', + 'BitLinear', 'FusedBitLinear', + 'FusedCrossEntropyLoss', 'FusedLinearCrossEntropyLoss', 'FusedKLDivLoss', + 'GroupNorm', 'GroupNormLinear', 'LayerNorm', 'LayerNormLinear', 'RMSNorm', 'RMSNormLinear', + 'FusedLayerNormGated', 'FusedLayerNormSwishGate', 'FusedLayerNormSwishGateLinear', + 'FusedRMSNormGated', 'FusedRMSNormSwishGate', 'FusedRMSNormSwishGateLinear', + 'GatedMLP', + 'RotaryEmbedding' +] diff --git a/fla/modules/activations.py b/fla/modules/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..92f4212581aa5a27b52113deed9f3f685d4346af --- /dev/null +++ b/fla/modules/activations.py @@ -0,0 +1,471 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Tri Dao, Yu Zhang, Songlin Yang. + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, get_multiprocessor_count, input_guard + +sigmoid_fwd_codestring = """ +template T sigmoid_fwd(T x) { + return 1.0f / (1.0f + ::exp(-float(x))); +} +""" +sigmoid_bwd_codestring = """ +template T sigmoid_bwd(T x, T g) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + return float(g) * x_sigmoid * (1.0f - x_sigmoid); +} +""" + +sigmoid_fwd_jit_fn = torch.cuda.jiterator._create_jit_fn(sigmoid_fwd_codestring) +sigmoid_bwd_jit_fn = torch.cuda.jiterator._create_jit_fn(sigmoid_bwd_codestring) + + +@torch.compiler.disable +def sigmoid_fwd(x): + return sigmoid_fwd_jit_fn(x) + + +@torch.compiler.disable +def sigmoid_bwd(x, g): + return sigmoid_bwd_jit_fn(x, g) + + +class SigmoidFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return sigmoid_fwd(x) + + @staticmethod + def backward(ctx, dout): + x, = ctx.saved_tensors + return sigmoid_bwd(x, dout) + + +sigmoid = SigmoidFunction.apply + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=['D'] +) +@triton.jit +def logsigmoid_fwd_kernel( + x, + y, + temperature, + T: tl.constexpr, + D: tl.constexpr, + B: tl.constexpr +): + i = tl.program_id(0) + o_i = i * B + tl.arange(0, B) + m_i = o_i < T + + b_x = tl.load(x + o_i, mask=m_i, other=0.).to(tl.float32) + b_m = tl.minimum(0., b_x) + b_z = 1. + exp(-tl.abs(b_x)) + b_y = (b_m - log(b_z)) / temperature + tl.store(y + o_i, b_y.to(y.dtype.element_ty), mask=m_i) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=['D'] +) +@triton.jit +def logsigmoid_bwd_kernel( + x, + dx, + dy, + temperature, + T: tl.constexpr, + D: tl.constexpr, + B: tl.constexpr +): + i = tl.program_id(0) + o_i = i * B + tl.arange(0, B) + m_i = o_i < T + + b_x = tl.load(x + o_i, mask=m_i, other=0.).to(tl.float32) + b_dy = tl.load(dy + o_i, mask=m_i, other=0.).to(tl.float32) + b_dx = b_dy * (1. - tl.sigmoid(b_x)) / temperature + tl.store(dx + o_i, b_dx.to(dx.dtype.element_ty), mask=m_i) + + +def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1.) -> torch.Tensor: + T, D = x.numel(), x.shape[-1] + B = triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index))) + y = torch.empty_like(x) + logsigmoid_fwd_kernel[(triton.cdiv(T, B),)]( + x=x, + y=y, + temperature=temperature, + T=T, + D=D, + B=B + ) + return y + + +def logsigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, temperature: float = 1.) -> torch.Tensor: + T, D = x.numel(), x.shape[-1] + B = triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index))) + dx = torch.empty_like(x) + logsigmoid_bwd_kernel[(triton.cdiv(T, B),)]( + x=x, + dx=dx, + dy=dy, + temperature=temperature, + T=T, + D=D, + B=B + ) + return dx + + +class LogSigmoidFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward(ctx, x, temperature): + ctx.save_for_backward(x,) + ctx.temperature = temperature + return logsigmoid_fwd(x, temperature) + + @staticmethod + @input_guard + def backward(ctx, dy): + x, = ctx.saved_tensors + return logsigmoid_bwd(x, dy, ctx.temperature), None + + +def logsigmoid(x: torch.Tensor, temperature: float = 1.) -> torch.Tensor: + return LogSigmoidFunction.apply(x, temperature) + + +swish_fwd_codestring = """ +template T swish_fwd(T x) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + return float(x) * x_sigmoid; +} +""" +swish_bwd_codestring = """ +template T swish_bwd(T x, T g) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + return float(g) * x_sigmoid * (1.0f - float(x) * x_sigmoid + float(x)); +} +""" + +swish_fwd_jit_fn = torch.cuda.jiterator._create_jit_fn(swish_fwd_codestring) +swish_bwd_jit_fn = torch.cuda.jiterator._create_jit_fn(swish_bwd_codestring) + + +@torch.compiler.disable +def swish_fwd(x): + return swish_fwd_jit_fn(x) + + +@torch.compiler.disable +def swish_bwd(x, g): + return swish_bwd_jit_fn(x, g) + + +class SwishFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return swish_fwd(x) + + @staticmethod + def backward(ctx, dout): + x, = ctx.saved_tensors + return swish_bwd(x, dout) + + +swish = SwishFunction.apply + +# 1/sqrt(2*pi)-> 0.3989423 +# 1/sqrt(2) -> 0.70710678 +# sqrt(2/pi) -> 0.79788456 + + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.compile +def bias_gelu(y, bias): + x = bias + y + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=y.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.compile +def bias_gelu_bwd(g, y, bias): + """Assume that y has shape (B, D) and bias has shape (D)""" + x = bias + y + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + grad_y = ff * g + return grad_y.to(dtype=y.dtype), grad_y.sum(dim=(0), dtype=bias.dtype) + + +class GeLUFunction(torch.autograd.Function): + + @staticmethod + # bias is an optional argument + def forward(ctx, input, bias): + ctx.save_for_backward(input, bias) + return bias_gelu(input, bias) + + @staticmethod + def backward(ctx, grad_output): + input, bias = ctx.saved_tensors + tmp = bias_gelu_bwd(grad_output, input, bias) + return tmp, tmp + + +bias_gelu_impl = GeLUFunction.apply + + +# this function is tanh approximation of gelu +# actual gelu is: +# x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) +@torch.compile +def gelu_fwd(x): + return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype) + + +# gradient of tanh approximation of gelu +# gradient of actual gelu is: +# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) +@torch.compile +def gelu_bwd(g, x): + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + return (ff * g).to(dtype=x.dtype) + + +class FastGeLUFunction(torch.autograd.Function): + @staticmethod + # bias is an optional argument + def forward(ctx, input): + ctx.save_for_backward(input) + return gelu_fwd(input) + + @staticmethod + def backward(ctx, grad_output): + (input,) = ctx.saved_tensors + tmp = gelu_bwd(grad_output, input) + return tmp + + +fast_gelu_impl = FastGeLUFunction.apply + + +@torch.compile +def relu_bwd(g, x): + return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype) + + +@torch.compile +def sqrelu_fwd(x): + r = F.relu(x.float()) + return (r * r).to(dtype=x.dtype) + + +@torch.compile +def sqrelu_bwd(g, x): + return (2.0 * g * F.relu(x.float())).to(dtype=x.dtype) + + +class SquaredReLUFunction(torch.autograd.Function): + + @staticmethod + def forward(ctx, input): + ctx.save_for_backward(input) + return sqrelu_fwd(input) + + @staticmethod + def backward(ctx, grad_output): + input, = ctx.saved_tensors + return sqrelu_bwd(grad_output, input) + + +sqrelu = SquaredReLUFunction.apply + + +swiglu_fwd_codestring = """ +template T swiglu_fwd(T x, T y) { + return float(x) * float(y) / (1.0f + ::exp(-float(x))); +} +""" +swiglu_bwd_codestring = """ +template T swiglu_bwd(T x, T y, T g, T& dx, T& dy) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); + dy = float(x) * x_sigmoid * float(g); +} +""" + +swiglu_fwdbwd_codestring = """ +template T swiglu_fwdbwd(T x, T y, T g, T& dx, T& dy, T& z) { + float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x))); + float x_swish = float(x) * x_sigmoid; + dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y); + dy = x_swish * float(g); + z = x_swish * float(y); +} +""" + + +swiglu_fwd_jit_fn = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring) +swiglu_bwd_jit_fn = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2) +swiglu_fwdbwd_jit_fn = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_fwdbwd_codestring, num_outputs=3) + + +@torch.compiler.disable +def swiglu_fwd(x, y): + return swiglu_fwd_jit_fn(x, y) + + +@torch.compiler.disable +def swiglu_bwd(x, y, g): + return swiglu_bwd_jit_fn(x, y, g) + + +@torch.compiler.disable +def swiglu_fwdbwd(x, y, g): + return swiglu_fwdbwd_jit_fn(x, y, g) + + +@torch.compile +def swiglu_fwd_torch(x, y): + return (F.silu(x.float()) * y).to(x.dtype) + + +@torch.compile +def swiglu_bwd_torch(x, y, g): + dtype = x.dtype + x, y, g = x.float(), y.float(), g.float() + x_sigmoid = x.sigmoid() + x_swish = x * x_sigmoid + dx = x_sigmoid * (1 + x * (1.0 - x_sigmoid)) * g * y + dy = x_swish * g + return dx.to(dtype), dy.to(dtype) + + +@torch.compile +def swiglu_fwdbwd_torch(x, y, g): + dtype = x.dtype + x, y, g = x.float(), y.float(), g.float() + x_sigmoid = x.sigmoid() + x_swish = x * x_sigmoid + dx = x_sigmoid * (1 + x * (1.0 - x_sigmoid)) * g * y + dy = x_swish * g + z = x_swish * y + return dx.to(dtype), dy.to(dtype), z.to(dtype) + + +class SwiGLUFunction(torch.autograd.Function): + r""" + Swish-Gated Linear Unit (SwiGLU) function. + + .. math:: + \text{SwiGLU}(x, y) = swish(x) * y = \frac{x}{1 + \exp(-x)} * y + """ + + @staticmethod + def forward(ctx, x, y): + ctx.save_for_backward(x, y) + if torch.compiler.is_compiling() or isinstance(x, torch.distributed.tensor.DTensor): + return swiglu_fwd_torch(x, y) + else: + return swiglu_fwd(x, y) + + @staticmethod + def backward(ctx, dout): + x, y = ctx.saved_tensors + if torch.compiler.is_compiling() or isinstance(x, torch.distributed.tensor.DTensor): + return swiglu_bwd_torch(x, y, dout) + else: + return swiglu_bwd(x, y, dout) + + +class SwiGLULinearFunction(torch.autograd.Function): + r""" + Swish-Gated Linear Unit (SwiGLU) function followed by a linear transformation. + + .. math:: + \text{SwiGLULinear}(x, y, W, b) = (swish(x) * y) W + b + + This simple wrap discards the intermediate results of SwiGLU(x, y) to save memory. + """ + + @staticmethod + @autocast_custom_fwd + def forward(ctx, x, y, weight, bias): + with torch.no_grad(): + if torch.compiler.is_compiling() or isinstance(x, torch.distributed.tensor.DTensor): + z = swiglu_fwd_torch(x, y) + else: + z = swiglu_fwd(x, y) + out = F.linear(z, weight, bias) + # We don't store z, will be recomputed in the backward pass to save memory + ctx.save_for_backward(x, y, weight) + ctx.linear_bias_is_none = bias is None + return out + + @staticmethod + @autocast_custom_bwd + def backward(ctx, dout, *args): + x, y, weight = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dz = F.linear(dout, weight.t()).view_as(x) + with torch.no_grad(): + if torch.compiler.is_compiling() or isinstance(x, torch.distributed.tensor.DTensor): + dx, dy, z = swiglu_fwdbwd_torch(x, y, dz) + else: + dx, dy, z = swiglu_fwdbwd(x, y, dz) + dlinear_weight = torch.einsum("bo,bi->oi", dout, z.reshape(-1, z.shape[-1])) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + return dx, dy, dlinear_weight, dlinear_bias + + +swiglu = SwiGLUFunction.apply + + +swiglu_linear = SwiGLULinearFunction.apply + + +ACT2FN = { + 'relu': F.relu, + 'sigmoid': sigmoid, + 'logsigmoid': logsigmoid, + 'silu': swish, + 'swish': swish, + 'sqrelu': sqrelu, + 'gelu': fast_gelu_impl, + 'bias_gelu': bias_gelu_impl, +} diff --git a/fla/modules/convolution.py b/fla/modules/convolution.py new file mode 100644 index 0000000000000000000000000000000000000000..b29186e7094576b4867cc9c934979eb98820cdc3 --- /dev/null +++ b/fla/modules/convolution.py @@ -0,0 +1,434 @@ +# -*- coding: utf-8 -*- + +# from https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/convolution.py + +import math +import warnings +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + +from fla.modules.activations import ACT2FN +from fla.ops.common.utils import prepare_position_ids, prepare_sequence_ids +from fla.utils import checkpoint, input_guard + +try: + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update +except ImportError: + causal_conv1d_fn = None + causal_conv1d_update = None + + +def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None): + seqlen = u.shape[-1] + fft_size = 2 * seqlen + k_f = torch.fft.rfft(k, n=fft_size) / fft_size + if k_rev is not None: + k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size + k_f = k_f + k_rev_f.conj() + u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size) + + if len(u.shape) > 3: + k_f = k_f.unsqueeze(1) + y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen] + + out = y + u + if gelu: + out = F.gelu(out) + if dropout_mask is not None: + return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype) + else: + return out.to(dtype=u.dtype) + + +@checkpoint +def proj_then_conv1d( + x: torch.Tensor, + proj_weight: torch.Tensor, + conv1d_weight: torch.Tensor, + conv1d_bias: Optional[torch.Tensor] = None, + cache: Optional[torch.Tensor] = None +) -> torch.Tensor: + # We do matmul and transpose BLH -> HBL at the same time + x = rearrange(proj_weight @ rearrange(x, "b t d -> d (b t)"), "d (b t) -> b d t", t=x.shape[-2]) + + if causal_conv1d_fn is None: + raise ImportError("`causal_conv1d_fn` is not available. Please install `causal-conv1d` first.") + if cache is None: + x = causal_conv1d_fn( + x=x, + weight=rearrange(conv1d_weight, "d 1 w -> d w"), + bias=conv1d_bias, + activation="silu", + ).transpose(1, 2) + else: + assert x.shape[-1] == 1, "Only support decoding with 1 token at a time for now" + x = x.squeeze(-1) + x = causal_conv1d_update( + x=x, + weight=rearrange(conv1d_weight, "d 1 w -> d w"), + bias=conv1d_bias, + cache=cache, + activation="silu", + ) + return x + + +@triton.jit +def causal_conv1d_varlen_states_fwd_kernel( + x, + cache, + offsets, + D, + W, + BD: tl.constexpr, + BW: tl.constexpr +): + i_d, i_w, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) + eos = tl.load(offsets + i_n + 1) + bos = tl.maximum(tl.load(offsets + i_n), eos - W) + o_t = eos - (i_w + 1) * BW + tl.arange(0, BW) + o_d = i_d * BD + tl.arange(0, BD) + o_w = W - (i_w + 1) * BW + tl.arange(0, BW) + + b_x = tl.load(x + o_t * D + o_d[:, None], mask=(o_t >= bos) & (o_d[:, None] < D), other=0) + tl.store(cache + i_n * D*W + o_d[:, None] * W + o_w, b_x, mask=(o_d[:, None] < D) & (o_w >= 0)) + + +@input_guard +def causal_conv1d_varlen_states_fwd( + x: torch.Tensor, + cache: torch.Tensor, + cu_seqlens: torch.Tensor, + state_len: int +) -> torch.Tensor: + N, D, W = len(cu_seqlens) - 1, x.shape[-1], state_len + cache = torch.empty(N, D, W, dtype=x.dtype, device=x.device) if cache is None else cache + BD = min(triton.next_power_of_2(D), 256) + BW = min(triton.next_power_of_2(state_len), 16) + grid = (triton.cdiv(D, BD), triton.cdiv(W, BW), N) + with torch.cuda.device(x.device.index): + causal_conv1d_varlen_states_fwd_kernel[grid]( + x=x, + cache=cache, + offsets=cu_seqlens, + D=D, + W=W, + BW=BW, + BD=BD + ) + return cache + + +class ShortConvolution(nn.Conv1d): + """ + Simple wrapper around `nn.Conv1d` that accepts dimension last. + """ + + def __init__( + self, + hidden_size: int, + kernel_size: int, + bias: bool = False, + activation: Optional[str] = 'silu', + use_fast_conv1d: Optional[bool] = True, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + super().__init__( + in_channels=hidden_size, + out_channels=hidden_size, + kernel_size=kernel_size, + groups=hidden_size, + bias=bias, + padding=kernel_size - 1, + device=device, + dtype=dtype, + ) + + self.hidden_size = hidden_size + self.activation = None + if activation is not None: + assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet." + self.activation = activation + + if causal_conv1d_fn is None: + if use_fast_conv1d: + raise RuntimeError( + "Please either install `causal-conv1d>=1.4.0` to enable fast causal short convolution CUDA kernel " + "or set `use_fast_conv1d` to False" + ) + else: + warnings.warn( + "The naive Pytorch verison is very slow in practice, " + "please run `pip install causal-conv1d>=1.4.0` to install fast causal short convolution CUDA kernel", + category=ImportWarning + ) + self.use_fast_conv1d = use_fast_conv1d + + def extra_repr(self): + s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}' + ', stride={stride}') + if self.padding != (0,) * len(self.padding): + s += ', padding={padding}' + if self.dilation != (1,) * len(self.dilation): + s += ', dilation={dilation}' + if self.output_padding != (0,) * len(self.output_padding): + s += ', output_padding={output_padding}' + if self.groups != 1: + s += ', groups={groups}' + if self.bias is None: + s += ', bias=False' + if self.padding_mode != 'zeros': + s += ', padding_mode={padding_mode}' + if self.activation is not None: + s += ', activation={activation}' + if not self.use_fast_conv1d: + s += ', use_fast_conv1d={use_fast_conv1d}' + return s.format(**self.__dict__) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + cache: Optional[torch.Tensor] = None, + output_final_state: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (`torch.Tensor`): + Tensor of shape `[B, T, D]`. + If `seq_idx` is provided, `B` must be 1. + mask (`Optional[torch.Tensor]`): + Attention mask dealing with padded positions. + cache (`Optional[torch.Tensor]`): + Previous cache tensor of shape `[N, D, W]`, where `W` is the kernel size. + If provided, the cache is updated **inplace**. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, D, W]`. Default: `False`. + cu_seqlens (Optional[torch.LongTensor]): + Cumulative sequence lengths for each batch. Used for varlen. Default: `None`. + Shape: [B+1] + + Returns: + Tensor of shape `[B, T, D]`. + """ + + B, T, D, W = *x.shape, self.kernel_size[0] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + if mask is not None: + if cu_seqlens is not None: + raise ValueError("`mask` and `cu_seqlens` cannot be provided at the same time") + x = x.mul_(mask.unsqueeze(-1)) + if output_final_state and cache is None: + cache = x.new_zeros(N, D, W) + # during the decoding phase, we assume the batch is composed of sequences of length 1 + if cache is not None and B * T == N: + return self.step(x, cache, cu_seqlens) + + if cache is not None: + if cu_seqlens is not None: + cache = causal_conv1d_varlen_states_fwd(x, cache, cu_seqlens, W) + else: + cache[:, :, -min(W, T):].copy_(rearrange(x[..., -min(W, T):, :], 'n w d -> n d w')) + + x = rearrange(x, 'b t d -> b d t') + if self.use_fast_conv1d: + # Sequence index for each token. Used for varlen. + # Suppose a batch consists of two sequences with lengths 3 and 4, + # seq_idx=[0, 0, 0, 1, 1, 1, 1] for this batch. + # NOTE: No need to provide this arg if `cu_seqlens` is passed. + # This arg is just for BC, and will be removed in the future. + # [B, T] + seq_idx = kwargs.get('seq_idx', None) + if cu_seqlens is not None and seq_idx is None: + seq_idx = prepare_sequence_ids(prepare_position_ids(cu_seqlens)).to(torch.int32).unsqueeze(0) + x = causal_conv1d_fn( + x=x, + weight=rearrange(self.weight, "d 1 w -> d w"), + bias=self.bias, + activation=self.activation, + seq_idx=seq_idx, + ) + else: + if cu_seqlens is not None: + raise ValueError("`cu_seqlens` is not supported for the naive Pytorch version") + x = self._conv_forward(x, self.weight, self.bias)[..., :x.shape[-1]] + if self.activation is not None: + x = ACT2FN[self.activation](x) + return rearrange(x, "b d t -> b t d"), cache + + def step( + self, + x: torch.Tensor, + cache: torch.Tensor, + cu_seqlens: Optional[torch.LongTensor] = None + ): + shape = x.shape + x = x.squeeze(0) if cu_seqlens is not None else x.squeeze(1) + if self.use_fast_conv1d: + x = causal_conv1d_update( + x=x, + conv_state=cache, + weight=rearrange(self.weight, "d 1 w -> d w"), + bias=self.bias, + activation=self.activation, + ) + else: + dtype = x.dtype + # we follow the fast mode that updates the cache in-place + cache.copy_(cache.roll(shifts=-1, dims=-1)) + cache[:, :, -1] = x + x = torch.sum(cache * rearrange(self.weight, "d 1 w -> d w"), dim=-1) + if self.bias is not None: + x = x + self.bias + if self.activation is not None: + x = ACT2FN[self.activation](x).to(dtype=dtype) + return x.view(shape), cache + + @property + def state_size(self) -> int: + return self.hidden_size * self.kernel_size + + +class LongConvolution(nn.Module): + """ + LongConvolution applies a convolution operation on the input tensor using a fixed + filter of length max_len. + The filter is learned during training and is applied using FFT convolution. + Args: + hidden_size (int): The number of expected features in the input and output. + max_len (int): The maximum sequence length. + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + + def __init__( + self, + hidden_size: int, + max_len: int, + **kwargs, + ): + """ + Initializes the LongConvolution module. + Args: + hidden_size (int): The number of expected features in the input and output. + max_len (int): The maximum sequence length. + """ + super().__init__() + self.hidden_size = hidden_size + self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True) + + def forward(self, x: torch.Tensor, *args, **kwargs): + """ + Applies the LongConvolution operation on the input tensor. + Args: + x: [batch_size, seq_len, hidden_size] tensor + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + x = x.transpose(1, 2) + y = fft_conv(x, self.filter, dropout_mask=None, gelu=False) + y = y.transpose(1, 2) + return y.to(dtype=x.dtype) + + +class PositionalEmbedding(nn.Module): + def __init__(self, emb_dim: int, seq_len: int, **kwargs): + """Complex exponential positional embeddings for implicit long convolution filters.""" + super().__init__() + + self.seq_len = seq_len + # The time embedding fed to the filteres is normalized so that t_f = 1 + t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1 + + if emb_dim > 1: + bands = (emb_dim - 1) // 2 + # To compute the right embeddings we use the "proper" linspace + t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None] + w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1 + + f = torch.linspace(1e-4, bands - 1, bands)[None, None] + z = torch.exp(-1j * f * w) + z = torch.cat([t, z.real, z.imag], dim=-1) + self.z = nn.Parameter(z, requires_grad=False) + + def forward(self, L): + return self.z[:, :L] + + +class ImplicitLongConvolution(nn.Module): + """ + Long convolution with implicit filter parameterized by an MLP. + + Args: + hidden_size (int): + The number of expected features in the input and output. + max_len (int): + The maximum sequence length. + d_emb (Optional[int]): + The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine). + Defaults to 3. + d_hidden (Optional[int]): + The number of features in the hidden layer of the MLP. Defaults to 16. + + Attributes: + pos_emb (`PositionalEmbedding`): The positional embedding layer. + mlp (`nn.Sequential`): The MLP that parameterizes the implicit filter. + + """ + + def __init__( + self, + hidden_size: int, + max_len: int, + d_emb: int = 3, + d_hidden: int = 16, + **kwargs, + ): + """ + Long convolution with implicit filter parameterized by an MLP. + + + """ + super().__init__() + self.hidden_size = hidden_size + self.d_emb = d_emb + + assert ( + d_emb % 2 != 0 and d_emb >= 3 + ), "d_emb must be odd and greater or equal to 3 (time, sine and cosine)" + self.pos_emb = PositionalEmbedding(d_emb, max_len) + + # final linear layer + self.mlp = nn.Sequential( + nn.Linear(d_emb, d_hidden), + torch.nn.ReLU(), + nn.Linear(d_hidden, hidden_size), + ) + + def filter(self, seq_len: int, *args, **kwargs): + k = self.mlp(self.pos_emb(seq_len)) + + return k.transpose(1, 2) + + def forward(self, x: torch.Tensor, *args, **kwargs): + """ + Args: + x: [batch_size, seq_len, hidden_size] tensor + Returns: + y: [batch_size, seq_len, hidden_size] tensor + """ + x = x.transpose(1, 2) + k = self.filter(x.shape[-1]) + y = fft_conv(x, k, dropout_mask=None, gelu=False) + + y = y.transpose(1, 2) + return y.to(dtype=x.dtype) diff --git a/fla/modules/feature_map.py b/fla/modules/feature_map.py new file mode 100644 index 0000000000000000000000000000000000000000..6af81e74d3975f67b8df23c1dfa60cd01b5a4950 --- /dev/null +++ b/fla/modules/feature_map.py @@ -0,0 +1,300 @@ +# -*- coding: utf-8 -*- + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from fla.modules.activations import fast_gelu_impl, sigmoid, sqrelu, swish +from fla.modules.layernorm import layer_norm +from fla.utils import checkpoint + + +@checkpoint +def flatten_diag_outer_product(x, y): + z = torch.einsum("...i,...j->...ij", x, y) + N = z.size(-1) + indicies = torch.triu_indices(N, N) + return z[..., indicies[0], indicies[1]] + + +@checkpoint +def flatten_diag_outer_product_off1(x, y): + z = torch.einsum("...i,...j->...ij", x, y) + N = z.size(-1) + indicies = torch.triu_indices(N, N, 1) + indices2 = torch.arange(0, N) + return z[..., indicies[0], indicies[1]], z[..., indices2, indices2] + + +def is_power_of_2(n): + return (n & (n - 1) == 0) and n != 0 + + +class HedgehogFeatureMap(nn.Module): + + r""" + Hedgehog feature map as introduced in + `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry `_ + """ + + def __init__( + self, + head_dim: int + ) -> HedgehogFeatureMap: + super().__init__() + # Trainable map + self.layer = nn.Linear(head_dim, head_dim) + self.init_weights_() + + def init_weights_(self): + """Initialize trainable map as identity""" + with torch.no_grad(): + identity = torch.eye(*self.layer.weight.shape[-2:], dtype=torch.float) + self.layer.weight.copy_(identity.to(self.layer.weight)) + nn.init.zeros_(self.layer.bias) + + def forward(self, x: torch.Tensor): + x = self.layer(x) # shape b, h, l, d + return torch.cat([2*x, -2*x], dim=-1).softmax(-1) + + +class T2RFeatureMap(nn.Module): + + r""" + Simple linear mapping feature map as in + `Finetuning Pretrained Transformers into RNNs `_ + """ + + def __init__( + self, + head_dim: int, + dot_dim: int = None, + bias: Optional[bool] = False + ) -> T2RFeatureMap: + super().__init__() + # Trainable map + if dot_dim is None: + dot_dim = head_dim + + self.head_dim = head_dim + self.dot_dim = dot_dim + self.bias = bias + + self.layer = nn.Linear(head_dim, dot_dim, bias=bias) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(head_dim={self.head_dim}, dot_dim={self.dot_dim}, bias={self.bias})" + + def forward(self, x: torch.Tensor): + return self.layer(x).relu() + + +class DPFPFeatureMap(nn.Module): + + r""" + Deterministic Parameter-Free Projection (DPFP) feature map in + `Linear Transformers Are Secretly Fast Weight Programmers `_ + """ + + def __init__( + self, + head_dim: int, + nu: int = 4 + ) -> DPFPFeatureMap: + super().__init__() + self.nu = nu + + def forward(self, x: torch.Tensor): + x = torch.cat([x.relu(), -x.relu()], dim=-1) + x_rolled = torch.cat([x.roll(shifts=j, dims=-1) for j in range(1, self.nu+1)], dim=-1) + x_repeat = torch.cat([x] * self.nu, dim=-1) + return x_repeat * x_rolled + + +class HadamardFeatureMap(nn.Module): + def __init__( + self, + head_dim: int + ) -> HadamardFeatureMap: + super().__init__() + # Trainable map + self.layer1 = nn.Linear(head_dim, head_dim) + self.layer2 = nn.Linear(head_dim, head_dim) + + def forward(self, x: torch.Tensor): + return self.layer1(x) * self.layer2(x) + + +class LearnableOuterProductFeatureMap(nn.Module): + def __init__( + self, + head_dim: int, + feature_dim: int + ) -> LearnableOuterProductFeatureMap: + super().__init__() + # Trainable map + self.layer1 = nn.Linear(head_dim, feature_dim, bias=False) + self.layer2 = nn.Linear(head_dim, feature_dim, bias=False) + self.normalizer = feature_dim ** -0.5 + + def forward(self, x: torch.Tensor): + return flatten_diag_outer_product(self.layer1(x), self.layer2(x)) + + +class LearnablePolySketchNonNegativeFeatureMap(nn.Module): + + def __init__( + self, + head_dim: int, + sketch_size: Optional[int] = None, + degree: Optional[int] = 2 + ) -> LearnablePolySketchNonNegativeFeatureMap: + super().__init__() + + assert is_power_of_2(degree) and degree >= 2, f"The degree {degree} must be a power of 2" + + self.head_dim = head_dim + self.sketch_size = sketch_size if sketch_size is not None else head_dim + self.degree = degree + + self.gamma = nn.Parameter(torch.ones(head_dim)) + self.beta = nn.Parameter(torch.zeros(head_dim)) + # NOTE: the sketch layers defined here are quite different from the original paper + # currently we simply use linear layers without any non-linear activations + self.sketches1 = nn.ModuleList([ + nn.Linear(head_dim, sketch_size, bias=False), + *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)] + ]) + self.sketches2 = nn.ModuleList([ + nn.Linear(head_dim, sketch_size, bias=False), + *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)] + ]) + + def forward(self, x: torch.Tensor): + # Section 2.1 + x = layer_norm(x, self.gamma, self.beta) + # first map the input to sketch size with learnable parameters + x = self.sketches1[0](x) * self.sketches2[0](x) * self.head_dim ** -0.5 + for i in range(1, int(math.log2(self.degree)) - 1): + x = self.sketches1[i](x) * self.sketches2[i](x) * self.head_dim ** -0.5 + # do sketch mapping for log2(p) - 1 times in total + # do p=2 mapping to ensure non-negativity + return flatten_diag_outer_product(x, x) + + +class TaylorFeatureMap(nn.Module): + def __init__( + self, + head_dim: int + ) -> TaylorFeatureMap: + super().__init__() + self.head_dim = head_dim + self.r2 = math.sqrt(2) + self.rd = math.sqrt(self.head_dim) + self.rrd = math.sqrt(self.rd) + + def forward(self, x: torch.Tensor): + x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) + return torch.cat([torch.ones_like(x[..., 0:1]), x / self.rrd, x2_2 / (self.rd * self.r2), x2_1 / self.rd], dim=-1) + + +class RebasedFeatureMap(nn.Module): + + def __init__( + self, + head_dim: int, + use_gamma: Optional[bool] = True, + use_beta: Optional[bool] = True, + normalize: Optional[bool] = True + ) -> RebasedFeatureMap: + super().__init__() + + self.head_dim = head_dim + self.use_gamma = use_gamma + self.use_beta = use_beta + self.normalize = normalize + + self.gamma = None + self.beta = None + if use_gamma: + self.gamma = nn.Parameter(torch.ones(head_dim)) + if use_beta: + self.beta = nn.Parameter(torch.zeros(head_dim)) + + def forward(self, x: torch.Tensor, flatten: Optional[bool] = True): + if self.use_beta and self.use_gamma and self.normalize: + x = layer_norm(x, self.gamma, self.beta) + elif self.normalize: + x = F.layer_norm(x, (self.head_dim,), self.gamma, self.beta) + elif self.use_gamma and self.use_beta: + x = torch.addcmul(self.beta, x, self.gamma) + elif self.use_gamma: + x = x.mul(self.gamma) + else: + raise RuntimeError(f"Not supported combination of `use_gamma`, `use_beta` and `normalize`, " + f"which is currentlt set as (`{self.use_gamma}`, `{self.use_beta}`, `{self.normalize}`)") + if not flatten: + return x + x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) + # rebased use learnable parameters to approximate any quadratic function + return torch.cat([x2_2 * self.head_dim ** -0.5, x2_1 * (2 / self.head_dim) ** 0.5], dim=-1) + + +class ReLUFeatureMap(nn.Module): + + def __init__( + self, + ) -> ReLUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return F.relu(x) + + +class SquaredReLUFeatureMap(nn.Module): + + def __init__( + self, + ) -> SquaredReLUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return sqrelu(x) + + +class GELUFeatureMap(nn.Module): + + def __init__( + self, + ) -> GELUFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return fast_gelu_impl(x) + + +class SwishFeatureMap(nn.Module): + + def __init__( + self, + ) -> SwishFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return swish(x) + + +class SigmoidFeatureMap(nn.Module): + + def __init__( + self, + ) -> SigmoidFeatureMap: + super().__init__() + + def forward(self, x: torch.Tensor): + return sigmoid(x) diff --git a/fla/modules/fused_bitlinear.py b/fla/modules/fused_bitlinear.py new file mode 100644 index 0000000000000000000000000000000000000000..d05928eaf6b721f29a4d15967cae3a8e014e7c9c --- /dev/null +++ b/fla/modules/fused_bitlinear.py @@ -0,0 +1,638 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# Implementations of BitLinear layer with fused LayerNorm and quantized Linear layer. +# [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764) +# [Scalable MatMul-free Language Modeling](https://arxiv.org/abs/2406.02528) + +# Code adapted from https://github.com/ridgerchu/matmulfreellm/ + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.modules.layernorm import RMSNorm +from fla.utils import get_multiprocessor_count, input_guard, require_version + + +def activation_quant(x): + """ + Per-token quantization to 8 bits. No grouping is needed for quantization. + + Args: + x: An activation tensor with shape [n, d]. + + Returns: + A quantized activation tensor with shape [n, d]. + """ + # Compute the scale factor + scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) + # Quantize and then de-quantize the tensor + y = (x * scale).round().clamp_(-128, 127) / scale + return y + + +def weight_quant(w): + """ + Per-tensor quantization to 1.58 bits. No grouping is needed for quantization. + + Args: + w: A weight tensor with shape [d, k]. + + Returns: + A quantized weight tensor with shape [d, k]. + """ + # Compute the scale factor + scale = 1.0 / w.abs().mean().clamp_(min=1e-5) + # Quantize and then de-quantize the tensor + u = (w * scale).round().clamp_(-1, 1) / scale + return u + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + triton.Config({}, num_warps=16), + triton.Config({}, num_warps=32), + ], + key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS"], +) +@triton.jit +def layer_norm_fwd_kernel_quant( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + RESIDUAL, # pointer to the residual + RESIDUAL_OUT, # pointer to the residual + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_res_row, + stride_res_out_row, + N, # number of columns in X + eps, # epsilon to avoid division by zero + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + X += row * stride_x_row + Y += row * stride_y_row + if HAS_RESIDUAL: + RESIDUAL += row * stride_res_row + if STORE_RESIDUAL_OUT: + RESIDUAL_OUT += row * stride_res_out_row + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) + x += residual + if STORE_RESIDUAL_OUT: + tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + + y = x_hat * w if HAS_WEIGHT else x_hat + if HAS_BIAS: + y = y + b + + # Aply quantization to the output + scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5) + # Quantize and then de-quantize the tensor + y = tl.extra.cuda.libdevice.round(y * scale) + y = tl.maximum(tl.minimum(y, 127), -128) / scale + + # Write output + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_fwd_quant( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False +): + if residual is not None: + residual_dtype = residual.dtype + M, N = x.shape + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + residual_out = torch.empty(M, N, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + layer_norm_fwd_kernel_quant[(M,)]( + x, + y, + weight, + bias, + residual, + residual_out, + mean, + rstd, + x.stride(0), + y.stride(0), + residual.stride(0) if residual is not None else 0, + residual_out.stride(0) if residual_out is not None else 0, + N, + eps, + is_rms_norm, + BLOCK_N, + residual is not None, + residual_out is not None, + weight is not None, + bias is not None, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +@triton.heuristics({ + "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + triton.Config({}, num_warps=4), + triton.Config({}, num_warps=8), + triton.Config({}, num_warps=16), + triton.Config({}, num_warps=32), + ], + key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS"], +) +@triton.jit +def layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DRESIDUAL, + DRESIDUAL_IN, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dres_row, + stride_dres_in_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + rows_per_program, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + row_start = row_block_id * rows_per_program + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += row_start * stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += row_start * stride_dres_in_row + DY += row_start * stride_dy_row + DX += row_start * stride_dx_row + if RECOMPUTE_OUTPUT: + Y += row_start * stride_y_row + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.0) + if RECOMPUTE_OUTPUT: + y = xhat * w if HAS_WEIGHT else xhat + if HAS_BIAS: + y = y + b + + # Aply quantization to the output + scale = 127.0 / tl.maximum(tl.max(tl.abs(y), 0), 1e-5) + # Quantize and then de-quantize the tensor + y = tl.extra.cuda.libdevice.round(y * scale) + y = tl.maximum(tl.minimum(y, 127), -128) / scale + + tl.store(Y + cols, y, mask=mask) + wdy = dy + if HAS_WEIGHT: + wdy = dy * w + dw += dy * xhat + if HAS_BIAS: + db += dy + if not IS_RMS_NORM: + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + c1 = tl.sum(xhat * wdy, axis=0) / N + dx = (wdy - xhat * c1) * rstd + if HAS_DRESIDUAL: + dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) + dx += dres + # Write dx + if STORE_DRESIDUAL: + tl.store(DRESIDUAL_IN + cols, dx, mask=mask) + tl.store(DX + cols, dx, mask=mask) + + X += stride_x_row + if HAS_DRESIDUAL: + DRESIDUAL += stride_dres_row + if STORE_DRESIDUAL: + DRESIDUAL_IN += stride_dres_in_row + if RECOMPUTE_OUTPUT: + Y += stride_y_row + DY += stride_dy_row + DX += stride_dx_row + if HAS_WEIGHT: + tl.store(DW + row_block_id * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * N + cols, db, mask=mask) + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + dresidual: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, +): + M, N = x.shape + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device) + dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + sm_count = get_multiprocessor_count(x.device.index) + _dw = torch.empty((sm_count, N), dtype=torch.float32, device=weight.device) if weight is not None else None + _db = torch.empty((sm_count, N), dtype=torch.float32, device=bias.device) if bias is not None else None + rows_per_program = math.ceil(M / sm_count) + grid = (sm_count,) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + y, + dy, + dx, + _dw, + _db, + dresidual, + dresidual_in, + mean, + rstd, + x.stride(0), + 0 if not recompute_output else y.stride(0), + dy.stride(0), + dx.stride(0), + dresidual.stride(0) if dresidual is not None else 0, + dresidual_in.stride(0) if dresidual_in is not None else 0, + M, + N, + eps, + rows_per_program, + is_rms_norm, + BLOCK_N, + dresidual is not None, + dresidual_in is not None, + weight is not None, + bias is not None, + ) + dw = _dw.sum(0).to(weight.dtype) if weight is not None else None + db = _db.sum(0).to(bias.dtype) if bias is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype: + dresidual_in = dx + return (dx, dw, db, dresidual_in) if not recompute_output else (dx, dw, db, dresidual_in, y) + + +class LayerNormLinearQuantFn(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + ): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None) + y, mean, rstd, residual_out = layer_norm_fwd_quant( + x, + norm_weight, + norm_bias, + eps, + residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = weight_quant(linear_weight).to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd( + dy, + x, + norm_weight, + norm_bias, + ctx.eps, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y) + return ( + dx.reshape(ctx.x_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +def layer_norm_linear_quant_fn( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-6, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, +): + return LayerNormLinearQuantFn.apply( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + ) + + +def rms_norm_linear_quant( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False +): + return layer_norm_linear_quant_fn( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True + ) + + +@require_version("triton>=3.0", "Triton >= 3.0 is required to do online quantization.") +def bit_linear(x, weight, bias=None, norm_weight=None, norm_bias=None, eps=1e-8): + """ + A functional version of BitLinear that applies quantization to activations and weights. + + Args: + x: Input tensor with shape [n, d]. + weight: Weight tensor with shape [out_features, in_features]. + bias: Bias tensor with shape [out_features] (optional). + norm_weight: Weight tensor for RMS normalization with shape [in_features]. + norm_bias: Bias tensor for RMS normalization with shape [in_features]. + eps: A small constant for numerical stability in normalization. + + Returns: + Output tensor with shape [n, out_features]. + """ + return layer_norm_linear_quant_fn( + x, + norm_weight, + norm_bias, + weight, + bias, + is_rms_norm=True + ) + + +class BitLinear(nn.Linear): + """ + A custom linear layer that applies quantization on both activations and weights. + This is primarily for training; kernel optimization is needed for efficiency in deployment. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + norm_eps: float = 1e-8 + ): + """ + Initializes the BitLinear layer. + + Args: + in_features: Size of each input sample. + out_features: Size of each output sample. + bias: If set to False, the layer will not learn an additive bias. Default: True. + """ + # Initialize the superclass nn.Linear with the given parameters + super(BitLinear, self).__init__(in_features, out_features, bias=bias) + + self.norm = RMSNorm(in_features, eps=norm_eps) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({super().extra_repr()}, norm_eps={self.norm.eps})" + + def forward(self, x): + """ + Overrides the forward pass to include quantization. + + Args: + x: An input tensor with shape [n, d]. + + Returns: + An output tensor with shape [n, d]. + """ + # Weight tensor + w = self.weight + + # Apply RMS normalization to the input + x_norm = self.norm(x) + + # Apply quantization to both activations and weights + # Uses Straight-Through Estimator (STE) trick with .detach() for gradient flow + x_quant = x_norm + (activation_quant(x_norm) - x_norm).detach() + w_quant = w + (weight_quant(w) - w).detach() + # Perform linear operation with quantized values + y = F.linear(x_quant, w_quant) + + return y + + +class FusedBitLinear(BitLinear): + """ + A custom linear layer that applies quantization on both activations and weights. + This is primarily for training; kernel optimization is needed for efficiency in deployment. + """ + + def __init__(self, in_features, out_features, bias=False): + """ + Initializes the BitLinear layer. + + Args: + in_features: Size of each input sample. + out_features: Size of each output sample. + bias: If set to False, the layer will not learn an additive bias. Default: True. + """ + # Initialize the superclass nn.Linear with the given parameters + super(FusedBitLinear, self).__init__(in_features, out_features, bias=bias) + + def forward(self, x): + return layer_norm_linear_quant_fn( + x, + self.norm.weight, + self.norm.bias, + self.weight, + self.bias, + is_rms_norm=True + ) diff --git a/fla/modules/fused_cross_entropy.py b/fla/modules/fused_cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..f85091f66fe5539d4d6c68ca801b3b51ac8b94e4 --- /dev/null +++ b/fla/modules/fused_cross_entropy.py @@ -0,0 +1,419 @@ +# -*- coding: utf-8 -*- + +# Copyright (c) 2023, Tri Dao. + +from typing import Any, Tuple + +import torch +import torch.nn as nn +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import input_guard + +# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for +# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent +# version of PyTorch. The following 2 lines are for backward compatibility with +# older PyTorch. +if "all_gather_into_tensor" not in dir(torch.distributed): + torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base + + +@triton.heuristics({ + "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0, +}) +@triton.jit +def cross_entropy_fwd_kernel( + loss_ptr, # data ptrs + lse_ptr, + z_loss_ptr, + logits_ptr, + labels_ptr, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + n_rows, + logits_row_stride, # strides + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, + # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE + SPLIT: tl.constexpr, +): + row_idx = tl.program_id(0) + col_block_idx = tl.program_id(1) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + label_idx = tl.load(labels_ptr + row_idx) + logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")) + logits = logits.to(tl.float32) * logit_scale + max_logits = tl.max(logits, 0) + if HAS_SMOOTHING: + sum_logits = tl.sum(tl.where(col_offsets < n_cols, logits, 0.0), 0) + lse = log(tl.sum(exp(logits - max_logits), 0)) + max_logits + tl.store(lse_ptr + col_block_idx * n_rows + row_idx, lse) + if label_idx == ignore_index: + loss = 0.0 + z_loss = 0.0 + else: + label_idx -= class_start_idx + if label_idx >= col_block_idx * BLOCK_SIZE and label_idx < min( + n_cols, (col_block_idx + 1) * BLOCK_SIZE + ): + logits_label = tl.load(logits_ptr + label_idx) * logit_scale + if HAS_SMOOTHING: + loss = ( + (lse if not SPLIT else 0.0) + - label_smoothing * sum_logits / total_classes + - (1 - label_smoothing) * logits_label + ) + else: + loss = (lse if not SPLIT else 0.0) - logits_label + else: + # If label is out of bounds, we set the CE loss to 0.0. But we still want the label_smoothing loss + if HAS_SMOOTHING: + loss = label_smoothing * ((lse if not SPLIT else 0.0) - sum_logits / total_classes) + else: + loss = 0.0 + if not SPLIT: + z_loss = lse_square_scale * lse * lse + loss += z_loss + else: + z_loss = 0.0 + tl.store(loss_ptr + col_block_idx * n_rows + row_idx, loss) + if not SPLIT: + tl.store(z_loss_ptr + col_block_idx * n_rows + row_idx, z_loss) + + +@triton.heuristics({ + "HAS_SMOOTHING": lambda args: args["label_smoothing"] > 0.0, +}) +@triton.jit +def cross_entropy_bwd_kernel( + dlogits_ptr, # data ptrs + dloss_ptr, + logits_ptr, + lse_ptr, + labels_ptr, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, # Useful for tensor parallel when each rank only has a subset of classes + n_cols, # shapes + logits_row_stride, # strides + dlogits_row_stride, + dloss_row_stride, + BLOCK_SIZE: tl.constexpr, + HAS_SMOOTHING: tl.constexpr, +): + row_idx = tl.program_id(0) + col_block_idx = tl.program_id(1) + logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64) + dlogits_ptr = dlogits_ptr + row_idx * dlogits_row_stride.to(tl.int64) + col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + label_idx = tl.load(labels_ptr + row_idx) + if label_idx != ignore_index: + dloss = tl.load(dloss_ptr + row_idx * dloss_row_stride) + else: + dloss = 0.0 + logits = tl.load(logits_ptr + col_offsets, mask=col_offsets < n_cols, other=-float("inf")).to( + tl.float32 + ) * logit_scale + lse = tl.load(lse_ptr + row_idx) + probs = exp(logits - lse) + probs += 2.0 * lse_square_scale * lse * probs + label_idx -= class_start_idx + if HAS_SMOOTHING: + smooth_negative = label_smoothing / total_classes + probs = tl.where(col_offsets == label_idx, probs - (1 - label_smoothing), probs) - smooth_negative + else: + probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) + tl.store(dlogits_ptr + col_offsets, (dloss * logit_scale) * probs, mask=col_offsets < n_cols) + + +def fused_cross_entropy_forward( + logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + ignore_index: int = -100, + process_group=None, +): + n_rows, n_cols = logits.shape + assert target.shape == (n_rows,) + world_size = 1 if process_group is None else torch.distributed.get_world_size(process_group) + total_classes = world_size * n_cols + rank = 0 if process_group is None else torch.distributed.get_rank(process_group) + class_start_idx = rank * n_cols + + if logits.stride(-1) != 1: + logits = logits.contiguous() + # Set these similar to https://github.com/openai/triton/blob/main/python/tutorials/02-fused-softmax.py + MAX_BLOCK_SIZE = 64 * 1024 + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), MAX_BLOCK_SIZE) + num_warps = ( + 4 + if BLOCK_SIZE < 2048 + else (8 if BLOCK_SIZE < 8192 else (16 if BLOCK_SIZE < 128 * 1024 else 32)) + ) + # We may split the lse computation across multiple blocks, then do a reduction + # lse(local_lse) to get the final LSE. This is faster for large n_cols (e.g., > 64k) + # where having just one thread block processing more than 64k elements is slow. + split = world_size > 1 or n_cols > MAX_BLOCK_SIZE + n_splits = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE + loss_shape = (n_splits, n_rows) if n_splits > 1 else (n_rows,) + losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + lse = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + z_losses = torch.empty(*loss_shape, dtype=torch.float, device=logits.device) + + cross_entropy_fwd_kernel[(n_rows, n_splits)]( + losses, # data ptrs + lse, + z_losses, + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + total_classes, + class_start_idx, + n_cols, # shapes + n_rows, + logits.stride(0), # strides + BLOCK_SIZE=BLOCK_SIZE, # constants + num_warps=num_warps, + SPLIT=split + ) + + if split: + # If there's no label_smoothing, if target are in the vocab of this partition, losses contains + # - predicted logit, and 0 otherwise. + # If there's label_smoothing=0.1, for target in the vocab of this partition, losses contains + # -0.9 * predicted logit - 0.1 * sum logit / total_classes. + # For target not in the vocab of this partition, losses contains + # -0.1 * sum logit / total_classes. + if n_splits > 1: + lse = torch.logsumexp(lse, dim=0) + losses = losses.sum(dim=0) + if world_size > 1: + lse_allgather = torch.empty(world_size, n_rows, dtype=lse.dtype, device=lse.device) + torch.distributed.all_gather_into_tensor(lse_allgather, lse, group=process_group) + handle_losses = torch.distributed.all_reduce( + losses, op=torch.distributed.ReduceOp.SUM, group=process_group, async_op=True + ) + lse = torch.logsumexp(lse_allgather, dim=0) + handle_losses.wait() + # After the allreduce, if there's no label_smoothing, the total losses are - predicted_logit, + # we just have to add the (global) lse. + # If there's label_smoothing=0.1, the total losses are + # -0.9 * predicted_logit - 0.1 * sum logit / total_classes. + # Again, we just have to add the (global) lse. + losses += lse + if lse_square_scale != 0.0: + z_losses = lse_square_scale * lse.square() + z_losses.masked_fill_(target == ignore_index, 0.0) + losses += z_losses + else: + z_losses = torch.zeros_like(losses) + losses.masked_fill_(target == ignore_index, 0.0) + + return losses, z_losses, lse, total_classes, class_start_idx + + +class CrossEntropyLossFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + logits, + target, + label_smoothing=0.0, + logit_scale=1.0, + lse_square_scale=0.0, + ignore_index=-100, + inplace_backward=False, + process_group=None, + ): + losses, z_losses, lse, total_classes, class_start_idx = fused_cross_entropy_forward( + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + process_group, + ) + ctx.save_for_backward(logits, lse, target) + ctx.mark_non_differentiable(z_losses) + ctx.label_smoothing = label_smoothing + ctx.logit_scale = logit_scale + ctx.lse_square_scale = lse_square_scale + ctx.ignore_index = ignore_index + ctx.total_classes = total_classes + ctx.class_start_idx = class_start_idx + ctx.inplace_backward = inplace_backward + + return losses, z_losses + + @staticmethod + @input_guard + def backward(ctx, grad_losses, grad_z_losses): + del grad_z_losses # z_losses are only for logging. + + logits, lse, target = ctx.saved_tensors + dlogits = logits if ctx.inplace_backward else torch.empty_like(logits) + n_rows, n_cols = logits.shape + BLOCK_SIZE = min(triton.next_power_of_2(n_cols), 4 * 1024) + num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16) + def grid(META): return (n_rows, triton.cdiv(n_cols, META["BLOCK_SIZE"])) # noqa + cross_entropy_bwd_kernel[grid]( + dlogits, # data ptrs + grad_losses, + logits, + lse, + target, + ctx.label_smoothing, + ctx.logit_scale, + ctx.lse_square_scale, + ctx.ignore_index, + ctx.total_classes, + ctx.class_start_idx, + n_cols, # shapes + logits.stride(0), # strides + dlogits.stride(0), + grad_losses.stride(0), + BLOCK_SIZE=BLOCK_SIZE, # constants + num_warps=num_warps, + ) + return dlogits, None, None, None, None, None, None, None, None + + +def cross_entropy_loss( + logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + ignore_index=-100, + inplace_backward: bool = False, + process_group=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Arguments: + logits: [batch, vocab_size] + target: [batch,] + label_smoothing: float + logit_scale: float. + Multiply logits by this scale before calculating the loss. + lse_square_scale: float. + If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. + This is also referred to as "z-loss". + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + inplace_backward: bool. + If True, we do the backward pass in-place by modifying the logits. + This saves memory. + process_group: + if not None, we're doing Tensor Parallel: each process is responsible for + one part of the vocab. The loss will be aggregated across processes. + Returns: + losses: [batch,], float + z_losses: [batch,], float + """ + return CrossEntropyLossFunction.apply( + logits, + target, + label_smoothing, + logit_scale, + lse_square_scale, + ignore_index, + inplace_backward, + process_group, + ) + + +class FusedCrossEntropyLoss(nn.Module): + def __init__( + self, + ignore_index: int = -100, + reduction: str = "mean", + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + lse_square_scale: float = 0.0, + inplace_backward: bool = False, + process_group: Any = None, + return_z_loss: bool = False, + ): + """ + Arguments: + ignore_index: int. If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + lse_square_scale: float. If > 0, we add lse_square_scale * lse(logits) ^ 2 to the loss. + This is also referred to as "z-loss". + inplace_backward: bool. If True, we do the backward pass in-place by modifying the logits. + This saves memory. + process_group: if not None, we're doing Tensor Parallel: each process is responsible for + one part of the vocab. The loss will be aggregated across processes. + return_z_loss: bool. If True, we return the component of the loss contributed by + the lse_square_scale value. This value is only for logging and does not support + backprop. + """ + super().__init__() + if reduction not in ["mean", "none", "sum"]: + raise NotImplementedError("Only support reduction = 'mean' or 'none' or 'sum'") + self.ignore_index = ignore_index + self.reduction = reduction + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.lse_square_scale = lse_square_scale + self.inplace_backward = inplace_backward + self.process_group = process_group + self.return_z_loss = return_z_loss + + def forward(self, input, target): + """ + Arguments: + input: (batch, vocab_size) + target: (batch,) + Returns: + losses: (batch,) if reduction is 'none', else (1,), dtype float + z_loss: (batch,) if reduction is 'none', else (1,), dtype float (if self.return_z_loss) + """ + assert input.is_cuda and target.is_cuda, "Only support CUDA tensors" + loss, z_loss = cross_entropy_loss( + input, + target, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + lse_square_scale=self.lse_square_scale, + ignore_index=self.ignore_index, + inplace_backward=self.inplace_backward, + process_group=self.process_group, + ) + if self.reduction == "mean": + loss = loss.sum() / (target != self.ignore_index).sum() + elif self.reduction == "sum": + loss = loss.sum() + else: + loss = loss + + if not self.return_z_loss: + return loss + + if self.reduction == "mean": + z_loss = z_loss.sum() / (target != self.ignore_index).sum() + elif self.reduction == "sum": + z_loss = z_loss.sum() + else: + z_loss = z_loss + + return loss, z_loss diff --git a/fla/modules/fused_kl_div.py b/fla/modules/fused_kl_div.py new file mode 100644 index 0000000000000000000000000000000000000000..5e49269dec9e4c09d058c0ac0d5e6e059c6240b8 --- /dev/null +++ b/fla/modules/fused_kl_div.py @@ -0,0 +1,323 @@ +# -*- coding: utf-8 -*- + +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import input_guard + +# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 +# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 +# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling +# The optimal maximum block size depends on your hardware, your kernel, and your dtype +MAX_FUSED_SIZE = 65536 // 2 + + +@triton.jit +def kl_div_kernel( + logits, + target_logits, + loss, + s_logits, + s_loss, + reduction: tl.constexpr, + N: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr +): + # https://github.com/triton-lang/triton/issues/1058 + # If N*V is too large, i_n * stride will overflow out of int32, so we convert to int64 + i_n = tl.program_id(0).to(tl.int64) + + logits += i_n * s_logits + target_logits += i_n * s_logits + + # m is the max value. use the notation from the paper + sm = float('-inf') + tm = float('-inf') + # d is the sum. use the notation from the paper + sd, td = 0.0, 0.0 + + NV = tl.cdiv(V, BV) + for iv in range(0, NV): + o_x = iv * BV + tl.arange(0, BV) + # for student + b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf')) + b_sm = tl.max(b_sl) + m_new = tl.maximum(sm, b_sm) + sd = sd * exp(sm - m_new) + tl.sum(exp(b_sl - m_new)) + sm = m_new + # for teacher + b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf')) + b_tm = tl.max(b_tl) + m_new = tl.maximum(tm, b_tm) + td = td * exp(tm - m_new) + tl.sum(exp(b_tl - m_new)) + tm = m_new + + b_loss = 0. + # KL(y_true || y) = exp(y_true) * (log(y_true) - log(y)) + for iv in range(0, NV): + o_x = iv * BV + tl.arange(0, BV) + b_sl = tl.load(logits + o_x, mask=o_x < V, other=float('-inf')) + b_tl = tl.load(target_logits + o_x, mask=o_x < V, other=float('-inf')) + b_sp_log = b_sl - sm - log(sd) + b_tp_log = b_tl - tm - log(td) + b_sp = exp(b_sp_log) + b_tp = exp(b_tp_log) + b_kl = tl.where(o_x < V, b_tp * (b_tp_log - b_sp_log), 0) + b_dl = -b_tp + b_sp + b_loss += tl.sum(b_kl) + if reduction == 'batchmean': + b_dl = b_dl / N + tl.store(logits + o_x, b_dl, mask=o_x < V) + + # Normalize the loss by the number of elements if reduction is 'batchmean' + if reduction == 'batchmean': + b_loss = b_loss / N + + tl.store(loss + i_n * s_loss, b_loss) + + +@triton.jit +def elementwise_mul_kernel( + x, + g, + N: tl.constexpr, + B: tl.constexpr +): + """ + This function multiplies each element of the tensor pointed by x with the value pointed by g. + The multiplication is performed in-place on the tensor pointed by x. + + Parameters: + x: + Pointer to the input tensor. + g: + Pointer to the gradient output value. + N (int): + The number of columns in the input tensor. + B (int): + The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + i_x = tl.program_id(0).to(tl.int64) + o_x = i_x * B + tl.arange(0, B) + + # Load the gradient output value + b_g = tl.load(g) + b_x = tl.load(x + o_x, mask=o_x < N) + tl.store(x + o_x, b_x * b_g, mask=o_x < N) + + +def fused_kl_div_forward( + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str = 'batchmean' +): + device = x.device + + # ideally, we would like to achieve the same memory consumption as [N, H], + # so the expected chunk size should be: + # NC = ceil(V / H) + # C = ceil(N / NC) + # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048 + N, H, V = *x.shape, weight.shape[0] + BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + # TODO: in real cases, we may need to limit the number of chunks NC to + # ensure the precisions of accumulated gradients + NC = min(8, triton.cdiv(V, H)) + C = triton.next_power_of_2(triton.cdiv(N, NC)) + NC = triton.cdiv(N, C) + + dx = torch.zeros_like(x, device=device) + dw = torch.zeros_like(weight, device=device) if weight is not None else None + # we use fp32 for loss accumulator + loss = torch.zeros(N, dtype=torch.float32, device=device) + + for ic in range(NC): + start, end = ic * C, min((ic + 1) * C, N) + # [C, N] + c_sx = x[start:end] + c_tx = target_x[start:end] + # when doing matmul, use the original precision + # [C, V] + c_sl = F.linear(c_sx, weight) + c_tl = F.linear(c_tx, target_weight) + + # unreduced loss + c_loss = loss[start:end] + + # Here we calculate the gradient of c_sx in place so we can save memory. + kl_div_kernel[(c_sx.shape[0],)]( + logits=c_sl, + target_logits=c_tl, + loss=c_loss, + s_logits=c_sl.stride(-2), + s_loss=c_loss.stride(-1), + reduction=reduction, + N=N, + V=V, + BV=BV, + num_warps=32 + ) + + # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V + # thus dx[start: end] should be of shape: C x H + # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only + # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens. + # Thus, we need an additional scaling factor of (n_non_ignore/total) to scale the gradients. + # [C, H] + + dx[start:end] = torch.mm(c_sl, weight) + + if weight is not None: + torch.addmm(input=dw, mat1=c_sl.t(), mat2=c_sx, out=dw) + + loss = loss.sum() + return loss, dx, dw + + +def fused_kl_div_backward( + do: torch.Tensor, + dx: torch.Tensor, + dw: torch.Tensor +): + # If cross entropy is the last layer, do is 1.0. Skip the mul to save time + if torch.ne(do, torch.tensor(1.0, device=do.device)): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + N, H = dx.shape + B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) + + elementwise_mul_kernel[(triton.cdiv(N * H, B),)]( + x=dx, + g=do, + N=N*H, + B=B, + num_warps=32, + ) + + # handle dw + if dw is not None: + V, H = dw.shape + elementwise_mul_kernel[(triton.cdiv(V * H, B),)]( + x=dw, + g=do, + N=V*H, + B=B, + num_warps=32, + ) + + return dx, dw + + +class FusedKLDivLossFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str + ): + loss, dx, dw = fused_kl_div_forward( + x=x, + target_x=target_x, + weight=weight, + target_weight=target_weight, + reduction=reduction + ) + ctx.save_for_backward(dx, dw) + return loss + + @staticmethod + @input_guard + def backward(ctx, do): + dx, dw = ctx.saved_tensors + dx, dw = fused_kl_div_backward(do, dx, dw) + return dx, None, dw, None, None + + +def fused_kl_div_loss( + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor, + reduction: str = 'batchmean' +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target_x (torch.Tensor): [batch_size * seq_len, hidden_size] + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + target_weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + reduction: + Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'. + Returns: + loss + """ + return FusedKLDivLossFunction.apply( + x, + target_x, + weight, + target_weight, + reduction + ) + + +class FusedKLDivLoss(nn.Module): + + def __init__( + self, + reduction: str = 'batchmean' + ): + """ + Args: + reduction: + Specifies the reduction to apply to the output: 'batchmean'. Default: 'batchmean'. + """ + super().__init__() + + assert reduction in ['batchmean'], f"reduction: {reduction} is not supported" + + self.reduction = reduction + + def forward( + self, + x: torch.Tensor, + target_x: torch.Tensor, + weight: torch.Tensor, + target_weight: torch.Tensor + ): + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target_x (torch.Tensor): [batch_size * seq_len, hidden_size] + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + target_weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + Returns: + loss + """ + loss = fused_kl_div_loss( + x=x, + target_x=target_x, + weight=weight, + target_weight=target_weight, + reduction=self.reduction + ) + return loss diff --git a/fla/modules/fused_linear_cross_entropy.py b/fla/modules/fused_linear_cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..a18566fdbffccdd0d2b3bbe3586fdb38f18720e4 --- /dev/null +++ b/fla/modules/fused_linear_cross_entropy.py @@ -0,0 +1,570 @@ +# -*- coding: utf-8 -*- + +# Code adapted from +# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py + +from functools import partial +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.ops.utils import logsumexp_fwd +from fla.ops.utils.op import exp +from fla.utils import input_guard + +# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 +# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 +# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling +# The optimal maximum block size depends on your hardware, your kernel, and your dtype +MAX_FUSED_SIZE = 65536 // 2 + + +@triton.jit +def cross_entropy_kernel( + logits, + lse, + target, + loss, + total, + ignore_index, + label_smoothing: tl.constexpr, + logit_scale: tl.constexpr, + reduction: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr +): + """ + This kernel computes both cross entropy loss and the gradient of the input. + We only consider hard label + mean reduction for now. + Please refer to https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html for the math. + + Args: + logits: + Pointer to logits tensor. + lse: + Pointer to logsumexp tensor. + target: Pointer to target tensor. + loss: + Pointer to tensor to store the loss. + V (int): + The number of columns in the input tensor. + total (int): + The number of non-ignored classes. + ignore_index (int): + The index to ignore in the target. + label_smoothing (float): + The amount of smoothing when computing the loss, where 0.0 means no smoothing. + reduction (str): + The string for the reduction to apply + BV (int): + The block size for vocab. + """ + + # https://github.com/triton-lang/triton/issues/1058 + # If B*T*V is too large, i_n * stride will overflow out of int32, so we convert to int64 + i_n = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + + # 1. Load target first because if the target is ignore_index, we can return right away + b_y = tl.load(target + i_n) + + # 2. locate the start index + logits += i_n * V + + if b_y == ignore_index: + # set all x as 0 + for i in range(0, V, BV): + o_v = i + tl.arange(0, BV) + tl.store(logits + o_v, 0.0, mask=o_v < V) + return + + # Online softmax: 2 loads + 1 store (compared with 3 loads + 1 store for the safe softmax) + # Refer to Algorithm 3 in the paper: https://arxiv.org/pdf/1805.02867 + + # 3. [Online softmax] first pass: compute logsumexp + # we did this in anouter kernel + b_l = tl.load(logits + b_y) * logit_scale + b_lse = tl.load(lse + i_n) + + # 4. Calculate the loss + # loss = lse - logits_l + b_loss = b_lse - b_l + + # Label smoothing is a general case of normal cross entropy + # See the full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issue-2503665310 + b_z = 0.0 + eps = label_smoothing / V + + # We need tl.debug_barrier() as mentioned in + # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/ops/cross_entropy.py#L34 + tl.debug_barrier() + + # 5. [Online Softmax] Second pass: compute gradients + # For 'mean' reduction, gradients are normalized by number of non-ignored elements + # dx_y = (softmax(x_y) - 1) / N + # dx_i = softmax(x_i) / N, i != y + # For label smoothing: + # dx_i = (softmax(x_y) - label_smoothing / V) / N, i != y + # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N + # = dx_i - (1 - label_smoothing) / N + for iv in range(0, NV): + o_v = iv * BV + tl.arange(0, BV) + b_logits = tl.load(logits + o_v, mask=o_v < V, other=float('-inf')) * logit_scale + if label_smoothing > 0: + # scale X beforehand to avoid overflow + b_z += tl.sum(tl.where(o_v < V, -eps * b_logits, 0.0)) + b_p = (exp(b_logits - b_lse) - eps) * logit_scale + if reduction == "mean": + b_p = b_p / total + tl.store(logits + o_v, b_p, mask=o_v < V) + + tl.debug_barrier() + + # Orginal loss = H(q, p), with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps + # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p) + # = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i)) + # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as: + # = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd)) + # Refer to H(q', p) in section 7 of the paper: + # https://arxiv.org/pdf/1512.00567 + # pytorch: + # https://github.com/pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L516 + # See full derivation at https://github.com/linkedin/Liger-Kernel/pull/198#issuecomment-2333753087 + if label_smoothing > 0: + b_loss = b_loss * (1 - label_smoothing) + (b_z + label_smoothing * b_lse) + + # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N` + b_l = tl.load(logits + b_y) + + # Normalize the loss by the number of non-ignored elements if reduction is "mean" + if reduction == 'mean': + b_loss = b_loss / total + b_l += (label_smoothing - 1) / total * logit_scale + else: + b_l += (label_smoothing - 1) * logit_scale + + tl.store(loss + i_n, b_loss) + tl.store(logits + b_y, b_l) + + +@triton.jit +def elementwise_mul_kernel( + x, + g, + N: tl.constexpr, + B: tl.constexpr +): + """ + This function multiplies each element of the tensor pointed by x with the value pointed by g. + The multiplication is performed in-place on the tensor pointed by x. + + Parameters: + x: + Pointer to the input tensor. + g: + Pointer to the gradient output value. + N (int): + The number of columns in the input tensor. + B (int): + The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + i_x = tl.program_id(0).to(tl.int64) + o_x = i_x * B + tl.arange(0, B) + + # Load the gradient output value + b_g = tl.load(g) + b_x = tl.load(x + o_x, mask=o_x < N) + tl.store(x + o_x, b_x * b_g, mask=o_x < N) + + +def fused_linear_cross_entropy_forward( + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" +): + device = x.device + # inputs have shape: [N, H] + # materialized activations will have shape: [N, V] + # the increase in memory = [N, V] + # reduction can be achieved by partitioning the number of tokens N into smaller chunks. + + # ideally, we would like to achieve the same memory consumption as [N, H], + # so the expected chunk size should be: + # NC = ceil(V / H) + # C = ceil(N / NC) + # for ex: N = 4096*4, V = 32000, H = 4096 ==> NC = 8, C = ceil(N / NC) = 2048 + N, H, V = *x.shape, weight.shape[0] + BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + # TODO: in real cases, we may need to limit the number of chunks NC to + # ensure the precisions of accumulated gradients + NC = min(num_chunks, triton.cdiv(V, H)) + C = triton.next_power_of_2(triton.cdiv(N, NC)) + NC = triton.cdiv(N, C) + + # [N, H] + dx = torch.zeros_like(x, device=device) + # [V, H] + dw = torch.zeros_like(weight, device=device, dtype=torch.float) if weight is not None else None + # [V] + db = torch.zeros_like(bias, device=device, dtype=torch.float) if bias is not None else None + # [N] + loss = torch.zeros(N, device=device, dtype=torch.float) + + total = target.ne(ignore_index).sum().item() + + for ic in range(NC): + start, end = ic * C, min((ic + 1) * C, N) + # [C, N] + c_x = x[start:end] + # when doing matmul, use the original precision + # [C, V] + c_logits = F.linear(c_x, weight, bias) + c_target = target[start:end] + # [C] + # keep lse in fp32 to maintain precision + c_lse = logsumexp_fwd(c_logits, scale=logit_scale, dtype=torch.float) + + # unreduced loss + c_loss = loss[start:end] + + # Here we calculate the gradient of c_logits in place so we can save memory. + cross_entropy_kernel[(c_logits.shape[0],)]( + logits=c_logits, + lse=c_lse, + target=c_target, + loss=c_loss, + total=total, + ignore_index=ignore_index, + label_smoothing=label_smoothing, + logit_scale=logit_scale, + reduction=reduction, + V=V, + BV=BV, + num_warps=32 + ) + + # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V + # thus dx should be of shape: C x H + dx[start:end] = torch.mm(c_logits, weight) + + # keep dw in fp32 to maintain precision + if weight is not None: + dw += c_logits.t() @ c_x + + if bias is not None: + torch.add(input=db, other=c_logits.sum(0), out=db) + + loss = loss.sum() + if dw is not None: + dw = dw.to(weight) + if db is not None: + db = db.to(bias) + return loss, dx, dw, db + + +def fused_linear_cross_entropy_backward( + do: torch.Tensor, + dx: torch.Tensor, + dw: torch.Tensor, + db: torch.Tensor +): + # If cross entropy is the last layer, do is 1.0. Skip the mul to save time + if torch.ne(do, torch.tensor(1.0, device=do.device)): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + N, H = dx.shape + B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) + + elementwise_mul_kernel[(triton.cdiv(N * H, B),)]( + x=dx, + g=do, + N=N*H, + B=B, + num_warps=32, + ) + + # handle dw + if dw is not None: + V, H = dw.shape + elementwise_mul_kernel[(triton.cdiv(V * H, B),)]( + x=dw, + g=do, + N=V*H, + B=B, + num_warps=32, + ) + + if db is not None: + V = db.shape[0] + elementwise_mul_kernel[(triton.cdiv(V, B),)]( + x=db, + g=do, + N=V, + B=B, + num_warps=32, + ) + return dx, dw, db + + +class FusedLinearCrossEntropyFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" + ): + """ + Fusing the last linear layer with cross-entropy loss + Reference: https://github.com/mgmalek/efficient_cross_entropy + + Handle the forward and backward pass of the final linear layer via cross-entropy loss by avoiding + the materialization of the large logits tensor. Since Cross Entropy Loss is the last layer, we can + compute the gradient at the forward pass. By doing so, we don't have to store the x and target + for the backward pass. + + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target (torch.LongTensor): [batch_size * seq_len] + where each value is in [0, vocab_size). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + ignore_index: + the index to ignore in the target. + label_smoothing: + the amount of smoothing when computing the loss, where 0.0 means no smoothing. + logit_scale: float = 1.0, + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + """ + loss, dx, dw, db = fused_linear_cross_entropy_forward( + x, + target, + weight, + bias, + ignore_index, + label_smoothing, + logit_scale, + num_chunks, + reduction + ) + # downcast to dtype and store for backward + ctx.save_for_backward( + dx.detach(), + dw.detach() if weight is not None else None, + db.detach() if bias is not None else None, + ) + return loss + + @staticmethod + @input_guard + def backward(ctx, do): + dx, dw, db = ctx.saved_tensors + dx, dw, db = fused_linear_cross_entropy_backward(do, dx, dw, db) + return dx, None, dw, db, None, None, None, None, None + + +def fused_linear_cross_entropy_loss( + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target (torch.LongTensor): [batch_size * seq_len] + where each value is in [0, vocab_size). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + Returns: + losses: [batch,], float + """ + return FusedLinearCrossEntropyFunction.apply( + x, + target, + weight, + bias, + ignore_index, + label_smoothing, + logit_scale, + num_chunks, + reduction + ) + + +class FusedLinearCrossEntropyLoss(nn.Module): + + def __init__( + self, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" + ): + """ + Args: + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + """ + super().__init__() + + assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" + + self.ignore_index = ignore_index + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.num_chunks = num_chunks + self.reduction = reduction + + @torch.compiler.disable + def forward( + self, + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None + ): + """ + Args: + x (torch.Tensor): [batch_size, seq_len, hidden_size] + target (torch.LongTensor): [batch_size, seq_len] + where each value is in [0, V). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + Returns: + loss + """ + loss = fused_linear_cross_entropy_loss( + x.view(-1, x.shape[-1]), + target.view(-1), + weight=weight, + bias=bias, + ignore_index=self.ignore_index, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + num_chunks=self.num_chunks, + reduction=self.reduction + ) + return loss + + +class LinearLossParallel(ParallelStyle): + def __init__( + self, + *, + sequence_dim: int = 1, + use_local_output: bool = False, + ): + super().__init__() + + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + x, target, weight, bias = inputs + + if not isinstance(x, DTensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + x = DTensor.from_local(x, device_mesh, sequence_sharding) + if x.placements != sequence_sharding: + x = x.redistribute(placements=sequence_sharding, async_op=True) + if not isinstance(target, DTensor): + target = DTensor.from_local(target, device_mesh, [Replicate()]) + if target.placements != sequence_sharding: + target = target.redistribute(placements=sequence_sharding, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, [Replicate()]) + if weight.placements != [Replicate()]: + # we replicate the weight/bias in FLCE + weight = weight.redistribute(placements=[Replicate()], async_op=True) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, [Replicate()]) + if bias is not None and bias.placements != [Replicate()]: + bias = bias.redistribute(placements=[Replicate()], async_op=True) + + return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias + + @staticmethod + def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.sequence_sharding), + output_fn=partial(self._prepare_output_fn, self.use_local_output) + ) diff --git a/fla/modules/fused_linear_listnet_loss.py b/fla/modules/fused_linear_listnet_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1dcc820f06ab925e6f596904dd0b147209eca6 --- /dev/null +++ b/fla/modules/fused_linear_listnet_loss.py @@ -0,0 +1,427 @@ +# -*- coding: utf-8 -*- + +# Code adapted from +# https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py + +from functools import partial +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.ops.utils import logsumexp_fwd +from fla.ops.utils.op import exp +from fla.utils import input_guard + +# The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 +# https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 +# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling +# The optimal maximum block size depends on your hardware, your kernel, and your dtype +MAX_FUSED_SIZE = 65536 // 2 + +@triton.jit +def listnet_kernel( + logits, + targets, # Now full target distributions + lse_logits, + lse_targets, + loss, + total, + ignore_index, + logit_scale: tl.constexpr, + reduction: tl.constexpr, + V: tl.constexpr, + BV: tl.constexpr +): + i_n = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + + # Pointers to current token's data + logits_ptr = logits + i_n * V + targets_ptr = targets + i_n * V + loss_ptr = loss + i_n + + # Compute prediction softmax + b_lse_logits = tl.load(lse_logits + i_n) + b_lse_targets = tl.load(lse_targets + i_n) + b_loss = 0.0 + + # Compute gradient: softmax(pred) - softmax(target) + for iv in range(0, NV): + o_v = iv * BV + tl.arange(0, BV) + mask = o_v < V + + # Load target and compute softmax + t_val = tl.load(targets_ptr + o_v, mask=mask, other=0.0) + p_target = tl.exp(t_val - b_lse_targets) + + # Load logits and compute softmax + l_val = tl.load(logits_ptr + o_v, mask=mask, other=0.0) * logit_scale + l_val_minus_lse = l_val - b_lse_logits + p_pred = tl.exp(l_val_minus_lse) + + # Gradient calculation + grad_val = p_pred - p_target + if reduction == "mean": + grad_val = grad_val / total + grad_val = tl.where(b_lse_targets == float('inf'), 0.0, grad_val) + tl.store(logits_ptr + o_v, grad_val, mask=mask) + + # Cross-entropy loss + # instead of: b_loss -= tl.sum(p_target * tl.log(p_pred), axis=0) + b_loss -= tl.sum(p_target * l_val_minus_lse, axis=0) + + tl.store(loss_ptr, b_loss) + +@triton.jit +def elementwise_mul_kernel( + x, + g, + N: tl.constexpr, + B: tl.constexpr +): + """ + This function multiplies each element of the tensor pointed by x with the value pointed by g. + The multiplication is performed in-place on the tensor pointed by x. + + Parameters: + x: + Pointer to the input tensor. + g: + Pointer to the gradient output value. + N (int): + The number of columns in the input tensor. + B (int): + The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + i_x = tl.program_id(0).to(tl.int64) + o_x = i_x * B + tl.arange(0, B) + + # Load the gradient output value + b_g = tl.load(g) + b_x = tl.load(x + o_x, mask=o_x < N) + tl.store(x + o_x, b_x * b_g, mask=o_x < N) + + +def fused_linear_listnet_forward( + x: torch.Tensor, + targets: torch.Tensor, # Float tensor [N, V] + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" +): + N, H, V = *x.shape, weight.shape[0] + BV = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + NC = min(num_chunks, triton.cdiv(V, H)) + C = triton.next_power_of_2(triton.cdiv(N, NC)) + NC = triton.cdiv(N, C) + + # Initialize outputs + dx = torch.zeros_like(x) + dw = torch.zeros_like(weight, dtype=torch.float) if weight is not None else None + db = torch.zeros_like(bias, dtype=torch.float) if bias is not None else None + loss = torch.zeros(N, device=x.device, dtype=torch.float) + total = N # All tokens considered + + for ic in range(NC): + start, end = ic * C, min((ic + 1) * C, N) + c_x = x[start:end] + c_logits = F.linear(c_x, weight, bias) + c_targets = targets[start:end] + c_lse_logits = logsumexp_fwd(c_logits, scale=logit_scale, dtype=torch.float) + c_lse_targets = logsumexp_fwd(c_targets, dtype=torch.float).nan_to_num(nan=float("inf")) + c_loss = loss[start:end] + + # Call ListNet kernel + listnet_kernel[(c_logits.shape[0],)]( + logits=c_logits, + targets=c_targets, # Full target distributions + lse_logits=c_lse_logits, + lse_targets=c_lse_targets, + loss=c_loss, + total=total, + ignore_index=ignore_index, + logit_scale=logit_scale, + reduction=reduction, + V=V, + BV=BV, + num_warps=32 + ) + + # Backward through linear layer + dx[start:end] = torch.mm(c_logits, weight) + if weight is not None: + dw += c_logits.t() @ c_x + if bias is not None: + db += c_logits.sum(0) + + loss = loss.sum() + if reduction == "mean": + loss = loss / total + + return loss, dx, dw, db + + +def fused_linear_listnet_backward( + do: torch.Tensor, + dx: torch.Tensor, + dw: torch.Tensor, + db: torch.Tensor +): + # If cross entropy is the last layer, do is 1.0. Skip the mul to save time + if torch.ne(do, torch.tensor(1.0, device=do.device)): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + N, H = dx.shape + B = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) + + elementwise_mul_kernel[(triton.cdiv(N * H, B),)]( + x=dx, + g=do, + N=N*H, + B=B, + num_warps=32, + ) + + # handle dw + if dw is not None: + V, H = dw.shape + elementwise_mul_kernel[(triton.cdiv(V * H, B),)]( + x=dw, + g=do, + N=V*H, + B=B, + num_warps=32, + ) + + if db is not None: + V = db.shape[0] + elementwise_mul_kernel[(triton.cdiv(V, B),)]( + x=db, + g=do, + N=V, + B=B, + num_warps=32, + ) + return dx, dw, db + + +class FusedLinearListNetFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + targets: torch.Tensor, # Float targets + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" + ): + loss, dx, dw, db = fused_linear_listnet_forward( + x, targets, weight, bias, ignore_index, + logit_scale, num_chunks, reduction + ) + ctx.save_for_backward(dx, dw, db) + return loss + + @staticmethod + def backward(ctx, do): + dx, dw, db = ctx.saved_tensors + dx, dw, db = fused_linear_listnet_backward(do, dx, dw, db) + return dx, None, dw, db, None, None, None, None + + +def fused_linear_listnet_loss( + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + x (torch.Tensor): [batch_size * seq_len, hidden_size] + target (torch.LongTensor): [batch_size * seq_len] + where each value is in [0, vocab_size). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + Returns: + losses: [batch,], float + """ + return FusedLinearListNetFunction.apply( + x, + target, + weight, + bias, + ignore_index, + logit_scale, + num_chunks, + reduction + ) + + +class FusedLinearListNetLoss(nn.Module): + + def __init__( + self, + ignore_index: int = -100, + label_smoothing: float = 0.0, + logit_scale: float = 1.0, + num_chunks: int = 8, + reduction: str = "mean" + ): + """ + Args: + ignore_index: int. + If target == ignore_index, the loss is set to 0.0. + label_smoothing: float + logit_scale: float + A scaling factor applied to the logits. Default: 1.0 + num_chunks: int + The number of chunks to split the input tensor into for processing. + This can help optimize memory usage and computation speed. + Default: 8 + reduction: + Specifies the reduction to apply to the output: 'mean' | 'sum'. + 'mean': the weighted mean of the output is taken, + 'sum': the output will be summed. + Default: 'mean'. + """ + super().__init__() + + assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" + + self.ignore_index = ignore_index + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.num_chunks = num_chunks + self.reduction = reduction + + @torch.compiler.disable + def forward( + self, + x: torch.Tensor, + target: torch.LongTensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None + ): + """ + Args: + x (torch.Tensor): [batch_size, seq_len, hidden_size] + target (torch.LongTensor): [batch_size, seq_len] + where each value is in [0, V). + weight (torch.Tensor): [vocab_size, hidden_size] + where `vocab_size` is the number of classes. + bias (Optional[torch.Tensor]): [vocab_size] + where `vocab_size` is the number of classes. + Returns: + loss + """ + loss = fused_linear_listnet_loss( + x.view(-1, x.shape[-1]), + target.view(-1, target.shape[-1]), + weight=weight, + bias=bias, + ignore_index=self.ignore_index, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + num_chunks=self.num_chunks, + reduction=self.reduction + ) + return loss + + +class LinearLossParallel(ParallelStyle): + def __init__( + self, + *, + sequence_dim: int = 1, + use_local_output: bool = False, + ): + super().__init__() + + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + x, target, weight, bias = inputs + + if not isinstance(x, DTensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + x = DTensor.from_local(x, device_mesh, sequence_sharding) + if x.placements != sequence_sharding: + x = x.redistribute(placements=sequence_sharding, async_op=True) + if not isinstance(target, DTensor): + target = DTensor.from_local(target, device_mesh, [Replicate()]) + if target.placements != sequence_sharding: + target = target.redistribute(placements=sequence_sharding, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, [Replicate()]) + if weight.placements != [Replicate()]: + # we replicate the weight/bias in FLCE + weight = weight.redistribute(placements=[Replicate()], async_op=True) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, [Replicate()]) + if bias is not None and bias.placements != [Replicate()]: + bias = bias.redistribute(placements=[Replicate()], async_op=True) + + return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias + + @staticmethod + def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.sequence_sharding), + output_fn=partial(self._prepare_output_fn, self.use_local_output) + ) + +# Naive ListNet loss function implementation +def list_net_loss(y_pred, y_true): + """ + ListNet loss introduced in "Learning to Rank: From Pairwise Approach to Listwise Approach". + :param y_pred: predictions from the model, shape [*, slate_length] + :param y_true: ground truth labels, shape [*, slate_length] + :return: loss value, a torch.Tensor + """ + return torch.mean(-torch.sum(F.softmax(y_true, dim=-1).nan_to_num(nan=0) * F.log_softmax(y_pred, dim=-1), dim=-1)) \ No newline at end of file diff --git a/fla/modules/fused_norm_gate.py b/fla/modules/fused_norm_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..13807390f8138c998ba6b5d0eca226e74f7efde3 --- /dev/null +++ b/fla/modules/fused_norm_gate.py @@ -0,0 +1,995 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.utils import get_multiprocessor_count, input_guard + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8, 16, 32] + for num_stages in [2, 3, 4] + ], + key=['N', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM', 'HAS_BIAS'], +) +@triton.jit +def layer_norm_gated_fwd_kernel( + X, # pointer to the input + G, # pointer to the gate + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + RESIDUAL, # pointer to the residual + RESIDUAL_OUT, # pointer to the residual + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + N, # number of columns in X + eps, # epsilon to avoid division by zero + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + X += row * N + Y += row * N + G += row * N + if HAS_RESIDUAL: + RESIDUAL += row * N + if STORE_RESIDUAL_OUT: + RESIDUAL_OUT += row * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) + x += residual + if STORE_RESIDUAL_OUT: + tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w if HAS_WEIGHT else x_hat + if HAS_BIAS: + y = y + b + + # Swish output gate + g = tl.load(G + cols, mask=cols < N, other=0.0).to(tl.float32) + if ACTIVATION == 'swish': + y = y * g * tl.sigmoid(g) + elif ACTIVATION == 'silu': + y = y * g * tl.sigmoid(g) + elif ACTIVATION == 'sigmoid': + y = y * tl.sigmoid(g) + + # Write output + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_gated_fwd( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = 'swish', + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False +): + if residual is not None: + residual_dtype = residual.dtype + M, N = x.shape + if residual is not None: + assert residual.shape == (M, N) + if weight is not None: + assert weight.shape == (N,) + if bias is not None: + assert bias.shape == (N,) + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + residual_out = torch.empty(M, N, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = torch.empty((M,), dtype=torch.float, device=x.device) if not is_rms_norm else None + rstd = torch.empty((M,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + + layer_norm_gated_fwd_kernel[(M,)]( + x, + g, + y, + weight, + bias, + residual, + residual_out, + mean, + rstd, + N, + eps, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + BLOCK_N=BLOCK_N, + HAS_RESIDUAL=residual is not None, + STORE_RESIDUAL_OUT=residual_out is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +@triton.heuristics({ + 'RECOMPUTE_OUTPUT': lambda args: args["Y"] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8, 16, 32] + for num_stages in [2, 3, 4] + ], + key=['N', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM', 'HAS_BIAS'], +) +@triton.jit +def layer_norm_gated_bwd_kernel( + X, # pointer to the input + G, # pointer to the gate + W, # pointer to the weights + B, # pointer to the biases + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DG, # pointer to the gate gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DRESIDUAL, + DRESIDUAL_IN, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + rows_per_program, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + row_start = row_block_id * rows_per_program + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * N + G += row_start * N + if HAS_DRESIDUAL: + DRESIDUAL += row_start * N + if STORE_DRESIDUAL: + DRESIDUAL_IN += row_start * N + DY += row_start * N + DX += row_start * N + DG += row_start * N + if RECOMPUTE_OUTPUT: + Y += row_start * N + if HAS_WEIGHT: + w = tl.load(W + cols, mask=mask).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + g = tl.load(G + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.0) + + y = xhat * w if HAS_WEIGHT else xhat + if HAS_BIAS: + y = y + b + if RECOMPUTE_OUTPUT: + tl.store(Y + cols, y, mask=mask) + + sigmoid_g = tl.sigmoid(g) + if ACTIVATION == 'swish': + dg = dy * y * (sigmoid_g + g * sigmoid_g * (1 - sigmoid_g)) + dy = dy * g * sigmoid_g + elif ACTIVATION == 'silu': + dg = dy * y * (sigmoid_g + g * sigmoid_g * (1 - sigmoid_g)) + dy = dy * g * sigmoid_g + elif ACTIVATION == 'sigmoid': + dg = dy * y * sigmoid_g * (1 - sigmoid_g) + dy = dy * sigmoid_g + wdy = dy + if HAS_WEIGHT: + wdy = dy * w + dw += dy * xhat + if HAS_BIAS: + db += dy + if not IS_RMS_NORM: + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + c1 = tl.sum(xhat * wdy, axis=0) / N + dx = (wdy - xhat * c1) * rstd + if HAS_DRESIDUAL: + dres = tl.load(DRESIDUAL + cols, mask=mask, other=0).to(tl.float32) + dx += dres + # Write dx + if STORE_DRESIDUAL: + tl.store(DRESIDUAL_IN + cols, dx, mask=mask) + tl.store(DX + cols, dx, mask=mask) + tl.store(DG + cols, dg, mask=mask) + + X += N + G += N + if HAS_DRESIDUAL: + DRESIDUAL += N + if STORE_DRESIDUAL: + DRESIDUAL_IN += N + if RECOMPUTE_OUTPUT: + Y += N + DY += N + DX += N + DG += N + if HAS_WEIGHT: + tl.store(DW + row_block_id * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * N + cols, db, mask=mask) + + +def layer_norm_gated_bwd( + dy: torch.Tensor, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = 'swish', + eps: float = 1e-5, + mean: torch.Tensor = None, + rstd: torch.Tensor = None, + dresidual: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, +): + M, N = x.shape + assert dy.shape == (M, N) + if dresidual is not None: + assert dresidual.shape == (M, N) + if weight is not None: + assert weight.shape == (N,) + if bias is not None: + assert bias.shape == (N,) + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device) + dg = torch.empty_like(g) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device) + dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + sm_count = get_multiprocessor_count(x.device.index) + dw = torch.empty((sm_count, N), dtype=torch.float, device=weight.device) if weight is not None else None + db = torch.empty((sm_count, N), dtype=torch.float, device=bias.device) if bias is not None else None + rows_per_program = math.ceil(M / sm_count) + grid = (sm_count,) + layer_norm_gated_bwd_kernel[grid]( + x, + g, + weight, + bias, + y, + dy, + dx, + dg, + dw, + db, + dresidual, + dresidual_in, + mean, + rstd, + M, + N, + eps, + rows_per_program, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + BLOCK_N=BLOCK_N, + HAS_DRESIDUAL=dresidual is not None, + STORE_DRESIDUAL=dresidual_in is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + ) + dw = dw.sum(0).to(weight.dtype) if weight is not None else None + db = db.sum(0).to(bias.dtype) if bias is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype: + dresidual_in = dx + return (dx, dg, dw, db, dresidual_in) if not recompute_output else (dx, dg, dw, db, dresidual_in, y) + + +class LayerNormGatedFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + ): + x_shape_og = x.shape + g_shape_og = g.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + g = g.reshape(-1, g.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float if residual_in_fp32 else None) + ) + y, mean, rstd, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=weight, + bias=bias, + activation=activation, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm + ) + ctx.save_for_backward(residual_out, g, weight, bias, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.g_shape_og = g_shape_og + ctx.activation = activation + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dy, *args): + x, g, weight, bias, mean, rstd = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dg, dw, db, dresidual_in = layer_norm_gated_bwd( + dy=dy, + x=x, + g=g, + weight=weight, + bias=bias, + activation=ctx.activation, + eps=ctx.eps, + mean=mean, + rstd=rstd, + dresidual=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + ) + return ( + dx.reshape(ctx.x_shape_og), + dg.reshape(ctx.g_shape_og), + dw, + db, + None, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +class LayerNormGatedLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + ): + x_shape_og = x.shape + g_shape_og = g.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + g = g.reshape(-1, g.shape[-1]) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape(-1, residual.shape[-1]) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float if residual_in_fp32 else None) + ) + y, mean, rstd, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=norm_weight, + bias=norm_bias, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = linear_weight.to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, g, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.g_shape_og = g_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, g, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, dresidual.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dg, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_gated_bwd( + dy=dy, + x=x, + g=g, + norm_weight=norm_weight, + norm_bias=norm_bias, + eps=ctx.eps, + mean=mean, + rstd=rstd, + dresidual=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y) + return ( + dx.reshape(ctx.x_shape_og), + dg.reshape(ctx.g_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + ) + + +def layer_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = 'swish', + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6 +): + return LayerNormGatedFunction.apply( + x, + g, + weight, + bias, + activation, + residual, + eps, + prenorm, + residual_in_fp32, + False + ) + + +def rms_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = 'swish', + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6 +): + return LayerNormGatedFunction.apply( + x, + g, + weight, + bias, + activation, + residual, + eps, + prenorm, + residual_in_fp32, + True + ) + + +def layer_norm_swish_gate_linear( + x: torch.Tensor, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6 +): + return LayerNormGatedLinearFunction.apply( + x, + g, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + False + ) + + +def rms_norm_swish_gate_linear( + x, + g: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6 +): + return LayerNormGatedLinearFunction.apply( + x, + g, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + True + ) + + +class FusedLayerNormGated(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + activation: str = 'swish', + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedLayerNormGated: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ['swish', 'silu', 'sigmoid']: + raise ValueError(f"Unsupported activation: {self.activation}") + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += f", activation={self.activation}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False + ) -> torch.Tensor: + return layer_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32 + ) + + +class FusedRMSNormGated(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + activation: str = 'swish', + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedRMSNormGated: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ['swish', 'silu', 'sigmoid']: + raise ValueError(f"Unsupported activation: {self.activation}") + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += f", activation={self.activation}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False + ) -> torch.Tensor: + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32 + ) + + +class FusedLayerNormSwishGate(FusedLayerNormGated): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedLayerNormSwishGate: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + bias=bias, + eps=eps, + device=device, + dtype=dtype + ) + + +class FusedRMSNormSwishGate(FusedRMSNormGated): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedRMSNormSwishGate: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype + ) + + +class FusedLayerNormGatedLinear(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedLayerNormGatedLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + weight: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False + ) -> torch.Tensor: + return layer_norm_swish_gate_linear( + x, + g, + self.weight, + self.bias, + weight, + bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32 + ) + + +class FusedLayerNormSwishGateLinear(FusedLayerNormGatedLinear): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedLayerNormSwishGateLinear: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype + ) + + +class FusedRMSNormGatedLinear(nn.Module): + + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedRMSNormGatedLinear: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward( + self, + x: torch.Tensor, + g: torch.Tensor, + weight: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + residual: Optional[torch.Tensor] = None, + prenorm: bool = False, + residual_in_fp32: bool = False + ) -> torch.Tensor: + return rms_norm_swish_gate_linear( + x, + g, + self.weight, + self.bias, + weight, + bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32 + ) + + +class FusedRMSNormSwishGateLinear(FusedRMSNormGatedLinear): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> FusedRMSNormSwishGateLinear: + super().__init__( + hidden_size=hidden_size, + elementwise_affine=elementwise_affine, + eps=eps, + device=device, + dtype=dtype + ) diff --git a/fla/modules/grpo.py b/fla/modules/grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..dd1155557fd45321311258f75b1cbf8810f79c99 --- /dev/null +++ b/fla/modules/grpo.py @@ -0,0 +1,396 @@ +# -*- coding: utf-8 -*- + +# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py +""" +# Get the per-token log probabilities for the completions for the model and the reference model + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits + logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + + input_ids = input_ids[:, -logits_to_keep:] + # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. + # See https://github.com/huggingface/trl/issues/2770 + logits = logits[:, -logits_to_keep:] + return selective_log_softmax(logits, input_ids) # compute logprobs for the input tokens + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + # Compute the per-token log probabilities for the model + + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + + # x - x.detach() allows for preserving gradients from x + advantages = inputs["advantages"] + per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + per_token_loss = -(per_token_loss - self.beta * per_token_kl) + loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + + # Log the metrics + completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() + self._metrics["completion_length"].append(completion_length) + + mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + + return loss +""" + + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log +from fla.utils import input_guard + + +@triton.autotune( + [triton.Config({'BLOCK_SIZE': BLOCK_SIZE}, num_warps=NUM_WARPS, num_stages=NUM_STAGES) + for BLOCK_SIZE in [1024, 2048, 4096, 8192] + for NUM_WARPS in [8, 16, 32] + for NUM_STAGES in [1, 2, 4] + ], key=['B', 'N'] +) +@triton.jit +def grpo_fwd_kernel( + logits_ptr, + ref_logp_ptr, + input_ids_ptr, + advantages_ptr, + completion_mask_ptr, + loss_ptr, + lse_ptr, + beta, + save_kl: tl.constexpr, + B, + M, + N, + L, + start_idx, + BLOCK_SIZE: tl.constexpr +): + row_idx = tl.program_id(0) + + off_b = row_idx // L + N = tl.cast(N, tl.int64) + + loss_ptr += row_idx + + completion_mask_ptr += row_idx + not_skip = tl.load(completion_mask_ptr).to(tl.int1) + if not_skip == 1: + ref_logp_ptr += row_idx + lse_ptr += row_idx + advantages_ptr += off_b + logits_ptr += N * (row_idx + off_b) + input_ids_ptr += row_idx + (off_b+1) * start_idx + base_cols = tl.arange(0, BLOCK_SIZE) + + m_i = -float("inf") + l_i = 0.0 + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + logits = tl.load(logits_ptr+cols, mask=mask, other=-float('inf')).to(tl.float32) + m_ij = tl.max(logits) + new_m_i = tl.maximum(m_i, m_ij) + l_i = l_i * exp(m_i - new_m_i) + tl.sum(exp(logits - new_m_i)) + m_i = new_m_i + lse = log(l_i) + m_i + + idx = tl.load(input_ids_ptr) + x = tl.load(logits_ptr+idx).to(tl.float32) + advantage = tl.load(advantages_ptr).to(tl.float32) + ref_logp = tl.load(ref_logp_ptr) + logp = x - lse + diff = ref_logp - logp + kl = exp(diff) - diff - 1 + loss = kl * beta - advantage + + tl.store(loss_ptr, loss.to(loss_ptr.dtype.element_ty)) + tl.store(lse_ptr, lse.to(lse_ptr.dtype.element_ty)) + if save_kl: + tl.store(loss_ptr+M, kl.to(loss_ptr.dtype.element_ty)) + else: + # store 0 + tl.store(loss_ptr, 0.0) + if save_kl: + tl.store(loss_ptr+M, 0.0) + + +@triton.autotune( + [triton.Config({'BLOCK_SIZE': BLOCK_SIZE}, num_warps=NUM_WARPS, num_stages=NUM_STAGES) + for BLOCK_SIZE in [1024, 2048, 4096, 8192] + for NUM_WARPS in [8, 16, 32] + for NUM_STAGES in [1, 2, 4] + ], key=['B', 'N'] +) +@triton.jit +def grpo_bwd_kernel( + dloss_ptr, + dlogits_ptr, + logits_ptr, + ref_logp_ptr, + input_ids_ptr, + advantages_ptr, + completion_mask_ptr, + lse_ptr, + beta, + B, + N, + L, + start_idx, + BLOCK_SIZE: tl.constexpr +): + + row_idx = tl.program_id(0) # B*L + off_b = row_idx // L + + N = tl.cast(N, tl.int64) + + dlogits_ptr += N * (row_idx + off_b) + base_cols = tl.arange(0, BLOCK_SIZE) + completion_mask_ptr += row_idx + not_skip = tl.load(completion_mask_ptr).to(tl.int1) + + if not_skip == 1: + lse_ptr += row_idx + dloss_ptr += row_idx + advantages_ptr += off_b + ref_logp_ptr += row_idx + logits_ptr += N * (row_idx + off_b) + input_ids_ptr += row_idx + (off_b+1) * start_idx + dloss = tl.load(dloss_ptr).to(tl.float32) + lse = tl.load(lse_ptr).to(tl.float32) + idx = tl.load(input_ids_ptr) + x = tl.load(logits_ptr+idx).to(tl.float32) + advantage = tl.load(advantages_ptr).to(tl.float32) + ref_logp = tl.load(ref_logp_ptr) + logp = x - lse + + dlogp = (beta * (-1.0 * exp(ref_logp - logp) + 1) + - advantage) * dloss + + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + logits = tl.load(logits_ptr+cols, mask=mask, other=-float('inf')).to(tl.float32) + probs = exp(logits - lse) + dlogits = tl.where(cols == idx, 1-probs, -probs) * dlogp + + tl.store(dlogits_ptr+cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask) + else: + dlogits = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + for start_n in tl.range(0, N, BLOCK_SIZE): + cols = start_n + base_cols + mask = cols < N + + tl.store(dlogits_ptr+cols, dlogits.to(dlogits_ptr.dtype.element_ty), mask=mask) + + +class GrpoLoss(torch.autograd.Function): + + @input_guard + @staticmethod + def forward(ctx, logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl): + ctx.input_shape = logits.shape + B, L_ADD_1, N = ctx.input_shape + L = L_ADD_1 - 1 + M = B * L + input_ids_start_index = input_ids.size(1) - L + + if not save_kl: + loss = torch.empty(B, L, device=logits.device, dtype=torch.float32) + else: + loss = torch.empty(B*2, L, device=logits.device, dtype=torch.float32) + + lse = torch.empty(B, L, device=logits.device, dtype=torch.float32) + + if completion_mask is None: + completion_mask = torch.ones(B, L, device=logits.device, dtype=torch.int32) + else: + loss[:B].masked_fill_(completion_mask.logical_not(), 0.0) + + grpo_fwd_kernel[(M,)]( + logits_ptr=logits, + ref_logp_ptr=ref_logp, + input_ids_ptr=input_ids, + advantages_ptr=advantages, + completion_mask_ptr=completion_mask, + loss_ptr=loss, + lse_ptr=lse, + beta=beta, + save_kl=save_kl, + B=B, M=M, N=N, L=L, + start_idx=input_ids_start_index, + ) + ctx.beta = beta + ctx.save_for_backward(lse, logits, input_ids, advantages, completion_mask) + ctx.ref_logp = ref_logp + return loss + + @input_guard + @staticmethod + def backward(ctx, dloss): + # The grad of logits comes from two parts, the reward part and the kl part + lse, logits, input_ids, advantages, completion_mask = ctx.saved_tensors + B, L_ADD_1, N = ctx.input_shape + L = L_ADD_1 - 1 + M = B * L + + input_ids_start_index = input_ids.size(1) - L + + dlogits = torch.empty_like(logits) # B, L_ADD_1, N + + grpo_bwd_kernel[(M,)]( + dloss_ptr=dloss, + dlogits_ptr=dlogits, + logits_ptr=logits, + ref_logp_ptr=ctx.ref_logp, + input_ids_ptr=input_ids, + advantages_ptr=advantages, + completion_mask_ptr=completion_mask, + lse_ptr=lse, + beta=ctx.beta, + B=B, N=N, L=L, + start_idx=input_ids_start_index, + ) + # The last token in the completion is not used in the loss computation + # and therefore its gradient should be set to 0 + dlogits[:, -1, :].fill_(0.0) + return dlogits.view(*ctx.input_shape), None, None, None, None, None, None + + +def fused_grpo_loss(logits, ref_logp, input_ids, advantages, beta=0.1, completion_mask=None, save_kl=False) -> torch.Tensor: + ''' + compute grpo loss, save memory(no addition usage) and fast speed(6X for A800) + + Args: + logtits: Tensor, [B, L+1, vocab_size], the origin output of model, it's not logits[:, :-1] + ref_logp: Tensor, [B, L], the origin output of model, it's not ref_logits[:, :-1] + input_ids: Tensor, [B, K+L], it's prompt_completion_id, it contains the prompt ids and output ids + advantages: Tensor, [B], the advantages of each prompt + beta: float, the weight of kl loss + completion_mask: Tensor, loss mask + save_kl: bool, if true will save kl + + Retutn: + loss: Tensor, [B, L], the loss of grpo, it contains the advantage part and kl part + + NOTE: logits(ref_logits) is computed by these steps + logits_to_keep = completion_ids.size(1) + + def get_per_token_logits(model, input_ids, attention_mask, logits_to_keep): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model( + input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1 + ).logits + return logits + + logits = get_per_token_logits(model, prompt_completion_ids, attention_mask, logits_to_keep) + ''' + out = GrpoLoss.apply(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl) + if not save_kl: + return out + else: + return out.chunk(2, axis=0) + + +def grpo_loss_torch(logits, ref_logp, input_ids, advantages, beta=0.1, completion_mask=None, save_kl=False): + def get_log_probs(logits, input_ids): + per_token_logps = [] + for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):]): + log_probs = logits_row.log_softmax(dim=-1) + token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) + per_token_logps.append(token_log_prob) + return torch.stack(per_token_logps) + + logits = logits[:, :-1] + per_token_logps = get_log_probs(logits, input_ids) + ref_per_token_logps = ref_logp + per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + + per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + per_token_loss = -(per_token_loss - beta * per_token_kl) + if completion_mask is not None: + per_token_loss *= completion_mask + if save_kl: + per_token_kl *= completion_mask + return per_token_loss if not save_kl else (per_token_loss, per_token_kl) + + +@torch.compile(fullgraph=True) +def grpo_loss_with_old_logps( + logps: torch.Tensor, + ref_logps: torch.Tensor, + old_logps: torch.Tensor, + pad_mask: torch.Tensor, + logits_to_keep: int, + rewards: torch.Tensor, + beta: float = 0.2, + epsilon: float = 0.2 +): + """ + Compute the GRPO (Group Relative Policy Optimization) loss. + + Args: + logps (torch.Tensor): [Batch, Token_length] Log probabilities of the current policy. + ref_logps (torch.Tensor):[Batch, Token_length] Log probabilities of the reference policy. + old_logps (torch.Tensor): [Batch, Token_length] Log probabilities of the old policy. + completion_ids (torch.Tensor): [Batch, Token_length] Completion token IDs (bool). + pad_token_id: Pad token ID. + logits_to_keep (int): Number of logits to keep for masking. + rewards (torch.Tensor): [Batch] Rewards for each generation. + beta (float) = 0.2: A hyperparameter for weighting the KL divergence term. + epsilon (float) = 0.2: An float hyperparameter for clipping the importance weights. + + Returns: + torch.Tensor: The computed GRPO loss. + """ + B = logps.shape[0] + assert B > 1, "Batch * Num generations should be greater than 1" + + rewards_shaped = rewards.view(-1, B) # B,num_generations + advantages = (rewards_shaped - rewards_shaped.mean(dim=1, keepdim=True)) / \ + (rewards_shaped.std(dim=1, keepdim=True) + 1e-8) + advantages = advantages.view(-1) # B*num_generations + # Calculate the per - token KL divergence + per_token_kl = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1 + + # Calculate the ratio of probabilities (importance weights) + # Importance weights are calculated as exp(log_pi_theta - log_pi_theta_old) + importance_weights = torch.exp(logps - old_logps) + + # Clip the importance weights to the range [1 - epsilon, 1 + epsilon] + importance_weights_clipped = torch.clamp(importance_weights, 1 - epsilon, 1 + epsilon) + + # Create a completion mask. It checks which positions are valid based on logits_to_keep + completion_mask = torch.arange(logits_to_keep, device=logps.device)[None, :] >= 0 + + # Combine the completion mask and padding mask + completion_mask = completion_mask & pad_mask # Ensure matching shape + + # Add an extra dimension to advantages to match the shape for element - wise multiplication + advantages = advantages.unsqueeze(1) + + # Calculate the per - token loss. It takes the minimum of the unclipped and clipped importance weights + # and subtracts the KL divergence term weighted by beta, then multiplies by the completion mask + token_loss = -(torch.min(advantages * importance_weights, advantages * + importance_weights_clipped) - beta * per_token_kl) * completion_mask + + # Calculate the final loss by summing the token losses and normalizing by the number of valid tokens + loss = -token_loss.sum() / completion_mask.sum() + + return loss diff --git a/fla/modules/l2norm.py b/fla/modules/l2norm.py new file mode 100644 index 0000000000000000000000000000000000000000..bb400f1ba5b18f8c20091be76325fea0b042cc4f --- /dev/null +++ b/fla/modules/l2norm.py @@ -0,0 +1,176 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from fla.utils import input_guard + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=['N'] +) +@triton.jit +def l2norm_fwd_kernel( + X, + Y, + N, + eps, + BLOCK_N: tl.constexpr, +): + i_m = tl.program_id(0) + X += i_m * N + Y += i_m * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + mask = cols < N + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + xbar = tl.where(mask, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) + rstd = 1 / tl.sqrt(var + eps) + # tl.store(Rstd + i_m, rstd) + # Normalize and apply linear transformation + y = x * rstd + # Write output + tl.store(Y + cols, y, mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=['N'] +) +@triton.jit +def l2norm_bwd_kernel( + X, + DY, + DX, + N, + eps, + BLOCK_N: tl.constexpr, +): + i_m = tl.program_id(0) + X += i_m * N + DX += i_m * N + DY += i_m * N + + # Y += i_m * stride_y_row + cols = tl.arange(0, BLOCK_N) + mask = cols < N + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + x = tl.where(mask, x, 0.0) + var = tl.sum(x * x) + rstd = 1 / tl.sqrt(var + eps) + # tl.store(Rstd + i_m, rstd) + # Normalize and apply linear transformation + # y = x * rstd + dy = tl.load(DY + cols, mask=mask, other=0.0).to(tl.float32) + dy = tl.where(mask, dy, 0.0) + dx = dy * rstd - tl.sum(dy * x) * (1 / (var+eps)) * rstd * x + tl.store(DX + cols, dx, mask=mask) + + +def l2norm_fwd( + x: torch.Tensor, + eps: float = 1e-6, + output_dtype: Optional[torch.dtype] = None +): + x_shape_og = x.shape + x = x.reshape(-1, x.shape[-1]) + # allocate output + if output_dtype is None: + y = torch.empty_like(x) + else: + y = torch.empty_like(x, dtype=output_dtype) + assert y.stride(-1) == 1 + N = x.shape[-1] + M = x.shape[0] + # rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + l2norm_fwd_kernel[(M,)]( + x, + y, + N, + eps, + BLOCK_N, + ) + return y.reshape(x_shape_og) + + +def l2norm_bwd( + x: torch.Tensor, + dy: torch.Tensor, + eps: float = 1e-5 +): + x_shape_og = x.shape + x = x.reshape(-1, dy.shape[-1]) + dy = dy.reshape(-1, dy.shape[-1]) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + # allocate output + dx = torch.empty_like(x) + M = x.shape[0] + N = x.shape[-1] + # rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + l2norm_bwd_kernel[(M,)]( + x, + dy, + dx, + N, + eps, + BLOCK_N, + ) + return dx.reshape(x_shape_og) + + +class L2NormFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + eps=1e-6, + output_dtype=None + ): + y = l2norm_fwd(x, eps, output_dtype) + ctx.eps = eps + ctx.x_dtype = x.dtype + ctx.save_for_backward(x) + return y + + @staticmethod + @input_guard + def backward(ctx, dy): + x, = ctx.saved_tensors + dx = l2norm_bwd(x, dy, ctx.eps) + return dx, None, None + + +def l2_norm( + x: torch.Tensor, + eps: float = 1e-6, + output_dtype: Optional[torch.dtype] = None +) -> torch.Tensor: + return L2NormFunction.apply(x, eps, output_dtype) diff --git a/fla/modules/layernorm.py b/fla/modules/layernorm.py new file mode 100644 index 0000000000000000000000000000000000000000..5a1f7fff62502f24e5ed1ee888ceebeefba5e504 --- /dev/null +++ b/fla/modules/layernorm.py @@ -0,0 +1,1196 @@ +# -*- coding: utf-8 -*- + +# Copyright (c) 2023, Tri Dao. +# https://github.com/state-spaces/mamba/blob/fb7b5310fa865dbd62aa059b1e26f2b431363e2a/mamba_ssm/ops/triton/layernorm.py +# Implement residual + layer_norm / rms_norm. + +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +from __future__ import annotations + +from functools import partial + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.utils import get_multiprocessor_count, input_guard + + +def layer_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( + dtype + ) + return out if not prenorm else (out, x) + + +def rms_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = out.to(dtype) + return out if not prenorm else (out, x) + + +def group_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + residual: torch.Tensor = None, + eps: float = 1e-5, + is_rms_norm: bool = False, + prenorm: bool = False, + upcast: bool = False +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + residual = x + x, weight = [ + rearrange(data, "... (g d) -> ... g d", g=num_groups) for data in (x, weight) + ] + if bias is not None: + bias = rearrange(bias, '... (g d) -> ... g d', g=num_groups) + if not is_rms_norm: + mean = x.mean(dim=-1, keepdim=True) + x = x - mean + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = rearrange(out, "... g d -> ... (g d)") + out = out.to(dtype) + return out if not prenorm else (out, residual) + + +class GroupNormRef(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False + ) -> GroupNormRef: + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False): + return group_norm_ref( + x, + self.weight, + self.bias, + num_groups=self.num_groups, + residual=residual, + eps=self.eps, + is_rms_norm=self.is_rms_norm, + prenorm=prenorm, + upcast=True + ) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8, 16, 32] + for num_stages in [2, 3, 4] + ], + key=["N", "HAS_RESIDUAL", "STORE_RESIDUAL_OUT", "IS_RMS_NORM", "HAS_BIAS"], +) +@triton.jit +def layer_norm_fwd_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + RESIDUAL, # pointer to the residual + RESIDUAL_OUT, # pointer to the residual + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + N, # number of columns in X + G, # number of groups + eps, # epsilon to avoid division by zero + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + group = row % G + X += row * N + Y += row * N + if HAS_RESIDUAL: + RESIDUAL += row * N + if STORE_RESIDUAL_OUT: + RESIDUAL_OUT += row * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32) + x += residual + if STORE_RESIDUAL_OUT: + tl.store(RESIDUAL_OUT + cols, x, mask=cols < N) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.0) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + if HAS_WEIGHT: + w = tl.load(W + group * N + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + group * N + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + + y = tl.fma(x_hat, w, b) if HAS_WEIGHT and HAS_BIAS else \ + x_hat * w if HAS_WEIGHT else \ + x_hat + b if HAS_BIAS else x_hat + # Write output + y = tl.cast(y, dtype=Y.dtype.element_ty, fp_downcast_rounding="rtne") + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, + num_groups: int = 1 +): + if residual is not None: + residual_dtype = residual.dtype + M, N, G = *x.shape, num_groups + if residual is not None: + assert residual.shape == (M, N) + if weight is not None: + assert weight.shape == (G * N,) + if bias is not None: + assert bias.shape == (G * N,) + # allocate output + y = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype) + if residual is not None or (residual_dtype is not None and residual_dtype != x.dtype): + residual_out = torch.empty(M, N, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = torch.empty((M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + layer_norm_fwd_kernel[(M,)]( + x, + y, + weight, + bias, + residual, + residual_out, + mean, + rstd, + N, + G, + eps, + is_rms_norm, + BLOCK_N, + residual is not None, + residual_out is not None, + weight is not None, + bias is not None, + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +@triton.heuristics({ + "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8, 16, 32] + for num_stages in [2, 3, 4] + ], + key=["N", "HAS_DRESIDUAL", "STORE_DRESIDUAL", "IS_RMS_NORM", "HAS_BIAS"], +) +@triton.jit +def layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DRESIDUAL, + DRESIDUAL_IN, + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + M, # number of rows in X + N, # number of columns in X + G, # number of groups + rows_per_program, + programs_per_group, + IS_RMS_NORM: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_DRESIDUAL: tl.constexpr, + STORE_DRESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, +): + row_block_id = tl.program_id(0) + group_id, program_id_in_group = row_block_id // programs_per_group, row_block_id % programs_per_group + + row_start = group_id + program_id_in_group * G * rows_per_program + row_end = min(row_start + G * rows_per_program, M) + + cols = tl.arange(0, BLOCK_N) + mask = cols < N + + if HAS_WEIGHT: + w = tl.load(W + group_id * N + cols, mask=mask).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b = tl.load(B + group_id * N + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + + for row in range(row_start, row_end, G): + # Load data to SRAM + x = tl.load(X + row * N + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + row * N + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.0) + if RECOMPUTE_OUTPUT: + y = xhat * w if HAS_WEIGHT else xhat + if HAS_BIAS: + y = y + b + tl.store(Y + row * N + cols, y, mask=mask) + wdy = dy + if HAS_WEIGHT: + wdy = dy * w + dw += dy * xhat + if HAS_BIAS: + db += dy + if not IS_RMS_NORM: + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + c1 = tl.sum(xhat * wdy, axis=0) / N + dx = (wdy - xhat * c1) * rstd + if HAS_DRESIDUAL: + dres = tl.load(DRESIDUAL + row * N + cols, mask=mask, other=0).to(tl.float32) + dx += dres + # Write dx + dx = tl.cast(dx, dtype=DX.dtype.element_ty, fp_downcast_rounding="rtne") + if STORE_DRESIDUAL: + tl.store(DRESIDUAL_IN + row * N + cols, dx, mask=mask) + tl.store(DX + row * N + cols, dx, mask=mask) + + if HAS_WEIGHT: + tl.store(DW + row_block_id * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * N + cols, db, mask=mask) + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + dresidual: torch.Tensor = None, + has_residual: bool = False, + is_rms_norm: bool = False, + x_dtype: torch.dtype = None, + recompute_output: bool = False, + num_groups: int = 1 +): + M, N, G = *x.shape, num_groups + assert dy.shape == (M, N) + if dresidual is not None: + assert dresidual.shape == (M, N) + if weight is not None: + assert weight.shape == (G * N,) + if bias is not None: + assert bias.shape == (G * N,) + # allocate output + dx = torch.empty_like(x) if x_dtype is None else torch.empty(M, N, dtype=x_dtype, device=x.device) + dresidual_in = torch.empty_like(x) if has_residual and dx.dtype != x.dtype else None + y = torch.empty(M, N, dtype=dy.dtype, device=dy.device) if recompute_output else None + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + if N > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # each program handles one group only + S = triton.cdiv(get_multiprocessor_count(x.device.index), G) * G + dw = torch.empty((S, N), dtype=torch.float32, device=weight.device) if weight is not None else None + db = torch.empty((S, N), dtype=torch.float32, device=bias.device) if bias is not None else None + rows_per_program = triton.cdiv(M, S) + programs_per_group = S // G + grid = (S,) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + y, + dy, + dx, + dw, + db, + dresidual, + dresidual_in, + mean, + rstd, + M, + N, + G, + rows_per_program, + programs_per_group, + is_rms_norm, + BLOCK_N, + dresidual is not None, + dresidual_in is not None, + weight is not None, + bias is not None, + ) + dw = dw.view(G, -1, N).sum(1).to(weight).view_as(weight) if weight is not None else None + db = db.view(G, -1, N).sum(1).to(bias).view_as(bias) if bias is not None else None + # Don't need to compute dresidual_in separately in this case + if has_residual and dx.dtype == x.dtype: + dresidual_in = dx + return (dx, dw, db, dresidual_in) if not recompute_output else (dx, dw, db, dresidual_in, y) + + +class LayerNormFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + weight, + bias, + residual=None, + eps=1e-5, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + num_groups=1 + ): + x_shape_og = x.shape + + if x.shape[-1] % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + # reshape input data into 2D tensor + x = x.reshape(-1, (x.shape[-1] // num_groups)) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape_as(x) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + y, mean, rstd, residual_out = layer_norm_fwd( + x, + weight, + bias, + eps, + residual, + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + num_groups=num_groups + ) + ctx.save_for_backward(residual_out, weight, bias, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.num_groups = num_groups + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dy, *args): + x, weight, bias, mean, rstd = ctx.saved_tensors + dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, x.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dw, db, dresidual_in = layer_norm_bwd( + dy, + x, + weight, + bias, + ctx.eps, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + num_groups=ctx.num_groups + ) + return ( + dx.reshape(ctx.x_shape_og), + dw, + db, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + None + ) + + +def layer_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm + ) + + +def group_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1 +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + num_groups + ) + + +def rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False +): + return LayerNormFunction.apply( + x, + weight, + bias, + residual, + eps, + prenorm, + residual_in_fp32, + True + ) + + +def layer_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1 +): + return LayerNormLinearFunction.apply( + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual, + eps, + prenorm, + residual_in_fp32, + is_rms_norm, + num_groups + ) + + +def rms_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False +): + return layer_norm_linear( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True + ) + + +def group_norm_linear( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + linear_weight: torch.Tensor, + linear_bias: torch.Tensor, + residual: torch.Tensor = None, + eps: float = 1e-5, + prenorm: bool = False, + residual_in_fp32: bool = False, + is_rms_norm: bool = False, + num_groups: int = 1 +): + return layer_norm_linear( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + linear_weight=linear_weight, + linear_bias=linear_bias, + residual=residual, + eps=eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=is_rms_norm, + num_groups=num_groups + ) + + +class LayerNorm(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5 + ) -> LayerNorm: + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32 + ) + + +class GroupNorm(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False + ) -> GroupNorm: + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return group_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=self.is_rms_norm, + num_groups=self.num_groups + ) + + +class RMSNorm(nn.Module): + + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5 + ) -> RMSNorm: + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): + return rms_norm( + x, + self.weight, + self.bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) + + +class LayerNormLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-5, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + num_groups=1 + ): + x_shape_og = x.shape + + if x.shape[-1] % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + # reshape input data into 2D tensor + x = x.reshape(-1, (x.shape[-1] // num_groups)) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape_as(x) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float32 if residual_in_fp32 else None) + ) + y, mean, rstd, residual_out = layer_norm_fwd( + x, + norm_weight, + norm_bias, + eps, + residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype(), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + num_groups=num_groups + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else y.dtype + linear_weight = linear_weight.to(dtype) + linear_bias = linear_bias.to(dtype) if linear_bias is not None else None + out = F.linear(y.to(linear_weight.dtype), linear_weight, linear_bias) + # We don't store y, will be recomputed in the backward pass to save memory + ctx.save_for_backward(residual_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.is_rms_norm = is_rms_norm + ctx.num_groups = num_groups + ctx.has_residual = residual is not None + ctx.prenorm = prenorm + ctx.x_dtype = x.dtype + ctx.linear_bias_is_none = linear_bias is None + return out if not prenorm else (out, residual_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + def backward(ctx, dout, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(dout, linear_weight.t()) + dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dresidual = args[0] + dresidual = dresidual.reshape(-1, x.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd( + dy, + x, + norm_weight, + norm_bias, + ctx.eps, + mean, + rstd, + dresidual, + ctx.has_residual, + ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + num_groups=ctx.num_groups + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y.view(-1, linear_weight.shape[-1])) + return ( + dx.reshape(ctx.x_shape_og), + dnorm_weight, + dnorm_bias, + dlinear_weight, + dlinear_bias, + dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, + None, + None, + None, + None, + None + ) + + +class LayerNormLinear(nn.Module): + + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5 + ) -> LayerNormLinear: + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=False + ) + + +class GroupNormLinear(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False + ) -> GroupNormLinear: + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=self.is_rms_norm, + num_groups=self.num_groups + ) + + +class RMSNormLinear(nn.Module): + + def __init__( + self, + hidden_size, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5 + ) -> RMSNormLinear: + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + nn.init.ones_(self.weight) + if self.bias is not None: + nn.init.zeros_(self.bias) + + def __repr__(self) -> str: + s = f"{self.__class__.__name__}({self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): + return layer_norm_linear( + x=x, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + is_rms_norm=True + ) + + +class NormParallel(ParallelStyle): + + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False): + super().__init__() + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + def _replicate_module_fn( + self, name: str, module: nn.Module, device_mesh: DeviceMesh + ): + for p_name, param in module.named_parameters(): + # simple replication with fixed ones_ init from LayerNorm/RMSNorm, which allow + # us to simply just use from_local + replicated_param = torch.nn.Parameter( + DTensor.from_local(param, device_mesh, [Replicate()], run_check=False) + ) + module.register_parameter(p_name, replicated_param) + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + input_tensor = inputs[0] + if isinstance(input_tensor, DTensor): + # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it + if input_tensor.placements != sequence_sharding: + input_tensor = input_tensor.redistribute( + placements=sequence_sharding, async_op=True + ) + return input_tensor + elif isinstance(input_tensor, torch.Tensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + return DTensor.from_local( + input_tensor, device_mesh, sequence_sharding, run_check=False + ) + else: + raise ValueError( + f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}" + ) + + @staticmethod + def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + self._replicate_module_fn, + partial(self._prepare_input_fn, self.sequence_sharding), + partial(self._prepare_output_fn, self.use_local_output), + ) diff --git a/fla/modules/layernorm_gated.py b/fla/modules/layernorm_gated.py new file mode 100644 index 0000000000000000000000000000000000000000..1a72ff839dc021e484d55486cc11a9c9f85863fe --- /dev/null +++ b/fla/modules/layernorm_gated.py @@ -0,0 +1,528 @@ +# Copyright (c) 2024, Tri Dao. +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. +# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. + +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + +from fla.utils import get_multiprocessor_count, input_guard + + +def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True): + dtype = x.dtype + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + z = z.float() if z is not None else z + if z is not None and not norm_before_gate: + x = x * F.silu(z) + if group_size is None: + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + else: + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) + out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight + if bias is not None: + out = out + bias + if z is not None and norm_before_gate: + out *= F.silu(z) + return out.to(dtype) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["B"] is not None, + "HAS_Z": lambda args: args["Z"] is not None, +}) +@triton.jit +def layer_norm_fwd_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_y_row, + stride_z_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_N: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + group = tl.program_id(1) + X += row * stride_x_row + group * N + Y += row * stride_y_row + group * N + if HAS_Z: + Z += row * stride_z_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + if HAS_BIAS: + B += group * N + # Compute mean and variance + cols = tl.arange(0, BLOCK_N) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=cols < N).to(tl.float32) + x *= z * tl.sigmoid(z) + if not IS_RMS_NORM: + mean = tl.sum(x, axis=0) / N + tl.store(Mean + row, mean) + xbar = tl.where(cols < N, x - mean, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + else: + xbar = tl.where(cols < N, x, 0.) + var = tl.sum(xbar * xbar, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + mask = cols < N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if HAS_BIAS: + b = tl.load(B + cols, mask=mask).to(tl.float32) + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + y = x_hat * w + b if HAS_BIAS else x_hat * w + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask).to(tl.float32) + y *= z * tl.sigmoid(z) + # Write output + tl.store(Y + cols, y, mask=mask) + + +def layer_norm_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + z: torch.Tensor = None, + out: torch.Tensor = None, + group_size: int = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + assert x.stride(-1) == 1 + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + if out is not None: + assert out.shape == x.shape + else: + out = torch.empty_like(x) + assert out.stride(-1) == 1 + mean = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_N // 256, 1), 8) + grid = (M, ngroups) + layer_norm_fwd_kernel[grid]( + x, + out, + weight, + bias, + z, + mean, + rstd, + x.stride(0), + out.stride(0), + z.stride(0) if z is not None else 0, + M, + group_size, + eps, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps + ) + return out, mean, rstd + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["B"] is not None, + "HAS_Z": lambda args: args["Z"] is not None, + "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None, +}) +@triton.jit +def layer_norm_bwd_kernel( + X, # pointer to the input + W, # pointer to the weights + B, # pointer to the biases + Z, # pointer to the other branch + Y, # pointer to the output to be recomputed + DY, # pointer to the output gradient + DX, # pointer to the input gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + DZ, # pointer to the other branch + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride_x_row, # how much to increase the pointer when moving by 1 row + stride_z_row, + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dz_row, + stride_dw_row, + stride_db_row, + M, # number of rows in X + N, # number of columns in X + eps, # epsilon to avoid division by zero + rows_per_program, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_Z: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # Map the program id to the elements of X, DX, and DY it should compute. + row_block_id = tl.program_id(0) + group = tl.program_id(1) + row_start = row_block_id * rows_per_program + cols = tl.arange(0, BLOCK_N) + mask = cols < N + X += row_start * stride_x_row + group * N + if HAS_Z: + Z += row_start * stride_z_row + group * N + DZ += row_start * stride_dz_row + group * N + DY += row_start * stride_dy_row + group * N + DX += row_start * stride_dx_row + group * N + if RECOMPUTE_OUTPUT: + Y += row_start * stride_y_row + group * N + if not IS_RMS_NORM: + Mean += group * M + Rstd += group * M + W += group * N + w = tl.load(W + cols, mask=mask).to(tl.float32) + if (RECOMPUTE_OUTPUT or HAS_Z) and HAS_BIAS: + B += group * N + b = tl.load(B + cols, mask=mask, other=0.).to(tl.float32) + dw = tl.zeros((BLOCK_N,), dtype=tl.float32) + if HAS_BIAS: + db = tl.zeros((BLOCK_N,), dtype=tl.float32) + row_end = min((row_block_id + 1) * rows_per_program, M) + for row in range(row_start, row_end): + # Load data to SRAM + x = tl.load(X + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(DY + cols, mask=mask, other=0).to(tl.float32) + if not IS_RMS_NORM: + mean = tl.load(Mean + row) + if HAS_Z and not NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) + x_og = x + x = x_og * z * tl.sigmoid(z) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd + xhat = tl.where(mask, xhat, 0.) + if HAS_Z and NORM_BEFORE_GATE: + z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) + z_sigmoid = tl.sigmoid(z) + y = xhat * w + b if HAS_BIAS else xhat * w + if RECOMPUTE_OUTPUT: + tl.store(Y + cols, y * z * z_sigmoid, mask=mask) + dz = dy * y * z_sigmoid * (1 + z * (1 - z_sigmoid)) + tl.store(DZ + cols, dz, mask=mask) + dy *= z * z_sigmoid + else: + if RECOMPUTE_OUTPUT: + y = xhat * w + b if HAS_BIAS else xhat * w + tl.store(Y + cols, y, mask=mask) + wdy = w * dy + c1 = tl.sum(xhat * wdy, axis=0) / N + if not IS_RMS_NORM: + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + else: + dx = (wdy - xhat * c1) * rstd + dw += dy * xhat + if HAS_BIAS: + db += dy + if HAS_Z and not NORM_BEFORE_GATE: + z_sigmoid = tl.sigmoid(z) + dz = dx * x_og * z_sigmoid * (1 + z * (1 - z_sigmoid)) + tl.store(DZ + cols, dz, mask=mask) + dx *= z * z_sigmoid + # Write dx + tl.store(DX + cols, dx, mask=mask) + + X += stride_x_row + if HAS_Z: + Z += stride_z_row + DZ += stride_dz_row + if RECOMPUTE_OUTPUT: + Y += stride_y_row + DY += stride_dy_row + DX += stride_dx_row + tl.store(DW + row_block_id * stride_dw_row + group * N + cols, dw, mask=mask) + if HAS_BIAS: + tl.store(DB + row_block_id * stride_db_row + group * N + cols, db, mask=mask) + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + z: torch.Tensor = None, + group_size: int = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, + recompute_output: bool = False, + dz: torch.Tensor = None, + out: torch.Tensor = None, +): + M, N = x.shape + if group_size is None: + group_size = N + assert N % group_size == 0 + ngroups = N // group_size + assert x.stride(-1) == 1 + assert dy.stride(-1) == 1 + assert dy.shape == (M, N) + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (M, N) + assert weight.shape == (N,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (N,) + # allocate output + dx = torch.empty_like(x) + if dz is not None: + assert z is not None + assert dz.shape == z.shape + assert dz.stride(-1) == 1 + else: + dz = torch.empty_like(z) if z is not None else None + if recompute_output: + if out is None: + out = torch.empty_like(x) + assert out.shape == x.shape + + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BLOCK_N: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + # heuristics for number of warps + num_warps = min(max(BLOCK_N // 256, 1), 8) + sm_count = get_multiprocessor_count(x.device.index) + # If group size is small (e.g., 64), we're only using 1 warp. So having just 108 programs + # would limit the occupancy. + nrow_groups = math.ceil(sm_count * math.ceil(4 / num_warps) / ngroups) + _dw = torch.empty((nrow_groups, N), dtype=torch.float32, device=weight.device) + _db = torch.empty((nrow_groups, N), dtype=torch.float32, device=bias.device) if bias is not None else None + rows_per_program = math.ceil(M / nrow_groups) + grid = (nrow_groups, ngroups) + layer_norm_bwd_kernel[grid]( + x, + weight, + bias, + z, + out if recompute_output else None, + dy, + dx, + _dw, + _db, + dz, + mean, + rstd, + x.stride(0), + z.stride(0) if z is not None else 0, + 0 if not recompute_output else out.stride(0), + dy.stride(0), + dx.stride(0), + dz.stride(0) if dz is not None else 0, + _dw.stride(0), + _db.stride(0) if _db is not None else 0, + M, group_size, eps, + rows_per_program, + BLOCK_N=BLOCK_N, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps + ) + dw = _dw.sum(0).to(weight.dtype) + db = _db.sum(0).to(bias.dtype) if bias is not None else None + return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out) + + +class LayerNormFn(torch.autograd.Function): + + @input_guard + @staticmethod + def forward(ctx, x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, + is_rms_norm=False): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + assert z.shape == x_shape_og + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y, mean, rstd = layer_norm_fwd( + x, + weight, + bias, + eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + ) + ctx.save_for_backward(x, weight, bias, mean, rstd, z) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.group_size = group_size + ctx.norm_before_gate = norm_before_gate + ctx.is_rms_norm = is_rms_norm + return y.reshape(x_shape_og) + + @input_guard + @staticmethod + def backward(ctx, dy): + x, weight, bias, mean, rstd, z = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + dx, dw, db, dz = layer_norm_bwd( + dy, + x, + weight, + bias, + ctx.eps, + mean, + rstd, + z, + ctx.group_size, + ctx.norm_before_gate, + ctx.is_rms_norm + ) + dx = dx.reshape(ctx.x_shape_og) + dz = dz.reshape(ctx.x_shape_og) if dz is not None else None + return dx, dw, db, dz, None, None, None, None + + +def layernorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, is_rms_norm=False): + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm) + + +def rmsnorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True): + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, True) + + +class LayerNormGated(nn.Module): + + def __init__( + self, + hidden_size, + eps: float = 1e-5, + group_size: Optional[int] = None, + norm_before_gate: bool = True, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + torch.nn.init.zeros_(self.bias) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + return layernorm_fn(x, self.weight, self.bias, z=z, group_size=self.group_size, eps=self.eps, + norm_before_gate=self.norm_before_gate) + + +class RMSNormGated(nn.Module): + + def __init__( + self, + hidden_size, + eps: float = 1e-5, + group_size: Optional[int] = None, + norm_before_gate: bool = False, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x, z=None): + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) + """ + return rmsnorm_fn(x, self.weight, self.bias, z=z, eps=self.eps, group_size=self.group_size, + norm_before_gate=self.norm_before_gate) diff --git a/fla/modules/mlp.py b/fla/modules/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..b5f35aa6910ad143eb632fc28684043704e8741a --- /dev/null +++ b/fla/modules/mlp.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from __future__ import annotations + +from functools import partial +from typing import TYPE_CHECKING, Any, Optional + +import torch +import torch.nn as nn +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor, Placement, Replicate, Shard, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle + +from fla.modules.activations import swiglu, swiglu_linear + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +class GatedMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + hidden_ratio: Optional[int] = None, + intermediate_size: Optional[int] = None, + hidden_act: str = 'swish', + fuse_swiglu: bool = True + ) -> GatedMLP: + super().__init__() + + self.hidden_size = hidden_size + # the final number of params is `hidden_ratio * hidden_size^2` + # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio` + if hidden_ratio is None: + hidden_ratio = 4 + if intermediate_size is None: + intermediate_size = int(hidden_size * hidden_ratio * 2 / 3) + intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256) + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.fuse_swiglu = fuse_swiglu + + if hidden_act != 'swish': + raise ValueError(f'Unsupported hidden_act: {hidden_act}') + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + if self.fuse_swiglu: + self.swiglu_linear = SwiGLULinear() + + def forward( + self, + x: torch.Tensor, + **kwargs: Unpack[Any] + ) -> torch.Tensor: + gate, y = self.gate_proj(x), self.up_proj(x) + if self.fuse_swiglu: + return self.swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias) + else: + return self.down_proj(swiglu(gate, y)) + + +class SwiGLULinear(nn.Module): + + def forward(self, x, y, weight, bias): + return swiglu_linear(x, y, weight, bias) + + +class SwiGLULinearParallel(ParallelStyle): + def __init__( + self, + *, + input_layouts: Optional[Placement] = None, + output_layouts: Optional[Placement] = None, + use_local_output: bool = True, + ): + super().__init__() + self.input_layouts = (input_layouts or Shard(-1),) + self.output_layouts = (output_layouts or Replicate(),) + self.desired_input_layouts = (Shard(-1),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn( + input_layouts, desired_input_layouts, mod, inputs, device_mesh + ): + x, y, weight, bias = inputs + if not isinstance(x, DTensor): + x = DTensor.from_local(x, device_mesh, input_layouts, run_check=False) + if x.placements != desired_input_layouts: + x = x.redistribute(placements=desired_input_layouts, async_op=True) + + if not isinstance(y, DTensor): + y = DTensor.from_local(y, device_mesh, input_layouts, run_check=False) + if y.placements != desired_input_layouts: + y = y.redistribute(placements=desired_input_layouts, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, (Shard(1),)) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, (Replicate(),)) + + return x, y, weight, bias + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + # Rowwise sharding produces partial output, depending on output layouts: + # 1. to replicate -> allreduce + # 2. to shard -> reduce_scatter + if outputs.placements != output_layouts: + outputs = outputs.redistribute(placements=output_layouts, async_op=True) + # back to local tensor if use_local_output is True + return outputs.to_local() if use_local_output else outputs + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts), + output_fn=partial(self._prepare_output_fn, self.output_layouts, self.use_local_output) + ) diff --git a/fla/modules/parallel.py b/fla/modules/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..3561476adeed63be51edbee270106291d6a646ac --- /dev/null +++ b/fla/modules/parallel.py @@ -0,0 +1,37 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch.nn as nn +from torch.distributed import DeviceMesh +from torch.distributed.tensor import DTensor, distribute_module +from torch.distributed.tensor.parallel import ParallelStyle +from torch.distributed.tensor.placement_types import Placement + + +class PrepareModuleWeight(ParallelStyle): + def __init__(self, *, layouts: Optional[Placement] = None): + super().__init__() + self.layouts = layouts + + def _replicate_module_fn( + self, + name: str, + module: nn.Module, + device_mesh: DeviceMesh + ): + for p_name, param in module.named_parameters(): + replicated_param = nn.Parameter( + DTensor.from_local(param, device_mesh, [self.layouts], run_check=False) + ) + module.register_parameter(p_name, replicated_param) + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=self._replicate_module_fn, + input_fn=None, + output_fn=None + ) diff --git a/fla/modules/rotary.py b/fla/modules/rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..9781207fa4e014fb06701c2ca9260259bdf2858f --- /dev/null +++ b/fla/modules/rotary.py @@ -0,0 +1,512 @@ +# -*- coding: utf-8 -*- + +# Copyright (c) 2023, Tri Dao. +# https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py + +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from einops import rearrange, repeat + +from fla.utils import get_multiprocessor_count, input_guard + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2) + + +def rotary_embedding_ref(x, cos, sin, interleaved=False): + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat(cos, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + sin = repeat(sin, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], -1) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8, 16, 32] + for num_stages in [2, 3, 4] + ], + key=['B', 'H', 'D', 'INTERLEAVED'], +) +@triton.jit +def rotary_embedding_kernel( + x, + cos, + sin, + y, + cu_seqlens, + seq_offsets, # this could be int or a pointer + # Matrix dimensions + B: tl.constexpr, + T: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + TR: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, + IS_VARLEN: tl.constexpr, + INTERLEAVED: tl.constexpr, + CONJUGATE: tl.constexpr +): + i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + if not IS_VARLEN: + x = x + i_b * T*H*D + i_h * D + y = y + i_b * T*H*D + i_h * D + else: + bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1) + T = eos - bos + x = x + bos * H*D + i_h * D + y = y + bos * H*D + i_h * D + + if i_t * BT >= T: + return + + o_t = i_t * BT + tl.arange(0, BT) + if not IS_SEQLEN_OFFSETS_TENSOR: + o_cs = o_t + seq_offsets + else: + o_cs = o_t + tl.load(seq_offsets + i_b) + + if not INTERLEAVED: + # Load the 1st and 2nd halves of x, do calculation, then store to 1st and 2nd halves of out + o_r = tl.arange(0, BD // 2) + p_x = x + o_t[:, None] * H*D + o_r[None, :] + p_cos = cos + (o_cs[:, None] * R + o_r[None, :]) + p_sin = sin + (o_cs[:, None] * R + o_r[None, :]) + mask = (o_t[:, None] >= 0) & (o_t[:, None] < T) & (o_r[None, :] < R) + + b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32) + b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32) + b_x0 = tl.load(p_x, mask=mask, other=0.0).to(tl.float32) + b_x1 = tl.load(p_x + R, mask=mask, other=0.0).to(tl.float32) + if CONJUGATE: + b_sin = -b_sin + b_o0 = b_x0 * b_cos - b_x1 * b_sin + b_o1 = b_x0 * b_sin + b_x1 * b_cos + # write back result + p_y = y + (o_t[:, None] * H*D + o_r[None, :]) + tl.store(p_y, b_o0, mask=mask) + tl.store(p_y + R, b_o1, mask=mask) + else: + # We don't want to load x[0, 2, 4, ...] and x[1, 3, 5, ...] separately since both are slow. + # Instead, we load x0 = x[0, 1, 2, 3, ...] and x1 = x[1, 0, 3, 2, ...]. + # Loading x0 will be fast but x1 will be slow. + # Then we load cos = cos[0, 0, 1, 1, ...] and sin = sin[0, 0, 1, 1, ...]. + # Then we do the calculation and use tl.where to pick put the right outputs for the even + # and for the odd indices. + o_d = tl.arange(0, BD) + o_d_swap = o_d + ((o_d + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ... + o_d_repeat = tl.arange(0, BD) // 2 + p_x0 = x + o_t[:, None] * H*D + o_d[None, :] + p_x1 = x + o_t[:, None] * H*D + o_d_swap[None, :] + p_cos = cos + (o_cs[:, None] * R + o_d_repeat[None, :]) + p_sin = sin + (o_cs[:, None] * R + o_d_repeat[None, :]) + mask = (o_cs[:, None] >= 0) & (o_cs[:, None] < TR) & (o_d_repeat[None, :] < R) + + b_cos = tl.load(p_cos, mask=mask, other=1.0).to(tl.float32) + b_sin = tl.load(p_sin, mask=mask, other=0.0).to(tl.float32) + b_x0 = tl.load(p_x0, mask=mask, other=0.0).to(tl.float32) + b_x1 = tl.load(p_x1, mask=mask, other=0.0).to(tl.float32) + if CONJUGATE: + b_sin = -b_sin + b_o0 = b_x0 * b_cos + b_o1 = b_x1 * b_sin + b_y = tl.where(o_d[None, :] % 2 == 0, b_o0 - b_o1, b_o0 + b_o1) + p_y = y + (o_t[:, None] * H*D + o_d[None, :]) + tl.store(p_y, b_y, mask=mask) + + +def rotary_embedding_fwdbwd( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + seqlen_offsets: Union[int, torch.Tensor] = 0, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + interleaved: bool = False, + inplace: bool = False, + conjugate: bool = False +) -> torch.Tensor: + """ + Args: + x: [B, T, H, D]. + cos: [TR, R / 2] + sin: [TR, R / 2] + seqlen_offsets: integer or integer tensor of size (N,) + cu_seqlens: (N + 1,) or None + max_seqlen: int + + Returns: + y: [B, T, H, D] + """ + is_varlen = cu_seqlens is not None + + B, T, H, D = x.shape + if not is_varlen: + N = B + else: + assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed" + N, T = cu_seqlens.shape[0] - 1, max_seqlen + TR, R = cos.shape + assert sin.shape == cos.shape + R2 = R * 2 + + assert D <= 256, "Only support D <= 256" + assert TR >= T, "TR must be >= T" + + assert cos.dtype == sin.dtype, f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}" + assert x.dtype == cos.dtype, f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}" + + if isinstance(seqlen_offsets, torch.Tensor): + assert seqlen_offsets.shape == (N,) + assert seqlen_offsets.dtype in [torch.int32, torch.int64] + else: + assert seqlen_offsets + T <= TR + + y = torch.empty_like(x) if not inplace else x + if R2 < D and not inplace: + y[..., R2:].copy_(x[..., R2:]) + + BD = triton.next_power_of_2(R2) + BT = min(128, triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index)))) + + def grid(meta): return (triton.cdiv(T, meta['BT']), N, H) # noqa + rotary_embedding_kernel[grid]( + x, + cos, + sin, + y, + cu_seqlens, + seqlen_offsets, + B=B, + T=T, + H=H, + D=D, + R=R, + TR=TR, + BT=BT, + BD=BD, + IS_SEQLEN_OFFSETS_TENSOR=isinstance(seqlen_offsets, torch.Tensor), + IS_VARLEN=is_varlen, + INTERLEAVED=interleaved, + CONJUGATE=conjugate + ) + return y + + +class RotaryEmbeddingFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + cos, + sin, + interleaved=False, + inplace=False, + seqlen_offsets: Union[int, torch.Tensor] = 0, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + ): + y = rotary_embedding_fwdbwd( + x, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + interleaved=interleaved, + inplace=inplace, + ) + if isinstance(seqlen_offsets, int): + # Can't save int with save_for_backward + ctx.save_for_backward(cos, sin, cu_seqlens) + ctx.seqlen_offsets = seqlen_offsets + else: + ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) + ctx.seqlen_offsets = None + ctx.interleaved = interleaved + ctx.inplace = inplace + ctx.max_seqlen = max_seqlen + return y if not inplace else x + + @staticmethod + @input_guard + def backward(ctx, do): + seqlen_offsets = ctx.seqlen_offsets + if seqlen_offsets is None: + cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors + else: + cos, sin, cu_seqlens = ctx.saved_tensors + # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with + # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works. + if not ctx.interleaved and not ctx.inplace: + do = do.clone() + dx = rotary_embedding_fwdbwd( + do, + cos, + sin, + seqlen_offsets=seqlen_offsets, + cu_seqlens=cu_seqlens, + max_seqlen=ctx.max_seqlen, + interleaved=ctx.interleaved, + inplace=ctx.inplace, + conjugate=True, + ) + return dx, None, None, None, None, None, None, None + + +def rotary_embedding( + x, + cos, + sin, + interleaved=False, + inplace=False, + seqlen_offsets: Union[int, torch.Tensor] = 0, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, +): + """ + Args: + x: [B, T, H, D] + cos, sin: [TR, R//2] + interleaved: + If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). + inplace: + If True, apply rotary embedding in-place. + seqlen_offsets: [N,] or int. + Each sequence in x is shifted by this amount. + Most commonly used in inference when we have KV cache. + cu_seqlens: [N + 1,] or None + max_seqlen: int + + Returns: + out: [B, T, H, D] + """ + return RotaryEmbeddingFunction.apply( + x, + cos, + sin, + interleaved, + inplace, + seqlen_offsets, + cu_seqlens, + max_seqlen + ) + + +class RotaryEmbedding(nn.Module): + """ + The rotary position embeddings from RoFormer_ (Su et. al). + A crucial insight from the method is that the query and keys are + transformed by rotation matrices which depend on the relative positions. + + Other implementations are available in the Rotary Transformer repo_ and in + GPT-NeoX_, GPT-NeoX was an inspiration + + .. _RoFormer: https://arxiv.org/abs/2104.09864 + .. _repo: https://github.com/ZhuiyiTechnology/roformer + .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox + + If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). + A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 + Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py + """ + + def __init__( + self, + dim: int, + base: float = 10000.0, + scale_base: Optional[float] = None, + interleaved: bool = False, + pos_idx_in_fp32: bool = True, + device: Optional[torch.device] = None, + ): + """ + interleaved: + If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). + pos_idx_in_fp32: + If True, the position indices [0.0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. + This option was added because previously (before 2023-07-02), when we construct + the position indices, we use the dtype of self.inv_freq. + In most cases this would be fp32, but if the model is trained in pure bf16 (not mixed precision), then + self.inv_freq would be bf16, and the position indices are also in bf16. + Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the + embeddings for some positions will coincide. + To maintain compatibility with models previously trained in pure bf16, we add this option. + """ + super().__init__() + + self.dim = dim + self.base = float(base) + self.scale_base = scale_base + self.interleaved = interleaved + self.pos_idx_in_fp32 = pos_idx_in_fp32 + self.device = device + + # Generate and save the inverse frequency buffer (non trainable) + self.register_buffer("inv_freq", torch.empty(-(dim // -2), dtype=torch.float32, device=device), persistent=False) + + scale = None + if scale_base is not None: + scale = torch.empty(-(dim // -2), dtype=torch.float32, device=device) + self.register_buffer("scale", scale, persistent=False) + + self._seq_len_cached = 0 + self._cos_cached = None + self._sin_cached = None + self._cos_k_cached = None + self._sin_k_cached = None + + self.reset_parameters() + + def reset_parameters(self): + with torch.no_grad(): + self.inv_freq.copy_(self._compute_inv_freq(device=self.inv_freq.device)) + if self.scale_base is not None: + self.scale.copy_(self._compute_scale(device=self.scale.device)) + + def __repr__(self): + s = f"{self.__class__.__name__}(" + s += f"dim={self.dim}, " + s += f"base={self.base}, " + s += f"interleaved={self.interleaved}, " + if self.scale_base is not None: + s += f"scale_base={self.scale_base}, " + s += f"pos_idx_in_fp32={self.pos_idx_in_fp32})" + return s + + def _compute_inv_freq(self, device=None): + return 1.0 / ( + self.base + ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) + ) + + def _compute_scale(self, device=None): + return (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + 0.4 * self.dim) / (1.4 * self.dim) + + def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): + # Reset the tables if the sequence length has changed, + # if we're on a new device (possibly due to tracing for instance), + # or if we're switching from inference mode to training + if ( + seqlen > self._seq_len_cached + or self._cos_cached is None + or self._cos_cached.device != device + or self._cos_cached.dtype != dtype + or (self.training and self._cos_cached.is_inference()) + ): + self._seq_len_cached = seqlen + # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 + # And the output of arange can be quite large, so bf16 would lose a lot of precision. + # However, for compatibility reason, we add an option to use the dtype of self.inv_freq. + if self.pos_idx_in_fp32: + t = torch.arange(seqlen, device=device, dtype=torch.float32) + # We want fp32 here as well since inv_freq will be multiplied with t, and the output + # will be large. Having it in bf16 will lose a lot of precision and cause the + # cos & sin output to change significantly. + # We want to recompute self.inv_freq if it was not loaded in fp32 + if self.inv_freq.dtype != torch.float32: + inv_freq = self._compute_inv_freq(device=device) + else: + inv_freq = self.inv_freq + else: + t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) + inv_freq = self.inv_freq + # Don't do einsum, it converts fp32 to fp16 under AMP + # freqs = torch.einsum("i,j->ij", t, self.inv_freq) + freqs = torch.outer(t, inv_freq) + if self.scale is None: + self._cos_cached = torch.cos(freqs).to(dtype) + self._sin_cached = torch.sin(freqs).to(dtype) + else: + power = ( + torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) + - seqlen // 2 + ) / self.scale_base + scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") + # We want the multiplication by scale to happen in fp32 + self._cos_cached = (torch.cos(freqs) * scale).to(dtype) + self._sin_cached = (torch.sin(freqs) * scale).to(dtype) + self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) + self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + seqlen_offset: Union[int, torch.Tensor] = 0, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + q: [B, T, H, D] + k: [B, T, H, D] + seqlen_offset: + (N,) or int. Each sequence in x is shifted by this amount. + Most commonly used in inference when we have KV cache. + If it's a tensor of shape (N,), then to update the cos / sin cache, one + should pass in max_seqlen, which will update the cos / sin cache up to that length. + cu_seqlens: (N + 1,) or None + max_seqlen: int + """ + if max_seqlen is not None: + self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype) + elif isinstance(seqlen_offset, int): + self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype) + if self.scale is None: + q = rotary_embedding( + q, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen + ) + k = rotary_embedding( + k, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen + ) + + else: + q = rotary_embedding( + q, + self._cos_cached, + self._sin_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen + ) + k = rotary_embedding( + k, + self._cos_k_cached, + self._sin_k_cached, + interleaved=self.interleaved, + seqlen_offsets=seqlen_offset, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen + ) + + return q, k diff --git a/fla/modules/seq_to_dsmtp.py b/fla/modules/seq_to_dsmtp.py new file mode 100644 index 0000000000000000000000000000000000000000..f39176ec27ff6b9c59db779527ced0e2c80a53d6 --- /dev/null +++ b/fla/modules/seq_to_dsmtp.py @@ -0,0 +1,25 @@ +import torch + +def seq_to_dsmtp( + input_ids: torch.Tensor, + labels: torch.Tensor, + model_seq_len: int, + n_future_tokens: int +) -> torch.Tensor: + if labels is None: + return input_ids.unsqueeze(1), None + + B, total_len = labels.shape + assert total_len >= model_seq_len + n_future_tokens, \ + "long_input_ids must be at least model_seq_len + n_future_tokens long." + + windows = labels.unfold(dimension=1, size=n_future_tokens + 1, step=1) + + all_targets = windows[:, :, 1:] + + output_targets = all_targets[:, :model_seq_len, :] + + all_inputs = all_targets[:, :model_seq_len, :-1] + all_inputs = torch.concat([input_ids.unsqueeze(-1), all_inputs], dim=-1) + + return all_inputs.transpose(1, 2), output_targets.transpose(1, 2) \ No newline at end of file diff --git a/fla/modules/seq_to_top.py b/fla/modules/seq_to_top.py new file mode 100644 index 0000000000000000000000000000000000000000..b698ea66b570dc89d86bb0d403cffddf3c949a13 --- /dev/null +++ b/fla/modules/seq_to_top.py @@ -0,0 +1,115 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE_V': block_size_v}, num_warps=num_warp) + for block_size_v in [256, 512, 1024, 2048] + for num_warp in [1, 2, 4, 8] + ], + key=['V'], +) +@triton.jit +def _seq_to_top_kernel( + seq_ptr, + output_ptr, + B, + T_total, + T, + V, + pad_token_id, + window_size, + T_val, + stride_seq_b, + stride_seq_t, + stride_out_b, + stride_out_t, + stride_out_v, + BLOCK_SIZE_V: tl.constexpr, +): + b = tl.program_id(0) + v_block = tl.program_id(1) + + v_start = v_block * BLOCK_SIZE_V + v_end = tl.minimum(v_start + BLOCK_SIZE_V, V) + v_idx = tl.arange(0, BLOCK_SIZE_V) + v = v_start + v_idx + mask = v < V + + next_occurrence = tl.full((BLOCK_SIZE_V,), T_val, dtype=tl.int64) + + for t in range(T_total - 1, -1, -1): + token = tl.load(seq_ptr + b * stride_seq_b + t * stride_seq_t) + + token_valid = (token != pad_token_id) & (token >= 0) & (token < V) + in_block = (token >= v_start) & (token < v_end) + + if token_valid: + if in_block: + local_v = token - v_start + next_occurrence = tl.where(v_idx == local_v, t, next_occurrence) + + if t < T: + distance = next_occurrence - t + valid = (distance < window_size) + value = tl.where(valid, window_size - distance, float('-inf')) + + output_offset = ( + b * stride_out_b + + t * stride_out_t + + v * stride_out_v + ) + tl.store(output_ptr + output_offset, value, mask=mask) + +def seq_to_top( + seq: torch.Tensor, + ctx_len: int, + vocab_size: int, + window_size: int, + pad_token_id: int +) -> torch.Tensor: + """ + Triton-optimized top sequence processing with autotuned block size. + + :param seq: Input sequence of shape (B, T + window_size) + :param vocab_size: Size of the vocabulary + :param window_size: How far ahead to look for next occurrences + :param pad_token_id: Padding token ID + :return: Tensor of shape (B, T, V) with window_size - distance for tokens in window, else -inf + """ + if pad_token_id is None: + pad_token_id = -100 + + B, T_total = seq.shape + T = ctx_len + + assert T_total > window_size, "T_total must be greater than window_size to produce valid output." + + output = torch.empty((B, T, vocab_size), device=seq.device, dtype=torch.float16) + if not output.is_contiguous(): + output = output.contiguous() + + # Let autotune select the best BLOCK_SIZE_V based on vocab_size + grid = (B, triton.cdiv(vocab_size, 128)) # Start with minimum block size + + _seq_to_top_kernel[grid]( + seq, + output, + B, + T_total, + T, + vocab_size, + pad_token_id, + window_size, + T_total, + seq.stride(0), + seq.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + ) + + return output \ No newline at end of file diff --git a/fla/ops/__init__.py b/fla/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..93d8ccf71454a261c2fcc22a2757fa1457941504 --- /dev/null +++ b/fla/ops/__init__.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- + +from .abc import chunk_abc +from .attn import parallel_attn +from .based import fused_chunk_based, parallel_based +from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule +from .forgetting_attn import parallel_forgetting_attn +from .gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from .generalized_delta_rule import ( + chunk_dplr_delta_rule, + chunk_iplr_delta_rule, + fused_recurrent_dplr_delta_rule, + fused_recurrent_iplr_delta_rule +) +from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla +from .gsa import chunk_gsa, fused_recurrent_gsa +from .hgrn import fused_recurrent_hgrn +from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn +from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn +from .nsa import parallel_nsa +from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention +from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 +from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 +from .simple_gla import chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla + +__all__ = [ + 'chunk_abc', + 'parallel_attn', + 'fused_chunk_based', 'parallel_based', + 'chunk_delta_rule', 'fused_chunk_delta_rule', 'fused_recurrent_delta_rule', + 'parallel_forgetting_attn', + 'chunk_gated_delta_rule', 'fused_recurrent_gated_delta_rule', + 'chunk_dplr_delta_rule', 'chunk_iplr_delta_rule', + 'fused_recurrent_dplr_delta_rule', 'fused_recurrent_iplr_delta_rule', + 'chunk_gla', 'fused_chunk_gla', 'fused_recurrent_gla', + 'chunk_gsa', 'fused_recurrent_gsa', + 'fused_recurrent_hgrn', + 'chunk_lightning_attn', 'fused_recurrent_lightning_attn', + 'chunk_linear_attn', 'fused_chunk_linear_attn', 'fused_recurrent_linear_attn', + 'parallel_nsa', + 'chunk_retention', 'fused_chunk_retention', 'fused_recurrent_retention', 'parallel_retention', + 'chunk_rwkv6', 'fused_recurrent_rwkv6', + 'chunk_rwkv7', 'fused_recurrent_rwkv7', + 'chunk_simple_gla', 'fused_recurrent_simple_gla', 'parallel_simple_gla', +] diff --git a/fla/ops/delta_rule/README.md b/fla/ops/delta_rule/README.md new file mode 100644 index 0000000000000000000000000000000000000000..607b0d583c7ec2904c18c0f1d86fb0ec2dfdf583 --- /dev/null +++ b/fla/ops/delta_rule/README.md @@ -0,0 +1,90 @@ +# Chunkwise-form Parallelism of DeltaNet + +This section expands on the formulation presented in Appendix B of the DeltaNet paper.[^1] + +To reduce notational clutter, we focus on the first chunk, denoting $\mathbf{S}^r=\mathbf{S}_{[1]}^r$. By partially expanding the recurrence, we have: +```math +\begin{equation} +\begin{aligned} +\mathbf{S}^r &= \underbrace{\left(\prod_{i=1}^r \mathbf{I} - \beta^i \boldsymbol{k}^i \boldsymbol{k}^{i\top} \right)}_{:= \mathbf{P}^r} \cdot\mathbf{S}^{0} + \overbrace{\sum_{i=1}^{r} \underbrace{\left(\prod_{j=i+1}^r \mathbf{I} - \beta^j \boldsymbol{k}^j \boldsymbol{k}^{j\top} \right)}_{:= \mathbf{P}_{i+1}^r}\beta^i \boldsymbol{k}^i\boldsymbol{v}^{i\top}}^{:=\mathbf{H}^r} \\ +&=\mathbf{P}^r \cdot \mathbf{S}^{0} + \mathbf{H}^r +\end{aligned} +\end{equation} +``` + +where $\mathbf{P}_i^r$ involves cumulative products of generalized Householder matrices. +We abbreviate $\mathbf{P}_1^r$ as $\mathbf{P}^r$. +This can be optimized using the classical WY representation: +```math +\begin{equation} +\mathbf{P}^{r} = \mathbf{I} - \sum_{i=1}^{r}\boldsymbol{k}^i\boldsymbol{w}^{i\top} \in \mathbb{R}^{d_k \times d_k};\qquad +\boldsymbol{w}^r = \beta^r \left(\boldsymbol{k}^r - \sum_{i=1}^{r-1} \left(\boldsymbol{k}^{r\top}\boldsymbol{k}^i \right)\boldsymbol{w}^i \right) \in \mathbb{R}^{d_k} +\end{equation} +``` + +We prove this by induction: +```math +\begin{align*} +\mathbf{P}^{r} &= \prod_{i=1}^r \mathbf{I} - \beta^i \boldsymbol{k}^i \boldsymbol{k}^{i\top} \\ +&= \left(\mathbf{I} - \beta^r \boldsymbol{k}^r \boldsymbol{k}^{r\top}\right)\mathbf{P}^{r-1} \\ +&= \left(\mathbf{I} - \beta^r \boldsymbol{k}^r \boldsymbol{k}^{r\top}\right)\left(\mathbf{I} - \sum_{i=1}^{r-1}\boldsymbol{k}^i\boldsymbol{w}^{i\top}\right) \\ +&= \mathbf{I} - \sum_{i=1}^{r-1}\boldsymbol{k}^i\boldsymbol{w}^{i\top} - \beta^r \boldsymbol{k}^r \boldsymbol{k}^{r\top} + \beta^r\boldsymbol{k}^r \boldsymbol{k}^{r\top} \left(\sum_{i=1}^{r-1}\boldsymbol{k}^i\boldsymbol{w}^{i\top}\right) \\ +&= \mathbf{I} - \sum_{i=1}^{r-1}\boldsymbol{k}^i\boldsymbol{w}^{i\top} - \beta^r \boldsymbol{k}^r \left(\boldsymbol{k}^{r} - \left(\sum_{i=1}^{r-1}\left(\boldsymbol{k}^{r\top} \boldsymbol{k}^i\right)\boldsymbol{w}^{i}\right) \right)^\top \\ +&= \mathbf{I} - \sum_{i=1}^{r}\boldsymbol{k}^i\boldsymbol{w}^{i\top} +\end{align*} +``` + +Similarly, $\mathbf{H}^r$ can be represented as: +```math +\begin{equation} +\mathbf{H}^{r} = \sum_{i=1}^{r} \boldsymbol{k}^i \boldsymbol{u}^{i\top} \in \mathbb{R}^{d_k \times d_v};\qquad \boldsymbol{u}^r = \beta^r \left(\boldsymbol{v}^r - \sum_{i=1}^{r-1} \left(\boldsymbol{k}^{r\top}\boldsymbol{k}^i\right) \boldsymbol{u}^i \right)\in \mathbb{R}^{d_v} +\end{equation} +``` + +This can also be proven by induction: +```math +\begin{align*} +\mathbf{H}^{r} &= \sum_{i=1}^{r} \mathbf{P}_{i+1}^r \beta^i \boldsymbol{k}^i \boldsymbol{v}^{i\top}\\ +&= \left(\mathbf{I} - \beta^r \boldsymbol{k}^r \boldsymbol{k}^{r\top}\right) \mathbf{H}^{r-1} + \beta^r \boldsymbol{k}^r \boldsymbol{v}^{r\top}\\ +&= \sum_{i=1}^{r-1}\boldsymbol{k}^i \boldsymbol{u}^{i\top} - \beta^r \boldsymbol{k}^r \boldsymbol{k}^{r\top} \sum_{i=1}^{r-1}\boldsymbol{k}^i \boldsymbol{u}^{i\top} +\beta^r \boldsymbol{k}^r \boldsymbol{v}^{r\top}\\ +&= \sum_{i=1}^{r-1}\boldsymbol{k}^i \boldsymbol{u}^{i\top} + \boldsymbol{k}^r \left(\beta^r \boldsymbol{v}^{r\top}-\beta^r \boldsymbol{k}^{r\top} \sum_{i=1}^{r-1}\boldsymbol{k}^i \boldsymbol{u}^{i\top}\right) \\ +&= \sum_{i=1}^{r-1}\boldsymbol{k}^i \boldsymbol{u}^{i\top} + \boldsymbol{k}^r \beta^r\left(\boldsymbol{v}^{r}-\sum_{i=1}^{r-1}\left(\boldsymbol{k}^{r\top}\boldsymbol{k}^{i}\right)\boldsymbol{u}^{i} \right)^\top \\ +&=\sum_{i=1}^{r} \boldsymbol{k}^i \boldsymbol{u}^{i\top} +\end{align*} +``` + +In matrix form, $\mathbf{P}$ and $\mathbf{H}$ can be written as: +```math +\begin{equation} +\mathbf{P}=\mathbf{I}-\mathbf{K}^\top\mathbf{W} \in \mathbb{R}^{d_k \times d_k}, \qquad\mathbf{H}=\mathbf{K}^\top\mathbf{U} \in \mathbb{R}^{d_k\times d_v} +\end{equation} +``` + +Now we can derive the matrix form of $\mathbf{W}$ and $\mathbf{U}$: +```math +\begin{align*} +\mathbf{W} &= \mathrm{diag}(\beta) \mathbf{K} - \mathrm{tril}(\mathrm{diag}(\beta) \mathbf{K}\mathbf{K}^\top, -1)\mathbf{W}\\ +\left(\mathbf{I} + \mathrm{tril}(\mathrm{diag}(\beta) \mathbf{K}\mathbf{K}^\top, -1)\right) \mathbf{W} &= \mathrm{diag}(\beta) \mathbf{K} +\end{align*} +``` +A similar process holds for $\mathbf{U}$. We can further write $\mathbf{W}$ and $\mathbf{U}$ in matrix form: +```math +\begin{align*} +\mathbf{T} &= \left(\mathbf{I} + \mathrm{tril}\left(\mathrm{diag}(\beta)\mathbf{K} \mathbf{K}^\top,-1\right)\right)^{-1}\mathrm{diag}\left(\beta\right)\in \mathbb{R}^{C \times C}\\ +\mathbf{W} &= \mathbf{T} \mathbf{K}\in \mathbb{R}^{C \times d_k}\\ +\mathbf{U} &= \mathbf{T}\mathbf{V}\in \mathbb{R}^{C \times d_v} +\end{align*} +``` + +Substituting these back into the original equations yields a hardware-efficient chunkwise algorithm for DeltaNet that leverages matrix multiplications, enabling tensor core based GPU optimization: +```math +\begin{equation} +\begin{aligned} +\mathbf{S} &= \mathbf{P}\cdot\mathbf{S}^0 + \mathbf{H} \\ +&= \mathbf{S}^0 + \mathbf{K}^\top (\mathbf{U} -\mathbf{W} \mathbf{S}^0) \in \mathbb{R}^{d_k \times d_v}\\ +\mathbf{O} &= \mathbf{Q} \mathbf{S}^0 + (\mathbf{Q} \mathbf{K}^{\top} \odot \mathbf{M}) \left(\mathbf{U} - \mathbf{W} \mathbf{S}^0\right) \in \mathbb{R}^{C \times d_v} +\end{aligned} +\end{equation} +``` + +[^1]: https://arxiv.org/abs/2406.06484 diff --git a/fla/ops/delta_rule/__init__.py b/fla/ops/delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e0acb6a7d0e4eec9a8dc697615604783b8858d13 --- /dev/null +++ b/fla/ops/delta_rule/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- + +from .chunk import chunk_delta_rule +from .fused_chunk import fused_chunk_delta_rule +from .fused_recurrent import fused_recurrent_delta_rule + +__all__ = [ + 'fused_chunk_delta_rule', + 'fused_recurrent_delta_rule', + 'chunk_delta_rule' +] diff --git a/fla/ops/delta_rule/fused_chunk.py b/fla/ops/delta_rule/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..6347fb9af47d3d9f82c03ea9aedbfa09fc1bfbc1 --- /dev/null +++ b/fla/ops/delta_rule/fused_chunk.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- + +def fused_chunk_delta_rule( + **kwargs +): + raise NotImplementedError("fused_chunk_delta_rule is deprecated. Please use chunk_delta_rule instead.") diff --git a/fla/ops/delta_rule/fused_recurrent.py b/fla/ops/delta_rule/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..ef57cc4c2ee6e51deeb138073007854f67a62b71 --- /dev/null +++ b/fla/ops/delta_rule/fused_recurrent.py @@ -0,0 +1,607 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from einops import rearrange + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_delta_rule_fwd_kernel( + q, + k, + v, + u, + beta, + o, + h0, + ht, + offsets, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if USE_OFFSETS: + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + if HEAD_FIRST: + p_q = q + i_nh * T*K + i_k * BK + tl.arange(0, BK) + p_k = k + i_nh * T*K + i_k * BK + tl.arange(0, BK) + p_v = v + i_nh * T*V + i_v * BV + tl.arange(0, BV) + p_u = u + i_nh * T*V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + i_nh * T*V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + i_nh * T + p_o = o + (i_k * B*H + i_nh) * T*V + i_v * BV + tl.arange(0, BV) + else: + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_u = u + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + bos * H + i_h + p_o = o + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + + mask_k = (i_k * BK + tl.arange(0, BK)) < K + mask_v = (i_v * BV + tl.arange(0, BV)) < V + mask_h = mask_k[None, :] & mask_v[:, None] + + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_v_minus = tl.sum(b_h * b_k[None, :], axis=1) + b_v -= b_v_minus + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + tl.store(p_u, b_v.to(p_v.dtype.element_ty), mask=mask_v) + b_v *= b_beta + b_h += b_k[None, :] * b_v[:, None] + b_o = b_h * b_q[None, :] + b_o = tl.sum(b_o, axis=1) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += K if HEAD_FIRST else H*K + p_k += K if HEAD_FIRST else H*K + p_o += V if HEAD_FIRST else H*V + p_v += V if HEAD_FIRST else H*V + p_u += V if HEAD_FIRST else H*V + p_beta += (1 if HEAD_FIRST else H) * (V if IS_BETA_HEADWISE else 1) + + if STORE_FINAL_STATE: + p_ht = ht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_delta_rule_bwd_kernel( + q, + k, + v, + beta, + h0, + dh0, + dht, + do, + dq, + dk, + dv, + db, + offsets, + scale, + B: tl.constexpr, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NK: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar + USE_INITIAL_STATE: tl.constexpr, # whether to use dh0 + USE_FINAL_STATE_GRADIENT: tl.constexpr, # whether to use dht + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_h = i_nh // H, i_nh % H + if USE_OFFSETS: + bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + mask_k = i_k * BK + tl.arange(0, BK) < K + mask_v = i_v * BV + tl.arange(0, BV) < V + + if HEAD_FIRST: + p_q = q + i_nh * T*K + i_k * BK + tl.arange(0, BK) + (T - 1) * K + p_k = k + i_nh * T*K + i_k * BK + tl.arange(0, BK) + (T - 1) * K + p_v = v + i_nh * T*V + i_v * BV + tl.arange(0, BV) + (T - 1) * V + p_do = do + i_nh * T*V + i_v * BV + tl.arange(0, BV) + (T - 1) * V + p_dk = dk + (i_v * B*H + i_nh) * T*K + i_k * BK + tl.arange(0, BK) + (T - 1) * K + p_dv = dv + (i_k * B*H + i_nh) * T*V + i_v * BV + tl.arange(0, BV) + (T - 1) * V + if IS_BETA_HEADWISE: + p_beta = beta + i_nh * T*V + i_v * BV + tl.arange(0, BV) + (T - 1) * V + p_dbeta = db + (i_v * NK*B*H + i_k * B*H + i_nh) * T*V + tl.arange(0, BV) + (T - 1) * V + else: + p_beta = beta + i_nh * T + T - 1 + p_dbeta = db + (i_v * B*H + i_nh) * T + T - 1 + else: + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K + p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + if IS_BETA_HEADWISE: + p_beta = beta + (bos + T - 1) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) + p_dbeta = db + ((i_v * NK + i_k) * all + bos + T - 1) * H*V + i_h * V + tl.arange(0, BV) + else: + p_beta = beta + (bos + T - 1) * H + i_h + p_dbeta = db + (i_v * all + bos + T - 1) * H + i_h + + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + if USE_FINAL_STATE_GRADIENT: + p_ht = dht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + b_dh += tl.load(p_ht, mask=mask_k[:, None] & mask_v[None, :], other=0).to(tl.float32) + + for _ in range(T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_dh += b_q[:, None] * b_do[None, :] + b_dk = tl.sum(b_dh * (b_v * b_beta)[None, :], axis=1) + b_dv = tl.sum(b_dh * b_k[:, None], axis=0) + + b_db = b_dv * b_v if IS_BETA_HEADWISE else tl.sum(b_dv * b_v) + b_dv = b_dv * b_beta + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) + if IS_BETA_HEADWISE: + tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty), mask=mask_v) + else: + tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty)) + + b_dh -= b_k[:, None] * b_dv[None, :] + + p_q -= K if HEAD_FIRST else H*K + p_k -= K if HEAD_FIRST else H*K + p_v -= V if HEAD_FIRST else H*V + p_do -= V if HEAD_FIRST else H*V + p_dk -= K if HEAD_FIRST else H*K + p_dv -= V if HEAD_FIRST else H*V + p_dbeta -= (1 if HEAD_FIRST else H) * (V if IS_BETA_HEADWISE else 1) + p_beta -= (1 if HEAD_FIRST else H) * (V if IS_BETA_HEADWISE else 1) + + if USE_INITIAL_STATE: + p_dh0 = dh0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_k[:, None] & mask_v[None, :]) + + tl.debug_barrier() + + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + if HEAD_FIRST: + p_q = q + i_nh * T*K + i_k * BK + tl.arange(0, BK) + p_k = k + i_nh * T*K + i_k * BK + tl.arange(0, BK) + p_v = v + i_nh * T*V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + i_nh * T*V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + i_nh * T + p_do = do + i_nh * T*V + i_v * BV + tl.arange(0, BV) + p_dq = dq + (i_v * B*H + i_nh) * T*K + i_k * BK + tl.arange(0, BK) + p_dk = dk + (i_v * B*H + i_nh) * T*K + i_k * BK + tl.arange(0, BK) + p_dv = dv + (i_k * B*H + i_nh) * T*V + i_v * BV + tl.arange(0, BV) + else: + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + if IS_BETA_HEADWISE: + p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + else: + p_beta = beta + bos * H + i_h + p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + p_dq = dq + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + + if USE_INITIAL_STATE: + mask_h = mask_k[:, None] & mask_v[None, :] + p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in range(0, T): + b_dk = tl.load(p_dk, mask=mask_k, other=0).to(tl.float32) + b_dv = tl.load(p_dv, mask=mask_v, other=0).to(tl.float32) + b_dk -= tl.sum(b_dv[None, :] * b_h, axis=1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + + b_h += b_k[:, None] * b_v[None, :] + b_dq = b_h * b_do[None, :] + d_q = tl.sum(b_dq, axis=1) * scale + tl.store(p_dq, d_q.to(p_dq.dtype.element_ty), mask=mask_k) + + p_k += K if HEAD_FIRST else H*K + p_v += V if HEAD_FIRST else H*V + p_do += V if HEAD_FIRST else H*V + p_dq += K if HEAD_FIRST else H*K + p_dk += K if HEAD_FIRST else H*K + p_dv += V if HEAD_FIRST else H*V + p_beta += (1 if HEAD_FIRST else H) * (V if IS_BETA_HEADWISE else 1) + + +def fused_recurrent_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + if head_first: + B, H, T, K, V = *k.shape, v.shape[-1] + else: + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if offsets is None else len(offsets) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 1 + num_warps = 1 + + o = q.new_empty(NK, *v.shape) + if output_final_state: + final_state = q.new_empty(N, H, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NV, NK, N * H) + u = torch.empty_like(v) + fused_recurrent_delta_rule_fwd_kernel[grid]( + q, + k, + v, + u, + beta, + o, + initial_state, + final_state, + offsets, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + HEAD_FIRST=head_first, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o, u, final_state + + +def fused_recurrent_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + dht: torch.Tensor, + do: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if head_first: + B, H, T, K, V = *k.shape, v.shape[-1] + else: + B, T, H, K, V = *k.shape, v.shape[-1] + N = B if offsets is None else len(offsets) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 1 + num_warps = 2 + + beta_vector = beta.ndim == v.ndim + + dq = q.new_empty(NV, *q.shape) + dk = q.new_empty(NV, *k.shape) + dv = q.new_empty(NK, *v.shape) + if beta_vector: + db = q.new_empty(NV, NK, B, H, T, V) if head_first else q.new_empty(NV, NK, B, T, H, V) + else: + db = q.new_empty(NV, B, H, T) if head_first else q.new_empty(NV, B, T, H) + grid = (NV, NK, N * H) + + if initial_state is not None and initial_state.requires_grad: + dh0 = torch.empty_like(initial_state, dtype=torch.float32) + else: + dh0 = None + + fused_recurrent_delta_rule_bwd_kernel[grid]( + q, + k, + v, + beta, + initial_state, + dh0, + dht, + do, + dq, + dk, + dv, + db, + offsets, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + NK=NK, + IS_BETA_HEADWISE=beta_vector, + HEAD_FIRST=head_first, + num_warps=num_warps, + num_stages=num_stages + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + db = db.sum((0, 1)) if beta_vector else db.sum(0) + + return dq, dk, dv, db, dh0 + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True, + use_qk_l2norm_in_kernel: bool = False + ): + q_orig = q + k_orig = k + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + o, u, final_state = fused_recurrent_delta_rule_fwd( + q=q, + k=k, + v=v, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + offsets=offsets, + head_first=head_first + ) + + ctx.save_for_backward(q_orig, k_orig, u, beta, initial_state) + ctx.scale = scale + ctx.offsets = offsets + ctx.head_first = head_first + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + q, k, v, beta, initial_state = ctx.saved_tensors + if ctx.use_qk_l2norm_in_kernel: + q, q_orig = l2norm_fwd(q), q + k, k_orig = l2norm_fwd(k), k + dq, dk, dv, db, dh0 = fused_recurrent_delta_rule_bwd( + q=q, + k=k, + v=v, + beta=beta, + dht=dht, + do=do, + scale=ctx.scale, + initial_state=initial_state, + offsets=ctx.offsets, + head_first=ctx.head_first + ) + if ctx.use_qk_l2norm_in_kernel: + dq, dk = l2norm_bwd(q_orig, dq), l2norm_bwd(k_orig, dk) + return dq.to(q), dk.to(k), dv.to(v), db.to(beta), None, dh0, None, None, None, None + + +@torch.compiler.disable +def fused_recurrent_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = True, + use_qk_l2norm_in_kernel: bool = False +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + beta (torch.Tensor): + betas of shape `[B, T, H]` if `head_first=False` else `(B, H, T)`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format, which is not supported for variable-length inputs. + Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.delta_rule import fused_recurrent_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> beta = torch.rand(B, T, H, device='cuda').sigmoid() + >>> h0 = torch.randn(B, H, K, V, device='cuda') + >>> o, ht = fused_recurrent_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_recurrent_delta_rule( + q, k, v, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert ht.allclose(ht_var) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if head_first: + raise RuntimeError( + "Sequences with variable lengths are not supported for head-first mode" + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + if beta is None: + beta = torch.ones_like(q[..., 0]) + if head_first: + q, k, v = map(lambda x: rearrange(x, 'b h t d -> b t h d'), (q, k, v)) + beta = rearrange(beta, 'b h t -> b t h') + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + False, + use_qk_l2norm_in_kernel + ) + if head_first: + o = rearrange(o, 'b t h v -> b h t v') + return o, final_state diff --git a/fla/ops/delta_rule/parallel.py b/fla/ops/delta_rule/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..722f2dec76271a08398dc105d7d0b8a33917b62e --- /dev/null +++ b/fla/ops/delta_rule/parallel.py @@ -0,0 +1,394 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Tuple + +import torch +import triton +import triton.language as tl +from einops import rearrange + +from fla.ops.delta_rule.wy_fast import fwd_prepare_T +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BT', 'K', 'V'], +) +@triton.jit(do_not_specialize=['T']) +def chunk_transform_qk_fwd_kernel( + q, + k, + v, + beta, + o, + A, + q_new, + k_new, + A_local, + scale, + T, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + BT: tl.constexpr, + OUTPUT_ATTENTIONS: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(p_q.dtype.element_ty) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + p_T = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_T = tl.load(p_T, boundary_check=(0, 1)) + + o_i = tl.arange(0, BT) + m_t = o_i[:, None] >= o_i[None, :] + b_qk = tl.where(m_t, tl.dot(b_q, tl.trans(b_k), allow_tf32=False), 0).to(b_q.dtype) + m_t = o_i[:, None] > o_i[None, :] + b_kk = tl.where(m_t, tl.dot(b_k, tl.trans(b_k), allow_tf32=False), 0).to(b_k.dtype) + + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) + b_beta = tl.load(p_beta, boundary_check=(0, )) + b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) + + b_qkT = tl.dot(b_qk, b_T, allow_tf32=False).to(b_k.dtype) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_a, b_qkT.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + b_kkT = tl.dot(b_kk, b_T, allow_tf32=False).to(b_k.dtype) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + tl.store(p_o, tl.dot(b_qkT, b_v).to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + p_q_new = tl.make_block_ptr(q_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + tl.store(p_q_new, (b_q - tl.dot(b_qkT, b_k_beta, allow_tf32=False)).to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + + p_k_new = tl.make_block_ptr(k_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + b_k_new = b_k - tl.dot(tl.trans(b_kkT), b_k_beta, allow_tf32=False) + tl.store(p_k_new, b_k_new.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_transform_qk_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + chunk_size: int, + output_attentions: bool +): + B, H, T, K = k.shape + BT = chunk_size + q_new = torch.empty_like(q) + k_new = torch.empty_like(k) + o = torch.empty_like(v) + grid = (triton.cdiv(T, BT), B*H) + V = v.shape[-1] + A_local = torch.empty_like(A) if output_attentions else None + chunk_transform_qk_fwd_kernel[grid]( + q, + k, + v, + beta, + o, + A, + q_new, + k_new, + A_local, + scale=scale, + T=T, + K=K, + V=V, + BT=BT, + BK=triton.next_power_of_2(K), + BV=triton.next_power_of_2(V), + OUTPUT_ATTENTIONS=output_attentions + ) + return q_new, k_new, o, A_local + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=1), + triton.Config({}, num_warps=2), + ], + key=['BT'], +) +@triton.jit(do_not_specialize=['T']) +def save_intra_chunk_attn( + A, + A_local, + T, + BT: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + p_A = tl.make_block_ptr(A + i_bh * T * T, (T, T), (T, 1), (i_t * BT, i_t * BT), (BT, BT), (1, 0)) + p_A_local = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A_local = tl.load(p_A_local, boundary_check=(0, 1)) + tl.store(p_A, b_A_local.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def parallel_delta_rule_fwd_kernel( + q, + k, + k2, # original k + v, + beta, + o, + o_new, + attn, + T, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + OUTPUT_ATTENTIONS: tl.constexpr +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [BT, BK] + b_q = tl.zeros([BT, BK], dtype=tl.float32) + b_q += tl.load(p_q, boundary_check=(0, 1)) + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + + # As opposed to Flashattention, this kernel requires scanning the KV blocks from right to left + # Q block and K block have overlap. + # masks required + for offset in range((i_t + 1) * BT - 2 * BS, i_t * BT - BS, -BS): + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) + p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS] + b_beta = tl.load(p_beta, boundary_check=(0,)) + # [BT, BS] + m_s = tl.arange(0, BT) >= (offset - i_t*BT + BS) + b_s = tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False) + b_s = tl.where(m_s[:, None], b_s, 0) + + b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + # Q block and K block have no overlap + # no need for mask, thereby saving flops + for offset in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) + p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) + + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS] + b_beta = tl.load(p_beta, boundary_check=(0,)) + # [BT, BS] + b_s = (tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False)) + # [BT, BV] + b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False).to(b_q.dtype) + + if OUTPUT_ATTENTIONS: + p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) + tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + + p_o_new = tl.make_block_ptr(o_new + i_bh * T*V, (T, V), (V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + tl.store(p_o_new, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +class ParallelDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, beta, scale, output_attentions): + B, H, T, K, V = *k.shape, v.shape[-1] + assert q.shape[-1] <= 128, 'The maximum supported sequence length is 128.' + BT, BS = 128, 32 + BK = triton.next_power_of_2(k.shape[-1]) + BV = triton.next_power_of_2(v.shape[-1]) + assert BT % BS == 0 + + A = fwd_prepare_T(k, beta, BS) + attn = q.new_zeros(B, H, T, T) if output_attentions else None + q_new, k_new, o, A_local = chunk_transform_qk_fwd( + q, + k, + v, + beta, + A, + scale, + BS, + output_attentions + ) + + num_stages = 3 if K <= 64 else 2 + num_warps = 4 + grid = (triton.cdiv(T, BT), B * H) + o_new = torch.empty_like(o) + + parallel_delta_rule_fwd_kernel[grid]( + q=q_new, + k=k_new, + k2=k, + v=v, + beta=beta, + o=o, + o_new=o_new, + attn=attn, + T=T, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + num_stages=num_stages, + num_warps=num_warps + ) + + if output_attentions: + grid = (triton.cdiv(T, BS), B * H) + save_intra_chunk_attn[grid]( + A=attn, + A_local=A_local, + T=T, + BT=BS + ) + return o_new.to(q.dtype), attn + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, d_attn=None): + raise NotImplementedError('Backward pass is not implemented. Stay tuned!') + + +def parallel_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + output_attentions: bool = False, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + beta (torch.Tensor): + betas of shape `[B, H, T]` if `head_first=True` else `[B, T, H]`. + scale (Optional[int]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + output_attentions (bool): + Whether to output the materialized attention scores of shape [B, H, T, T]. Default: `False`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. + Default: `True`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + attn (torch.Tensor): + Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. + """ + if not head_first: + q, k, v, beta = map(lambda x: x.transpose(1, 2), (q, k, v, beta)) + o, attn = ParallelDeltaRuleFunction.apply(q, k, v, beta, scale, output_attentions) + if not head_first: + o = o.transpose(1, 2) + return o, attn + + +def naive_delta_rule_parallel(q, k, v, beta, BM=128, BN=32): + b, h, l, d_k = q.shape + q = q * (d_k ** -0.5) + v = v * beta[..., None] + k_beta = k * beta[..., None] + # compute (I - tri(diag(beta) KK^T))^{-1} + q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BN), [q, k, v, k_beta]) + mask = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=0) + T = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, BN): + T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) + T = T + torch.eye(BN, dtype=q.dtype, device=q.device) + + mask2 = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=1) + A_local = (q @ k.transpose(-1, -2)).masked_fill(mask2, 0) @ T + o_intra = A_local @ v + + # apply cumprod transition matrices on k to the last position within the chunk + k = k - ((k @ k.transpose(-1, -2)).masked_fill(mask, 0) @ T).transpose(-1, -2) @ k_beta + # apply cumprod transition matrices on q to the first position within the chunk + q = q - A_local @ k_beta + o_intra = A_local @ v + + A = torch.zeros(b, h, l, l, device=q.device) + + q, k, v, k_beta, o_intra = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, v, k_beta, o_intra]) + o = torch.empty_like(v) + for i in range(0, l, BM): + q_i = q[:, :, i:i+BM] + o_i = o_intra[:, :, i:i+BM] + # intra block + for j in range(i + BM - 2 * BN, i-BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + mask = torch.arange(i, i+BM) >= (j + BN) + A_ij = A_ij.masked_fill_(~mask[:, None].to(A_ij.device), 0) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + # inter block + for j in range(i - BN, -BN, -BN): + k_j = k[:, :, j:j+BN] + A_ij = q_i @ k_j.transpose(-1, -2) + A[:, :, i:i+BM, j:j+BN] = A_ij + q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] + o_i += A_ij @ v[:, :, j:j+BN] + o[:, :, i:i+BM] = o_i + + for i in range(0, l//BN): + A[:, :, i*BN:i*BN+BN, i*BN:i*BN+BN] = A_local[:, :, i] + + return o, A diff --git a/fla/ops/forgetting_attn/parallel.py b/fla/ops/forgetting_attn/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..88fea7f29b238e488848711ed894cb6cae7ea91b --- /dev/null +++ b/fla/ops/forgetting_attn/parallel.py @@ -0,0 +1,708 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +import triton +import triton.language as tl +from einops import rearrange, reduce + +from fla.ops.common.utils import prepare_chunk_indices +from fla.ops.utils import chunk_global_cumsum, chunk_local_cumsum +from fla.ops.utils.op import div, exp, log +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, input_guard + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else []) + for num_stages in [2, 3, 4, 5] + ], + key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV'], +) +@triton.jit +def parallel_forgetting_attn_fwd_kernel( + q, + k, + v, + g, + o, + lse, + scale, + offsets, + indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT,] + b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + # [BT, BV] + b_o = tl.zeros([BT, BV], dtype=tl.float32) + + b_m = tl.full([BT], float('-inf'), dtype=tl.float32) + b_acc = tl.zeros([BT], dtype=tl.float32) + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_gk = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_k = i_s + tl.arange(0, BS) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS,] + b_gk = tl.load(p_gk, boundary_check=(0,)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) + b_gq[:, None] - b_gk[None, :] + b_s = tl.where(o_q[:, None] >= o_k[None, :], b_s, float('-inf')) + + # [BT] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [BT, BS] + b_p = exp(b_s - b_m[:, None]) + # [BT] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [BT, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + + for i_s in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_gk = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS,] + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + + b_gn = tl.load(g + (bos + min(i_s + BS, T) - 1) * HQ + i_hq).to(tl.float32) + b_gp = tl.load(g + (bos + i_s - 1) * HQ + i_hq).to(tl.float32) if i_s % BT > 0 else 0. + # [BT, BS] + b_s = tl.dot(b_q, b_k) + b_gq[:, None] + (b_gn - b_gk)[None, :] + + b_gq += b_gn - b_gp + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [BT, BS] + b_p = exp(b_s - b_m[:, None]) + # [BT] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [BT, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + + b_o = div(b_o, b_acc[:, None]) + b_m += log(b_acc) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) + + +@triton.jit +def parallel_forgetting_attn_bwd_kernel_preprocess( + o, + do, + delta, + B: tl.constexpr, + V: tl.constexpr +): + i_n = tl.program_id(0) + o_d = tl.arange(0, B) + m_d = o_d < V + + b_o = tl.load(o + i_n * V + o_d, mask=m_d, other=0) + b_do = tl.load(do + i_n * V + o_d, mask=m_d, other=0).to(tl.float32) + b_delta = tl.sum(b_o * b_do) + + tl.store(delta + i_n, b_delta.to(delta.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4] + ([8] if check_shared_mem('hopper') else []) + for num_stages in [2, 3, 4] + ], + key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_forgetting_attn_bwd_kernel_dq( + q, + k, + v, + g, + lse, + delta, + do, + dq, + dg, + scale, + offsets, + indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_g = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos * HQ + i_hq), (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BT] + b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + # [BT, BK] + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + # [BT] + b_dg = tl.zeros([BT,], dtype=tl.float32) + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + p_gk = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_k = i_s + tl.arange(0, BS) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS,] + b_gk = tl.load(p_gk, boundary_check=(0,)) + # [BT, BS] + b_s = tl.dot(b_q, b_k) + (b_gq - b_lse)[:, None] - b_gk[None, :] + b_p = exp(tl.where(o_q[:, None] >= o_k[None, :], b_s, float('-inf'))) + + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + # [BT] + b_dg += tl.sum(b_ds, 1) + + for i_s in range(i_t * BT - BS, -BS, -BS): + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + p_gk = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BS,] + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + + b_gn = tl.load(g + (bos + min(i_s + BS, T) - 1) * HQ + i_hq).to(tl.float32) + b_gp = tl.load(g + (bos + i_s - 1) * HQ + i_hq).to(tl.float32) if i_s % BT > 0 else 0. + # [BT, BS] + b_s = tl.dot(b_q, b_k) + (b_gq - b_lse)[:, None] + (b_gn - b_gk)[None, :] + b_p = exp(b_s) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp - b_delta[:, None]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + # [BT] + b_dg += tl.sum(b_ds, 1) + + b_gq += b_gn - b_gp + + b_dq *= scale + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['B', 'H', 'G', 'K', 'V', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_forgetting_attn_bwd_kernel_dkv( + q, + k, + v, + g, + lse, + delta, + do, + dk, + dv, + dg, + offsets, + indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + i_n = i_b + bos, eos = i_n * T, i_n * T + T + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gk = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos * HQ + i_hq), (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BT, BV], dtype=tl.float32) + # [BT] + b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + b_dg = tl.zeros([BT,], dtype=tl.float32) + + o_k = i_t * BT + tl.arange(0, BT) + m_k = o_k < T + for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_gq = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_q = i_s + tl.arange(0, BS) + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BS] + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + + m_q = o_q < T + m_s = (o_k[:, None] <= o_q[None, :]) & m_k[:, None] & m_q[None, :] + # [BT, BS] + b_s = tl.dot(b_k, tl.trans(b_q)) - b_gk[:, None] + (b_gq - b_lse)[None, :] + b_p = tl.where(m_s, exp(b_s), 0) + # [BT, BS] @ [BS, BV] -> [BT, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BT, BS] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + # [BT] + b_dg -= tl.sum(b_ds, 1) + + b_gk -= tl.load(g + (bos + min((i_t + 1) * BT, T) - 1) * HQ + i_hq).to(tl.float32) + for i_s in range((i_t + 1) * BT, T, BS): + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) + p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + p_gq = tl.make_block_ptr(g + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) + + # [BS] + o_q = i_s + tl.arange(0, BS) + # [BS, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BS, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [BS] + b_lse = tl.load(p_lse, boundary_check=(0,)) + b_delta = tl.load(p_delta, boundary_check=(0,)) + b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + + b_gn = tl.load(g + (bos + min(i_s + BS, T) - 1) * HQ + i_hq).to(tl.float32) + b_gp = tl.load(g + (bos + i_s - 1) * HQ + i_hq).to(tl.float32) if i_s % BT > 0 else 0. + # [BT, BS] + b_s = tl.dot(b_k, tl.trans(b_q)) - (b_gk + b_gp)[:, None] + (b_gq - b_lse)[None, :] + b_p = exp(b_s) + # [BT, BS] @ [BS, BV] -> [BT, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BT, BV] @ [BV, BS] -> [BT, BS] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BT, BS] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BT, BS] @ [BS, BK] -> [BT, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + # [BT] + b_dg -= tl.sum(b_ds, 1) + + b_gk -= b_gn - b_gp + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +def parallel_forgetting_attn_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float, + chunk_size: int = 128, + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + BT = chunk_size + BK = max(16, triton.next_power_of_2(K)) + assert V <= 256, "V must be less than or equal to 256" + if check_shared_mem('hopper'): + BS = min(64, max(16, triton.next_power_of_2(T))) + else: + BS = min(32, max(16, triton.next_power_of_2(T))) + BV = min(256, max(16, triton.next_power_of_2(V))) + NV = triton.cdiv(V, BV) + NT = triton.cdiv(T, BT) if offsets is None else len(indices) + + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + grid = (NV, NT, B * HQ) + parallel_forgetting_attn_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + o=o, + lse=lse, + scale=scale, + offsets=offsets, + indices=indices, + B=B, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + ) + return o, lse + + +def parallel_forgetting_attn_bwd_preprocess( + o: torch.Tensor, + do: torch.Tensor +): + V = o.shape[-1] + delta = torch.empty_like(o[..., 0], dtype=torch.float) + parallel_forgetting_attn_bwd_kernel_preprocess[(delta.numel(),)]( + o=o, + do=do, + delta=delta, + B=triton.next_power_of_2(V), + V=V, + ) + return delta + + +def parallel_forgetting_attn_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + scale: float = None, + chunk_size: int = 128, + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + BT = chunk_size + BS = min(32, max(16, triton.next_power_of_2(T))) + BK = max(16, triton.next_power_of_2(K)) + BV = max(16, triton.next_power_of_2(V)) + NV = triton.cdiv(V, BV) + NT = triton.cdiv(T, BT) if offsets is None else len(indices) + + delta = parallel_forgetting_attn_bwd_preprocess(o, do) + dq = q.new_empty(B, T, HQ, K, dtype=q.dtype) + dk = q.new_empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float) + dv = q.new_empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float) + dg = q.new_empty(g.shape, dtype=torch.float) + # NOTE: the original `dg` can be destroyed during autotuning + # this is [a known triton issue](https://github.com/triton-lang/triton/issues/5082), which will be fixed in 3.3 (?) + # so we need to make a copy of `dg` + dg2 = q.new_empty(g.shape, dtype=torch.float) + grid = (NV, NT, B * HQ) + parallel_forgetting_attn_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + g=g, + lse=lse, + delta=delta, + do=do, + dq=dq, + dg=dg, + offsets=offsets, + indices=indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV + ) + parallel_forgetting_attn_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + g=g, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + dg=dg2, + offsets=offsets, + indices=indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV + ) + dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') + dv = reduce(dv, 'b t (h g) v -> b t h v', g=G, reduction='sum') + dg = dg.add_(dg2) + return dq, dk, dv, dg + + +@torch.compile +class ParallelForgettingAttentionFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, g, scale, offsets): + ctx.dtype = q.dtype + if check_shared_mem('hopper'): + chunk_size = min(128, max(16, triton.next_power_of_2(q.shape[1]))) + else: + chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) + # 2-d indices denoting the offsets of chunks in each sequence + # for example, if the passed `offsets` is [0, 100, 356] and `chunk_size` is 64, + # then there are 2 and 4 chunks in the 1st and 2nd sequences respectively, and `indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + indices = prepare_chunk_indices(offsets, chunk_size) if offsets is not None else None + + g = chunk_local_cumsum(g, chunk_size, offsets=offsets, indices=indices, head_first=False) + o, lse = parallel_forgetting_attn_fwd( + q=q, + k=k, + v=v, + g=g, + scale=scale, + chunk_size=chunk_size, + offsets=offsets, + indices=indices + ) + ctx.save_for_backward(q, k, v, g, o, lse) + ctx.chunk_size = chunk_size + ctx.offsets = offsets + ctx.indices = indices + ctx.scale = scale + return o.to(q.dtype) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do): + q, k, v, g, o, lse = ctx.saved_tensors + dq, dk, dv, dg = parallel_forgetting_attn_bwd( + q=q, + k=k, + v=v, + g=g, + o=o, + lse=lse, + do=do, + scale=ctx.scale, + chunk_size=ctx.chunk_size, + offsets=ctx.offsets, + indices=ctx.indices + ) + dg = chunk_global_cumsum(dg, reverse=True, head_first=False, offsets=ctx.offsets) + return dq.to(q), dk.to(k), dv.to(v), dg.to(g), None, None, None, None, None, None, None, None + + +def parallel_forgetting_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: Optional[float] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]` if `head_first=False` else `[B, HQ, T, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + GQA will be applied if HQ is divisible by H. + v (torch.Tensor): + values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + g (torch.Tensor): + Forget gates (in **log space**) of shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. + scale (Optional[int]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`. + """ + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + if g is not None: + g = g.float() + if head_first: + q, k, v = map(lambda x: rearrange(x, 'b h t d -> b t h d'), (q, k, v)) + g = rearrange(g, 'b h t -> b t h') + o = ParallelForgettingAttentionFunction.apply(q, k, v, g, scale, cu_seqlens) + if head_first: + o = rearrange(o, 'b t h d -> b h t d') + return o diff --git a/fla/ops/generalized_delta_rule/README.md b/fla/ops/generalized_delta_rule/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f96c22f44a51ad3e6fdeb824eb2aded660223600 --- /dev/null +++ b/fla/ops/generalized_delta_rule/README.md @@ -0,0 +1,37 @@ +# Generalized Delta Rule + +In delta rule we have the recurrence: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I}-\beta_t \mathbf{k}_t\mathbf{k}_t^T) + \beta_t \mathbf{v}_t\mathbf{k}_t^T +``` + +This repository implements a delta rule variant where $\mathbf{I}$ is not necessarily an identity matrix; $\mathbf{k}_t$ in $\mathbf{I} - \beta_t \mathbf{k}_t\mathbf{k}_t^T$ might be different from input $\mathbf{k}_t$ in $\mathbf{v}_t\mathbf{k}_t^T$. + +## IPLR (Identity Plus Low Rank) + +The first variant is IPLR, where we have: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I}+\mathbf{a}_t\mathbf{b}_t^T) + \mathbf{v}_t\mathbf{k}_t^T +``` + +When $\mathbf{a}_t = -\beta_t \mathbf{k}_t$, $\mathbf{b}_t = \mathbf{k}_t$, $\mathbf{v}_t= \beta_t \mathbf{v}_t$, we recover the original delta rule. Since here the transition matrix is identity-plus-low-rank, we refer to this variant as IPLR. + +### Numerical Stability + +$\mathbf{a}_t$ and $\mathbf{b}_t$ must be in opposite directions, that is, $\mathbf{b}_t = \lambda_t \mathbf{a}_t$ where $\lambda_t < 0$. For an understanding of why this is necessary, you can derive the eigenvalues of the transition matrix. + +## DPLR (Diagonal Plus Low Rank) + +The second variant is DPLR, where we have: + +```math +\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{D}_t+\mathbf{a}_t\mathbf{b}_t^T) + \mathbf{v}_t\mathbf{k}_t^T +``` + +Here, $\mathbf{I}$ is replaced by a diagonal matrix $\mathbf{D}_t$. This transition matrix structure has been utilized in RWKV7. + +## Efficient Chunkwise Implementation + +For detailed information about efficient chunkwise implementation, please refer to our [technical note](https://drive.google.com/file/d/1rJbO3dU4fe7OKG3w7Yg058z_BNIuavNF/view?usp=sharing). diff --git a/fla/ops/generalized_delta_rule/__init__.py b/fla/ops/generalized_delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f4b4155a215ca8c44ea45d6b151b1e584872ed6c --- /dev/null +++ b/fla/ops/generalized_delta_rule/__init__.py @@ -0,0 +1,9 @@ +from .dplr import chunk_dplr_delta_rule, fused_recurrent_dplr_delta_rule +from .iplr import chunk_iplr_delta_rule, fused_recurrent_iplr_delta_rule + +__all__ = [ + 'chunk_dplr_delta_rule', + 'fused_recurrent_dplr_delta_rule', + 'chunk_iplr_delta_rule', + 'fused_recurrent_iplr_delta_rule' +] diff --git a/fla/ops/gsa/__init__.py b/fla/ops/gsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ed8a88014ddfc3143e67d3a48c38a54b75d7f3d6 --- /dev/null +++ b/fla/ops/gsa/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- + +from .chunk import chunk_gsa +from .fused_recurrent import fused_recurrent_gsa + +__all__ = [ + 'chunk_gsa', + 'fused_recurrent_gsa' +] diff --git a/fla/ops/gsa/chunk.py b/fla/ops/gsa/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..950b76cd9c0631fd785bbc85c39c8786b55dfcee --- /dev/null +++ b/fla/ops/gsa/chunk.py @@ -0,0 +1,1264 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from einops import reduce + +from fla.ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from fla.ops.gla.chunk import chunk_gla_bwd, chunk_gla_fwd +from fla.ops.utils import chunk_local_cumsum, softmax_bwd, softmax_fwd +from fla.ops.utils.op import exp, safe_exp +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BT'] +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_k_kernel_inter( + q, + k, + h, + g, + o, + A, + offsets, + indices, + scale, + T, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_bg = i_bh // NG + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if USE_OFFSETS: + i_tg = i_t + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + if HEAD_FIRST: + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bg * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h + (i_bg * NT + i_t) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + else: + p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BK, BV] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_o += tl.dot(b_q, b_h) + # [BT, BT] + b_A += tl.dot(b_q, b_k) + if HEAD_FIRST: + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + else: + p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o = b_o * exp(b_g) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + # [BT, BT] + b_A = tl.where(m_s, b_A, 0.) + if i_v == 0: + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_k_kernel_intra( + v, + g, + o, + A, + offsets, + indices, + T, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_bg = i_bh // NG + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i = i_c // NC, i_c % NC + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC > T: + return + + if HEAD_FIRST: + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + min(i_t * BT + i_i * BC, T) * V + o_v, BV), BV) + else: + p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = g + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_o = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(0, i_i): + if HEAD_FIRST: + p_A = tl.make_block_ptr(A + i_bh * T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + p_v = tl.make_block_ptr(v + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_gv = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + else: + p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[None, :] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_o += tl.dot(b_A, b_vg) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_o *= exp(b_g - b_gn[None, :]) + + o_i = tl.arange(0, BC) + if HEAD_FIRST: + o_A = i_bh * T*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC + else: + o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + if HEAD_FIRST: + p_v = tl.max_contiguous(tl.multiple_of(v + i_bg * T*V + (i_t * BT + i_i * BC + j) * V + o_v, BV), BV) + p_gv = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (i_t * BT + i_i * BC + j) * V + o_v, BV), BV) + else: + p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + # [BC,] + b_A = tl.load(A + o_A + j, mask=m_A, other=0) + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC, BV] + b_vg = b_v[None, :] * exp(b_g - b_gv[None, :]) + # avoid 0 * inf = inf + b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) + if HEAD_FIRST: + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + else: + p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + b_o += tl.load(p_o, boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8] + ], + key=["BT"] +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_dA( + v, + g, + do, + dA, + indices, + offsets, + scale, + T, + B: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_bg = i_bh // NG + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC > T: + return + + if HEAD_FIRST: + p_dA = tl.make_block_ptr(dA+(i_v*B*H+i_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + else: + p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) + + # [BC, BC] + b_dA = tl.zeros([BC, BC], dtype=tl.float32) + if i_i > i_j: + if HEAD_FIRST: + p_v = tl.make_block_ptr(v + i_bg * T*V, (V, T), (1, V), (i_v * BV, i_t * BT + i_j * BC), (BV, BC), (0, 1)) + p_gv = tl.make_block_ptr(g + i_bg * T*V, (V, T), (1, V), (i_v * BV, i_t * BT + i_j * BC), (BV, BC), (0, 1)) + p_gn = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (i_t * BT + i_i * BC) * V + o_v, BV), BV) + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + else: + p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + p_gn = g + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0.) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g - b_gn[None, :]) * scale).to(b_do.dtype) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_vg = (b_v * exp(b_gn[:, None] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_dA = tl.dot(b_do, b_vg) + elif i_i == i_j: + if HEAD_FIRST: + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_v = tl.max_contiguous(tl.multiple_of(v + i_bg * T*V + (i_t * BT + i_j * BC) * V + o_v, BV), BV) + p_gv = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (i_t * BT + i_j * BC) * V + o_v, BV), BV) + else: + p_g = tl.make_block_ptr(g + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + m_v = o_v < V + + o_i = tl.arange(0, BC) + # [BC, BC] + m_dA = o_i[:, None] >= o_i[None, :] + for j in range(0, min(BC, T - i_t * BT - i_j * BC)): + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC,] + b_dAj = tl.sum(b_do * b_v[None, :] * exp(b_g - b_gv[None, :]), 1) + b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) + + p_v += (1 if HEAD_FIRST else H) * V + p_gv += (1 if HEAD_FIRST else H) * V + b_dA = tl.where(m_dA, b_dA, 0.) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['BT'] +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_dqkvg( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + offsets, + indices, + scale, + T, + B: tl.constexpr, + HQ: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_bg = i_bh // NG + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + if USE_OFFSETS: + i_tg = i_t + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + all = T + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_i = tl.arange(0, BT) + o_t = min(i_t * BT + BT, T) + m_s = o_i[:, None] >= o_i[None, :] + + if HEAD_FIRST: + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bg * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + (i_k*B*H+i_bh) * T*BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + else: + p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + if HEAD_FIRST: + p_v = tl.make_block_ptr(v + i_bg * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gn = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (o_t - 1) * V + o_v, BV), BV) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_k*B*H+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dgv = tl.make_block_ptr(dgv + (i_k*B*H+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + i_bg * NT*K*V + i_t * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + else: + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + m_v = o_v < V + + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_gv = exp(b_gn[None, :] - b_g) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (b_do * exp(b_g) * scale).to(b_do.dtype) + # [BK, BV] + b_dh = tl.load(p_dh, boundary_check=(0, 1)) + # [BV] + b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp(b_gn) + + b_dh = b_dh.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) + b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) + # [BT, BV] + b_dv = tl.dot(b_k, b_dh) * b_gv + # [BV] + b_dg += tl.sum(b_dv * b_v, 0) + + if i_k == 0: + b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] + else: + b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] + + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + if HEAD_FIRST: + p_dA = tl.make_block_ptr(dA + i_bh * T*BT, (T, BT, ), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + else: + p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT, BT] + b_dA = tl.load(p_dA, boundary_check=(0, 1)) + # [BT, BK] + b_dq += tl.dot(b_dA, b_k) + b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_k_kernel_intra_dvg( + v, + g, + o, + A, + do, + dv, + dg, + offsets, + indices, + T, + HQ: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + NG: tl.constexpr, + USE_OFFSETS: tl.constexpr, + HEAD_FIRST: tl.constexpr +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_bg = i_bh // NG + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // NG + i_t, i_i = i_c // NC, i_c % NC + if USE_OFFSETS: + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + if i_t * BT + i_i * BC > T: + return + + if HEAD_FIRST: + p_gv = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (min(i_t * BT + i_i * BC + BC, T) - 1) * V + o_v, BV), BV) + else: + p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + if HEAD_FIRST: + p_g = tl.make_block_ptr(g + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_A = tl.make_block_ptr(A + i_bh * T*BT, (BT, T), (1, BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + else: + p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (BT, T), (1, HQ*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) + # [BC, BV] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) * safe_exp(b_g - b_gn[None, :]) + # [BC, BC] + b_A = tl.load(p_A, boundary_check=(0, 1)) + # [BC, BV] + b_dv += tl.dot(b_A, b_do.to(b_A.dtype)) + b_dv *= exp(b_gn[None, :] - b_gv) + + o_i = tl.arange(0, BC) + o_c = i_i * BC + tl.arange(0, BC) + + if HEAD_FIRST: + p_g = tl.max_contiguous(tl.multiple_of(g + i_bg * T*V + (i_t * BT + i_i * BC) * V + o_v, BV), BV) + p_A = tl.max_contiguous(tl.multiple_of(A + i_bh * T*BT + (i_t * BT + i_i * BC) * BT + o_c, BC), BC) + p_do = tl.max_contiguous(tl.multiple_of(do + i_bh * T*V + (i_t * BT + i_i * BC) * V + o_v, BV), BV) + else: + p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v + p_A = A + (bos + i_t*BT + i_i*BC) * HQ*BT + i_hq * BT + o_c + p_do = do + (bos + i_t*BT + i_i*BC) * HQ*V + i_hq * V + o_v + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_A = tl.load(p_A) + # [BV,] + b_g = tl.load(p_g, mask=m_v, other=0) + b_do = tl.load(p_do, mask=m_v, other=0) + # [BC, BV] + m_i = o_i[:, None] <= j + b_dv += tl.where(m_i, exp(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) + + p_g += (1 if HEAD_FIRST else H) * V + p_A += (1 if HEAD_FIRST else HQ) * BT + p_do += (1 if HEAD_FIRST else HQ) * V + if HEAD_FIRST: + p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_v = tl.make_block_ptr(v + i_bg * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + else: + p_o = tl.make_block_ptr(o + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + + b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) + b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) + b_dg = b_o * b_do - b_v * b_dv + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gsa_fwd_v( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float = 1., + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + _, A, h, ht, o = chunk_gla_fwd( + q=q, + k=k, + v=v, + g=None, + g_cumsum=g, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + return A, h, ht, o + + +def chunk_gsa_fwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h0: Optional[torch.Tensor] = None, + output_final_state: bool = False, + scale: float = 1., + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if head_first: + B, H, T, K, V = *k.shape, v.shape[-1] + else: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + HQ = q.shape[1] if head_first else q.shape[2] + NT = triton.cdiv(T, BT) if offsets is None else len(indices) + NC = triton.cdiv(BT, BC) + NG = HQ // H + + h, ht = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=output_final_state, + offsets=offsets, + head_first=head_first, + chunk_size=BT, + states_in_fp32=False + ) + o = v.new_empty(B, *((HQ, T) if head_first else (T, HQ)), V) + A = q.new_empty(B, *((HQ, T) if head_first else (T, HQ)), BT) + def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) + chunk_gsa_fwd_k_kernel_inter[grid]( + q, + k, + h, + g, + o, + A, + offsets=offsets, + indices=indices, + scale=scale, + T=T, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + NG=NG, + HEAD_FIRST=head_first + ) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + chunk_gsa_fwd_k_kernel_intra[grid]( + v, + g, + o, + A, + offsets=offsets, + indices=indices, + T=T, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + HEAD_FIRST=head_first, + num_warps=4, + num_stages=2 + ) + return A, h, ht, o + + +def chunk_gsa_bwd_v( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h0: torch.Tensor, + h: torch.Tensor, + A: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dg: torch.Tensor, + scale: float = 1., + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +): + dq, dk, dv, dg, dh0 = chunk_gla_bwd( + q=q, + k=k, + v=v, + g=None, + g_cumsum=g, + scale=scale, + initial_state=h0, + h=h, + A=A, + do=do, + dht=dht, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + return dq, dk, dv, dg, dh0 + + +def chunk_gsa_bwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + h0: torch.Tensor, + o: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dg: torch.Tensor, + scale: float = 1., + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +): + if head_first: + B, H, T, K, V = *k.shape, v.shape[-1] + else: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = min(chunk_size, max(16, triton.next_power_of_2(T))) + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + HQ = q.shape[1] if head_first else q.shape[2] + NT = triton.cdiv(T, BT) if offsets is None else len(indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + NG = HQ // H + + if h is None: + h, _ = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=False, + offsets=offsets, + head_first=head_first, + chunk_size=BT, + states_in_fp32=False + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=None, + gk=None, + gv=g, + do=do, + h0=h0, + dht=dht, + scale=scale, + offsets=offsets, + head_first=head_first, + chunk_size=BT, + states_in_fp32=True + ) + dA = q.new_empty(NV, B, *((HQ, T) if head_first else (T, HQ)), BT) + grid = (NV, NT * NC * NC, B * HQ) + chunk_gsa_bwd_k_kernel_dA[grid]( + v, + g, + do, + dA, + offsets=offsets, + indices=indices, + scale=scale, + T=T, + B=B, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + HEAD_FIRST=head_first + ) + dA = dA.sum(0, dtype=dA.dtype) + + A = do.new_empty(NK, B, *((HQ, T) if head_first else (T, HQ)), BT) + dq = torch.empty_like(q) + dk = k.new_empty(B, *((HQ, T) if head_first else (T, HQ)), K) + dv = v.new_empty(NK, B, *((HQ, T) if head_first else (T, HQ)), V) + dgv = g.new_empty(NK, B, *((HQ, T) if head_first else (T, HQ)), V, dtype=torch.float) + grid = (NK, NT, B * HQ) + chunk_gsa_bwd_k_kernel_dqkvg[grid]( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + offsets=offsets, + indices=indices, + scale=scale, + T=T, + B=B, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + NG=NG, + HEAD_FIRST=head_first + ) + A = A.sum(0, dtype=A.dtype) + dv = dv.sum(0, dtype=dv.dtype) + dgv = dgv.sum(0, dtype=dgv.dtype) + + def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + chunk_gsa_bwd_k_kernel_intra_dvg[grid]( + v, + g, + o, + A, + do, + dv, + dg, + offsets=offsets, + indices=indices, + T=T, + HQ=HQ, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + NG=NG, + HEAD_FIRST=head_first, + num_warps=4, + num_stages=2 + ) + dg = dgv.add_(chunk_local_cumsum(dg, chunk_size=BT, reverse=True, offsets=offsets, indices=indices, head_first=head_first)) + + return dq, dk, dv, dg, dh0 + + +def chunk_gsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + output_final_state: bool = False, + scale: float = 1., + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + Ak, hk, hkt, ok = chunk_gsa_fwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + output_final_state=output_final_state, + scale=scale, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + + # p is kept in fp32 for safe softmax backward + p = softmax_fwd(ok, dtype=torch.float) + + qv = p.to(q.dtype) + Av, hv, hvt, ov = chunk_gsa_fwd_v( + q=qv, + k=s, + v=v, + g=g, + scale=1., + initial_state=hv0, + output_final_state=output_final_state, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + return Ak, hk, hkt, ok, p, Av, hv, hvt, ov + + +def chunk_gsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + ok: torch.Tensor, + p: torch.Tensor, + A: Tuple[torch.Tensor, torch.Tensor], + h: Tuple[torch.Tensor, torch.Tensor], + initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]], + scale: float, + do: torch.Tensor, + dht: Tuple[torch.Tensor, torch.Tensor], + offsets: Optional[torch.LongTensor] = None, + indices: Optional[torch.LongTensor] = None, + head_first: bool = True, + chunk_size: int = 64 +): + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + + _, Av = A + hk, hv = h + dhkt, dhvt = dht + + qv = p.to(q.dtype) + dqv, dsv, dv, dg, dhv0 = chunk_gsa_bwd_v( + q=qv, + k=s, + v=v, + g=g, + h0=hv0, + h=hv, + A=Av, + do=do, + dht=dhvt, + dg=None, + scale=1., + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + + # softmax gradient, equivalent to: + # dok = qv * (dqv - (qv * dqv).sum(-1, True)) + dok = softmax_bwd(p, dqv, dtype=ok.dtype) + + dq, dk, dsk, dg, dhk0 = chunk_gsa_bwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + h=hk, + o=ok, + do=dok, + dht=dhkt, + dg=dg, + scale=scale, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + + ds = dsv.add_(dsk) + if q.shape[1] != k.shape[1]: + dk, dv, ds, dg = map(lambda x: reduce(x, 'b (h g) ... -> b h ...', 'sum', h=k.shape[1]), (dk, dv, ds, dg)) + dg = dg.to(s.dtype) + return dq, dk, dv, ds, dg, dhk0, dhv0 + + +class ChunkGSAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + scale: float, + hk0: Optional[torch.Tensor], + hv0: Optional[torch.Tensor], + output_final_state: bool, + checkpoint_level: int, + offsets: Optional[torch.LongTensor], + head_first: bool = True + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + T = q.shape[2] if head_first else q.shape[1] + chunk_size = min(64, max(16, triton.next_power_of_2(T))) + + # 2-d indices denoting the offsets of chunks in each sequence + # for example, if the passed `offsets` is [0, 100, 356] and `chunk_size` is 64, + # then there are 2 and 4 chunks in the 1st and 2nd sequences respectively, and `indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + indices = None + if offsets is not None: + indices = torch.cat([torch.arange(n) for n in triton.cdiv(offsets[1:] - offsets[:-1], chunk_size).tolist()]) + indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(offsets) + g_org, g = g, chunk_local_cumsum(g, chunk_size, offsets=offsets, indices=indices, head_first=head_first) + Ak, hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + + if checkpoint_level >= 1: + del g + g = g_org + if checkpoint_level > 1: + del hk + del hv + hk, hv = None, None + else: + hk0, hv0 = None, None + + ctx.save_for_backward(q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv) + ctx.checkpoint_level = checkpoint_level + ctx.scale = scale + ctx.offsets = offsets + ctx.indices = indices + ctx.head_first = head_first + ctx.chunk_size = chunk_size + return ov, hkt, hvt + + @staticmethod + @input_guard + def backward(ctx, dov, dhkt=None, dhvt=None): + q, k, v, s, g, ok, p, Av, hk0, hv0, hk, hv = ctx.saved_tensors + scale = ctx.scale + offsets = ctx.offsets + indices = ctx.indices + head_first = ctx.head_first + chunk_size = ctx.chunk_size + + if ctx.checkpoint_level >= 1: + g = chunk_local_cumsum(g, chunk_size, offsets=offsets, indices=indices, head_first=head_first) + dq, dk, dv, ds, dg, dhk0, dhv0 = chunk_gsa_bwd( + q=q, + k=k, + v=v, + s=s, + g=g, + ok=ok, + p=p, + A=(None, Av), + h=(hk, hv), + initial_state=(hk0, hv0), + scale=scale, + do=dov, + dht=(dhkt, dhvt), + offsets=offsets, + indices=indices, + head_first=head_first, + chunk_size=chunk_size + ) + return dq, dk, dv, ds, dg, None, dhk0, dhv0, None, None, None, None + + +@torch.compiler.disable +def chunk_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: Optional[torch.Tensor] = None, + scale: Optional[int] = None, + initial_state: Optional[Tuple[torch.Tensor]] = None, + output_final_state: Optional[bool] = False, + checkpoint_level: Optional[int] = 2, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: Optional[bool] = True +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, HQ, T, K]` if `head_first=True` else `[B, T, HQ, K]`. + k (torch.Tensor): + keys of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + GQA is performed if `H` is not equal to `HQ`. + v (torch.Tensor): + values of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + s (torch.Tensor): + slot representations of shape `[B, H, T, M]` if `head_first=True` else `[B, T, H, M]`. + g (torch.Tensor): + Forget gates of shape `[B, H, T, M]` applied to keys. + If not provided, this function is equivalent to vanilla ABC. + scale (Optional[int]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[Tuple[torch.Tensor]]): + Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state tuple, having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. + Default: `False`. + checkpoint_level (Optional[int]): + Checkpointing level; higher values will save more memories and do more recomputations during backward. + Default: `2`: + - Level `0`: no memory saved, no recomputation. + - Level `1`: recompute the fp32 cumulative values during backward. + - Level `2`: recompute the fp32 cumulative values and forward hidden states during backward. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format, which is not supported for variable-length inputs. + Default: `True`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + final_state (Tuple[torch.Tensor]): + Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` if `output_final_state=True`. + `None` otherwise. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gsa import fused_recurrent_gsa + # inputs with equal lengths + >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> s = torch.randn(B, T, H, M, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) + >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) + >>> o, (hk, hv) = chunk_gsa(q, k, v, s, g, + initial_state=h0, + output_final_state=True, + head_first=False) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, (hk_var, hv_var) = chunk_gsa(q, k, v, s, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + head_first=False) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert hk.allclose(hk_var) + >>> assert hv.allclose(hv_var) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError(f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.") + if head_first: + raise RuntimeError("Sequences with variable lengths are not supported for head-first mode") + if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: + raise ValueError(f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.") + assert checkpoint_level in [0, 1, 2] + if g is None: + # TODO: this 3 steps took huge amount of time, ought to be optimized + z = s.float().logcumsumexp(2) + g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z + s = torch.exp(s - z).to(k.dtype) + if scale is None: + scale = q.shape[-1] ** -0.5 + + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + o, *final_state = ChunkGSAFunction.apply( + q, + k, + v, + s, + g, + scale, + hk0, + hv0, + output_final_state, + checkpoint_level, + cu_seqlens, + head_first + ) + return o, final_state diff --git a/fla/ops/gsa/fused_recurrent.py b/fla/ops/gsa/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..a4eacfda688efee1fb8ebe9543234f53bef85ade --- /dev/null +++ b/fla/ops/gsa/fused_recurrent.py @@ -0,0 +1,564 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2024, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from fla.ops.common.fused_recurrent import fused_recurrent_bwd_kernel, fused_recurrent_fwd_kernel +from fla.ops.utils import chunk_global_cumsum +from fla.ops.utils.op import exp +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.jit +def fused_recurrent_gsa_inference_kernel( + q, + k, + v, + s, + g, + o, + hk0, + hv0, + hkt, + hvt, + scale, + K: tl.constexpr, + V: tl.constexpr, + M: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NG: tl.constexpr +): + i_bh = tl.program_id(0) + i_bg = i_bh // NG + + b_s = tl.load(s + i_bg * M + tl.arange(0, M)).to(tl.float32) + b_g = tl.load(g + i_bg * M + tl.arange(0, M)).to(tl.float32) + b_g = exp(b_g) + + b_ok = tl.zeros([M], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + + p_hk0 = hk0 + i_bg * K * M + (o_k[None, :]) * M + tl.arange(0, M)[:, None] + # [BK,] + mask_k = o_k < K + # [M, BK] + mask_hk = (tl.arange(0, M) < M)[:, None] & mask_k[None, :] + # [M, BK] + b_hk = tl.load(p_hk0, mask=mask_hk, other=0.).to(tl.float32) + # [BK,] + b_q = tl.load(q + i_bh * K + o_k, mask=mask_k, other=0.).to(tl.float32) * scale + b_k = tl.load(k + i_bg * K + o_k, mask=mask_k, other=0.).to(tl.float32) + b_hk = b_hk * b_g[:, None] + b_k[None, :] * b_s[:, None] + b_ok += tl.sum(b_hk * b_q[None, :], axis=1) + + if i_bh % NG == 0: + p_hkt = hkt + i_bg * K * M + o_k[None, :] * M + tl.arange(0, M)[:, None] + tl.store(p_hkt, b_hk.to(p_hkt.dtype.element_ty), mask=mask_hk) + + b_qv = tl.softmax(b_ok) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + + p_hv0 = hv0 + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] + # [BV,] + mask_v = o_v < V + # [BV, M] + mask_hv = mask_v[:, None] & (tl.arange(0, M) < M)[None, :] + # [BV, M] + b_hv = tl.load(p_hv0, mask=mask_hv, other=0).to(tl.float32) + # [BV,] + b_v = tl.load(v + i_bg * V + o_v, mask=mask_v, other=0).to(tl.float32) + b_hv = b_hv * b_g[None, :] + b_s[None, :] * b_v[:, None] + b_ov = tl.sum(b_hv * b_qv[None, :], axis=1) + + tl.store(o + i_bh * V + o_v, b_ov.to(o.dtype.element_ty), mask=mask_v) + + if i_bh % NG == 0: + p_hvt = hvt + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] + tl.store(p_hvt, b_hv.to(p_hvt.dtype.element_ty), mask=mask_hv) + + +def fused_recurrent_gsa_inference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + output_final_state: bool = False, + scale: float = 1., + head_first: bool = True +) -> torch.Tensor: + if head_first: + B, H, T, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + else: + B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + HQ = q.shape[1] if head_first else q.shape[2] + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NG = HQ // H + + if initial_state != (None, None) and initial_state is not None: + hk0, hv0 = initial_state + else: + hk0, hv0 = q.new_zeros(B, H, K, M, dtype=torch.float), q.new_zeros(B, H, M, V, dtype=torch.float) + + hkt, hvt = None, None + if output_final_state: + if NG == 1: + hkt, hvt = hk0, hv0 + else: + hkt, hvt = q.new_empty(B, H, K, M, dtype=torch.float), q.new_empty(B, H, M, V, dtype=torch.float) + + o = v.new_empty(B, HQ, T, V) if head_first else v.new_empty(B, T, HQ, V) + grid = (B * HQ,) + fused_recurrent_gsa_inference_kernel[grid]( + q, + k, + v, + s, + g, + o, + hk0, + hv0, + hkt, + hvt, + scale=scale, + K=K, + V=V, + M=M, + BK=BK, + BV=BV, + NG=NG + ) + return o, (hkt, hvt) + + +def fused_recurrent_gsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + output_final_state: bool = False, + scale: float = 1., + reverse: bool = False, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]: + if head_first: + B, H, T, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + else: + B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] + N = B if offsets is None else len(offsets) - 1 + HQ = q.shape[1] if head_first else q.shape[2] + if HQ != H: + raise ValueError("GQA not supported yet.") + + BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(M, 64) + NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) + + hk0, hv0 = None, None + if initial_state != (None, None) and initial_state is not None: + hk0, hv0 = initial_state + hkt, hvt = None, None + if output_final_state: + hkt, hvt = q.new_empty(N, H, K, M, dtype=torch.float), q.new_empty(N, H, M, V, dtype=torch.float) + + ok = q.new_empty(NK, *s.shape, dtype=torch.float) + gk, gv = None, g + grid = (NM, NK, N * H) + fused_recurrent_fwd_kernel[grid]( + q=q, + k=k, + v=s, + g=None, + gk=gk, + gv=gv, + o=ok, + h0=hk0, + ht=hkt, + offsets=offsets, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=M, + BK=BK, + BV=BM, + USE_G=False, + USE_GK=False, + USE_GV=True, + REVERSE=reverse, + HEAD_FIRST=head_first + ) + ok = ok.sum(0) + + qv = ok.softmax(-1, dtype=torch.float) + ov = q.new_empty(NM, *v.shape, dtype=torch.float) + gk, gv = g, None + grid = (NV, NM, N * H) + fused_recurrent_fwd_kernel[grid]( + q=qv, + k=s, + v=v, + g=None, + gk=gk, + gv=gv, + o=ov, + h0=hv0, + ht=hvt, + offsets=offsets, + scale=1., + B=B, + T=T, + H=H, + K=M, + V=V, + BK=BM, + BV=BV, + USE_G=False, + USE_GK=True, + USE_GV=False, + REVERSE=reverse, + HEAD_FIRST=head_first + ) + ov = ov.sum(0) + return ok, hkt, qv, ov, hvt + + +def fused_recurrent_gsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + qv: torch.Tensor, + hk0: Optional[torch.Tensor] = None, + hv0: Optional[torch.Tensor] = None, + ok: Optional[torch.Tensor] = None, + do: Optional[torch.Tensor] = None, + dhkt: Optional[torch.Tensor] = None, + dhvt: Optional[torch.Tensor] = None, + scale: float = 1., + reverse: bool = False, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor]: + if head_first: + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + else: + B, T, H, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + N = B if offsets is None else len(offsets) - 1 + + BK, BV, BM = min(K, 64), min(V, 64), min(M, 64) + NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) + + if head_first: + dqv = q.new_empty(NV, B, H, T, M, dtype=torch.float) + dsv = q.new_empty(NV, B, H, T, M, dtype=torch.float) + dv = q.new_empty(NM, B, H, T, V, dtype=torch.float) + else: + dqv = q.new_empty(NV, B, T, H, M, dtype=torch.float) + dsv = q.new_empty(NV, B, T, H, M, dtype=torch.float) + dv = q.new_empty(NM, B, T, H, V, dtype=torch.float) + dhk0 = torch.empty_like(hk0)if hk0 is not None else None + dhv0 = torch.empty_like(hv0)if hv0 is not None else None + + gk, gv = g, None + grid = (NV, NM, N * H) + fused_recurrent_bwd_kernel[grid]( + q=qv, + k=s, + v=v, + g=None, + gk=gk, + gv=gv, + h0=hv0, + do=do, + dq=dqv, + dk=dsv, + dv=dv, + dht=dhvt, + dh0=dhv0, + offsets=offsets, + scale=1., + B=B, + T=T, + H=H, + K=M, + V=V, + BK=BM, + BV=BV, + USE_G=False, + USE_GK=True, + USE_GV=False, + REVERSE=reverse, + HEAD_FIRST=head_first + ) + dqv = dqv.sum(0) + dsv = dsv.sum(0) + dv = dv.sum(0) + dgk = chunk_global_cumsum(dqv * qv.float() - dsv * s.float(), + reverse=not reverse, + offsets=offsets, + head_first=head_first) + + dok = qv * (dqv - (qv * dqv).sum(-1, True)) + if head_first: + dq = q.new_empty(NM, B, H, T, K, dtype=torch.float) + dk = q.new_empty(NM, B, H, T, K, dtype=torch.float) + dsk = q.new_empty(NK, B, H, T, M, dtype=torch.float) + else: + dq = q.new_empty(NM, B, T, H, K, dtype=torch.float) + dk = q.new_empty(NM, B, T, H, K, dtype=torch.float) + dsk = q.new_empty(NK, B, T, H, M, dtype=torch.float) + gk, gv = None, g + grid = (NM, NK, N * H) + fused_recurrent_bwd_kernel[grid]( + q=q, + k=k, + v=s, + g=None, + gk=gk, + gv=gv, + h0=hk0, + do=dok, + dq=dq, + dk=dk, + dv=dsk, + dht=dhkt, + dh0=dhk0, + offsets=offsets, + scale=scale, + B=B, + T=T, + H=H, + K=K, + V=M, + BK=BK, + BV=BM, + USE_G=False, + USE_GK=False, + USE_GV=True, + REVERSE=reverse, + HEAD_FIRST=head_first + ) + dq = dq.sum(0) + dk = dk.sum(0) + dsk = dsk.sum(0) + + dgv = chunk_global_cumsum(dok.float() * ok.float() - dsk * s.float(), + reverse=not reverse, + offsets=offsets, + head_first=head_first) + + ds = dsk.add_(dsv) + dg = dgk.add_(dgv) + + return dq, dk, dv, ds, dg, dhk0, dhv0 + + +class FusedRecurrentGSAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + scale: Optional[float] = None, + hk0: Optional[torch.Tensor] = None, + hv0: Optional[torch.Tensor] = None, + output_final_state: bool = False, + reverse: bool = False, + offsets: Optional[torch.LongTensor] = None, + head_first: bool = True + ) -> Tuple[torch.Tensor, Tuple[torch.Tensor]]: + T = q.shape[2] if head_first else q.shape[1] + if T == 1 and not q.requires_grad: + o, (hkt, hvt) = fused_recurrent_gsa_inference( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + head_first=head_first + ) + return o, hkt, hvt + ok, hkt, qv, ov, hvt = fused_recurrent_gsa_fwd( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=output_final_state, + scale=scale, + reverse=reverse, + offsets=offsets, + head_first=head_first + ) + ctx.save_for_backward(q, k, v, s, g, qv, hk0, hv0, ok) + ctx.scale = scale + ctx.reverse = reverse + ctx.offsets = offsets + ctx.head_first = head_first + return ov.to(q.dtype), hkt, hvt + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dhkt=None, dhvt=None): + q, k, v, s, g, qv, hk0, hv0, ok = ctx.saved_tensors + scale = ctx.scale + reverse = ctx.reverse + offsets = ctx.offsets + head_first = ctx.head_first + + # not supported yet. + if dhkt is not None or dhvt is not None: + if g is not None: + assert g.requires_grad is False, "Cannot load final state gradient and use gates at the same time" + dq, dk, dv, ds, dg, dhk0, dhv0 = fused_recurrent_gsa_bwd( + q=q, + k=k, + v=v, + s=s, + g=g, + qv=qv, + hk0=hk0, + hv0=hv0, + ok=ok, + do=do, + dhkt=dhkt, + dhvt=dhvt, + scale=scale, + reverse=reverse, + offsets=offsets, + head_first=head_first + ) + return dq.to(q), dk.to(k), dv.to(v), ds.to(s), dg.to(g), None, dhk0, dhv0, None, None, None, None + + +def fused_recurrent_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: Optional[torch.Tensor] = None, + scale: Optional[int] = None, + initial_state: Optional[Tuple[torch.Tensor]] = None, + output_final_state: Optional[bool] = False, + reverse: Optional[bool] = False, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + s (torch.Tensor): + slot representations of shape `[B, H, T, M]` if `head_first=True` else `[B, T, H, M]`. + g (torch.Tensor): + Forget gates of shape `[B, H, T, M]` applied to keys. + scale (Optional[int]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[Tuple[torch.Tensor]]): + Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]` and `[N, H, M, V]`. + Default: `False`. + reverse (Optional[bool]): + If `True`, process the state passing in reverse order. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (Optional[bool]): + Whether the inputs are in the head-first format, which is not supported for variable-length inputs. + Default: `True`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + final_state (Tuple[torch.Tensor]): + Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gsa import fused_recurrent_gsa + # inputs with equal lengths + >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = torch.randn(B, T, H, K, device='cuda') + >>> v = torch.randn(B, T, H, V, device='cuda') + >>> s = torch.randn(B, T, H, M, device='cuda') + >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) + >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) + >>> o, (hk, hv) = fused_recurrent_gsa(q, k, v, s, g, + initial_state=h0, + output_final_state=True, + head_first=False) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, (hk_var, hv_var) = fused_recurrent_gsa(q, k, v, s, g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + head_first=False) + >>> assert o.allclose(o_var.view(o.shape)) + >>> assert hk.allclose(hk_var) + >>> assert hv.allclose(hv_var) + """ + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError(f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.") + if head_first: + raise RuntimeError("Sequences with variable lengths are not supported for head-first mode") + if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: + raise ValueError(f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.") + if scale is None: + scale = k.shape[-1] ** -0.5 + if initial_state is None: + initial_state = (None, None) + o, *final_state = FusedRecurrentGSAFunction.apply( + q, + k, + v, + s, + g, + scale, + *initial_state, + output_final_state, + reverse, + cu_seqlens, + head_first + ) + return o, final_state diff --git a/fla/ops/gsa/naive.py b/fla/ops/gsa/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..699b2a4d1e5b4b8415a98c27da142a6130797685 --- /dev/null +++ b/fla/ops/gsa/naive.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- + +from typing import Optional + +import torch +from einops import repeat + + +def naive_recurrent_gsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: Optional[torch.Tensor] = None, + scale: Optional[int] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: Optional[bool] = False +) -> torch.Tensor: + dtype = q.dtype + + NG = q.shape[1]//k.shape[1] + # [batch_size, n_heads, seq_len, n_slots] + if g is None: + z = s.float().logcumsumexp(2) + g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z + s = torch.exp(s - z) + q, k, v, s, g = map(lambda x: x.float(), (q, k, v, s, g)) + k, v, s, g = map(lambda x: repeat(x, 'b h t d -> b (h g) t d', g=NG), (k, v, s, g)) + if initial_state is not None: + initial_state = tuple(map(lambda x: repeat(x, 'b h k v -> b (h g) k v', g=NG), initial_state)) + + B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] + + hk = torch.zeros(B, H, K, M, dtype=torch.float, device=q.device) + ok = torch.zeros_like(s) + + if scale is None: + scale = q.shape[-1] ** -0.5 + + final_state = None + if initial_state is not None: + hk += initial_state[0] + + for i in range(T): + q_i = q[:, :, i] * scale + k_i = k[:, :, i] + v_i = s[:, :, i] + g_i = g[:, :, i].exp() + hk = hk * g_i[..., None, :] + k_i[..., None] * v_i[..., None, :] + ok[:, :, i] = (q_i[..., None] * hk).sum(-2) + + qv = ok.softmax(-1) + hv = torch.zeros(B, H, M, V, dtype=torch.float, device=q.device) + ov = torch.zeros_like(v) + if initial_state is not None: + hv += initial_state[1] + + for i in range(T): + q_i = qv[:, :, i] + k_i = s[:, :, i] + v_i = v[:, :, i] + g_i = g[:, :, i].exp() + hv = hv * g_i[..., :, None] + k_i[..., None] * v_i[..., None, :] + ov[:, :, i] = (q_i[..., None] * hv).sum(-2) + + if output_final_state: + final_state = (hk.view(B, -1, NG, K, M)[:, :, 0], hv.view(B, -1, NG, M, V)[:, :, 0]) + return ov.to(dtype), final_state diff --git a/fla/ops/linear_attn/fused_chunk.py b/fla/ops/linear_attn/fused_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..bfcc1212a534aa3debb5bb2d1cdbbce5f95f06e4 --- /dev/null +++ b/fla/ops/linear_attn/fused_chunk.py @@ -0,0 +1,318 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from packaging import version + +from fla.ops.linear_attn.utils import normalize_output +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.jit +def fused_chunk_linear_attn_fwd_kernel( + q, # query [B, H, T, K] + k, # key [B, H, T, V] + v, # value [B, H, T, V] + o, # output [B, H, T, V] + h0, + ht, + scale, + B, # batch size + H, # H + T, # T + K: tl.constexpr, # K + V: tl.constexpr, # V + BT: tl.constexpr, # BLOCK SIZE along the sequence dimension, a.k.a. chunk size + BK: tl.constexpr, # BLOCK SIZE along the K dimension + BV: tl.constexpr, # BLOCK SIZE along the V dimension + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + CHECK: tl.constexpr +): + # indices + i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + o_i = tl.arange(0, BT) + + # [BT, BT] + m_s = o_i[:, None] >= o_i[None, :] + # [BK, BV] + b_h = tl.zeros([BK, BV], dtype=tl.float32) + + # make block pointers + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BT, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BT), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) + p_o = tl.make_block_ptr(o + (i_bh+i_k*B*H) * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) + + if USE_INITIAL_STATE: + p_h0 = tl.make_block_ptr(h0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + + for i in range(0, tl.cdiv(T, BT)): + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # [BT, BT] + b_s = tl.dot(b_q, b_k, allow_tf32=False) + b_s = tl.where(m_s, b_s, 0) + # [BT, BV] + b_o = tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + if CHECK and i == 0: + b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False) + b_h = b_h + tl.dot(b_k, b_v, allow_tf32=False) + else: + b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False) + b_h = b_h + tl.dot(b_k, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_q = tl.advance(p_q, (BT, 0)) + p_k = tl.advance(p_k, (0, BT)) + p_v = tl.advance(p_v, (BT, 0)) + p_o = tl.advance(p_o, (BT, 0)) + + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit +def fused_chunk_linear_attn_bwd_kernel( + q, # query [B, H, T, K] + k, # key [B, H, T, V] + v, # value [B, H, T, V] + do, # gradient of output [B, H, T, V] + dq, # gradient of query [NV, B, H, T, K] + dk, # gradient of key [NV, B, H, T, K] + dv, # gradient of value [NK, B, H, T, V] + h0, # initial state of the chunk [B, H, K, V] + scale, # K ** -0.5 + B, # B + H, # H + T, # T + K: tl.constexpr, # K + V: tl.constexpr, # V + BT: tl.constexpr, # BLOCK SIZE along the sequence dimension, a.k.a. chunk size + BK: tl.constexpr, # BLOCK SIZE along the K dimension + BV: tl.constexpr, # BLOCK SIZE along the V dimension + USE_INITIAL_STATE: tl.constexpr, + CHECK: tl.constexpr +): + i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + o_i = tl.arange(0, BT) + + m_s = o_i[:, None] >= o_i[None, :] + # [BV, BK] + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + p_h = tl.make_block_ptr(h0 + i_bh * K * V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + + for i in range(0, tl.cdiv(T, BT)): + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i * BT), (BV, BT), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i * BT, i_v * BV), (BT, BV), (1, 0)) + p_dq = tl.make_block_ptr(dq + (i_bh + i_v*B*H) * T*K, (T, K), (K, 1), (i*BT, i_k*BK), (BT, BK), (1, 0)) + + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [V, BT] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, V] + b_do = tl.load(p_do, boundary_check=(0, 1)) + + # [BT, BT] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + b_ds = tl.where(m_s, b_ds, 0) + # [BT, BK] + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k, allow_tf32=False) + # [BV, BK] + if CHECK and i == 0: + b_dq += tl.dot(b_do, b_h.to(b_do.dtype), allow_tf32=False) + b_h = b_h + tl.dot(b_v, b_k, allow_tf32=False) + else: + b_dq += tl.dot(b_do, b_h.to(b_do.dtype), allow_tf32=False) + b_h = b_h + tl.dot(b_v, b_k, allow_tf32=False) + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + # sync threads + b_h = None + tl.debug_barrier() + # [BK, BV] + b_dh = tl.zeros([BK, BV], dtype=tl.float32) + m_s = o_i[:, None] <= o_i[None, :] + for i in range(1, tl.cdiv(T, BT) + 1): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, T - i * BT), (BK, BT), (0, 1)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (T - i * BT, i_k * BK), (BT, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (T - i * BT, i_v * BV), (BT, BV), (1, 0)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (T - i * BT, i_v * BV), (BT, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_bh+i_v*B*H) * T*K, (T, K), (K, 1), (T - i*BT, i_k*BK), (BT, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_bh+i_k*B*H) * T*V, (T, V), (V, 1), (T - i*BT, i_v*BV), (BT, BV), (1, 0)) + # [BK, BT] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)) + + # [BT, BT] + b_s = tl.dot(b_k, b_q, allow_tf32=False) + b_s = tl.where(m_s, b_s, 0).to(b_q.dtype) + # [BT, BT] + b_ds = tl.dot(b_v, tl.trans(b_do), allow_tf32=False) + b_ds = tl.where(m_s, b_ds, 0).to(b_q.dtype) + # [BT, BK] + b_dk = tl.dot(b_ds, tl.trans(b_q), allow_tf32=False) + # [BT, BV] + b_dv = tl.dot(b_s, b_do, allow_tf32=False) + if CHECK and i == 1: + b_dk += tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype), allow_tf32=False) + b_dv += tl.dot(b_k, b_dh.to(b_k.dtype), allow_tf32=False) + b_dh += tl.dot(b_q, b_do, allow_tf32=False) + else: + b_dk += tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype), allow_tf32=False) + b_dv += tl.dot(b_k, b_dh.to(b_k.dtype), allow_tf32=False) + b_dh += tl.dot(b_q, b_do, allow_tf32=False) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +class FusedChunkLinearAttentionFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale, initial_state, output_final_state): + B, H, T, K, V = *k.shape, v.shape[-1] + BT = 64 + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + num_warps = 4 + num_stages = 1 + + o = q.new_empty(NK, B, H, T, V) + final_state = q.new_empty(B, H, K, V, dtype=torch.float) if output_final_state else None + # the bug still exists even for Triton 2.2 on H100 GPUs + # so we always enable initial checks + CHECK = True + if version.parse(triton.__version__) < version.parse('2.2.0'): + import warnings + warnings.warn( + "Triton<2.2.0 detected for running this kernel, " + "which is known to have some weird compiler issues (refer to https://github.com/openai/triton/issues/2852) " + "that lead to significant precision loss. " + "We've add some initial condition checks to resolve this, sadly at the sacrifice of the speed. " + "For optimal performance, it is recommended to install Triton>=2.2.0 (if possible)." + ) + CHECK = True + + grid = (NV, NK, B * H) + fused_chunk_linear_attn_fwd_kernel[grid]( + q, k, v, o, initial_state, final_state, + scale, + B=B, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, + USE_INITIAL_STATE=initial_state is not None, + STORE_FINAL_STATE=output_final_state, + CHECK=CHECK, + num_warps=num_warps, + num_stages=num_stages + ) + o = o.sum(0) if NK > 1 else o[0] + + ctx.save_for_backward(q, k, v, initial_state) + ctx.scale = scale + ctx.CHECK = CHECK + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht=None): + q, k, v, initial_state = ctx.saved_tensors + B, H, T, K, V = *k.shape, v.shape[-1] + scale = ctx.scale + + BT = 64 + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + num_warps = 4 + num_stages = 1 + + dq = q.new_empty(NV, B, H, T, K) + dk = q.new_empty(NV, B, H, T, K) + dv = q.new_empty(NK, B, H, T, V) + grid = (NV, NK, B * H) + + fused_chunk_linear_attn_bwd_kernel[grid]( + q, k, v, do, dq, dk, dv, initial_state, + scale, + B=B, H=H, T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, + USE_INITIAL_STATE=initial_state is not None, + CHECK=ctx.CHECK, + num_warps=num_warps, + num_stages=num_stages + ) + dq = dq.sum(0) + dk = dk.sum(0) + dv = dv.sum(0) + return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), None, None, None + + +def fused_chunk_linear_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: Optional[float] = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + normalize: bool = True, + head_first: bool = True +) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]` + k (torch.Tensor): + keys of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]` + v (torch.Tensor): + values of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]` + scale (Optional[int]): + Scale factor for linear attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[B, H, K, V]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. + normalize (bool): + Whether to normalize the output. Default: `True`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `True`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]` + final_state (torch.Tensor): + Final state of shape `[B, H, K, V]` if `output_final_state=True` else `None` + """ + if scale is None: + scale = q.shape[-1] ** -0.5 + if not head_first: + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + o, final_state = FusedChunkLinearAttentionFunction.apply(q, k, v, scale, initial_state, output_final_state) + if normalize: + o = normalize_output(q * scale, k, o) + if not head_first: + o = o.transpose(1, 2) + return o, final_state diff --git a/fla/ops/nsa/__init__.py b/fla/ops/nsa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..941a1be41e1650961af0d28e64837421826ffab2 --- /dev/null +++ b/fla/ops/nsa/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- + +from .naive import naive_nsa +from .parallel import parallel_nsa + +__all__ = [ + 'naive_nsa', + 'parallel_nsa' +] diff --git a/fla/ops/nsa/naive.py b/fla/ops/nsa/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..79433b80ef024874816f3b8e4fff47fe93a2578d --- /dev/null +++ b/fla/ops/nsa/naive.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch +from einops import rearrange, repeat + + +def naive_nsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + indices: torch.LongTensor, + block_size: int = 64, + scale: Optional[float] = None, + head_first: bool = False, + cu_seqlens: Optional[torch.LongTensor] = None +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, HQ, T, K]` if `head_first=True` else `[B, T, HQ, K]`. + k (torch.Tensor): + keys of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16. + v (torch.Tensor): + values of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + indices (torch.LongTensor): + Block indices of shape `[B, T, H, S]` if `head_first=True` else `[B, T, H, S]`. + `S` is the number of selected blocks for each query token, which is set to 16 in the paper. + block_size (int): + Selected block size. Default: 64. + scale (Optional[int]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, HQ, T, V]` if `head_first=True` else `[B, T, HQ, V]`. + """ + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + if head_first: + raise RuntimeError("Sequences with variable lengths are not supported for head-first mode") + if head_first: + q, k, v, indices = map(lambda x: rearrange(x, 'b h t d -> b t h d'), (q, k, v, indices)) + + dtype = q.dtype + G = q.shape[2] // k.shape[2] + BS = block_size + k, v, indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, indices)) + q, k, v = map(lambda x: x.float(), (q, k, v)) + + o = torch.zeros_like(v) + varlen = True + if cu_seqlens is None: + varlen = False + B, T = q.shape[:2] + cu_seqlens = torch.cat([indices.new_tensor(range(0, B*T, T)), indices.new_tensor([B*T])]) + + for i in range(len(cu_seqlens) - 1): + if not varlen: + q_b, k_b, v_b, i_b = q[i], k[i], v[i], indices[i] + else: + T = cu_seqlens[i+1] - cu_seqlens[i] + q_b, k_b, v_b, i_b = map(lambda x: x[0][cu_seqlens[i]:cu_seqlens[i+1]], (q, k, v, indices)) + + i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS)) + # [T, S*BS, HQ] + i_b = i_b.view(T, indices.shape[2], -1).transpose(1, 2) + for i_q in range(T): + # [HQ, D] + q_i = q_b[i_q] * scale + # [S*BS, HQ] + i_i = i_b[i_q] + # [S*BS, HQ, -1] + k_i, v_i = map(lambda x: x.gather(0, i_i.clamp(0, T-1).unsqueeze(-1).expand(*i_i.shape, x.shape[-1])), (k_b, v_b)) + # [S*BS, HQ] + attn = torch.einsum('h d, n h d -> n h', q_i, k_i).masked_fill(i_i > i_q, float('-inf')).softmax(0) + if not varlen: + o[i, i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) + else: + o[0][cu_seqlens[i]+i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) + + if head_first: + o = rearrange(o, 'b t h d -> b h t d') + return o.to(dtype) diff --git a/fla/ops/nsa/parallel.py b/fla/ops/nsa/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..7e89d964c7357ceeabaaeb9500849ce6cbdecfad --- /dev/null +++ b/fla/ops/nsa/parallel.py @@ -0,0 +1,1435 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import warnings +from typing import Optional, Union + +import torch +import triton +import triton.language as tl +from einops import rearrange + +from fla.ops.common.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices +from fla.ops.nsa.utils import _bitonic_merge +from fla.ops.utils import mean_pooling +from fla.ops.utils.op import exp, log +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func +except ImportError: + warnings.warn( + "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", + category=ImportWarning + ) + flash_attn_func = None + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit +def parallel_nsa_compression_fwd_kernel( + q, + k, + v, + o, + lse, + scale, + offsets, + token_indices, + chunk_offsets, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr, +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + + p_o = tl.make_block_ptr(o + (bos + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + # [G, BV] + b_o = tl.zeros([G, BV], dtype=tl.float32) + # max scores for the current block + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + # lse = log(acc) + m + b_acc = tl.zeros([G], dtype=tl.float32) + + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c, i_v * BV), (BC, BV), (1, 0)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BC] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + + # [G, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + if NC == 0: + b_lse = tl.zeros([G], dtype=tl.float32) + else: + b_o = b_o / b_acc[:, None] + b_lse = b_m + log(b_acc) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + if i_v == 0: + tl.store(lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G), b_lse.to(lse.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_compression_bwd_kernel_dq( + q, + k, + v, + lse, + delta, + do, + dq, + scale, + offsets, + token_indices, + chunk_offsets, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + q += (bos + i_t) * HQ*K + do += (bos + i_t) * HQ*V + lse += (bos + i_t) * HQ + delta += (bos + i_t) * HQ + dq += (i_v * B * T + bos + i_t) * HQ*K + + p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + i_h * G + tl.arange(0, G) + p_delta = delta + i_h * G + tl.arange(0, G) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + + # [G, BK] + b_dq = tl.zeros([G, BK], dtype=tl.float32) + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (V, TC), (1, H*V), (i_v * BV, i_c), (BV, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BC] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_p = exp(b_s - b_lse[:, None]) + b_p = tl.where((o_c < NC)[None, :], b_p, 0) + + # [G, BV] @ [BV, BC] -> [G, BC] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [G, BC] @ [BC, BK] -> [G, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq *= scale + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_compression_bwd_kernel_dkv( + q, + k, + v, + lse, + delta, + do, + dk, + dv, + offsets, + chunk_indices, + chunk_offsets, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int32), tl.load(chunk_indices + i_c * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_v * B*T*H + boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_v * B*T*H + boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) + + # [BC, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BC, BK], dtype=tl.float32) + # [BC, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + + for i in range(i_c * BC * BS, T): + o_c = i_c * BC + tl.arange(0, BC) + + p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) + p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + # [BC, G] + b_s = tl.dot(b_k, tl.trans(b_q)) + b_p = exp(b_s - b_lse[None, :]) + b_p = tl.where((i >= max(0, (o_c + 1) * BS - 1))[:, None], b_p, 0) + # [BC, G] @ [G, BV] -> [BC, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BC, BV] @ [BV, G] -> [BC, G] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BC, G] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BC, G] @ [G, BK] -> [BC, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK'], +) +@triton.jit +def parallel_nsa_kernel_topk( + q, + k, + lse, + scale, + block_indices, + offsets, + token_indices, + chunk_offsets, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + S: tl.constexpr, + BC: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + USE_OFFSETS: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + boc = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_b * T, i_b * T + T + boc = i_b * tl.cdiv(T, BS) + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + # the number of compression representations in total + TC = tl.cdiv(T, BS) + # the number of compression representations required to iterate over + # incomplete compression blocks are not included + NC = (i_t + 1) // BS + ################################ + # 1. lse computation + ################################ + if lse is not None: + b_lse = tl.load(lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G)) + else: + # max scores for the current block + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + # lse = log(acc) + m + b_acc = tl.zeros([G], dtype=tl.float32) + for i_c in range(0, NC, BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BC] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + + b_mp = b_m + if NC == 0: + b_lse = tl.zeros([G], dtype=tl.float32) + else: + b_lse = b_m + log(b_acc) + + ################################ + # 2. topk selection + ################################ + # [BC] + b_i = tl.full([BC], -1, dtype=tl.float32) + o_i = tl.zeros([BC], dtype=tl.int32) + m_i = tl.arange(0, BC) < BC//2 + for i_c in range(0, i_t // BS + 1, BC): + o_c = i_c + tl.arange(0, BC) + + p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + # [BK, BC] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [G, BC] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((i_t // BS > o_c)[None, :], b_s, float('-inf')) + # [G, BC] + b_p = tl.where((i_t // BS == o_c)[None, :], float(1.0), exp(b_s - b_lse[:, None])) + # the importance scores of the current block + # [BC] + b_i, b_ip = tl.sum(b_p, 0), b_i + o_i, o_ip = tl.where(o_c <= i_t // BS, o_c + 1, 0), o_i + + n_dims: tl.constexpr = tl.standard._log2(b_i.shape[0]) + for i in tl.static_range(1, n_dims): + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), i, 2, n_dims) + + if i_c != 0: + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, False, n_dims) + b_i_new = b_ip * m_i + b_i * (1 - m_i) + o_i_new = o_ip * m_i + o_i * (1 - m_i) + b_i, o_i = _bitonic_merge(b_i_new, o_i_new.to(tl.int32), n_dims, True, n_dims) + else: + b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, True, n_dims) + + m_top = tl.arange(0, BC//S) == 0 + b_top = tl.sum(m_top[:, None] * tl.reshape(o_i - 1, [BC//S, S]), 0) + + p_b = tl.make_block_ptr(block_indices + (bos + i_t) * H*S, (H*S,), (1,), (i_h * S,), (S,), (0,)) + tl.store(p_b, b_top.to(p_b.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None, + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit +def parallel_nsa_fwd_kernel( + q, + k, + v, + o, + lse, + scale, + block_indices, + block_counts, + offsets, + token_indices, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + block_indices += (bos + i_t) * H*S + i_h * S + + if USE_BLOCK_COUNTS: + NS = tl.load(block_counts + (bos + i_t) * H + i_h) + else: + NS = S + + p_q = tl.make_block_ptr(q + (bos + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # the Q block is kept in the shared memory throughout the whole kernel + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_o = tl.make_block_ptr(o + (bos + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i_t) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_o = tl.zeros([G, BV], dtype=tl.float32) + + b_m = tl.full([G], float('-inf'), dtype=tl.float32) + b_acc = tl.zeros([G], dtype=tl.float32) + for i in range(NS): + i_s = tl.load(block_indices + i).to(tl.int32) * BS + if i_s <= i_t and i_s >= 0: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [G, BS] + b_s = tl.dot(b_q, b_k) + b_s = tl.where((i_t >= (i_s + tl.arange(0, BS)))[None, :], b_s, float('-inf')) + + # [G] + b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m + b_r = exp(b_mp - b_m) + # [G, BS] + b_p = exp(b_s - b_m[:, None]) + # [G] + b_acc = b_acc * b_r + tl.sum(b_p, 1) + # [G, BV] + b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) + + b_mp = b_m + b_o = b_o / b_acc[:, None] + b_m += log(b_acc) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor) +}) +@triton.jit +def parallel_nsa_kernel_mask( + block_indices, + block_counts, + block_mask, + T: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + NS: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr +): + i_t, i_b, i_hs = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_h, i_s = i_hs // S, i_hs % S + + b_i = tl.load(block_indices + i_b * T * H * S + i_t * H * S + i_h * S + i_s) + if USE_BLOCK_COUNTS: + b_m = b_i * BS <= i_t and i_s < tl.load(block_counts + i_b * T * H + i_t * H + i_h) + else: + b_m = b_i * BS <= i_t + + if b_i < NS and b_i >= 0: + tl.store(block_mask + i_b * T * H * NS + i_t * H * NS + i_h * NS + b_i, b_m.to(block_mask.dtype.element_ty)) + + +@triton.jit +def parallel_nsa_bwd_kernel_preprocess( + o, + do, + delta, + B: tl.constexpr, + V: tl.constexpr +): + i_n = tl.program_id(0) + o_d = tl.arange(0, B) + m_d = o_d < V + + b_o = tl.load(o + i_n * V + o_d, mask=m_d, other=0) + b_do = tl.load(do + i_n * V + o_d, mask=m_d, other=0).to(tl.float32) + b_delta = tl.sum(b_o * b_do) + + tl.store(delta + i_n, b_delta.to(delta.dtype.element_ty)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None, + 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor) +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_bwd_kernel_dq( + q, + k, + v, + lse, + delta, + do, + dq, + scale, + block_indices, + block_counts, + offsets, + token_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + S: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr, + USE_BLOCK_COUNTS: tl.constexpr +): + i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + q += (bos + i_t) * HQ*K + do += (bos + i_t) * HQ*V + lse += (bos + i_t) * HQ + delta += (bos + i_t) * HQ + dq += (i_v * B * T + bos + i_t) * HQ*K + block_indices += (bos + i_t) * H*S + i_h * S + + if USE_BLOCK_COUNTS: + NS = tl.load(block_counts + (bos + i_t) * H + i_h) + else: + NS = S + + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + + p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + i_h * G + tl.arange(0, G) + p_delta = delta + i_h * G + tl.arange(0, G) + + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + + # [G, BK] + b_dq = tl.zeros([G, BK], dtype=tl.float32) + for i in range(NS): + i_s = tl.load(block_indices + i).to(tl.int32) * BS + if i_s <= i_t and i_s >= 0: + p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) + p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + # [BK, BS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # [G, BS] + b_s = tl.dot(b_q, b_k) + b_p = exp(b_s - b_lse[:, None]) + b_p = tl.where((i_t >= (i_s + tl.arange(0, BS)))[None, :], b_p, 0) + + # [G, BV] @ [BV, BS] -> [G, BS] + b_dp = tl.dot(b_do, b_v) + b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) + # [G, BS] @ [BS, BK] -> [G, BK] + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq *= scale + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'USE_OFFSETS': lambda args: args['offsets'] is not None +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BS', 'BK', 'BV'], +) +@triton.jit(do_not_specialize=['T']) +def parallel_nsa_bwd_kernel_dkv( + q, + k, + v, + lse, + delta, + do, + dk, + dv, + block_mask, + offsets, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + M: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_OFFSETS: tl.constexpr +): + i_v, i_s, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if USE_OFFSETS: + i_n, i_s = tl.load(chunk_indices + i_s * 2).to(tl.int32), tl.load(chunk_indices + i_s * 2 + 1).to(tl.int32) + bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s * BS, 0), (BS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s * BS, i_v * BV), (BS, BV), (1, 0)) + p_dk = tl.make_block_ptr(dk + (i_v * B*T*H + bos * H + i_h) * K, (T, K), (H*K, 1), (i_s * BS, 0), (BS, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s * BS, i_v * BV), (BS, BV), (1, 0)) + + # [BS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dk = tl.zeros([BS, BK], dtype=tl.float32) + # [BS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_dv = tl.zeros([BS, BV], dtype=tl.float32) + + for i in range(i_s * BS, T): + b_m = tl.load(block_mask + (bos + i) * H*M + i_h * M + i_s) + if b_m: + p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + # [G, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + + p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) + p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) + # [G, BV] + b_do = tl.load(p_do, boundary_check=(0, 1)) + # [G] + b_lse = tl.load(p_lse) + b_delta = tl.load(p_delta) + # [BS, G] + b_s = tl.dot(b_k, tl.trans(b_q)) + b_p = exp(b_s - b_lse[None, :]) + b_p = tl.where((i >= (i_s * BS + tl.arange(0, BS)))[:, None], b_p, 0) + # [BS, G] @ [G, BV] -> [BS, BV] + b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + # [BS, BV] @ [BV, G] -> [BS, G] + b_dp = tl.dot(b_v, tl.trans(b_do)) + # [BS, G] + b_ds = b_p * (b_dp - b_delta[None, :]) + # [BS, G] @ [G, BK] -> [BS, BK] + b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + +def parallel_nsa_compression_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int, + scale: float, + offsets: Optional[torch.LongTensor] = None, + token_indices: Optional[torch.LongTensor] = None, +): + B, T, HQ, K, V = *q.shape, v.shape[-1] + H = k.shape[2] + G = HQ // H + BC = BS = block_size + if check_shared_mem('hopper', q.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + else: + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, "The key dimension can not be larger than 256" + + chunk_offsets = prepare_chunk_offsets(offsets, BS) if offsets is not None else None + + grid = (T, NV, B * H) + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + parallel_nsa_compression_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + lse=lse, + scale=scale, + offsets=offsets, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV, + ) + return o, lse + + +def parallel_nsa_compression_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + block_size: int = 64, + scale: float = None, + offsets: Optional[torch.LongTensor] = None, + token_indices: Optional[torch.LongTensor] = None, +): + B, T, HQ, K, V = *q.shape, v.shape[-1] + H = k.shape[2] + G = HQ // H + BC = BS = block_size + BK = triton.next_power_of_2(K) + BV = min(128, triton.next_power_of_2(v.shape[-1])) + NV = triton.cdiv(V, BV) + if offsets is not None: + lens = prepare_lens(offsets) + chunk_indices = torch.cat([torch.arange(n) for n in triton.cdiv(triton.cdiv(lens, BS), BC).tolist()]) + chunk_indices = torch.stack([chunk_indices.eq(0).cumsum(0) - 1, chunk_indices], 1).to(offsets) + chunk_offsets = prepare_chunk_offsets(offsets, BS) + NC = len(chunk_indices) + else: + chunk_indices, chunk_offsets = None, None + NC = triton.cdiv(triton.cdiv(T, BS), BC) + + delta = parallel_nsa_bwd_preprocess(o, do) + + dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) + grid = (T, NV, B * H) + parallel_nsa_compression_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dq=dq, + scale=scale, + offsets=offsets, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV + ) + dq = dq.sum(0) + + dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) + dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) + + grid = (NV, NC, B * H) + parallel_nsa_compression_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + offsets=offsets, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BC=BC, + BS=BS, + BK=BK, + BV=BV + ) + dk = dk.sum(0) + return dq, dk, dv + + +class ParallelNSACompressionFunction(torch.autograd.Function): + + @staticmethod + @contiguous + @autocast_custom_fwd + def forward( + ctx, + q, + k, + v, + block_size, + scale, + offsets + ): + ctx.dtype = q.dtype + + # 2-d sequence indices denoting the offsets of tokens in each sequence + # for example, if the passed `offsets` is [0, 2, 6], + # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + token_indices = prepare_token_indices(offsets) if offsets is not None else None + + o, lse = parallel_nsa_compression_fwd( + q=q, + k=k, + v=v, + block_size=block_size, + scale=scale, + offsets=offsets, + token_indices=token_indices + ) + ctx.save_for_backward(q, k, v, o, lse) + ctx.offsets = offsets + ctx.token_indices = token_indices + ctx.block_size = block_size + ctx.scale = scale + return o.to(q.dtype), lse + + @staticmethod + @contiguous + @autocast_custom_bwd + def backward(ctx, do, *args): + q, k, v, o, lse = ctx.saved_tensors + dq, dk, dv = parallel_nsa_compression_bwd( + q=q, + k=k, + v=v, + o=o, + lse=lse, + do=do, + block_size=ctx.block_size, + scale=ctx.scale, + offsets=ctx.offsets, + token_indices=ctx.token_indices + ) + return dq.to(q), dk.to(k), dv.to(v), None, None, None + + +def parallel_nsa_topk( + q: torch.Tensor, + k: torch.Tensor, + lse: torch.Tensor, + block_counts: Union[torch.LongTensor, int], + block_size: int = 64, + scale: float = None, + offsets: Optional[torch.LongTensor] = None, +) -> torch.LongTensor: + B, T, HQ, K = q.shape + H = k.shape[2] + G = HQ // H + S = block_counts if isinstance(block_counts, int) else block_counts.max().item() + S = triton.next_power_of_2(S) + # here we set BC = BS, but beware that they are actually decoupled + BC = BS = block_size + BK = triton.next_power_of_2(K) + + block_indices = torch.zeros(B, T, H, S, dtype=torch.int32, device=q.device) + token_indices = prepare_token_indices(offsets) if offsets is not None else None + chunk_offsets = prepare_chunk_offsets(offsets, BS) if offsets is not None else None + grid = (T, B * H) + parallel_nsa_kernel_topk[grid]( + q=q, + k=k, + lse=lse, + scale=scale, + block_indices=block_indices, + offsets=offsets, + token_indices=token_indices, + chunk_offsets=chunk_offsets, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + S=S, + BC=BC, + BS=BS, + BK=BK + ) + return block_indices + + +def parallel_nsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_indices: torch.LongTensor, + block_counts: Union[torch.LongTensor, int], + block_size: int, + scale: float, + offsets: Optional[torch.LongTensor] = None, + token_indices: Optional[torch.LongTensor] = None, +): + B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] + HQ = q.shape[2] + G = HQ // H + BS = block_size + if check_shared_mem('hopper', q.device.index): + BK = min(256, triton.next_power_of_2(K)) + BV = min(256, triton.next_power_of_2(V)) + else: + BK = min(128, triton.next_power_of_2(K)) + BV = min(128, triton.next_power_of_2(V)) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + assert NK == 1, "The key dimension can not be larger than 256" + + grid = (T, NV, B * H) + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + + parallel_nsa_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + lse=lse, + scale=scale, + block_indices=block_indices, + block_counts=block_counts, + offsets=offsets, + token_indices=token_indices, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + S=S, + BS=BS, + BK=BK, + BV=BV, + ) + return o, lse + + +def parallel_nsa_block_mask( + block_indices: torch.LongTensor, + block_counts: Union[torch.LongTensor, int], + offsets: torch.LongTensor, + block_size: int, +): + B, T, H, S = block_indices.shape + BS = block_size + if offsets is not None: + NS = triton.cdiv(prepare_lens(offsets).max().item(), BS) + else: + NS = triton.cdiv(T, BS) + block_mask = torch.zeros(B, T, H, NS, dtype=torch.bool, device=block_indices.device) + + parallel_nsa_kernel_mask[(T, B, H*S)]( + block_indices=block_indices, + block_counts=block_counts, + block_mask=block_mask, + T=T, + H=H, + S=S, + BS=BS, + NS=NS + ) + return block_mask + + +def parallel_nsa_bwd_preprocess( + o: torch.Tensor, + do: torch.Tensor +): + V = o.shape[-1] + delta = torch.empty_like(o[..., 0], dtype=torch.float32) + parallel_nsa_bwd_kernel_preprocess[(delta.numel(),)]( + o=o, + do=do, + delta=delta, + B=triton.next_power_of_2(V), + V=V, + ) + return delta + + +def parallel_nsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + block_indices: torch.Tensor, + block_counts: Union[torch.LongTensor, int], + block_size: int = 64, + scale: float = None, + offsets: Optional[torch.LongTensor] = None, + token_indices: Optional[torch.LongTensor] = None, +): + B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] + HQ = q.shape[2] + G = HQ // H + BS = block_size + BK = triton.next_power_of_2(K) + BV = min(128, triton.next_power_of_2(v.shape[-1])) + NV = triton.cdiv(V, BV) + + delta = parallel_nsa_bwd_preprocess(o, do) + + dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) + grid = (T, NV, B * H) + parallel_nsa_bwd_kernel_dq[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dq=dq, + block_indices=block_indices, + block_counts=block_counts, + offsets=offsets, + token_indices=token_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + S=S, + BS=BS, + BK=BK, + BV=BV + ) + dq = dq.sum(0) + + if offsets is not None: + chunk_indices = prepare_chunk_indices(offsets, BS) + NS = len(chunk_indices) + else: + chunk_indices = None + NS = triton.cdiv(T, BS) + + # [B, T, H, M] + block_mask = parallel_nsa_block_mask(block_indices, block_counts, offsets, block_size) + dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) + dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) + + grid = (NV, NS, B * H) + parallel_nsa_bwd_kernel_dkv[grid]( + q=q, + k=k, + v=v, + lse=lse, + delta=delta, + do=do, + dk=dk, + dv=dv, + block_mask=block_mask, + offsets=offsets, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + M=block_mask.shape[-1], + BS=BS, + BK=BK, + BV=BV + ) + dk = dk.sum(0) + return dq, dk, dv + + +@torch.compile +class ParallelNSAFunction(torch.autograd.Function): + + @staticmethod + @contiguous + @autocast_custom_fwd + def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, offsets): + ctx.dtype = q.dtype + + # 2-d sequence indices denoting the offsets of tokens in each sequence + # for example, if the passed `offsets` is [0, 2, 6], + # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be + # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] + token_indices = prepare_token_indices(offsets) if offsets is not None else None + + o, lse = parallel_nsa_fwd( + q=q, + k=k, + v=v, + block_indices=block_indices, + block_counts=block_counts, + block_size=block_size, + scale=scale, + offsets=offsets, + token_indices=token_indices + ) + ctx.save_for_backward(q, k, v, o, lse) + ctx.block_indices = block_indices + ctx.block_counts = block_counts + ctx.offsets = offsets + ctx.token_indices = token_indices + ctx.block_size = block_size + ctx.scale = scale + return o.to(q.dtype) + + @staticmethod + @contiguous + @autocast_custom_bwd + def backward(ctx, do): + q, k, v, o, lse = ctx.saved_tensors + dq, dk, dv = parallel_nsa_bwd( + q=q, + k=k, + v=v, + o=o, + lse=lse, + do=do, + block_indices=ctx.block_indices, + block_counts=ctx.block_counts, + block_size=ctx.block_size, + scale=ctx.scale, + offsets=ctx.offsets, + token_indices=ctx.token_indices + ) + return dq.to(q), dk.to(k), dv.to(v), None, None, None, None, None, None, None, None + + +def parallel_nsa_compression( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int = 64, + scale: float = None, + offsets: Optional[torch.LongTensor] = None +): + return ParallelNSACompressionFunction.apply( + q, + k, + v, + block_size, + scale, + offsets + ) + + +def parallel_nsa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cmp: torch.Tensor, + g_slc: torch.Tensor, + g_swa: torch.Tensor, + block_indices: Optional[torch.LongTensor] = None, + block_counts: Union[torch.LongTensor, int] = 16, + block_size: int = 64, + window_size: int = 0, + scale: Optional[float] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False +) -> torch.Tensor: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, HQ, K]` if `head_first=False` else `[B, HQ, T, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`. + GQA is enforced here. The ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >=16. + v (torch.Tensor): + values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. + g_cmp (torch.Tensor): + Gate score for compressed attention of shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. + g_slc (torch.Tensor): + Gate score for selected attention of shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. + g_swa (torch.Tensor): + Gate score for sliding attentionof shape `[B, T, HQ]` if `head_first=False` else `[B, HQ, T]`. + block_indices (torch.LongTensor): + Block indices of shape `[B, T, H, S]` if `head_first=False` else `[B, H, T, S]`. + `S` is the number of selected blocks for each query token, which is set to 16 in the paper. + If `g_cmp` is provided, the passed `block_indices` will be ignored. + block_counts (Optional[Union[torch.LongTensor, int]]): + Number of selected blocks for each query. + If a tensor is provided, with shape `[B, T, H]` if `head_first=False` else `[B, H, T]`, + each query can select the same number of blocks. + If not provided, it will default to 16. + block_size (int): + Selected block size. Default: 64. + window_size (int): + Sliding window size. Default: 0. + scale (Optional[int]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + head_first (Optional[bool]): + Whether the inputs are in the head-first format. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]` if `head_first=False` else `[B, HQ, T, V]`. + """ + assert block_counts is not None, "block counts must be provided for selection" + if scale is None: + scale = k.shape[-1] ** -0.5 + if cu_seqlens is not None: + assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" + if head_first: + q, k, v = map(lambda x: rearrange(x, 'b h t d -> b t h d'), (q, k, v)) + g_cmp, g_slc, g_swa = map(lambda x: rearrange(x, 'b h t -> b t h') if x is not None else None, (g_cmp, g_slc, g_swa)) + if not isinstance(block_counts, int): + block_counts = rearrange(block_counts, 'b h t -> b t h') + assert q.shape[2] % (k.shape[2] * 16) == 0, "Group size must be a multiple of 16 in NSA" + + k_cmp, v_cmp = mean_pooling(k, block_size, cu_seqlens), mean_pooling(v, block_size, cu_seqlens) + o_cmp, lse_cmp = None, None + if g_cmp is not None: + o_cmp, lse_cmp = parallel_nsa_compression( + q=q, + k=k_cmp, + v=v_cmp, + block_size=block_size, + scale=scale, + offsets=cu_seqlens + ) + if block_indices is not None: + warnings.warn("`block_indices` will be ignored when `g_cmp` is provided") + block_indices = parallel_nsa_topk( + q=q, + k=k_cmp, + lse=lse_cmp, + block_counts=block_counts, + block_size=block_size, + scale=scale, + offsets=cu_seqlens + ) + o_slc = ParallelNSAFunction.apply(q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens) + o = o_slc * g_slc.unsqueeze(-1) + if o_cmp is not None: + o = torch.addcmul(o, o_cmp, g_cmp.unsqueeze(-1)) + if window_size > 0: + if cu_seqlens is not None: + max_seqlen = q.shape[1] + o_swa = flash_attn_varlen_func( + q.squeeze(0), k.squeeze(0), v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=True, + window_size=(window_size-1, 0) + ).unsqueeze(0) + else: + o_swa = flash_attn_func( + q, k, v, + causal=True, + window_size=(window_size-1, 0) + ) + o = torch.addcmul(o, o_swa, g_swa.unsqueeze(-1)) + if head_first: + o = rearrange(o, 'b t h d -> b h t d') + return o diff --git a/fla/ops/rebased/__init__.py b/fla/ops/rebased/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ec6a0cb31f7f635aa528cad753d5e19196a2028 --- /dev/null +++ b/fla/ops/rebased/__init__.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- + +from .parallel import parallel_rebased + +__all__ = [ + 'parallel_rebased' +] diff --git a/fla/ops/rebased/naive.py b/fla/ops/rebased/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..a70242eb3c2dcc5918503af8b03a15b5740e4c2a --- /dev/null +++ b/fla/ops/rebased/naive.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch + + +def naive_parallel_rebased( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: Optional[float] = None, + use_norm: bool = True, +) -> torch.Tensor: + if scale is None: + scale = q.shape[-1] ** -0.5 + q = q * scale + attn = q @ k.transpose(-2, -1) + attn = attn ** 2 + attn.masked_fill_(~torch.tril(torch.ones(q.shape[-2], q.shape[-2], dtype=torch.bool, device=q.device)), 0) + o = attn @ v + if use_norm: + z = attn.sum(-1) + return o / (z[..., None] + 1e-6) + else: + return o diff --git a/fla/ops/rebased/parallel.py b/fla/ops/rebased/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..54de25ffb88a2b82672661f62f42257bac5c5fa3 --- /dev/null +++ b/fla/ops/rebased/parallel.py @@ -0,0 +1,466 @@ + +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + +# Rebased: Linear Transformers with Learnable Kernel Functions are Better In-Context Models +# https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/ops/triton/rebased_fast/parallel.py + + +@triton.jit(do_not_specialize=['T']) +def parallel_rebased_fwd_kernel( + q, + k, + v, + o, + z, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # i_c: chunk index. used for sequence parallelism + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % (NV) + + p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, 0), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v*BV), (BTS, BV), (1, 0)) + + # [BQ, BD] block Q, in the shared memory throughout the whole kernel + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + b_o = tl.zeros([BTL, BV], dtype=tl.float32) + b_z = tl.zeros([BTL], dtype=tl.float32) + + # Q block and K block have no overlap + # no need for mask, thereby saving flops + for _ in range(0, i_c*BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_s = tl.dot(b_q, (b_k), allow_tf32=False) + b_s = b_s * b_s + b_z += tl.sum(b_s, axis=1) + + # [BQ, BD] + b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + + # # rescale interchunk output + tl.debug_barrier() + o_q = tl.arange(0, BTL) + # # sync threads, easy for compiler to optimize + # tl.debug_barrier() + + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, i_c*BTL), (BK, BTS), (0, 1)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTS, BV), (1, 0)) + # Q block and K block have overlap. masks required + for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + # [BK, BTS] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BTS, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_s = tl.dot(b_q, b_k, allow_tf32=False) + b_s = b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_z += tl.sum(b_s, axis=1) + # [BTL, BV] + b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + p_k = tl.advance(p_k, (0, BTS)) + p_v = tl.advance(p_v, (BTS, 0)) + o_k += BTS + + p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_z = z + (i_bh + B * H * i_k) * T + i_c*BTL + tl.arange(0, BTL) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c*BTL + tl.arange(0, BTL)) < T)) + + +@triton.jit(do_not_specialize=['T']) +def _parallel_rebased_bwd_dq( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dq, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr +): + p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_q = (b_q * scale).to(b_q.dtype) + b_dq = tl.zeros([BTL, BK], dtype=tl.float32) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k*BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, 0), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i_c*BTL + tl.arange(0, BTL) + b_dz = tl.load(p_dz, mask=(i_c*BTL + tl.arange(0, BTL)) < T) + + for _ in range(0, i_c*BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + # [BQ, BD] + b_dq += tl.dot((2 * b_ds * b_s).to(b_v.dtype), b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + + b_dq *= scale + o_q = tl.arange(0, BTL) + o_k = tl.arange(0, BTS) + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTS, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, i_c*BTL), (BV, BTS), (0, 1)) + # Q block and K block have overlap. masks required + for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + # [BTS, BK] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BTS] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BTL, BTS] + m_s = o_q[:, None] >= o_k[None, :] + b_ds = tl.dot(b_do, b_v, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[:, None] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) + b_s = tl.where(m_s, b_s, 0) + # [BTL, BK] + b_dq += tl.dot((2 * b_ds * b_s).to(b_k.dtype), + b_k, allow_tf32=False) + p_k = tl.advance(p_k, (BTS, 0)) + p_v = tl.advance(p_v, (0, BTS)) + o_k += BTS + p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit(do_not_specialize=['T']) +def _parallel_rebased_bwd_dkv( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + # compute dk dv + p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) + b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros( + [BTL, BV], dtype=tl.float32) + + for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + # [BK, BTS] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BV, BTS] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + # [BTL, BTS] + b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale + b_s2 = b_s * b_s + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale + if i_v == 0: + b_ds += b_dz[None, :] * scale + else: + b_ds = b_ds + b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + + tl.debug_barrier() + o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) + for i in range(i_c*BTL, (i_c+1)*BTL, BTS): + p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) + p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + p_dz = dz + i_bh * T + i + tl.arange(0, BTS) + b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] + b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) + # [BK, BQ] + m_s = o_k[:, None] <= o_q[None, :] + b_s = tl.dot(b_k, b_q, allow_tf32=False) * scale + b_s2 = b_s * b_s + b_s = tl.where(m_s, b_s, 0) + b_s2 = tl.where(m_s, b_s2, 0) + + b_ds = tl.dot(b_v, b_do, allow_tf32=False) + if i_v == 0: + b_ds += b_dz[None, :] + else: + b_ds = b_ds + b_ds = tl.where(m_s, b_ds, 0) * scale + # [BK, BD] + b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + o_q += BTS + + p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) + p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + return + + +@triton.jit(do_not_specialize=['T']) +def parallel_rebased_bwd_kernel( + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BTL: tl.constexpr, + BTS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr +): + i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_k = i_kv // (NV) + i_v = i_kv % (NV) + i_h = i_bh % H + _parallel_rebased_bwd_dq( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dq, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV + ) + tl.debug_barrier() + _parallel_rebased_bwd_dkv( + i_bh, + i_c, + i_k, + i_v, + i_h, + q, + k, + v, + do, + dz, + dk, + dv, + scale, + B=B, + H=H, + T=T, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV + ) + + +class ParallelBasedFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale): + BTL, BTS = 128, 32 + assert BTL % BTS == 0 + # assert q.shape[-1] % 16 == 0 + BK = min(128, triton.next_power_of_2(k.shape[-1])) + BV = min(128, triton.next_power_of_2(v.shape[-1])) + BK, BV = max(BK, 16), max(BV, 16) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not." + + o = torch.empty(NK, B, H, T, V, device=q.device) + z = torch.empty(NK, B, H, T, device=q.device) + parallel_rebased_fwd_kernel[grid]( + q, + k, + v, + o, + z, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages + ) + ctx.save_for_backward(q, k, v) + ctx.scale = scale + return o.sum(0).to(q.dtype), z.sum(0).to(q.dtype) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dz): + q, k, v = ctx.saved_tensors + scale = ctx.scale + BTL, BTS = 64, 32 + assert BTL % BTS == 0 + BK = min(128, triton.next_power_of_2(k.shape[-1])) + BV = min(128, triton.next_power_of_2(v.shape[-1])) + BK, BV = max(BK, 16), max(BV, 16) + B, H, T, K, V = *k.shape, v.shape[-1] + num_stages = 2 + num_warps = 4 + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + grid = (NK * NV, triton.cdiv(T, BTL), B * H) + + assert NK == 1, "will encounter some synchronization issue if not" + + dq = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dk = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) + dv = torch.empty(NK, B, H, T, V, dtype=q.dtype, device=q.device) + + parallel_rebased_bwd_kernel[grid]( + q, + k, + v, + do, + dz, + dq, + dk, + dv, + scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BTL=BTL, + BTS=BTS, + BK=BK, + BV=BV, + num_warps=num_warps, + num_stages=num_stages + ) + + return dq.sum(0).to(q.dtype), dk.sum(0).to(k.dtype), dv.sum(0).to(v.dtype), None + + +def parallel_rebased( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + eps: float = 1e-5, + use_scale: bool = True, + use_normalize: bool = True, + return_both: bool = False, + head_first: bool = True +): + assert q.shape[-1] <= 128, "only support feature dim up to 128" + if use_scale: + scale = q.shape[-1] ** -0.5 + else: + scale = 1 + if not head_first: + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + o, z = ParallelBasedFunction.apply(q, k, v, scale) + if return_both: + return o, z + if use_normalize: + o = o / (z[..., None] + eps) + if not head_first: + o = o.transpose(1, 2) + return o.to(q.dtype) diff --git a/fla/ops/rwkv7/__init__.py b/fla/ops/rwkv7/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5efa2f31d0fa68361a4ff19e65d797157e03c83d --- /dev/null +++ b/fla/ops/rwkv7/__init__.py @@ -0,0 +1,9 @@ +# -*- coding: utf-8 -*- + +from .chunk import chunk_rwkv7 +from .fused_recurrent import fused_recurrent_rwkv7 + +__all__ = [ + 'chunk_rwkv7', + 'fused_recurrent_rwkv7' +] diff --git a/fla/ops/rwkv7/channel_mixing.py b/fla/ops/rwkv7/channel_mixing.py new file mode 100644 index 0000000000000000000000000000000000000000..991ea426f086859b8eaf0623f5acf059f9bddc5c --- /dev/null +++ b/fla/ops/rwkv7/channel_mixing.py @@ -0,0 +1,323 @@ +import logging + +import torch +import triton +import triton.language as tl + +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_pytorch_version, input_guard, use_cuda_graph + +logger = logging.getLogger(__name__) + +if not check_pytorch_version('2.4'): + logger.warning('PyTorch < 2.4 detected - computations may be slower due to lack of optimizations') + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': block_size}) + for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] + ], + key=['hidden_dim'], + use_cuda_graph=use_cuda_graph, +) +@triton.jit +def rwkv_seq_mix_kernel( + x_ptr, + x_prev_ptr, + mix_k_ptr, + output_ptr, + batch_size: tl.constexpr, + token_length, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr +): + block_start = tl.program_id(0) * BLOCK_SIZE + block_idx = block_start + tl.arange(0, BLOCK_SIZE)[:] + + total_seq_dim = token_length * hidden_dim + batch_idx = block_idx // total_seq_dim + seq_and_feat = block_idx % total_seq_dim + seq_idx = seq_and_feat // hidden_dim + feat_idx = seq_and_feat % hidden_dim + + is_valid = (batch_idx < batch_size) & (seq_idx < token_length) + + x_idx = batch_idx * total_seq_dim + seq_idx * hidden_dim + feat_idx + + curr_x = tl.load(x_ptr + x_idx, mask=is_valid, other=0.0).to(tl.float32) + k_value = tl.load(mix_k_ptr + feat_idx).to(tl.float32) + + is_first = seq_idx < 1 + prev_state_idx = batch_idx * hidden_dim + feat_idx + prev_state = tl.load(x_prev_ptr + prev_state_idx, + mask=(is_first & is_valid), + other=0.0).to(tl.float32) + + prev_x_idx = x_idx - hidden_dim + prev_x = tl.load(x_ptr + prev_x_idx, + mask=(~is_first & is_valid), + other=0.0).to(tl.float32) + + prev_value = tl.where(is_first, prev_state, prev_x) + state_diff = prev_value - curr_x + mixed = state_diff * k_value + result = tl.cast(curr_x + mixed, dtype=output_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(output_ptr + x_idx, result, mask=is_valid) + + +@triton.jit +def rwkv_channel_mixing_pow_and_relu( + in_ptr, + out_ptr, + BLOCK_SIZE: tl.constexpr +): + """Fused ReLU and Power operation: x = ReLU(x)^2""" + xoffset = tl.program_id(0) * BLOCK_SIZE + xindex = xoffset + tl.arange(0, BLOCK_SIZE) + x0 = xindex + x = tl.load(in_ptr + (x0), None) + x = tl.maximum(x, 0.0).to(tl.float32) + x = tl.cast(x * x, dtype=out_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(out_ptr + (x0), x, None) + + +def rwkv_mix_torch(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor): + if x_prev.dim() == 2: + x_prev = x_prev.unsqueeze(1) # (batch_size, 1, hidden_dim) + xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x + k = x.addcmul(xx, x_k) + return k + + +def rwkv_relu_and_square_torch(x: torch.Tensor): + return torch.relu(x) ** 2 + + +def rwkv_mix_fwd(x, x_prev, x_k): + has_batch = x.dim() == 3 + + if has_batch: + batch_size, token_length, hidden_dim = x.shape + else: + token_length, hidden_dim = x.shape + batch_size = 1 + x = x.unsqueeze(0) + x_prev = x_prev.unsqueeze(0) + + token_length = x.shape[1] + hidden_dim = x.shape[2] + total_elements = batch_size * token_length * hidden_dim + + output = torch.empty_like(x) + + def grid(meta): return ( + (total_elements + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 + 1, # grid_1 + 1 # grid_2 + ) + + rwkv_seq_mix_kernel[grid]( + x.contiguous(), + x_prev.contiguous(), + x_k.squeeze(), + output, + batch_size=batch_size, + token_length=token_length, + hidden_dim=hidden_dim, + ) + if not has_batch: + output = output.squeeze(0) + return output + + +def rwkv_relu_and_square_fwd(x: torch.Tensor, inplace: bool = True): + """ + Triton implementation of RWKV's ReLU and square operation + Args: + x: Input tensor + Returns: + Tensor after ReLU and square operations + """ + x = x.contiguous() + output = x if inplace else torch.empty_like(x) + + def grid(meta): return ( + (output.numel() + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 + 1, # grid_1 + 1 # grid_2 + ) + rwkv_channel_mixing_pow_and_relu[grid]( + x, + output, + BLOCK_SIZE=4096, + ) + + return output + + +@triton.jit +def relu_square_bwd_kernel( + out_ptr, + forward_input_ptr, + BLOCK_SIZE: tl.constexpr +): + """ReLU(x)^2 backward kernel + grad_input = grad_output * 2 * x if x > 0 else 0 + """ + pid = tl.program_id(0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + x = tl.load(forward_input_ptr + offsets).to(tl.float32) + grad = tl.load(out_ptr + offsets).to(tl.float32) + + x = tl.maximum(x, 0.0) + + grad_input = grad * 2 * x + + tl.store(out_ptr + offsets, grad_input.to(out_ptr.dtype.element_ty)) + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': block_size}) + for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] + ], + key=['hidden_dim'], + use_cuda_graph=use_cuda_graph, +) +@triton.jit +def rwkv_mix_bwd_kenel( + dk1_ptr0, + xk_ptr, + dx_ptr, + dx_prev_ptr, + batch_size, + token_length, + hidden_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + + batch_idx = offsets // (token_length * hidden_dim) + seq_feat = offsets % (token_length * hidden_dim) + seq_idx = seq_feat // hidden_dim + feat_idx = seq_feat % hidden_dim + + is_valid = offsets < (batch_size * token_length * hidden_dim) + + dk1 = tl.load(dk1_ptr0 + offsets, mask=is_valid) + xk = tl.load(xk_ptr + feat_idx, mask=is_valid) + prod = dk1 * xk + + mask_next = seq_idx < (token_length - 1) + next_offset = offsets + hidden_dim + dk1_next = tl.load(dk1_ptr0 + next_offset, mask=mask_next & is_valid, other=0.0) + prod_next = dk1_next * xk + dx_val = dk1 - prod + tl.where(mask_next, prod_next, 0.0) + dx_val = tl.cast(dx_val, dtype=dx_ptr.dtype.element_ty, fp_downcast_rounding='rtne') + tl.store(dx_ptr + offsets, dx_val, mask=is_valid) + + dx_prev_offset = batch_idx * hidden_dim + feat_idx + is_first_step = seq_idx == 0 + + tl.store( + dx_prev_ptr + dx_prev_offset, + tl.cast(prod, dtype=dx_prev_ptr.dtype.element_ty), + mask=is_first_step + ) + + +@torch.compile(fullgraph=True) +def compute_x_k_grad(dk1, x, x_prev): + """ + Args: + dk1: (batch*seq_len, hidden_dim) + x: (batch, seq_len, hidden_dim) + x_prev: (batch, hidden_dim) or (batch, 1, hidden_dim) + """ + + if x_prev.dim() == 2: + x_prev = x_prev.unsqueeze(1) # (batch, 1, hidden_dim) + xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x # (batch, seq_len, hidden_dim) + + # (hidden_dim,) --> (1, 1, hidden_dim) + grad_x_k = (dk1 * xx.reshape(-1, x.shape[2])).sum(dim=0).view(1, 1, -1) + return grad_x_k + + +def rwkv_channel_mixing_bwd(grad_output, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, inplace=True): + batch_size = x.shape[0] if x.dim() == 3 else 1 + seq_len, n_embd = x.shape[-2], x.shape[-1] + + dV = k.transpose(-2, -1) @ grad_output + dk = grad_output @ value_weight.transpose(-2, -1) + + BLOCK_SIZE = 4096 + grid = ((dk.numel() + BLOCK_SIZE - 1) // BLOCK_SIZE,) + relu_square_bwd_kernel[grid]( + dk, + k1_K, + BLOCK_SIZE=BLOCK_SIZE + ) + + dK = k1.transpose(-2, -1) @ dk + dk1 = dk @ key_weight.transpose(-2, -1) + dk1 = dk1.view(-1, n_embd).contiguous() + + dk_reduced = compute_x_k_grad(dk1, x, x_prev) + dx_prev = torch.empty_like(x_prev) if not inplace else x_prev + dx = torch.empty_like(x) if not inplace else x + + def grid(meta): return ((batch_size * seq_len * n_embd + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], 1, 1) + rwkv_mix_bwd_kenel[grid]( + dk1, + x_k.squeeze(), + dx, + dx_prev, + batch_size, + seq_len, + n_embd, + ) + # dx_prev.shape batch_size, seq_len, n_embd + return dx, dx_prev, dk_reduced, dK, dV + + +class Rwkv7ChannelMixing(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, x, x_prev, x_k, key_weight, value_weight, inplace: bool = True): + k1 = rwkv_mix_fwd(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_fwd(k1_K, inplace=True) + ctx.save_for_backward(x, x_prev, x_k, key_weight, value_weight) + ctx.inplace = inplace + return k @ value_weight + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dkv): + x, x_prev, x_k, key_weight, value_weight = ctx.saved_tensors + k1 = rwkv_mix_fwd(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_fwd(k1_K, inplace=False) + dx, dx_prev, dk_reduced, dK, dV = rwkv_channel_mixing_bwd( + dkv, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, ctx.inplace) + return dx, dx_prev, dk_reduced, dK, dV, None + + +def channel_mixing_rwkv7(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor, + key_weight: torch.Tensor, value_weight: torch.Tensor, inplace: bool = True): + assert x.dim() == 3 + + return Rwkv7ChannelMixing.apply(x, x_prev, x_k, key_weight, value_weight, inplace), x[-1, :] + + +def channel_mixing_rwkv7_torch(x, x_prev, x_k, key_weight, value_weight): + k1 = rwkv_mix_torch(x, x_prev, x_k) + k1_K = k1 @ key_weight + k = rwkv_relu_and_square_torch(k1_K) + return k @ value_weight, x[-1, :] diff --git a/fla/ops/rwkv7/chunk.py b/fla/ops/rwkv7/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..956c458974316d1f83121bd71ef1ec433cb6cdde --- /dev/null +++ b/fla/ops/rwkv7/chunk.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch + +from fla.ops.generalized_delta_rule import chunk_dplr_delta_rule + + +def chunk_rwkv7( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = 1.0, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False +): + """ + Args: + r (torch.Tensor): + r of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + w (torch.Tensor): + log decay of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + k (torch.Tensor): + k of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + v (torch.Tensor): + v of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + scale (float): + scale of the attention. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (bool): + whether to use head first. Recommended to be False to avoid extra transposes. + """ + return chunk_dplr_delta_rule( + q=r, + k=k, + v=v, + a=a, + b=b, + gk=w, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + head_first=head_first + ) diff --git a/fla/ops/rwkv7/fused_recurrent.py b/fla/ops/rwkv7/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..0ce2d15aec9c6995f2df26992c89a29182f0169d --- /dev/null +++ b/fla/ops/rwkv7/fused_recurrent.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2024-2025, Songlin Yang, Yu Zhang + +from typing import Optional + +import torch + +from fla.ops.generalized_delta_rule import fused_recurrent_dplr_delta_rule + + +def fused_recurrent_rwkv7( + r: torch.Tensor, + w: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + scale: float = 1.0, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False +): + """ + Args: + r (torch.Tensor): + r of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + w (torch.Tensor): + log decay of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + k (torch.Tensor): + k of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + v (torch.Tensor): + v of shape `[B, H, T, V]` if `head_first=True` else `[B, T, H, V]`. + a (torch.Tensor): + a of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + b (torch.Tensor): + b of shape `[B, H, T, K]` if `head_first=True` else `[B, T, H, K]`. + scale (float): + scale of the attention. + initial_state (torch.Tensor): + initial state of shape `[B, H, K, V]` if cu_seqlens is None else `[N, H, K, V]` where N = len(cu_seqlens) - 1. + output_final_state (bool): + whether to output the final state. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + head_first (bool): + whether to use head first. Recommended to be False to avoid extra transposes. + """ + return fused_recurrent_dplr_delta_rule( + q=r, + k=k, + v=v, + a=a, + b=b, + gk=w, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + head_first=head_first + ) diff --git a/fla/ops/simple_gla/README.md b/fla/ops/simple_gla/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2a64f3dcdee7ff9863089a6b47ef694f6234ab8f --- /dev/null +++ b/fla/ops/simple_gla/README.md @@ -0,0 +1,10 @@ +# Simple GLA + +Gating mechanism in [Gated RFA](https://arxiv.org/abs/2103.02143), [Mamba2](https://arxiv.org/abs/2405.21060) and [YOCO](https://arxiv.org/abs/2405.05254) (a.k.a., Gated RetNet). + +Compared to GLA, the gating is head-wise instead of elementwise. +As a result, we can adapt the RetNet kernel for training using matmul w/o numerical instability. +It is faster than GLA but has less expressive power. +I will use it as a baseline for the GLA. + +$S_{t+1} = g_{t+1} \odot S_{t} + K_{t+1} V_{t+1}^{\top}$ where $g$ is a scalar. diff --git a/fla/ops/simple_gla/__init__.py b/fla/ops/simple_gla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..561e3afbf81e8ab1b0fe738e5c5e8d1e1626868e --- /dev/null +++ b/fla/ops/simple_gla/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- + +from .chunk import chunk_simple_gla +from .fused_recurrent import fused_recurrent_simple_gla +from .parallel import parallel_simple_gla + +__all__ = [ + 'chunk_simple_gla', + 'fused_recurrent_simple_gla', + 'parallel_simple_gla' +] diff --git a/fla/ops/simple_gla/naive.py b/fla/ops/simple_gla/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..5fcc96ebeb720cc8b9699793ee6bdf8d3d39fdaa --- /dev/null +++ b/fla/ops/simple_gla/naive.py @@ -0,0 +1,54 @@ +# -*- coding: utf-8 -*- + +import torch +from einops import rearrange + + +def torch_simple_gla(q, k, v, g, chunk_size=64, scale=None): + if scale is None: + scale = (q.shape[-1] ** -0.5) + q = rearrange(q, 'b h (n c) d -> b h n c d', c=chunk_size) * scale + k = rearrange(k, 'b h (n c) d -> b h n c d', c=chunk_size) + v = rearrange(v, 'b h (n c) d -> b h n c d', c=chunk_size) + g = rearrange(g, 'b h (n c) -> b h n c', c=chunk_size) + g = g.cumsum(-1) + kv = k.transpose(-1, -2) @ (v * (-g + g[:, :, :, -1, None]).exp()[..., None]) + S = torch.zeros_like(kv) + + for i in range(1, g.shape[-2]): + S[:, :, i] = S[:, :, i-1].clone() * g[:, :, i-1, -1, None, None].exp() + kv[:, :, i-1] + + inter = (q * g[..., None].exp()) @ S + attn = q @ k.transpose(-1, -2) + attn = attn * (g[..., None] - g[..., None, :]).exp() + attn = attn.masked_fill(torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), 0) + intra = attn @ v + o = inter + intra + return rearrange(o, 'b h n c d -> b h (n c) d') + + +def torch_simple_gla_recurrent(q, k, v, g, scale=None, initial_state=None, output_final_state=True): + B, H, T, DK = q.shape + original_dtype = q.dtype + q, k, v, g = q.float(), k.float(), v.float(), g.float() + if scale is None: + scale = DK ** -0.5 + q = q * scale + _, _, _, DV = v.shape + if initial_state is None: + S = torch.zeros(B, H, DK, DV) + else: + S = initial_state + o = torch.zeros(B, H, T, DV).to(q) + for i in range(T): + gate = g[:, :, i].exp() + key = k[:, :, i] + value = v[:, :, i] + kv = key.unsqueeze(-1) * value.unsqueeze(-2) + S = S.clone() * gate.unsqueeze(-1).unsqueeze(-1) + kv + q_i = q[:, :, i, :] + o_i = (q_i.unsqueeze(-1) * S).sum(-2) + o[:, :, i] = o_i + if not output_final_state: + S = None + return o.to(original_dtype), S diff --git a/fla/ops/utils/asm.py b/fla/ops/utils/asm.py new file mode 100644 index 0000000000000000000000000000000000000000..c4a96bad2cecf24733832b6817f8d4b855685f05 --- /dev/null +++ b/fla/ops/utils/asm.py @@ -0,0 +1,17 @@ +# -*- coding: utf-8 -*- + +from fla.utils import device_platform + + +def fp32_to_tf32_asm() -> str: + """ + Get the assembly code for converting FP32 to TF32. + """ + ASM_DICT = { + 'nvidia': 'cvt.rna.tf32.f32 $0, $1;' + } + if device_platform in ASM_DICT: + return ASM_DICT[device_platform] + else: + # return empty string if the device is not supported + return "" diff --git a/fla/ops/utils/logcumsumexp.py b/fla/ops/utils/logcumsumexp.py new file mode 100644 index 0000000000000000000000000000000000000000..7476d3220599aa78e8f8ae5b10d0d15297cc47b4 --- /dev/null +++ b/fla/ops/utils/logcumsumexp.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import triton +import triton.language as tl + +from fla.ops.utils.op import exp, log + + +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for BT in [16, 32, 64] + for num_warps in [2, 4, 8] + ], + key=['S'] +) +@triton.jit(do_not_specialize=['T']) +def logcumsumexp_fwd_kernel( + s, + z, + T, + S: tl.constexpr, + BT: tl.constexpr +): + i_bh = tl.program_id(0) + o_i = tl.arange(0, BT) + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) + + b_mp = tl.full([S,], float('-inf'), dtype=tl.float32) + b_zp = tl.zeros([S,], dtype=tl.float32) + for i_t in range(tl.cdiv(T, BT)): + p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) + p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) + + # [BT, S] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + # [S,] + b_mc = tl.max(b_s, 0) + b_mc = tl.maximum(b_mp, b_mc) + b_zp = b_zp * exp(b_mp - b_mc) + # [BT, S] + b_s = exp(b_s - b_mc) + b_z = tl.dot(m_s, b_s, allow_tf32=False) + b_zp + # [S,] + b_zc = tl.max(b_z, 0) + b_mp = b_mc + b_zp = b_zc + # [BT, BS] + # small eps to prevent underflows + b_z = log(tl.where(b_z != 0, b_z, 1e-20)) + b_mc + tl.store(p_z, b_z.to(p_z.dtype.element_ty), boundary_check=(0, 1)) diff --git a/fla/ops/utils/testing.py b/fla/ops/utils/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..6f4fb01202e6bfbea3351defe8a424ca648ce7d1 --- /dev/null +++ b/fla/ops/utils/testing.py @@ -0,0 +1,26 @@ +import os + +compiled_mode = os.getenv("COMPILER_MODE") == "1" +ci_env = os.getenv("CI_ENV") == "1" + + +def get_abs_err(x, y): + return (x.detach()-y.detach()).flatten().abs().max().item() + + +def get_err_ratio(x, y): + err = (x-y).flatten().square().mean().sqrt().item() + base = (x).flatten().square().mean().sqrt().item() + return err / (base + 1e-15) + + +def assert_close(prefix, ref, tri, ratio, warning=False): + msg = f"{prefix} diff: {get_abs_err(ref, tri):.6f} ratio: {get_err_ratio(ref, tri):.6f}" + print(msg) + error_rate = get_err_ratio(ref, tri) + if warning or str(prefix).strip().lower() == "dh0" or (ci_env and error_rate < 0.01): + if error_rate > ratio: + import warnings + warnings.warn(msg) + else: + assert error_rate < ratio, msg diff --git a/flame/components/__init__.py b/flame/components/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/flame/components/checkpoint.py b/flame/components/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..d7a21fbbcaf9503c1f0f8d965acce420b223201b --- /dev/null +++ b/flame/components/checkpoint.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass, field +from datetime import timedelta +from io import BytesIO +from typing import Any, Dict, List + +import torch +from torch.distributed.checkpoint.stateful import Stateful + + +@dataclass +class TrainState(Stateful): + step: int = 0 + skipped_step: int = 0 + token: int = 0 + elapsed: timedelta = timedelta(0) + global_avg_losses: List[float] = field(default_factory=list) + global_max_losses: List[float] = field(default_factory=list) + log_steps: List[int] = field(default_factory=list) + + def state_dict(self) -> Dict[str, Any]: + # Only checkpoint global_avg_losses and global_max_losses per log frequency + # to avoid sync overhead in every iteration. + global_avg_losses_bytes = BytesIO() + torch.save(self.global_avg_losses, global_avg_losses_bytes) + global_max_losses_bytes = BytesIO() + torch.save(self.global_max_losses, global_max_losses_bytes) + log_steps_bytes = BytesIO() + torch.save(self.log_steps, log_steps_bytes) + return { + "step": torch.tensor(self.step, dtype=torch.int32), + "skipped_step": torch.tensor(self.skipped_step, dtype=torch.int32), + "token": torch.tensor(self.token, dtype=torch.int64), + "elapsed": self.elapsed, + "global_avg_losses": global_avg_losses_bytes, + "global_max_losses": global_max_losses_bytes, + "log_steps": log_steps_bytes, + } + + def load_state_dict(self, state_dict) -> None: + self.step = state_dict["step"].item() + self.skipped_step = state_dict.get("skipped_step", 0).item() + self.token = state_dict["token"].item() + self.elapsed = state_dict["elapsed"] + state_dict["global_avg_losses"].seek(0) + self.global_avg_losses = torch.load( + state_dict["global_avg_losses"], weights_only=False + ) + state_dict["global_max_losses"].seek(0) + self.global_max_losses = torch.load( + state_dict["global_max_losses"], weights_only=False + ) + state_dict["log_steps"].seek(0) + self.log_steps = torch.load(state_dict["log_steps"], weights_only=False) diff --git a/flame/models/activation_offloading.py b/flame/models/activation_offloading.py new file mode 100644 index 0000000000000000000000000000000000000000..80012c714ab205dd5529176c3c2e2ab18263ad63 --- /dev/null +++ b/flame/models/activation_offloading.py @@ -0,0 +1,447 @@ +# Adapted from https://github.com/pytorch/torchtune/blob/main/torchtune/training/_activation_offloading.py +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import contextlib +from typing import Union +from warnings import warn + +import psutil +import torch +from torch import nn +from torch.autograd.graph import saved_tensors_hooks + +from torchtitan.tools.logging import logger + +try: + import torchao + from torchao.dtypes.nf4tensor import NF4Tensor +except ImportError: + torchao = None + NF4Tensor = None + logger.warning("torchao not found. ") + +# from torchtune.modules import TiedLinear + + +class OffloadActivations(saved_tensors_hooks): + """Context manager under which activation tensors created in the forward pass will be offloaded. + + Enable the memory efficiency technique of activation offloading, where activations bigger than + min_offload_size bytes will be offloaded to CPU in the forward and brought back in the backward. + This is in contrast to maintaining the activation on GPU VRAM throughout the program. + + This manager contains the option of using one additional CUDA stream to handle the communication + between CUDA and CPU, which is intended to overlap with the default computation stream to improve + runtime. We designed synchronization with a few heuristics for optimizing the tradeoff between + runtime vs memory usage. + + Args: + use_pin_memory (bool): Whether or not the offloaded Tensor will be placed in pinned + memory on the CPU. Pinned memory allows the Tensor to be moved back onto GPU more quickly + but is a limited resource. Default: True. + + use_streams (bool): Whether or not to use streams for performance optimization where + the communications get overlapped with the computation. Requires a torch build + after torch-2.5.0.]. Default: True. + + max_fwd_stash_size (int): The maximum size of the forward stash, or the maximum number of + consecutive activations to keep alive during the forward pass. This number must be at + least 1. Keeping alive more activations will potentially allow more overlap between the + communication and compute streams at the cost of increasing memory usage. Keeping alive + fewer activations will conserve memory, but may cause poor overlap between the streams, + increasing runtime. Default: 5. + + min_offload_size (int): The minimum number of bytes a Tensor must be in order to qualify + for offloading. If the tensor is too small, we do not want to waste bandwidth and resources + moving it to CPU and back. Default: 1024 bytes. + + Raises: + ValueError: if max_fwd_stash_size is not at least 1. + + Example: + >>> with OffloadActivations(): + >>> logits = model(inputs) + >>> loss = ... + >>> loss.backward() + """ + + def __init__( + self, + use_pin_memory: bool = True, + use_streams: bool = True, + max_fwd_stash_size: int = 5, + min_offload_size: int = 1024, + ) -> None: + + self.use_streams: bool = use_streams + + self.min_tensor_size_bytes = ( + min_offload_size # we don't want to bother with small tensors + ) + self.tracker = ( + {} + ) # tensor_id => (new_tensor, if_modified) ---> track what saved/offloaded tensors are where + self.tensor_id: int = 0 + self.is_first_forward_call = True + self.is_first_backward_call = True + self.is_first_forward_pass = True + + # managing cpu memory + self.use_pin_memory: bool = use_pin_memory + self.virtual_memory_safe_pct = ( + 60 # we should not exceed this percentage of memory + ) + + self.s0 = torch.cuda.default_stream() # comp stream + + # for streaming + if self.use_streams: + self.s1 = torch.cuda.Stream() # comms stream + self.fwd_stash = {} # tensor_id => (activation, ev1) + if max_fwd_stash_size < 1: + raise ValueError( + f"max_fwd_stash_size should be at least 1 but is {max_fwd_stash_size}" + ) + self.max_fwd_stash_size = max_fwd_stash_size + self.bwd_tensor_stash = {} # tensor_id => activation + self.bwd_ev_stash = {} # tensor_id => ev0 + self.curr_graph_id = None + self.curr_autograd_node = None + + # -------- platform util functions -------- # + def verify_sufficient_virtual_memory(): + curr_pct = get_cpu_ram_pct() + if curr_pct > self.virtual_memory_safe_pct: + warn( + f"***** WARNING: {curr_pct=}% > {self.virtual_memory_safe_pct=}% of virtual memory used" + ) + + def get_cpu_ram_pct() -> float: + # get the percentage of memory used by the system + return psutil.virtual_memory().percent + + def get_tensor_id() -> int: + # create a unique id for each tensor we are managing + self.tensor_id += 1 + return self.tensor_id + + def get_num_bytes_tensor(x: torch.Tensor) -> int: + # get the number of bytes in a tensor, for memory management purposes + return ( + x.element_size() * x.nelement() + ) # x.element_size() * x._base_storage().nbytes() + + # -------- core pack / unpack work -------- # + def pack_tensor(activation: torch.Tensor) -> int: + # activations are passed in during forward pass - from here we take over and return a unique id + if self.is_first_forward_call: + assert ( + len(self.tracker) == 0 + ), "backward pass should have cleared tracker of all tensors" + + # set training phase trackers + self.is_first_forward_call = False + self.is_first_backward_call = True + + # query for basic tensor info + num_bytes = get_num_bytes_tensor(activation) + tensor_id = get_tensor_id() + + # only offload hefty bois if they're activations on CUDA (our heuristic + # for that is to check if they're not params or buffers)! + if ( + activation.is_cuda + and num_bytes >= self.min_tensor_size_bytes + and ( + not isinstance(activation, torch.nn.Parameter) + and not isinstance(activation, torch.nn.Buffer) + ) + ): + if self.use_streams: + # First, sync back and dereference previously offloaded tensors + # as the offloading should be done sufficiently long ago. + for id in [k for k in self.fwd_stash.keys()]: + if id <= tensor_id - self.max_fwd_stash_size: + _, ev = self.fwd_stash[id] + self.s0.wait_event(ev) + del self.fwd_stash[id] + else: + break + + # Sync in, offload, and add an event to sync back later + self.s1.wait_stream(self.s0) + + stream = self.s1 if self.use_streams else self.s0 + with torch.cuda.stream(stream): + try: + cpu_tensor = torch.empty_like( + activation, pin_memory=self.use_pin_memory, device="cpu" + ) + except NotImplementedError as e: + if ( + isinstance(activation, NF4Tensor) + and torchao.__version__ < "0.6.0.dev20240917" + ): + raise RuntimeError( + "Offloading NF4Tensors requires torchao-0.6.0.dev20240917 or later" + ) from e + raise e + cpu_tensor.copy_(activation, non_blocking=True) + self.tracker[tensor_id] = ( + cpu_tensor, + True, + ) # True = (in future) modified + + if self.use_streams: + event = self.s1.record_event() + + # Stash to keep activation alive til s1 is done + self.fwd_stash[tensor_id] = (activation, event) + else: + self.tracker[tensor_id] = ( + activation, + False, + ) # False = not modified, tensor is as is + + return tensor_id + + def unpack_tensor_single_stream(unpack_tensor_id: int) -> torch.Tensor: + # backward pass - we are called with the tensor_id, which + # we will use to retrieve the saved/offloaded tensor + if self.is_first_backward_call: + if self.is_first_forward_pass: + self.is_first_forward_pass = False + if self.use_pin_memory: + verify_sufficient_virtual_memory() + + self.is_first_backward_call = False + self.is_first_forward_call = True + + assert ( + unpack_tensor_id in self.tracker + ), f"untracked tensor with id {unpack_tensor_id}" + + maybe_gpu_tensor, modified = self.tracker[unpack_tensor_id] + if modified: + gpu_tensor = maybe_gpu_tensor.to("cuda", non_blocking=True) + maybe_gpu_tensor = gpu_tensor + + # clear tensor from tracking + del self.tracker[unpack_tensor_id] + return maybe_gpu_tensor + + def unpack_tensor_with_streams(unpack_tensor_id: int) -> torch.Tensor: + # backward pass - we are called with the tensor_id, which + # we will use to retrieve the saved/offloaded tensor + if self.is_first_backward_call: + self.curr_graph_id = torch._C._current_graph_task_id() + + def wait_and_del_remaining_references() -> None: + for id in [k for k in self.bwd_tensor_stash.keys()]: + event = self.bwd_ev_stash[id] + self.s1.wait_event(event) + del self.bwd_tensor_stash[id] + + # Register a callback to the end of autograd to clean everything up + torch.autograd.variable.Variable._execution_engine.queue_callback( + wait_and_del_remaining_references + ) + + if self.is_first_forward_pass: + self.is_first_forward_pass = False + if self.use_pin_memory: + verify_sufficient_virtual_memory() + + self.is_first_backward_call = False + self.is_first_forward_call = True + + assert ( + unpack_tensor_id in self.tracker + ), f"untracked tensor with id {unpack_tensor_id}" + + maybe_gpu_tensor, modified = self.tracker[unpack_tensor_id] + if modified: + # Get data on the current autograd node + graph_id = torch._C._current_graph_task_id() + node = torch._C._current_autograd_node() + prev_node_ids = [] + + # If we're on a new node, mark prev node's tensors to be freed later + if graph_id == self.curr_graph_id and self.curr_autograd_node != node: + self.curr_autograd_node = node + prev_node_ids = [id for id in self.bwd_tensor_stash.keys()] + + brought_back_from_cpu = True + if unpack_tensor_id in self.fwd_stash: + maybe_gpu_tensor = self.fwd_stash[unpack_tensor_id][0] + brought_back_from_cpu = False + else: + # Kick off the process to bring tensors back + with torch.cuda.stream(self.s1): + gpu_tensor = maybe_gpu_tensor.to("cuda", non_blocking=True) + maybe_gpu_tensor = gpu_tensor + + # Tell comp stream to wait for the info to be loaded before executing + self.s0.wait_stream(self.s1) + + # Stash the tensor to keep memory alive until compute stream is complete + self.bwd_tensor_stash[unpack_tensor_id] = maybe_gpu_tensor + + # Note: [Track views of the unpacked] + # Why do we get the use count of the unpacked tensor here? We want an + # initial count to compare to later, during the post-hook of the + # backward node, when we need to decide whether we're allowed to free + # the tensor yet. In what obscure cases must we delay freeing the + # tensor (and thus call record_stream)? + # 1. Any of the outputs of the backward node is a view of the unpacked + # tensor. + # 2. In the case that this unpacked tensor will be used in a + # checkpointed region, if one of the recomputed saved tensors ends + # up as a view of the unpacked tensor. + # 3. The user abuses the system somehow and manually relies on the + # unpacked tensor to exist after the backward node has executed. + storage_refcount = torch._C._storage_Use_Count( + maybe_gpu_tensor.untyped_storage()._cdata + ) + + def hook(outputs, inputs): + # create events for the current node inputs/outputs if they were streamed in + if brought_back_from_cpu: + # See Note: [Track views of the unpacked] + # IF any of the outputs is a view of the tensor, OR if a view of + # the tensor has been saved as a part of checkpoint's recompute + # process, OR the user has abusedly incurred a reference on the + # unpacked tensor, THEN the tensor might be used later and we + # cannot presume to delete it after only the current node is + # done! So we use our frenemy, record_stream, to ensure the + # Tensor stays unmessed with until it's done getting used in the + # compute stream (s0 here). Note that the con here is we introduce + # non-deterministic (thus higher) memory usage, but this case + # should not happen often. + unpacked_tensor = self.bwd_tensor_stash[unpack_tensor_id] + if ( + torch._C._storage_Use_Count( + unpacked_tensor.untyped_storage()._cdata + ) + > storage_refcount + ): + unpacked_tensor.record_stream(self.s0) + del self.bwd_tensor_stash[unpack_tensor_id] + else: + event = self.s0.record_event() + self.bwd_ev_stash[unpack_tensor_id] = event + + # if there are still things in the fwd_stash, get rid of them as we're in bwd now + for id in [k for k in self.fwd_stash.keys()]: + _, ev = self.fwd_stash[id] + self.s0.wait_event(ev) + del self.fwd_stash[id] + + # wait on prev node's events and del those + for id in prev_node_ids: + event = self.bwd_ev_stash[id] + self.s1.wait_event(event) + del self.bwd_tensor_stash[id] + + return outputs + + node.register_hook(hook) + + # clear tensor from tracking + del self.tracker[unpack_tensor_id] + return maybe_gpu_tensor + + unpack_tensor = ( + unpack_tensor_with_streams + if self.use_streams + else unpack_tensor_single_stream + ) + super().__init__(pack_tensor, unpack_tensor) + + +class NoOpManager(saved_tensors_hooks): + """ + A saved_tensors_hook manager used to disable any other saved_tensors_hook manager + applied before. This relies on the behavior that only the most recently registered + saved_tensors_hook will run. + + One example usage is to opt a local region of code out of activations offloading, + which is usually applied globally to best track state. + """ + + def __init__(self) -> None: + def noop(tensor): + return tensor + + super().__init__(noop, noop) + + +def get_act_offloading_ctx_manager( + model: nn.Module, enable_activation_offloading: bool +) -> Union[OffloadActivations, contextlib.nullcontext]: + """Returns the activation offloading context manager for the model, which will be + a null context if enable_activation_offloading is False. + + If activation offloading is enabled, we return the OffloadActivations context manager. + If activation offloading is disabled, we return a NoOpManager context manager. + + Args: + model (nn.Module): the model to wrap with the activation offloading context manager. + enable_activation_offloading (bool): whether or not to enable activation offloading + for the model. + + Returns: + contextlib.ContextDecorator: the activation offloading context manager for the model. + + Raises: + NotImplementedError: If the model is a multimodal model and activation offloading is enabled. + """ + if enable_activation_offloading: + activations_handling_ctx = OffloadActivations() + + # Below is our hack to disable offloading the last output Linear in every + # step, as the cost for offloading the activation and then soon after bringing + # it back is expensive. Moreover, due to heuristics in our streaming API, + # we actually use more memory if we offload it as it interferes with chunkedCE. + output_head_detected = False + noop_ctx = NoOpManager() + + if hasattr(model, "output"): + if isinstance(model.output, nn.Module): + model.output.register_forward_pre_hook( + lambda *args: noop_ctx.__enter__() + ) + model.output.register_forward_hook( + lambda *args: noop_ctx.__exit__(), always_call=True + ) + print("registering hooks for model.output ============ ") + output_head_detected = True + # ================================ + # ! TODO[flame] check if we need to detal with TiedLinear + # The following code appears in `torchtune` + # elif isinstance(model.output, TiedLinear): + # model.output.linear.register_forward_pre_hook( + # lambda *args: noop_ctx.__enter__() + # ) + # model.output.linear.register_forward_hook( + # lambda *args: noop_ctx.__exit__(), always_call=True + # ) + # output_head_detected = True + + if not output_head_detected: + logger.warning( + "During activation offloading, no output head was detected. " + "If your model has an output head, it will be offloaded. " + "This usually greatly slows training, given the large vocabulary size. " + "To change this behavior, set your output head as model.output and make it " + "an nn.Module." + ) + + else: + activations_handling_ctx = contextlib.nullcontext() + + return activations_handling_ctx diff --git a/flame/models/fla.toml b/flame/models/fla.toml new file mode 100644 index 0000000000000000000000000000000000000000..bbc11346144d0bfa3d8c6cf5b17a03208028ea46 --- /dev/null +++ b/flame/models/fla.toml @@ -0,0 +1,68 @@ +[model] +config = "fla-hub/transformer-1.3B-100B" +tokenizer_path = "fla-hub/transformer-1.3B-100B" + +[job] +dump_folder = "exp" +print_args = true + +[training] +batch_size = 32 +seq_len = 2048 +context_len = 2048 +gradient_accumulation_steps = 1 +steps = 20480 +max_norm = 1.0 +skip_nan_inf = true +data_parallel_replicate_degree = 1 +data_parallel_shard_degree = -1 +tensor_parallel_degree = 1 +compile = false +dataset = "HuggingFaceFW/fineweb-edu" +dataset_name = "default" +num_workers = 32 +pin_memory = false +persistent_workers = false +prefetch_factor = 2 +seed = 42 +varlen = false +dataset_mode = "pretrain" + +[optimizer] +name = "AdamW" +eps = 1e-15 +lr = 3e-4 + +[lr_scheduler] +warmup_steps = 1024 +decay_type = "cosine" +lr_min = 0.1 + +[checkpoint] +enable_checkpoint = true +folder = "checkpoint" +interval_type = "steps" +interval = 2048 +model_weights_only = false +export_dtype = "float32" +async_mode = "disabled" # ["disabled", "async", "async_with_pinned_mem"] + +[profiling] +enable_profiling = true +save_traces_folder = "profile_trace" +profile_freq = 512 + +[metrics] +log_freq = 32 +enable_wandb = true + +[experimental] +context_parallel_degree = 1 +pipeline_parallel_degree = 1 + +[float8] +enable_fsdp_float8_all_gather = false +precompute_float8_dynamic_scale_for_fsdp = false + +[activation_checkpoint] +mode = "none" diff --git a/flame/utils/__init__.py b/flame/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/flame/utils/__pycache__/__init__.cpython-312.pyc b/flame/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..babff5aed6ba1370459571da0ba606508254a944 Binary files /dev/null and b/flame/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/flame/utils/__pycache__/convert_hf_to_dcp.cpython-312.pyc b/flame/utils/__pycache__/convert_hf_to_dcp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cedb437b7f4d5f783345a207a37f87923fcefd8d Binary files /dev/null and b/flame/utils/__pycache__/convert_hf_to_dcp.cpython-312.pyc differ diff --git a/flame/utils/convert_dcp_to_hf.py b/flame/utils/convert_dcp_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..2240578d54b6be51241ddad0e253c548cb362492 --- /dev/null +++ b/flame/utils/convert_dcp_to_hf.py @@ -0,0 +1,66 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import argparse +import io +import os +import tempfile +from datetime import timedelta + +import torch +import torch.serialization +from torch.distributed.checkpoint.format_utils import dcp_to_torch_save +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import fla # noqa +from torchtitan.tools.logging import init_logger, logger + + +@torch.inference_mode() +def save_pretrained( + path: str, + step: int, + config: str, + tokenizer: str +): + logger.info(f"Loading the config from {config}") + config = AutoConfig.from_pretrained(config, trust_remote_code=True) + + logger.info(f"Saving the config to {path}") + config.save_pretrained(path) + logger.info(f"Loading the tokenizer from {tokenizer}") + tokenizer = AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=True) + logger.info(f"Saving the tokenizer to {path}") + tokenizer.save_pretrained(path) + + with tempfile.TemporaryDirectory() as tmpdir: + # base_checkpoint_dir = os.path.dirname(path) + base_checkpoint_dir = path + checkpoint = os.path.join(base_checkpoint_dir, f'checkpoint/step-{step}') + checkpoint_path = os.path.join(tmpdir, 'checkpoint.pt') + logger.info(f"Saving the distributed checkpoint to {checkpoint_path}") + dcp_to_torch_save(checkpoint, checkpoint_path) + + logger.info(f"Initializing the model from config\n{config}") + model = AutoModelForCausalLM.from_config(config) + logger.info(model) + logger.info("Loading state dict from the checkpoint") + + # Add datetime.timedelta and io.BytesIO to safe globals + torch.serialization.add_safe_globals([timedelta, io.BytesIO]) + # torch.load now with default weights_only=True will work + model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model']) + + logger.info(f"Saving the model to {path}") + model.save_pretrained(path) + + +if __name__ == "__main__": + init_logger() + parser = argparse.ArgumentParser("Convert DCP format model weights to huggingface-style.") + parser.add_argument("--path", type=str, required=True) + parser.add_argument("--step", type=int, required=True) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--tokenizer", type=str, required=True) + args = parser.parse_args() + save_pretrained(args.path, args.step, args.config, args.tokenizer) diff --git a/model-00001-of-00006.safetensors b/model-00001-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..00d13e3ddbaceeecffc82016a2953164738d810c --- /dev/null +++ b/model-00001-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a571148850e1f6648c9d10afe1b2fbfe83929dbc0f48f404654be9f3abf134a +size 4987213192 diff --git a/model-00002-of-00006.safetensors b/model-00002-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b53ad4bec7c14a61d7d7c1b6764297394c709516 --- /dev/null +++ b/model-00002-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b8e20408833bcde4d5c5cbf01b75bd3967a656c1e3f5fce04153ad1fba7dbf05 +size 4899132712 diff --git a/model-00003-of-00006.safetensors b/model-00003-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a5931c69b1ff418121db6abee28b9a6e4c1d6316 --- /dev/null +++ b/model-00003-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fba920cb6746046a0dad5ab159e4448969fb3b5ab8b6a483a477ebc407c34a30 +size 4999812880 diff --git a/model-00004-of-00006.safetensors b/model-00004-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8f2f5cb511c8cc68e4d0482808b2a2ab66db03c9 --- /dev/null +++ b/model-00004-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c73c5e3a28a7cc8d6d6d20383aae47960bdad7899d35b3818a384b5b07f105f +size 4999812880 diff --git a/model-00005-of-00006.safetensors b/model-00005-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..207f1b672e2be2bf428c64373f2836c367000ed6 --- /dev/null +++ b/model-00005-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5914c27f39f3042eb82108e04463e0ef14ee09ef4ea49bc2f14f59b83d4f2c1d +size 4832007296 diff --git a/model-00006-of-00006.safetensors b/model-00006-of-00006.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..00a333d926c4cec5a2428d9384c01a30ecd8a65f --- /dev/null +++ b/model-00006-of-00006.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e827fd7dc26508f18129998b03925e7c3c4a7c814b691adcd617bfe0fb754129 +size 3028371824 diff --git a/tb/20260114-0839/wandb/run-20260114_083920--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834/run--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834.wandb b/tb/20260114-0839/wandb/run-20260114_083920--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834/run--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834.wandb new file mode 100644 index 0000000000000000000000000000000000000000..5a46b21b13eaf01531efa6671accb23b9ea13c3a --- /dev/null +++ b/tb/20260114-0839/wandb/run-20260114_083920--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834/run--top.code.7B.batch8.seqlen4096.context4096.warmup400.update2.steps40000.lr2e-5.cosine-202601140834.wandb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12e2a66decd80b8ffaba806a9c3649430a67c9d706a1c1eaae249ea13c2c7ad3 +size 254502683 diff --git a/torchtitan/components/__pycache__/float8.cpython-312.pyc b/torchtitan/components/__pycache__/float8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36d492432a5f0082d56f9d505435d97230e7d935 Binary files /dev/null and b/torchtitan/components/__pycache__/float8.cpython-312.pyc differ diff --git a/torchtitan/components/__pycache__/metrics.cpython-312.pyc b/torchtitan/components/__pycache__/metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d1f23fa02a0905849360203836d7f6f20e9f9f7 Binary files /dev/null and b/torchtitan/components/__pycache__/metrics.cpython-312.pyc differ diff --git a/torchtitan/experiments/deepseek_v3/LICENSE-CODE b/torchtitan/experiments/deepseek_v3/LICENSE-CODE new file mode 100644 index 0000000000000000000000000000000000000000..d84f527e101b2cdd171e2b14253f84ea4fedabe9 --- /dev/null +++ b/torchtitan/experiments/deepseek_v3/LICENSE-CODE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/torchtitan/experiments/deepseek_v3/checkpoint.py b/torchtitan/experiments/deepseek_v3/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..535ac7fe069a88555841181dddc1e870c2d30934 --- /dev/null +++ b/torchtitan/experiments/deepseek_v3/checkpoint.py @@ -0,0 +1,154 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import json +import logging +import os +from typing import Dict, Optional, Set, Tuple + +import torch +from safetensors import safe_open + +from transformers.utils import cached_file + + +logger = logging.getLogger(__name__) + +_DEFAULT_SAFETENSOR_FILE_NAME = "model.safetensors.index.json" + + +def read_weights_from_json(file_path: str) -> Optional[Dict[str, str]]: + try: + with open(file_path, "r") as file: + data = json.load(file) + + if "weight_map" in data and isinstance(data["weight_map"], dict): + return data["weight_map"] + else: + logger.info("No 'weight_map' dictionary found in the JSON file.") + return None + except (json.JSONDecodeError, Exception) as e: + logger.info(f"An error occurred while reading the JSON file: {str(e)}") + return None + + +def get_hf_weight_map_and_path( + model_id: str, +) -> Tuple[Dict[str, str], str]: + """Get the weight map for a given HF model id and also the cache path for loading the weights""" + try: + index_file = cached_file(model_id, _DEFAULT_SAFETENSOR_FILE_NAME) + except Exception as e: + logger.error( + f"Model `{model_id}` not found in HF cache. " + f"You can download the model using `python download.py {model_id}" + ) + raise e + + weight_map = read_weights_from_json(index_file) + weight_path = os.path.dirname(index_file) + logger.info(f"Loading weights from: {weight_path}") + return weight_map, weight_path + + +def get_needed_files( + state_dict: Dict[str, torch.Tensor], weight_map: Dict[str, str] +) -> Set[str]: + needed_files = set() + for param in state_dict.keys(): + file = weight_map.get(param) + if file: + needed_files.add(file) + elif param.endswith("weight"): + raise ValueError( + f"Parameter {param} not found in weight map, please check..." + ) + logger.info(f"Needed files: {needed_files}") + return needed_files + + +def load_safetensor_file( + full_path: str, device: torch.device +) -> Dict[str, torch.Tensor]: + tensors = {} + with safe_open(full_path, framework="pt", device=device) as f: + for k in f.keys(): + tensors[k] = f.get_tensor(k) + logger.info(f"Loaded {len(tensors)} tensors from {full_path}") + return tensors + + +def load_safetensor_weights( + model: torch.nn.Module, + weight_map: Dict[str, str], + file_location: str, + device: torch.device, +): + """ + Load safetensor weights into a `nn.Module`. + + Args: + model (Module): The PyTorch module to load weights into. It may be a + model chunk or a full model. + weight_map (Dict[str, str]): Mapping of model parameters to file names. + file_location (str): Directory containing the weight files. + device (torch.device): The device to load tensors onto. + """ + model_state_dict = model.state_dict() + needed_files = get_needed_files(model_state_dict, weight_map) + updated_states: Set[str] = set() + + for file in needed_files: + full_path = os.path.join(file_location, file) + try: + checkpoint = load_safetensor_file(full_path, "cpu") + except FileNotFoundError: + logger.error(f"File not found: {full_path}") + except Exception as e: + logger.error(f"Error during checkpoint processing of {full_path}: {str(e)}") + + matched_keys = set(checkpoint.keys()) & set(model_state_dict.keys()) + for key in matched_keys: + # Check shape + if model_state_dict[key].shape != checkpoint[key].shape: + raise ValueError( + f"Shape mismatch for {key}: " + f"model needs {model_state_dict[key].shape}, but " + f"checkpoint has {checkpoint[key].shape}" + ) + model_state_dict[key] = checkpoint[key].to(device) + + updated_states.update(matched_keys) + + missing_keys = set(model_state_dict.keys()) - updated_states + if missing_keys: + raise RuntimeError( + f"Partially updated state dict. Missing parameters: {missing_keys}" + ) + + model.load_state_dict(model_state_dict, strict=False, assign=True) + logger.info(f"Successfully loaded {len(updated_states)} weights into model") + + +def load_weights_from_hf( + model: torch.nn.Module, + distribution: str, + device: torch.device, +): + """ + Load the weights from Hugging Face format (index file + multiple safetensor + files), and fill into `model`. Model config is needed b/c we permute + wq and wk weights based on attn heads. + """ + + weight_map, weight_path = get_hf_weight_map_and_path(distribution) + + load_safetensor_weights( + model, + weight_map, + weight_path, + device, + ) diff --git a/torchtitan/experiments/deepseek_v3/generate.py b/torchtitan/experiments/deepseek_v3/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..6d7302f1078ec42d7a30dd7d92e113c9affdd650 --- /dev/null +++ b/torchtitan/experiments/deepseek_v3/generate.py @@ -0,0 +1,308 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# torchrun --standalone --nproc-per-node 4 generate.py + +# use inference.sh "Your Question Here?" to run inference with a single prompt. + +import sys +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from checkpoint import load_weights_from_hf +from model import DeepseekForCausalLM +from model_config import deepseek_config_registry +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.pipelining import PipelineStage, ScheduleGPipe +from torchtitan.tools.utils import Color +from transformers import AutoTokenizer + +# Uncomment the model you want to run. +model_id, mesh_shape = "deepseek-ai/DeepSeek-V2-Lite-Chat", (1, 4) +# model_id, mesh_shape = "deepseek-ai/deepseek-v3", (8, 4) + + +def colorize_chat(text, user_color=None, assistant_color=None, output_color=None): + """Parse and colorize chat output with optional colors for each role.""" + lines = text.split("\n") + result = [] + + current_role = None + current_content = [] + + def _process_current_content(): + if not current_role or not current_content: + return None + + content = "\n".join(current_content) + if current_role == "output": + return ( + f"Output: {output_color}{content}{color.reset}" + if output_color + else f"Output: {content}" + ) + else: + try: + prefix, rest = current_content[0].split(":", 1) + role_color = user_color if current_role == "user" else assistant_color + if role_color: + formatted = f"{prefix}:{role_color}{rest}{color.reset}" + if len(current_content) > 1: + formatted += ( + f"{role_color}\n" + + "\n".join(current_content[1:]) + + f"{color.reset}" + ) + return formatted + except ValueError: + pass + return content + + for line in lines: + if line.startswith("Output:"): + if processed := _process_current_content(): + result.append(processed) + current_role = "output" + content = line[len("Output:") :].strip() + if output_color: + content = f"Output: {output_color}{content}{color.reset}" + else: + content = f"Output: {content}" + result.append(content) + current_content = [] + + elif line.startswith("User:"): + if processed := _process_current_content(): + result.append(processed) + current_role = "user" + current_content = [line] + + elif line.startswith("Assistant:"): + if processed := _process_current_content(): + result.append(processed) + current_role = "assistant" + current_content = [line] + + else: + if current_content: + current_content.append(line) + elif line.strip() and current_role is None: + # Handle system message at the beginning + current_role = "output" + if output_color: + result.append(f"Output: {output_color}{line.strip()}{color.reset}") + else: + result.append(f"Output: {line.strip()}") + + # Process the last segment + if processed := _process_current_content(): + result.append(processed) + + return "\n".join(result) + + +color = Color() + + +@dataclass +class DistConfig: + mesh: DeviceMesh + pp_mesh: DeviceMesh + ep_mesh: DeviceMesh + pp_size: int + ep_size: int + ep_rank: int + pp_rank: int + device: torch.device + + +def create_model(dist_config: DistConfig): + model_args = deepseek_config_registry[model_id] + model_args.ep_size = dist_config.ep_size + model_args.num_stages = dist_config.pp_size + model_args.stage_idx = dist_config.pp_rank + model_args.max_seq_len = 16384 + + with dist_config.device, dist_config.mesh: + model = DeepseekForCausalLM(model_args) + load_weights_from_hf(model, model_id, dist_config.device) + model.eval() + model.setup_symm_mem(torch.bfloat16, dist_config.device) + + stage = PipelineStage( + model, + dist_config.pp_rank, + dist_config.pp_size, + dist_config.device, + group=dist_config.pp_mesh.get_group(), + ) + pp_schedule = ScheduleGPipe(stage, dist_config.pp_size) + return model, pp_schedule + + +def create_dist_config(mesh: DeviceMesh): + rank = dist.get_rank() + device_count = torch.cuda.device_count() + device = torch.device("cuda", rank % device_count) + + dist_config = DistConfig( + mesh=mesh, + pp_mesh=mesh["pp"], + ep_mesh=mesh["ep"], + pp_rank=mesh["pp"].get_local_rank(), + pp_size=mesh["pp"].size(), + ep_size=mesh["ep"].size(), + ep_rank=mesh["ep"].get_local_rank(), + device=device, + ) + return dist_config + + +def decode(tokenizer, x): + output = tokenizer.decode(x[0]) + # Clean up the output by removing special tokens + bos = tokenizer.bos_token + output = output.replace(bos, "") + # Truncate at end of sentence token + eos_token = tokenizer.eos_token + if eos_token and eos_token in output: + output = output.split(eos_token)[0] + colored_output = colorize_chat( + output, + user_color=color.green, + assistant_color=color.cyan, + output_color=color.blue, + ) + return colored_output + + +@torch.inference_mode() +def generate( + model, + pp_schedule, + tokenizer, + dist_config, + messages: list[dict], + n_tokens: int = 50, +): + rank = dist.get_rank() + device = dist_config.device + x = tokenizer.apply_chat_template( + [messages] * dist_config.pp_size, + add_generation_prompt=True, + return_tensors="pt", + ) + next_idx = x.shape[-1] + x = torch.cat([x, torch.zeros(x.shape[0], n_tokens, dtype=torch.int64)], dim=-1) + x = x.to(device) + + for _ in range(n_tokens): + if dist_config.pp_size > 1: + if dist_config.pp_rank == 0: + pp_schedule.step(x) + torch.distributed.broadcast( + x, + group=dist_config.pp_mesh.get_group(), + group_src=dist_config.pp_size - 1, + ) + elif dist_config.pp_rank == dist_config.pp_size - 1: + preds = pp_schedule.step() + next_token = torch.argmax(preds[:, next_idx - 1], dim=-1) + x[:, next_idx] = next_token + torch.distributed.broadcast( + x, + group=dist_config.pp_mesh.get_group(), + group_src=dist_config.pp_size - 1, + ) + else: + pp_schedule.step() + torch.distributed.broadcast( + x, + group=dist_config.pp_mesh.get_group(), + group_src=dist_config.pp_size - 1, + ) + + next_idx += 1 + else: + preds = model(x) + next_token = torch.argmax(preds[:, next_idx - 1], dim=-1) + x[:, next_idx] = next_token + next_idx += 1 + + if rank == 0: + colored_output = decode(tokenizer, x) + print(f"Without CUDA Graph:\n{colored_output}") + + +@torch.inference_mode() +def generate_with_cuda_graph( + model, + tokenizer, + dist_config, + messages: list[dict], + n_tokens: int = 10, +): + rank = dist.get_rank() + device = dist_config.device + x = tokenizer.apply_chat_template( + [messages] * dist_config.pp_size, + add_generation_prompt=True, + return_tensors="pt", + ) + next_idx = x.shape[-1] + x = torch.cat([x, torch.zeros(x.shape[0], n_tokens, dtype=torch.int64)], dim=-1) + x = x.to(device) + + torch.cuda.synchronize() + + # Create CUDA graph + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + preds = model(x) + + # Run CUDA graph + for _ in range(n_tokens): + g.replay() + next_token = torch.argmax(preds[:, next_idx - 1], dim=-1) + x[:, next_idx] = next_token + next_idx += 1 + + if rank == 0: + colored_output = decode(tokenizer, x) + print(f"With CUDA Graph:\n{colored_output}") + + +if __name__ == "__main__": + # Get user prompt from command line arguments + user_prompt = "What is 2+2?" # Default prompt + if len(sys.argv) > 1: + user_prompt = sys.argv[1] + + mesh = dist.init_device_mesh("cuda", mesh_shape, mesh_dim_names=("pp", "ep")) + rank = dist.get_rank() + if rank == 0: + print( + f"{color.yellow}Running inference with {model_id} on {mesh_shape} mesh{color.reset}" + ) + + dist_config = create_dist_config(mesh) + model, pp_schedule = create_model(dist_config) + tokenizer = AutoTokenizer.from_pretrained(model_id) + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": user_prompt}, + ] + + generate(model, pp_schedule, tokenizer, dist_config, messages) + generate_with_cuda_graph(model, tokenizer, dist_config, messages) + + if rank == 0: + print(f"\n{color.yellow}Closing inference mesh...{color.reset}") + + dist.destroy_process_group() diff --git a/torchtitan/experiments/kernels/triton_mg_group_gemm/torchao_pr/unit_test_backwards.py b/torchtitan/experiments/kernels/triton_mg_group_gemm/torchao_pr/unit_test_backwards.py new file mode 100644 index 0000000000000000000000000000000000000000..becb761d83bb8d55ef4cef95ee291262de6e4761 --- /dev/null +++ b/torchtitan/experiments/kernels/triton_mg_group_gemm/torchao_pr/unit_test_backwards.py @@ -0,0 +1,174 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe +import logging +import unittest +from typing import Tuple + +import torch +import torch.nn as nn + +from mg_grouped_gemm import ( + grouped_gemm_backward, + grouped_gemm_dw_tma, + grouped_gemm_dx_tma, + grouped_gemm_forward, + mg_grouped_gemm, +) + +from reference_utils import ( + analyze_tensor_differences, + compute_reference_backward, + compute_reference_forward, +) + + +class TestMG_GroupedGEMM_Backward(unittest.TestCase): + def setUp(self) -> None: + torch.manual_seed(2020) # Set seed for reproducibility + + def _run_grouped_gemm_backward_test( + self, + shape: Tuple[int, int, int, int], + device: torch.device, + dtype: torch.dtype = torch.bfloat16, + atol: float = 1e-5, + rtol: float = 1.6e-2, + ) -> None: + G, M, N, K = shape + # Set up inputs for forward pass + # In M*G grouping, input is [M*G, K] and weights are [N, K] + a = torch.randn(M * G, K, dtype=dtype, device=device, requires_grad=True) + b = torch.randn(N, K, dtype=dtype, device=device, requires_grad=True) + + # Create equal-sized groups for simplicity + m_size = M + m_sizes = torch.full((G,), m_size, device=device, dtype=torch.int32) + + # Run forward pass with our implementation + result = grouped_gemm_forward(a, b, m_sizes) + # Ensure result has correct shape + self.assertTrue(result.shape == (M * G, N)) + + # Compute expected result using reference implementation + expected_result = compute_reference_forward(a, b, m_sizes) + + # Verify forward pass correctness + forward_close = analyze_tensor_differences( + result, expected_result, "Forward output" + ) + self.assertTrue(forward_close) + + # Create a gradient for backpropagation + grad_output = torch.randn_like(result) + + # Compute gradients using our custom backward implementation + grad_a, grad_b = grouped_gemm_backward(grad_output, a, b, m_sizes) + + # Compute expected gradients using reference implementation + expected_grad_a, expected_grad_b = compute_reference_backward( + a, b, m_sizes, grad_output + ) + + # Verify gradient correctness + grad_a_close = analyze_tensor_differences(grad_a, expected_grad_a, "grad_x") + grad_b_close = analyze_tensor_differences(grad_b, expected_grad_b, "grad_w") + + self.assertTrue(grad_a_close) + self.assertTrue(grad_b_close) + + def test_MG_grouped_gemm_backward_bf16(self) -> None: + for G in (1, 8, 16): + for M in (512, 1024): + print(f"Testing BF16 M*G GroupGeMM Backward with G={G}, M={M}") + self._run_grouped_gemm_backward_test( + (G, M, 1024, 1024), + torch.device("cuda"), + dtype=torch.float16, + atol=1e-2, + rtol=1e-2, + ) + + def test_MG_grouped_gemm_backward_deepseek_shapes(self) -> None: + """Test backward pass with shapes from Deepseek model.""" + deepseek_shapes = [ + (4, 2048, 4096, 7168), # G, M, N, K + (4, 2048, 7168, 2048), + (8, 512, 4096, 7168), + (8, 512, 7168, 2048), + ] + + device = torch.device("cuda") + + for shape in deepseek_shapes: + G, M, N, K = shape + print( + f"Testing BF16 M*G Deepseek Backward shape: G={G}, M={M}, N={N}, K={K}" + ) + self._run_grouped_gemm_backward_test( + shape, device, dtype=torch.float16, atol=1e-2, rtol=1e-2 + ) + + def test_MG_dx(self) -> None: + """Test specifically the dx (gradient w.r.t. input) computation.""" + G, M, N, K = 4, 512, 1024, 2048 + device = torch.device("cuda") + dtype = torch.bfloat16 + + # Set up inputs + a = torch.randn(M * G, K, dtype=dtype, device=device, requires_grad=True) + b = torch.randn(N, K, dtype=dtype, device=device, requires_grad=True) + + # Create equal-sized groups + m_size = M + m_sizes = torch.full((G,), m_size, device=device, dtype=torch.int32) + + # Forward pass + result = grouped_gemm_forward(a, b, m_sizes) + + # Create gradient for backward + grad_output = torch.randn_like(result) + + # Compute gradient using our optimized function + grad_a, _ = grouped_gemm_backward(grad_output, a, b, m_sizes) + + # Compute expected gradient using reference implementation + expected_grad_a, _ = compute_reference_backward(a, b, m_sizes, grad_output) + + # Verify gradient + dx_close = analyze_tensor_differences(grad_a, expected_grad_a, "grad_a (dx)") + self.assertTrue(dx_close) + + def test_MG_dw(self) -> None: + """Test specifically the dw (gradient w.r.t. weights) computation.""" + G, M, N, K = 4, 512, 1024, 2048 + device = torch.device("cuda") + dtype = torch.bfloat16 + + # Set up inputs + a = torch.randn(M * G, K, dtype=dtype, device=device, requires_grad=True) + b = torch.randn(N, K, dtype=dtype, device=device, requires_grad=True) + + # Create equal-sized groups + m_size = M + m_sizes = torch.full((G,), m_size, device=device, dtype=torch.int32) + + # Forward pass + result = grouped_gemm_forward(a, b, m_sizes) + + # Create gradient for backward + grad_output = torch.randn_like(result) + + # Compute gradient using our optimized function + _, grad_b = grouped_gemm_backward(grad_output, a, b, m_sizes) + + # Compute expected gradient using reference implementation + _, expected_grad_b = compute_reference_backward(a, b, m_sizes, grad_output) + + # Verify gradient + dw_close = analyze_tensor_differences(grad_b, expected_grad_b, "grad_b (dw)") + self.assertTrue(dw_close) diff --git a/torchtitan/experiments/simple_fsdp/simple_fsdp.py b/torchtitan/experiments/simple_fsdp/simple_fsdp.py new file mode 100644 index 0000000000000000000000000000000000000000..34d255e199f83d41e19b7000d7d0b79fcb61455d --- /dev/null +++ b/torchtitan/experiments/simple_fsdp/simple_fsdp.py @@ -0,0 +1,194 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn as nn + +from torch.distributed._tensor import ( + distribute_tensor, + DTensor, + Partial, + Replicate, + Shard, +) +from torch.utils.checkpoint import ( + checkpoint, + CheckpointPolicy, + create_selective_checkpoint_contexts, +) + + +_active_parametrization = True + + +@contextmanager +def disable_data_parallel(): + global _active_parametrization + try: + _active_parametrization = False + yield + finally: + _active_parametrization = True + + +@dataclass(frozen=True) +class MixedPrecisionPolicy: + param_dtype: Optional[torch.dtype] = None + reduce_dtype: Optional[torch.dtype] = None + + +def fsdp_policy(): + def _fsdp_recomp_policy(): + def _custom_policy(ctx, func, *args, **kwargs): + to_recompute = func in { + torch.ops._c10d_functional.all_gather_into_tensor.default, + torch.ops._c10d_functional.wait_tensor.default, + torch.ops.aten._to_copy.default, # for dtype cast in FSDP + } + return ( + CheckpointPolicy.MUST_RECOMPUTE + if to_recompute + else CheckpointPolicy.MUST_SAVE + ) + + return _custom_policy + + return create_selective_checkpoint_contexts(_fsdp_recomp_policy()) + + +class ReplicateComputation(torch.nn.Module): + def __init__(self, device_mesh, param_sharding, mode, regional_ac, mp_policy): + super().__init__() + self.device_mesh = device_mesh + self.param_sharding = param_sharding + self.mode = mode + self.compute_placements = [Replicate()] * self.device_mesh.ndim + self.grad_placements = [Partial(reduce_op="avg")] * self.device_mesh.ndim + self.regional_ac = regional_ac + mp_policy = mp_policy or MixedPrecisionPolicy() + self.param_dtype = mp_policy.param_dtype + self.reduce_dtype = mp_policy.reduce_dtype + + def replicate_compute(self, x): + # data parallel runtime replicate parameters and do local compute + # the gradients are partial tensors that needs to perform reduction + # (i.e. DDP: allreduce, FSDP: reduce_scatter, HSDP: mix of both) + + # NOTE: specifying mixed precision is only available in pytorch_intern24 + # https://github.com/tianyu-l/pytorch_intern24/pull/20 + # support for FSDP + TP (assuming TP shards the inner-most dim) + if self.mode == "fully_shard" and x._spec.mesh.ndim == 2: + dp_placement, tp_placement = x._spec.placements + dp_mesh, tp_mesh = self.device_mesh, x._spec.mesh["tp"] + + # re-wrap 2D DTensor to 1D DTensor on dp_mesh for efficient FSDP all-gather + # TODO: we should consider merging this logic into DTensor redistribute API + sharded_local_tensor = x.to_local() + sharded_dtensor = DTensor.from_local( + sharded_local_tensor, dp_mesh, self.param_sharding + ) + + # the actuall FSDP all-gather on dp_mesh + # TODO(ruisizhang123): enable mixed-precision training here + # add the forward_dtype and backward_dtype back after landing changes in PyTorch DTensor + replicated_dtensor = sharded_dtensor.redistribute( + placements=self.compute_placements, + # forward_dtype=self.param_dtype, + # backward_dtype=self.reduce_dtype, + ) + + # re-wrap 1D all-gathered DTensor on dp_mesh to 1D DTensor on tp_mesh + # TODO: DTensor should support this mesh collasping operation + replicated_local_tensor = replicated_dtensor.to_local( + grad_placements=self.grad_placements + ) + output = DTensor.from_local( + replicated_local_tensor, tp_mesh, (tp_placement,) + ) + else: + output = x.redistribute( + placements=self.compute_placements, + # forward_dtype=self.param_dtype, + # backward_dtype=self.reduce_dtype, + ).to_local(grad_placements=self.grad_placements) + + return output + + def forward(self, x): + global _active_parametrization + # This should never be set to true during forward, only outside for model + # inspection / debugging / initialization + # model initialization can be done now through + # with disable_data_parallel(): + # model.init_weights() + if not _active_parametrization: + return x + + if self.regional_ac and self.mode in ("fully_shard", "hybrid_shard"): + # apply checkpointing to implement reshard_after_forward + output = checkpoint( + self.replicate_compute, x, use_reentrant=False, context_fn=fsdp_policy + ) + else: + output = self.replicate_compute(x) + + return output + + +def data_parallel( + model, + device_mesh, + mode="replicate", + ac_mode: str = "none", + mp_policy: Optional[MixedPrecisionPolicy] = None, +): + if mode == "replicate": + param_sharding = (Replicate(),) + elif mode == "fully_shard": + param_sharding = (Shard(0),) + elif mode == "hybrid_shard": + # replicate inter-host, fully shard intra-host + param_sharding = (Replicate(), Shard(0)) + assert ( + device_mesh.ndim == 2 + ), "hybrid sharded data parallel requires 2D DeviceMesh" + else: + raise ValueError(f"Unsupported mode {mode}") + + modules = list(model.modules()) + + # apply regional ac (with fsdp_policy) if no global ac is to be applied + regional_ac = ac_mode == "none" + + for mod in modules: + params_dict = dict(mod.named_parameters(recurse=False)) + for p_name, p in params_dict.items(): + if p is not None and p.numel() > 0: + mod.register_parameter( + p_name, + # NOTE: for 2D we need to distribute_tensor a DTensor + # which requires latest change in pytorch_intern24 + # https://github.com/tianyu-l/pytorch_intern24/pull/25 + nn.Parameter(distribute_tensor(p, device_mesh, param_sharding)), + ) + nn.utils.parametrize.register_parametrization( + mod, + p_name, + ReplicateComputation( + device_mesh, + param_sharding, + mode, + regional_ac, + mp_policy=mp_policy, + ), + unsafe=True, + ) + + return model diff --git a/torchtitan/models/llama3/train_configs/llama3_8b.toml b/torchtitan/models/llama3/train_configs/llama3_8b.toml new file mode 100644 index 0000000000000000000000000000000000000000..5594b0fe33763fd289ab0e5a298858ec1869b23b --- /dev/null +++ b/torchtitan/models/llama3/train_configs/llama3_8b.toml @@ -0,0 +1,63 @@ +# torchtitan Config.toml +# NOTE: this toml config is a preset for 64 A100 GPUs. + +[job] +dump_folder = "./outputs" +description = "Llama 3 8B training" + +[profiling] +enable_profiling = true +save_traces_folder = "profile_trace" +profile_freq = 100 + +[metrics] +log_freq = 10 +enable_tensorboard = true +save_tb_folder = "tb" + +[model] +name = "llama3" +flavor = "8B" +norm_type = "rmsnorm" # layernorm / np_layernorm / rmsnorm +tokenizer_path = "./assets/tokenizer/original/tokenizer.model" +# converters = "float8" + +[optimizer] +name = "AdamW" +lr = 3e-4 +eps = 1e-8 + +[lr_scheduler] +warmup_steps = 200 # lr scheduler warm up + +[training] +batch_size = 1 +seq_len = 8192 +max_norm = 1.0 # grad norm clipping +steps = 1000 +compile = false +dataset = "c4" + +[parallelism] +data_parallel_replicate_degree = 1 +data_parallel_shard_degree = -1 +tensor_parallel_degree = 1 +pipeline_parallel_degree = 1 +context_parallel_degree = 1 + +[checkpoint] +enable_checkpoint = false +folder = "checkpoint" +interval = 500 +model_weights_only = false +export_dtype = "float32" +async_mode = "disabled" # ["disabled", "async", "async_with_pinned_mem"] + +[activation_checkpoint] +mode = 'selective' # ['none', 'selective', 'full'] +selective_ac_option = 'op' # 'int' = ac every positive int layer or 'op', ac based on ops policy + +[float8] +enable_fsdp_float8_all_gather = false +precompute_float8_dynamic_scale_for_fsdp = false +filter_fqns = "output"