feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

来源:
  1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
     - inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
     - inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
     - contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
     - contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
     - csrc/include/ixformer/: C++ kernel headers + cmake

  2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
     - npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
     - npu_torch/qwen3_5_gated_delta_net.cpp/.h
     - npu_torch/qwen3_next_*.cpp/.h (6 files)
     - npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
     - models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
     - models/vlm/qwen3_5.h

调用链完整性:
  ixformer_sdk/inference/functions/vllm.py
    → ops.infer.moe_topk_softmax() (C++ 层)
    → 这就是 base 镜像 libixformer.so 里的实现

  upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
    → ixformer::infer::topk_softmax() (直接 C++ 调用)
    → ixformer::infer::group_gemm() → 完整 7-step MoE pipeline

ixformer_sdk/train/speedformer/layers/baichuan/attention.py

@@ -0,0 +1,162 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input 
+from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig
+from ixformer.train.speedformer.models.baichuan.modeling_baichuan import Attention
+
+from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb
+from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+
+class FlashAttention(Attention):
+    # 这个类主要的改进包含：1. apply_rotary_pos_emb；2. flash-attn 代替 native attention
+    def __init__(self, config: BaichuanConfig):
+        super().__init__(config)
+        self.rotary_emb = RotaryEmbedding(self.head_dim)
+        
+        
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value: Optional[Tuple[torch.Tensor]] = None,
+            output_attentions: bool = False,
+            use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        proj = self.W_pack(hidden_states)
+        proj = proj.unflatten(-1, (3, self.hidden_size)).unsqueeze(0).transpose(0, -2).squeeze(-2)
+        
+        # fused_apply_rotary_pos_emb need qk to be in "sbhd", v stay in "bshd"
+        query_states = proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous()
+        key_states = proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous()
+        value_states = proj[2].view(bsz, q_len, self.num_heads, self.head_dim)
+
+        kv_seq_len = key_states.shape[0]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[0]
+            
+        # fused_apply_rotary_pos_emb need emb in float32
+        emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32)
+        query_states = fused_apply_rotary_pos_emb(query_states, emb)
+        key_states = fused_apply_rotary_pos_emb(key_states, emb)
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=0)
+            value_states = torch.cat([past_key_value[1], value_states], dim=0)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa
+        if attention_mask is None: # flash-attn
+            query_states = query_states.transpose(0, 1).contiguous()
+            key_states   = key_states.transpose(0, 1).contiguous()
+        else:  # sdpa
+            query_states = query_states.permute(1, 2, 0, 3).contiguous()
+            key_states   = key_states.permute(1, 2, 0, 3).contiguous()
+            value_states = value_states.transpose(1, 2).contiguous()
+
+        '''
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]   # bsz, q_len, self.num_heads, self.head_dim
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, q_len
+            )
+            
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=0.0,
+                softmax_scale=None,
+                causal=True,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, 0.0, softmax_scale=None, causal=True
+            )
+        '''
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, q_len, attention_mask, dropout=0.0
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+        
+        return attn_output, attn_weights, past_key_value
+    
+    def _flash_attention_forward(
+        self, 
+        query_states: torch.Tensor, 
+        key_states: torch.Tensor, 
+        value_states: torch.Tensor, 
+        query_length: int, 
+        attention_mask: Optional[torch.Tensor] = None, 
+        dropout=0.0, 
+        softmax_scale=None
+    ):
+        if attention_mask is not None:
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+                is_causal=query_length > 1,
+            )
+            attn_output = attn_output.transpose(1, 2).contiguous()
+            
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal
+            )
+
+        return attn_output
+        
+class BaichuanAttention(FlashAttention):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "BaichuanAttention is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native BaichuanAttention module to LlamaAttention module provided above."
+        )
+        
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        LazyInitContext.materialize(module)
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        config = getattr(module, "config")
+
+        attention = FlashAttention(
+            config=config,
+        )
+    
+        attention.W_pack.weight = module.W_pack.weight
+        attention.o_proj.weight = module.o_proj.weight
+
+        return attention

ixformer_sdk/train/speedformer/layers/baichuan/baichuan_model.py

@@ -0,0 +1,141 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig
+from ixformer.train.speedformer.models.baichuan.modeling_baichuan import BaichuanModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.utils import logging, ContextManagers
+
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+logger = logging.get_logger(__name__)
+
+
+class IXFBaichuanModel(BaichuanModel):
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError(
+                "You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (
+                    layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )

ixformer_sdk/train/speedformer/layers/baichuan/mlp.py

@@ -0,0 +1,53 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+import ixformer.train.functions as F
+from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig
+from ixformer.train.speedformer.models.baichuan.modeling_baichuan import MLP
+from transformers.utils import logging
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+
+class BaseMLP(MLP):
+    """
+    这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x)))
+    """
+
+    def __init__(self, hidden_size, intermediate_size, hidden_act):
+        super().__init__(hidden_size, intermediate_size, hidden_act)
+        self.gate_up = nn.Linear(
+            hidden_size, intermediate_size * 2, bias=False)
+        del self.gate_proj, self.up_proj
+        del self.act_fn
+
+    def forward(self, x):
+        res = self.gate_up(x)
+        down_proj = self.down_proj(F.swiglu(res))
+        return down_proj
+
+
+class IXFBaichuanMLP(BaseMLP):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "IXFLlamaMLP is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to IXFLlamaMLP module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+        hidden_size, intermediate_size = module.gate_proj.in_features, module.gate_proj.out_features
+        hidden_act = "silu"
+
+        mlp = BaseMLP(hidden_size=hidden_size,
+                      intermediate_size=intermediate_size, hidden_act=hidden_act)
+
+        mlp.gate_up.weight.data = torch.concat(
+            (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0)
+        mlp.down_proj.weight.data = module.down_proj.weight.data
+
+        return mlp

ixformer_sdk/train/speedformer/layers/bloom/attention.py

@@ -0,0 +1,160 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input 
+from ixformer.train.speedformer.models.bloom.modeling_bloom import BloomAttention, dropout_add
+from ixformer.train.speedformer.models.bloom.configuration_bloom import BloomConfig
+
+from apex.transformer.functional.fused_rope import fused_apply_rotary_pos_emb_cached
+from apex.transformer.functional.fused_rope import FusedRoPEFunc
+
+
+class FlashAttention(BloomAttention):
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        alibi: torch.Tensor,
+        attention_mask: torch.Tensor,
+        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+    ):
+        fused_qkv = self.query_key_value(hidden_states)
+        (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)  # 3 x [batch_size, seq_length, num_heads, head_dim]
+        batch_size, q_length, _, _ = query_layer.shape
+        
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key_layer = torch.cat((past_key, key_layer), dim=1)
+            value_layer = torch.cat((past_value, value_layer), dim=1)
+            
+        present = (key_layer, value_layer) if use_cache else None
+        # if attention_mask is not None:
+        if False:
+            query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_layer, key_layer, value_layer, attention_mask, q_length
+            )
+            
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+            attn_output_unpad = flash_attn_varlen_func(
+                query_layer,
+                key_layer,
+                value_layer,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=0.0,
+                softmax_scale=None,
+                causal=True,
+                use_alibi=True,
+            )
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_length)
+        else:
+            attn_output = flash_attn_func(
+                query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True, use_alibi=True,
+            )
+
+        attn_output = attn_output.reshape(batch_size, q_length, attn_output.shape[2]*attn_output.shape[3]).contiguous()
+        output_tensor = self.dense(attn_output)
+        
+        output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
+        
+        outputs = (output_tensor, present, None)
+
+        return outputs
+        
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        
+        def _get_unpad_data(attention_mask):
+            seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+            max_seqlen_in_batch = seqlens_in_batch.max().item()
+            cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+            return (
+                indices,
+                cu_seqlens,
+                max_seqlen_in_batch,
+            )
+            
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+        
+        
+class BloomFlashAttention(FlashAttention):
+
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "BloomAttention is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native BloomAttention module to FlashAttention module provided above."
+        )
+        
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        new_config = BloomConfig()
+        new_config.pretraining_tp = module.pretraining_tp
+        new_config.slow_but_exact = module.slow_but_exact
+        new_config.hidden_size = module.hidden_size
+        new_config.n_head = module.num_heads
+        new_config.hidden_size = module.split_size
+        new_config.hidden_dropout = module.hidden_dropout
+        new_config.attention_dropout = module.attention_dropout.p
+
+        attention = FlashAttention(
+            config=new_config,
+        )
+    
+        attention.query_key_value.weight = module.query_key_value.weight
+        attention.query_key_value.bias = module.query_key_value.bias
+
+        attention.dense.weight = module.dense.weight
+        attention.dense.bias = module.dense.bias
+
+        return attention

ixformer_sdk/train/speedformer/layers/chatglm/attention.py

@@ -0,0 +1,199 @@
+import math
+import os
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ixformer.train.speedformer.models.chatglm.modeling_chatglm import (
+    CoreAttention,
+    SelfAttention,
+    split_tensor_along_last_dim,
+    apply_rotary_pos_emb
+)
+from ixformer.train.speedformer.models.chatglm.configuration_chatglm import ChatGLMConfig
+
+from transformers.utils import is_flash_attn_2_available
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+class FlashCoreAttention(CoreAttention):
+
+    def forward(self, query_layer, key_layer, value_layer, attention_mask):
+        if int(os.environ.get("USE_FLASH_ATTN", 0)):
+            query_layer, key_layer, value_layer = [
+                k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]]
+            batch_size, query_length, _, _ = query_layer.shape
+
+            if attention_mask is not None:
+                batch_size = query_layer.shape[0]
+                query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                    query_layer, key_layer, value_layer, attention_mask, query_length
+                )
+
+                cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+                max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_layer,
+                    key_layer,
+                    value_layer,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
+                    dropout_p=0.0,
+                    softmax_scale=None,
+                    causal=True,
+                )
+                attn_output = pad_input(
+                    attn_output_unpad, indices_q, batch_size, query_length)
+                context_layer = attn_output.permute(1, 0, 2, 3)
+            else:
+                attn_output = flash_attn_func(
+                    query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True
+                )
+                context_layer = attn_output.permute(1, 0, 2, 3)
+
+        if attention_mask is not None:
+            if query_layer.shape[2] != key_layer.shape[2]:
+                num_group = query_layer.shape[2] // key_layer.shape[2]
+                final_shape = (*key_layer.shape[:2], *query_layer.shape[2:])
+                key_layer = key_layer.unsqueeze(-2)
+                key_layer = key_layer.expand(
+                    -1, -1, -1, num_group, -1
+                )
+                key_layer = key_layer.contiguous().view(
+                    final_shape
+                )
+                value_layer = value_layer.unsqueeze(-2)
+                value_layer = value_layer.expand(
+                    -1, -1, -1, num_group, -1
+                )
+                value_layer = value_layer.contiguous().view(
+                    final_shape
+                )
+
+            query_layer, key_layer, value_layer = [
+                k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]  # bhsd
+            attention_mask = ~attention_mask
+            context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
+                                                                             attention_mask)
+            context_layer = context_layer.permute(2, 0, 1, 3)
+
+        else:
+            query_layer, key_layer, value_layer = [
+                k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]]  # bshd
+            context_layer = flash_attn_func(
+                query_layer, key_layer, value_layer, 0, softmax_scale=None, causal=True
+            )  # bshd
+            context_layer = context_layer.permute(1, 0, 2, 3)
+
+        context_layer = context_layer.reshape(
+            context_layer.size(0), context_layer.size(1), -1)
+
+        return context_layer
+
+
+class FlashSelfAttention(SelfAttention):
+
+    def __init__(self, config: ChatGLMConfig, layer_number, device=None):
+        super().__init__(config, layer_number, device=device)
+        self.core_attention = FlashCoreAttention(config, self.layer_number)
+
+    def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True):
+        mixed_x_layer = self.query_key_value(hidden_states)
+        if self.multi_query_attention:
+            (query_layer, key_layer, value_layer) = mixed_x_layer.split(
+                [
+                    self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
+                    self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
+                    self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
+                ],
+                dim=-1,
+            )
+            query_layer = query_layer.view(
+                query_layer.size()[
+                    :-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
+            )
+            key_layer = key_layer.view(
+                key_layer.size()[
+                    :-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
+            )
+            value_layer = value_layer.view(
+                value_layer.size()[:-1]
+                + (self.num_multi_query_groups_per_partition,
+                   self.hidden_size_per_attention_head)
+            )
+        else:
+            new_tensor_shape = mixed_x_layer.size()[:-1] + \
+                (self.num_attention_heads_per_partition,
+                 3 * self.hidden_size_per_attention_head)
+            mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
+
+            # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
+            (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(
+                mixed_x_layer, 3)
+
+        if rotary_pos_emb is not None:
+            query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
+            key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
+
+        # adjust key and value for inference
+        if kv_cache is not None:
+            cache_k, cache_v = kv_cache
+            key_layer = torch.cat((cache_k, key_layer), dim=0)
+            value_layer = torch.cat((cache_v, value_layer), dim=0)
+        if use_cache:
+            kv_cache = (key_layer, value_layer)
+        else:
+            kv_cache = None
+
+        # 这里省略了 kv "sbhd" -> "sb(h*num_multi-group)d" 的过程，因为flash-attn支持 MGA
+        # ==================================
+        # core attention computation
+        # ==================================
+
+        context_layer = self.core_attention(
+            query_layer, key_layer, value_layer, attention_mask)
+
+        # =================
+        # Output. [sq, b, h]
+        # =================
+
+        output = self.dense(context_layer)
+
+        return output, kv_cache
+
+
+class ChatglmFlashAttention(FlashSelfAttention):
+
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "BloomAttention is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native BloomAttention module to FlashAttention module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+        # 这个原实现没有在类中保存config，所以需要初始化一个config
+        layer_number = getattr(module, "layer_number")
+        config = getattr(module, "config")
+        attention = FlashSelfAttention(
+            config=config,
+            layer_number=layer_number,
+        )
+
+        attention.query_key_value.weight.data = module.query_key_value.weight.data
+        attention.dense.weight.data = module.dense.weight.data
+        if getattr(attention.query_key_value, "bias") is not None:
+            attention.query_key_value.bias.data = module.query_key_value.bias.data
+        if getattr(attention.dense, "bias") is not None:
+            attention.dense.bias.data = module.dense.bias.data
+
+        return attention

ixformer_sdk/train/speedformer/layers/chatglm/attributions.py

@@ -0,0 +1,9 @@
+from ixformer.train.speedformer.models.chatglm.modeling_chatglm import RotaryEmbedding
+from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
+
+
+class ChatglmRotaryEmbedding(RotaryEmbedding):
+    def from_native_attr(attr_class, *args, **kwargs):
+        dim = attr_class.dim
+        rote = RotaryEmbedding(dim=dim)
+        return rote

ixformer_sdk/train/speedformer/layers/chatglm/methods.py

@@ -0,0 +1,127 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ixformer.train.speedformer.models.chatglm.modeling_chatglm import ChatGLMModel
+
+
+def ChatGLMModel_forward():
+    from transformers.modeling_outputs import BaseModelOutputWithPast
+    from transformers.utils import logging, is_flash_attn_2_available
+
+    def forward(
+            self,
+            input_ids,
+            position_ids: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.BoolTensor] = None,
+            full_attention_mask: Optional[torch.BoolTensor] = None,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            use_cache: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ):
+        def is_lower_triangular(mask):
+            """
+            ixdnn 虽然支持2种causal mask, 如下图:
+            mode0:
+                if seqlen_q < seqlen_k
+                    1 0 0 0 0
+                    1 1 0 0 0 
+                if seqlen_k < seqlen_q
+                    1 0
+                    1 1
+                    1 1
+                    1 1
+                    1 1
+            mode1:
+                if seqlen_q < seqlen_k
+                    1 1 1 1 0
+                    1 1 1 1 1
+                if seqlen_k < seqlen_q
+                    0 0
+                    0 0
+                    0 0
+                    1 0
+                    1 1
+
+            但 flash-attn 目前只支持 mode1, 所以下面需要判断一下传入的mask是不是mode1这种模式
+            """
+            batch_size, _, rows, cols = mask.shape
+
+            # 创建一个mode1的下三角矩阵
+            if rows <= cols:
+                part = torch.ones(rows, cols - rows,
+                                  dtype=torch.bool, device=mask.device)
+                gt = ~torch.triu(torch.ones(
+                    rows, rows, dtype=torch.bool, device=mask.device), diagonal=1)
+                gt = torch.cat((part, gt), dim=1)
+            else:
+                part = torch.zeros(
+                    rows-cols, cols, dtype=torch.bool, device=mask.device)
+                gt = ~torch.triu(torch.ones(
+                    cols, cols, dtype=torch.bool, device=mask.device), diagonal=1)
+                gt = torch.cat((part, gt), dim=0)
+            gt = gt[None, None, :, :].expand(batch_size, -1, -1, -1)
+
+            # 检查所有的元素是不是都一样
+            check = (gt == mask).all()
+
+            return check
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size, seq_length = input_ids.shape
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embedding(input_ids)
+
+        if self.pre_seq_len is not None:
+            if past_key_values is None:
+                past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
+                                                  dtype=inputs_embeds.dtype)
+            if attention_mask is not None:
+                attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
+                                            attention_mask], dim=-1)
+
+        if full_attention_mask is None:
+            if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
+                full_attention_mask = self.get_masks(
+                    input_ids, past_key_values, padding_mask=attention_mask)
+
+        # Rotary positional embeddings
+        rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
+        if position_ids is not None:
+            rotary_pos_emb = rotary_pos_emb[position_ids]
+        else:
+            rotary_pos_emb = rotary_pos_emb[None, :seq_length]
+        rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
+
+        # Run encoder.
+        attn_mask = None
+        if full_attention_mask is not None:
+            if not is_lower_triangular(full_attention_mask):
+                attn_mask = full_attention_mask
+        hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
+            inputs_embeds, attn_mask, rotary_pos_emb=rotary_pos_emb,
+            kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
+        )
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+    return forward

ixformer_sdk/train/speedformer/layers/cross_entropy_loss.py

@@ -0,0 +1,370 @@
+import time
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+from packaging.version import Version
+if Version(triton.__version__) >= Version("3.0.0"):
+    from triton.language.extra import libdevice
+    triton_tanh = libdevice.tanh
+else:
+    import triton.language as tl
+    triton_tanh = tl.math.tanh
+
+
+def calculate_settings(n):
+    BLOCK_SIZE = triton.next_power_of_2(n)
+    if BLOCK_SIZE > MAX_FUSED_SIZE:
+        raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "
+                           f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
+    num_warps = 4
+    if BLOCK_SIZE >= 32768:
+        num_warps = 32
+    elif BLOCK_SIZE >= 8192:
+        num_warps = 16
+    elif BLOCK_SIZE >= 2048:
+        num_warps = 8
+    return BLOCK_SIZE, num_warps
+
+
+@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], })
+@triton.jit
+def _cross_entropy_forward(
+    logits_ptr, logits_row_stride,
+    loss_ptr,
+    logsumexp_ptr,
+    labels_ptr,
+    VOCAB_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    DO_SOFTCAPPING: tl.constexpr,
+    SOFTCAP: tl.constexpr,
+):
+    """
+        Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
+        Pi = exp(xi) / sum(exp(xi))
+        CE_i = -y log(p) = -y log[ exp(x) / sum(exp(x)) ]
+             = -y [ x - log[sum(exp(x))] ]
+             = y * (log[sum(exp(x))] - x)
+        If y == 0: CE_i = 0
+        If y == 1: CE_i = logsumexp - x
+
+        logsumexp is also stable
+        Take    y =         log[sum(exp(x))]
+           exp(y) =             sum(exp(x))
+           exp(y) =             sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x
+           exp(y) =      exp(c)*sum(exp(x - c))
+               y  = log(exp(c)*sum(exp(x - c)))
+               y  = c + log[sum(exp(x - c))]
+        This means we can set c = max(x) to make sure
+        exp(x - c) always is exp(x - max(x)).
+        This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1.
+    """
+    row_idx = tl.program_id(0)
+    logits_ptr += row_idx * logits_row_stride.to(tl.int64)
+    loss_ptr += row_idx
+    logsumexp_ptr += row_idx
+    labels_ptr += row_idx
+
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < VOCAB_SIZE
+
+    label_idx = tl.load(labels_ptr).to(tl.int32)
+    logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf"))
+    # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+    if DO_SOFTCAPPING:
+        logits = SOFTCAP * triton_tanh(logits / SOFTCAP)
+
+    logits = logits.to(tl.float32)
+    c = tl.max(logits, 0)
+    logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
+
+    if label_idx != -100:
+        x = tl.load(logits_ptr + label_idx)
+        # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+        if DO_SOFTCAPPING:
+            x = SOFTCAP * triton_tanh(x / SOFTCAP)
+        loss = logsumexp - x.to(tl.float32)
+    else:
+        loss = 0.0
+    tl.store(logsumexp_ptr, logsumexp)
+    tl.store(loss_ptr, loss)
+
+
+@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], })
+@triton.jit
+def _chunked_cross_entropy_forward(
+    logits_ptr, logits_row_stride,
+    loss_ptr,
+    logsumexp_ptr,
+    labels_ptr,
+    VOCAB_SIZE: tl.constexpr,
+    N_CHUNKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    DO_SOFTCAPPING: tl.constexpr,
+    SOFTCAP: tl.constexpr,
+):
+    """
+        256K vocab divided in 4 chunks
+
+        |-65536-| |-65536-| |-65536-| |-65536-|
+        |-------| |-------| |-------| |-------|
+        |-------| |-------| |-------| |-------|
+
+        If y == 0: CE_i = 0
+        If y == 1: CE_i = logsumexp - x
+
+        Notice we can do logsumexp for each chunk and then
+        logsumexp[chunk_sum(logsumexp)] == logsumexp
+
+        chunk_sum = log[chunk_sum(logsumexp)]
+                  = log[exp(logsumexp(a)) + ... + exp(logsumexp(z))]
+                  = log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])]
+                  = log[sum(exp(a)) + ... + sum(exp(z))]
+                  = logsumexp(x)
+
+        This means we can perform a logsumexp for each chunk, then do a
+        final logsumexp reduction!
+
+        Ie do: logsumexp(chunked_logsumexp) - x
+    """
+    row_idx = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    logits_ptr += row_idx * logits_row_stride.to(tl.int64)
+    loss_ptr += row_idx
+    logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
+    labels_ptr += row_idx
+
+    col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < VOCAB_SIZE
+
+    label_idx = tl.load(labels_ptr).to(tl.int32)
+    logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf"))
+    # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+    if DO_SOFTCAPPING:
+        logits = SOFTCAP * triton_tanh(logits / SOFTCAP)
+
+    logits = logits.to(tl.float32)
+    c = tl.max(logits, 0)
+    logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
+
+    if chunk_idx == 0:
+        # logsumexp(chunked_logsumexp) - x
+        # Do the -x separately
+        if label_idx != -100:
+            x = tl.load(logits_ptr + label_idx).to(tl.float32)
+            # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+            if DO_SOFTCAPPING:
+                x = SOFTCAP * triton_tanh(x / SOFTCAP)
+            loss = -1.0 * x.to(tl.float32)
+        else:
+            loss = 0.0
+        tl.store(loss_ptr, loss)
+
+    tl.store(logsumexp_ptr, logsumexp)
+
+
+@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], })
+@triton.jit
+def _cross_entropy_backward(
+    logits_ptr, logits_row_stride,
+    dloss_ptr,   dloss_row_stride,
+    logsumexp_ptr,
+    labels_ptr,
+    VOCAB_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    DO_SOFTCAPPING: tl.constexpr,
+    SOFTCAP: tl.constexpr,
+):
+    """
+        CE_i = -y log(P) = y * (log[sum(exp(x))] - x)
+        dC/dx = d/dx (y * log[sum(exp(x))] - x * y)
+
+        From https://en.wikipedia.org/wiki/LogSumExp
+        d/dx logsumexp = exp(x) / sum(exp(x)) = softmax(x)
+
+        dC/dx = y * exp(x) / sum(exp(x)) - d/dx (x * y)
+        dC/dx = y * exp[ log[exp(x) / sum(exp(x))] ] using x = exp(log(x)) trick
+        dC/dx = y * exp[x - logsumexp] - d/dx (x * y)
+
+        If y == 0: dC/dx = 0
+        If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1
+        If y == 1 and x != label: dC/dx     = exp[x - logsumexp]
+    """
+    row_idx = tl.program_id(0)
+    block_idx = tl.program_id(1)
+
+    logits_ptr += row_idx * logits_row_stride.to(tl.int64)
+    dloss_ptr += row_idx * dloss_row_stride
+    col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < VOCAB_SIZE
+    label_idx = tl.load(labels_ptr + row_idx).to(tl.int32)
+
+    if label_idx != -100:
+        dloss = tl.load(dloss_ptr)
+    else:
+        dloss = 0.0
+
+    x = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf"))
+    # Do logit softcapping for Gemma 2: t * tanh(1/t * x)
+    if DO_SOFTCAPPING:
+        # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x)
+        partial = triton_tanh(x / SOFTCAP)
+        x = SOFTCAP * partial
+
+    logsumexp = tl.load(logsumexp_ptr + row_idx)
+    y = tl.exp(x.to(tl.float32) - logsumexp)
+    y = tl.where(
+        col_offsets == label_idx,
+        y - 1.0,  # exp(x - logsumexp) - 1
+        y,       # exp(x - logsumexp)
+    )
+
+    if DO_SOFTCAPPING:
+        # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x)
+        y = y * (1.0 - partial*partial)
+
+    # If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0.
+    tl.store(logits_ptr + col_offsets, dloss * y, mask=mask)
+
+
+MAX_FUSED_SIZE = 65536  # 2**16
+
+
+class Fast_CrossEntropyLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, logits, labels, logit_softcapping=0):
+        n_rows, vocab_size = logits.shape
+
+        div, mod = divmod(vocab_size, MAX_FUSED_SIZE)
+        n_chunks = div + (mod != 0)
+        losses = torch.empty(n_rows, dtype=torch.float32, device=logits.device)
+
+        DO_SOFTCAPPING = (logit_softcapping != 0)
+
+        if n_chunks == 1:
+            # For small vocabs <= 65336 like Llama, Mistral
+            BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
+            logsumexp = torch.empty(
+                n_rows, dtype=torch.float32, device=logits.device)
+
+            _cross_entropy_forward[(n_rows,)](
+                logits, logits.stride(0),
+                losses,
+                logsumexp,
+                labels,
+                VOCAB_SIZE=vocab_size,
+                BLOCK_SIZE=BLOCK_SIZE,
+                DO_SOFTCAPPING=DO_SOFTCAPPING,
+                SOFTCAP=logit_softcapping,
+                num_warps=num_warps,
+            )
+        else:
+            # For large vocabs > 65336 like Gemma 256K
+            logsumexp = torch.empty(
+                (n_rows, n_chunks,), dtype=torch.float32, device=logits.device)
+
+            _chunked_cross_entropy_forward[(n_rows, n_chunks,)](
+                logits, logits.stride(0),
+                losses,
+                logsumexp,
+                labels,
+                VOCAB_SIZE=vocab_size,
+                N_CHUNKS=n_chunks,
+                BLOCK_SIZE=MAX_FUSED_SIZE,
+                DO_SOFTCAPPING=DO_SOFTCAPPING,
+                SOFTCAP=logit_softcapping,
+                num_warps=32,
+            )
+            # logsumexp(chunked_logsumexp) - x
+            # Do the -x separately
+            logsumexp = torch.logsumexp(logsumexp, dim=1)  # Row sum
+            losses += logsumexp
+            # Don't forget to mask padding out!
+            losses.masked_fill_(labels == -100, 0)
+
+        ctx.save_for_backward(logits, logsumexp, labels)
+        ctx.DO_SOFTCAPPING = DO_SOFTCAPPING
+        ctx.logit_softcapping = logit_softcapping
+        return losses
+
+    @staticmethod
+    def backward(ctx, dlosses):
+        logits, logsumexp, labels = ctx.saved_tensors
+        n_rows, vocab_size = logits.shape
+
+        BLOCK_SIZE = 4096
+        div, mod = divmod(vocab_size, BLOCK_SIZE)
+        n_blocks = div + (mod != 0)
+
+        _cross_entropy_backward[(n_rows, n_blocks,)](
+            logits,   logits.stride(0),
+            dlosses, dlosses.stride(0),
+            logsumexp,
+            labels,
+            VOCAB_SIZE=vocab_size,
+            BLOCK_SIZE=BLOCK_SIZE,
+            DO_SOFTCAPPING=ctx.DO_SOFTCAPPING,
+            SOFTCAP=ctx.logit_softcapping,
+            num_warps=8,
+        )
+        return logits, None, None,
+
+
+@torch._disable_dynamo
+def fast_cross_entropy_loss(logits, labels, logit_softcapping=0):
+    """
+    Arguments:
+        logits: (batch, seq_len, vocab_size)
+        labels: (batch, seq_len,)
+    Returns:
+        losses: float
+    """
+    assert len(logits.size()) == 2 or len(logits.size()) == 3
+    if len(logits.size()) == 3:
+        batch, seq_len, d = logits.shape
+        assert (labels.shape == (batch, seq_len))
+        logits = logits.view(batch*seq_len, d)
+        labels = labels.view(-1)
+
+    loss = Fast_CrossEntropyLoss.apply(
+        logits,
+        labels,
+        logit_softcapping,
+    )
+    n_items = torch.count_nonzero(labels != -100)
+    return loss.sum() / n_items
+
+
+if __name__ == "__main__":
+    shift_logits_numpy = np.random.randn(4096, 32000).astype(np.float32)
+    shift_labels_numpy = np.random.randint(0, 32000, (4096, )).astype(np.int64)
+
+    shift_logits = torch.from_numpy(shift_logits_numpy).cuda()
+    shift_labels = torch.from_numpy(shift_labels_numpy).cuda()
+
+    shift_logits_ref = torch.from_numpy(shift_logits_numpy).cuda()
+    shift_labels_ref = torch.from_numpy(shift_labels_numpy).cuda()
+
+    shift_logits.requires_grad = True
+    shift_logits_ref.requires_grad = True
+
+    # test accuracy
+    loss = fast_cross_entropy_loss(shift_logits, shift_labels)
+    loss_ref = torch.nn.CrossEntropyLoss()(shift_logits_ref, shift_labels_ref)
+    loss.backward()
+    loss_ref.backward()
+
+    torch.testing.assert_close(loss, loss_ref)
+    torch.testing.assert_close(shift_logits.grad, shift_logits_ref.grad)
+
+    start = time.time()
+    for i in range(1000):
+        loss = fast_cross_entropy_loss(shift_logits, shift_labels)
+        loss.backward()
+    print("triton:", time.time() - start)
+
+    start = time.time()
+    for i in range(1000):
+        loss_ref = torch.nn.CrossEntropyLoss()(shift_logits, shift_labels)
+        loss_ref.backward()
+    print("torch:", time.time() - start)

ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora.py

@@ -0,0 +1,305 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ixformer.train.speedformer.layers.fast_lora.swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
+import torch
+from ixformer.train.speedformer.layers.fast_lora.utils import (
+    fast_dequantize,
+    QUANT_STATE,
+    get_lora_parameters,
+    matmul_lora,
+    torch_amp_custom_fwd,
+    torch_amp_custom_bwd,
+)
+
+
+class LoRA_MLP(torch.autograd.Function):
+    """
+    ### LoRA weights
+    G = G + Ag @ Bg
+    U = U + Au @ Bu
+    W = W + Aw @ Bw
+
+    ### SwiGLU(X)
+    e = X @ G
+    f = e * sigmoid(e)
+    g = X @ U
+    h = f * g
+    i = h @ W
+
+    ### Backpropagation chain rule
+    See our blog post for more details
+
+    df = sigmoid(e) * (1 - f) + f
+    dC/dW = h.T @ dY
+    dC/dU = X.T @ (D @ W.T * f)
+    dC/dG = X.T @ (D @ W.T * df * g)
+
+    ### Down projection LoRA weights
+    dC/dAw = dC/dW @ B.T
+    dC/dBw = A.T @ dC/dW
+    dC/dAw =       h.T @ dY @ B.T
+    dC/dBw = A.T @ h.T @ dY
+
+    ### Up projection LoRA weights
+    dC/dAu =       X.T @ (D @ W.T * f) @ B.T
+    dC/dBu = A.T @ X.T @ (D @ W.T * f)
+
+    ### Gate projection LoRA weights
+    dC/dAg =       X.T @ (D @ W.T * df * g) @ B.T
+    dC/dBg = A.T @ X.T @ (D @ W.T * df * g)
+
+    Don't forget to see our blog post for more details!
+    """
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(ctx, X: torch.Tensor,
+                gateW, gateW_quant, gateA, gateB, gateS,
+                upW,   upW_quant, upA,   upB,   upS,
+                downW, downW_quant, downA, downB, downS,
+                _forward_function, _backward_function,):
+        dtype = X.dtype
+
+        e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS)
+        g = matmul_lora(X,   upW,   upW_quant,   upA,   upB,   upS)
+        h = _forward_function(e, g)
+        i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
+
+        ctx.custom_saved_tensors = (
+            gateW, gateW_quant, gateS,
+            upW, upW_quant, upS,
+            downW, downW_quant, downS,
+            _backward_function,
+        )
+        ctx.save_for_backward(gateA, gateB, upA, upB, downA, downB,
+                              X, e, g)
+        return i
+    pass
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(ctx, dY: torch.Tensor):
+        gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, \
+            _backward_function = ctx.custom_saved_tensors
+        gateA, gateB, upA, upB, downA, downB, \
+            X, e, g = ctx.saved_tensors
+
+        gateA, gateB, upA, upB, downA, downB = \
+            gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t()
+
+        batch, seq_len, hd = X.shape
+        dY = dY.view(-1, dY.shape[-1])
+        X = X .view(-1, X .shape[-1])
+        e = e .view(-1, e .shape[-1])
+        g = g .view(-1, g .shape[-1])
+        dtype = X.dtype
+
+        DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
+        DW, e, g = _backward_function(DW, e, g)
+        h, df, de = DW, e, g
+
+        # Down projection LoRA weights
+        d_downA = h.t() @ (dY @ downB.t())
+        d_downB = (downA.t() @ h.t()) @ dY
+        d_downA *= downS
+        d_downB *= downS
+
+        # Up projection LoRA weights
+        d_upA = X.t() @ (df @ upB.t())
+        d_upB = (upA.t() @ X.t()) @ df
+        d_upA *= upS
+        d_upB *= upS
+
+        # Gate projection LoRA weights
+        d_gateA = X.t() @ (de @ gateB.t())
+        d_gateB = (gateA.t() @ X.t()) @ de
+        d_gateA *= gateS
+        d_gateB *= gateS
+
+        # dX  = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
+        # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
+        upW = fast_dequantize(upW.t(), upW_quant)
+        dX = torch.matmul(df, upW.t(), out=X)
+        del upW
+        dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t())
+
+        gateW = fast_dequantize(gateW.t(), gateW_quant)
+        dX += de @ gateW.t()
+        del gateW
+        dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t())
+
+        # gateW, gateW_quant, gateA, gateB, gateS,
+        #  upW,    upW_quant,   upA,   upB,   upS,
+        # downW, downW_quant, downA, downB, downS,
+        return dX.view(batch, seq_len, hd), \
+            None, None, d_gateA.t(), d_gateB.t(), None, \
+            None, None,   d_upA.t(),   d_upB.t(), None, \
+            None, None, d_downA.t(), d_downB.t(), None, \
+            None, None,  # _backward and _forward
+    pass
+
+
+pass
+
+
+def apply_lora_mlp_swiglu(self, X):
+    gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(
+        self.gate_proj)
+    upW,     upW_quant,   upA,   upB,   upS = get_lora_parameters(
+        self.  up_proj)
+    downW, downW_quant, downA, downB, downS = get_lora_parameters(
+        self.down_proj)
+
+    out = LoRA_MLP.apply(X,
+                         gateW, gateW_quant, gateA, gateB, gateS,
+                         upW,     upW_quant, upA,   upB,   upS,
+                         downW, downW_quant, downA, downB, downS,
+                         swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,)
+    return out
+
+
+pass
+
+
+class LoRA_FUSEMLP(torch.autograd.Function):
+    """
+    ### LoRA weights
+    G = G + Ag @ Bg
+    U = U + Au @ Bu
+    W = W + Aw @ Bw
+
+    ### SwiGLU(X)
+    e = X @ G
+    f = e * sigmoid(e)
+    g = X @ U
+    h = f * g
+    i = h @ W
+
+    ### Backpropagation chain rule
+    See our blog post for more details
+
+    df = sigmoid(e) * (1 - f) + f
+    dC/dW = h.T @ dY
+    dC/dU = X.T @ (D @ W.T * f)
+    dC/dG = X.T @ (D @ W.T * df * g)
+
+    ### Down projection LoRA weights
+    dC/dAw = dC/dW @ B.T
+    dC/dBw = A.T @ dC/dW
+    dC/dAw =       h.T @ dY @ B.T
+    dC/dBw = A.T @ h.T @ dY
+
+    ### Up projection LoRA weights
+    dC/dAu =       X.T @ (D @ W.T * f) @ B.T
+    dC/dBu = A.T @ X.T @ (D @ W.T * f)
+
+    ### Gate projection LoRA weights
+    dC/dAg =       X.T @ (D @ W.T * df * g) @ B.T
+    dC/dBg = A.T @ X.T @ (D @ W.T * df * g)
+
+    Don't forget to see our blog post for more details!
+    """
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(ctx, X: torch.Tensor,
+                gateupW, gateupW_quant, gateupA, gateupB, gateupS,
+                downW, downW_quant, downA, downB, downS,
+                _forward_function, _backward_function,):
+        dtype = X.dtype
+
+        res_gateup_proj = matmul_lora(
+            X, gateupW, gateupW_quant, gateupA, gateupB, gateupS)
+        # e, g = torch.chunk(res_gateup_proj, 2, dim=-1)
+        e, g = torch.split(
+            res_gateup_proj, res_gateup_proj.size(-1)//2, dim=-1)
+        h = _forward_function(e, g)
+        i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
+
+        ctx.custom_saved_tensors = (
+            gateupW, gateupW_quant, gateupS,
+            downW, downW_quant, downS,
+            _backward_function,
+        )
+        ctx.save_for_backward(gateupA, gateupB, downA, downB, X, e, g)
+        return i
+    pass
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(ctx, dY: torch.Tensor):
+        gateupW, gateupW_quant, gateupS, downW, downW_quant, downS, \
+            _backward_function = ctx.custom_saved_tensors
+        gateupA, gateupB, downA, downB, \
+            X, e, g = ctx.saved_tensors
+
+        gateupA, gateupB, downA, downB = \
+            gateupA.t(), gateupB.t(), downA.t(), downB.t()
+
+        batch, seq_len, hd = X.shape
+        dY = dY.view(-1, dY.shape[-1])
+        X = X .view(-1, X .shape[-1])
+        e = e .view(-1, e .shape[-1])
+        g = g .view(-1, g .shape[-1])
+        dtype = X.dtype
+
+        DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
+        DW, e, g = _backward_function(DW, e, g)
+        h, df, de = DW, e, g
+
+        # Down projection LoRA weights
+        d_downA = h.t() @ (dY @ downB.t())
+        d_downB = (downA.t() @ h.t()) @ dY
+        d_downA *= downS
+        d_downB *= downS
+
+        # Gate_up projection LoRA weights
+        d_gateupA = X.t() @ (de @ gateupB.t())
+        d_gateupB = (gateupA.t() @ X.t()) @ de
+        d_gateupA *= gateupS
+        d_gateupB *= gateupS
+
+        # dX  = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
+        # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
+        gateupW = fast_dequantize(gateupW.t(), gateupW_quant)
+        dX = de @ gateupW.t()
+        del gateupW
+        dX += de @ gateupB.to(dtype).t() @ (gateupS * gateupA.to(dtype).t())
+
+        # gateW, gateW_quant, gateA, gateB, gateS,
+        #  upW,    upW_quant,   upA,   upB,   upS,
+        # downW, downW_quant, downA, downB, downS,
+        return dX.view(batch, seq_len, hd), \
+            None, None, d_gateupA.t(), d_gateupB.t(), None, \
+            None, None, d_downA.t(), d_downB.t(), None, \
+            None, None,  # _backward and _forward
+    pass
+
+
+pass
+
+
+def apply_lora_fuse_mlp_swiglu(self, X):
+    gateupW, gateupW_quant, gateupA, gateupB, gateupS = get_lora_parameters(
+        self.gate_up)
+    downW, downW_quant, downA, downB, downS = get_lora_parameters(
+        self.down_proj)
+
+    out = LoRA_FUSEMLP.apply(X,
+                             gateupW, gateupW_quant, gateupA, gateupB, gateupS,
+                             downW, downW_quant, downA, downB, downS,
+                             swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,)
+    return out
+
+
+pass

ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora_.py

@@ -0,0 +1,148 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from .utils import (
+    fast_dequantize,
+    QUANT_STATE,
+    get_lora_parameters,
+    matmul_lora,
+    torch_amp_custom_fwd,
+    torch_amp_custom_bwd,
+)
+
+
+class LoRA_MLP(torch.autograd.Function):
+    """
+    ### LoRA weights
+    G = G + Ag @ Bg
+    U = U + Au @ Bu
+    W = W + Aw @ Bw
+
+    ### SwiGLU(X)
+    e = X @ G
+    f = e * sigmoid(e)
+    g = X @ U
+    h = f * g
+    i = h @ W
+
+    ### Backpropagation chain rule
+    See our blog post for more details
+
+    df = sigmoid(e) * (1 - f) + f
+    dC/dW = h.T @ dY
+    dC/dU = X.T @ (D @ W.T * f)
+    dC/dG = X.T @ (D @ W.T * df * g)
+
+    ### Down projection LoRA weights
+    dC/dAw = dC/dW @ B.T
+    dC/dBw = A.T @ dC/dW
+    dC/dAw =       h.T @ dY @ B.T
+    dC/dBw = A.T @ h.T @ dY
+
+    ### Up projection LoRA weights
+    dC/dAu =       X.T @ (D @ W.T * f) @ B.T
+    dC/dBu = A.T @ X.T @ (D @ W.T * f)
+
+    ### Gate projection LoRA weights
+    dC/dAg =       X.T @ (D @ W.T * df * g) @ B.T
+    dC/dBg = A.T @ X.T @ (D @ W.T * df * g)
+
+    Don't forget to see our blog post for more details!
+    """
+    @staticmethod
+    @torch_amp_custom_fwd
+    def forward(ctx, X : torch.Tensor,
+                gateupW, gateupW_quant, gateupA, gateupB, gateupS,
+                downW, downW_quant, downA, downB, downS,
+                _forward_function, _backward_function,):
+        dtype = X.dtype
+
+        res_gateup_proj = matmul_lora(X, gateupW, gateupW_quant, gateupA, gateupB, gateupS)
+        res_swiglu = _forward_function(res_gateup_proj)
+        res_mlp = matmul_lora(res_swiglu, downW, downW_quant, downA, downB, downS)
+
+        ctx.custom_saved_tensors = (
+            gateupW, gateupW_quant, gateupS,
+            downW, downW_quant, downS,
+            _backward_function,
+        )
+        ctx.save_for_backward(gateupA, gateupB, downA, downB, X, res_gateup_proj, res_mlp)
+        return res_mlp
+    pass
+
+
+    @staticmethod
+    @torch_amp_custom_bwd
+    def backward(ctx, dY : torch.Tensor):
+        gateupW, gateupW_quant, gateupS, downW, downW_quant, downS, \
+            _backward_function = ctx.custom_saved_tensors
+        gateupA, gateupB, downA, downB, \
+            X, res_gateup_proj, res_mlp = ctx.saved_tensors
+
+        gateupA, gateupB, downA, downB = \
+            gateupA.t(), gateupB.t(), downA.t(), downB.t()
+
+        batch, seq_len, hd = X.shape
+        dY = dY.view(-1, dY.shape[-1])
+        X  = X .view(-1, X .shape[-1])
+        res_gateup_proj  = res_gateup_proj.view(-1, res_gateup_proj.shape[-1])
+        dtype = X.dtype
+
+        D_swiglu = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS)
+        DW, e, g = _backward_function(D_swiglu, res_gateup_proj)
+        h, df, de = DW, e, g
+
+        # Down projection LoRA weights
+        d_downA = h.t() @ (dY @ downB.t())
+        d_downB = (downA.t() @ h.t()) @ dY
+        d_downA *= downS
+        d_downB *= downS
+
+        # Gate_up projection LoRA weights
+        d_gateupA = X.t() @ (de @ gateupB.t())
+        d_gateupB = (gateupA.t() @ X.t()) @ de
+        d_gateupA *= gateupS
+        d_gateupB *= gateupS
+
+        # dX  = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS)
+        # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS)
+
+        gateupW = fast_dequantize(gateupW.t(), gateupW_quant)
+        dX = de @ gateupW.t()
+        del gateupW
+        dX += de @ gateupB.to(dtype).t() @ (gateupS * gateupA.to(dtype).t())
+
+        # gateW, gateW_quant, gateA, gateB, gateS,
+        #  upW,    upW_quant,   upA,   upB,   upS,
+        # downW, downW_quant, downA, downB, downS,
+        return dX.view(batch, seq_len, hd), \
+            None, None, d_gateupA.t(), d_gateupB.t(), None, \
+            None, None, d_downA.t(), d_downB.t(), None, \
+            None, None, # _backward and _forward
+    pass
+pass
+
+
+from .swiglu_ import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
+def apply_lora_mlp_swiglu(self, X):
+    gateupW, gateupW_quant, gateupA, gateupB, gateupS = get_lora_parameters(self.gate_up)
+    downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
+
+    out = LoRA_MLP.apply(X,
+                         gateupW, gateupW_quant, gateupA, gateupB, gateupS,
+                         downW, downW_quant, downA, downB, downS,
+                         swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,)
+    return out
+pass

ixformer_sdk/train/speedformer/layers/fast_lora/swiglu.py

@@ -0,0 +1,106 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def _fg_kernel(e, g, h, n_elements, BLOCK_SIZE: tl.constexpr,):
+    block_idx = tl.program_id(0)
+    offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32)
+    g_row = tl.load(g + offsets, mask=mask, other=0)  # .to(tl.float32)
+
+    # f = e * sigmoid(e)
+    f_row = e_row * tl.sigmoid(e_row)  # e_row / (1 + tl.exp(-e_row))
+    f_row = f_row.to(g_row.dtype)  # Exact copy from HF
+    # h = f * g
+    h_row = f_row * g_row
+
+    # Store h
+    tl.store(h + offsets, h_row, mask=mask)
+
+
+pass
+
+
+def swiglu_fg_kernel(e, g):
+    batch, seq_len, hd = e.shape
+    n_elements = e.numel()
+    h = torch.empty((batch, seq_len, hd), dtype=e.dtype, device="cuda:0")
+    def grid(meta): return (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
+    _fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE=1024,)
+    return h
+
+
+pass
+
+
+@triton.jit
+def _DWf_DW_dfg_kernel(DW, e, g, n_elements, BLOCK_SIZE: tl.constexpr,):
+    """
+    e = e.float()
+    se = 1.0 / (1.0 + torch.exp(-e))
+    f = (se * e).to(dtype)
+    h = f * g
+    df = DW * f
+    dg = DW * g
+    de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
+    """
+    block_idx = tl.program_id(0)
+    offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    DW_row = tl.load(DW + offsets, mask=mask, other=0)  # .to(tl.float32)
+    e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32)
+    g_row = tl.load(g + offsets, mask=mask, other=0)  # .to(tl.float32)
+
+    # e = e.float()
+    # se = 1.0 / (1.0 + torch.exp(-e))
+    se_row = tl.sigmoid(e_row)  # 1.0 / (1.0 + tl.exp(-e_row))
+    # f = (se * e).to(dtype)
+    f_row = se_row * e_row
+    f_row = f_row.to(DW_row.dtype)
+    # h = f * g
+    h_row = f_row * g_row
+    # df = DW * f
+    df_row = DW_row * f_row
+    # dg = DW * g
+    dg_row = DW_row * g_row
+    # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
+    de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row))
+    de_row = de_row.to(DW_row.dtype)
+
+    # Store derivatives in buffers
+    tl.store(DW + offsets, h_row,  mask=mask)  # h  = f * g
+    tl.store(e + offsets, df_row, mask=mask)  # df = DW * f
+    tl.store(g + offsets, de_row, mask=mask)  # de
+
+
+pass
+
+
+def swiglu_DWf_DW_dfg_kernel(DW, e, g):
+    batch_seq_len, hd = e.shape
+    n_elements = e.numel()
+    def grid(meta): return (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
+    _DWf_DW_dfg_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE=1024,)
+    return DW, e, g
+
+
+pass

ixformer_sdk/train/speedformer/layers/fast_lora/swiglu_.py

@@ -0,0 +1,102 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def _fg_kernel(x, h, hd, BLOCK_SIZE : tl.constexpr,):
+    block_idx = tl.program_id(0)
+    offsets0 = block_idx*2*hd + tl.arange(0, BLOCK_SIZE)
+    offsets1 = block_idx*2*hd + hd + tl.arange(0, BLOCK_SIZE)
+    mask = offsets0 < hd
+
+    e_row = tl.load(x + offsets0, mask = mask, other = 0).to(tl.float32)
+    g_row = tl.load(x + offsets1, mask = mask, other = 0)#.to(tl.float32)
+
+    # f = e * sigmoid(e)
+    f_row = e_row * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row))
+    f_row = f_row.to(g_row.dtype) # Exact copy from HF
+    # h = f * g
+    h_row = f_row * g_row
+
+    # Store h
+    tl.store(h + offsets0, h_row, mask = mask)
+pass
+
+
+def swiglu_fg_kernel(x):
+    batch, seq_len, hdx2 = x.shape
+    hd = hdx2 // 2
+    n_rows = batch * seq_len
+    BLOCK_SIZE = triton.next_power_of_2(hd)
+    h = torch.empty((batch, seq_len, hd), dtype = x.dtype, device = "cuda:0")
+
+    _fg_kernel[n_rows,](x, h, hd, BLOCK_SIZE=BLOCK_SIZE)
+    return h
+pass
+
+
+@triton.jit
+def _DWf_DW_dfg_kernel(DW, x, hd, BLOCK_SIZE : tl.constexpr,):
+    """
+    e = e.float()
+    se = 1.0 / (1.0 + torch.exp(-e))
+    f = (se * e).to(dtype)
+    h = f * g
+    df = DW * f
+    dg = DW * g
+    de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
+    """
+    block_idx = tl.program_id(0)
+    offsets0 = block_idx*hd*2 + tl.arange(0, BLOCK_SIZE)
+    offsets1 = block_idx*hd*2 + hd + tl.arange(0, BLOCK_SIZE)
+    mask = BLOCK_SIZE < hd
+
+    DW_row = tl.load(DW + offsets0, mask = mask, other = 0)#.to(tl.float32)
+    e_row  = tl.load(x  + offsets0, mask = mask, other = 0).to(tl.float32)
+    g_row  = tl.load(x  + offsets1, mask = mask, other = 0)#.to(tl.float32)
+
+    # e = e.float()
+    # se = 1.0 / (1.0 + torch.exp(-e))
+    se_row = tl.sigmoid(e_row) # 1.0 / (1.0 + tl.exp(-e_row))
+    # f = (se * e).to(dtype)
+    f_row = se_row * e_row
+    f_row = f_row.to(DW_row.dtype)
+    # h = f * g
+    h_row  =  f_row * g_row
+    # df = DW * f
+    df_row = DW_row * f_row
+    # dg = DW * g
+    dg_row = DW_row * g_row
+    # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype)
+    de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row))
+    de_row = de_row.to(DW_row.dtype)
+
+    # Store derivatives in buffers
+    tl.store(DW + offsets0, h_row,  mask = mask) # h  = f * g
+    tl.store(x  + offsets0, df_row, mask = mask) # df = DW * f
+    tl.store(x  + offsets1, de_row, mask = mask) # de
+pass
+
+
+def swiglu_DWf_DW_dfg_kernel(DW, x):
+    batch_seq_len, hdx2 = x.shape
+    hd = hdx2 // 2
+    BLOCK_SIZE = triton.next_power_of_2(hd)
+    _DWf_DW_dfg_kernel[batch_seq_len, ](DW, x, hd, BLOCK_SIZE=BLOCK_SIZE,)
+    return DW, x
+pass

ixformer_sdk/train/speedformer/layers/fast_lora/utils.py

@@ -0,0 +1,195 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ctypes
+import bitsandbytes as bnb
+from packaging.version import Version
+import torch
+import triton
+MAX_FUSED_SIZE = 65536
+next_power_of_2 = triton.next_power_of_2
+
+# torch.cuda.amp.custom_fwd is deprecated >= 2.4
+if Version(torch.__version__) < Version("2.4.0"):
+    torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
+    torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
+else:
+    torch_amp_custom_fwd = torch.amp.custom_fwd(device_type="cuda")
+    torch_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda")
+pass
+
+
+# tl.math.tanh now is libdevice.tanh
+if Version(triton.__version__) >= Version("3.0.0"):
+    from triton.language.extra import libdevice
+    triton_tanh = libdevice.tanh
+else:
+    import triton.language as tl
+    triton_tanh = tl.math.tanh
+pass
+
+
+def calculate_settings(n):
+    BLOCK_SIZE = next_power_of_2(n)
+    if BLOCK_SIZE > MAX_FUSED_SIZE:
+        raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "
+                           f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
+    num_warps = 4
+    if BLOCK_SIZE >= 32768:
+        num_warps = 32
+    elif BLOCK_SIZE >= 8192:
+        num_warps = 16
+    elif BLOCK_SIZE >= 2048:
+        num_warps = 8
+    return BLOCK_SIZE, num_warps
+
+
+pass
+
+
+get_ptr = bnb.functional.get_ptr
+cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
+cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
+cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
+
+
+def QUANT_STATE(W):
+    return getattr(W, "quant_state", None)
+
+
+pass
+
+
+def get_lora_parameters(proj):
+    # For DPO or disabled adapters
+    base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj)
+    W = base_layer.weight
+
+    if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
+        return W, QUANT_STATE(W), None, None, None
+    pass
+
+    active_adapter = proj.active_adapters[0] if \
+        hasattr(proj, "active_adapters") else proj.active_adapter
+    A = proj.lora_A[active_adapter].weight
+    B = proj.lora_B[active_adapter].weight
+    s = proj.scaling[active_adapter]
+    return W, QUANT_STATE(W), A, B, s
+
+
+pass
+
+
+def get_lora_parameters_bias(proj):
+    # For DPO or disabled adapters
+    base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj)
+    W = base_layer.weight
+    bias = base_layer.bias
+
+    if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
+        return W, QUANT_STATE(W), None, None, None, bias
+    pass
+
+    active_adapter = proj.active_adapters[0] if \
+        hasattr(proj, "active_adapters") else proj.active_adapter
+    A = proj.lora_A[active_adapter].weight
+    B = proj.lora_B[active_adapter].weight
+    s = proj.scaling[active_adapter]
+    return W, QUANT_STATE(W), A, B, s, bias
+
+
+pass
+
+
+def fast_dequantize(W, quant_state=None, out=None):
+    if quant_state is None:
+        return W
+    if type(quant_state) is not list:
+        # New quant_state as a class
+        # https://github.com/TimDettmers/bitsandbytes/pull/763/files
+        absmax = quant_state.absmax
+        shape = quant_state.shape
+        dtype = quant_state.dtype
+        blocksize = quant_state.blocksize
+        offset = quant_state.offset
+        state2 = quant_state.state2
+        absmax2 = state2.absmax
+        code2 = state2.code
+        blocksize2 = state2.blocksize
+    else:
+        # Old quant_state as a list of lists
+        absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state
+        offset, state2 = compressed_stats
+        absmax2, code2, blocksize2, _, _, _, _ = state2
+    pass
+
+    # Create weight matrix
+    if out is None:
+        out = torch.empty(shape, dtype=dtype, device="cuda:0")
+    else:
+        assert (out.shape == shape)
+        assert (out.dtype == dtype)
+
+    # NF4 dequantization of statistics
+    n_elements_absmax = absmax.numel()
+    out_absmax = torch.empty(
+        n_elements_absmax, dtype=torch.float32, device="cuda:0")
+
+    # Do dequantization
+    ptr_out_absmax = get_ptr(out_absmax)
+    cdequantize_blockwise_fp32(
+        get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax,
+        ctypes.c_int(blocksize2), ctypes.c_int(n_elements_absmax)
+    )
+    out_absmax += offset
+
+    fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \
+        cdequantize_blockwise_bf16_nf4
+    fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out),
+       ctypes.c_int(blocksize), ctypes.c_int(out.numel()))
+
+    # Careful returning transposed data
+    is_transposed = (True if W.shape[0] == 1 else False)
+    return out.t() if is_transposed else out
+
+
+pass
+
+
+def matmul_lora(X, W, W_quant, A, B, s, out=None):
+    dtype = X.dtype
+    W = fast_dequantize(W.t(), W_quant)
+
+    if X.dim() == 3:
+        batch, seq_len, d = X.shape
+        X = X.view(-1, X.shape[-1])
+        reshape = True
+    else:
+        reshape = False
+    pass
+
+    out = torch.matmul(X, W, out=out)
+    if W_quant is not None:
+        del W
+
+    if A is not None:
+        # LoRA is enabled
+        A, B = A.t(), B.t()
+        out += (X @ A.to(dtype)) @ (s * B.to(dtype))
+    pass
+
+    return out.view(batch, seq_len, -1) if reshape else out
+
+
+pass

ixformer_sdk/train/speedformer/layers/gpt2/attention.py

@@ -0,0 +1,45 @@
+import torch
+import os
+from einops import rearrange
+from flash_attn import flash_attn_varlen_func
+
+
+@staticmethod
+def replace_flash_attn_forward(self, q, k, v, attention_mask, query_length, dropout=0.0, softmax_scale=None):
+
+    # flash-attn(ixdnn)存在gpt2(118M,338M,738M) shape没适配，只能采用普通版本
+    assert os.getenv('ENABLE_FLASH_ATTENTION_WITH_IXDNN', "1") == '0', "flash-attn should not be use ixdnn version, please set variables" \
+        " in shell \"export ENABLE_FLASH_ATTENTION_WITH_IXDNN=0 \" "
+    assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v)))
+    assert all((i.is_cuda for i in (q, k, v)))
+
+    batch_size, seqlen_q = q.shape[0], q.shape[1]
+    seqlen_k = k.shape[1]
+
+    q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]]
+    cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32,
+                                device=q.device)
+
+    if query_length != 1:
+        # during training q,k,v always have same seqlen
+        assert seqlen_k == seqlen_q
+
+        is_causal = self.is_causal
+        cu_seqlens_k = cu_seqlens_q
+        dropout_p = dropout
+    else:
+        # turn off FA causal mask after first inference autoregressive iteration
+        # only on first autoregressive step q,k,v have same seqlen
+        is_causal = seqlen_q == seqlen_k
+        cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32,
+                                    device=q.device)
+        dropout_p = 0
+
+    output = flash_attn_varlen_func(
+        q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k,
+        dropout_p,
+        softmax_scale=softmax_scale, causal=is_causal
+    )
+    # print(f"{output}")
+    output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+    return output

ixformer_sdk/train/speedformer/layers/lazy/__init__.py

@@ -0,0 +1,6 @@
+from .lazy_init import LazyInitContext, LazyTensor
+
+__all__ = [
+    "LazyInitContext",
+    "LazyTensor",
+]

ixformer_sdk/train/speedformer/layers/lazy/construction.py

@@ -0,0 +1,87 @@
+from contextlib import contextmanager
+from typing import Callable, Dict, Tuple
+
+import torch
+
+__all__ = [
+    "_LEGACY_TENSOR_CONSTRUCTOR",
+    "_NO_META_FACTORY",
+    "_NORMAL_FACTORY",
+    "ConstructorManager",
+]
+
+# reference: https://pytorch.org/cppdocs/notes/tensor_creation.html
+_NORMAL_FACTORY = [
+    "arange",
+    "full",
+    "empty",
+    "linspace",
+    "logspace",
+    "ones",
+    "rand",
+    "randn",
+    "randint",
+    "randperm",
+    "zeros",
+    "tensor",
+]
+
+# factory function that does not support meta tensor backend
+_NO_META_FACTORY = [
+    "eye",
+]
+
+_LEGACY_TENSOR_CONSTRUCTOR = {
+    "FloatTensor": torch.float,
+    "DoubleTensor": torch.double,
+    "HalfTensor": torch.half,
+    "BFloat16Tensor": torch.bfloat16,
+    "ByteTensor": torch.uint8,
+    "CharTensor": torch.int8,
+    "ShortTensor": torch.short,
+    "IntTensor": torch.int,
+    "LongTensor": torch.long,
+    "BoolTensor": torch.bool,
+}
+
+
+class ConstructorManager:
+    # function name: (new, old)
+    overwrites: Dict[str, Tuple[Callable, Callable]] = {}
+    changed: bool = False
+
+    @staticmethod
+    def apply(overwrites: Dict[Callable, Callable]):
+        ConstructorManager.overwrites.clear()
+        ConstructorManager.overwrites.update(overwrites)
+        ConstructorManager.redo()
+
+    @staticmethod
+    def undo():
+        assert ConstructorManager.changed, "No constructor change to undo"
+        for name, (new, old) in ConstructorManager.overwrites.items():
+            setattr(torch, name, old)
+        ConstructorManager.changed = False
+
+    @staticmethod
+    def redo():
+        assert not ConstructorManager.changed, "Constructor already changed"
+        for name, (new, old) in ConstructorManager.overwrites.items():
+            setattr(torch, name, new)
+        ConstructorManager.changed = True
+
+    @staticmethod
+    @contextmanager
+    def disable():
+        enabled = ConstructorManager.changed
+        if enabled:
+            ConstructorManager.undo()
+        yield
+        if enabled:
+            ConstructorManager.redo()
+
+    @staticmethod
+    def clear():
+        if ConstructorManager.changed:
+            ConstructorManager.undo()
+        ConstructorManager.overwrites.clear()

ixformer_sdk/train/speedformer/layers/lazy/lazy_init.py

@@ -0,0 +1,669 @@
+from types import MethodType
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+from packaging import version
+from torch import Tensor
+from torch.nn import Parameter
+from torch.utils._pytree import tree_map
+
+from ixformer.train.speedformer.layers.lazy.construction import ConstructorManager
+from ixformer.train.speedformer.layers.lazy.pretrained import PretrainedManager
+
+# reference: https://pytorch.org/cppdocs/notes/tensor_creation.html
+_NORMAL_FACTORY = [
+    "arange",
+    "full",
+    "empty",
+    "linspace",
+    "logspace",
+    "ones",
+    "rand",
+    "randn",
+    "randint",
+    "randperm",
+    "zeros",
+    "tensor",
+]
+
+# factory function that does not support meta tensor backend
+_NO_META_FACTORY = [
+    "eye",
+]
+
+_EARLY_MATERIALIZED_OPS = ["__getitem__", "split"]
+
+# If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset)
+# without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block.
+# These ops cannot be unwrapped using .data
+_CHANGE_META_OPS = ["_cudnn_rnn_flatten_weight",
+                    "requires_grad_", "__get__", "__set__", "numel", "size", "dim"]
+
+# These ops is not related to tensor value and should not be rerun
+_NO_RERUN_OPS = ["__get__", "numel", "size", "dim"]
+
+_LEGACY_TENSOR_CONSTRUCTOR = {
+    "FloatTensor": torch.float,
+    "DoubleTensor": torch.double,
+    "HalfTensor": torch.half,
+    "BFloat16Tensor": torch.bfloat16,
+    "ByteTensor": torch.uint8,
+    "CharTensor": torch.int8,
+    "ShortTensor": torch.short,
+    "IntTensor": torch.int,
+    "LongTensor": torch.long,
+    "BoolTensor": torch.bool,
+}
+
+# These ops have at least one lazy tensor argument and maybe a scalar argument
+# scalar value should be converted to meta tensor
+# this is a hack for torch 2.0
+_EXPAND_SCALAR_OPS = [
+    "where",
+    "clamp",
+    "clamp_min",
+    "clamp_max",
+    "clamp_",
+    "clamp_min_",
+    "clamp_max_",
+]
+_old_tensor_factory = torch.tensor
+
+_EMPTY_DATA = torch.empty(0)
+
+
+class _MyTensor(Tensor):
+    """This class is only for correctness verification."""
+
+    _pre_op_fn: Callable[["LazyTensor"], None] = lambda *args: None
+
+    default_device: Optional[torch.device] = None
+
+    def __new__(cls, func, *args, concrete_data=None, **kwargs) -> "_MyTensor":
+        cls._pre_op_fn()
+        if concrete_data is not None:
+            # uniform api as LazyTensor
+            data = concrete_data
+        else:
+            kwargs["device"] = cls.default_device
+            data = func(*args, **kwargs)
+        return Tensor._make_subclass(cls, data, require_grad=data.requires_grad)
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        cls._pre_op_fn()
+        return super().__torch_function__(func, types, args, kwargs)
+
+
+def _data_tolist(tensor: torch.Tensor) -> list:
+    """tolist() method is not allowed for a subclass of tensor. Tensor.data returns a Tensor."""
+    return tensor.data.tolist()
+
+
+def _convert_cls(tensor: "LazyTensor", target: torch.Tensor) -> torch.Tensor:
+    """Convert a lazy tensor's class to target's class, with target's data.
+
+    The reason why we change the class of a lazy tensor in-place is that this can easily handle shared modules/parameters, which is common in huggingface models.
+    If we create a new tensor and update the module by ``setattr(module, name, param)``, the shared parameters will not be updated. And we have to track all shared parameters and update them manually.
+
+    Args:
+        tensor (LazyTensor): the LazyTensor to be converted
+        target (torch.Tensor): target tensor
+
+    Returns:
+        torch.Tensor: the converted tensor
+    """
+    cls_to_become = Parameter if isinstance(
+        tensor, Parameter) else torch.Tensor
+    tensor.__class__ = cls_to_become
+    if cls_to_become is Parameter:
+        # to fit UninitializedParameter
+        delattr(tensor, "_is_param")
+    tensor.data = target
+    tensor.requires_grad = target.requires_grad
+    # subclass of torch.Tensor does not have tolist() method
+    # overwrite this method after materialization or distribution
+    tensor.tolist = MethodType(_data_tolist, tensor)
+    return tensor
+
+
+class LazyTensor(torch.Tensor):
+    """A naive implementation of LazyTensor (https://arxiv.org/pdf/2102.13267.pdf).
+
+    Usage:
+        1. Use ``LazyTensor`` instead of ``torch.Tensor``.
+        >>> x = LazyTensor(torch.zeros, 2, 3)
+        >>> x += 1
+        >>> y = x * x
+        >>> y = y.cuda().half()
+        >>> y[0, 0] = 0
+        >>> y = y.materialize()     # materialize the tensor
+        >>> print(y)
+        tensor([[0., 1., 1.],
+                [1., 1., 1.]], device='cuda:0', dtype=torch.float16)
+
+    Warnings:
+        1. Cases that ``LazyTensor`` can't deal with.
+        >>> x = LazyTensor(torch.ones, 2, 3)
+        >>> x[0, 0] = -x[0, 0]    # this will cause infinite recursion
+        >>> y = x.clone()
+        >>> x.add_(1) # modifying origin tensor after cloning leads to wrong materialization
+        >>> z = x.tolist()
+        >>> x.zeros_() # modifying origin tensor after cloning tolist is not allowed
+        >>> nn.utils.weight_norm(self.conv, name="weight", dim=2) # applying weight norm on a lazy tensor is not allowed
+
+
+        2. Cases that ``LazyTensor`` becomes eager (early materialization).
+        >>> b = a[:, 2:]  # get a slice of a lazy tensor triggers early materialization
+        >>> chunks = a.split(3)  # this also triggers early materialization
+        >>> x.data = torch.rand(2, 3) # directly setting data of a lazy tensor triggers early materialization
+
+    """
+
+    _repr = True
+    _meta_data: Optional[torch.Tensor] = None  # shape, dtype, device
+    _pre_op_fn: Callable[["LazyTensor"], None] = lambda *args: None
+
+    default_device: Optional[torch.device] = None
+    _device: torch.device  # fake device of mate tensor
+
+    @staticmethod
+    def __new__(cls, func, *args, meta_data=None, concrete_data=None, **kwargs):
+        # tips for torch 2.0:
+        # torch 2.0 disables torch dispatch for subclass of tensor
+        # MetaTensor is cannot be used
+        # Now lazy tensor contains device injection and meta tensor
+        if concrete_data is not None:
+            # some ops don't support meta backend and should have concrete data
+            elem = concrete_data
+        else:
+            if meta_data is None:
+                with ConstructorManager.disable():
+                    # to disable create lazy tensor in inner ops, this is a hack for torch 2.0
+                    meta_data = func(*args, **{**kwargs, "device": "meta"})
+            elem = meta_data
+        # As a meta tensor cannot be modified __class__ to torch.Tensor, we should use an empty real tensor here
+        r = torch.Tensor._make_subclass(
+            cls, _EMPTY_DATA, require_grad=elem.requires_grad)
+        r._meta_data = meta_data
+
+        return r
+
+    def __init__(self, func, *args, meta_data=None, concrete_data=None, **kwargs):
+        self._device = torch.device(kwargs.get("device", None) or "cpu")
+        if func.__name__ in _NORMAL_FACTORY:
+            kwargs = {**kwargs, "device": LazyTensor.default_device}
+        self._factory_method = (func, args, kwargs)  # (func, args, kwargs)
+        self._op_buffer = []  # (func, args, kwargs, replace)
+        # materialized data
+        self._materialized_data: Optional[torch.Tensor] = concrete_data
+
+    @property
+    def device(self) -> torch.device:
+        return self._materialized_data.device if self._materialized_data is not None else self._device
+
+    def __repr__(self):
+        return f"LazyTensor(..., size={tuple(self.shape)}, device='{self.device}', dtype={self.dtype})"
+
+    def materialize(self) -> torch.Tensor:
+        """Materialize the ``LazyTensor`` to ``torch.Tensor`` by modifying __class__ (inplace).
+
+        Returns:
+            torch.Tensor: The materialized tensor (self).
+        """
+        target = self._materialize_data()
+        self.clean()
+        return _convert_cls(self, target)
+
+    def clean(self) -> None:
+        """Clean all stored operations, meta data and materialized data, which prevents memory leaking. This should be called after all tensors are materialized."""
+        delattr(self, "_factory_method")
+        delattr(self, "_op_buffer")
+        delattr(self, "_materialized_data")
+        delattr(self, "_meta_data")
+
+    @staticmethod
+    def _replace_with_materialized(x):
+        if isinstance(x, LazyTensor):
+            return x._materialize_data()
+        return x
+
+    def _materialize_data(self) -> torch.Tensor:
+        # self._materialized_data should be generated after the first call of this function
+        if self._materialized_data is None:
+            # apply factory method
+            func, args, kwargs = self._factory_method
+            # apply cached sequence
+            self._pre_op_fn()
+
+            init_val = func(
+                *tree_map(self._replace_with_materialized, args), **tree_map(self._replace_with_materialized, kwargs)
+            )
+
+            self._materialized_data = self._rerun_ops(init_val)
+        return self._materialized_data
+
+    def _rerun_ops(self, target=None) -> torch.Tensor:
+        """Do lazy execution by rerunning all (stored) related operations.
+
+        Args:
+            target (torc.Tensor, optional): Intial value of the target tensor (self). Defaults to None.
+        """
+
+        def replace(x):
+            if x is self:
+                return target
+            elif isinstance(x, LazyTensor):
+                return x._materialize_data()
+            return x
+
+        packed = None
+
+        for func, args, kwargs in self._op_buffer:
+            if func == torch.Tensor.requires_grad_:
+                packed = func, args, kwargs  # requires grad should be set at last
+            else:
+                self._pre_op_fn()
+                o = func(*tree_map(replace, args), **tree_map(replace, kwargs))
+                # if func returns non-Tensor, discard the value
+                target = o if isinstance(o, torch.Tensor) else target
+
+        # super-dainiu: set requires_grad after all inplace-ops are done
+        if packed is not None:
+            func, args, kwargs = packed
+            func(*tree_map(replace, args), **tree_map(replace, kwargs))
+
+        return target
+
+    # cache everything with __torch_function__
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+        if func.__name__ in _EARLY_MATERIALIZED_OPS:
+            # These OPs cannot be lazy and related tensors should be early materialized
+            tree_map(cls._replace_with_materialized, args)
+            tree_map(cls._replace_with_materialized, kwargs)
+        is_inplace: bool = (
+            func.__name__.endswith("_")
+            and not (func.__name__.endswith("__"))
+            or func.__name__ in ("__setitem__", "__set__")
+        )
+
+        is_change_meta_op: bool = func.__name__ in _CHANGE_META_OPS
+
+        if isinstance(func, torch._C.ScriptMethod):
+            # FIXME(ver217): torch script functions are not verified
+
+            target = None
+
+            def unwrap(x):
+                if isinstance(x, LazyTensor):
+                    return x._meta_data
+                return x
+
+            target: LazyTensor = args[0].clone()
+            target._op_buffer.append((func, args, kwargs))
+            target._meta_data = getattr(target._meta_data, func.name)(
+                *tree_map(unwrap, args[1:]), **tree_map(unwrap, kwargs)
+            )
+            return target
+        else:
+            meta_to_lazy = {}
+
+            def unwrap(x):
+                if isinstance(x, LazyTensor):
+                    if x._materialized_data is not None:
+                        # for early materialized tensor, use its materialized data directly
+                        return x._materialized_data if is_change_meta_op else x._materialized_data.data
+                    t = x if is_inplace else x.clone()
+                    if func.__name__ not in _NO_RERUN_OPS:
+                        t._op_buffer.append((func, args, kwargs))
+                    meta = x._meta_data if is_change_meta_op else x._meta_data.data
+                    meta_to_lazy[meta] = t
+                    return meta
+                elif (
+                    version.parse(torch.__version__) >= version.parse("2.0.0")
+                    and func.__name__ in _EXPAND_SCALAR_OPS
+                    and not isinstance(x, torch.Tensor)
+                ):
+                    return _old_tensor_factory(x, device="meta")
+                return x
+
+            def wrap(y, i=None):
+                if isinstance(y, torch.Tensor):
+                    if y.is_meta:
+                        if y in meta_to_lazy:
+                            # inplace op, just return origin lazy tensor
+                            return meta_to_lazy[y]
+                        else:
+                            # out of place op, create new lazy tensor
+                            fn = lambda *a, **kw: func(*a, **
+                                                       kw) if i is None else func(*a, **kw)[i]
+                            fn.__name__ = func.__name__
+                            lazy_y = LazyTensor(
+                                fn, *args, meta_data=y, **kwargs)
+                            return lazy_y
+                    else:
+                        # for early materialized tensor
+                        return LazyTensor(lambda: None, concrete_data=y)
+                return y
+
+            cls._pre_op_fn()
+            with ConstructorManager.disable():
+                # to disable create lazy tensor in inner ops, this is a hack for torch 2.0
+                o = func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))
+            if isinstance(o, (tuple, list)):
+                return type(o)(wrap(y, i=i) for i, y in enumerate(o))
+            return wrap(o)
+
+    def to(self, *args, **kwargs) -> torch.Tensor:
+        if self._materialized_data is not None:
+            return LazyTensor(lambda: None, concrete_data=self._materialized_data.to(*args, **kwargs))
+
+        device = None
+
+        def replace(x):
+            nonlocal device
+            if isinstance(x, (str, int, torch.device)) and not isinstance(x, bool):
+                device = x
+                return torch.device("meta")
+            return x
+
+        meta_data = self._meta_data.to(
+            *tree_map(replace, args), **tree_map(replace, kwargs))
+
+        if meta_data is self._meta_data and device == self.device:
+            return self
+
+        def factory_fn(t: torch.Tensor, **kw):
+            return t.to(*args, **kwargs)
+
+        return LazyTensor(factory_fn, self, meta_data=meta_data, device=device)
+
+    def cpu(self, memory_format: torch.memory_format = torch.preserve_format):
+        return self.to(device=torch.device("cpu"), memory_format=memory_format)
+
+    def cuda(self, device=None, non_blocking=False, memory_format: torch.memory_format = torch.preserve_format):
+        device = torch.device(device or "cuda")
+        return self.to(device=device, non_blocking=non_blocking, memory_format=memory_format)
+
+    def clone(self) -> "LazyTensor":
+        def factory_fn(t: torch.Tensor, **kw):
+            # if self is materialized, return self
+            return t.clone()
+
+        target = LazyTensor(factory_fn, self, meta_data=self._meta_data)
+
+        return target
+
+    def detach(self) -> Tensor:
+        return self
+
+    def __deepcopy__(self, memo):
+        if not self.is_leaf:
+            raise RuntimeError(
+                "Only Tensors created explicitly by the user "
+                "(graph leaves) support the deepcopy protocol at the moment"
+            )
+        if id(self) in memo:
+            return memo[id(self)]
+
+        def factory_fn(t: torch.Tensor, **kw):
+            # if self is materialized, return self
+            return _copy_tensor(t, t.requires_grad)
+
+        if self._materialized_data is not None:
+            # self is early materialized
+            copied = _copy_tensor(self._materialized_data, self.requires_grad)
+            target = LazyTensor(lambda: None, concrete_data=copied)
+        else:
+            target = LazyTensor(factory_fn, self, meta_data=self._meta_data)
+
+        if isinstance(self, Parameter):
+            # hack isinstance check of parameter
+            target._is_param = True
+
+        memo[id(self)] = target
+        return target
+
+    @property
+    def data(self):
+        return self
+
+    @data.setter
+    def data(self, other: "LazyTensor"):
+        """This is sightly different from oringinal `data` setter.
+
+        E.g.:
+            >>> a = torch.randn(3, 3) # a is a Tensor
+            >>> b = torch.rand(2, 2)
+            >>> a.data = b
+            >>> b.add_(1)   # this will affect a
+            >>> x = torch.randn(3, 3) # x is a LazyTensor
+            >>> y = torch.rand(2, 2) # y is a LazyTensor
+            >>> x.data = y
+            >>> y.add_(1)   # this will not affect x
+
+        """
+        if other is self:
+            return
+
+        def replace(x):
+            if x is other:
+                return self
+            return x
+
+        for func, args, kwargs in [other._factory_method, *other._op_buffer]:
+            self._op_buffer.append(
+                (func, tree_map(replace, args), tree_map(replace, kwargs)))
+
+    def tolist(self) -> list:
+        # Though self.__class__ is modified to torch.Tensor, in C++ side, it is still a subclass of torch.Tensor
+        # And subclass of torch.Tensor does not have tolist() method
+        t = self._materialize_data()
+        return t.tolist()
+
+    def __hash__(self):
+        return id(self)
+
+    def __rpow__(self, other):
+        dtype = torch.result_type(self, other)
+        return torch.tensor(other, dtype=dtype, device=self.device) ** self
+
+
+class LazyInitContext:
+    """Context manager for lazy initialization. Enables initializing the model without allocating real memory.
+
+    Args:
+        tensor_cls (Union[_MyTensor, LazyTensor], optional): This is only for test. Defaults to LazyTensor.
+        default_device (Optional[Union[torch.device, str, int]], optional): Defalt device for initialization.
+            If it's cuda, initilization will be accelerated, but cuda memory will be allocated. By default, it's cpu.
+            Defaults to None.
+    """
+
+    _replaced: bool = False
+
+    def __init__(
+        self,
+        tensor_cls: Union[_MyTensor, LazyTensor] = LazyTensor,
+        default_device: Optional[Union[torch.device, str, int]] = None,
+    ):
+        assert tensor_cls is LazyTensor or tensor_cls is _MyTensor
+        self.tensor_cls = tensor_cls
+        self.old_default_device = LazyTensor.default_device
+        self.default_device = default_device
+
+    def __enter__(self):
+        if LazyInitContext._replaced:
+            raise RuntimeError(f"LazyInitContext is not reentrant")
+        LazyInitContext._replaced = True
+        self.old_default_device = self.tensor_cls.default_device
+        self.tensor_cls.default_device = self.default_device
+
+        def wrap_factory_method(target):
+            # factory functions (eg. torch.empty())
+            def wrapper(*args, **kwargs):
+                return self.tensor_cls(target, *args, **kwargs)
+
+            return wrapper, target
+
+        def wrap_factory_like_method(orig_target, target):
+            # factory_like functions (eg. torch.empty_like())
+            def wrapper(*args, **kwargs):
+                orig_t = args[0]
+                return self.tensor_cls(
+                    orig_target, *orig_t.shape, *args[1:], device=orig_t.device, dtype=orig_t.dtype, **kwargs
+                )
+
+            return wrapper, target
+
+        def wrap_legacy_constructor(target, dtype):
+            # legacy constructor (e.g. torch.LongTensor())
+            def wrapper(*args, **kwargs):
+                if len(args) == 1 and isinstance(args[0], torch.Tensor):
+                    # (Tensor other)
+                    return args[0]
+                elif len(args) == 1:
+                    # (object data, *, torch.device device)
+                    kwargs = {**kwargs, "dtype": dtype}
+                    replaced, orig = self.overrides["tensor"]
+                    return replaced(*args, **kwargs)
+                elif _is_int_tuple(args):
+                    # (tuple of ints size, *, torch.device device)
+                    kwargs = {**kwargs, "dtype": dtype}
+                    replaced, orig = self.overrides["empty"]
+                    return replaced(*args, **kwargs)
+                else:
+                    raise TypeError(
+                        f"new() received an invalid combination of arguments - got {tuple(type(x) for x in args)}, but expected one of:\n * (Tensor other)\n * (tuple of ints size, *, torch.device device)\n * (object data, *, torch.device device)"
+                    )
+
+            return wrapper, target
+
+        def wrap_no_meta_factory(target):
+            # factory functions which don't support meta tensor backend
+            def wrapper(*args, **kwargs):
+                tensor = target(*args, **kwargs)
+                return self.tensor_cls(lambda: None, concrete_data=tensor)
+
+            return wrapper, target
+
+        overrides = {
+            target: wrap_factory_method(getattr(torch, target))
+            for target in _NORMAL_FACTORY
+            if callable(getattr(torch, target, None))
+        }
+
+        overrides.update(
+            {
+                target + "_like": wrap_factory_like_method(getattr(torch, target), getattr(torch, target + "_like"))
+                for target in _NORMAL_FACTORY
+                if callable(getattr(torch, target + "_like", None))
+            }
+        )
+
+        overrides.update(
+            {
+                target: wrap_legacy_constructor(getattr(torch, target), dtype)
+                for target, dtype in _LEGACY_TENSOR_CONSTRUCTOR.items()
+                if callable(getattr(torch, target, None))
+            }
+        )
+
+        overrides.update(
+            {
+                target: wrap_no_meta_factory(getattr(torch, target))
+                for target in _NO_META_FACTORY
+                if callable(getattr(torch, target, None))
+            }
+        )
+
+        ConstructorManager.apply(overrides)
+        PretrainedManager.inject()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.tensor_cls.default_device = self.old_default_device
+        LazyInitContext._replaced = False
+        ConstructorManager.clear()
+        PretrainedManager.recover()
+
+    @staticmethod
+    def materialize(module: nn.Module, verbose: bool = False) -> nn.Module:
+        """Initialize all ``Parameter`` from ``LazyTensor``. This function will modify the module in-place.
+
+        Args:
+            module (nn.Module): Target ``nn.Module``
+            verbose (bool): Whether to print lazy initialization rate. Defaults to False.
+        """
+
+        def apply_fn(name: str, p: LazyTensor):
+            p.materialize()
+
+        return _apply_to_lazy_module(module, apply_fn, verbose)
+
+
+def _apply_to_lazy_module(
+    module: nn.Module, apply_fn: Callable[[str, torch.Tensor], None], verbose: bool = False
+) -> nn.Module:
+    if verbose:
+        # verbose info
+        param_cnt = 0
+        param_lazy_cnt = 0
+        buf_cnt = 0
+        buf_lazy_cnt = 0
+        total_numel = 0
+        non_lazy_numel = 0
+
+    for name, p in module.named_parameters():
+        if verbose:
+            param_cnt += 1
+            total_numel += p.numel()
+            if getattr(p, "_materialized_data", False) is None:
+                # if no _materialized_data attr, the tensor is not lazy
+                param_lazy_cnt += 1
+            else:
+                non_lazy_numel += p.numel()
+        if isinstance(p, LazyTensor):
+            apply_fn(name, p)
+
+    for name, buf in module.named_buffers():
+        if verbose:
+            buf_cnt += 1
+            total_numel += buf.numel()
+            if getattr(buf, "_materialized_data", False) is None:
+                # if no _materialized_data attr, the tensor is not lazy
+                buf_lazy_cnt += 1
+            else:
+                non_lazy_numel += buf.numel()
+        if isinstance(buf, LazyTensor):
+            apply_fn(name, buf)
+
+    # if verbose:
+    #     non_lazy_numel_ratio = non_lazy_numel / total_numel * 100 if non_lazy_numel != 0 else 0
+    #     logger = get_dist_logger()
+    #     logger.info(f"Param lazy rate: {param_lazy_cnt}/{param_cnt}", ranks=[0])
+    #     logger.info(f"Buffer lazy rate: {buf_lazy_cnt}/{buf_cnt}", ranks=[0])
+    #     logger.info(
+    #         f"Non lazy numel: {non_lazy_numel} ({non_lazy_numel/1024**2:.3f} M), ratio: {non_lazy_numel_ratio}%",
+    #         ranks=[0],
+    #     )
+
+    return module
+
+
+def _is_int_tuple(args) -> bool:
+    if not isinstance(args, tuple):
+        return False
+    for x in args:
+        if not isinstance(x, int):
+            return False
+    return True
+
+
+def _copy_tensor(tensor: Tensor, requires_grad: bool) -> Tensor:
+    copied = tensor.data.clone()
+    copied.requires_grad = requires_grad
+    return copied

ixformer_sdk/train/speedformer/layers/lazy/pretrained.py

@@ -0,0 +1,318 @@
+import os
+from typing import Callable, Optional, Union
+
+import torch
+from torch.nn import Module
+
+
+class PretrainedManager:
+    old_from_pretrained: Optional[Callable] = None
+
+    @staticmethod
+    def inject() -> None:
+        try:
+            from transformers.modeling_utils import PreTrainedModel
+        except ImportError:
+            return
+        # recover bound method to plain function
+        PretrainedManager.old_from_pretrained = PreTrainedModel.from_pretrained.__func__
+        PreTrainedModel.from_pretrained = new_from_pretrained
+
+    @staticmethod
+    def recover() -> None:
+        try:
+            from transformers.modeling_utils import PreTrainedModel
+        except ImportError:
+            return
+        # convert plain function to class method
+        PreTrainedModel.from_pretrained = classmethod(
+            PretrainedManager.old_from_pretrained)
+        PretrainedManager.old_from_pretrained = None
+
+
+@classmethod
+def new_from_pretrained(
+    cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs
+) -> Module:
+    from transformers import GenerationConfig
+    from transformers.configuration_utils import PretrainedConfig
+    from transformers.modeling_utils import (
+        ContextManagers,
+        _add_variant,
+        cached_file,
+        download_url,
+        has_file,
+        is_offline_mode,
+        is_remote_url,
+        no_init_weights,
+    )
+    from transformers.utils import (
+        SAFE_WEIGHTS_INDEX_NAME,
+        SAFE_WEIGHTS_NAME,
+        WEIGHTS_INDEX_NAME,
+        WEIGHTS_NAME,
+        is_safetensors_available,
+        logging,
+    )
+
+    logger = logging.get_logger(__name__)
+
+    config = kwargs.pop("config", None)
+    cache_dir = kwargs.pop("cache_dir", None)
+    force_download = kwargs.pop("force_download", False)
+    resume_download = kwargs.pop("resume_download", False)
+    proxies = kwargs.pop("proxies", None)
+    local_files_only = kwargs.pop("local_files_only", False)
+    use_auth_token = kwargs.pop("use_auth_token", None)
+    revision = kwargs.pop("revision", None)
+    _ = kwargs.pop("mirror", None)
+    from_pipeline = kwargs.pop("_from_pipeline", None)
+    from_auto_class = kwargs.pop("_from_auto", False)
+    _fast_init = kwargs.pop("_fast_init", True)
+    torch_dtype = kwargs.pop("torch_dtype", None)
+    subfolder = kwargs.pop("subfolder", "")
+    commit_hash = kwargs.pop("_commit_hash", None)
+    variant = kwargs.pop("variant", None)
+    use_safetensors = kwargs.pop(
+        "use_safetensors", None if is_safetensors_available() else False)
+
+    if len(kwargs) > 0:
+        logger.warning(f"Below kwargs may be ignored: {list(kwargs.keys())}")
+
+    from_pt = True
+
+    user_agent = {"file_type": "model", "framework": "pytorch",
+                  "from_auto_class": from_auto_class}
+    if from_pipeline is not None:
+        user_agent["using_pipeline"] = from_pipeline
+
+    if is_offline_mode() and not local_files_only:
+        logger.info("Offline mode: forcing local_files_only=True")
+        local_files_only = True
+
+    # Load config if we don't provide a configuration
+    if not isinstance(config, PretrainedConfig):
+        config_path = config if config is not None else pretrained_model_name_or_path
+        config, model_kwargs = cls.config_class.from_pretrained(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            use_auth_token=use_auth_token,
+            revision=revision,
+            subfolder=subfolder,
+            _from_auto=from_auto_class,
+            _from_pipeline=from_pipeline,
+            **kwargs,
+        )
+    else:
+        model_kwargs = kwargs
+
+    if commit_hash is None:
+        commit_hash = getattr(config, "_commit_hash", None)
+
+    # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the
+    # index of the files.
+
+    if pretrained_model_name_or_path is not None:
+        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+        is_local = os.path.isdir(pretrained_model_name_or_path)
+        if is_local:
+            if use_safetensors is not False and os.path.isfile(
+                os.path.join(pretrained_model_name_or_path, subfolder,
+                             _add_variant(SAFE_WEIGHTS_NAME, variant))
+            ):
+                # Load from a safetensors checkpoint
+                archive_file = os.path.join(
+                    pretrained_model_name_or_path, subfolder, _add_variant(
+                        SAFE_WEIGHTS_NAME, variant)
+                )
+            elif use_safetensors is not False and os.path.isfile(
+                os.path.join(pretrained_model_name_or_path, subfolder,
+                             _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant))
+            ):
+                # Load from a sharded safetensors checkpoint
+                archive_file = os.path.join(
+                    pretrained_model_name_or_path, subfolder, _add_variant(
+                        SAFE_WEIGHTS_INDEX_NAME, variant)
+                )
+            elif os.path.isfile(
+                os.path.join(pretrained_model_name_or_path,
+                             subfolder, _add_variant(WEIGHTS_NAME, variant))
+            ):
+                # Load from a PyTorch checkpoint
+                archive_file = os.path.join(
+                    pretrained_model_name_or_path, subfolder, _add_variant(
+                        WEIGHTS_NAME, variant)
+                )
+            elif os.path.isfile(
+                os.path.join(pretrained_model_name_or_path, subfolder,
+                             _add_variant(WEIGHTS_INDEX_NAME, variant))
+            ):
+                # Load from a sharded PyTorch checkpoint
+                archive_file = os.path.join(
+                    pretrained_model_name_or_path, subfolder, _add_variant(
+                        WEIGHTS_INDEX_NAME, variant)
+                )
+            else:
+                raise EnvironmentError(
+                    f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory"
+                    f" {pretrained_model_name_or_path}."
+                )
+        elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
+            archive_file = pretrained_model_name_or_path
+            is_local = True
+        elif is_remote_url(pretrained_model_name_or_path):
+            filename = pretrained_model_name_or_path
+            resolved_archive_file = download_url(pretrained_model_name_or_path)
+        else:
+            # set correct filename
+            if use_safetensors is not False:
+                filename = _add_variant(SAFE_WEIGHTS_NAME, variant)
+            else:
+                filename = _add_variant(WEIGHTS_NAME, variant)
+
+            try:
+                # Load from URL or cache if already cached
+                cached_file_kwargs = {
+                    "cache_dir": cache_dir,
+                    "force_download": force_download,
+                    "proxies": proxies,
+                    "resume_download": resume_download,
+                    "local_files_only": local_files_only,
+                    "use_auth_token": use_auth_token,
+                    "user_agent": user_agent,
+                    "revision": revision,
+                    "subfolder": subfolder,
+                    "_raise_exceptions_for_missing_entries": False,
+                    "_commit_hash": commit_hash,
+                }
+                resolved_archive_file = cached_file(
+                    pretrained_model_name_or_path, filename, **cached_file_kwargs)
+
+                # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None
+                # result when internet is up, the repo and revision exist, but the file does not.
+                if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):
+                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
+                    resolved_archive_file = cached_file(
+                        pretrained_model_name_or_path,
+                        _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),
+                        **cached_file_kwargs,
+                    )
+                    if resolved_archive_file is not None:
+                        pass
+                    elif use_safetensors:
+                        raise EnvironmentError(
+                            f" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`."
+                        )
+                    else:
+                        # This repo has no safetensors file of any kind, we switch to PyTorch.
+                        filename = _add_variant(WEIGHTS_NAME, variant)
+                        resolved_archive_file = cached_file(
+                            pretrained_model_name_or_path, filename, **cached_file_kwargs
+                        )
+                if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):
+                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
+                    resolved_archive_file = cached_file(
+                        pretrained_model_name_or_path,
+                        _add_variant(WEIGHTS_INDEX_NAME, variant),
+                        **cached_file_kwargs,
+                    )
+                    if resolved_archive_file is not None:
+                        pass
+                if resolved_archive_file is None:
+                    # Otherwise, maybe there is a TF or Flax model file.  We try those to give a helpful error
+                    # message.
+                    has_file_kwargs = {
+                        "revision": revision,
+                        "proxies": proxies,
+                        "use_auth_token": use_auth_token,
+                    }
+                    if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):
+                        raise EnvironmentError(
+                            f"{pretrained_model_name_or_path} does not appear to have a file named"
+                            f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant"
+                            f" {variant}. Use `variant=None` to load this model from those weights."
+                        )
+                    else:
+                        raise EnvironmentError(
+                            f"{pretrained_model_name_or_path} does not appear to have a file named"
+                            f" {_add_variant(WEIGHTS_NAME, variant)}"
+                        )
+            except EnvironmentError:
+                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
+                # to the original exception.
+                raise
+            except Exception:
+                # For any other exception, we throw a generic error.
+                raise EnvironmentError(
+                    f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
+                    " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
+                    f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
+                    f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}."
+                )
+
+        if is_local:
+            logger.info(f"loading weights file {archive_file}")
+            resolved_archive_file = archive_file
+        else:
+            logger.info(
+                f"loading weights file {filename} from cache at {resolved_archive_file}")
+    else:
+        resolved_archive_file = None
+
+    if from_pt:
+        # set dtype to instantiate the model under:
+        # 1. If torch_dtype is not None, we use that dtype
+        dtype_orig = None
+
+        if torch_dtype is not None:
+            if not isinstance(torch_dtype, torch.dtype):
+                raise ValueError(
+                    f"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}")
+            dtype_orig = cls._set_default_torch_dtype(torch_dtype)
+
+    config.name_or_path = pretrained_model_name_or_path
+
+    # Instantiate model.
+    init_contexts = [no_init_weights(_enable=_fast_init)]
+
+    with ContextManagers(init_contexts):
+        model = cls(config, *model_args, **model_kwargs)
+
+    if from_pt:
+        # restore default dtype
+        if dtype_orig is not None:
+            torch.set_default_dtype(dtype_orig)
+
+    # make sure token embedding weights are still tied if needed
+    model.tie_weights()
+
+    # Set model in evaluation mode to deactivate DropOut modules by default
+    model.eval()
+
+    # If it is a model with generation capabilities, attempt to load the generation config
+    if model.can_generate():
+        try:
+            model.generation_config = GenerationConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
+                revision=revision,
+                subfolder=subfolder,
+                _from_auto=from_auto_class,
+                _from_pipeline=from_pipeline,
+                **kwargs,
+            )
+        except (OSError, TypeError):
+            logger.info(
+                "Generation config file not found, using a generation config created from the model config.")
+
+    return model

ixformer_sdk/train/speedformer/layers/llama/attention.py

@@ -0,0 +1,186 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig
+from ixformer.train.speedformer.models.llama.modeling_llama import LlamaFlashAttention2
+from transformers import Cache
+from transformers.utils import logging
+
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+
+from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb
+from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
+
+
+class BaseLlamaAttention(LlamaFlashAttention2):
+    """ 
+    加这个层的原因：1.当原模型中使用的是torch nvtive的attention，强制替换成flash_attn; 2.优化rope 
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.config.rope_scaling is None:
+            self.rotary_emb = RotaryEmbedding(self.head_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+        bsz, q_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # fused_apply_rotary_pos_emb need qk to be in "sbhd"
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim).transpose(1, 0).contiguous()
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 0).contiguous()
+        value_states = value_states.view(
+            bsz, q_len, self.num_heads, self.head_dim)
+
+        kv_seq_len = key_states.shape[0]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[0]
+
+        emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32)
+        query_states = fused_apply_rotary_pos_emb(query_states, emb)
+        key_states = fused_apply_rotary_pos_emb(key_states, emb)
+
+        # kv cache staff
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=0)
+            value_states = torch.cat([past_key_value[1], value_states], dim=0)
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        dropout_rate = self.attention_dropout if self.training else 0.0
+
+        # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa
+        if attention_mask is None:  # flash-attn
+            query_states = query_states.transpose(0, 1).contiguous()
+            key_states = key_states.transpose(0, 1).contiguous()
+        else:  # sdpa
+            query_states = query_states.permute(1, 2, 0, 3).contiguous()
+            key_states = key_states.permute(1, 2, 0, 3).contiguous()
+            value_states = value_states.transpose(1, 2).contiguous()
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            # Handle the case where the model is quantized
+            if hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(
+            bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        for now, if attention_mask is none, flash-attn has better performance than torch.nn.functional.scaled_dot_product_attention;
+        if attention_mask is not none, torch.nn.functional.scaled_dot_product_attention works better
+        so sdpa and flash-attn is perfered according to attention_mask
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        # Contains at least one padding token in the sequence
+        # if attention_mask is not None:
+        if attention_mask is not None:
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+                is_causal=self.is_causal and attention_mask is None and query_length > 1,
+            )
+            attn_output = attn_output.transpose(1, 2).contiguous()
+
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal
+            )
+
+        return attn_output
+
+
+class LlamaAttention(BaseLlamaAttention):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "LlamaAttention is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to LlamaAttention module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        # LazyInitContext.materialize(module)
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        config = getattr(module, "config")
+        layer_idx = getattr(module, "layer_idx", None)
+
+        attention = BaseLlamaAttention(
+            config=config,
+            layer_idx=layer_idx,
+        )
+
+        attention.q_proj.weight = module.q_proj.weight
+        attention.k_proj.weight = module.k_proj.weight
+        attention.v_proj.weight = module.v_proj.weight
+        attention.o_proj.weight = module.o_proj.weight
+
+        if config.attention_bias:
+            attention.q_proj.bias = module.q_proj.bias
+            attention.k_proj.bias = module.k_proj.bias
+            attention.v_proj.bias = module.v_proj.bias
+            attention.o_proj.bias = module.o_proj.bias
+        return attention

ixformer_sdk/train/speedformer/layers/llama/llama_method.py

@@ -0,0 +1,224 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ixformer.train.speedformer.models.llama.modeling_llama import LlamaModel
+from ixformer.train.speedformer.models.llama.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa
+from ixformer.train.speedformer.layers.cross_entropy_loss import fast_cross_entropy_loss as CrossEntropyLoss
+from transformers.utils import logging
+from transformers.cache_utils import Cache, DynamicCache
+
+logger = logging.get_logger(__name__)
+
+
+def LlamaModel_forward():
+    from transformers.modeling_outputs import BaseModelOutputWithPast
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError(
+                "You have to specify either input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        past_key_values_length = 0
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(
+                    past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(
+                seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if attention_mask is not None:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache(
+            ) if use_legacy_cache else next_decoder_cache
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    return forward
+
+
+def LlamaForCausalLM_forward():
+    from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings
+    from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, CausalLMOutputWithPast, _CONFIG_FOR_DOC
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(
+                self.vocab_size // self.config.pretraining_tp, dim=0)
+            logits = [F.linear(hidden_states, lm_head_slices[i])
+                      for i in range(self.config.pretraining_tp)]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    return forward

ixformer_sdk/train/speedformer/layers/llama/mlp.py

@@ -0,0 +1,55 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+import ixformer.train.functions as F
+from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig
+from ixformer.train.speedformer.models.llama.modeling_llama import LlamaMLP
+from transformers import Cache
+from transformers.utils import logging
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+
+class BaseLlamaMLP(LlamaMLP):
+    """
+    这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x)))
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gate_up = nn.Linear(
+            self.hidden_size, self.intermediate_size * 2, bias=False)
+        del self.gate_proj, self.up_proj
+        del self.act_fn
+
+    def forward(self, x):
+        res = self.gate_up(x)
+        down_proj = self.down_proj(F.swiglu(res))
+        return down_proj
+
+
+class IXFLlamaMLP(BaseLlamaMLP):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "IXFLlamaMLP is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to IXFLlamaMLP module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        LazyInitContext.materialize(module)
+
+        config = getattr(module, "config")
+
+        mlp = BaseLlamaMLP(config=config)
+
+        mlp.gate_up.weight.data = torch.concat(
+            (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0)
+        mlp.down_proj.weight.data = module.down_proj.weight.data
+
+        return mlp

ixformer_sdk/train/speedformer/layers/normalization.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+import warnings
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+import ixformer.functions as ixff
+from ixformer.train.functions import FusedRMSNorm as ixf_FusedRMSNorm
+from apex.normalization.fused_layer_norm import FusedRMSNorm as apex_FusedRMSNorm
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+
+class BaseLayerNorm(ABC):
+    @abstractmethod
+    def from_native_module(module: nn.Module, sp_partial_derived: bool = False):
+        """
+        Convert a native PyTorch layer normalization module to a specific layer normalization module,
+        and optionally mark parameters for gradient aggregation.
+
+        Args:
+            module (nn.Module): The native PyTorch layer normalization module to be converted.
+            sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism.
+
+        Returns:
+            nn.Module: The specific layer normalization module.
+
+        Raises:
+            AssertionError: If the provided module is not an instance of the supported layer normalization type.
+        """
+
+
+class IXFFusedRMSNorm(BaseLayerNorm):
+    """
+    This is a wrapper around the apex fused rms norm implementation. It is meant to be used only with the from_native_module interface.
+    """
+
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "FusedRMSNorm is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native RMSNorm module to FusedRMSNorm module provided by apex."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+        r"""
+        Convert a native RMSNorm module module to FusedRMSNorm module provided by ixformer,
+        and optionally marking parameters for gradient aggregation.
+
+        Args:
+            module (nn.LayerNorm): The native PyTorch LayerNorm module to be converted.
+            sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism.
+
+        Returns:
+            nn.Module: FusedRMSNorm module.
+        """
+
+        LazyInitContext.materialize(module)
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        normalized_shape = getattr(
+            module, "normalized_shape", module.weight.shape[0])
+        eps = module.variance_epsilon if hasattr(
+            module, "variance_epsilon") else module.eps
+        elementwise_affine = getattr(module, "elementwise_affine", True)
+
+        rmsnorm = ixf_FusedRMSNorm(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+        )
+
+        rmsnorm.weight = module.weight
+
+        return rmsnorm
+
+
+class APEXFusedRMSNorm(BaseLayerNorm):
+    """
+    This is a wrapper around the apex fused rms norm implementation. It is meant to be used only with the from_native_module interface.
+    """
+
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "FusedRMSNorm is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native RMSNorm module to FusedRMSNorm module provided by apex."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+        r"""
+        Convert a native RMSNorm module module to FusedRMSNorm module provided by ixformer,
+        and optionally marking parameters for gradient aggregation.
+
+        Args:
+            module (nn.LayerNorm): The native PyTorch LayerNorm module to be converted.
+            sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism.
+
+        Returns:
+            nn.Module: FusedRMSNorm module.
+        """
+
+        LazyInitContext.materialize(module)
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        normalized_shape = getattr(
+            module, "normalized_shape", module.weight.shape[0])
+        eps = module.variance_epsilon if hasattr(
+            module, "variance_epsilon") else module.eps
+        elementwise_affine = getattr(module, "elementwise_affine", True)
+
+        rmsnorm = apex_FusedRMSNorm(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+        )
+
+        rmsnorm.weight = module.weight
+
+        return rmsnorm
+
+
+# 替换torch LayerNorm 的forward
+@staticmethod
+def replace_layernorm_forward(self, input: torch.Tensor) -> torch.Tensor:
+
+    output = torch.empty_like(input)
+
+    return ixff.layernorm_train(input, self.weight, self.bias, self.normalized_shape, output, True)

ixformer_sdk/train/speedformer/layers/qwen2/attention.py

@@ -0,0 +1,263 @@
+import math
+import warnings
+import inspect
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ixformer.train.speedformer.models.qwen2.configuration_qwen2 import Qwen2Config
+from ixformer.train.speedformer.models.qwen2.modeling_qwen2 import Qwen2FlashAttention2
+from transformers import Cache
+from transformers.utils import logging
+
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+
+from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+
+from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb
+from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+_flash_supports_window_size = "window_size" in list(
+    inspect.signature(flash_attn_func).parameters)
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class BaseQwenAttention(Qwen2FlashAttention2):
+    """ 
+    加这个层的原因：1.当原模型中使用的是torch nvtive的attention，强制替换成flash_attn; 2.优化rope 
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        out_dim = self.num_heads * self.head_dim + \
+            self.num_key_value_heads * self.head_dim * 2
+        self.qkv_proj = nn.Linear(self.hidden_size, out_dim, bias=True)
+        del self.q_proj, self.k_proj, self.v_proj
+        self.rotary_emb = RotaryEmbedding(self.head_dim, self.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+        qkv = self.qkv_proj(hidden_states)
+        q_dim = self.num_heads * self.head_dim
+        kv_dim = self.num_key_value_heads * self.head_dim
+        query_states, key_states, value_states = torch.split(
+            qkv, (q_dim, kv_dim, kv_dim), dim=-1)
+        # fused_apply_rotary_pos_emb need qk to be in "sbhd", v stay "bshd"
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous()
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(0, 1).contiguous()
+        value_states = value_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        kv_seq_len = key_states.shape[0]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value[0].shape[0]
+
+        emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32)
+        query_states = fused_apply_rotary_pos_emb(query_states, emb)
+        key_states = fused_apply_rotary_pos_emb(key_states, emb)
+        use_sliding_windows = (
+            _flash_supports_window_size
+            and getattr(self.config, "sliding_window", None) is not None
+            and kv_seq_len > self.config.sliding_window
+            and self.config.use_sliding_window
+        )
+
+        if not _flash_supports_window_size:
+            logger.warning_once(
+                "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
+                " make sure to upgrade flash-attn library."
+            )
+
+        # for now, attention with sliding_windows have not test, so if use_sliding_windows throw error
+        if use_sliding_windows:
+            raise KeyError("use_sliding_windows not support for now")
+
+        # kv cache staff
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=0)
+            value_states = torch.cat([past_key_value[1], value_states], dim=0)
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # if attention mask is None, use flashattn which support GQA
+        if attention_mask is not None:
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in float16 just to be sure everything works as expected.
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa
+        if attention_mask is None:  # flash-attn
+            query_states = query_states.transpose(0, 1).contiguous()
+            key_states = key_states.transpose(0, 1).contiguous()
+        else:  # sdpa
+            query_states = query_states.permute(1, 2, 0, 3).contiguous()
+            key_states = key_states.permute(1, 2, 0, 3).contiguous()
+            value_states = value_states.transpose(1, 2).contiguous()
+
+        attn_output = self._attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=dropout_rate,
+            use_sliding_windows=use_sliding_windows,
+        )
+
+        attn_output = attn_output.reshape(
+            bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _attention_forward(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        query_length,
+        dropout=0.0,
+        softmax_scale=None,
+        use_sliding_windows=False,
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+            use_sliding_windows (`bool`, *optional*):
+                Whether to activate sliding window attention.
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+                is_causal=causal,
+            )
+            attn_output = attn_output.transpose(1, 2).contiguous()
+        else:
+            attn_output = flash_attn_func(
+                query_states,
+                key_states,
+                value_states,
+                dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+        return attn_output
+
+
+class QwenAttention(BaseQwenAttention):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "LlamaAttention is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to LlamaAttention module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        LazyInitContext.materialize(module)
+
+        # try to get normalized_shape, eps, elementwise_affine from the module
+        config = getattr(module, "config")
+        layer_idx = getattr(module, "layer_idx", None)
+
+        attention = BaseQwenAttention(
+            config=config,
+            layer_idx=layer_idx,
+        )
+
+        attention.qkv_proj.weight.data = torch.cat(
+            (module.q_proj.weight.data, module.k_proj.weight.data, module.v_proj.weight.data), dim=0)
+        attention.qkv_proj.bias.data = torch.cat(
+            (module.q_proj.bias.data, module.k_proj.bias.data, module.v_proj.bias.data), dim=0)
+
+        attention.o_proj.weight.data = module.o_proj.weight.data
+
+        return attention

ixformer_sdk/train/speedformer/layers/qwen2/mlp.py

@@ -0,0 +1,55 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+import ixformer.train.functions as F
+from ixformer.train.speedformer.models.qwen2.configuration_qwen2 import Qwen2Config
+from ixformer.train.speedformer.models.qwen2.modeling_qwen2 import Qwen2MLP
+from transformers import Cache
+from transformers.utils import logging
+
+from ixformer.train.speedformer.layers.lazy import LazyInitContext
+
+
+class BaseQwen2MLP(Qwen2MLP):
+    """
+    这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x)))
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gate_up = nn.Linear(
+            self.hidden_size, self.intermediate_size * 2, bias=False)
+        del self.gate_proj, self.up_proj
+        del self.act_fn
+
+    def forward(self, x):
+        res = self.gate_up(x)
+        down_proj = self.down_proj(F.swiglu(res))
+        return down_proj
+
+
+class IXFQwen2MLP(BaseQwen2MLP):
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "IXFQwen2MLP is not implemented as a physical class. "
+            "It is meant to be used only with the from_native_module interface to Convert a native Qwen2MLP module to BaseQwen2MLP module provided above."
+        )
+
+    @staticmethod
+    def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module:
+
+        LazyInitContext.materialize(module)
+
+        config = getattr(module, "config")
+
+        mlp = BaseQwen2MLP(config=config)
+
+        mlp.gate_up.weight.data = torch.concat(
+            (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0)
+        mlp.down_proj.weight.data = module.down_proj.weight.data
+
+        return mlp

ixformer_sdk/train/speedformer/layers/rotary_pos_embedding.py

@@ -0,0 +1,55 @@
+import importlib.util
+import torch
+
+from torch import einsum, nn
+
+__all__ = ['RotaryEmbedding']
+
+
+# RotaryEmbedding and apply_rotary_pos_emb are copy from http://bitbucket.iluvatar.ai:7990/projects/PSR/repos/megatron-deepspeed/browse/megatron/model/rotary_pos_embedding.py
+# for now RotaryEmbedding is used, apply_rotary_pos_emb can be replaced by fused_apply_rotary_pos_emb from ixformer for better performance
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, base=10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        if importlib.util.find_spec('einops') is None:
+            raise RuntimeError("einops is required for Rotary Embedding")
+
+    def forward(self, max_seq_len, offset=0):
+        seq = torch.arange(max_seq_len, device=self.inv_freq.device) + offset
+        freqs = einsum(
+            'i , j -> i j', seq.type_as(self.inv_freq), self.inv_freq)
+        # first part even vector components, second part odd vector components,
+        #  2 * dim in dimension size
+        emb = torch.cat((freqs, freqs), dim=-1)
+        # emb [seq_length, .., dim]
+        from einops import rearrange
+        return rearrange(emb, 'n d -> n 1 1 d')
+
+
+def _rotate_half(x):
+    """
+    change sign so the last dimension becomes [-odd, +even]
+    """
+    from einops import rearrange
+    x = rearrange(x, '... (j d) -> ... j d', j=2)
+    x1, x2 = x.unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(t, freqs):
+    """
+    input tensor t is of shape [seq_length, ..., dim]
+    rotary positional embeding tensor freqs is of shape [seq_length, ..., dim]
+    check https://kexue.fm/archives/8265 for detailed formulas
+    """
+    rot_dim = freqs.shape[-1]
+    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
+    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+
+    # first part is cosine component
+    # second part is sine component, need to change signs with _rotate_half method
+    t = (t * freqs.cos()) + (_rotate_half(t) * freqs.sin())
+    return torch.cat((t, t_pass), dim=-1)