From 1b882cd236a0fa9fe665b9155a08077a485f91f2 Mon Sep 17 00:00:00 2001 From: szm Date: Wed, 20 May 2026 00:29:39 +0800 Subject: [PATCH] Backport fix for Qwen3.5 context parallel from #3529 When integrating mcore with ms-swift/llamafactory for Qwen3.5 context parallel (CP), the core_dev_r0.16.0 branch encounters the same bug resolved in #3529. This commit backports the fix from main to core_dev_r0.16.0 to unblock Qwen3.5 CP training. --- megatron/core/transformer/attention.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index bc5e4e2ee0d..739450c0738 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -1491,6 +1491,14 @@ def get_query_key_value_tensors( if output_gate: # Gate [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] gate = gate.reshape(*gate.shape[:2], -1, self.hidden_size_per_attention_head) + if self.config.num_query_groups < self.world_size: + idx = get_tensor_model_parallel_rank() % ( + self.world_size // self.config.num_query_groups + ) + size = self.num_attention_heads_per_partition // ( + self.world_size // self.config.num_query_groups + ) + gate = gate[:, :, idx * size : (idx + 1) * size, :] return query, key, value, gate return query, key, value