fix

2025-09-02 17:46:42 +00:00 · 2024-09-25 18:56:18 +08:00
parent 65c8297710
commit 6fb1322db1
3 changed files with 13 additions and 9 deletions
--- a/colossalai/shardformer/modeling/llama.py
+++ b/colossalai/shardformer/modeling/llama.py
@@ -564,6 +564,7 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        tp_group = shard_config.tensor_parallel_process_group
+
        if sp_mode == "ring_attn":
            attn_output = RingAttention.attention(
                query_states,