fix

hpcaitech · Aug 21, 2024 · 698c8b9 · 698c8b9
1 parent 6aface9
commit 698c8b9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py
@@ -538,6 +538,7 @@ def forward(
 
  # sp: all-to-all comminucation when introducing sequence parallel
  if sp_mode == "all_to_all":
+ query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication)
  key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication)
  value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication)
  bsz, q_len, _ = query_states.size()