hpcaitech
diff --git a/‎.github/workflows/build_on_pr.yml
Lines changed: 2 additions & 2 deletions b/‎.github/workflows/build_on_pr.yml
Lines changed: 2 additions & 2 deletions
diff --git a/‎colossalai/inference/modeling/models/glide_llama.py
Lines changed: 20 additions & 59 deletions b/‎colossalai/inference/modeling/models/glide_llama.py
Lines changed: 20 additions & 59 deletions
diff --git a/‎colossalai/inference/modeling/models/nopadding_llama.py
Lines changed: 3 additions & 3 deletions b/‎colossalai/inference/modeling/models/nopadding_llama.py
Lines changed: 3 additions & 3 deletions
diff --git a/‎colossalai/inference/spec/drafter.py
Lines changed: 4 additions & 2 deletions b/‎colossalai/inference/spec/drafter.py
Lines changed: 4 additions & 2 deletions
diff --git a/‎colossalai/lazy/pretrained.py
Lines changed: 2 additions & 2 deletions b/‎colossalai/lazy/pretrained.py
Lines changed: 2 additions & 2 deletions
diff --git a/‎colossalai/shardformer/modeling/bert.py
Lines changed: 84 additions & 1 deletion b/‎colossalai/shardformer/modeling/bert.py
Lines changed: 84 additions & 1 deletion
@@ -34,7 +34,7 @@ jobs:
       anyExtensionFileChanged: ${{ steps.find-extension-change.outputs.any_changed }}
       changedLibraryFiles: ${{ steps.find-lib-change.outputs.all_changed_files }}
       anyLibraryFileChanged: ${{ steps.find-lib-change.outputs.any_changed }}
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted,ubuntu-latest]
     concurrency:
       group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-detect-change
       cancel-in-progress: true
@@ -87,7 +87,7 @@ jobs:
     name: Build and Test Colossal-AI
     needs: detect
     if: needs.detect.outputs.anyLibraryFileChanged == 'true'
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted,ubuntu-latest]
     container:
       image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0
       options: --gpus all --shm-size=2g --rm -v /dev/shm -v /data/scratch:/data/scratch
 
@@ -6,19 +6,16 @@
 
 import torch
 import torch.nn as nn
-from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.cache_utils import DynamicCache
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.models.llama.modeling_llama import (
     LlamaAttention,
     LlamaConfig,
     LlamaDecoderLayer,
-    LlamaDynamicNTKScalingRotaryEmbedding,
     LlamaForCausalLM,
-    LlamaLinearScalingRotaryEmbedding,
     LlamaMLP,
     LlamaModel,
     LlamaRMSNorm,
-    LlamaRotaryEmbedding,
 )
 
 from colossalai.inference.spec import GlideInput
@@ -156,31 +153,29 @@ def glide_llama_model_forward(
     if inputs_embeds is None:
         inputs_embeds = self.embed_tokens(input_ids)
 
-    past_seen_tokens = 0
-    if use_cache:  # kept for BC (cache positions)
-        if not isinstance(past_key_values, StaticCache):
-            past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-            past_seen_tokens = past_key_values.get_seq_length()
+    if use_cache and past_key_values is None:
+        past_key_values = DynamicCache()
 
     if cache_position is None:
-        if isinstance(past_key_values, StaticCache):
-            raise ValueError("cache_position is a required argument when using StaticCache.")
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
         cache_position = torch.arange(
             past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
         )
 
     if position_ids is None:
         position_ids = cache_position.unsqueeze(0)
 
-    attention_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
+    attention_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values)
+    if hasattr(glide_input, "n_spec_tokens"):
+        position_ids = position_ids + glide_input.n_spec_tokens
 
     # embed positions
     hidden_states = inputs_embeds
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
     # 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:
@@ -189,9 +184,9 @@ def glide_llama_model_forward(
         # GlideLlamaDecoderLayer
         layer_outputs = decoder_layer(
             hidden_states,
+            position_embeddings=position_embeddings,
             glide_input=glide_input,
             attention_mask=attention_mask,
-            position_ids=position_ids,
             past_key_value=past_key_values,
             output_attentions=output_attentions,
             use_cache=use_cache,
@@ -200,9 +195,6 @@ def glide_llama_model_forward(
 
         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],)
 
@@ -212,16 +204,11 @@ def glide_llama_model_forward(
     if output_hidden_states:
         all_hidden_states += (hidden_states,)
 
-    next_cache = None
-    if use_cache:
-        next_cache = (
-            next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, 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 tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None)
     return BaseModelOutputWithPast(
         last_hidden_state=hidden_states,
-        past_key_values=next_cache,
+        past_key_values=past_key_values,
         hidden_states=all_hidden_states,
         attentions=all_self_attns,
     )
@@ -267,41 +254,17 @@ def __init__(self, config: GlideLlamaConfig):
 
         self.q_proj = nn.Linear(self.hidden_size, self.large_num_heads * self.large_head_dim, bias=False)
         self.o_proj = nn.Linear(self.large_num_heads * self.large_head_dim, self.hidden_size, bias=False)
-        self._init_rope()
-
-    def _init_rope(self):
-        if self.config.rope_scaling is None:
-            self.rotary_emb = LlamaRotaryEmbedding(
-                self.large_head_dim,
-                max_position_embeddings=self.max_position_embeddings,
-            )
-        else:
-            scaling_type = self.config.rope_scaling["type"]
-            scaling_factor = self.config.rope_scaling["factor"]
-            if scaling_type == "linear":
-                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
-                    self.large_head_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                )
-            elif scaling_type == "dynamic":
-                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
-                    self.large_head_dim,
-                    max_position_embeddings=self.max_position_embeddings,
-                    scaling_factor=scaling_factor,
-                )
-            else:
-                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
 
     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
 
     def forward(
         self,
         hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        position_ids: Optional[torch.LongTensor] = None,
         glide_input: GlideInput = None,  # Used for glimpsing main model's KV caches
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
     ) -> Optional[torch.Tensor]:
@@ -319,8 +282,7 @@ def forward(
         query_states = query_states.view(bsz, -1, self.large_num_heads, self.large_head_dim).transpose(1, 2)
 
         # for RoPE
-        position_ids = position_ids + glide_input.n_spec_tokens
-        cos, sin = self.rotary_emb(query_states, position_ids)
+        cos, sin = position_embeddings
         query_states = apply_single_rotary_pos_emb(query_states, cos, sin, position_ids)
         query_states = query_states.transpose(1, 2)
         query_states = query_states.reshape(-1, self.large_num_heads, self.large_head_dim)
@@ -367,9 +329,10 @@ def from_native_module(module: LlamaDecoderLayer, *args, **kwargs) -> "GlideLlam
     def forward(
         self,
         hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor = None,
+        position_ids: Optional[torch.LongTensor] = None,
         glide_input: GlideInput = None,
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
@@ -399,10 +362,10 @@ def forward(
         hidden_states = self.input_layernorm(hidden_states)
 
         # Self Attention
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+        hidden_states, self_attn_weights = self.self_attn(
             hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
             attention_mask=attention_mask,
-            position_ids=position_ids,
             past_key_value=past_key_value,
             output_attentions=output_attentions,
             use_cache=use_cache,
@@ -425,9 +388,10 @@ def forward(
 
             hidden_states = self.cross_attn(
                 hidden_states=hidden_states,
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
                 glide_input=glide_input,
                 attention_mask=attention_mask,
-                position_ids=position_ids,
                 output_attentions=output_attentions,
                 use_cache=True,
             )
@@ -441,9 +405,6 @@ def forward(
 
         outputs = (hidden_states,)
 
-        if use_cache:
-            outputs += (present_key_value,)
-
         return outputs
 
 
 
@@ -478,9 +478,9 @@ def from_native_module(
             attn_oproj=attn_oproj,
             process_group=process_group,
             model_shard_infer_config=model_shard_infer_config,
-            num_heads=module.num_heads,
-            hidden_size=module.hidden_size,
-            num_key_value_heads=module.num_key_value_heads,
+            num_heads=module.config.num_attention_heads,
+            hidden_size=module.config.hidden_size,
+            num_key_value_heads=module.config.num_key_value_heads,
         )
 
         return attn_layer
 
@@ -3,6 +3,7 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedTokenizer
+from transformers.cache_utils import DynamicCache
 
 from colossalai.utils import get_current_device
 
@@ -93,9 +94,8 @@ def speculate(
 
         for _ in range(n_spec_tokens):
             # update past key values
-            kwargs["past_key_values"] = past_key_values
 
-            outputs = self._drafter_model(input_ids, **kwargs)
+            outputs = self._drafter_model(input_ids, past_key_values=past_key_values, **kwargs)
             next_token_logits = outputs.logits[:, -1, :]
 
             # NOTE Only use greedy search for speculating.
@@ -114,6 +114,8 @@ def speculate(
         speculated_length = len(token_ids)  # For now, only support bsz 1
         logits = torch.concat(logits, dim=0)
         token_ids = torch.concat(token_ids, dim=-1)
+        if isinstance(past_key_values, DynamicCache):
+            past_key_values = past_key_values.to_legacy_cache()
 
         out = DrafterOutput(
             speculated_length=speculated_length, logits=logits, next_tokens=token_ids, past_key_values=past_key_values
 
@@ -69,7 +69,7 @@ def new_from_pretrained(
     _ = 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)
+    kwargs.pop("_fast_init", True)
     torch_dtype = kwargs.pop("torch_dtype", None)
     subfolder = kwargs.pop("subfolder", "")
     commit_hash = kwargs.pop("_commit_hash", None)
@@ -286,7 +286,7 @@ def new_from_pretrained(
     config.name_or_path = pretrained_model_name_or_path
 
     # Instantiate model.
-    init_contexts = [no_init_weights(_enable=_fast_init)]
+    init_contexts = [no_init_weights()]
 
     with ContextManagers(init_contexts):
         model = cls(config, *model_args, **model_kwargs)
 
@@ -58,7 +58,7 @@ def bert_model_forward(
         hidden_states: Optional[torch.FloatTensor] = None,  # this is from the previous stage
         stage_index: Optional[List[int]] = None,
         shard_config: ShardConfig = None,
-    ):
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
         # TODO(jianghai): add explaination of the output here.
         r"""
         encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
@@ -1037,6 +1037,89 @@ def forward(self: BertOutput, hidden_states: torch.Tensor, input_tensor: torch.T
     return forward
 
 
+# Fix the tgt_len size in sequence parallel attention:
+# same with the one in BertSdpaSelfAttention forward in v4.51.3 transformers except the
+def get_bert_sequence_parallel_attention_forward(shard_config: ShardConfig):
+    from transformers.models.bert.modeling_bert import BertSdpaSelfAttention
+
+    def forward(
+        self: BertSdpaSelfAttention,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        query_layer = self.transpose_for_scores(self.query(hidden_states))
+
+        # If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention
+        # mask needs to be such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        current_states = encoder_hidden_states if is_cross_attention else hidden_states
+        attention_mask = encoder_attention_mask if is_cross_attention else attention_mask
+
+        # Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning
+        if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]:
+            key_layer, value_layer = past_key_value
+        else:
+            key_layer = self.transpose_for_scores(self.key(current_states))
+            value_layer = self.transpose_for_scores(self.value(current_states))
+            if past_key_value is not None and not is_cross_attention:
+                key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+                value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_layer, value_layer)
+
+        # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom
+        # attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0.
+        # Reference: https://github.com/pytorch/pytorch/issues/112577
+        if self.require_contiguous_qkv and query_layer.device.type == "cuda" and attention_mask is not None:
+            query_layer = query_layer.contiguous()
+            key_layer = key_layer.contiguous()
+            value_layer = value_layer.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+        # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create
+        # a causal mask in case tgt_len == 1.
+        is_causal = (
+            True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False
+        )
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_layer,
+            key_layer,
+            value_layer,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout_prob if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2)
+        _, _, tgt_len, _ = query_layer.shape
+        attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size)
+
+        outputs = (attn_output,)
+        if self.is_decoder:
+            outputs = outputs + (past_key_value,)
+        return outputs
+
+    return forward
+
+
 def bert_sequence_parallel_forward_fn(shard_config: ShardConfig):
     def forward(
         self,