[docker] backport output-gate slicing when query groups < TP

[LLMShark]

Summary

• Backport [main] fix(moe): fix the bug where gate was not sliced when kv_head < tp_size. NVIDIA/Megatron-LM#3575 to the Megatron revision pinned by the slime image.
• Run the existing Qwen3.5-0.8B E2E test with TP=4 so it covers output-gated attention when num_query_groups < tensor_model_parallel_size.

Problem

slime currently pins Megatron to 1dcf0dafa884ad52ffb243625717a3471643e087. In that revision, the GQA projection is gathered when the number of query groups is smaller than the TP world size, and query is narrowed back to the heads owned by the local TP rank. The matching gate tensor is reshaped but not narrowed.

The reported Qwen3.5-397B-A17B failure used TP=8 with num_attention_heads=32, num_query_groups=2, attention_output_gate=true, and head dimension 256:

query heads per TP rank = 32 / 8 = 4
gate heads after the GQA group gather/select = 32 / 2 = 16
partitions = 16 / 4 = 4
partition = tp_rank % 4
gate slice = gate[:, :, partition * 4 : (partition + 1) * 4, :]

Before the fix, the reported gate shape was [6144, 1, 16, 256]; the TP-local query/output owns 4 heads, so the matching gate must be narrowed to [6144, 1, 4, 256] before _apply_output_gate. The unsliced gate causes the invalid view failure.

PR #1662 previously reduced a Qwen3.5 configuration from TP=4 to TP=2 because the model has two query groups. This backport fixes the underlying gated-attention partition mismatch instead of requiring TP to stay at or below the query-group count.

Fix

The added hunk is the production change from NVIDIA/Megatron-LM@16a8cdb: use the same TP-local slice for gate that Megatron already applies to query.

The existing Qwen3.5-0.8B test has 8 attention heads, 2 query groups, and output gating enabled. Moving it from TP=1 to TP=4 exercises the failing num_query_groups < TP path without increasing its existing four-GPU allocation.

Validation

• Applied the complete updated docker/patch/latest/megatron.patch with git apply --3way to a clean checkout of the pinned Megatron commit; every hunk applied cleanly and no conflict markers were produced.
• git diff --check
• Python syntax parsing for tests/test_qwen3.5_0.8B_gsm8k_short.py

The full GPU training test was not run locally. Because this changes the image's Megatron patch, run-ci-image should be used to validate the TP=4 Qwen3.5 E2E path on a rebuilt test image.