Update new model

modeling_phi.py

@@ -24,11 +24,14 @@ try:
     from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
     from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
     from flash_attn.ops.fused_dense import FusedDense
-except:
+    print("Using Flash Attention!")
+except Exception as exc:
+    print(exc)
     pad_input, unpad_input = None, None
     FlashRotaryEmbedding = None
     FlashSelfAttention, FlashCrossAttention = None, None
     FusedDense = None
+    print("Not using Flash Attention!")
 
 
 @dataclass
@@ -525,7 +528,7 @@ class MHA(nn.Module):
         softmax_scale: Optional[float] = None,
         layer_idx: Optional[int] = None,
         return_residual: bool = False,
-        checkpointing: bool = False,
+        checkpointing: bool = True,
     ) -> None:
         super().__init__()
 
@@ -607,7 +610,7 @@ class MHA(nn.Module):
 
             if self.checkpointing:
                 attn_output = torch.utils.checkpoint.checkpoint(
-                    self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
+                    self.inner_attn, qkv, None, cu_seqlens, max_seqlen, use_reentrant=False
                 )
             else:
                 attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
@@ -616,7 +619,7 @@ class MHA(nn.Module):
             return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
 
         if self.checkpointing:
-            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
+            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, None, key_padding_mask, use_reentrant=False)
 
         return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
 
@@ -669,11 +672,12 @@ class MHA(nn.Module):
                     self.inner_cross_attn,
                     q,
                     kv,
-                    causal=causal,
-                    cu_seqlens=cu_seqlens_q,
-                    max_seqlen=max_seqlen_q,
-                    cu_seqlens_k=cu_seqlens_k,
-                    max_seqlen_k=max_seqlen_k,
+                    causal,
+                    cu_seqlens_q,
+                    max_seqlen_q,
+                    cu_seqlens_k,
+                    max_seqlen_k,
+                    use_reentrant=False,
                 )
             else:
                 attn_output = self.inner_cross_attn(
@@ -697,8 +701,9 @@ class MHA(nn.Module):
                 self.inner_cross_attn,
                 q,
                 kv,
-                key_padding_mask=key_padding_mask,
-                causal=causal,
+                causal,
+                key_padding_mask, 
+                use_reentrant=False,
             )
 
         return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
@@ -835,7 +840,7 @@ class PhiPreTrainedModel(PreTrainedModel):
 
     config_class = PhiConfig
     base_model_prefix = "transformer"
-    supports_gradient_checkpointing = False
+    supports_gradient_checkpointing = True
     _no_split_modules = ["ParallelBlock"]
 
     def __init__(self, *inputs, **kwargs) -> None:
@@ -862,20 +867,20 @@ class PhiPreTrainedModel(PreTrainedModel):
         attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
         **kwargs,
     ) -> Dict[str, Any]:
-        if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
-            past_key_values = InferenceParams(
-                max_seqlen=self.config.n_positions,
-                max_batch_size=input_ids.shape[0],
-                seqlen_offset=0,
-                batch_size_offset=0,
-                key_value_memory_dict={},
-                lengths_per_sample=None,
-            )
-        else:
-            # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
-            past_key_values.seqlen_offset = input_ids.shape[1] - 1
-            input_ids = input_ids[:, -1].unsqueeze(-1)
-            attention_mask = attention_mask[:, -1].unsqueeze(-1)
+        # if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
+        #     past_key_values = InferenceParams(
+        #         max_seqlen=self.config.n_positions,
+        #         max_batch_size=input_ids.shape[0],
+        #         seqlen_offset=0,
+        #         batch_size_offset=0,
+        #         key_value_memory_dict={},
+        #         lengths_per_sample=None,
+        #     )
+        # else:
+        #     # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
+        #     past_key_values.seqlen_offset = input_ids.shape[1] - 1
+        #     input_ids = input_ids[:, -1].unsqueeze(-1)
+        #     attention_mask = attention_mask[:, -1].unsqueeze(-1)
 
         return {
             "input_ids": input_ids,
@@ -891,17 +896,19 @@ class PhiModel(PhiPreTrainedModel):
     _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
 
     def __init__(self, config: PhiConfig) -> None:
+        config.flash_attn = True
+        config.flash_rotary = True
         super().__init__(config)
 
         self.embd = Embedding(config)
         self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
-        self.gradient_checkpointing = False
+        self.gradient_checkpointing = True
         self.post_init()
 
-    def get_input_embeddings(self):
-        return self.embd
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embd.wte
 
-    def set_input_embeddings(self, new_embeddings) -> None:
+    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
         self.embd.wte = new_embeddings
 
     def forward(
@@ -919,11 +926,20 @@ class PhiModel(PhiPreTrainedModel):
             raise ValueError("You have to specify either input_ids or inputs_embeds")
 
         for layer in self.h:
-            hidden_states = layer(
-                hidden_states,
-                past_key_values=past_key_values,
-                attention_mask=attention_mask,
-            )
+            if self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    layer.__call__,
+                    hidden_states,
+                    past_key_values,
+                    attention_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states = layer(
+                    hidden_states,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                )
 
         return hidden_states
 
@@ -947,10 +963,10 @@ class PhiForCausalLM(PhiPreTrainedModel):
 
         self.post_init()
 
-    def get_output_embeddings(self):
-        return self.lm_head
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.linear
 
-    def set_output_embeddings(self, new_embeddings) -> None:
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
         self.lm_head.linear = new_embeddings
 
     def forward(