[bugfix] LoRA + TP>1 + sequence_parallel: sum grads of the replicated LoRA factor over the TP group (exported adapter had last-layer linear_proj.lora_B == 0)

[hojuna]

Summary

With tensor_model_parallel_size > 1 and sequence_parallel=True, the replicated LoRA factor created by LoraParallelLinear.update_layer — lora_B for RowParallel targets (linear_proj, linear_fc2) and lora_A for ColumnParallel targets (linear_qkv, linear_fc1) — only receives the gradient of its own TP rank's sequence shard, and that gradient is never summed over the TP group. Each TP rank therefore trains a different copy of the factor, and export_weights saves rank 0's copy only, so the exported adapter is not the model that was trained.

• For RowParallel targets lora_A (TERowParallelLinear) reduce-scatters its output, so the local lora_B (built with _build_local_te_linear) sees a different sequence shard on every TP rank.
• For ColumnParallel targets the local lora_A consumes the sequence-sharded input.
• Megatron sums such gradients in finalize_model_grads only for params flagged param.sequence_parallel = True (as it does for layernorm weights). LoraParallelLinear never sets the flag (tuners/patcher.py only marks the non-parallel LoraLinear).

Not affected: tensor_model_parallel_size=1, or sequence_parallel=False (the replicated factor then sees identical inputs on all TP ranks).

Observed

ms-swift 4.4.2 megatron sft, mcore-bridge 1.5.2 (code path unchanged on current main), Qwen3.5-35B-A3B-FP8, TP4 / EP2, --sequence_parallel true, LoRA r=8 on linear_qkv linear_proj, --save_safetensors true.

• The saved adapter has layers.39.self_attn.o_proj.lora_B.weight == 0 (exactly zero) at every checkpoint (100 … 541 steps) while its lora_A keeps training normally on all 4 TP shards.
• Printing the parameter on every rank right before save_checkpoint after 3 iterations (TP ranks 0-3 of DP rank 0):

tensor	rank0	rank1	rank2	rank3
layers.39.linear_proj.lora_B	0.000000	0.000000	0.000000	0.007745
layers.35.linear_proj.lora_B	0.008371	0.007328	—	0.007253

For the last layer, rank 0-2's shards (system prompt + image tokens, loss-masked) get exactly zero gradient, hence the all-zero tensor; deeper layers are silently different per rank rather than zero. With TP=1 / EP=8 the same run trains layers.39.o_proj.lora_B to 0.0075.

Fix

Set p.sequence_parallel = True on the parameters of the replicated factor (lora_b if is_parallel_a else lora_a) when tp_size > 1 and sequence parallelism is enabled (router LoRA excluded). Megatron's finalize_model_grads then all-reduces (SUM) the gradient over the TP group.

Verification

Same 3-iteration run with the fix: all 8 ranks hold identical values (layers.35.linear_proj.lora_B = 0.007253, layers.39.linear_proj.lora_B = 0.007726 on every rank), the exported adapter has no all-zero tensors, and layers.39.o_proj.lora_B matches the TP=1 result (0.0077 vs 0.0075).