From a96334289220b3489bb131508d9c59d8b834e6ce Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Sat, 10 Oct 2026 15:46:39 +0800 Subject: [PATCH 1/2] fix ple cp --- src/mcore_bridge/model/modules/ple.py | 11 +- src/mcore_bridge/utils/megatron_utils.py | 9 ++ tests/test_ple_cp_grad.py | 170 +++++++++++++++++++++++ 3 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 tests/test_ple_cp_grad.py diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 5cb75a6..1f33278 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from megatron.core import parallel_state from megatron.core.extensions.transformer_engine import TELinear +from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing from megatron.core.tensor_parallel import VocabParallelEmbedding from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region) @@ -614,7 +615,15 @@ def forward( # 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 = reconstruct_tensor_cp(hidden_states, psp_for_cp, dim=0) + # The causal conv mixes tokens across the shard boundary and the output + # is split back below, so a token's gradient partly lands on other CP + # ranks: the gather must reduce-scatter it back (reconstruct_tensor_cp + # keeps only the local slice). + 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]: diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index aaceea8..03c0863 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -131,6 +131,15 @@ def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int, cp_partition_mode Returns: torch.Tensor: Full-sequence tensor with the same shape as ``tensor`` except the size at ``dim`` is multiplied by ``cp_size``. + + Gradient: only the local shard keeps its autograd graph; the slices gathered + from other ranks are constants. That is correct when the tensor needs no + grad, or when every rank computes the same full-sequence loss from the + result (each rank then owns the gradient of its own shard). It drops + gradient when the consumer mixes tokens across shards and splits its output + back to the local shard; use a differentiable gather there + (``gather_from_sequence_parallel_region(..., tensor_parallel_output_grad=True, + group=cp_group)``), as PLE does. """ cp_size = mpu.get_context_parallel_world_size() diff --git a/tests/test_ple_cp_grad.py b/tests/test_ple_cp_grad.py new file mode 100644 index 0000000..1e7b8bd --- /dev/null +++ b/tests/test_ple_cp_grad.py @@ -0,0 +1,170 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""GPU regression test: PLE under context parallelism must match cp=1 in forward and backward. + +PLE gathers the CP shards, runs a causal conv over the full sequence and splits the output back +to the local shard, so the gradient of a token near a shard boundary is partly produced on another +CP rank. The test drives the real ``Qwen4ExpTextPLELayer.forward`` on real NCCL ranks and compares +the local input gradient against a cp=1 run of the same layer. A negative control switches the CP +gather back to a local-slice backward (the ``reconstruct_tensor_cp`` semantics) and must fail. +""" +import itertools +import os +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from megatron.core import parallel_state +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer import TransformerConfig + +from mcore_bridge.model.modules import ple as ple_module +from mcore_bridge.utils.megatron_utils import split_cp_inputs + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA') +requires_nccl = pytest.mark.skipif(not dist.is_nccl_available(), reason='requires NCCL') + + +def _config(cp_size): + config = TransformerConfig( + num_layers=1, + hidden_size=8, + num_attention_heads=1, + context_parallel_size=cp_size, + params_dtype=torch.float32, + perform_initialization=False, # every parameter is overwritten in _layer + ) + # PLE fields normally filled in by the qwen4_exp config parser. + config.hc_count = 2 + config.ple_embed_dim = 16 + config.ple_conv_kernel_size = 4 + config.ngram_size = 3 # also the conv dilation: the receptive field spans 9 previous tokens + config.heads_per_ngram = 2 + config.eos_token_id = 0 + config.split_ngram_parts = 1 + config.ple_seed = 1234 + config.padded_vocab_size = 64 + config.ngram_vocab_size_base = 31 + config.make_ngram_vocab_size_divisible_by = 4 + return config + + +def _layer(cp_size, device): + layer = ple_module.Qwen4ExpTextPLELayer(_config(cp_size), ple_layer_index=0).to(device) + generator = torch.Generator(device='cpu').manual_seed(2026) + with torch.no_grad(): + for name, param in layer.named_parameters(): + # conv1d is zero-initialized; a zero conv would hide the cross-shard gradient. + if 'norm' in name: + value = 0.5 + torch.rand(param.shape, generator=generator) + else: + value = 0.5 * torch.randn(param.shape, generator=generator) + param.copy_(value.to(param.dtype)) + return layer + + +def _packed_params(lengths, device): + boundaries = torch.tensor([0, *itertools.accumulate(lengths)], dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=boundaries, + cu_seqlens_kv=boundaries, + max_seqlen_q=max(lengths), + max_seqlen_kv=max(lengths), + ) + + +def _inputs(layout, cp_size, device): + generator = torch.Generator(device='cpu').manual_seed(7) + width = 16 # hc_count * hidden_size + if layout == 'thd': + packed = _packed_params([4 * cp_size, 12 * cp_size], device) + total, batch = 16 * cp_size, 1 + else: + packed, total, batch = None, 16 * cp_size, 2 + hidden = torch.randn(total, batch, width, generator=generator).to(device) + probe = torch.randn(total, batch, width, generator=generator).to(device) + input_ids = torch.randint(0, 64, (batch, total), generator=generator).to(device) + return hidden, probe, input_ids, packed + + +def _run_cp(layer, hidden, probe, input_ids, packed): + cu = None if packed is None else packed.cu_seqlens_q + local_hidden = split_cp_inputs(hidden, cu, dim=0).detach().clone().requires_grad_(True) + local_ids = split_cp_inputs(input_ids, cu, dim=1) + layer.zero_grad(set_to_none=True) + local_out = layer(local_hidden, local_ids, packed) + (local_out * split_cp_inputs(probe, cu, dim=0)).sum().backward() + param_grads = {name: param.grad.detach().clone() for name, param in layer.named_parameters()} + for grad in param_grads.values(): + dist.all_reduce(grad, group=parallel_state.get_context_parallel_group()) + return local_out.detach(), local_hidden.grad, param_grads + + +def _local_slice_backward_gather(input_, tensor_parallel_output_grad=True, group=None, **kwargs): + return _megatron_gather(input_, tensor_parallel_output_grad=False, group=group, **kwargs) + + +_megatron_gather = ple_module.gather_from_sequence_parallel_region + + +def _cp_grad_worker(rank, world_size, init_file): + # TE runs fp32 GEMMs in TF32 by default (~3e-4 relative on the projection grads); pin true fp32 + # before the first cuBLAS handle so the cp=1 vs cp=N comparison can use tight tolerances. + os.environ['NVIDIA_TF32_OVERRIDE'] = '0' + try: + torch.cuda.set_device(rank) + device = torch.device('cuda', rank) + dist.init_process_group(backend='nccl', init_method=f'file://{init_file}', rank=rank, world_size=world_size) + parallel_state.initialize_model_parallel(context_parallel_size=world_size) + for layout, fused in itertools.product(('sbhd', 'thd'), ('1', '0')): + os.environ['PLE_FUSED_KERNEL'] = fused + case = f'layout={layout} fused={fused} cp={world_size}' + layer = _layer(world_size, device) + hidden, probe, input_ids, packed = _inputs(layout, world_size, device) + cu = None if packed is None else packed.cu_seqlens_q + + # cp=1 reference: the same layer on the full sequence. + full_hidden = hidden.detach().clone().requires_grad_(True) + layer.zero_grad(set_to_none=True) + full_out = layer._forward_impl(full_hidden, input_ids, packed) + (full_out * probe).sum().backward() + ref_grads = {name: param.grad.detach().clone() for name, param in layer.named_parameters()} + ref_out = split_cp_inputs(full_out.detach(), cu, dim=0) + ref_hidden_grad = split_cp_inputs(full_hidden.grad, cu, dim=0) + + out, hidden_grad, param_grads = _run_cp(layer, hidden, probe, input_ids, packed) + torch.testing.assert_close(out, ref_out, atol=1e-5, rtol=1e-5, msg=lambda m: f'{case} output: {m}') + torch.testing.assert_close( + hidden_grad, ref_hidden_grad, atol=1e-5, rtol=1e-5, msg=lambda m: f'{case} hidden grad: {m}') + assert param_grads.keys() == ref_grads.keys() + for name in param_grads: + torch.testing.assert_close( + param_grads[name], ref_grads[name], atol=1e-4, rtol=1e-4, msg=lambda m: f'{case} {name}: {m}') + + # Negative control: a local-slice gather backward must lose the cross-shard gradient. + ple_module.gather_from_sequence_parallel_region = _local_slice_backward_gather + try: + _, bad_hidden_grad, _ = _run_cp(layer, hidden, probe, input_ids, packed) + finally: + ple_module.gather_from_sequence_parallel_region = _megatron_gather + error = (bad_hidden_grad - ref_hidden_grad).abs().max() + dist.all_reduce(error, op=dist.ReduceOp.MAX) + assert error > 1e-2, f'{case}: negative control did not fail (max error {error.item():.3e})' + if rank == 0: + print(f'{case}: ok, negative-control max hidden-grad error {error.item():.3e}') + finally: + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@requires_cuda +@requires_nccl +@pytest.mark.parametrize('world_size', [2, 4]) +def test_ple_cp_backward_matches_cp1(tmp_path, world_size): + """Real NCCL CP ranks: PLE output, input gradient and summed parameter gradients match cp=1.""" + if torch.cuda.device_count() < world_size: + pytest.skip(f'requires {world_size} CUDA devices') + init_file = str(tmp_path / f'ple-cp{world_size}-init') + mp.spawn(_cp_grad_worker, args=(world_size, init_file), nprocs=world_size, join=True) From 0c583d3b92d76eb9b490ea6863da50be023856af Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Sat, 10 Oct 2026 22:57:54 +0800 Subject: [PATCH 2/2] lint --- src/mcore_bridge/model/modules/ple.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 1f33278..0c5a60c 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -615,15 +615,10 @@ def forward( # 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 - # The causal conv mixes tokens across the shard boundary and the output - # is split back below, so a token's gradient partly lands on other CP - # ranks: the gather must reduce-scatter it back (reconstruct_tensor_cp - # keeps only the local slice). 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) + 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]: