[Bug] Qwen3.8 PLE gradients are incorrectly scaled by TP size when sequence parallelism is enabled

Checklist / 检查清单

I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

With Qwen3.8/Qwen4Exp PLE, sequence_parallel=True and TP > 1, the replicated PLE parameters already receive the full-sequence gradient during backward. Marking these parameters sequence_parallel=True causes Megatron's gradient finalizer to SUM the already-complete gradients again across TP. In a TP=8 regression, the resulting parameter gradients are 8 times the full-sequence reference.

This affects training updates and global gradient clipping, not just gradient-norm reporting. Model-level SP should remain enabled; the issue is the gradient-reduction metadata on the replicated PLE parameters.

This is a silent gradient correctness issue; no exception or crash is required to reproduce it.

Relevant code and backward path
The constructor marks key_proj, value_proj, norm_key, norm_query, norm_conv, and conv1d parameters for SP gradient reduction:

mcore-bridge/src/mcore_bridge/model/modules/ple.py

Lines 500 to 504 in 9c43ac4

    
           replicated_prefixes = ('key_proj', 'value_proj', 'norm_key', 'norm_query', 'norm_conv', 'conv1d') 
        
           for name, param in self.named_parameters(): 
        
               if name.startswith(replicated_prefixes) or name.endswith('norm.weight'): 
        
                   # Replicated across TP; reduce grads across TP when SP is on. 
        
                   setattr(param, 'sequence_parallel', config.sequence_parallel)

setattr(param, 'sequence_parallel', config.sequence_parallel)
However, PLE's SP forward path gathers the full sequence, computes PLE, and scatters its output back to SP ranks:

mcore-bridge/src/mcore_bridge/model/modules/ple.py

Lines 610 to 647 in 9c43ac4

    
           if sp_on: 
        
               hidden_states = gather_from_sequence_parallel_region( 
        
                   hidden_states, tensor_parallel_output_grad=False, group=self.pg_collection.tp) 
        
           if cp_on: 
        
               # For packed (thd) inputs the data pipeline (swift get_batch_on_this_cp_rank) 
        
               # zigzag-splits hidden/input_ids per sample via cu_seqlens_q while keeping 
        
               # packed_seq_params itself global, so the undo uses the same cu. 
        
               psp_for_cp = packed_seq_params if thd else None 
        
               hidden_states = gather_from_sequence_parallel_region( 
        
                   hidden_states, tensor_parallel_output_grad=True, group=parallel_state.get_context_parallel_group()) 
        
               hidden_states = _undo_attention_load_balancing(hidden_states, 
        
                                                              parallel_state.get_context_parallel_world_size(), psp_for_cp) 
        
               # The data pipeline may hand us either a CP-sharded or a full copy of 
        
               # input_ids; re-align only when the lengths disagree. 
        
               if input_ids.shape[-1] != hidden_states.shape[0]: 
        
                   input_ids = reconstruct_tensor_cp(input_ids, psp_for_cp, dim=1) 
        
           elif input_ids.shape[-1] != hidden_states.shape[0]: 
        
               raise ValueError(f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered ' 
        
                                f'sequence length {hidden_states.shape[0]} under sequence parallelism; ' 
        
                                'gpt_model is expected to pass the full, unsharded input_ids.') 
        
           if thd: 
        
               # SP keeps one global cu_seqlens copy per rank; normalize padded/offset 
        
               # forms against the gathered total so cu indexes the full sequence. 
        
               # copy.copy (not dataclasses.replace) preserves dynamically attached 
        
               # fields such as `num_samples` that the data pipeline relies on. 
        
               cu = self._normalize_cu_seqlens(getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0]) 
        
               psp = copy.copy(packed_seq_params) 
        
               psp.cu_seqlens_q = cu 
        
               packed_seq_params = psp 
        
           out = self._forward_impl(hidden_states, input_ids, packed_seq_params) 
        
           # ---- scatter the additive output back to the caller's shard ---- 
        
           if cp_on: 
        
               out = split_cp_inputs(out, getattr(packed_seq_params, 'cu_seqlens_q', None) if thd else None, dim=0) 
        
           if sp_on: 
        
               out = scatter_to_sequence_parallel_region(out, group=self.pg_collection.tp)

The output scatter's backward gathers the output gradients, so each TP rank's replicated PLE parameters already receive the complete gradient. Megatron's _allreduce_non_tensor_model_parallel_grads subsequently SUMs those gradients because of the sequence_parallel parameter attribute. For TP=8 this changes g into 8*g.

How to Reproduce / 如何复现

Versions and runtime environment
Reproduced on 16 NVIDIA GPUs using TP=8, EP=16, ETP=1, CP=1, with the PLE n-gram embedding frozen.
Tested mcore-bridge base: cef925c.
Tested Megatron-LM: core_v0.19.0, commit 5be9626709af2722333bf54797c954c09edeada3.
On 2026-10-15, source inspection confirmed that both main (9c43ac4) and v1.6.5 still retain the relevant parameter flags and SP gather/scatter path. We have not rerun the GPU regression on those newer revisions.
Reproduction and observed results
The test used the real Qwen4ExpTextPLELayer.forward and Megatron gradient finalizer, rather than only a synthetic all-reduce:

Construct a small PLE layer with TP=8 and SP enabled. Freeze the n-gram table, synchronize replicated parameters across TP, and initialize the convolution weight to nonzero values so its default zero initialization does not hide discrepancies.
Compute a full-sequence reference with layer._forward_impl(full_hidden, input_ids) and backpropagate a fixed upstream gradient.
Split the same hidden states and upstream gradient along the sequence dimension. Run layer(local_hidden, input_ids) and backward on each TP rank.
Verify local outputs and local input gradients against the corresponding reference slices. Verify replicated parameter gradients against the full-sequence reference before finalization.
Copy each parameter's gradient into main_grad and invoke the real _allreduce_non_tensor_model_parallel_grads.
With the original flags, check elementwise that main_grad == 8 * reference_grad. With the proposed flags below, check that main_grad == reference_grad.
All 16 ranks passed the original-behavior control and the corrected-behavior checks. Recorded corrected gradient-norm ratios from rank 0:

Parameter Corrected gradient norm / reference gradient norm
key_proj.weight 1.00000012
value_proj.weight 1.00000000
norm_key.weight 1.00000012
norm_query.weight 1.00000012
norm_conv.weight 1.00000000
conv1d.weight 1.00000000
The correction left input gradients unchanged. The regression uses FP32 to isolate the reduction error from BF16 rounding.

Additional Information / 补充信息

Proposed fix
For the same replicated-parameter selection in the PLE constructor:

setattr(param, 'sequence_parallel', config.sequence_parallel)

setattr(param, 'sequence_parallel', False)
setattr(param, 'average_gradients_across_tp_domain', config.sequence_parallel)
This preserves model-level SP and the existing gather/scatter computation. It replaces the redundant SUM of complete replicated gradients with AVG synchronization across TP.

This was validated with the Megatron version above; compatibility with other supported Megatron versions and combined CP+SP should be reviewed separately.

thanks for your report, this issue has been fixed in #234

Navigation Menu

Repository navigation

[Bug] Qwen3.8 PLE gradients are incorrectly scaled by TP size when sequence parallelism is enabled #233

Checklist / 检查清单

Bug Description / Bug 描述

How to Reproduce / 如何复现

Additional Information / 补充信息

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	replicated_prefixes = ('key_proj', 'value_proj', 'norm_key', 'norm_query', 'norm_conv', 'conv1d')
	for name, param in self.named_parameters():
	if name.startswith(replicated_prefixes) or name.endswith('norm.weight'):
	# Replicated across TP; reduce grads across TP when SP is on.
	setattr(param, 'sequence_parallel', config.sequence_parallel)

	if sp_on:
	hidden_states = gather_from_sequence_parallel_region(
	hidden_states, tensor_parallel_output_grad=False, group=self.pg_collection.tp)
	if cp_on:
	# For packed (thd) inputs the data pipeline (swift get_batch_on_this_cp_rank)
	# zigzag-splits hidden/input_ids per sample via cu_seqlens_q while keeping
	# packed_seq_params itself global, so the undo uses the same cu.
	psp_for_cp = packed_seq_params if thd else None
	hidden_states = gather_from_sequence_parallel_region(
	hidden_states, tensor_parallel_output_grad=True, group=parallel_state.get_context_parallel_group())
	hidden_states = _undo_attention_load_balancing(hidden_states,
	parallel_state.get_context_parallel_world_size(), psp_for_cp)
	# The data pipeline may hand us either a CP-sharded or a full copy of
	# input_ids; re-align only when the lengths disagree.
	if input_ids.shape[-1] != hidden_states.shape[0]:
	input_ids = reconstruct_tensor_cp(input_ids, psp_for_cp, dim=1)
	elif input_ids.shape[-1] != hidden_states.shape[0]:
	raise ValueError(f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered '
	f'sequence length {hidden_states.shape[0]} under sequence parallelism; '
	'gpt_model is expected to pass the full, unsharded input_ids.')

	if thd:
	# SP keeps one global cu_seqlens copy per rank; normalize padded/offset
	# forms against the gathered total so cu indexes the full sequence.
	# copy.copy (not dataclasses.replace) preserves dynamically attached
	# fields such as `num_samples` that the data pipeline relies on.
	cu = self._normalize_cu_seqlens(getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0])
	psp = copy.copy(packed_seq_params)
	psp.cu_seqlens_q = cu
	packed_seq_params = psp

	out = self._forward_impl(hidden_states, input_ids, packed_seq_params)

	# ---- scatter the additive output back to the caller's shard ----
	if cp_on:
	out = split_cp_inputs(out, getattr(packed_seq_params, 'cu_seqlens_q', None) if thd else None, dim=0)
	if sp_on:
	out = scatter_to_sequence_parallel_region(out, group=self.pg_collection.tp)

[Bug] Qwen3.8 PLE gradients are incorrectly scaled by TP size when sequence parallelism is enabled #233

Description

Checklist / 检查清单

Bug Description / Bug 描述

How to Reproduce / 如何复现

Additional Information / 补充信息

Activity

hjh0119 commented on Oct 15, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions