[Bugs] Fix dispatch attn bug (#829)

* fix collate bug

* fix dispatch attn bugs

* fix rotary_seq_len bug

xtuner/dataset/collate_fns/default_collate_fn.py

@@ -56,8 +56,8 @@ def default_collate_fn(instances: Sequence[Dict],
  # Some tokenizers have the same eos token and pad token, so input_ids
  # cannot be masked directly based on the pad token id.
  attention_mask = torch.zeros_like(input_ids).bool()
- for i in ori_length:
- attention_mask[:i] = True
+ for i, length in enumerate(ori_length):
+ attention_mask[i, :length] = True
 
  bs, seq_len = input_ids.shape
  position_ids = torch.arange(seq_len).unsqueeze(0).long().repeat(bs, 1)

xtuner/model/modules/dispatch/cohere.py

@@ -105,17 +105,22 @@ def cohere_attn_forward(
  query_states, key_states, value_states = \
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states)
+ # self.num_heads is used in self._upad_input method
+ # num_heads has been changed because of sequence parallel
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  attn_output = self._flash_attention_forward(
  query_states,
  key_states,
  value_states,
  attention_mask,
- q_len,
+ query_states.shape[1],
  dropout=dropout_rate)
 
  if enable_sequence_parallel:
  attn_output = post_process_for_sequence_parallel_attn(attn_output)
+ self.num_heads = ori_num_head
 
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
  attn_output = self.o_proj(attn_output)

xtuner/model/modules/dispatch/deepseek_v2.py

@@ -128,6 +128,10 @@ def deepseek_attn_forward(
  query_states, key_states, value_states = \
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states)
+ # self.num_heads is used in self._upad_input method
+ # num_heads has been changed because of sequence parallel
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  attn_output = self._flash_attention_forward(
  query_states,
@@ -141,6 +145,7 @@ def deepseek_attn_forward(
 
  if enable_sequence_parallel:
  attn_output = post_process_for_sequence_parallel_attn(attn_output)
+ self.num_heads = ori_num_head
 
  if self.q_head_dim != self.v_head_dim:
  attn_output = attn_output[:, :, :, :self.v_head_dim]

xtuner/model/modules/dispatch/internlm2.py

@@ -149,6 +149,10 @@ def internlm2_attn_forward(
  query_states, key_states, value_states = \
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states)
+ # self.num_heads is used in self._upad_input method
+ # num_heads has been changed because of sequence parallel
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  dropout_rate = 0.0
  attn_output = self._flash_attention_forward(
@@ -161,6 +165,7 @@ def internlm2_attn_forward(
 
  if enable_sequence_parallel:
  attn_output = post_process_for_sequence_parallel_attn(attn_output)
+ self.num_heads = ori_num_head
 
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
  attn_output = self.wo(attn_output)

xtuner/model/modules/dispatch/mistral.py

@@ -214,6 +214,11 @@ def mistral_attn_forward(
  query_states, key_states, value_states = \
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states)
+ # num_heads has been changed because of sequence parallel
+ # `self.num_heads`` is not used in self._flash_attention_forward
+ # in mistral/mixtral, we are doing this to avoid some unnecessary risk
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  attn_output = self._flash_attention_forward(
  query_states,
@@ -227,6 +232,7 @@ def mistral_attn_forward(
 
  if enable_sequence_parallel:
  attn_output = post_process_for_sequence_parallel_attn(attn_output)
+ self.num_heads = ori_num_head
 
  attn_output = attn_output.reshape(bsz, q_len,
  self.hidden_size).contiguous()
@@ -311,7 +317,7 @@ def mistral_varlen_attn_forward(
  value_states = value_states.transpose(1, 2)
  # Because the input can be padded, the absolute sequence length
  # depends on the max position id.
- rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
+ rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
  cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
  query_states, key_states = apply_rotary_pos_emb(
  query_states, key_states, cos, sin, position_ids)

xtuner/model/modules/dispatch/phi3.py

@@ -233,6 +233,11 @@ def phi3_attn_forward(
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states,
  scatter_dim=2, gather_dim=1)
+ # num_heads has been changed because of sequence parallel
+ # `self.num_heads`` is not used in self._flash_attention_forward
+ # in mistral/mixtral, we are doing this to avoid some unnecessary risk
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  attn_output = self._flash_attention_forward(
  query_states,
@@ -248,6 +253,7 @@ def phi3_attn_forward(
  # (b, s, nd // sp_world_size, dim) -> (b, s // sp_world_size, nd, dim)
  attn_output = post_process_for_sequence_parallel_attn(
  attn_output, scatter_dim=1, gather_dim=2)
+ self.num_heads = ori_num_head
 
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
  attn_output = self.o_proj(attn_output)
@@ -333,7 +339,7 @@ def phi3_varlen_attn_forward(
  self.layer_idx)
 
  assert position_ids is not None
- rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
+ rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
  cos, sin = self.rotary_emb(
  value_states, position_ids, seq_len=rotary_seq_len)
 

xtuner/model/modules/dispatch/qwen2.py

@@ -151,6 +151,11 @@ def qwen2_attn_forward(
  query_states, key_states, value_states = \
  pre_process_for_sequence_parallel_attn(
  query_states, key_states, value_states)
+ # num_heads has been changed because of sequence parallel
+ # `self.num_heads`` is not used in self._flash_attention_forward
+ # in mistral/mixtral, we are doing this to avoid some unnecessary risk
+ ori_num_head = self.num_heads
+ self.num_heads = query_states.shape[-2]
 
  attn_output = self._flash_attention_forward(
  query_states,
@@ -164,6 +169,7 @@ def qwen2_attn_forward(
 
  if enable_sequence_parallel:
  attn_output = post_process_for_sequence_parallel_attn(attn_output)
+ self.num_heads = ori_num_head
 
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
  attn_output = self.o_proj(attn_output)
@@ -227,7 +233,7 @@ def qwen2_varlen_attn_forward(
  self.layer_idx)
 
  assert position_ids is not None
- rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
+ rotary_seq_len = max(kv_seq_len, position_ids.max().item() + 1)
  cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
 
  query_states, key_states = apply_rotary_pos_emb(query_states, key_states,