Fix train_mem for the upstream changes
diff --git a/fastchat/train/llama_flash_attn_monkey_patch.py b/fastchat/train/llama_flash_attn_monkey_patch.py
index d16a001ed2..f87b76e19e 100644
--- a/fastchat/train/llama_flash_attn_monkey_patch.py
+++ b/fastchat/train/llama_flash_attn_monkey_patch.py
@@ -14,8 +14,9 @@
 def forward(
     self,
     hidden_states: torch.Tensor,
-    past_key_value: Optional[Tuple[torch.Tensor]] = None,
     attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
@@ -26,30 +27,20 @@ def forward(
     """
     bsz, q_len, _ = hidden_states.size()
 
-    query_states = self.q_proj(hidden_states).view(
-        bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-    key_states = self.k_proj(hidden_states).view(
-        bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-    value_states = self.v_proj(hidden_states).view(
-        bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
     # [bsz, q_len, nh, hd]
     # [bsz, nh, q_len, hd]
 
     kv_seq_len = key_states.shape[-2]
-    offset = 0
-    if past_key_value is not None:
-        offset = past_key_value[0].shape[-2]
-        kv_seq_len += offset
+    assert past_key_value is None, "past_key_value is not supported"
+
     cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-    query_states, key_states = apply_rotary_pos_emb(query_states,
-                                                    key_states,
-                                                    cos,
-                                                    sin,
-                                                    offset=offset)
+    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
     # [bsz, nh, t, hd]
     assert not output_attentions, "output_attentions is not supported"
     assert not use_cache, "use_cache is not supported"
-    assert past_key_value is None, "past_key_value is not supported"
 
     # Flash attention codes from
     # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py

This is to fix the training code due to the upstream change for the `LlamaAttention` module.

Fixes #187 

Tested:
- [x] `sky launch -c vicuna-7b ./scripts/train-vicuna.yaml --cloud gcp --env WANDB_API_KEY --env MODEL_SIZE=7 --cloud gcp -s`