From 79d355946267f9f46f8022e184bf969dc760ede9 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Wed, 5 Aug 2026 03:41:20 +1000 Subject: [PATCH 1/7] feat(megatron): backport FLA chunkwise CP for GDN (Task 32 v1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compatibility layer only: FLA 0.4.1 -> 0.4.2 and a selective backport of Megatron-LM `5139086e` (NVIDIA/Megatron-LM#3282) onto the pinned MCore `85bced0a`. No Relax routing changes -- `relax/` is untouched, `auto` never selects chunkwise, and every existing recipe runs the same path as before. RFC: redai-infra/Relax#213, reworked per the 2026-08-05 review decisions: * the GDN CP mode is **static** for the whole process. There is no per-call override; `linear_cp_mode` is read by both the construction-time head check and by `GatedDeltaNet.forward`, and nothing in a forward writes to `self` or the shared config. Dynamic CP varies only `cp_group` / `local_cp_size`. * **v1 depends on #3282 only.** Nothing from the still-open #5664 is included. # ⭐ Feature ## MCore backport (docker/patch/megatron/20260805-85bced0ae.patch) - New `megatron/core/context_parallel_layout.py`, **byte-identical to `5139086e` below the module docstring**: zigzag <-> contiguous THD/SBHD partitions and a single-all-to-all swap between them. The THD swap rebuilds its routing from `cu_seqlens` per call, which is upstream behaviour. - `packed_seq_params.py`: `resolve_cp_group()` only. The dataclass field list is untouched. - `transformer_config.py`: `linear_cp_mode` with headwise `% (tp*cp)` vs chunkwise `% tp` head divisibility. Default is `headwise`, NOT upstream's `chunkwise`, so upgrading the image cannot silently reroute a recipe. `all_gather` is accepted as a third declared value using the TP-only rule, so the declared config equals the resolved `--gdn-cp-mode` rather than declaring one mode while running another. Unknown values, including an unresolved `auto`, assert at construction. - `gated_delta_net.py`: `_resolve_cp_routing()` gives the whole CP group to exactly one of headwise / chunkwise and `None` to the other; validates `local_cp_size == cp_group.size()`; never creates a process group; short circuits on `cp_size == 1` before the mode is read so a CP=1 micro-batch is legal under any declared mode; raises if `all_gather` reaches MCore's forward with cp>1 (the Relax wrapper was not installed). Plus zigzag<->contiguous conversion around conv + scan and `cp_context` for both FLA kernels. - Backwards compatible by construction: `cp_context=` is only passed when chunkwise is active, so with it off the FLA call is byte-identical to before and still works against FLA 0.4.1; `_prepare_qkv_for_gated_delta_rule` gains an optional argument so Relax's all-gather fallback keeps calling it unchanged; both existing Relax GDN fixes are preserved verbatim. ## Dependency - `docker/Dockerfile`: `flash-linear-attention==0.4.2` (first release carrying `fla.ops.cp`), plus build-time capability assertions after the FLA install and after the patch apply. One of them asserts `linear_cp_mode` is **absent** from `GatedDeltaNet.forward`, so re-introducing a per-call override fails the build. --- # ✅ Tests ## tests/backends/megatron/test_gdn_chunkwise_cp_layout.py (45 CPU tests) - Both partitions cover every token exactly once for CP in {1,2,4,8} and are permutations of each other. - MCore's zigzag partition is token-for-token identical to Relax's `slice_with_cp` and `gdn_cp_slice`. - Construction gate: default is `headwise`; chunkwise and `all_gather` use the TP-only head rule; `auto` and other unresolved values are rejected. ## tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py (7 NCCL tests) - FLA `causal_conv1d` / `chunk_gated_delta_rule` under `cp_context` vs no CP, in fp32 and bf16, including `dweight`/`dbias`. - Full `GatedDeltaNet` CP=2 vs CP=1 in fp32 and bf16, with headwise run side by side as the control: in fp32 the two CP algorithms' deviation from CP=1 agrees to 1.00x-1.03x per tensor. - zigzag -> contiguous -> zigzag over a real CP group is token-exact, for packed THD with unequal-length samples and for SBHD. - Illegal combinations fail fast: `all_gather` reaching MCore's forward, mismatched `local_cp_size`, chunkwise + deterministic, chunkwise + inference. - `state_dict` / `sharded_state_dict` keys and shard dims identical across CP=1 / headwise / chunkwise. ## tests/backends/megatron/gdn_cp_numeric_probe.py - Cross-image probe for RFC 3.4-2. Old image vs new image on CP=1 / headwise / all-gather: 41 of 45 tensors bitwise identical and **0 outside tolerance**. The four that differ are at relative RMS 1e-9..6e-7 with cosine 1.0000000000 -- FLA 0.4.2 reorders a few backward reductions. All-gather is 18 of 18 bitwise identical. --- # 📝 Documentation - `docker/patch/megatron/TASK32-BACKPORT.md`: file-level and hunk-level record of what came from `5139086e`, which Relax adaptations were made, which existing Relax GDN fixes are preserved, and what was excluded. Includes two mechanical commands a reviewer can run to confirm the new module matches upstream byte-for-byte and that no #5664 content is present. Co-Authored-By: Claude Opus 5 (1M context) --- docker/Dockerfile | 28 +- docker/patch/latest/megatron.patch | 2 +- .../patch/megatron/20260805-85bced0ae.patch | 2414 +++++++++++++++++ docker/patch/megatron/TASK32-BACKPORT.md | 134 + .../backends/megatron/gdn_cp_numeric_probe.py | 300 ++ .../megatron/test_gdn_chunkwise_cp_gpu.py | 689 +++++ .../megatron/test_gdn_chunkwise_cp_layout.py | 214 ++ 7 files changed, 3779 insertions(+), 2 deletions(-) create mode 100644 docker/patch/megatron/20260805-85bced0ae.patch create mode 100644 docker/patch/megatron/TASK32-BACKPORT.md create mode 100644 tests/backends/megatron/gdn_cp_numeric_probe.py create mode 100644 tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py create mode 100644 tests/backends/megatron/test_gdn_chunkwise_cp_layout.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 13a7dd615..c57768f82 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -37,8 +37,22 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29 FROM base as train +# flash-linear-attention 0.4.2 (pulls fla-core 0.4.2) is the first release carrying +# `fla.ops.cp`, which GDN chunkwise context parallel needs. Fail the build here if the +# capability is missing, so a silent wheel/index change can never produce an image whose +# GDN chunkwise path only breaks at training time. NOTE: `fla-core` depends on an +# unpinned `torch`, so pip re-resolves torch's own looser `nvidia-cudnn-cu12` pin here; +# the explicit `nvidia-cudnn-cu12==9.16.0.29` install further down runs after this layer +# and is what the final image keeps. RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir && \ - pip install --no-cache-dir flash-linear-attention==0.4.1 && \ + pip install --no-cache-dir flash-linear-attention==0.4.2 && \ + python -c "import inspect; \ +from fla.ops.cp import build_cp_context; \ +from fla.modules.convolution import causal_conv1d; \ +from fla.ops.gated_delta_rule import chunk_gated_delta_rule; \ +assert 'cp_context' in inspect.signature(causal_conv1d).parameters, 'causal_conv1d lacks cp_context'; \ +assert 'cp_context' in inspect.signature(chunk_gated_delta_rule).parameters, 'chunk_gated_delta_rule lacks cp_context'; \ +print('FLA chunkwise-CP capability OK')" && \ pip install --no-cache-dir tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ # FA3 (Hopper flash-attention), built from source. This commit's _flash_attn_forward carries @@ -135,6 +149,18 @@ RUN cd Megatron-LM && \ exit 1; \ fi && \ rm megatron.patch && \ + python -c "import inspect; \ +import megatron.core.context_parallel_layout as cpl; \ +from megatron.core.packed_seq_params import resolve_cp_group; \ +from megatron.core.ssm.gated_delta_net import GatedDeltaNet; \ +from megatron.core.transformer.transformer_config import TransformerConfig; \ +assert hasattr(cpl, 'zigzag_to_contiguous_chunks') and hasattr(cpl, 'contiguous_to_zigzag_chunks'); \ +assert hasattr(cpl, 'get_thd_context_parallel_rank_indices'); \ +assert hasattr(GatedDeltaNet, '_resolve_cp_routing') and hasattr(GatedDeltaNet, '_build_chunkwise_cp_context'); \ +assert 'linear_cp_mode' not in inspect.signature(GatedDeltaNet.forward).parameters, \ + 'linear_cp_mode must be static config, not a per-forward argument'; \ +assert TransformerConfig.linear_cp_mode == 'headwise', TransformerConfig.linear_cp_mode; \ +print('MCore chunkwise-CP backport OK')" && \ apt update && apt install -y jq && ln -s /usr/local/lib/python3.12/dist-packages/torch_memory_saver_hook_mode_preload_cu12.abi3.so \ /usr/local/lib/python3.12/dist-packages/torch_memory_saver_hook_mode_preload.abi3.so diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch index ec9557dc5..b8992ec40 120000 --- a/docker/patch/latest/megatron.patch +++ b/docker/patch/latest/megatron.patch @@ -1 +1 @@ -../megatron/20260506-85bced0ae.patch \ No newline at end of file +../megatron/20260805-85bced0ae.patch \ No newline at end of file diff --git a/docker/patch/megatron/20260805-85bced0ae.patch b/docker/patch/megatron/20260805-85bced0ae.patch new file mode 100644 index 000000000..ec77afb0f --- /dev/null +++ b/docker/patch/megatron/20260805-85bced0ae.patch @@ -0,0 +1,2414 @@ +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +index 7c73cc1..fe47876 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +@@ -187,8 +187,15 @@ class KimiK25VLBridge(MegatronModelBridge): + for fqn, tensor in converted_weights_dict.items(): + if self._is_quantized_expert_key(fqn): + base = fqn[:-7] if fqn.endswith(".weight") else fqn +- # Preserve the original scale dtype from the HF checkpoint + orig_scale_key = f"{base}.weight_scale" ++ # When the source HF checkpoint has been pre-cast to BF16 (no ++ # `weight_scale` triplet present), passthrough instead of re- ++ # quantizing — downstream sglang loads BF16 and would reject ++ # the INT4 export names with "not found in params_dict". ++ if orig_scale_key not in hf_state_dict: ++ result[fqn] = tensor ++ continue ++ # Preserve the original scale dtype from the HF checkpoint + scale_dtype = ( + hf_state_dict[orig_scale_key].dtype if orig_scale_key in hf_state_dict else torch.bfloat16 + ) +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +index 11b883e..8c6e346 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +@@ -57,6 +57,11 @@ class KimiK25VLModelProvider(MLAModelProvider): + pad_token_id: int = 163839 + ignore_index: int = -100 + ++ # Split vision encoder workload across TP ranks (data-parallel over TP). ++ # Each TP rank processes a chunk of images, then all-reduce gathers the ++ # full embedding. Reduces per-GPU peak memory for the vision encoder. ++ vision_dp_when_tp: bool = False ++ + # Freeze options for fine-tuning scenarios + freeze_language_model: bool = False + freeze_vision_model: bool = False +diff --git a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +index 0e4b120..3c606ac 100644 +--- a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py ++++ b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +@@ -16,6 +16,8 @@ import logging + from typing import List, Optional + + import torch ++import torch.distributed ++from megatron.core import parallel_state as mpu + from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region + from megatron.core.transformer.module import MegatronModule +@@ -23,6 +25,7 @@ from torch import Tensor + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + from megatron.bridge.models.gpt_provider import GPTModelProvider ++from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import preprocess_packed_seqs + from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync + + +@@ -136,11 +139,15 @@ class KimiK25VLModel(MegatronModule): + if not hasattr(MoonViT3dEncoder, "use_deterministic_attn"): + MoonViT3dEncoder.use_deterministic_attn = False + +- # transformers >=5.5 strictly validates `attn_implementation` at +- # __init__ and selects `flash_attention_2` by default when flash-attn +- # is installed. MoonViT3dPretrainedModel doesn't declare flash-attn-2 +- # support, so force eager attention before construction. +- self.vision_tower_config._attn_implementation = "eager" ++ # MoonViT3dPretrainedModel declares `_supports_flash_attn_2 = True` (old ++ # transformers API). transformers >= 5.0 unified flash-attn support under ++ # `_supports_flash_attn` (default False); bridge the old flag so the strict ++ # init-time check in _flash_attn_can_dispatch() does not raise ValueError. ++ if not getattr(MoonViT3dPretrainedModel, "_supports_flash_attn", False): ++ MoonViT3dPretrainedModel._supports_flash_attn = getattr( ++ MoonViT3dPretrainedModel, "_supports_flash_attn_2", False ++ ) ++ self.vision_tower_config._attn_implementation = "flash_attention_2" + self.vision_tower = MoonViT3dPretrainedModel(self.vision_tower_config) + self.mm_projector = PatchMergerMLP(self.projector_config) # TODO: support different types of mm projector + # Ensure HF visual tower params are marked for TP grad sync and future assignments are hooked. +@@ -317,8 +324,79 @@ class KimiK25VLModel(MegatronModule): + + return final_embedding, final_attention_mask, final_labels, position_ids + ++ def _vision_forward_tp_split( ++ self, ++ pixel_values: torch.Tensor, ++ grid_thws: torch.Tensor, ++ ) -> List[torch.Tensor]: ++ """Run vision encoder + projector with workload split across TP ranks. ++ ++ Each TP rank processes a subset of images determined by splitting ++ ``grid_thws``, then the partial feature tensors are all-reduced so ++ every rank holds the complete result. ++ """ ++ tp_rank = mpu.get_tensor_model_parallel_rank() ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ ++ num_images = grid_thws.shape[0] ++ merge_h, merge_w = self.vision_tower.merge_kernel_size ++ param_dtype = next(self.vision_tower.parameters()).dtype ++ text_hidden = self.projector_config.hidden_size ++ ++ pixel_counts = grid_thws.prod(dim=-1) ++ out_token_counts = (grid_thws[:, 1] // merge_h) * (grid_thws[:, 2] // merge_w) ++ total_out_tokens = out_token_counts.sum().item() ++ ++ chunk_indices = list(range(num_images)) ++ chunks = [chunk_indices[i::tp_size] for i in range(tp_size)] ++ my_indices = chunks[tp_rank] if tp_rank < len(chunks) else [] ++ ++ out_buffer = torch.zeros( ++ (total_out_tokens, text_hidden), ++ device=pixel_values.device, ++ dtype=param_dtype, ++ ) ++ ++ if my_indices: ++ pixel_cumsum = pixel_counts.cumsum(dim=0) ++ pv_parts = [] ++ grid_parts = [] ++ for idx in my_indices: ++ px_start = 0 if idx == 0 else pixel_cumsum[idx - 1].item() ++ px_end = pixel_cumsum[idx].item() ++ pv_parts.append(pixel_values[px_start:px_end]) ++ grid_parts.append(grid_thws[idx : idx + 1]) ++ ++ local_pv = torch.cat(pv_parts, dim=0) ++ local_grid = torch.cat(grid_parts, dim=0) ++ ++ local_vit_out = self.vision_tower(local_pv, local_grid) ++ local_features = self.mm_projector(local_vit_out) ++ ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for feat_i, img_idx in enumerate(my_indices): ++ out_start = 0 if img_idx == 0 else out_cumsum[img_idx - 1].item() ++ out_end = out_cumsum[img_idx].item() ++ out_buffer[out_start:out_end] = local_features[feat_i].to(param_dtype) ++ ++ tp_group = mpu.get_tensor_model_parallel_group() ++ torch.distributed.all_reduce(out_buffer, group=tp_group) ++ ++ result = [] ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for i in range(num_images): ++ start = 0 if i == 0 else out_cumsum[i - 1].item() ++ end = out_cumsum[i].item() ++ result.append(out_buffer[start:end]) ++ ++ return result ++ + def _extract_image_features(self, pixel_values, grid_thws): + """Extract and project image features.""" ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ if getattr(self.config, "vision_dp_when_tp", False) and tp_size > 1: ++ return self._vision_forward_tp_split(pixel_values, grid_thws) ++ + image_features = self.vision_tower(pixel_values, grid_thws) + return self.mm_projector(image_features) + +@@ -357,6 +435,12 @@ class KimiK25VLModel(MegatronModule): + 2. Dynamic expansion: input_ids has 1 placeholder per image, expands to N tokens. + """ + if self.pre_process: ++ # Save the caller-supplied per-sample attention mask before any rewrite — ++ # _merge_input_ids_with_image_features sets `attention_mask = None` on the ++ # vision path, but the THD repack below needs the original [B, T] mask to ++ # know each sample's valid length. ++ saved_attention_mask = attention_mask ++ + if inputs_embeds is None: + inputs_embeds = self.language_model.embedding( + input_ids=input_ids, position_ids=None +@@ -392,8 +476,30 @@ class KimiK25VLModel(MegatronModule): + # Don't need attention mask for causal attention. + attention_mask = None + +- # Transpose back to (T, B, D) for Megatron language model +- inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) ++ # When THD packed_seq_params is provided (VL+CP/SP), repack the raw ++ # padded [B, T_max, D] embedding into compact THD ++ # [sum(padded_seqlens)/cp, 1, D] using the saved per-sample attention ++ # mask. Mirrors the Qwen3VL bridge: without this, downstream MLA ++ # attention sees a tensor whose first dim does not match ++ # cu_seqlens_q_padded (sum=sum(padded_seqlens)), and SP scatter can ++ # hit `T_max % tp_size != 0` since T_max is raw batch-max. ++ # preprocess_packed_seqs also recomputes cu_seqlens with align64. ++ needs_thd_repack = ( ++ packed_seq_params is not None ++ and packed_seq_params.qkv_format == "thd" ++ and saved_attention_mask is not None ++ ) ++ if needs_thd_repack: ++ inputs_embeds, packed_seq_params = preprocess_packed_seqs( ++ inputs_embeds, # [B, T_max, D] ++ saved_attention_mask, ++ pre_process=True, ++ ) ++ # preprocess_packed_seqs returns [1, T_thd, D]; switch to (T_thd, 1, D) ++ inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() ++ else: ++ # Transpose back to (T, B, D) for Megatron language model ++ inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) + + if self.config.sequence_parallel: + inputs_embeds = scatter_to_sequence_parallel_region(inputs_embeds) +diff --git a/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +index 04afd67..a83cebf 100755 +--- a/megatron/bridge/models/qwen/qwen3_moe_bridge.py ++++ b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +@@ -69,6 +69,40 @@ class Qwen3MoEBridge(MegatronModelBridge): + + return provider + ++ def build_conversion_tasks(self, hf_pretrained, megatron_model): ++ """Inject virtual .weight keys so INT4 checkpoints (weight_packed/weight_scale/ ++ weight_zero_point) pass the hf_keys validation in the base class. ++ ++ When hf_checkpoint points to an INT4 compressed-tensors checkpoint, expert ++ weights are stored as weight_packed/weight_scale/weight_zero_point triplets ++ with no plain .weight key. The base build_conversion_tasks checks that each ++ mapped HF name exists in hf_keys and skips the param if not found, causing ++ all expert weights to be silently dropped. We patch get_all_keys() to return ++ synthetic .weight keys alongside the real packed keys so the check passes. ++ Downstream quantize_params in HfWeightIteratorBridge then converts the BF16 ++ output back to INT4 before sending to the rollout engine. ++ """ ++ if not (hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")): ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ ++ original_get_all_keys = hf_pretrained.state.source.get_all_keys ++ ++ def _get_all_keys_with_virtual(): ++ keys = original_get_all_keys() ++ all_keys_set = set(keys) ++ virtual_keys = [ ++ key[:-7] # "...weight_packed" -> "...weight" ++ for key in keys ++ if key.endswith("_packed") and f"{key[:-7]}_scale" in all_keys_set ++ ] ++ return keys + virtual_keys ++ ++ hf_pretrained.state.source.get_all_keys = _get_all_keys_with_virtual ++ try: ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ finally: ++ hf_pretrained.state.source.get_all_keys = original_get_all_keys ++ + def mapping_registry(self) -> MegatronMappingRegistry: + # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format + # First create simple 1:1 parameter mappings using a dictionary for readability +diff --git a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +index d067655..3a4e744 100644 +--- a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py ++++ b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +@@ -120,6 +120,7 @@ class Qwen3VLGPTModel(GPTModel): + *, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, ++ mtp_kwargs: Optional[dict] = None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, +@@ -210,6 +211,7 @@ class Qwen3VLGPTModel(GPTModel): + runtime_gather_output=runtime_gather_output, + extra_block_kwargs=extra_block_kwargs, + inference_context=inference_context, ++ mtp_kwargs=mtp_kwargs or {}, + ) + + if _shadow_embedding: +diff --git a/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py b/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py +index 7fcf295..7ac1134 100644 +--- a/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py ++++ b/megatron/bridge/models/qwen_vl/qwen35_vl_bridge.py +@@ -404,15 +404,6 @@ class Qwen35VLMoEBridge(MegatronModelBridge): + k="mtp.layers.*.self_attn.k_proj.weight", + v="mtp.layers.*.self_attn.v_proj.weight", + ), +- GatedMLPMapping( +- megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc1.weight*", +- gate="mtp.layers.*.mlp.experts.*.gate_proj.weight", +- up="mtp.layers.*.mlp.experts.*.up_proj.weight", +- ), +- AutoMapping( +- megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc2.weight*", +- hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight", +- ), + GatedMLPMapping( + megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.shared_experts.linear_fc1.weight", + gate="mtp.layers.*.mlp.shared_expert.gate_proj.weight", +@@ -429,6 +420,43 @@ class Qwen35VLMoEBridge(MegatronModelBridge): + ] + ) + ++ # Detect MTP MoE expert weight format: Qwen3.5 stores per-expert ++ # (mtp.layers.0.mlp.experts.{i}.gate_proj.weight), Qwen3.6 stores packed ++ # (mtp.layers.0.mlp.experts.gate_up_proj). Same architecture string, ++ # different storage — must inspect HF keys. ++ mtp_experts_packed = False ++ if hasattr(self.hf_pretrained, "state") and hasattr(self.hf_pretrained.state, "source"): ++ hf_keys = set(self.hf_pretrained.state.source.get_all_keys()) ++ if "mtp.layers.0.mlp.experts.gate_up_proj" in hf_keys: ++ mtp_experts_packed = True ++ ++ if mtp_experts_packed: ++ # Qwen3.6: packed format (same as main decoder) ++ mapping_list.extend([ ++ FusedGatedExpertMapping( ++ megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc1.weight*", ++ hf_param="mtp.layers.*.mlp.experts.gate_up_proj", ++ ), ++ FusedExpertMapping( ++ megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc2.weight*", ++ hf_param="mtp.layers.*.mlp.experts.down_proj", ++ transpose_on_export=True, ++ ), ++ ]) ++ else: ++ # Qwen3.5: per-expert format (current behavior) ++ mapping_list.extend([ ++ GatedMLPMapping( ++ megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc1.weight*", ++ gate="mtp.layers.*.mlp.experts.*.gate_proj.weight", ++ up="mtp.layers.*.mlp.experts.*.up_proj.weight", ++ ), ++ AutoMapping( ++ megatron_param="language_model.mtp.layers.*.mtp_model_layer.mlp.experts.linear_fc2.weight*", ++ hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight", ++ ), ++ ]) ++ + return MegatronMappingRegistry(*mapping_list) + + +diff --git a/megatron/bridge/peft/utils.py b/megatron/bridge/peft/utils.py +index c297a23..587e1de 100644 +--- a/megatron/bridge/peft/utils.py ++++ b/megatron/bridge/peft/utils.py +@@ -15,28 +15,36 @@ + import logging + import math + import re +-from dataclasses import dataclass ++from dataclasses import dataclass, fields ++from importlib import import_module + from importlib.metadata import version +-from typing import Callable, Dict, List, Optional, Tuple ++from pathlib import Path ++from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + + import packaging + import torch + import torch.nn as nn + from megatron.core import ModelParallelConfig, parallel_state + from megatron.core.dist_checkpointing.mapping import ShardedStateDict, ShardedTensor, ShardedTensorFactory ++from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear + from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, + ) + from megatron.core.transformer.mlp import apply_swiglu_sharded_factory ++from megatron.core.transformer.module import MegatronModule + from megatron.core.transformer.moe.router import TopKRouter + ++from megatron.bridge.utils.activation_map import str_to_dtype + from megatron.bridge.utils.import_utils import safe_import_from + + + logger = logging.getLogger(__name__) + ++ModelList = list[MegatronModule] ++CheckpointPath = str | Path ++ + + TEColumnParallelLinear, HAVE_TE_COL_LINEAR = safe_import_from( + "megatron.core.extensions.transformer_engine", "TEColumnParallelLinear" +@@ -1479,3 +1487,148 @@ class GroupedExpertLinearAdapter(nn.Module): + sharded_state_dict.update(linear_in_sd) + sharded_state_dict.update(linear_out_sd) + return sharded_state_dict ++ ++ ++def create_peft(config: Mapping[str, Any], dtype: torch.dtype | str | int | None = None) -> object | None: ++ """Create a Bridge PEFT object from a small config mapping.""" ++ kwargs = dict(config) ++ peft_type = kwargs.pop("type", "lora") ++ if "rank" in kwargs: ++ kwargs["dim"] = kwargs.pop("rank") ++ if kwargs.get("dim", 0) <= 0: ++ return None ++ ++ peft_cls = _import_peft_class(peft_type) ++ ++ peft_fields = {field.name for field in fields(peft_cls) if field.init} ++ config_dtype = kwargs.pop("dtype", None) ++ if "lora_dtype" not in kwargs: ++ kwargs["lora_dtype"] = config_dtype if config_dtype is not None else dtype ++ ++ if kwargs.get("lora_dtype") is None or "lora_dtype" not in peft_fields: ++ kwargs.pop("lora_dtype", None) ++ else: ++ kwargs["lora_dtype"] = str_to_dtype(str(kwargs["lora_dtype"]).lower()) ++ ++ kwargs = {key: value for key, value in kwargs.items() if key in peft_fields} ++ ++ return peft_cls(**kwargs) ++ ++ ++def _import_peft_class(peft_type: str) -> type[Any]: ++ peft_classes = { ++ "lora": ("megatron.bridge.peft.lora", "LoRA"), ++ "vlm_lora": ("megatron.bridge.peft.lora", "VLMLoRA"), ++ "canonical_lora": ("megatron.bridge.peft.canonical_lora", "CanonicalLoRA"), ++ "dora": ("megatron.bridge.peft.dora", "DoRA"), ++ } ++ if peft_type not in peft_classes: ++ supported_types = ", ".join(sorted(peft_classes)) ++ raise ValueError(f"Unsupported PEFT type {peft_type!r}. Supported types: {supported_types}.") ++ ++ module_name, class_name = peft_classes[peft_type] ++ try: ++ module = import_module(module_name) ++ except ImportError as err: ++ message = f"Failed to import PEFT type {peft_type!r} ({module_name}:{class_name})." ++ if peft_type in {"lora", "vlm_lora", "canonical_lora"}: ++ message += " Install Megatron Bridge with the [te] extra for Transformer Engine support." ++ raise ImportError(message) from err ++ ++ return getattr(module, class_name) ++ ++ ++def load_peft_adapter_checkpoint( ++ model: ModelList | MegatronModule, ++ adapter_checkpoint_path: CheckpointPath, ++ peft: object, ++ strict: bool = False, ++ model_sd_kwargs: Mapping[str, object] | None = None, ++ ckpt_format: str = "torch_dist", ++ pg_collection: ProcessGroupCollection | None = None, ++ fully_parallel_load: bool = True, ++ load_strategy: object | None = None, ++) -> None: ++ """Load a PEFT adapter checkpoint into an already transformed model.""" ++ from megatron.core import dist_checkpointing ++ from megatron.core.dist_checkpointing.serialization import get_default_load_sharded_strategy ++ from megatron.core.dist_checkpointing.strategies.fully_parallel import FullyParallelLoadStrategyWrapper ++ ++ from megatron.bridge.training.checkpointing import apply_peft_adapter_filter_to_state_dict ++ ++ model_chunks = _ensure_model_list(model) ++ sharded_state_dict = _model_state_dict( ++ model_chunks, ++ model_sd_kwargs, ++ ckpt_format, ++ pg_collection=pg_collection, ++ ) ++ sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, peft) ++ ++ checkpoint_path = str(adapter_checkpoint_path) ++ if load_strategy is None: ++ load_strategy = get_default_load_sharded_strategy(checkpoint_path) ++ if pg_collection is None and fully_parallel_load: ++ try: ++ pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["dp_cp"]) ++ except AssertionError: ++ pg_collection = None ++ dp_cp_group = _get_process_group(pg_collection, "dp_cp") ++ if fully_parallel_load and dp_cp_group is not None: ++ load_strategy = FullyParallelLoadStrategyWrapper(load_strategy, dp_cp_group) ++ ++ loaded_state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_path, load_strategy) ++ for vpp_rank, model_chunk in enumerate(model_chunks): ++ model_key = "model" if len(model_chunks) == 1 else f"model{vpp_rank}" ++ if model_key not in loaded_state_dict: ++ if len(model_chunks) == 1: ++ fallback_model_key = next((key for key in loaded_state_dict if key.startswith("model")), None) ++ if fallback_model_key is not None: ++ model_key = fallback_model_key ++ else: ++ raise KeyError( ++ "Expected adapter checkpoint to contain a top-level 'model' or 'model*' key, " ++ f"but found keys: {list(loaded_state_dict.keys())}" ++ ) ++ else: ++ expected_model_keys = [f"model{rank}" for rank in range(len(model_chunks))] ++ raise KeyError( ++ f"Expected adapter checkpoint to contain top-level key {model_key!r} " ++ f"for virtual pipeline model chunk {vpp_rank} " ++ f"(expected keys: {expected_model_keys}), " ++ f"but found keys: {list(loaded_state_dict.keys())}" ++ ) ++ model_chunk.load_state_dict(loaded_state_dict[model_key], strict=strict) ++ ++ ++def _model_state_dict( ++ model: ModelList, ++ model_sd_kwargs: Mapping[str, object] | None = None, ++ ckpt_format: str = "torch_dist", ++ pg_collection: ProcessGroupCollection | None = None, ++) -> dict[str, Any]: ++ """Generate Bridge model checkpoint sections for an external trainer.""" ++ from megatron.bridge.training.checkpointing import _generate_model_state_dict ++ ++ return _generate_model_state_dict( ++ model, ++ dict(model_sd_kwargs or {}), ++ ckpt_format, ++ pg_collection=pg_collection, ++ ) ++ ++ ++def _ensure_model_list(model: ModelList | MegatronModule) -> ModelList: ++ return model if isinstance(model, list) else [model] ++ ++ ++def _get_process_group(pg_collection: ProcessGroupCollection | None, *names: str) -> object | None: ++ """Return the first named process group available on a collection.""" ++ ++ if pg_collection is None: ++ return None ++ for name in names: ++ group = getattr(pg_collection, name, None) ++ if group is not None: ++ return group ++ return None +diff --git a/megatron/core/context_parallel_layout.py b/megatron/core/context_parallel_layout.py +new file mode 100644 +index 0000000..e3acdd9 +--- /dev/null ++++ b/megatron/core/context_parallel_layout.py +@@ -0,0 +1,324 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. ++ ++"""Context parallel tensor layout helpers. ++ ++Backported VERBATIM from Megatron-LM commit 5139086e (NVIDIA/Megatron-LM#3282, ++merged to dev) for Relax Task 32 v1. No Relax modifications -- if this file ever ++diverges from upstream, that is a bug. ++ ++Two layouts: ++ ++``zigzag`` ++ Megatron's attention load-balanced partition. Each sequence is cut into ++ ``2 * cp_size`` chunks and rank ``r`` owns chunks ``r`` and ++ ``2 * cp_size - r - 1``. ++``contiguous`` ++ The flattened packed buffer is cut into ``cp_size`` equal spans and rank ++ ``r`` owns span ``r``. This is the layout FLA's chunkwise CP kernels expect: ++ ``fla.ops.cp.build_cp_context`` derives each rank's local ``cu_seqlens`` ++ from exactly this partition. ++""" ++ ++from typing import List, Optional, Tuple ++ ++import torch ++ ++from megatron.core.tensor_parallel import all_to_all ++ ++ ++def get_thd_context_parallel_rank_indices( ++ cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, layout: str ++) -> torch.Tensor: ++ """Return global THD token indices owned by one CP rank in a layout. ++ ++ Args: ++ cu_seqlens: Global packed-sequence cumulative lengths before CP partitioning. ++ cp_size: Context-parallel group size. ++ cp_rank: Context-parallel rank. ++ layout: Either ``"zigzag"`` or ``"contiguous"``. ++ ++ The returned indices are ordered exactly as the rank-local THD tensor is stored. ++ ``"zigzag"`` follows Megatron's per-sequence load-balanced chunk order; ``"contiguous"`` ++ partitions the flattened packed THD buffer into rank-contiguous spans. ++ """ ++ if layout not in ("zigzag", "contiguous"): ++ raise ValueError(f"Unsupported context-parallel layout {layout!r}.") ++ if cp_size < 1: ++ raise ValueError(f"cp_size must be >= 1, got {cp_size}.") ++ if not 0 <= cp_rank < cp_size: ++ raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}.") ++ if cu_seqlens.dim() != 1: ++ raise ValueError(f"cu_seqlens must be 1-D, got shape {tuple(cu_seqlens.shape)}.") ++ ++ cu = cu_seqlens.to(dtype=torch.long) ++ if cu.numel() == 0 or cu[0].item() != 0: ++ raise ValueError(f"cu_seqlens must start at 0, got {cu_seqlens}.") ++ ++ if torch.any(torch.diff(cu) < 0): ++ raise ValueError(f"cu_seqlens must be nondecreasing, got {cu_seqlens}.") ++ ++ nonduplicate_boundaries = torch.ones(cu.numel(), device=cu.device, dtype=torch.bool) ++ nonduplicate_boundaries[1:] = cu[1:] != cu[:-1] ++ cu = cu[nonduplicate_boundaries] ++ ++ total_tokens = int(cu[-1].item()) ++ positions = torch.arange(total_tokens, device=cu.device, dtype=torch.long) ++ if total_tokens == 0: ++ return positions ++ ++ seq_lens = torch.diff(cu) ++ chunk_divisor = 2 * cp_size ++ if torch.any(seq_lens % chunk_divisor != 0): ++ raise ValueError( ++ "All packed sequence lengths must be divisible by " ++ f"2 * cp_size ({chunk_divisor}) for zigzag/contiguous CP layout conversion, " ++ f"got {seq_lens}." ++ ) ++ ++ if layout == "contiguous": ++ part_len = total_tokens // cp_size ++ rank_start = cp_rank * part_len ++ return positions[rank_start : rank_start + part_len] ++ ++ seq_idx = torch.bucketize(positions, cu[1:], right=True) ++ global_starts = cu[:-1] ++ pos_in_seq = positions - global_starts[seq_idx] ++ chunk_lens = (seq_lens // chunk_divisor)[seq_idx] ++ chunk = pos_in_seq // chunk_lens ++ offset = pos_in_seq - chunk * chunk_lens ++ ++ owner = torch.where(chunk < cp_size, chunk, 2 * cp_size - chunk - 1) ++ local_slot = torch.where(chunk < cp_size, torch.zeros_like(chunk), torch.ones_like(chunk)) ++ ++ local_starts = (global_starts // cp_size)[seq_idx] ++ local_pos = local_starts + local_slot * chunk_lens + offset ++ ++ rank_mask = owner == cp_rank ++ rank_positions = positions[rank_mask] ++ rank_local_pos = local_pos[rank_mask] ++ return rank_positions[torch.argsort(rank_local_pos)] ++ ++ ++def zigzag_to_contiguous_chunks( ++ x: torch.Tensor, ++ cp_group: torch.distributed.ProcessGroup, ++ seq_dim: int = 0, ++ cu_seqlens: Optional[torch.Tensor] = None, ++) -> torch.Tensor: ++ """Permute CP chunks from Megatron zigzag layout to contiguous-time layout. ++ ++ SBHD tensors have two equal chunks per rank along ``seq_dim`` and use a ++ chunk-level all-to-all. THD tensors pass global ``cu_seqlens`` and use one ++ packed-token all-to-all over the whole local THD tensor. ++ """ ++ if cu_seqlens is not None: ++ return _zigzag_contiguous_thd_swap( ++ x, cp_group, seq_dim, cu_seqlens, source_layout="zigzag", target_layout="contiguous" ++ ) ++ return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=True) ++ ++ ++def contiguous_to_zigzag_chunks( ++ x: torch.Tensor, ++ cp_group: torch.distributed.ProcessGroup, ++ seq_dim: int = 0, ++ cu_seqlens: Optional[torch.Tensor] = None, ++) -> torch.Tensor: ++ """Inverse of :func:`zigzag_to_contiguous_chunks`.""" ++ if cu_seqlens is not None: ++ return _zigzag_contiguous_thd_swap( ++ x, cp_group, seq_dim, cu_seqlens, source_layout="contiguous", target_layout="zigzag" ++ ) ++ return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=False) ++ ++ ++def _zigzag_contiguous_thd_swap( ++ x: torch.Tensor, ++ cp_group: Optional[torch.distributed.ProcessGroup], ++ seq_dim: int, ++ cu_seqlens: torch.Tensor, ++ source_layout: str, ++ target_layout: str, ++) -> torch.Tensor: ++ """Single-all-to-all THD permutation between zigzag and contiguous layouts. ++ ++ The packed THD tensor stays packed: we first group local tokens by their ++ target CP rank, exchange those groups once, then scatter received tokens ++ back into the target rank-local order. ++ """ ++ cp_size = cp_group.size() if cp_group is not None else 1 ++ if cp_size == 1: ++ return x ++ cp_rank = cp_group.rank() ++ ++ if seq_dim != 0: ++ x = x.movedim(seq_dim, 0) ++ x = x.contiguous() ++ ++ cu = cu_seqlens.to(device=x.device, dtype=torch.long) ++ # TODO: Let a future CP layout scheduler precompute this routing once per ++ # microbatch from immutable cu_seqlens and pass it through both THD swaps. ++ # Do not cache it across microbatches because packed sequence boundaries change. ++ source_by_rank = [ ++ get_thd_context_parallel_rank_indices(cu, cp_size, rank, source_layout) ++ for rank in range(cp_size) ++ ] ++ target_by_rank = [ ++ get_thd_context_parallel_rank_indices(cu, cp_size, rank, target_layout) ++ for rank in range(cp_size) ++ ] ++ ++ local_source_indices = source_by_rank[cp_rank] ++ local_target_indices = target_by_rank[cp_rank] ++ if x.size(0) != local_source_indices.numel(): ++ raise ValueError( ++ f"Local THD tensor length ({x.size(0)}) does not match {source_layout} " ++ f"rank-{cp_rank} partition length ({local_source_indices.numel()})." ++ ) ++ ++ total_tokens = int(cu[-1].item()) ++ target_owner = torch.empty(total_tokens, device=x.device, dtype=torch.long) ++ target_local_pos = torch.empty(total_tokens, device=x.device, dtype=torch.long) ++ for rank, indices in enumerate(target_by_rank): ++ target_owner[indices] = rank ++ target_local_pos[indices] = torch.arange(indices.numel(), device=x.device) ++ ++ local_target_owner = target_owner[local_source_indices] ++ local_target_pos = target_local_pos[local_source_indices] ++ ++ send_parts: List[torch.Tensor] = [] ++ input_split_sizes: List[int] = [] ++ for dst_rank in range(cp_size): ++ dst_mask = local_target_owner == dst_rank ++ dst_rows = dst_mask.nonzero(as_tuple=False).flatten() ++ if dst_rows.numel() > 0: ++ dst_rows = dst_rows[torch.argsort(local_target_pos[dst_rows])] ++ send_part = x.index_select(0, dst_rows) ++ else: ++ send_part = x.narrow(0, 0, 0) ++ send_parts.append(send_part) ++ input_split_sizes.append(send_part.size(0)) ++ send_buf = torch.cat(send_parts, dim=0).contiguous() ++ ++ output_split_sizes: List[int] = [] ++ recv_target_positions: List[torch.Tensor] = [] ++ for src_rank in range(cp_size): ++ src_indices = source_by_rank[src_rank] ++ src_to_this_rank = target_owner[src_indices] == cp_rank ++ recv_global_indices = src_indices[src_to_this_rank] ++ if recv_global_indices.numel() > 0: ++ recv_positions = target_local_pos[recv_global_indices] ++ recv_positions = recv_positions[torch.argsort(recv_positions)] ++ else: ++ recv_positions = local_target_indices.narrow(0, 0, 0) ++ recv_target_positions.append(recv_positions) ++ output_split_sizes.append(recv_positions.numel()) ++ ++ recv_buf = all_to_all(cp_group, send_buf, output_split_sizes, input_split_sizes) ++ ++ out_shape = (local_target_indices.numel(),) + tuple(x.shape[1:]) ++ out = x.new_empty(out_shape) ++ offset = 0 ++ for recv_positions in recv_target_positions: ++ recv_len = recv_positions.numel() ++ if recv_len > 0: ++ out[recv_positions] = recv_buf[offset : offset + recv_len] ++ offset += recv_len ++ ++ if seq_dim != 0: ++ out = out.movedim(0, seq_dim) ++ return out.contiguous() ++ ++ ++def _zigzag_contiguous_chunk_swap( ++ x: torch.Tensor, ++ cp_group: Optional[torch.distributed.ProcessGroup], ++ seq_dim: int, ++ to_contiguous: bool, ++) -> torch.Tensor: ++ """Single-all-to-all chunk permutation between zigzag and contiguous layouts. ++ ++ Each rank holds exactly two chunks along ``seq_dim``. The mapping from ++ local (rank, slot) to (rank, slot) in the target layout is deterministic ++ and depends only on ``cp_size`` and ``cp_rank``, so we pack send data in ++ destination-rank order and use one ``all_to_all_single`` with unequal ++ splits to route each chunk to its target rank. ++ """ ++ cp_size = cp_group.size() if cp_group is not None else 1 ++ if cp_size == 1: ++ return x ++ cp_rank = cp_group.rank() ++ ++ # Work with seq_dim at position 0. ++ if seq_dim != 0: ++ x = x.movedim(seq_dim, 0) ++ x = x.contiguous() ++ ++ seq_len_local = x.size(0) ++ assert seq_len_local % 2 == 0, ( ++ f"zigzag/contiguous chunk swap requires an even local sequence length, " ++ f"got {seq_len_local}." ++ ) ++ chunk_len = seq_len_local // 2 ++ ++ def _rank_to_chunks(rank: int, in_zigzag: bool) -> Tuple[int, int]: ++ """Global chunk indices at (slot 0, slot 1) for this rank.""" ++ if in_zigzag: ++ return (rank, 2 * cp_size - rank - 1) ++ return (2 * rank, 2 * rank + 1) ++ ++ def _chunk_to_dest(chunk_idx: int, target_zigzag: bool) -> Tuple[int, int]: ++ """Destination (rank, slot) for a given global chunk index in the target layout.""" ++ if target_zigzag: ++ if chunk_idx < cp_size: ++ return chunk_idx, 0 ++ return 2 * cp_size - chunk_idx - 1, 1 ++ return chunk_idx // 2, chunk_idx % 2 ++ ++ source_in_zigzag = to_contiguous ++ target_in_zigzag = not to_contiguous ++ ++ local_chunk_indices = _rank_to_chunks(cp_rank, source_in_zigzag) ++ local_dests = [_chunk_to_dest(c, target_in_zigzag) for c in local_chunk_indices] ++ ++ # Pack the send buffer so chunks are ordered by (dst_rank, dst_slot). ++ local_slot_order = sorted(range(2), key=lambda s: local_dests[s]) ++ local_chunks = [x[:chunk_len], x[chunk_len:]] ++ send_buf = torch.cat([local_chunks[s] for s in local_slot_order], dim=0).contiguous() ++ ++ input_split_chunks = [0] * cp_size ++ for dst_rank, _ in local_dests: ++ input_split_chunks[dst_rank] += 1 ++ ++ # Mirror every source rank's packing logic so we know which received chunk ++ # belongs in which local target slot. ++ output_split_chunks = [0] * cp_size ++ recv_dst_slots_per_source: List[List[int]] = [[] for _ in range(cp_size)] ++ for src in range(cp_size): ++ src_chunks = _rank_to_chunks(src, source_in_zigzag) ++ src_dests = [_chunk_to_dest(c, target_in_zigzag) for c in src_chunks] ++ src_slot_order = sorted(range(2), key=lambda s: src_dests[s]) ++ for s in src_slot_order: ++ dst_rank, dst_slot = src_dests[s] ++ if dst_rank == cp_rank: ++ output_split_chunks[src] += 1 ++ recv_dst_slots_per_source[src].append(dst_slot) ++ ++ input_split_sizes = [n * chunk_len for n in input_split_chunks] ++ output_split_sizes = [n * chunk_len for n in output_split_chunks] ++ ++ recv_buf = all_to_all(cp_group, send_buf, output_split_sizes, input_split_sizes) ++ ++ # Reassemble local chunks in target-layout slot order. ++ target_slots: List[Optional[torch.Tensor]] = [None, None] ++ offset = 0 ++ for src in range(cp_size): ++ for dst_slot in recv_dst_slots_per_source[src]: ++ target_slots[dst_slot] = recv_buf[offset : offset + chunk_len] ++ offset += chunk_len ++ assert all(t is not None for t in target_slots), "Incomplete chunk reassembly in CP swap" ++ ++ out = torch.cat(target_slots, dim=0) ++ if seq_dim != 0: ++ out = out.movedim(0, seq_dim) ++ return out.contiguous() +diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py +index 58e1e56..abe561d 100644 +--- a/megatron/core/dist_checkpointing/strategies/torch.py ++++ b/megatron/core/dist_checkpointing/strategies/torch.py +@@ -501,10 +501,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + def _validate_global_shapes(self, metadata, sharded_tensors): + for sh_ten in sharded_tensors: + if sh_ten.key not in metadata.state_dict_metadata: +- raise KeyError( +- f"{sh_ten.key} from model not in state dict:" +- f" {sorted(metadata.state_dict_metadata.keys())}" +- ) ++ # raise KeyError( ++ # f"{sh_ten.key} from model not in state dict:" ++ # f" {sorted(metadata.state_dict_metadata.keys())}" ++ # ) ++ print(f"{sh_ten.key} from model not in state dict, will skip") ++ continue + loaded_shape = metadata.state_dict_metadata[sh_ten.key].size + expected_shape = sh_ten.global_shape + if loaded_shape != expected_shape: +@@ -528,7 +530,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + tensor_metadata = self.metadata.state_dict_metadata + metadata_with_sizes = [ + (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor) +- for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() ++ for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata + ] + try: + # Temporarily set sizes to expected shapes +@@ -865,6 +867,7 @@ class TorchDistLoadShardedStrategy: + planner=MCoreLoadPlanner( + shapes_validation_sharded_tensors=flexible_shape_sharded_tensors, + allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors, ++ allow_partial_load=True, + flatten_state_dict=False, + flatten_sharded_tensors=False, + ), +diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py +index 2a82a1e..5441f93 100644 +--- a/megatron/core/extensions/transformer_engine.py ++++ b/megatron/core/extensions/transformer_engine.py +@@ -836,6 +836,7 @@ class TELinear(te.pytorch.Linear): + self.te_quant_params: Optional[TEQuantizationParams] = None + + for param in self.parameters(): ++ setattr(param, "parallel_mode", parallel_mode) + if is_expert: + # Reduce the gradient on the expert_data_parallel group for expert linear layers + setattr(param, "allreduce", not self.expert_parallel) +@@ -1671,6 +1672,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention): + + + if HAVE_TE and is_te_min_version("1.9.0.dev0"): ++ def ceil_div(x: int, y: int) -> int: ++ return (x + y - 1) // y ++ ++ class _FakeInt4QuantizationSTE(torch.autograd.Function): ++ @staticmethod ++ def forward(ctx, x, group_size): ++ m, n = x.shape ++ block_size_m, block_size_n = 1, group_size ++ ++ ++ m_padded = ceil_div(m, block_size_m) * block_size_m ++ n_padded = ceil_div(n, block_size_n) * block_size_n ++ ++ x_padded = torch.zeros( ++ (m_padded, n_padded), ++ dtype=x.dtype, device=x.device ++ ) ++ x_padded[:m, :n] = x ++ ++ x_view = x_padded.view( ++ m_padded // block_size_m, ++ block_size_m, ++ n_padded // block_size_n, ++ block_size_n ++ ) ++ ++ x_max = x_view.abs().float().amax(dim=(1, 3), keepdim=True) ++ q_max = 7 ++ x_scale = x_max / q_max ++ ++ x_scale = x_scale.clamp(min=1e-5) ++ ++ x_div = x_view / x_scale ++ x_round = torch.round(x_div) ++ ++ x_q_clamped = x_round.clamp(-q_max, q_max) ++ ++ x_dequant_view = x_q_clamped * x_scale ++ ++ x_dequant_full = x_dequant_view.view_as(x_padded) ++ x_out = x_dequant_full[:m, :n].contiguous().to(x.dtype) ++ ++ return x_out ++ ++ @staticmethod ++ def backward(ctx, grad_output): ++ return grad_output, None ++ ++ def fake_int4_quantization_ste(x, group_size): ++ x_out = _FakeInt4QuantizationSTE.apply(x, group_size) ++ ++ if hasattr(x, 'main_grad'): ++ x_out.main_grad = x.main_grad ++ ++ return x_out + + class TEGroupedLinear(te.pytorch.GroupedLinear): + """ +@@ -1913,6 +1969,7 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + "amax_history_bwd": torch.cat( + [state["amax_history_bwd"].view(-1, 1) for state in state_list], + dim=1, ++ + ).view(self.fp8_meta["recipe"].amax_history_len, -1), + } + ) +@@ -1990,6 +2047,20 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + return out + return out, None + ++ def _get_weight_tensors(self): ++ """Get the weight tensors of the module.""" ++ weight_tensors = super()._get_weight_tensors() ++ ++ if os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1": ++ group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128")) ++ ++ weight_tensors = [ ++ fake_int4_quantization_ste(w, group_size) ++ for w in weight_tensors ++ ] ++ ++ return weight_tensors ++ + def _encode_extra_state(self, state): + # TE 2.0 changed the format of extra_state to be a byte tensor + if is_te_min_version("2.0.0"): +diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +index 1fd5dcf..c9aeef1 100644 +--- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py ++++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +@@ -385,6 +385,7 @@ def rotary_fwd_kv_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -434,21 +435,27 @@ def rotary_fwd_kv_kernel( + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + +- KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads +- kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads +- mask = kv_off < head_num * stride_kv_nheads +- k_in_off = kv_off + tl.arange(0, k_dim)[None, :] +- v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] +- k = tl.load(KV_ptr + k_in_off, mask=mask) +- v = tl.load(KV_ptr + v_in_off, mask=mask) ++ KV_ptr = KV + pid_m * stride_kv_seq # + pid_head * BLOCK_H * stride_kv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ k_off = ki_range * stride_kv_nheads + kj_range ++ if v_dim > 0: ++ v_off = ki_range * stride_kv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ v = tl.load(KV_ptr + v_off, mask=mask_v) ++ else: ++ v = tl.zeros((BLOCK_H, 1), dtype=KV.dtype.element_ty) ++ k = tl.load(KV_ptr + k_off, mask=mask_k) + +- K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads +- V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads ++ K_ptr = O_KEY + pid_m * stride_k_seq # + pid_head * BLOCK_H * stride_k_nheads ++ V_ptr = O_VALUE + pid_m * stride_v_seq # + pid_head * BLOCK_H * stride_v_nheads + +- k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] +- v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] +- tl.store(K_ptr + k_out_off, k, mask=mask) +- tl.store(V_ptr + v_out_off, v, mask=mask) ++ k_out_off = ki_range * stride_k_nheads + kj_range ++ tl.store(K_ptr + k_out_off, k, mask=mask_k) ++ if v_dim > 0: ++ v_out_off = ki_range * stride_v_nheads + tl.arange(0, v_dim)[None, :] ++ tl.store(V_ptr + v_out_off, v, mask=mask_v) + + EMB = K_POS_EMB + pid_m * stride_emb_seq + # x1 = t[..., 0::2], x2 = t[..., 1::2] +@@ -460,14 +467,16 @@ def rotary_fwd_kv_kernel( + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + ++ x_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ mask_x = x_range < head_num + x_left_off = ( +- tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads ++ x_range * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 +- tl.store(K_ptr + x_left_off, x_left, mask=mask) +- tl.store(K_ptr + x_right_off, x_right, mask=mask) ++ tl.store(K_ptr + x_left_off, x_left, mask=mask_x) ++ tl.store(K_ptr + x_right_off, x_right, mask=mask_x) + + + @triton.autotune( +@@ -493,6 +502,7 @@ def rotary_bwd_kv_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -533,27 +543,32 @@ def rotary_bwd_kv_kernel( + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + +- dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads +- dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads +- mask = dkv_off < head_num * stride_dkv_nheads +- dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] +- dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] +- +- dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads +- dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads +- dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] +- dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] +- dk = tl.load(dK_ptr + dk_in_off, mask=mask) +- dv = tl.load(dV_ptr + dv_in_off, mask=mask) +- tl.store(dKV_ptr + dk_out_off, dk, mask=mask) +- tl.store(dKV_ptr + dv_out_off, dv, mask=mask) ++ dKV_ptr = dKV + pid_m * stride_dkv_seq # + pid_head * BLOCK_H * stride_dkv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ dk_out_off = ki_range * stride_dkv_nheads + kj_range ++ ++ dK_ptr = dK + pid_m * stride_dk_seq # + pid_head * BLOCK_H * stride_dk_nheads ++ dV_ptr = dV + pid_m * stride_dv_seq # + pid_head * BLOCK_H * stride_dv_nheads ++ dk_in_off = ki_range * stride_dk_nheads + kj_range ++ ++ dk = tl.load(dK_ptr + dk_in_off, mask=mask_k) ++ tl.store(dKV_ptr + dk_out_off, dk, mask=mask_k) ++ ++ if v_dim > 0: ++ dv_out_off = ki_range * stride_dkv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :] ++ dv = tl.load(dV_ptr + dv_in_off, mask=mask_v) ++ tl.store(dKV_ptr + dv_out_off, dv, mask=mask_v) + + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): +- dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads +- x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim ++ dK_ptr = dK + pid_m * stride_dk_seq # + i * BLOCK_H * stride_dk_nheads ++ x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim + i * BLOCK_H * stride_dk_nheads + mask = x_off < head_num * stride_dk_nheads + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 +@@ -632,6 +647,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): + + o_key = kv.new_empty(total_seqlen, nheads, emb_dim + k_dim) + o_value = kv.new_empty(total_seqlen, nheads, v_dim) ++ k_dim_ceil = triton.next_power_of_2(k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_kv_kernel[grid]( +@@ -643,6 +659,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): + sin, + emb_dim, + k_dim, ++ k_dim_ceil, + v_dim, + nheads, + batch_size, +@@ -700,6 +717,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): + + d_kv = dk.new_empty(total_seqlen, nheads, ctx.k_dim + ctx.v_dim) + d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) ++ k_dim_ceil = triton.next_power_of_2(ctx.k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_kv_kernel[grid]( +@@ -711,6 +729,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function): + sin, + ctx.emb_dim, + ctx.k_dim, ++ k_dim_ceil, + ctx.v_dim, + nheads, + batch_size, +diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py +index 3ff370f..21858ea 100644 +--- a/megatron/core/inference/contexts/dynamic_context.py ++++ b/megatron/core/inference/contexts/dynamic_context.py +@@ -57,7 +57,8 @@ except ImportError: + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # Commented out: breaks SGLang CUDA graph (requires hook_mode="preload") ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False +diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py +index 92d561c..d4f62cb 100755 +--- a/megatron/core/models/gpt/gpt_layer_specs.py ++++ b/megatron/core/models/gpt/gpt_layer_specs.py +@@ -189,6 +189,8 @@ def get_gpt_layer_with_transformer_engine_submodules( + enable_hyper_connection: bool = False, + mla_down_proj_fusion: bool = False, + dense_grouped_gemm: bool = False, ++ post_self_attn_layernorm: bool = False, ++ post_mlp_layernorm: bool = False, + ) -> TransformerLayerSubmodules: + """Use these submodules to use lower-level Transformer Engine modules (required for fp8 + training). +@@ -282,9 +284,11 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + ), + self_attn_bda=get_bias_dropout_add, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map=( + { + "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.", +@@ -314,10 +318,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + ) + else: + qk_norm = backend.layer_norm(for_qk=True) +@@ -339,10 +345,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map={ + "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight", + "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias", +diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py +index 19de0ed..52156a2 100644 +--- a/megatron/core/models/gpt/gpt_model.py ++++ b/megatron/core/models/gpt/gpt_model.py +@@ -31,6 +31,7 @@ from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + mtp_on_this_rank, + process_mtp_loss, ++ roll_tensor, + ) + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_block import TransformerBlock +@@ -506,6 +507,7 @@ class GPTModel(LanguageModule): + loss_mask: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + is_spec_decode: Optional[bool] = None, ++ mtp_kwargs: Optional[dict] = {}, + ) -> Tensor: + """Forward function of the GPT Model This function passes the input tensors + through the embedding layer, and then the decoder and finally into the post +@@ -585,6 +587,7 @@ class GPTModel(LanguageModule): + extra_block_kwargs=extra_block_kwargs, + inference_context=inference_context, + is_spec_decode=is_spec_decode, ++ mtp_kwargs=mtp_kwargs, + ) + + def _postprocess( +@@ -607,6 +610,7 @@ class GPTModel(LanguageModule): + extra_block_kwargs=None, + inference_context=None, + is_spec_decode=None, ++ mtp_kwargs={}, + ): + """Postprocesses decoder hidden states to generate logits or compute loss. + +@@ -631,7 +635,27 @@ class GPTModel(LanguageModule): + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() +- if mtp_in_postprocess and not (in_inference_mode or is_spec_decode): ++ if output_weight is not None: ++ output_weight = output_weight.detach() ++ mtp_labels = labels ++ if mtp_kwargs is not None and mtp_kwargs.get("mtp_labels", None) is not None: ++ mtp_labels = mtp_kwargs["mtp_labels"] ++ mtp_labels, _ = roll_tensor( ++ mtp_labels, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ if loss_mask is not None: ++ loss_mask, _ = roll_tensor( ++ loss_mask, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ if mtp_in_postprocess and not (in_inference_mode or is_spec_decode) and mtp_labels is not None: + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, +@@ -660,7 +684,7 @@ class GPTModel(LanguageModule): + # In training/eval, use the utility function for processing MTP loss/scaling. + hidden_states = process_mtp_loss( + hidden_states=hidden_states, +- labels=labels, ++ labels=mtp_labels, + loss_mask=loss_mask, + output_layer=self.output_layer, + output_weight=output_weight, +diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +index c87ccd5..8572db4 100644 +--- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py ++++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +@@ -249,6 +249,11 @@ class HybridDeviceOptimizer(torch.optim.Optimizer): + return cpu_optimizers + + def _get_sub_optimizer_param_groups(self, offload_fraction: float): ++ import warnings ++ warnings.warn( ++ "CPU offload optimizer init can be very slow (potentially minutes) for " ++ "large MoE models due to per-parameter pinned-memory allocation and H2D copies." ++ ) + params = [] + for group in self.param_groups: + params.extend(group["params"]) +diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py +index 4430a8c..084389b 100644 +--- a/megatron/core/optimizer/distrib_optimizer.py ++++ b/megatron/core/optimizer/distrib_optimizer.py +@@ -706,6 +706,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # TE FusedAdam will not accumulate step for empty param groups, so we need to + # align the step across param groups. + param_group["step"] = int(step) ++ if "step" in param_group and param_group["step"] is None: ++ del param_group["step"] + + # Grad scaler state. + if self.grad_scaler: +@@ -1771,6 +1773,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # separately via param_groups, not as part of the gradient buffer. + tensors[key] = LocalNonpersistentObject(tensors[key]) + continue ++ if key == 'step': ++ continue + assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), ( + tensors[key].shape, + gbuf_local_start, +diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py +index 322f12a..0be9e5e 100644 +--- a/megatron/core/packed_seq_params.py ++++ b/megatron/core/packed_seq_params.py +@@ -64,3 +64,20 @@ class PackedSeqParams: + .to(torch.int32) + .unsqueeze(0) # Add a batch dimension + ) ++ ++ ++def resolve_cp_group( ++ static_cp_group: dist.ProcessGroup, packed_seq_params: PackedSeqParams = None ++) -> dist.ProcessGroup: ++ """Return the dynamic CP group from packed_seq_params when available, else the static one. ++ ++ Dynamic CP assigns a per-microbatch CP group that may differ from the ++ process-group stored at model construction time. This helper centralises ++ the resolution logic so every consumer (layout conversion, GDN, loss and ++ recompute paths) derives size/rank from exactly the same group object. ++ ++ Backported from Megatron-LM 5139086e. ++ """ ++ if packed_seq_params is not None and packed_seq_params.cp_group is not None: ++ return packed_seq_params.cp_group ++ return static_cp_group +diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py +index 863b5d5..6c81e6e 100644 +--- a/megatron/core/parallel_state.py ++++ b/megatron/core/parallel_state.py +@@ -11,6 +11,7 @@ from typing import Callable, List, Optional + + import numpy as np + import torch ++import torch.distributed as dist + + from megatron.core.inference.symmetric_memory import SymmetricMemoryManager + +diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py +index 465e83f..232caef 100644 +--- a/megatron/core/pipeline_parallel/p2p_communication.py ++++ b/megatron/core/pipeline_parallel/p2p_communication.py +@@ -232,34 +232,26 @@ class P2PCommunicator: + group=self.pp_group, + ) + else: +- ops = [] +- if send_prev_shape_tensor is not None: +- send_prev_op = torch.distributed.P2POp( +- torch.distributed.isend, send_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(send_prev_op) +- if recv_prev_shape_tensor is not None: +- recv_prev_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(recv_prev_op) +- if send_next_shape_tensor is not None: +- send_next_op = torch.distributed.P2POp( +- torch.distributed.isend, send_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(send_next_op) +- if recv_next_shape_tensor is not None: +- recv_next_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(recv_next_op) +- if len(ops) > 0: +- reqs = torch.distributed.batch_isend_irecv(ops) +- for req in reqs: +- req.wait() ++ # PR #5271 (Megatron-LM): shape exchange MUST use _p2p_ops rather than ++ # batch_isend_irecv. batch_isend_irecv is one tagless NCCL group; when ++ # pp_group.size()==2, prev_rank == next_rank (single physical peer) and ++ # same-peer ops pair FIFO by enqueue order. Both ranks build ops in the ++ # same fixed order but hold opposite prev/next roles → recv_prev_shape ++ # and recv_next_shape get silently crossed. _p2p_ops handles size==2 ++ # via the group.WORLD split + even/odd ordering, correct for size>=4 too. ++ reqs = _p2p_ops( ++ tensor_send_prev=send_prev_shape_tensor, ++ tensor_recv_prev=recv_prev_shape_tensor, ++ tensor_send_next=send_next_shape_tensor, ++ tensor_recv_next=recv_next_shape_tensor, ++ group=self.pp_group, ++ prev_pipeline_rank=self.prev_rank, ++ next_pipeline_rank=self.next_rank, ++ ) ++ for req in reqs.values(): ++ req.wait() + +- # To protect against race condition when using batch_isend_irecv(). +- # should take this out once the bug with batch_isend_irecv is resolved. ++ # keep the CUDA sync as a defensive measure — cheap for 3-int64 tensors. + torch.cuda.synchronize() + + recv_prev_shape = [0, 0, 0] +@@ -371,6 +363,11 @@ class P2PCommunicator: + return [] + + p2p_func = _ring_exchange_wrapper ++ elif self.pp_group.size() == 2: ++ # PR #5271 (Megatron-LM): size==2 has same-peer (prev_rank == next_rank); ++ # batch_isend_irecv cannot pair same-peer bidirectional ops correctly (see ++ # #1450). Force _p2p_ops which handles this via WORLD split + even/odd order. ++ p2p_func = _p2p_ops + elif config.batch_p2p_comm: + assert wait_on_reqs + p2p_func = _batched_p2p_ops +@@ -381,10 +378,12 @@ class P2PCommunicator: + next_rank = self.next_rank + prev_rank = self.prev_rank + +- if config.use_ring_exchange_p2p or config.batch_p2p_comm: +- reqs = [] +- else: ++ # reqs init must match p2p_func return type: _p2p_ops returns dict, ++ # _batched_p2p_ops + _ring_exchange_wrapper return list. ++ if p2p_func is _p2p_ops: + reqs = {} ++ else: ++ reqs = [] + + tensor_recv_prev = None + tensor_recv_next = None +diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py +index 8df4df1..f39eff8 100644 +--- a/megatron/core/ssm/gated_delta_net.py ++++ b/megatron/core/ssm/gated_delta_net.py +@@ -14,12 +14,16 @@ import torch.nn as nn + import torch.nn.functional as F + from torch import Tensor + ++from megatron.core.context_parallel_layout import ( ++ contiguous_to_zigzag_chunks, ++ zigzag_to_contiguous_chunks, ++) + from megatron.core.dist_checkpointing import ShardedTensor + from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory + from megatron.core.fp8_utils import get_fp8_align_size + from megatron.core.inference.contexts import BaseInferenceContext + from megatron.core.jit import jit_fuser +-from megatron.core.packed_seq_params import PackedSeqParams ++from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.ssm.mamba_context_parallel import ( + _all_to_all_cp2hp, +@@ -52,6 +56,17 @@ except ImportError: + + HAVE_FLA = False + ++try: ++ # FLA >= 0.4.2. Only linear_cp_mode="chunkwise" needs it; every other GDN path ++ # works against 0.4.1, so a missing fla.ops.cp disables chunkwise and nothing else. ++ from fla.ops.cp import build_cp_context ++ ++ HAVE_FLA_CP = True ++except ImportError: ++ build_cp_context = None ++ ++ HAVE_FLA_CP = False ++ + + logger = logging.getLogger(__name__) + +@@ -233,6 +248,15 @@ class GatedDeltaNet(MegatronModule): + tp_group=self.pg_collection.tp, + ) + ++ # Cache for chunkwise-CP contexts consumed by FLA kernels. Rebuilding these ++ # per-forward is unsafe under CUDA graph capture because build_cp_context ++ # allocates fresh tensors whose pointers get baked into the captured graph. ++ # For non-packed (SBHD) input the cu_seqlens is fully determined by the ++ # (static) global sequence length and batch size, so the context is cached on ++ # both values. Packed (THD) input rebuilds per micro-batch because the packed ++ # boundaries change every step. ++ self._chunkwise_cp_context_cache = {} ++ + self.reset_parameters() + + def reset_parameters(self): +@@ -257,6 +281,62 @@ class GatedDeltaNet(MegatronModule): + ).uniform_(*self.A_init_range) + self.A_log.data.copy_(torch.log(A)) + ++ def _resolve_cp_routing(self, packed_seq_params, pg_collection): ++ """Split the resolved CP group into a headwise and a chunkwise group. ++ ++ The CP algorithm itself comes from ``self.config.linear_cp_mode``, which is ++ static for the whole process: dynamic context parallelism varies only the CP ++ group and its size, never the algorithm. Nothing here writes to ``self`` or to ++ the shared config, and no process group is ever created inside a forward. ++ ++ The two algorithms are mutually exclusive: whichever one runs owns the whole ++ CP group and the other gets ``None``, treated as size 1 everywhere below. ++ ++ Returns ``(headwise_group, headwise_size, chunkwise_group, chunkwise_size)``. ++ """ ++ base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp ++ cp_group = resolve_cp_group(base_cp_group, packed_seq_params) ++ cp_size = cp_group.size() if cp_group is not None else 1 ++ ++ # Dynamic CP must keep local_cp_size and the group it selected consistent; ++ # a mismatch means some collective would be issued on the wrong group. ++ if packed_seq_params is not None: ++ local_cp_size = getattr(packed_seq_params, "local_cp_size", None) ++ if local_cp_size is not None and local_cp_size != cp_size: ++ raise ValueError( ++ f"PackedSeqParams.local_cp_size ({local_cp_size}) does not match the " ++ f"resolved CP group size ({cp_size}). Dynamic CP must set both from the " ++ f"same group." ++ ) ++ ++ if cp_size == 1: ++ # No CP communication either way; keep the original (size-1) group so the ++ # existing helpers behave exactly as before. Checked before the mode so a ++ # CP=1 micro-batch never trips the mode validation below. ++ return cp_group, 1, None, 1 ++ ++ mode = self.config.linear_cp_mode ++ if mode == "headwise": ++ return cp_group, cp_size, None, 1 ++ if mode == "chunkwise": ++ if not HAVE_FLA_CP: ++ raise RuntimeError( ++ "linear_cp_mode='chunkwise' requires fla.ops.cp (flash-linear-attention " ++ ">= 0.4.2), which is not importable in this environment." ++ ) ++ return None, 1, cp_group, cp_size ++ if mode == "all_gather": ++ raise RuntimeError( ++ "linear_cp_mode='all_gather' is implemented by the Relax GatedDeltaNet " ++ "wrapper, not by Megatron. Reaching this forward with cp_size=" ++ f"{cp_size} means the wrapper was not installed." ++ ) ++ raise ValueError( ++ f"Unsupported linear_cp_mode {mode!r}; expected 'headwise', 'chunkwise' or " ++ "'all_gather'. 'auto' must be resolved to a concrete mode before the model is " ++ "constructed." ++ ) ++ + def forward( + self, + hidden_states: Tensor, +@@ -265,6 +345,7 @@ class GatedDeltaNet(MegatronModule): + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, ++ pg_collection: Optional[ProcessGroupCollection] = None, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ): +@@ -279,6 +360,8 @@ class GatedDeltaNet(MegatronModule): + packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format. + sequence_len_offset (Optional[int]): Sequence length offset used for + inference CUDA graphs. ++ pg_collection (Optional[ProcessGroupCollection]): Overrides the process ++ groups captured at construction time. + + Return: + (Tuple[Tensor, Tensor]) GDN output and bias. +@@ -288,8 +371,21 @@ class GatedDeltaNet(MegatronModule): + + inference_context = deprecate_inference_params(inference_context, inference_params) + ++ ( ++ cp_group_headwise, ++ cp_size_headwise, ++ cp_group_chunkwise, ++ cp_size_chunkwise, ++ ) = self._resolve_cp_routing(packed_seq_params, pg_collection) ++ + seq_len, batch, _ = hidden_states.shape +- seq_len = seq_len * self.sp_size * self.cp_size ++ # Sequence length seen by the conv / scan kernels: headwise CP turns the ++ # sequence split into a head split and so restores the full sequence; ++ # chunkwise CP keeps this rank's time slice. ++ seq_len_post_headwise = seq_len * self.sp_size * cp_size_headwise ++ # Sequence length described by the global (pre-CP-split) cu_seqlens. ++ seq_len_global = seq_len_post_headwise * cp_size_chunkwise ++ seq_len = seq_len_post_headwise + + if inference_context is not None: + assert ( +@@ -299,24 +395,30 @@ class GatedDeltaNet(MegatronModule): + # TODO: support inference + raise NotImplementedError("GDN does not support inference for now.") + +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': ++ is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' ++ if is_thd: + assert batch == 1, "Packed sequence expects batch dimension to be 1" + assert ( + not self.config.deterministic_mode + ), "Packed sequence does not support deterministic mode." + +- # Resolve cu_seqlens with alignment padding handling. ++ # Resolve cu_seqlens with alignment padding handling. These are the *global* ++ # (pre-CP-split) boundaries in both CP modes, taken from the PackedSeqParams ++ # this forward was handed -- for the Bridge/VLM unsplit path that is the ++ # post-embedding, post-repack object, so no stale boundary can leak in. + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, +- seq_len, ++ seq_len_global, + "cu_seqlens_q", ++ cp_size=cp_size_chunkwise, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, +- seq_len, ++ seq_len_global, + "cu_seqlens_kv", ++ cp_size=cp_size_chunkwise, + ) + assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( + "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " +@@ -331,21 +433,59 @@ class GatedDeltaNet(MegatronModule): + cu_seqlens_q = None + cu_seqlens_kv = None + ++ if cp_size_chunkwise > 1: ++ # Chunkwise CP runs through the FLA CP kernels only; the torch-native ++ # reference scan used by deterministic mode has no cp_context path. ++ if self.config.deterministic_mode: ++ raise ValueError( ++ "GDN chunkwise CP does not support deterministic_mode: the torch-native " ++ "reference scan has no CP context. Use linear_cp_mode='headwise'." ++ ) ++ if not is_thd and batch > 1: ++ raise ValueError( ++ "GDN chunkwise CP with SBHD inputs requires micro_batch_size == 1 because " ++ "the FLA gated delta rule backend needs a single batch dimension when " ++ "cp_context is used. Use packed THD input or micro_batch_size=1." ++ ) ++ ++ # Build the chunkwise CP context once per forward. It carries this rank's *local* ++ # cu_seqlens plus the neighbour bookkeeping the FLA kernels need in order to ++ # exchange conv boundary tokens and chunk-boundary states. ++ chunkwise_cp_context = self._build_chunkwise_cp_context( ++ cp_group_chunkwise, cp_size_chunkwise, cu_seqlens_q, seq_len_global, batch ++ ) ++ + # Input projection + nvtx_range_push(suffix="in_proj") + qkvzba, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + ++ # Chunkwise CP expects the contiguous-time layout (rank r owns global chunks ++ # [2r, 2r+1]) inside conv1d / chunk_gated_delta_rule, because that is the ++ # partition fla.ops.cp.build_cp_context assumes. Megatron and Relax hand us the ++ # zigzag attention-load-balanced layout (rank r owns [r, 2*cp-r-1]), so reshuffle ++ # over the CP group with a single all-to-all. No full-sequence gather, and the ++ # tensor stays sharded 1/cp throughout. ++ if cp_size_chunkwise > 1: ++ nvtx_range_push(suffix="zigzag_to_contiguous") ++ qkvzba = zigzag_to_contiguous_chunks( ++ qkvzba, ++ cp_group_chunkwise, ++ seq_dim=0, ++ cu_seqlens=cu_seqlens_q if is_thd else None, ++ ) ++ nvtx_range_pop(suffix="zigzag_to_contiguous") ++ + # CP All to All: CP to HP +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': +- unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // self.cp_size, dim=0) ++ if cp_size_headwise > 1 and is_thd: ++ unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // cp_size_headwise, dim=0) + outputs = [] + for qkvzba_i in unpacked_qkvzba: + qkvzba_i = tensor_a2a_cp2hp( + qkvzba_i, + seq_dim=0, + head_dim=-1, +- cp_group=self.pg_collection.cp, ++ cp_group=cp_group_headwise, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, +@@ -357,12 +497,12 @@ class GatedDeltaNet(MegatronModule): + ) + outputs.append(qkvzba_i) + qkvzba = torch.cat(outputs, dim=0) +- else: ++ elif cp_size_headwise > 1: + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, +- cp_group=self.pg_collection.cp, ++ cp_group=cp_group_headwise, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, +@@ -381,10 +521,10 @@ class GatedDeltaNet(MegatronModule): + qkv, gate, beta, alpha = torch.split( + qkvzba, + [ +- (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // self.cp_size, +- self.v_dim_local_tp // self.cp_size, +- self.num_value_heads // self.tp_size // self.cp_size, +- self.num_value_heads // self.tp_size // self.cp_size, ++ (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, ++ self.v_dim_local_tp // cp_size_headwise, ++ self.num_value_heads // self.tp_size // cp_size_headwise, ++ self.num_value_heads // self.tp_size // cp_size_headwise, + ], + dim=-1, + ) +@@ -403,14 +543,14 @@ class GatedDeltaNet(MegatronModule): + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, +- cp_group=self.pg_collection.cp, ++ cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, +- cp_group=self.pg_collection.cp, ++ cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias +@@ -425,7 +565,7 @@ class GatedDeltaNet(MegatronModule): + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, +- groups=self.conv_dim_local_tp // self.cp_size, ++ groups=self.conv_dim_local_tp // cp_size_headwise, + ) + qkv = self.act_fn(conv_out[..., :seq_len]) + qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d +@@ -439,22 +579,24 @@ class GatedDeltaNet(MegatronModule): + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens_q, ++ # Only pass cp_context when chunkwise is active, so with it off this call ++ # is byte-identical to before the backport (and still works on FLA 0.4.1). ++ **({} if chunkwise_cp_context is None else {"cp_context": chunkwise_cp_context}), + ) + nvtx_range_pop(suffix="conv1d") + + # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) + nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") +- query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( +- qkv, gate, beta, alpha, batch, seq_len +- ) ++ with torch._dynamo.config.patch(disable=True): ++ query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( ++ qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise=cp_size_headwise ++ ) + nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") + + # Calculate g and beta + nvtx_range_push(suffix="g_and_beta") +- A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) +- dt_bias_local_cp = get_parameter_local_cp( +- self.dt_bias, dim=0, cp_group=self.pg_collection.cp +- ) ++ A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) ++ dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) + nvtx_range_pop(suffix="g_and_beta") + +@@ -469,6 +611,7 @@ class GatedDeltaNet(MegatronModule): + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, ++ **({} if chunkwise_cp_context is None else {"cp_context": chunkwise_cp_context}), + ) + nvtx_range_pop(suffix="gated_delta_rule") + +@@ -482,19 +625,31 @@ class GatedDeltaNet(MegatronModule): + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + ++ # Inverse of the zigzag -> contiguous reshuffle done before conv1d, so the layers ++ # after GDN and the loss see the layout they expect. ++ if cp_size_chunkwise > 1: ++ nvtx_range_push(suffix="contiguous_to_zigzag") ++ norm_out = contiguous_to_zigzag_chunks( ++ norm_out, ++ cp_group=cp_group_chunkwise, ++ seq_dim=0, ++ cu_seqlens=cu_seqlens_q if is_thd else None, ++ ) ++ nvtx_range_pop(suffix="contiguous_to_zigzag") ++ + # CP all to all: HP to CP +- if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': ++ if cp_size_headwise > 1 and is_thd: + unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( +- norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp ++ norm_out_i, seq_dim=0, head_dim=-1, cp_group=cp_group_headwise + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) +- else: ++ elif cp_size_headwise > 1: + norm_out = tensor_a2a_hp2cp( +- norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp ++ norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group_headwise + ) + + # Output projection +@@ -504,6 +659,45 @@ class GatedDeltaNet(MegatronModule): + + return out, out_bias + ++ def _build_chunkwise_cp_context( ++ self, cp_group_chunkwise, cp_size_chunkwise, cu_seqlens_q, seq_len_global, batch ++ ): ++ """Build (or fetch from cache) the FLA chunkwise-CP context for this forward. ++ ++ ``build_cp_context`` takes the *global* cu_seqlens plus the CP group and derives, ++ for this rank: its local cu_seqlens, whether it holds the first / last chunk of ++ each sequence, how many ranks it must receive state from, and how many conv ++ boundary tokens to pull from the previous rank. That is the entire cross-rank ++ contract of chunkwise CP. ++ """ ++ if cp_size_chunkwise <= 1: ++ return None ++ ++ if cu_seqlens_q is not None: ++ return build_cp_context( ++ cu_seqlens=cu_seqlens_q, ++ group=cp_group_chunkwise, ++ conv1d_kernel_size=self.conv_kernel_dim, ++ ) ++ ++ # Non-packed (SBHD) input: cu_seqlens is fully determined by the static global ++ # sequence length and batch size, so cache the context. Rebuilding it every ++ # forward allocates fresh tensors, which breaks CUDA graph replay. ++ cache_key = (seq_len_global, batch) ++ cached = self._chunkwise_cp_context_cache.get(cache_key) ++ if cached is None: ++ cached_cu_seqlens = ( ++ torch.arange(batch + 1, device=torch.cuda.current_device(), dtype=torch.long) ++ * seq_len_global ++ ) ++ cached = build_cp_context( ++ cu_seqlens=cached_cu_seqlens, ++ group=cp_group_chunkwise, ++ conv1d_kernel_size=self.conv_kernel_dim, ++ ) ++ self._chunkwise_cp_context_cache[cache_key] = cached ++ return cached ++ + @jit_fuser + def _apply_gated_norm(self, x, gate): + # Output Norm +@@ -517,15 +711,23 @@ class GatedDeltaNet(MegatronModule): + return y + + @jit_fuser +- def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): ++ def _prepare_qkv_for_gated_delta_rule( ++ self, qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise: Optional[int] = None ++ ): + """ + Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. + Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. ++ ++ ``cp_size_headwise`` is how many CP ranks the heads were split across; it is 1 for ++ chunkwise CP, where all TP-local heads stay on the rank. Defaults to ++ ``self.cp_size`` so pre-existing callers keep working unchanged. + """ ++ if cp_size_headwise is None: ++ cp_size_headwise = self.cp_size + # Split qkv into query_key and value + query_key, value = torch.split( + qkv, +- [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], ++ [2 * self.qk_dim_local_tp // cp_size_headwise, self.v_dim_local_tp // cp_size_headwise], + dim=-1, + ) + +@@ -538,7 +740,7 @@ class GatedDeltaNet(MegatronModule): + query_key = l2norm(query_key.contiguous()) + + # Split query and key +- split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size ++ split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size_headwise + query, key = torch.split(query_key, [split_size, split_size], dim=2) + + # Expand query and key if needed (grouped query attention) +@@ -567,7 +769,9 @@ class GatedDeltaNet(MegatronModule): + beta = beta.sigmoid() + return g, beta + +- def _resolve_cu_seqlens(self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name): ++ def _resolve_cu_seqlens( ++ self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1 ++ ): + """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" + if cu_seqlens_padded is not None: + cu_seqlens = cu_seqlens_padded +@@ -582,6 +786,18 @@ class GatedDeltaNet(MegatronModule): + f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." + ) + ++ if cp_size > 1: ++ # Chunkwise CP slices the packed buffer per sequence; a sequence whose padded ++ # length is not a multiple of the CP size cannot be split evenly. (The zigzag ++ # layout needs 2*cp; that stricter check lives in the layout conversion, which ++ # reports the offending lengths.) ++ seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] ++ if (seq_lengths % cp_size != 0).any(): ++ raise ValueError( ++ f"GDN: all per-sequence lengths in {name} must be divisible by " ++ f"cp_size={cp_size}, but got lengths: {seq_lengths.tolist()}" ++ ) ++ + return cu_seqlens + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): +@@ -780,8 +996,10 @@ def get_parameter_local_cp( + torch.Tensor: The local parameter for the current context parallel rank. + """ + +- cp_size = cp_group.size() +- cp_rank = cp_group.rank() ++ # A ``None`` group means "this CP algorithm is not active on this forward" ++ # (chunkwise CP keeps every TP-local head), i.e. behave as size 1. ++ cp_size = cp_group.size() if cp_group is not None else 1 ++ cp_rank = cp_group.rank() if cp_group is not None else 0 + + # No need to split if CP size is 1. + if cp_size == 1: +@@ -800,7 +1018,7 @@ def get_parameter_local_cp( + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) +- param = param[slices] ++ param = param[tuple(slices)] + return param + + +diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py +index 610700f..1ef4fee 100644 +--- a/megatron/core/tensor_parallel/layers.py ++++ b/megatron/core/tensor_parallel/layers.py +@@ -358,7 +358,14 @@ class LinearWithFrozenWeight(torch.autograd.Function): + def backward(ctx, grad_output): + """Backward with frozen weight.""" + (weight,) = ctx.saved_tensors +- grad_input = grad_output.matmul(weight) ++ if grad_output.dim() > 2: ++ # Work around PyTorch matmul not folding some size-1 leading dims to mm. ++ # Remove this once https://github.com/pytorch/pytorch/issues/186148 is fixed. ++ grad_output_2d = grad_output.reshape(-1, grad_output.size(-1)) ++ grad_input = grad_output_2d.matmul(weight) ++ grad_input = grad_input.reshape(*grad_output.shape[:-1], weight.size(1)) ++ else: ++ grad_input = grad_output.matmul(weight) + + if ctx.allreduce_dgrad: + # All-reduce. Note: here async and sync are effectively the same. +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index d316d23..9bb6d2b 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -787,6 +787,9 @@ def topk_routing_with_score_function( + scores, topk, num_groups, group_topk, _compute_topk + ) + ++ from relax.utils.training.routing_replay import get_routing_replay_compute_topk ++ compute_topk = get_routing_replay_compute_topk(compute_topk) ++ + # Precision notes: + # - Logits are converted to fp32 for score functions. + # - All the intermediate calculations are in fp32. +diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py +index b675d33..0cf3e00 100644 +--- a/megatron/core/transformer/moe/router.py ++++ b/megatron/core/transformer/moe/router.py +@@ -216,6 +216,9 @@ class TopKRouter(Router): + if self.config.moe_enable_routing_replay: + self.router_replay = RouterReplay() + ++ from relax.utils.training.routing_replay import register_routing_replay ++ register_routing_replay(self) ++ + def _maintain_float32_expert_bias(self): + """ + Maintain the expert bias in float32. +diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py +index ba5018a..a900a7d 100755 +--- a/megatron/core/transformer/multi_token_prediction.py ++++ b/megatron/core/transformer/multi_token_prediction.py +@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union + + import torch + from torch import Tensor ++import warnings + + from megatron.core import InferenceParams, parallel_state, tensor_parallel + from megatron.core.dist_checkpointing.mapping import ShardedStateDict +@@ -609,6 +610,31 @@ class MTPLossAutoScaler(torch.autograd.Function): + MTPLossAutoScaler.main_loss_backward_scale = scale + + ++def _call_output_layer_for_mtp_loss( ++ output_layer: Callable, ++ hidden_states: Tensor, ++ output_weight: Optional[Tensor], ++ runtime_gather_output: Optional[bool], ++ **kwargs, ++): ++ """Call the output layer without letting MTP loss update lm_head parameters.""" ++ output_layer_params = { ++ name: param.detach() for name, param in output_layer.named_parameters() ++ } ++ output_layer_buffers = dict(output_layer.named_buffers()) ++ detached_output_weight = output_weight.detach() if output_weight is not None else None ++ return torch.func.functional_call( ++ output_layer, ++ {**output_layer_params, **output_layer_buffers}, ++ (hidden_states,), ++ { ++ "weight": detached_output_weight, ++ "runtime_gather_output": runtime_gather_output, ++ **kwargs, ++ }, ++ ) ++ ++ + def process_mtp_loss( + hidden_states: Tensor, + labels: Tensor, +@@ -646,12 +672,12 @@ def process_mtp_loss( + Returns: + Tensor: Updated hidden states after MTP loss processing (first chunk only). + """ +- hidden_states_list = torch.chunk(hidden_states, 1 + config.mtp_num_layers, dim=0) +- hidden_states = hidden_states_list[0] +- + if labels is None: + return hidden_states + ++ hidden_states_list = torch.chunk(hidden_states, 1 + config.mtp_num_layers, dim=0) ++ hidden_states = hidden_states_list[0] ++ + mtp_labels = labels.clone() + if loss_mask is None: + loss_mask = torch.ones_like(mtp_labels) +@@ -672,27 +698,30 @@ def process_mtp_loss( + loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params + ) + if fuse_linear_cross_entropy: +- mtp_loss = output_layer( ++ mtp_loss = _call_output_layer_for_mtp_loss( ++ output_layer, + hidden_states_list[mtp_layer_number + 1], +- weight=output_weight, ++ output_weight=output_weight, + runtime_gather_output=runtime_gather_output, + output_cross_entropy_loss=True, + labels=mtp_labels, + ) + else: +- mtp_logits, _ = output_layer( ++ mtp_logits, _ = _call_output_layer_for_mtp_loss( ++ output_layer, + hidden_states_list[mtp_layer_number + 1], +- weight=output_weight, ++ output_weight=output_weight, + runtime_gather_output=runtime_gather_output, + ) + if scale_logits_fn is not None: + mtp_logits = scale_logits_fn(mtp_logits) + mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits) + mtp_loss = loss_mask * mtp_loss ++ mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + if is_training: + # Safe divide without sync: mask numerator when num_tokens==0, divide by clamp(min=1) + mtp_loss_for_log = ( +- torch.sum(mtp_loss) * (num_tokens > 0).to(mtp_loss.dtype) ++ mtp_loss_scale * torch.sum(mtp_loss) * (num_tokens > 0).to(mtp_loss.dtype) + ) / num_tokens.clamp(min=1) + MTPLossLoggingHelper.save_loss_to_tracker( + mtp_loss_for_log, +@@ -700,7 +729,6 @@ def process_mtp_loss( + config.mtp_num_layers, + avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True), + ) +- mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + if config.calculate_per_token_loss: + # When calculate_per_token_loss is enabled, finalize_model_grads will + # divide all gradients by total_num_tokens (from main loss). +@@ -891,17 +919,19 @@ class MultiTokenPredictionLayer(MegatronModule): + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) +- position_ids, _ = roll_tensor( +- position_ids, +- shifts=-1, +- dims=-1, +- cp_group=self.cp_group, +- packed_seq_params=packed_seq_params, +- ) ++ if position_ids is not None: ++ position_ids, _ = roll_tensor( ++ position_ids, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.cp_group, ++ packed_seq_params=packed_seq_params, ++ ) + # embedding + decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) ++ decoder_input = decoder_input.detach() + +- hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) ++ hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=False) + + return input_ids, position_ids, decoder_input, hidden_states + +@@ -1059,6 +1089,51 @@ class MultiTokenPredictionLayer(MegatronModule): + return hidden_states + + def _checkpointed_forward(self, forward_func, *args, **kwargs): ++ """Wrap `forward_func` with activation checkpointing while only passing tensors. ++ ++ Non-tensor arguments (e.g., configuration objects, None) are captured via closure so ++ that checkpoint implementations never receive them directly, avoiding save_for_backward ++ issues with non-tensor inputs. ++ """ ++ ++ # TODO(jiajun): Is there any better implementation here? ++ positional_specs = [] ++ kw_specs = [] ++ tensor_args: List[torch.Tensor] = [] ++ ++ for arg in args: ++ if torch.is_tensor(arg): ++ positional_specs.append(('tensor', len(tensor_args))) ++ tensor_args.append(arg) ++ else: ++ positional_specs.append(('const', arg)) ++ ++ for key, value in kwargs.items(): ++ if torch.is_tensor(value): ++ kw_specs.append((key, ('tensor', len(tensor_args)))) ++ tensor_args.append(value) ++ else: ++ kw_specs.append((key, ('const', value))) ++ ++ def run(*flat_tensor_args): ++ rebuilt_args = [] ++ for spec_type, payload in positional_specs: ++ if spec_type == 'tensor': ++ rebuilt_args.append(flat_tensor_args[payload]) ++ else: ++ rebuilt_args.append(payload) ++ ++ rebuilt_kwargs = {} ++ for key, (spec_type, payload) in kw_specs: ++ if spec_type == 'tensor': ++ rebuilt_kwargs[key] = flat_tensor_args[payload] ++ else: ++ rebuilt_kwargs[key] = payload ++ ++ return forward_func(*rebuilt_args, **rebuilt_kwargs) ++ ++ tensor_args_tuple = tuple(tensor_args) ++ + def checkpoint_handler(): + """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`""" + if self.config.fp8: +@@ -1069,12 +1144,11 @@ class MultiTokenPredictionLayer(MegatronModule): + self.config.distribute_saved_activations, + tensor_parallel.random.get_cuda_rng_tracker, + parallel_state.get_tensor_model_parallel_group(), +- *args, +- **kwargs, ++ *tensor_args_tuple, + ) + else: + return tensor_parallel.checkpoint( +- forward_func, self.config.distribute_saved_activations, *args, *kwargs.values() ++ run, self.config.distribute_saved_activations, *tensor_args_tuple + ) + + if self.config.recompute_method == 'uniform': +diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py +index cac634f..bcca610 100644 +--- a/megatron/core/transformer/transformer_config.py ++++ b/megatron/core/transformer/transformer_config.py +@@ -244,6 +244,9 @@ class TransformerConfig(ModelParallelConfig): + attention_output_gate: bool = False + """Whether to apply output gate to the attention layers.""" + ++ post_self_attn_layernorm: bool = False ++ post_mlp_layernorm: bool = False ++ + test_mode: bool = False + """Whether to run real-time tests.""" + +@@ -314,6 +317,32 @@ class TransformerConfig(ModelParallelConfig): + linear_num_value_heads: Optional[int] = 32 + """Number of value and gate heads for the gated delta net.""" + ++ linear_cp_mode: Optional[str] = "headwise" ++ """Context-parallel execution mode for linear-attention layers (e.g. Gated Delta Net). ++ Independent of `cp_comm_type`, which only controls standard attention. ++ ++ This is a *static* setting: it is fixed before the model is constructed and is read by ++ both the construction-time head-divisibility check below and by `GatedDeltaNet.forward`. ++ It must not be rewritten during training -- dynamic context parallelism varies only the ++ CP group and its size, never the CP algorithm. ++ ++ "headwise": Ulysses-style. An all-to-all turns the sequence split into a head split, so ++ each rank scans the *full* sequence for `num_heads / (tp * cp)` heads. Requires ++ `linear_num_{key,value}_heads % (tp * cp) == 0`. ++ "chunkwise": Each rank keeps a contiguous *time* slice and all TP-local heads, and the ++ CP-aware FLA kernels (`causal_conv1d` / `chunk_gated_delta_rule` with an `fla.ops.cp` ++ context) exchange only conv boundary tokens and chunk-boundary state summaries. Requires ++ `linear_num_{key,value}_heads % tp == 0` only. ++ "all_gather": Relax-only. Declares that a Relax wrapper implements CP by gathering the ++ full sequence and running a duplicated scan on every CP rank. Megatron's own ++ `GatedDeltaNet.forward` does not implement it and raises if it is ever reached with ++ cp_size > 1, so a missing wrapper fails fast instead of running the wrong algorithm. ++ Uses the same TP-only head rule as "chunkwise", which is what makes non-divisible ++ geometries constructible. ++ ++ Upstream Megatron-LM defaults this to "chunkwise"; Relax keeps "headwise" so that ++ upgrading the image never silently reroutes an existing recipe.""" ++ + #################### + # initialization + #################### +@@ -1242,15 +1271,24 @@ class TransformerConfig(ModelParallelConfig): + f"linear_num_key_heads ({self.linear_num_key_heads})." + ) + +- # Check tensor parallelism compatibility +- tp_cp_size = self.tensor_model_parallel_size * self.context_parallel_size +- assert self.linear_num_key_heads % tp_cp_size == 0, ( ++ # Check tensor parallelism compatibility. Headwise CP splits linear-attention ++ # heads across CP ranks; chunkwise CP and the Relax all-gather fallback keep all ++ # TP-local heads on every CP rank, so they only need TP divisibility. ++ assert self.linear_cp_mode in ("headwise", "chunkwise", "all_gather"), ( ++ f"linear_cp_mode must be one of 'headwise', 'chunkwise' or 'all_gather', got " ++ f"{self.linear_cp_mode!r}. It is a static setting resolved before model " ++ f"construction; 'auto' must already have been resolved by the caller." ++ ) ++ linear_head_parallel_size = self.tensor_model_parallel_size ++ if self.context_parallel_size > 1 and self.linear_cp_mode == "headwise": ++ linear_head_parallel_size *= self.context_parallel_size ++ assert self.linear_num_key_heads % linear_head_parallel_size == 0, ( + f"{self.linear_num_key_heads=} must be a multiple of " +- f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." ++ f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." + ) +- assert self.linear_num_value_heads % tp_cp_size == 0, ( ++ assert self.linear_num_value_heads % linear_head_parallel_size == 0, ( + f"{self.linear_num_value_heads=} must be a multiple of " +- f"({self.tensor_model_parallel_size=} * {self.context_parallel_size=})." ++ f"{linear_head_parallel_size=} for {self.linear_cp_mode=}." + ) + elif self.experimental_attention_variant == "dsa": + pass +diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py +index ee20545..2ba4d06 100644 +--- a/megatron/core/transformer/transformer_layer.py ++++ b/megatron/core/transformer/transformer_layer.py +@@ -245,6 +245,7 @@ class TransformerLayerSubmodules: + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp + self_attention: Union[ModuleSpec, type] = IdentityOp + self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp + + pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp +@@ -255,6 +256,7 @@ class TransformerLayerSubmodules: + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp + mlp: Union[ModuleSpec, type] = IdentityOp + mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp + + # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method + sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) +@@ -352,6 +354,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + # [Module 3: BiasDropoutFusion] + self.self_attn_bda = build_module(submodules.self_attn_bda) + ++ self.post_self_attn_layernorm = build_module( ++ submodules.post_self_attn_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon, ++ ) ++ + # [Module 4: Post SelfAttention] Optional Layernorm after self-attn + self.pre_cross_attn_layernorm = submodules.pre_cross_attn_layernorm( + config=self.config, +@@ -418,6 +427,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + + self.is_moe_layer = isinstance(self.mlp, MoELayer) + ++ self.post_mlp_layernorm = build_module( ++ submodules.post_mlp_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon ++ ) ++ + self.recompute_input_layernorm = False + self.recompute_pre_mlp_layernorm = False + self.recompute_mlp = False +@@ -638,6 +654,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + attention_output_with_bias[0] + ) + ++ attention_output, attention_output_bias = attention_output_with_bias ++ attention_output = self.post_self_attn_layernorm(attention_output) ++ attention_output_with_bias = (attention_output, attention_output_bias) ++ + # TODO: could we move `bias_dropout_add_exec_handler` itself + # inside the module provided in the `bias_dropout_add_spec` module? + nvtx_range_push(suffix="self_attn_bda") +@@ -823,6 +843,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + self._set_fc2_residual(residual) + mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask) + ++ mlp_output, mlp_output_bias = mlp_output_with_bias ++ mlp_output = self.post_mlp_layernorm(mlp_output) ++ mlp_output_with_bias = (mlp_output, mlp_output_bias) ++ + nvtx_range_pop(suffix="mlp") + + if ( +diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py +index 62f6e44..c25969b 100644 +--- a/megatron/training/arguments.py ++++ b/megatron/training/arguments.py +@@ -1992,6 +1992,9 @@ def core_transformer_config_from_args(args, config_class=None): + + kw_args['inference_sampling_seed'] = args.seed + ++ kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm ++ kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm ++ + # handle quantization config + # NOTE: Kitchen arguments are only added to the namespace when + # Kitchen library is available. +@@ -2475,7 +2478,7 @@ def _add_network_size_args(parser): + '--position-embedding-type', + type=str, + default='learned_absolute', +- choices=['learned_absolute', 'rope', 'mrope', 'relative', 'none'], ++ choices=['learned_absolute', 'rope', 'yarn', 'mrope', 'relative', 'none'], + help='Position embedding type.', + ) + group.add_argument( +diff --git a/megatron/training/training.py b/megatron/training/training.py +index a0817e8..7cd094d 100644 +--- a/megatron/training/training.py ++++ b/megatron/training/training.py +@@ -222,7 +222,9 @@ from megatron.training.utils import ( + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # NOTE(wuhuan): keep the default hook mode; forcing "torch" triggers ++ # 'torch.AcceleratorError: CUDA error: invalid argument' on weight updates. ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False diff --git a/docker/patch/megatron/TASK32-BACKPORT.md b/docker/patch/megatron/TASK32-BACKPORT.md new file mode 100644 index 000000000..87e61213b --- /dev/null +++ b/docker/patch/megatron/TASK32-BACKPORT.md @@ -0,0 +1,134 @@ +# Task 32 v1 — GDN chunkwise context parallel: MCore backport manifest + +This file is the file-level and hunk-level record required by the Task 32 RFC +review: what was taken from upstream, what was deliberately left behind, and why. +It describes the delta between `20260506-85bced0ae.patch` (previous) and +`20260805-85bced0ae.patch` (current). Both patches apply to the same source tree, +so the whole difference below *is* the Task 32 change. + +**v1 depends on NVIDIA/Megatron-LM#3282 only.** NVIDIA/Megatron-LM#5664 is +explicitly out of scope per the RFC review and contributes nothing to this patch. + +## Base + +| Thing | Value | +| ------------------------------------------ | ---------------------------------------------------------------------------------------------- | +| Megatron-Bridge | `2faedbf6fe3c422835a44b2b360cadcb2a116a54` | +| Megatron-LM (`.dev.commit` of that Bridge) | `85bced0ae6ab46f61a0fd774074a3273daf6ae02` | +| Tree assembly | `cp -r Bridge/src/megatron` then `rsync` MCore's `megatron/` over it (see `docker/Dockerfile`) | + +## Upstream sources + +| Ref | State | Used for | +| ------------------------------------------------------------------------------------------ | --------------- | --------------------------------------------------------------------------------------------- | +| [Megatron-LM #3282](https://github.com/NVIDIA/Megatron-LM/pull/3282), merged as `5139086e` | merged to `dev` | The chunkwise CP feature. The only MCore source used by v1. | +| flash-linear-attention / fla-core `0.4.2` | released | `fla.ops.cp.build_cp_context`; `cp_context=` on `causal_conv1d` and `chunk_gated_delta_rule`. | + +`5139086e` is ~597 commits ahead of `85bced0a`, and it edits the same regions of +`gated_delta_net.py` / `transformer_config.py` that Relax's own patch edits, so a +cherry-pick is not possible. Everything below is a selective port onto the pinned +tree. + +## Design constraint from the RFC review + +The GDN CP mode is **static for the whole process**. `linear_cp_mode` is read by +both the construction-time head check and by `GatedDeltaNet.forward`; there is no +per-call override and nothing in a forward writes to `self` or to the shared +config. Dynamic context parallelism varies only `cp_group` / `local_cp_size`, +never the algorithm. + +## Files added / changed + +### `megatron/core/context_parallel_layout.py` — new, +324 + +**Byte-identical to `5139086e` below the module docstring.** Verified with +`diff` against the upstream file; the only change is a Relax provenance note in +the docstring. Contains: + +| Symbol | Purpose | +| ------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------- | +| `get_thd_context_parallel_rank_indices` | Reference description of both partitions. Tests assert Relax's `slice_with_cp` / `gdn_cp_slice` against it. | +| `zigzag_to_contiguous_chunks`, `contiguous_to_zigzag_chunks` | Public entry points; dispatch on whether `cu_seqlens` is given. | +| `_zigzag_contiguous_thd_swap` | Packed THD: one packed-token all-to-all. Builds its routing from `cu_seqlens` on each call. | +| `_zigzag_contiguous_chunk_swap` | SBHD: one chunk-level all-to-all. | + +### `megatron/core/packed_seq_params.py` — +17 + +- `resolve_cp_group()` from `5139086e`. The single place where "dynamic + per-micro-batch CP group, else the construction-time one" is decided, so every + consumer derives size/rank from the same group object. +- `local_cp_size` / `cp_group` already existed on `85bced0a`; not re-backported. +- The dataclass field list is **untouched**. + +### `megatron/core/transformer/transformer_config.py` — +50 / −16 + +- `linear_cp_mode` field and the head-divisibility rule from `5139086e`: headwise + needs `heads % (tp * cp) == 0`, chunkwise needs `heads % tp == 0`. +- **Relax adaptation:** the default is `"headwise"`, not upstream's `"chunkwise"`. + Upgrading the image must not silently change the algorithm an existing recipe + runs. +- **Relax adaptation:** `"all_gather"` is accepted as a third declared value, + using the same TP-only head rule as chunkwise. That is what makes a + non-divisible geometry constructible for Relax's all-gather fallback, and it + means the declared config equals the resolved `--gdn-cp-mode` instead of + declaring one mode while running another. +- Unknown values — including an unresolved `"auto"` — assert at construction. +- **Excluded:** `gdn_conv_pad_alignment`, `gdn_pre_gated_delta_rule_fusion` and + their interaction asserts. Neither field exists on this base. + +### `megatron/core/ssm/gated_delta_net.py` — +290 / −42 + +- `_resolve_cp_routing()`: resolve the CP group once via `resolve_cp_group`, then + give the whole group to exactly one of headwise / chunkwise and `None` to the + other (`None` is treated as size 1 everywhere downstream). No process group is + ever created in a forward. +- Validates `PackedSeqParams.local_cp_size == cp_group.size()`; a mismatch means + a collective would run on the wrong group. +- `cp_size == 1` short-circuits *before* the mode is read, so a CP=1 micro-batch + is legal under any declared mode. +- `linear_cp_mode="all_gather"` raises if MCore's own forward is reached with + `cp_size > 1`: that mode is implemented by the Relax wrapper, so arriving here + means the wrapper was not installed. +- zigzag ↔ contiguous conversion around the conv + scan; `cp_context` built once + per forward and passed to both FLA kernels. +- `_resolve_cu_seqlens` gains the `cp_size` divisibility check from `5139086e`. +- **Relax adaptation — backwards compatibility:** + - `_prepare_qkv_for_gated_delta_rule` takes `cp_size_headwise` as an *optional* + argument defaulting to `self.cp_size`, so Relax's existing all-gather + fallback (`relax/backends/megatron/model.py`) keeps calling it unchanged. + - `get_parameter_local_cp` accepts a `None` group as size 1. + - `cp_context=` is only passed when chunkwise is active, so with chunkwise off + the FLA call is byte-identical to before the backport — and still works + against FLA 0.4.1. +- **Preserved Relax fixes** (both from the previous patch, unchanged): + - `torch._dynamo.config.patch(disable=True)` around + `_prepare_qkv_for_gated_delta_rule` (Qwen3.6 `torch.compile` failure); + - `param[tuple(slices)]` in `get_parameter_local_cp` (multi-dim basic indexing). +- **Excluded from `5139086e`:** the `_forward_compute` split and `recompute_gdn` + selective-recompute wrapper; `gdn_pre_gated_delta_rule_fusion` / + `_fused_streamed_pre_gated_delta_rule`; `gdn_conv_pad_alignment` conv padding; + the `_a2a_cp_to_hp` / `_a2a_hp_to_cp` refactor of the headwise path; the rename + of `get_parameter_local_cp` to `get_parameter_local_cp_headwise`. All are + independent changes from the 597-commit gap and would alter existing paths' + structure or numerics for no Task 32 benefit. + +### Explicitly not in this patch + +| Thing | Why | +| ------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Anything from NVIDIA/Megatron-LM#5664 | Out of scope for v1 per the RFC review. No `cp_partition_mode`, no route tensors, no `prebuild_thd_cp_partition_routes`, no `pad_between_seqs`. The THD swap rebuilds its routing per call, which is `5139086e` behaviour. | +| `megatron/core/extensions/transformer_engine.py` | The THD output-length fix in the RFC was conditional on adopting `pad_between_seqs`. That representation is not backported and no test on this base reproduces the mismatch. | +| `pyproject.toml` / `uv.lock` | Relax installs FLA from `docker/Dockerfile`, not from upstream package metadata. | + +## Verifying this patch is exactly what it claims + +```bash +# byte-identity of the new module against upstream, below the docstring +curl -s https://raw.githubusercontent.com/NVIDIA/Megatron-LM/5139086e/megatron/core/context_parallel_layout.py \ + > /tmp/up.py +diff <(sed -n '/^from typing import/,$p' /megatron/core/context_parallel_layout.py) \ + <(sed -n '/^from typing import/,$p' /tmp/up.py) # must be empty + +# no #5664 content anywhere +grep -rn 'cp_partition_route\|prebuild_thd_cp\|pad_between_seqs' /megatron/ # must be empty +``` diff --git a/tests/backends/megatron/gdn_cp_numeric_probe.py b/tests/backends/megatron/gdn_cp_numeric_probe.py new file mode 100644 index 000000000..aa7f6cee3 --- /dev/null +++ b/tests/backends/megatron/gdn_cp_numeric_probe.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Cross-image numerical probe for the GDN context-parallel paths. + +RFC Task 32 acceptance item 3.4-2: the FLA 0.4.1 -> 0.4.2 upgrade and the MCore +backport must not move the numbers on any path that already existed, i.e. CP=1, +MCore headwise CP, and Relax's all-gather fallback. That cannot be a normal unit +test, because it compares *two images* -- it needs the same tensors produced +under the old dependency set and the new one. + +So this module is a runner, not a pytest file (hence the non-``test_`` name): + + # inside the OLD image + python -m tests.backends.megatron.gdn_cp_numeric_probe dump --mode headwise --out /out/old + # inside the NEW image + python -m tests.backends.megatron.gdn_cp_numeric_probe dump --mode headwise --out /out/new + # anywhere + python -m tests.backends.megatron.gdn_cp_numeric_probe compare --ref /out/old --cand /out/new + +Everything that could drift between images is pinned by hand: every parameter is +overwritten with a tensor drawn from a name-seeded CPU generator, and the inputs +come from a fixed-seed CPU generator too. So a difference in the report is a +difference in the *kernels*, not in initialisation order or RNG plumbing. + +Modes: + cp1 1 GPU, no CP at all. + headwise 2 GPUs, MCore's native cp2hp all-to-all path. + all_gather 2 GPUs, Relax's `_dcp_gdn_forward` fallback (head geometry chosen so + the dispatcher cannot use headwise), which also exercises the FLA + conv/scan kernels on the full gathered sequence. + chunkwise 2 GPUs, the newly backported path (candidate image only). +""" + +from __future__ import annotations + +import argparse +import json +import os +import zlib +from types import SimpleNamespace + +import torch +import torch.multiprocessing as mp + + +MODE_WORLD_SIZE = {"cp1": 1, "headwise": 2, "all_gather": 2, "chunkwise": 2} +# all_gather is only reachable when the heads do NOT divide tp * cp. +MODE_HEADS = { + "cp1": (4, 8), + "headwise": (4, 8), + "chunkwise": (4, 8), + "all_gather": (1, 2), +} +SEQ_LENS = [256, 128] +HIDDEN_SIZE = 512 + + +def _deterministic_fill_(module) -> None: + """Overwrite every parameter from a name-seeded generator. + + Makes the dump independent of Megatron's initialisation code, which is one + of the things the patch touches. + """ + with torch.no_grad(): + for name, p in sorted(module.named_parameters()): + gen = torch.Generator(device="cpu").manual_seed(zlib.crc32(name.encode()) & 0x7FFFFFFF) + values = torch.randn(p.shape, generator=gen, dtype=torch.float32) * 0.05 + p.copy_(values.to(device=p.device, dtype=p.dtype)) + + +def _build(mode, cp_size): + import torch.nn.functional as F + from megatron.core import parallel_state + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, + ) + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.transformer_config import TransformerConfig + + model_parallel_cuda_manual_seed(123) + num_key_heads, num_value_heads = MODE_HEADS[mode] + + if mode == "all_gather": + # Relax relaxes Megatron's headwise `% (tp*cp)` config gate down to `% tp` so a + # non-divisible geometry can even be constructed. v1 keeps this helper working + # untouched, and using it in both images keeps the A/B identical. (The backported + # `linear_cp_mode="all_gather"` is the eventual replacement; it is covered by + # test_gdn_chunkwise_cp_layout.py instead, so it cannot skew this comparison.) + from relax.backends.megatron.model import _relax_gdn_cp_config_assert + + _relax_gdn_cp_config_assert() + + config = TransformerConfig( + hidden_size=HIDDEN_SIZE, + num_layers=1, + num_attention_heads=8, + num_query_groups=2, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=1, + context_parallel_size=cp_size, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_conv_kernel_dim=4, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=num_key_heads, + linear_num_value_heads=num_value_heads, + transformer_impl="transformer_engine", + # The CP algorithm is static config, resolved once at launch. Only chunkwise needs + # to be declared here, and only the candidate image has the field at all -- for + # cp1 / headwise / all_gather both images must run literally the same code, which + # is the whole point of this probe. + **( + {"linear_cp_mode": "chunkwise"} + if mode == "chunkwise" and "linear_cp_mode" in TransformerConfig.__dataclass_fields__ + else {} + ), + ) + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + gdn = ( + GatedDeltaNet( + config, + submodules=get_experimental_attention_variant_module_spec(config=config).submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + ) + _deterministic_fill_(gdn) + return gdn, config + + +def _dump_worker(rank, mode, out_dir): + world_size = MODE_WORLD_SIZE[mode] + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + torch.cuda.set_device(rank) + import torch.distributed as dist + + dist.init_process_group("nccl", rank=rank, world_size=world_size) + from megatron.core import parallel_state + from megatron.core.packed_seq_params import PackedSeqParams + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + cp_size = world_size + gdn, config = _build(mode, cp_size) + + if mode == "all_gather": + # Install Relax's runtime GDN wrapper and satisfy its full-recompute gate. + from relax.backends.megatron import model as relax_model + + relax_model.get_args = lambda: SimpleNamespace(recompute_granularity="full") + relax_model._patch_gdn_for_dynamic_cp() + + device = torch.device("cuda", rank) + total = sum(SEQ_LENS) + cu = torch.tensor([0, SEQ_LENS[0], total], device=device, dtype=torch.int32) + local_total = total // cp_size + + gen = torch.Generator(device="cpu").manual_seed(20260805) + hidden_full = torch.randn(total, 1, HIDDEN_SIZE, generator=gen, dtype=torch.float32) + grad_full = torch.randn(total, 1, HIDDEN_SIZE, generator=gen, dtype=torch.float32) + + if cp_size == 1: + hidden_local = hidden_full + grad_local = grad_full + else: + from relax.backends.megatron.cp_utils import gdn_cp_slice + + cu_list = [0, SEQ_LENS[0], total] + hidden_local = gdn_cp_slice(hidden_full, cu_list, cp_size, rank) + grad_local = gdn_cp_slice(grad_full, cu_list, cp_size, rank) + assert hidden_local.shape[0] == local_total + + h = hidden_local.to(device=device, dtype=torch.bfloat16).clone().requires_grad_(True) + psp = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + max_seqlen_q=max(SEQ_LENS), + max_seqlen_kv=max(SEQ_LENS), + cp_group=parallel_state.get_context_parallel_group(), + local_cp_size=cp_size, + ) + out, _ = gdn(h, None, packed_seq_params=psp) + (out.float() * grad_local.to(device)).sum().backward() + + payload = {"out": out.detach().float().cpu(), "grad_hidden": h.grad.detach().float().cpu()} + for name, p in sorted(gdn.named_parameters()): + payload[f"grad::{name}"] = p.grad.detach().float().cpu() + + os.makedirs(out_dir, exist_ok=True) + torch.save(payload, os.path.join(out_dir, f"{mode}.rank{rank}.pt")) + if rank == 0: + import importlib.metadata as md + + meta = { + "mode": mode, + "world_size": world_size, + "fla": md.version("flash-linear-attention"), + "fla_core": md.version("fla-core"), + "torch": torch.__version__, + "heads": MODE_HEADS[mode], + "seq_lens": SEQ_LENS, + } + with open(os.path.join(out_dir, f"{mode}.meta.json"), "w") as fh: + json.dump(meta, fh, indent=2) + print(json.dumps(meta)) + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _cmp(a, b): + a, b = a.flatten().double(), b.flatten().double() + diff = (a - b).abs() + denom = b.abs().clamp_min(1e-12) + cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item() + rms = ((a - b).square().mean().sqrt() / (b.square().mean().sqrt() + 1e-12)).item() + return { + "max_abs": diff.max().item(), + "max_rel": (diff / denom).max().item(), + "rms_ratio": rms, + "cosine": cos, + "bitwise_equal": bool(torch.equal(a, b)), + } + + +def main(): + ap = argparse.ArgumentParser() + sub = ap.add_subparsers(dest="cmd", required=True) + + d = sub.add_parser("dump") + d.add_argument("--mode", choices=sorted(MODE_WORLD_SIZE), required=True) + d.add_argument("--out", required=True) + d.add_argument("--port", default="29601") + + c = sub.add_parser("compare") + c.add_argument("--ref", required=True) + c.add_argument("--cand", required=True) + c.add_argument("--mode", choices=sorted(MODE_WORLD_SIZE), required=True) + c.add_argument("--atol", type=float, default=2e-4) + c.add_argument("--rtol", type=float, default=2e-3) + c.add_argument("--cos", type=float, default=0.99999) + + args = ap.parse_args() + if args.cmd == "dump": + os.environ["MASTER_PORT"] = args.port + world_size = MODE_WORLD_SIZE[args.mode] + mp.spawn(_dump_worker, args=(args.mode, args.out), nprocs=world_size, join=True) + return + + world_size = MODE_WORLD_SIZE[args.mode] + failures, rows = [], [] + for rank in range(world_size): + ref = torch.load(os.path.join(args.ref, f"{args.mode}.rank{rank}.pt"), weights_only=True) + cand = torch.load(os.path.join(args.cand, f"{args.mode}.rank{rank}.pt"), weights_only=True) + assert set(ref) == set(cand), f"tensor sets differ: {set(ref) ^ set(cand)}" + for key in sorted(ref): + stats = _cmp(cand[key], ref[key]) + rows.append({"rank": rank, "tensor": key, **stats}) + # RFC section 5, "MCore GDN, CP=1: candidate image vs old image". + ok = stats["max_abs"] <= args.atol + args.rtol * abs(ref[key]).max().item() + ok = ok and stats["cosine"] >= args.cos + if not ok: + failures.append(rows[-1]) + + print(json.dumps(rows, indent=2)) + print( + f"\n{args.mode}: {len(rows)} tensors compared, " + f"{sum(r['bitwise_equal'] for r in rows)} bitwise identical, {len(failures)} outside tolerance" + ) + if failures: + print("FAIL") + raise SystemExit(1) + print("PASS") + + +if __name__ == "__main__": + main() diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py b/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py new file mode 100644 index 000000000..64bbb9576 --- /dev/null +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py @@ -0,0 +1,689 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Real-kernel / real-collective tests for the GDN chunkwise-CP backport. + +RFC Task 32 phase-1 acceptance items 3 and 4: + +* a minimal CP=2 chunkwise case must drive the *actual* FLA kernels and match a + CP=1 reference in forward and backward within tolerance; +* the GDN parameter keys and shard dimensions in ``state_dict`` / + ``sharded_state_dict`` must not move, i.e. GDN weights stay TP-only and + checkpoints are unaffected by the CP mode. + +Three layers are covered: + +1. the FLA kernels directly (``causal_conv1d`` / ``chunk_gated_delta_rule`` with + a ``cp_context``); +2. the whole MCore ``GatedDeltaNet`` module in fp32 -- the *algebraic* check. In + fp32 the only difference between CP=1 and CP=2 is float summation order, so + the tolerances can be tight enough to catch a genuinely wrong permutation or + a dropped boundary term; +3. the same module in bf16 -- the *production* check, at the dtype training + actually uses, where the achievable agreement is bounded by the storage + format rather than by the algorithm. + +A headwise CP=2 run is included at every layer as a control: if headwise and +chunkwise both drift the same way, the cause is shared plumbing, not the new +code. + +Run with 2 visible GPUs: + pytest tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py +""" + +from __future__ import annotations + +import os + +import pytest +import torch +import torch.multiprocessing as mp + + +WORLD_SIZE = 2 + +# --- tolerances ------------------------------------------------------------ +# bf16 module output / input gradient: RFC section 5, "MCore GDN, CP>1 vs CP=1". +ATOL_BF16 = 2e-3 +RTOL_BF16 = 1e-2 +MIN_COSINE = 0.9999 +# FLA kernels: RFC section 5, normalised RMS error thresholds. +CONV_RMS_RATIO = 1e-3 +GDN_RMS_RATIO = 2e-3 +# fp32 module run: both CP algorithms must land far below any bf16 threshold. +# 1e-3 is 2x below the bf16 element-wise atol of the RFC gate, i.e. "fp32 must be +# comfortably better than the dtype we actually ship". +RMS_RATIO_FP32 = 1e-3 +# fp32 kernel-level: measured ~1e-7, so 1e-5 is a real gate, not a rubber stamp. +KERNEL_RMS_RATIO_FP32 = 1e-5 +# chunkwise vs headwise. headwise is the already-shipped CP algorithm, so whatever +# CP-vs-no-CP disagreement it shows is the floor this environment imposes (reduced +# precision inside the Triton dots, changed summation order, bf16 storage) rather +# than anything about the algorithm. Requiring chunkwise to be no worse than that +# floor is the assertion that actually means something; a fixed atol on a bf16 +# token-sum gradient mostly measures rounding luck. +CHUNKWISE_VS_HEADWISE_RMS_FACTOR = 4.0 +# ...with a floor, so a headwise value that happens to land at or near zero on a given +# run cannot turn into an impossible budget. 1e-6 is still ~1000x tighter than the fp32 +# absolute gate, so the comparison keeps its teeth. +RMS_FLOOR_FP32 = 1e-6 +# ...applied only where headwise is not bit-exact. headwise hands each rank the +# whole sequence and 1/cp of the heads, so none of its reductions are +# repartitioned and it can land exactly on the CP=1 result; "4x zero" would be a +# budget no correct implementation could meet. Those tensors are covered by the +# absolute gates instead. + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < WORLD_SIZE, + reason=f"requires {WORLD_SIZE} CUDA devices", +) + + +def _has_backport() -> bool: + try: + import megatron.core.context_parallel_layout # noqa: F401 + from fla.ops.cp import build_cp_context # noqa: F401 + except ImportError: + return False + return True + + +needs_backport = pytest.mark.skipif(not _has_backport(), reason="requires patched Megatron-LM + FLA >= 0.4.2") + + +# --------------------------------------------------------------------------- +# comparison helpers (run inside the workers) +# --------------------------------------------------------------------------- +def _prep(name, got, want): + got32 = got.detach().float().flatten() + want32 = want.detach().float().flatten() + assert got32.shape == want32.shape, f"{name}: shape {got32.shape} vs {want32.shape}" + assert torch.isfinite(got32).all(), f"{name}: non-finite values in candidate" + assert torch.isfinite(want32).all(), f"{name}: non-finite values in reference" + return got32, want32 + + +def _stats(got32, want32): + diff = (got32 - want32).abs() + rms = (diff.square().mean().sqrt() / (want32.square().mean().sqrt() + 1e-12)).item() + cos = torch.nn.functional.cosine_similarity(got32, want32, dim=0).item() + return diff, rms, cos + + +def _report_elementwise(name, got, want, atol, rtol): + """Per-token tensors: every element within atol + rtol * |ref|.""" + got32, want32 = _prep(name, got, want) + diff, rms, cos = _stats(got32, want32) + worst = (diff - (atol + rtol * want32.abs())).max().item() + assert worst <= 0, ( + f"{name}: max |diff| {diff.max().item():.3e} exceeds atol({atol:.0e})+rtol({rtol:.0e})*|ref| " + f"by {worst:.3e} (rms {rms:.3e}, cosine {cos:.8f})" + ) + assert cos >= MIN_COSINE, f"{name}: cosine {cos:.8f} < {MIN_COSINE}" + + +def _report_rms(name, got, want, ratio): + """Whole-tensor normalised RMS error -- the metric FLA's own CP tests + use.""" + got32, want32 = _prep(name, got, want) + diff, rms, cos = _stats(got32, want32) + assert rms < ratio, ( + f"{name}: normalised RMS error {rms:.3e} >= {ratio:.1e} (max |diff| {diff.max().item():.3e}, cosine {cos:.8f})" + ) + assert cos >= MIN_COSINE, f"{name}: cosine {cos:.8f} < {MIN_COSINE} (rms {rms:.3e})" + + +def _zigzag_shard(full: torch.Tensor, cu, cp_size: int, cp_rank: int) -> torch.Tensor: + from relax.backends.megatron.cp_utils import gdn_cp_slice + + return gdn_cp_slice(full, cu, cp_size, cp_rank) + + +def _init_dist(rank, world_size): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29531") + torch.cuda.set_device(rank) + import torch.distributed as dist + + dist.init_process_group("nccl", rank=rank, world_size=world_size) + + +# --------------------------------------------------------------------------- +# worker: FLA kernel level +# --------------------------------------------------------------------------- +def _worker_fla_kernels(rank, world_size, dtype_name, _unused): + seq_lens = [256, 128] + dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype_name] + _init_dist(rank, world_size) + import torch.distributed as dist + from fla.modules.convolution import causal_conv1d + from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + device = torch.device("cuda", rank) + cp_group = dist.new_group(list(range(world_size))) + + H, DK, DV, W = 2, 64, 64, 4 + total = sum(seq_lens) + cu = torch.tensor([0] + torch.tensor(seq_lens).cumsum(0).tolist(), device=device, dtype=torch.int32) + part = total // world_size + lo, hi = rank * part, (rank + 1) * part + + gen = torch.Generator(device="cpu").manual_seed(1234) + + def _mk(*shape, dtype=dtype): + return torch.randn(*shape, generator=gen, dtype=torch.float32).to(device=device, dtype=dtype) + + conv_ratio = KERNEL_RMS_RATIO_FP32 if dtype is torch.float32 else CONV_RMS_RATIO + gdn_ratio = KERNEL_RMS_RATIO_FP32 if dtype is torch.float32 else GDN_RMS_RATIO + tag = f"[rank{rank}/{dtype_name}]" + + # ---- causal conv ---- + # weight/bias are leaves here on purpose: their gradients are sums over every + # token, which is exactly the quantity chunkwise CP repartitions. Checking only + # dx would leave that untested. + x = _mk(1, total, H * DK) + w0 = _mk(H * DK, W) + b0 = _mk(H * DK) + conv_grad = _mk(1, total, H * DK) + + x_ref = x.clone().requires_grad_(True) + w_ref = w0.clone().requires_grad_(True) + b_ref = b0.clone().requires_grad_(True) + out_ref, _ = causal_conv1d(x=x_ref, weight=w_ref, bias=b_ref, activation="silu", cu_seqlens=cu) + (out_ref.float() * conv_grad.float()).sum().backward() + + x_cp = x[:, lo:hi].clone().requires_grad_(True) + w_cp = w0.clone().requires_grad_(True) + b_cp = b0.clone().requires_grad_(True) + ctx = build_cp_context(cu_seqlens=cu, group=cp_group, conv1d_kernel_size=W) + out_cp, _ = causal_conv1d(x=x_cp, weight=w_cp, bias=b_cp, activation="silu", cu_seqlens=cu, cp_context=ctx) + _report_rms(f"{tag} conv fwd", out_cp, out_ref[:, lo:hi], conv_ratio) + (out_cp.float() * conv_grad[:, lo:hi].float()).sum().backward() + _report_rms(f"{tag} conv dx", x_cp.grad, x_ref.grad[:, lo:hi], conv_ratio) + + for pname, cp_leaf, ref_leaf in (("dweight", w_cp, w_ref), ("dbias", b_cp, b_ref)): + summed = cp_leaf.grad.detach().float().clone() + dist.all_reduce(summed, group=cp_group) + if dtype is torch.float32: + _report_rms(f"{tag} conv {pname}", summed, ref_leaf.grad, conv_ratio) + else: + # These are 384-token sums landing in bf16. The fp32 parametrisation of + # this very test pins the algebra at ~1e-7; in bf16 the achievable + # agreement is set by the storage format, so assert direction and report + # the size rather than pretend a sub-ULP threshold is meaningful. + got32, want32 = _prep(f"{tag} conv {pname}", summed, ref_leaf.grad) + _, rms, cos = _stats(got32, want32) + assert cos >= MIN_COSINE, f"{tag} conv {pname}: cosine {cos:.8f} (rms {rms:.3e})" + if rank == 0: + print(f" {tag} conv {pname}: rms {rms:.3e} cosine {cos:.10f}") + + # ---- gated delta rule ---- + # Inputs must look like what GatedDeltaNet actually feeds the kernel: + # * q/k are L2-normalised (the module sets use_qk_l2norm=True). Un-normalised + # q/k make the recurrent state diverge over hundreds of steps and the + # reference itself goes to NaN -- that would test nothing. + # * g is a log-domain decay built as -A.exp() * softplus(...), hence <= 0. + q = l2norm(_mk(1, total, H, DK).contiguous()) + k = l2norm(_mk(1, total, H, DK).contiguous()) + v = _mk(1, total, H, DV) + g = -_mk(1, total, H, dtype=torch.float32).abs() * 0.1 + beta = _mk(1, total, H, dtype=torch.float32).sigmoid() + + leaves_ref = [t.detach().clone().requires_grad_(True) for t in (q, k, v)] + o_ref, _ = chunk_gated_delta_rule( + *leaves_ref, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu, + ) + o_grad = _mk(1, total, H, DV) + (o_ref.float() * o_grad.float()).sum().backward() + + leaves_cp = [t.detach()[:, lo:hi].clone().requires_grad_(True) for t in (q, k, v)] + ctx2 = build_cp_context(cu_seqlens=cu, group=cp_group, conv1d_kernel_size=W) + o_cp, _ = chunk_gated_delta_rule( + *leaves_cp, + g=g[:, lo:hi], + beta=beta[:, lo:hi], + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu, + cp_context=ctx2, + ) + _report_rms(f"{tag} gdr fwd", o_cp, o_ref[:, lo:hi], gdn_ratio) + (o_cp.float() * o_grad[:, lo:hi].float()).sum().backward() + for name, a, b in zip("qkv", leaves_cp, leaves_ref): + _report_rms(f"{tag} gdr d{name}", a.grad, b.grad[:, lo:hi], gdn_ratio) + + dist.barrier() + dist.destroy_process_group() + + +# --------------------------------------------------------------------------- +# worker: full MCore GatedDeltaNet +# --------------------------------------------------------------------------- +def _build_gdn(cp_size, linear_cp_mode, dtype, num_key_heads=4, num_value_heads=8): + import torch.nn.functional as F + from megatron.core import parallel_state + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, + ) + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.transformer_config import TransformerConfig + + model_parallel_cuda_manual_seed(123) + config = TransformerConfig( + hidden_size=512, + num_layers=1, + num_attention_heads=8, + num_query_groups=2, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + activation_func=F.silu, + bf16=dtype is torch.bfloat16, + tensor_model_parallel_size=1, + context_parallel_size=cp_size, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_conv_kernel_dim=4, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=num_key_heads, + linear_num_value_heads=num_value_heads, + linear_cp_mode=linear_cp_mode, + transformer_impl="transformer_engine", + ) + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + gdn = GatedDeltaNet( + config, + submodules=get_experimental_attention_variant_module_spec(config=config).submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().to(dtype), config + + +def _run_gdn_once(gdn, hidden, psp, grad_out): + """One forward+backward; returns (out, d_hidden, {param: grad}).""" + gdn.zero_grad(set_to_none=True) + h = hidden.clone().requires_grad_(True) + out, _ = gdn(h, None, packed_seq_params=psp) + (out.float() * grad_out).sum().backward() + grads = {n: p.grad.detach().float().clone() for n, p in gdn.named_parameters()} + return out.detach().clone(), h.grad.detach().clone(), grads + + +def _worker_gdn_module(rank, world_size, dtype_name, _unused): + """CP=1 reference vs CP=N, for BOTH CP algorithms, in one process. + + Running headwise and chunkwise side by side is the point: it turns "is + chunkwise close enough to CP=1" (which needs an absolute threshold, and in + bf16 lands on the noise floor of the storage format) into "is chunkwise as + close to CP=1 as the algorithm we already ship" -- a comparison with no + free parameters to tune. + """ + dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype_name] + + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + from megatron.core.packed_seq_params import PackedSeqParams + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + device = torch.device("cuda", rank) + # The CP algorithm is static config now, so each mode needs its own module. They + # share weights, so the comparison is still like-for-like. + modules = {} + gdn, config = _build_gdn(world_size, "headwise", dtype) + for p in gdn.parameters(): + # Same weights on every rank so the CP=1 reference is rank-independent. + dist.broadcast(p.data, src=0) + modules["headwise"] = gdn + modules["chunkwise"], _ = _build_gdn(world_size, "chunkwise", dtype) + modules["chunkwise"].load_state_dict(gdn.state_dict()) + + cp_group = parallel_state.get_context_parallel_group() + # A per-rank size-1 group gives us the CP=1 reference *inside* the same + # process, driving the very same weights through the very same forward. + solo = [dist.new_group([r]) for r in range(world_size)][rank] + + seq_lens = [256, 128] + total = sum(seq_lens) + cu = torch.tensor([0, seq_lens[0], total], device=device, dtype=torch.int32) + + gen = torch.Generator(device="cpu").manual_seed(7) + hidden_full = torch.randn(total, 1, config.hidden_size, generator=gen, dtype=torch.float32).to( + device=device, dtype=dtype + ) + grad_seed = torch.randn(total, 1, config.hidden_size, generator=gen, dtype=torch.float32).to(device) + + def _psp(group, local_cp_size): + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + max_seqlen_q=max(seq_lens), + max_seqlen_kv=max(seq_lens), + cp_group=group, + local_cp_size=local_cp_size, + ) + + # CP=1 reference. A size-1 group short-circuits before the mode is read, so either + # module gives the same reference; use the headwise one. + out_ref, in_grad_ref, param_grads_ref = _run_gdn_once(gdn, hidden_full, _psp(solo, 1), grad_seed) + ref = { + "out": _zigzag_shard(out_ref, cu, world_size, rank), + "d_hidden": _zigzag_shard(in_grad_ref, cu, world_size, rank), + } + ref.update({f"grad {n}": g for n, g in param_grads_ref.items()}) + + shard = _zigzag_shard(hidden_full, cu, world_size, rank) + grad_shard = _zigzag_shard(grad_seed, cu, world_size, rank) + + metrics = {} + for mode in ("headwise", "chunkwise"): + out_cp, in_grad_cp, param_grads_cp = _run_gdn_once( + modules[mode], shard, _psp(cp_group, world_size), grad_shard + ) + got = {"out": out_cp, "d_hidden": in_grad_cp} + # Each CP rank holds a partial parameter gradient; the total is the CP sum. + for name, g in param_grads_cp.items(): + summed = g.clone() + dist.all_reduce(summed, group=cp_group) + got[f"grad {name}"] = summed + metrics[mode] = {} + for key, value in got.items(): + got32, want32 = _prep(f"[rank{rank}][{mode}/{dtype_name}] {key}", value, ref[key]) + diff, rms, cos = _stats(got32, want32) + metrics[mode][key] = (rms, cos, diff.max().item()) + assert cos >= MIN_COSINE, ( + f"[rank{rank}][{mode}/{dtype_name}] {key}: cosine {cos:.8f} < {MIN_COSINE} (rms {rms:.3e})" + ) + + # RFC section 5 absolute gate, on the per-token tensors, in the dtype the + # RFC specifies it for. Applied to headwise too, so a drift in the shared + # plumbing cannot hide behind the comparative check below. + if dtype is torch.bfloat16: + for key in ("out", "d_hidden"): + _report_elementwise( + f"[rank{rank}][{mode}/{dtype_name}] {key}", got[key], ref[key], ATOL_BF16, RTOL_BF16 + ) + # Parameter gradients are token-sum reductions stored in bf16. Bound them + # by the RFC's own relative tolerance for this comparison row (rtol=1e-2) + # applied to the whole tensor, plus the RFC's cosine floor. What actually + # pins the algebra is the fp32 parametrisation of this same test. + for key, (rms, cos, _) in metrics[mode].items(): + if key in ("out", "d_hidden"): + continue + assert rms < RTOL_BF16, ( + f"[rank{rank}][{mode}/bf16] {key}: relative RMS {rms:.3e} >= {RTOL_BF16:.0e} (cosine {cos:.8f})" + ) + else: + for key, (rms, _, _) in metrics[mode].items(): + assert rms < RMS_RATIO_FP32, f"[rank{rank}][{mode}/fp32] {key}: rms {rms:.3e} >= {RMS_RATIO_FP32:.0e}" + + # The comparative assertion -- fp32 only, on purpose. + # + # Its premise is "headwise's disagreement with CP=1 is the floor this environment + # imposes". That holds only while both algorithms perform the *same* reductions. + # They do not: headwise hands each rank the whole sequence and 1/cp of the heads, so + # a gradient like conv1d.weight / dt_bias / A_log (a sum over every token) is summed + # in one go exactly as at CP=1 and can come out bit-exact. Chunkwise splits the + # tokens, so that same sum really is partitioned and re-added. In fp32 the mantissa + # absorbs it and the two are directly comparable (observed 1.00x-1.10x). In bf16 the + # repartitioned sum sits on the format's ULP floor while headwise sits near zero, so + # their *ratio* measures the dtype, not the algorithm -- bf16 is covered by the + # absolute gates above instead. + worst = [] + for key, (rms_c, cos_c, max_c) in metrics["chunkwise"].items(): + rms_h = metrics["headwise"][key][0] + worst.append((rms_c / max(rms_h, 1e-12), key, rms_c, rms_h)) + if dtype is not torch.float32: + continue + assert rms_c <= CHUNKWISE_VS_HEADWISE_RMS_FACTOR * max(rms_h, RMS_FLOOR_FP32), ( + f"[rank{rank}][{dtype_name}] {key}: chunkwise rms {rms_c:.3e} exceeds " + f"{CHUNKWISE_VS_HEADWISE_RMS_FACTOR}x the shipped headwise rms {rms_h:.3e} " + f"(cosine {cos_c:.8f}, max |diff| {max_c:.3e})" + ) + worst.sort(reverse=True) + if rank == 0: + print(f"\n[{dtype_name}] chunkwise vs headwise, worst 6 by rms ratio:") + for ratio, key, rms_c, rms_h in worst[:6]: + shown = f"{ratio:6.2f}x" if rms_h > 0 else " n/a" + print(f" {key:38s} {shown} chunkwise {rms_c:.3e} headwise {rms_h:.3e}") + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _worker_layout_round_trip(rank, world_size, _spec, _unused): + """zigzag -> contiguous -> zigzag over a real CP group must be token-exact. + + This is the collective-level version of RFC 5.1: it drives the actual + ``all_to_all`` in ``context_parallel_layout``, for packed THD (several + unequal-length samples) and for SBHD. + """ + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core.context_parallel_layout import ( + contiguous_to_zigzag_chunks, + get_thd_context_parallel_rank_indices, + zigzag_to_contiguous_chunks, + ) + + device = torch.device("cuda", rank) + cp_group = dist.new_group(list(range(world_size))) + + # --- packed THD, three samples of different lengths --- + lengths = [2 * world_size * f for f in (5, 1, 3)] + cu = torch.tensor([0] + torch.tensor(lengths).cumsum(0).tolist(), device=device, dtype=torch.int32) + total = int(cu[-1]) + # Row t is (t, t+1e6, t+2e6): a permuted token is impossible to miss. + full = ( + torch.arange(total, dtype=torch.float64, device=device).unsqueeze(1) + + torch.arange(3, dtype=torch.float64, device=device).unsqueeze(0) * 1e6 + ) + + zig_idx = get_thd_context_parallel_rank_indices(cu, world_size, rank, "zigzag") + con_idx = get_thd_context_parallel_rank_indices(cu, world_size, rank, "contiguous") + local_zig = full[zig_idx] + + got_con = zigzag_to_contiguous_chunks(local_zig, cp_group, seq_dim=0, cu_seqlens=cu) + assert torch.equal(got_con, full[con_idx]), f"rank {rank}: THD zigzag->contiguous is wrong" + got_zig = contiguous_to_zigzag_chunks(got_con, cp_group=cp_group, seq_dim=0, cu_seqlens=cu) + assert torch.equal(got_zig, local_zig), f"rank {rank}: THD round trip is not identity" + + # --- SBHD (chunk-level swap, no cu_seqlens) --- + seq_local = 2 * world_size * 4 + sbhd = torch.arange(seq_local * 2 * 3, dtype=torch.float64, device=device).reshape(seq_local, 2, 3) + rank * 1e9 + swapped = zigzag_to_contiguous_chunks(sbhd, cp_group, seq_dim=0) + back = contiguous_to_zigzag_chunks(swapped, cp_group=cp_group, seq_dim=0) + assert torch.equal(back, sbhd), f"rank {rank}: SBHD round trip is not identity" + + dist.barrier() + dist.destroy_process_group() + + +def _worker_illegal_modes_fail_fast(rank, world_size, _spec, _unused): + """Illegal / unresolved GDN CP modes must raise, not silently pick an + algorithm.""" + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + from megatron.core.packed_seq_params import PackedSeqParams + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + device = torch.device("cuda", rank) + cp_group = parallel_state.get_context_parallel_group() + solo = [dist.new_group([r]) for r in range(world_size)][rank] + + seq_lens = [256, 128] + total = sum(seq_lens) + cu = torch.tensor([0, seq_lens[0], total], device=device, dtype=torch.int32) + + def _psp(group, local_cp_size): + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + max_seqlen_q=max(seq_lens), + max_seqlen_kv=max(seq_lens), + cp_group=group, + local_cp_size=local_cp_size, + ) + + gdn, config = _build_gdn(world_size, "all_gather", torch.bfloat16) + hidden = torch.randn(total // world_size, 1, config.hidden_size, device=device, dtype=torch.bfloat16) + + # 1. all_gather is a Relax-wrapper mode: reaching MCore's forward with cp>1 means the + # wrapper is missing, and that must be loud. + with pytest.raises(RuntimeError, match="implemented by the Relax"): + gdn(hidden, None, packed_seq_params=_psp(cp_group, world_size)) + + # 2. ...but a CP=1 micro-batch is legal under any declared mode: it needs no CP + # communication at all, so the mode is never consulted. + full = torch.randn(total, 1, config.hidden_size, device=device, dtype=torch.bfloat16) + gdn(full, None, packed_seq_params=_psp(solo, 1)) + + # 3. local_cp_size disagreeing with the group it selected means some collective would + # run on the wrong group. + gdn_hw, _ = _build_gdn(world_size, "headwise", torch.bfloat16) + with pytest.raises(ValueError, match="local_cp_size"): + gdn_hw(hidden, None, packed_seq_params=_psp(cp_group, world_size + 1)) + + # 4. deterministic mode has no CP-context scan. + gdn_cw, _ = _build_gdn(world_size, "chunkwise", torch.bfloat16) + gdn_cw.config.deterministic_mode = True + try: + with pytest.raises((ValueError, AssertionError)): + gdn_cw(hidden, None, packed_seq_params=_psp(cp_group, world_size)) + finally: + gdn_cw.config.deterministic_mode = False + + # 5. inference is not supported. + class _Ctx: + def is_static_batching(self): + return True + + with pytest.raises(NotImplementedError): + gdn_cw(hidden, None, inference_context=_Ctx(), packed_seq_params=_psp(cp_group, world_size)) + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _worker_state_dict_invariance(rank, world_size, _spec, _unused): + """GDN weights must stay TP-only: same keys and shard dims in every CP + mode.""" + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + + signatures = {} + for cp_size, mode in ((1, "headwise"), (world_size, "headwise"), (world_size, "chunkwise")): + gdn, _ = _build_gdn(cp_size, mode, torch.bfloat16) + sd = gdn.state_dict() + sharded = gdn.sharded_state_dict(prefix="mixer.") + signatures[(cp_size, mode)] = ( + {k: tuple(v.shape) for k, v in sd.items() if torch.is_tensor(v)}, + { + k: ( + tuple(getattr(v, "global_shape", ())), + tuple(getattr(v, "local_shape", ())), + getattr(v, "axis_fragmentations", None), + ) + for k, v in sorted(sharded.items()) + }, + ) + del gdn + + baseline = signatures[(1, "headwise")] + assert baseline[0], "state_dict is empty; the invariance check would be vacuous" + assert baseline[1], "sharded_state_dict is empty; the invariance check would be vacuous" + for key, sig in signatures.items(): + assert sig[0] == baseline[0], f"state_dict shapes changed for {key}" + assert set(sig[1]) == set(baseline[1]), f"sharded_state_dict keys changed for {key}" + for k in baseline[1]: + assert sig[1][k] == baseline[1][k], f"sharded shard dims changed for {key} at {k}" + + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +# --------------------------------------------------------------------------- +# pytest entry points +# --------------------------------------------------------------------------- +def _spawn(fn, spec, port, extra=None): + os.environ["MASTER_PORT"] = str(port) + mp.spawn(fn, args=(WORLD_SIZE, extra if extra is not None else spec, None), nprocs=WORLD_SIZE, join=True) + + +@needs_backport +@pytest.mark.parametrize("dtype_name,port", [("fp32", 29540), ("bf16", 29541)]) +def test_fla_cp_kernels_match_single_rank(dtype_name, port): + """FLA causal_conv1d / chunk_gated_delta_rule under cp_context vs no CP.""" + _spawn(_worker_fla_kernels, dtype_name, port) + + +@needs_backport +@pytest.mark.parametrize("dtype_name,port", [("fp32", 29542), ("bf16", 29543)]) +def test_gdn_cp_matches_cp1(dtype_name, port): + """RFC 3.4 item 3: CP=2 vs CP=1, forward and backward, both CP + algorithms.""" + _spawn(_worker_gdn_module, dtype_name, port) + + +@needs_backport +def test_layout_round_trip_over_real_cp_group(): + """RFC 5.1 at the collective level: the layout swap is a pure + permutation.""" + _spawn(_worker_layout_round_trip, "n/a", 29544) + + +@needs_backport +def test_illegal_gdn_cp_modes_fail_fast(): + """RFC emphasis 1: illegal combinations must fail fast, never pick + silently.""" + _spawn(_worker_illegal_modes_fail_fast, "n/a", 29545) + + +@needs_backport +def test_gdn_state_dict_invariant_across_cp_modes(): + """RFC 3.4 item 4: checkpoint keys and shard dims do not depend on the CP + mode.""" + _spawn(_worker_state_dict_invariance, "n/a", 29546) diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py new file mode 100644 index 000000000..49cd4912d --- /dev/null +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Unit tests for the GDN chunkwise-CP layout backport (Task 32, phase 1). + +Covers the pure-tensor half of the backported MCore capability: the two THD CP +partitions (``zigzag`` / ``contiguous``), the route tensors that implement the +all-to-all between them, and the construction-time ``linear_cp_mode`` gate. + +Everything here runs on CPU with no process group: a CP all-to-all is emulated +locally by decoding every rank's route and delivering the pieces by hand, which +is exactly what NCCL would do and lets us assert token-level identity for CP +sizes we do not have GPUs for. + +The real-kernel / real-collective half lives in +``test_gdn_chunkwise_cp_gpu.py``. +""" + +from __future__ import annotations + +import pytest +import torch + + +cpl = pytest.importorskip("megatron.core.context_parallel_layout", reason="requires the patched Megatron-LM") + +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group # noqa: E402 + +from relax.backends.megatron.cp_utils import gdn_cp_slice, slice_with_cp # noqa: E402 + + +def _cu(lengths: list[int]) -> torch.Tensor: + cu = [0] + for n in lengths: + cu.append(cu[-1] + n) + return torch.tensor(cu, dtype=torch.int64) + + +def _tagged_tokens(total: int, width: int = 3) -> torch.Tensor: + """[total, width] where row t is (t, t+1e6, t+2e6): token identity is + unambiguous.""" + base = torch.arange(total, dtype=torch.float64).unsqueeze(1) + return base + torch.arange(width, dtype=torch.float64).unsqueeze(0) * 1e6 + + +# --------------------------------------------------------------------------- +# Partition definitions +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("cp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("layout", ["zigzag", "contiguous"]) +@pytest.mark.parametrize("lengths_factor", [[1], [1, 2, 3], [3, 1, 1, 2]]) +def test_thd_rank_indices_partition_all_tokens_exactly_once(cp_size, layout, lengths_factor): + lengths = [2 * cp_size * f for f in lengths_factor] + cu = _cu(lengths) + owned = torch.cat([cpl.get_thd_context_parallel_rank_indices(cu, cp_size, r, layout) for r in range(cp_size)]) + assert owned.numel() == int(cu[-1]) + assert torch.equal(torch.sort(owned).values, torch.arange(int(cu[-1]))) + + +@pytest.mark.parametrize("cp_size", [2, 4, 8]) +def test_zigzag_rank_indices_match_relax_data_sharding(cp_size): + """MCore's zigzag partition must be token-for-token what Relax's data path + produces. + + If these ever disagree, chunkwise CP would silently permute tokens relative + to the all-gather fallback and the attention layers. + """ + lengths = [2 * cp_size * f for f in (1, 3, 2)] + cu = _cu(lengths) + full = _tagged_tokens(int(cu[-1])).reshape(-1, 1, 3) # [s, b=1, C] + + for rank in range(cp_size): + mcore_idx = cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, "zigzag") + mcore_shard = full[mcore_idx] + + # Relax data.py: per-sample slice_with_cp then concat. + relax_shard = torch.cat( + [ + slice_with_cp( + full[cu[i] : cu[i + 1]], + pad_value=0.0, + qkv_format="thd", + dynamic_cp_size=cp_size, + dynamic_cp_rank=rank, + ) + for i in range(len(lengths)) + ], + dim=0, + ) + assert torch.equal(mcore_shard, relax_shard) + + # Relax model.py (all-gather fallback) re-slices with gdn_cp_slice. + assert torch.equal(mcore_shard, gdn_cp_slice(full, cu, cp_size, rank)) + + +@pytest.mark.parametrize("cp_size", [2, 4, 8]) +@pytest.mark.parametrize("lengths_factor", [[1], [1, 2, 3], [3, 1, 1, 2]]) +def test_both_layouts_are_permutations_of_each_other(cp_size, lengths_factor): + """The two partitions must describe the same token set with the same per- + rank size. + + That is the precondition for the all-to-all between them to be a pure + permutation -- no token invented, dropped, or duplicated. The real collective + round trip is asserted in ``test_gdn_chunkwise_cp_gpu.py``. + """ + lengths = [2 * cp_size * f for f in lengths_factor] + cu = _cu(lengths) + total = int(cu[-1]) + for rank in range(cp_size): + zig = cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, "zigzag") + con = cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, "contiguous") + assert zig.numel() == con.numel() == total // cp_size + # contiguous is exactly this rank's span of the flattened buffer + assert torch.equal(con, torch.arange(rank * (total // cp_size), (rank + 1) * (total // cp_size))) + # zigzag is two equal chunks per sequence, in local storage order + assert torch.equal(torch.sort(zig).values, torch.sort(zig).values) + + +@pytest.mark.parametrize("cp_size", [2, 4]) +def test_rank_indices_reject_lengths_not_divisible_by_two_cp(cp_size): + bad = _cu([2 * cp_size, 2 * cp_size + 1]) + with pytest.raises(ValueError, match="divisible by"): + cpl.get_thd_context_parallel_rank_indices(bad, cp_size, 0, "zigzag") + + +def test_rank_indices_reject_unknown_layout(): + with pytest.raises(ValueError, match="Unsupported context-parallel layout"): + cpl.get_thd_context_parallel_rank_indices(_cu([16, 16]), 2, 0, "contiguous_ish") + + +# --------------------------------------------------------------------------- +# Dynamic CP group resolution +# --------------------------------------------------------------------------- +def test_resolve_cp_group_prefers_packed_seq_params(): + static = object() + dynamic = object() + assert resolve_cp_group(static, None) is static + assert resolve_cp_group(static, PackedSeqParams(qkv_format="thd")) is static + assert resolve_cp_group(static, PackedSeqParams(qkv_format="thd", cp_group=dynamic)) is dynamic + + +# --------------------------------------------------------------------------- +# Construction-time capability gate +# --------------------------------------------------------------------------- +def _gdn_config(**overrides): + import torch.nn.functional as F + from megatron.core.transformer.transformer_config import TransformerConfig + + kwargs = dict( + hidden_size=2048, + num_layers=1, + num_attention_heads=16, + num_query_groups=2, + normalization="RMSNorm", + use_cpu_initialization=True, + activation_func=F.silu, + bf16=True, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + ) + kwargs.update(overrides) + return TransformerConfig(**kwargs) + + +def test_config_default_mode_is_headwise(): + """Upgrading the image must not silently reroute an existing recipe.""" + assert _gdn_config().linear_cp_mode == "headwise" + + +def test_headwise_config_requires_heads_divisible_by_tp_times_cp(): + # 16 key heads, tp=2, cp=4 -> 16 % 8 == 0: fine. + _gdn_config(tensor_model_parallel_size=2, context_parallel_size=4) + # tp=2, cp=16 -> 16 % 32 != 0: the geometry headwise cannot express. + with pytest.raises(AssertionError, match="linear_num_key_heads"): + _gdn_config(tensor_model_parallel_size=2, context_parallel_size=16) + + +def test_chunkwise_config_only_requires_heads_divisible_by_tp(): + """This is what replaces Relax's temporary head-count rewrite hack.""" + cfg = _gdn_config(tensor_model_parallel_size=2, context_parallel_size=16, linear_cp_mode="chunkwise") + assert cfg.linear_num_key_heads == 16 and cfg.linear_num_value_heads == 32 + # ... but TP divisibility is still enforced: GDN weights stay TP-sharded. + with pytest.raises(AssertionError, match="linear_num_key_heads"): + _gdn_config( + tensor_model_parallel_size=8, + context_parallel_size=2, + linear_cp_mode="chunkwise", + num_query_groups=8, + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + + +def test_all_gather_config_uses_the_tp_only_head_rule(): + """`--gdn-cp-mode=all_gather` must be constructible on a non-divisible + geometry. + + Relax's all-gather fallback keeps GDN weights TP-only, so declaring it + should relax the head check exactly as chunkwise does. + """ + cfg = _gdn_config(tensor_model_parallel_size=2, context_parallel_size=16, linear_cp_mode="all_gather") + assert cfg.linear_num_key_heads == 16 and cfg.linear_num_value_heads == 32 + + +def test_config_rejects_unresolved_and_unknown_linear_cp_mode(): + """`auto` is resolved before construction; MCore must never see it.""" + for bad in ("auto", "allgather", "chunk", ""): + with pytest.raises(AssertionError, match="linear_cp_mode"): + _gdn_config(context_parallel_size=2, linear_cp_mode=bad) + with pytest.raises(AssertionError, match="linear_cp_mode"): + _gdn_config(context_parallel_size=4, tensor_model_parallel_size=2, linear_cp_mode=bad) From 31b38db2526a86bbe955f4521dcea80e92016987 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Fri, 7 Aug 2026 13:58:54 +1000 Subject: [PATCH 2/7] version v2 --- docker/Dockerfile | 25 +- .../patch/megatron/20260805-85bced0ae.patch | 262 ++++++------- docker/patch/megatron/TASK32-BACKPORT.md | 134 ------- .../backends/megatron/gdn_cp_numeric_probe.py | 300 --------------- .../megatron/test_gdn_chunkwise_cp_gpu.py | 347 +++++++++++++++++- .../megatron/test_gdn_chunkwise_cp_layout.py | 50 ++- 6 files changed, 488 insertions(+), 630 deletions(-) delete mode 100644 docker/patch/megatron/TASK32-BACKPORT.md delete mode 100644 tests/backends/megatron/gdn_cp_numeric_probe.py diff --git a/docker/Dockerfile b/docker/Dockerfile index c57768f82..4eb0ed441 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -38,21 +38,12 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29 FROM base as train # flash-linear-attention 0.4.2 (pulls fla-core 0.4.2) is the first release carrying -# `fla.ops.cp`, which GDN chunkwise context parallel needs. Fail the build here if the -# capability is missing, so a silent wheel/index change can never produce an image whose -# GDN chunkwise path only breaks at training time. NOTE: `fla-core` depends on an -# unpinned `torch`, so pip re-resolves torch's own looser `nvidia-cudnn-cu12` pin here; +# `fla.ops.cp`, which GDN chunkwise context parallel needs. NOTE: `fla-core` depends +# on an unpinned `torch`, so pip re-resolves torch's own looser `nvidia-cudnn-cu12` pin here; # the explicit `nvidia-cudnn-cu12==9.16.0.29` install further down runs after this layer # and is what the final image keeps. RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir && \ pip install --no-cache-dir flash-linear-attention==0.4.2 && \ - python -c "import inspect; \ -from fla.ops.cp import build_cp_context; \ -from fla.modules.convolution import causal_conv1d; \ -from fla.ops.gated_delta_rule import chunk_gated_delta_rule; \ -assert 'cp_context' in inspect.signature(causal_conv1d).parameters, 'causal_conv1d lacks cp_context'; \ -assert 'cp_context' in inspect.signature(chunk_gated_delta_rule).parameters, 'chunk_gated_delta_rule lacks cp_context'; \ -print('FLA chunkwise-CP capability OK')" && \ pip install --no-cache-dir tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ # FA3 (Hopper flash-attention), built from source. This commit's _flash_attn_forward carries @@ -149,18 +140,6 @@ RUN cd Megatron-LM && \ exit 1; \ fi && \ rm megatron.patch && \ - python -c "import inspect; \ -import megatron.core.context_parallel_layout as cpl; \ -from megatron.core.packed_seq_params import resolve_cp_group; \ -from megatron.core.ssm.gated_delta_net import GatedDeltaNet; \ -from megatron.core.transformer.transformer_config import TransformerConfig; \ -assert hasattr(cpl, 'zigzag_to_contiguous_chunks') and hasattr(cpl, 'contiguous_to_zigzag_chunks'); \ -assert hasattr(cpl, 'get_thd_context_parallel_rank_indices'); \ -assert hasattr(GatedDeltaNet, '_resolve_cp_routing') and hasattr(GatedDeltaNet, '_build_chunkwise_cp_context'); \ -assert 'linear_cp_mode' not in inspect.signature(GatedDeltaNet.forward).parameters, \ - 'linear_cp_mode must be static config, not a per-forward argument'; \ -assert TransformerConfig.linear_cp_mode == 'headwise', TransformerConfig.linear_cp_mode; \ -print('MCore chunkwise-CP backport OK')" && \ apt update && apt install -y jq && ln -s /usr/local/lib/python3.12/dist-packages/torch_memory_saver_hook_mode_preload_cu12.abi3.so \ /usr/local/lib/python3.12/dist-packages/torch_memory_saver_hook_mode_preload.abi3.so diff --git a/docker/patch/megatron/20260805-85bced0ae.patch b/docker/patch/megatron/20260805-85bced0ae.patch index ec77afb0f..e24e6852b 100644 --- a/docker/patch/megatron/20260805-85bced0ae.patch +++ b/docker/patch/megatron/20260805-85bced0ae.patch @@ -529,27 +529,10 @@ new file mode 100644 index 0000000..e3acdd9 --- /dev/null +++ b/megatron/core/context_parallel_layout.py -@@ -0,0 +1,324 @@ +@@ -0,0 +1,307 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + -+"""Context parallel tensor layout helpers. -+ -+Backported VERBATIM from Megatron-LM commit 5139086e (NVIDIA/Megatron-LM#3282, -+merged to dev) for Relax Task 32 v1. No Relax modifications -- if this file ever -+diverges from upstream, that is a bug. -+ -+Two layouts: -+ -+``zigzag`` -+ Megatron's attention load-balanced partition. Each sequence is cut into -+ ``2 * cp_size`` chunks and rank ``r`` owns chunks ``r`` and -+ ``2 * cp_size - r - 1``. -+``contiguous`` -+ The flattened packed buffer is cut into ``cp_size`` equal spans and rank -+ ``r`` owns span ``r``. This is the layout FLA's chunkwise CP kernels expect: -+ ``fla.ops.cp.build_cp_context`` derives each rank's local ``cu_seqlens`` -+ from exactly this partition. -+""" ++"""Context parallel tensor layout helpers.""" + +from typing import List, Optional, Tuple + @@ -1338,7 +1321,7 @@ diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_param index 322f12a..0be9e5e 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py -@@ -64,3 +64,20 @@ class PackedSeqParams: +@@ -64,3 +64,17 @@ class PackedSeqParams: .to(torch.int32) .unsqueeze(0) # Add a batch dimension ) @@ -1350,11 +1333,8 @@ index 322f12a..0be9e5e 100644 + """Return the dynamic CP group from packed_seq_params when available, else the static one. + + Dynamic CP assigns a per-microbatch CP group that may differ from the -+ process-group stored at model construction time. This helper centralises -+ the resolution logic so every consumer (layout conversion, GDN, loss and -+ recompute paths) derives size/rank from exactly the same group object. -+ -+ Backported from Megatron-LM 5139086e. ++ process-group stored at model construction time. This helper centralises ++ the resolution logic used by GPTModel, GatedDeltaNet, and MTP layers. + """ + if packed_seq_params is not None and packed_seq_params.cp_group is not None: + return packed_seq_params.cp_group @@ -1479,25 +1459,14 @@ index 8df4df1..f39eff8 100644 from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_context_parallel import ( _all_to_all_cp2hp, -@@ -52,6 +56,17 @@ except ImportError: - - HAVE_FLA = False - -+try: -+ # FLA >= 0.4.2. Only linear_cp_mode="chunkwise" needs it; every other GDN path -+ # works against 0.4.1, so a missing fla.ops.cp disables chunkwise and nothing else. +@@ -43,5 +47,6 @@ try: + from fla.modules.convolution import causal_conv1d + from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context -+ -+ HAVE_FLA_CP = True -+except ImportError: -+ build_cp_context = None -+ -+ HAVE_FLA_CP = False -+ - - logger = logging.getLogger(__name__) + from fla.ops.gated_delta_rule import chunk_gated_delta_rule -@@ -233,6 +248,15 @@ class GatedDeltaNet(MegatronModule): + HAVE_FLA = True +@@ -233,6 +238,15 @@ class GatedDeltaNet(MegatronModule): tp_group=self.pg_collection.tp, ) @@ -1513,96 +1482,70 @@ index 8df4df1..f39eff8 100644 self.reset_parameters() def reset_parameters(self): -@@ -257,6 +281,62 @@ class GatedDeltaNet(MegatronModule): - ).uniform_(*self.A_init_range) - self.A_log.data.copy_(torch.log(A)) - -+ def _resolve_cp_routing(self, packed_seq_params, pg_collection): -+ """Split the resolved CP group into a headwise and a chunkwise group. -+ -+ The CP algorithm itself comes from ``self.config.linear_cp_mode``, which is -+ static for the whole process: dynamic context parallelism varies only the CP -+ group and its size, never the algorithm. Nothing here writes to ``self`` or to -+ the shared config, and no process group is ever created inside a forward. -+ -+ The two algorithms are mutually exclusive: whichever one runs owns the whole -+ CP group and the other gets ``None``, treated as size 1 everywhere below. +@@ -265,6 +279,7 @@ class GatedDeltaNet(MegatronModule): + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, ++ pg_collection: Optional[ProcessGroupCollection] = None, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ): +@@ -286,10 +301,65 @@ class GatedDeltaNet(MegatronModule): + """ + # TODO: Deal with attention_mask +- + -+ Returns ``(headwise_group, headwise_size, chunkwise_group, chunkwise_size)``. -+ """ -+ base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp -+ cp_group = resolve_cp_group(base_cp_group, packed_seq_params) -+ cp_size = cp_group.size() if cp_group is not None else 1 + inference_context = deprecate_inference_params(inference_context, inference_params) +- + -+ # Dynamic CP must keep local_cp_size and the group it selected consistent; -+ # a mismatch means some collective would be issued on the wrong group. ++ # Bridge repacking may replace PackedSeqParams. Validate the object received ++ # by every GDN forward; group.size() is local metadata and needs no collective. + if packed_seq_params is not None: ++ dynamic_cp_group = getattr(packed_seq_params, "cp_group", None) + local_cp_size = getattr(packed_seq_params, "local_cp_size", None) -+ if local_cp_size is not None and local_cp_size != cp_size: ++ if (dynamic_cp_group is None) != (local_cp_size is None): ++ raise ValueError( ++ "PackedSeqParams.cp_group and local_cp_size must either both be set " ++ "or both be None." ++ ) ++ if dynamic_cp_group is not None and local_cp_size != dynamic_cp_group.size(): + raise ValueError( -+ f"PackedSeqParams.local_cp_size ({local_cp_size}) does not match the " -+ f"resolved CP group size ({cp_size}). Dynamic CP must set both from the " -+ f"same group." ++ f"PackedSeqParams.local_cp_size ({local_cp_size}) does not match " ++ f"cp_group.size() ({dynamic_cp_group.size()})." + ) ++ base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp ++ cp_group = resolve_cp_group(base_cp_group, packed_seq_params) ++ cp_size = cp_group.size() if cp_group is not None else 1 + ++ # Route the resolved group to exactly one CP algorithm. The unused path gets ++ # None/size 1, which makes its existing divisions and collectives no-ops. + if cp_size == 1: -+ # No CP communication either way; keep the original (size-1) group so the -+ # existing helpers behave exactly as before. Checked before the mode so a -+ # CP=1 micro-batch never trips the mode validation below. -+ return cp_group, 1, None, 1 -+ -+ mode = self.config.linear_cp_mode -+ if mode == "headwise": -+ return cp_group, cp_size, None, 1 -+ if mode == "chunkwise": -+ if not HAVE_FLA_CP: -+ raise RuntimeError( -+ "linear_cp_mode='chunkwise' requires fla.ops.cp (flash-linear-attention " -+ ">= 0.4.2), which is not importable in this environment." -+ ) -+ return None, 1, cp_group, cp_size -+ if mode == "all_gather": ++ cp_group_headwise = cp_group ++ cp_size_headwise = 1 ++ cp_group_chunkwise = None ++ cp_size_chunkwise = 1 ++ elif self.config.linear_cp_mode == "headwise": ++ cp_group_headwise = cp_group ++ cp_size_headwise = cp_size ++ cp_group_chunkwise = None ++ cp_size_chunkwise = 1 ++ elif self.config.linear_cp_mode == "chunkwise": ++ cp_group_headwise = None ++ cp_size_headwise = 1 ++ cp_group_chunkwise = cp_group ++ cp_size_chunkwise = cp_size ++ elif self.config.linear_cp_mode == "all_gather": + raise RuntimeError( + "linear_cp_mode='all_gather' is implemented by the Relax GatedDeltaNet " + "wrapper, not by Megatron. Reaching this forward with cp_size=" + f"{cp_size} means the wrapper was not installed." + ) -+ raise ValueError( -+ f"Unsupported linear_cp_mode {mode!r}; expected 'headwise', 'chunkwise' or " -+ "'all_gather'. 'auto' must be resolved to a concrete mode before the model is " -+ "constructed." -+ ) -+ - def forward( - self, - hidden_states: Tensor, -@@ -265,6 +345,7 @@ class GatedDeltaNet(MegatronModule): - packed_seq_params: Optional[PackedSeqParams] = None, - sequence_len_offset: Optional[int] = None, - *, -+ pg_collection: Optional[ProcessGroupCollection] = None, - inference_params: Optional[BaseInferenceContext] = None, - **kwargs, - ): -@@ -279,6 +360,8 @@ class GatedDeltaNet(MegatronModule): - packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format. - sequence_len_offset (Optional[int]): Sequence length offset used for - inference CUDA graphs. -+ pg_collection (Optional[ProcessGroupCollection]): Overrides the process -+ groups captured at construction time. - - Return: - (Tuple[Tensor, Tensor]) GDN output and bias. -@@ -288,8 +371,21 @@ class GatedDeltaNet(MegatronModule): - - inference_context = deprecate_inference_params(inference_context, inference_params) - -+ ( -+ cp_group_headwise, -+ cp_size_headwise, -+ cp_group_chunkwise, -+ cp_size_chunkwise, -+ ) = self._resolve_cp_routing(packed_seq_params, pg_collection) ++ else: ++ raise ValueError( ++ f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; expected " ++ "'headwise', 'chunkwise' or 'all_gather'. 'auto' must be resolved to a " ++ "concrete mode before the model is constructed." ++ ) + seq_len, batch, _ = hidden_states.shape - seq_len = seq_len * self.sp_size * self.cp_size @@ -1616,7 +1559,7 @@ index 8df4df1..f39eff8 100644 if inference_context is not None: assert ( -@@ -299,24 +395,30 @@ class GatedDeltaNet(MegatronModule): +@@ -299,24 +369,30 @@ class GatedDeltaNet(MegatronModule): # TODO: support inference raise NotImplementedError("GDN does not support inference for now.") @@ -1651,7 +1594,7 @@ index 8df4df1..f39eff8 100644 ) assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " -@@ -331,21 +433,59 @@ class GatedDeltaNet(MegatronModule): +@@ -331,21 +407,59 @@ class GatedDeltaNet(MegatronModule): cu_seqlens_q = None cu_seqlens_kv = None @@ -1673,7 +1616,7 @@ index 8df4df1..f39eff8 100644 + # Build the chunkwise CP context once per forward. It carries this rank's *local* + # cu_seqlens plus the neighbour bookkeeping the FLA kernels need in order to + # exchange conv boundary tokens and chunk-boundary states. -+ chunkwise_cp_context = self._build_chunkwise_cp_context( ++ cu_seqlens_q, chunkwise_cp_context = self._build_chunkwise_cp_context( + cp_group_chunkwise, cp_size_chunkwise, cu_seqlens_q, seq_len_global, batch + ) + @@ -1684,7 +1627,7 @@ index 8df4df1..f39eff8 100644 + # Chunkwise CP expects the contiguous-time layout (rank r owns global chunks + # [2r, 2r+1]) inside conv1d / chunk_gated_delta_rule, because that is the -+ # partition fla.ops.cp.build_cp_context assumes. Megatron and Relax hand us the ++ # partition fla.ops.cp.build_cp_context assumes. Megatron hands us the + # zigzag attention-load-balanced layout (rank r owns [r, 2*cp-r-1]), so reshuffle + # over the CP group with a single all-to-all. No full-sequence gather, and the + # tensor stays sharded 1/cp throughout. @@ -1714,7 +1657,7 @@ index 8df4df1..f39eff8 100644 split_sections=[ self.qk_dim_local_tp, self.qk_dim_local_tp, -@@ -357,12 +497,12 @@ class GatedDeltaNet(MegatronModule): +@@ -357,12 +471,12 @@ class GatedDeltaNet(MegatronModule): ) outputs.append(qkvzba_i) qkvzba = torch.cat(outputs, dim=0) @@ -1729,7 +1672,7 @@ index 8df4df1..f39eff8 100644 split_sections=[ self.qk_dim_local_tp, self.qk_dim_local_tp, -@@ -381,10 +521,10 @@ class GatedDeltaNet(MegatronModule): +@@ -381,10 +495,10 @@ class GatedDeltaNet(MegatronModule): qkv, gate, beta, alpha = torch.split( qkvzba, [ @@ -1744,7 +1687,7 @@ index 8df4df1..f39eff8 100644 ], dim=-1, ) -@@ -403,14 +543,14 @@ class GatedDeltaNet(MegatronModule): +@@ -403,14 +517,14 @@ class GatedDeltaNet(MegatronModule): conv1d_weight = get_parameter_local_cp( self.conv1d.weight, dim=0, @@ -1761,7 +1704,7 @@ index 8df4df1..f39eff8 100644 split_sections=qkv_channels_split_sections, ) if self.conv_bias -@@ -425,7 +565,7 @@ class GatedDeltaNet(MegatronModule): +@@ -425,7 +539,7 @@ class GatedDeltaNet(MegatronModule): stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, @@ -1770,13 +1713,11 @@ index 8df4df1..f39eff8 100644 ) qkv = self.act_fn(conv_out[..., :seq_len]) qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d -@@ -439,22 +579,24 @@ class GatedDeltaNet(MegatronModule): +@@ -439,22 +553,22 @@ class GatedDeltaNet(MegatronModule): initial_state=None, output_final_state=False, cu_seqlens=cu_seqlens_q, -+ # Only pass cp_context when chunkwise is active, so with it off this call -+ # is byte-identical to before the backport (and still works on FLA 0.4.1). -+ **({} if chunkwise_cp_context is None else {"cp_context": chunkwise_cp_context}), ++ cp_context=chunkwise_cp_context, ) nvtx_range_pop(suffix="conv1d") @@ -1802,15 +1743,15 @@ index 8df4df1..f39eff8 100644 g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") -@@ -469,6 +611,7 @@ class GatedDeltaNet(MegatronModule): +@@ -469,6 +583,7 @@ class GatedDeltaNet(MegatronModule): output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=cu_seqlens_q, -+ **({} if chunkwise_cp_context is None else {"cp_context": chunkwise_cp_context}), ++ cp_context=chunkwise_cp_context, ) nvtx_range_pop(suffix="gated_delta_rule") -@@ -482,19 +625,31 @@ class GatedDeltaNet(MegatronModule): +@@ -482,19 +597,31 @@ class GatedDeltaNet(MegatronModule): norm_out = norm_out.reshape(batch, seq_len, -1) norm_out = norm_out.transpose(0, 1).contiguous() @@ -1846,7 +1787,7 @@ index 8df4df1..f39eff8 100644 ) # Output projection -@@ -504,6 +659,45 @@ class GatedDeltaNet(MegatronModule): +@@ -504,6 +631,48 @@ class GatedDeltaNet(MegatronModule): return out, out_bias @@ -1862,18 +1803,18 @@ index 8df4df1..f39eff8 100644 + contract of chunkwise CP. + """ + if cp_size_chunkwise <= 1: -+ return None ++ return cu_seqlens_q, None + + if cu_seqlens_q is not None: -+ return build_cp_context( ++ return cu_seqlens_q, build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + + # Non-packed (SBHD) input: cu_seqlens is fully determined by the static global -+ # sequence length and batch size, so cache the context. Rebuilding it every -+ # forward allocates fresh tensors, which breaks CUDA graph replay. ++ # sequence length and batch size, so cache both the generated boundaries and ++ # context. Rebuilding either every forward breaks CUDA graph replay. + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: @@ -1881,10 +1822,13 @@ index 8df4df1..f39eff8 100644 + torch.arange(batch + 1, device=torch.cuda.current_device(), dtype=torch.long) + * seq_len_global + ) -+ cached = build_cp_context( -+ cu_seqlens=cached_cu_seqlens, -+ group=cp_group_chunkwise, -+ conv1d_kernel_size=self.conv_kernel_dim, ++ cached = ( ++ cached_cu_seqlens, ++ build_cp_context( ++ cu_seqlens=cached_cu_seqlens, ++ group=cp_group_chunkwise, ++ conv1d_kernel_size=self.conv_kernel_dim, ++ ), + ) + self._chunkwise_cp_context_cache[cache_key] = cached + return cached @@ -1892,7 +1836,7 @@ index 8df4df1..f39eff8 100644 @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm -@@ -517,15 +711,23 @@ class GatedDeltaNet(MegatronModule): +@@ -517,15 +686,23 @@ class GatedDeltaNet(MegatronModule): return y @jit_fuser @@ -1918,7 +1862,7 @@ index 8df4df1..f39eff8 100644 dim=-1, ) -@@ -538,7 +740,7 @@ class GatedDeltaNet(MegatronModule): +@@ -538,7 +715,7 @@ class GatedDeltaNet(MegatronModule): query_key = l2norm(query_key.contiguous()) # Split query and key @@ -1927,7 +1871,7 @@ index 8df4df1..f39eff8 100644 query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) -@@ -567,7 +769,9 @@ class GatedDeltaNet(MegatronModule): +@@ -567,7 +744,9 @@ class GatedDeltaNet(MegatronModule): beta = beta.sigmoid() return g, beta @@ -1938,7 +1882,7 @@ index 8df4df1..f39eff8 100644 """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" if cu_seqlens_padded is not None: cu_seqlens = cu_seqlens_padded -@@ -582,6 +786,18 @@ class GatedDeltaNet(MegatronModule): +@@ -582,6 +761,18 @@ class GatedDeltaNet(MegatronModule): f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." ) @@ -1957,7 +1901,7 @@ index 8df4df1..f39eff8 100644 return cu_seqlens def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): -@@ -780,8 +996,10 @@ def get_parameter_local_cp( +@@ -780,8 +971,10 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ @@ -1970,7 +1914,7 @@ index 8df4df1..f39eff8 100644 # No need to split if CP size is 1. if cp_size == 1: -@@ -800,7 +1018,7 @@ def get_parameter_local_cp( +@@ -800,7 +993,7 @@ def get_parameter_local_cp( slices = [slice(None)] * param.dim() dim_size = param.size(dim=dim) slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) @@ -1979,6 +1923,26 @@ index 8df4df1..f39eff8 100644 return param +@@ -935,7 +1128,8 @@ def torch_chunk_gated_delta_rule( + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=None, ++ cp_context=None, + ): + # pylint: disable=line-too-long + ''' +@@ -948,6 +1142,9 @@ def torch_chunk_gated_delta_rule( + assert ( + cu_seqlens is None + ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." ++ assert ( ++ cp_context is None ++ ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." +- ++ + initial_dtype = query.dtype + if use_qk_l2norm_in_kernel: diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 610700f..1ef4fee 100644 --- a/megatron/core/tensor_parallel/layers.py diff --git a/docker/patch/megatron/TASK32-BACKPORT.md b/docker/patch/megatron/TASK32-BACKPORT.md deleted file mode 100644 index 87e61213b..000000000 --- a/docker/patch/megatron/TASK32-BACKPORT.md +++ /dev/null @@ -1,134 +0,0 @@ -# Task 32 v1 — GDN chunkwise context parallel: MCore backport manifest - -This file is the file-level and hunk-level record required by the Task 32 RFC -review: what was taken from upstream, what was deliberately left behind, and why. -It describes the delta between `20260506-85bced0ae.patch` (previous) and -`20260805-85bced0ae.patch` (current). Both patches apply to the same source tree, -so the whole difference below *is* the Task 32 change. - -**v1 depends on NVIDIA/Megatron-LM#3282 only.** NVIDIA/Megatron-LM#5664 is -explicitly out of scope per the RFC review and contributes nothing to this patch. - -## Base - -| Thing | Value | -| ------------------------------------------ | ---------------------------------------------------------------------------------------------- | -| Megatron-Bridge | `2faedbf6fe3c422835a44b2b360cadcb2a116a54` | -| Megatron-LM (`.dev.commit` of that Bridge) | `85bced0ae6ab46f61a0fd774074a3273daf6ae02` | -| Tree assembly | `cp -r Bridge/src/megatron` then `rsync` MCore's `megatron/` over it (see `docker/Dockerfile`) | - -## Upstream sources - -| Ref | State | Used for | -| ------------------------------------------------------------------------------------------ | --------------- | --------------------------------------------------------------------------------------------- | -| [Megatron-LM #3282](https://github.com/NVIDIA/Megatron-LM/pull/3282), merged as `5139086e` | merged to `dev` | The chunkwise CP feature. The only MCore source used by v1. | -| flash-linear-attention / fla-core `0.4.2` | released | `fla.ops.cp.build_cp_context`; `cp_context=` on `causal_conv1d` and `chunk_gated_delta_rule`. | - -`5139086e` is ~597 commits ahead of `85bced0a`, and it edits the same regions of -`gated_delta_net.py` / `transformer_config.py` that Relax's own patch edits, so a -cherry-pick is not possible. Everything below is a selective port onto the pinned -tree. - -## Design constraint from the RFC review - -The GDN CP mode is **static for the whole process**. `linear_cp_mode` is read by -both the construction-time head check and by `GatedDeltaNet.forward`; there is no -per-call override and nothing in a forward writes to `self` or to the shared -config. Dynamic context parallelism varies only `cp_group` / `local_cp_size`, -never the algorithm. - -## Files added / changed - -### `megatron/core/context_parallel_layout.py` — new, +324 - -**Byte-identical to `5139086e` below the module docstring.** Verified with -`diff` against the upstream file; the only change is a Relax provenance note in -the docstring. Contains: - -| Symbol | Purpose | -| ------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------- | -| `get_thd_context_parallel_rank_indices` | Reference description of both partitions. Tests assert Relax's `slice_with_cp` / `gdn_cp_slice` against it. | -| `zigzag_to_contiguous_chunks`, `contiguous_to_zigzag_chunks` | Public entry points; dispatch on whether `cu_seqlens` is given. | -| `_zigzag_contiguous_thd_swap` | Packed THD: one packed-token all-to-all. Builds its routing from `cu_seqlens` on each call. | -| `_zigzag_contiguous_chunk_swap` | SBHD: one chunk-level all-to-all. | - -### `megatron/core/packed_seq_params.py` — +17 - -- `resolve_cp_group()` from `5139086e`. The single place where "dynamic - per-micro-batch CP group, else the construction-time one" is decided, so every - consumer derives size/rank from the same group object. -- `local_cp_size` / `cp_group` already existed on `85bced0a`; not re-backported. -- The dataclass field list is **untouched**. - -### `megatron/core/transformer/transformer_config.py` — +50 / −16 - -- `linear_cp_mode` field and the head-divisibility rule from `5139086e`: headwise - needs `heads % (tp * cp) == 0`, chunkwise needs `heads % tp == 0`. -- **Relax adaptation:** the default is `"headwise"`, not upstream's `"chunkwise"`. - Upgrading the image must not silently change the algorithm an existing recipe - runs. -- **Relax adaptation:** `"all_gather"` is accepted as a third declared value, - using the same TP-only head rule as chunkwise. That is what makes a - non-divisible geometry constructible for Relax's all-gather fallback, and it - means the declared config equals the resolved `--gdn-cp-mode` instead of - declaring one mode while running another. -- Unknown values — including an unresolved `"auto"` — assert at construction. -- **Excluded:** `gdn_conv_pad_alignment`, `gdn_pre_gated_delta_rule_fusion` and - their interaction asserts. Neither field exists on this base. - -### `megatron/core/ssm/gated_delta_net.py` — +290 / −42 - -- `_resolve_cp_routing()`: resolve the CP group once via `resolve_cp_group`, then - give the whole group to exactly one of headwise / chunkwise and `None` to the - other (`None` is treated as size 1 everywhere downstream). No process group is - ever created in a forward. -- Validates `PackedSeqParams.local_cp_size == cp_group.size()`; a mismatch means - a collective would run on the wrong group. -- `cp_size == 1` short-circuits *before* the mode is read, so a CP=1 micro-batch - is legal under any declared mode. -- `linear_cp_mode="all_gather"` raises if MCore's own forward is reached with - `cp_size > 1`: that mode is implemented by the Relax wrapper, so arriving here - means the wrapper was not installed. -- zigzag ↔ contiguous conversion around the conv + scan; `cp_context` built once - per forward and passed to both FLA kernels. -- `_resolve_cu_seqlens` gains the `cp_size` divisibility check from `5139086e`. -- **Relax adaptation — backwards compatibility:** - - `_prepare_qkv_for_gated_delta_rule` takes `cp_size_headwise` as an *optional* - argument defaulting to `self.cp_size`, so Relax's existing all-gather - fallback (`relax/backends/megatron/model.py`) keeps calling it unchanged. - - `get_parameter_local_cp` accepts a `None` group as size 1. - - `cp_context=` is only passed when chunkwise is active, so with chunkwise off - the FLA call is byte-identical to before the backport — and still works - against FLA 0.4.1. -- **Preserved Relax fixes** (both from the previous patch, unchanged): - - `torch._dynamo.config.patch(disable=True)` around - `_prepare_qkv_for_gated_delta_rule` (Qwen3.6 `torch.compile` failure); - - `param[tuple(slices)]` in `get_parameter_local_cp` (multi-dim basic indexing). -- **Excluded from `5139086e`:** the `_forward_compute` split and `recompute_gdn` - selective-recompute wrapper; `gdn_pre_gated_delta_rule_fusion` / - `_fused_streamed_pre_gated_delta_rule`; `gdn_conv_pad_alignment` conv padding; - the `_a2a_cp_to_hp` / `_a2a_hp_to_cp` refactor of the headwise path; the rename - of `get_parameter_local_cp` to `get_parameter_local_cp_headwise`. All are - independent changes from the 597-commit gap and would alter existing paths' - structure or numerics for no Task 32 benefit. - -### Explicitly not in this patch - -| Thing | Why | -| ------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Anything from NVIDIA/Megatron-LM#5664 | Out of scope for v1 per the RFC review. No `cp_partition_mode`, no route tensors, no `prebuild_thd_cp_partition_routes`, no `pad_between_seqs`. The THD swap rebuilds its routing per call, which is `5139086e` behaviour. | -| `megatron/core/extensions/transformer_engine.py` | The THD output-length fix in the RFC was conditional on adopting `pad_between_seqs`. That representation is not backported and no test on this base reproduces the mismatch. | -| `pyproject.toml` / `uv.lock` | Relax installs FLA from `docker/Dockerfile`, not from upstream package metadata. | - -## Verifying this patch is exactly what it claims - -```bash -# byte-identity of the new module against upstream, below the docstring -curl -s https://raw.githubusercontent.com/NVIDIA/Megatron-LM/5139086e/megatron/core/context_parallel_layout.py \ - > /tmp/up.py -diff <(sed -n '/^from typing import/,$p' /megatron/core/context_parallel_layout.py) \ - <(sed -n '/^from typing import/,$p' /tmp/up.py) # must be empty - -# no #5664 content anywhere -grep -rn 'cp_partition_route\|prebuild_thd_cp\|pad_between_seqs' /megatron/ # must be empty -``` diff --git a/tests/backends/megatron/gdn_cp_numeric_probe.py b/tests/backends/megatron/gdn_cp_numeric_probe.py deleted file mode 100644 index aa7f6cee3..000000000 --- a/tests/backends/megatron/gdn_cp_numeric_probe.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Cross-image numerical probe for the GDN context-parallel paths. - -RFC Task 32 acceptance item 3.4-2: the FLA 0.4.1 -> 0.4.2 upgrade and the MCore -backport must not move the numbers on any path that already existed, i.e. CP=1, -MCore headwise CP, and Relax's all-gather fallback. That cannot be a normal unit -test, because it compares *two images* -- it needs the same tensors produced -under the old dependency set and the new one. - -So this module is a runner, not a pytest file (hence the non-``test_`` name): - - # inside the OLD image - python -m tests.backends.megatron.gdn_cp_numeric_probe dump --mode headwise --out /out/old - # inside the NEW image - python -m tests.backends.megatron.gdn_cp_numeric_probe dump --mode headwise --out /out/new - # anywhere - python -m tests.backends.megatron.gdn_cp_numeric_probe compare --ref /out/old --cand /out/new - -Everything that could drift between images is pinned by hand: every parameter is -overwritten with a tensor drawn from a name-seeded CPU generator, and the inputs -come from a fixed-seed CPU generator too. So a difference in the report is a -difference in the *kernels*, not in initialisation order or RNG plumbing. - -Modes: - cp1 1 GPU, no CP at all. - headwise 2 GPUs, MCore's native cp2hp all-to-all path. - all_gather 2 GPUs, Relax's `_dcp_gdn_forward` fallback (head geometry chosen so - the dispatcher cannot use headwise), which also exercises the FLA - conv/scan kernels on the full gathered sequence. - chunkwise 2 GPUs, the newly backported path (candidate image only). -""" - -from __future__ import annotations - -import argparse -import json -import os -import zlib -from types import SimpleNamespace - -import torch -import torch.multiprocessing as mp - - -MODE_WORLD_SIZE = {"cp1": 1, "headwise": 2, "all_gather": 2, "chunkwise": 2} -# all_gather is only reachable when the heads do NOT divide tp * cp. -MODE_HEADS = { - "cp1": (4, 8), - "headwise": (4, 8), - "chunkwise": (4, 8), - "all_gather": (1, 2), -} -SEQ_LENS = [256, 128] -HIDDEN_SIZE = 512 - - -def _deterministic_fill_(module) -> None: - """Overwrite every parameter from a name-seeded generator. - - Makes the dump independent of Megatron's initialisation code, which is one - of the things the patch touches. - """ - with torch.no_grad(): - for name, p in sorted(module.named_parameters()): - gen = torch.Generator(device="cpu").manual_seed(zlib.crc32(name.encode()) & 0x7FFFFFFF) - values = torch.randn(p.shape, generator=gen, dtype=torch.float32) * 0.05 - p.copy_(values.to(device=p.device, dtype=p.dtype)) - - -def _build(mode, cp_size): - import torch.nn.functional as F - from megatron.core import parallel_state - from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - get_experimental_attention_variant_module_spec, - ) - from megatron.core.process_groups_config import ProcessGroupCollection - from megatron.core.ssm.gated_delta_net import GatedDeltaNet - from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed - from megatron.core.transformer.transformer_config import TransformerConfig - - model_parallel_cuda_manual_seed(123) - num_key_heads, num_value_heads = MODE_HEADS[mode] - - if mode == "all_gather": - # Relax relaxes Megatron's headwise `% (tp*cp)` config gate down to `% tp` so a - # non-divisible geometry can even be constructed. v1 keeps this helper working - # untouched, and using it in both images keeps the A/B identical. (The backported - # `linear_cp_mode="all_gather"` is the eventual replacement; it is covered by - # test_gdn_chunkwise_cp_layout.py instead, so it cannot skew this comparison.) - from relax.backends.megatron.model import _relax_gdn_cp_config_assert - - _relax_gdn_cp_config_assert() - - config = TransformerConfig( - hidden_size=HIDDEN_SIZE, - num_layers=1, - num_attention_heads=8, - num_query_groups=2, - normalization="RMSNorm", - use_cpu_initialization=True, - layernorm_zero_centered_gamma=True, - activation_func=F.silu, - bf16=True, - tensor_model_parallel_size=1, - context_parallel_size=cp_size, - experimental_attention_variant="gated_delta_net", - linear_attention_freq=[1], - linear_conv_kernel_dim=4, - linear_key_head_dim=64, - linear_value_head_dim=64, - linear_num_key_heads=num_key_heads, - linear_num_value_heads=num_value_heads, - transformer_impl="transformer_engine", - # The CP algorithm is static config, resolved once at launch. Only chunkwise needs - # to be declared here, and only the candidate image has the field at all -- for - # cp1 / headwise / all_gather both images must run literally the same code, which - # is the whole point of this probe. - **( - {"linear_cp_mode": "chunkwise"} - if mode == "chunkwise" and "linear_cp_mode" in TransformerConfig.__dataclass_fields__ - else {} - ), - ) - pg_collection = ProcessGroupCollection( - tp=parallel_state.get_tensor_model_parallel_group(), - cp=parallel_state.get_context_parallel_group(), - ) - gdn = ( - GatedDeltaNet( - config, - submodules=get_experimental_attention_variant_module_spec(config=config).submodules, - layer_number=1, - bias=False, - conv_bias=False, - conv_init=1.0, - use_qk_l2norm=True, - A_init_range=(1, 16), - pg_collection=pg_collection, - ) - .cuda() - .bfloat16() - ) - _deterministic_fill_(gdn) - return gdn, config - - -def _dump_worker(rank, mode, out_dir): - world_size = MODE_WORLD_SIZE[mode] - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - torch.cuda.set_device(rank) - import torch.distributed as dist - - dist.init_process_group("nccl", rank=rank, world_size=world_size) - from megatron.core import parallel_state - from megatron.core.packed_seq_params import PackedSeqParams - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - context_parallel_size=world_size, - ) - cp_size = world_size - gdn, config = _build(mode, cp_size) - - if mode == "all_gather": - # Install Relax's runtime GDN wrapper and satisfy its full-recompute gate. - from relax.backends.megatron import model as relax_model - - relax_model.get_args = lambda: SimpleNamespace(recompute_granularity="full") - relax_model._patch_gdn_for_dynamic_cp() - - device = torch.device("cuda", rank) - total = sum(SEQ_LENS) - cu = torch.tensor([0, SEQ_LENS[0], total], device=device, dtype=torch.int32) - local_total = total // cp_size - - gen = torch.Generator(device="cpu").manual_seed(20260805) - hidden_full = torch.randn(total, 1, HIDDEN_SIZE, generator=gen, dtype=torch.float32) - grad_full = torch.randn(total, 1, HIDDEN_SIZE, generator=gen, dtype=torch.float32) - - if cp_size == 1: - hidden_local = hidden_full - grad_local = grad_full - else: - from relax.backends.megatron.cp_utils import gdn_cp_slice - - cu_list = [0, SEQ_LENS[0], total] - hidden_local = gdn_cp_slice(hidden_full, cu_list, cp_size, rank) - grad_local = gdn_cp_slice(grad_full, cu_list, cp_size, rank) - assert hidden_local.shape[0] == local_total - - h = hidden_local.to(device=device, dtype=torch.bfloat16).clone().requires_grad_(True) - psp = PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu, - cu_seqlens_kv=cu, - cu_seqlens_q_padded=cu, - cu_seqlens_kv_padded=cu, - max_seqlen_q=max(SEQ_LENS), - max_seqlen_kv=max(SEQ_LENS), - cp_group=parallel_state.get_context_parallel_group(), - local_cp_size=cp_size, - ) - out, _ = gdn(h, None, packed_seq_params=psp) - (out.float() * grad_local.to(device)).sum().backward() - - payload = {"out": out.detach().float().cpu(), "grad_hidden": h.grad.detach().float().cpu()} - for name, p in sorted(gdn.named_parameters()): - payload[f"grad::{name}"] = p.grad.detach().float().cpu() - - os.makedirs(out_dir, exist_ok=True) - torch.save(payload, os.path.join(out_dir, f"{mode}.rank{rank}.pt")) - if rank == 0: - import importlib.metadata as md - - meta = { - "mode": mode, - "world_size": world_size, - "fla": md.version("flash-linear-attention"), - "fla_core": md.version("fla-core"), - "torch": torch.__version__, - "heads": MODE_HEADS[mode], - "seq_lens": SEQ_LENS, - } - with open(os.path.join(out_dir, f"{mode}.meta.json"), "w") as fh: - json.dump(meta, fh, indent=2) - print(json.dumps(meta)) - - dist.barrier() - parallel_state.destroy_model_parallel() - dist.destroy_process_group() - - -def _cmp(a, b): - a, b = a.flatten().double(), b.flatten().double() - diff = (a - b).abs() - denom = b.abs().clamp_min(1e-12) - cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item() - rms = ((a - b).square().mean().sqrt() / (b.square().mean().sqrt() + 1e-12)).item() - return { - "max_abs": diff.max().item(), - "max_rel": (diff / denom).max().item(), - "rms_ratio": rms, - "cosine": cos, - "bitwise_equal": bool(torch.equal(a, b)), - } - - -def main(): - ap = argparse.ArgumentParser() - sub = ap.add_subparsers(dest="cmd", required=True) - - d = sub.add_parser("dump") - d.add_argument("--mode", choices=sorted(MODE_WORLD_SIZE), required=True) - d.add_argument("--out", required=True) - d.add_argument("--port", default="29601") - - c = sub.add_parser("compare") - c.add_argument("--ref", required=True) - c.add_argument("--cand", required=True) - c.add_argument("--mode", choices=sorted(MODE_WORLD_SIZE), required=True) - c.add_argument("--atol", type=float, default=2e-4) - c.add_argument("--rtol", type=float, default=2e-3) - c.add_argument("--cos", type=float, default=0.99999) - - args = ap.parse_args() - if args.cmd == "dump": - os.environ["MASTER_PORT"] = args.port - world_size = MODE_WORLD_SIZE[args.mode] - mp.spawn(_dump_worker, args=(args.mode, args.out), nprocs=world_size, join=True) - return - - world_size = MODE_WORLD_SIZE[args.mode] - failures, rows = [], [] - for rank in range(world_size): - ref = torch.load(os.path.join(args.ref, f"{args.mode}.rank{rank}.pt"), weights_only=True) - cand = torch.load(os.path.join(args.cand, f"{args.mode}.rank{rank}.pt"), weights_only=True) - assert set(ref) == set(cand), f"tensor sets differ: {set(ref) ^ set(cand)}" - for key in sorted(ref): - stats = _cmp(cand[key], ref[key]) - rows.append({"rank": rank, "tensor": key, **stats}) - # RFC section 5, "MCore GDN, CP=1: candidate image vs old image". - ok = stats["max_abs"] <= args.atol + args.rtol * abs(ref[key]).max().item() - ok = ok and stats["cosine"] >= args.cos - if not ok: - failures.append(rows[-1]) - - print(json.dumps(rows, indent=2)) - print( - f"\n{args.mode}: {len(rows)} tensors compared, " - f"{sum(r['bitwise_equal'] for r in rows)} bitwise identical, {len(failures)} outside tolerance" - ) - if failures: - print("FAIL") - raise SystemExit(1) - print("PASS") - - -if __name__ == "__main__": - main() diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py b/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py index 64bbb9576..d8b23d229 100644 --- a/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py @@ -25,7 +25,7 @@ chunkwise both drift the same way, the cause is shared plumbing, not the new code. -Run with 2 visible GPUs: +Most tests need 2 visible GPUs; the TP2/CP2 matrix test needs 4: pytest tests/backends/megatron/test_gdn_chunkwise_cp_gpu.py """ @@ -226,14 +226,16 @@ def _mk(*shape, dtype=dtype): q = l2norm(_mk(1, total, H, DK).contiguous()) k = l2norm(_mk(1, total, H, DK).contiguous()) v = _mk(1, total, H, DV) - g = -_mk(1, total, H, dtype=torch.float32).abs() * 0.1 - beta = _mk(1, total, H, dtype=torch.float32).sigmoid() + g0 = -_mk(1, total, H, dtype=torch.float32).abs() * 0.1 + beta0 = _mk(1, total, H, dtype=torch.float32).sigmoid() leaves_ref = [t.detach().clone().requires_grad_(True) for t in (q, k, v)] + g_ref = g0.detach().clone().requires_grad_(True) + beta_ref = beta0.detach().clone().requires_grad_(True) o_ref, _ = chunk_gated_delta_rule( *leaves_ref, - g=g, - beta=beta, + g=g_ref, + beta=beta_ref, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, @@ -243,11 +245,13 @@ def _mk(*shape, dtype=dtype): (o_ref.float() * o_grad.float()).sum().backward() leaves_cp = [t.detach()[:, lo:hi].clone().requires_grad_(True) for t in (q, k, v)] + g_cp = g0.detach()[:, lo:hi].clone().requires_grad_(True) + beta_cp = beta0.detach()[:, lo:hi].clone().requires_grad_(True) ctx2 = build_cp_context(cu_seqlens=cu, group=cp_group, conv1d_kernel_size=W) o_cp, _ = chunk_gated_delta_rule( *leaves_cp, - g=g[:, lo:hi], - beta=beta[:, lo:hi], + g=g_cp, + beta=beta_cp, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, @@ -258,6 +262,8 @@ def _mk(*shape, dtype=dtype): (o_cp.float() * o_grad[:, lo:hi].float()).sum().backward() for name, a, b in zip("qkv", leaves_cp, leaves_ref): _report_rms(f"{tag} gdr d{name}", a.grad, b.grad[:, lo:hi], gdn_ratio) + _report_rms(f"{tag} gdr dg", g_cp.grad, g_ref.grad[:, lo:hi], gdn_ratio) + _report_rms(f"{tag} gdr dbeta", beta_cp.grad, beta_ref.grad[:, lo:hi], gdn_ratio) dist.barrier() dist.destroy_process_group() @@ -266,7 +272,15 @@ def _mk(*shape, dtype=dtype): # --------------------------------------------------------------------------- # worker: full MCore GatedDeltaNet # --------------------------------------------------------------------------- -def _build_gdn(cp_size, linear_cp_mode, dtype, num_key_heads=4, num_value_heads=8): +def _build_gdn( + cp_size, + linear_cp_mode, + dtype, + num_key_heads=4, + num_value_heads=8, + tp_size=1, + deterministic_mode=False, +): import torch.nn.functional as F from megatron.core import parallel_state from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( @@ -288,8 +302,9 @@ def _build_gdn(cp_size, linear_cp_mode, dtype, num_key_heads=4, num_value_heads= layernorm_zero_centered_gamma=True, activation_func=F.silu, bf16=dtype is torch.bfloat16, - tensor_model_parallel_size=1, + tensor_model_parallel_size=tp_size, context_parallel_size=cp_size, + deterministic_mode=deterministic_mode, experimental_attention_variant="gated_delta_net", linear_attention_freq=[1], linear_conv_kernel_dim=4, @@ -318,11 +333,20 @@ def _build_gdn(cp_size, linear_cp_mode, dtype, num_key_heads=4, num_value_heads= return gdn.cuda().to(dtype), config -def _run_gdn_once(gdn, hidden, psp, grad_out): +def _run_gdn_once(gdn, hidden, psp, grad_out, *, recompute=False, **forward_kwargs): """One forward+backward; returns (out, d_hidden, {param: grad}).""" gdn.zero_grad(set_to_none=True) h = hidden.clone().requires_grad_(True) - out, _ = gdn(h, None, packed_seq_params=psp) + if recompute: + from torch.utils.checkpoint import checkpoint + + out = checkpoint( + lambda x: gdn(x, None, packed_seq_params=psp, **forward_kwargs)[0], + h, + use_reentrant=False, + ) + else: + out, _ = gdn(h, None, packed_seq_params=psp, **forward_kwargs) (out.float() * grad_out).sum().backward() grads = {n: p.grad.detach().float().clone() for n, p in gdn.named_parameters()} return out.detach().clone(), h.grad.detach().clone(), grads @@ -478,6 +502,249 @@ def _psp(group, local_cp_size): dist.destroy_process_group() +def _worker_deterministic_reference(rank, world_size, _spec, _unused): + """The torch-native deterministic rule must accept cp_context=None and + preserve headwise CP correctness.""" + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + from megatron.core.process_groups_config import ProcessGroupCollection + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + device = torch.device("cuda", rank) + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + solo = [dist.new_group([r]) for r in range(world_size)][rank] + + gdn, config = _build_gdn( + world_size, + "headwise", + torch.float32, + deterministic_mode=True, + ) + assert gdn.gated_delta_rule.__name__ == "torch_chunk_gated_delta_rule" + for p in gdn.parameters(): + dist.broadcast(p.data, src=0, group=cp_group) + + total = 64 + cu = torch.tensor([0, total], device=device, dtype=torch.int32) + gen = torch.Generator(device="cpu").manual_seed(17) + hidden_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + grad_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + + solo_pg = ProcessGroupCollection(tp=tp_group, cp=solo) + out_ref, in_grad_ref, param_grads_ref = _run_gdn_once( + gdn, + hidden_full, + None, + grad_full, + pg_collection=solo_pg, + ) + + hidden_shard = _zigzag_shard(hidden_full, cu, world_size, rank) + grad_shard = _zigzag_shard(grad_full, cu, world_size, rank) + out_cp, in_grad_cp, param_grads_cp = _run_gdn_once(gdn, hidden_shard, None, grad_shard) + + _report_rms( + f"[rank{rank}] deterministic out", + out_cp, + _zigzag_shard(out_ref, cu, world_size, rank), + RMS_RATIO_FP32, + ) + _report_rms( + f"[rank{rank}] deterministic d_hidden", + in_grad_cp, + _zigzag_shard(in_grad_ref, cu, world_size, rank), + RMS_RATIO_FP32, + ) + for name, grad in param_grads_cp.items(): + summed = grad.clone() + dist.all_reduce(summed, group=cp_group) + _report_rms( + f"[rank{rank}] deterministic grad {name}", + summed, + param_grads_ref[name], + RMS_RATIO_FP32, + ) + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _worker_recompute_parity(rank, world_size, _spec, _unused): + """External activation checkpointing must replay chunkwise collectives + without changing outputs or gradients.""" + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + from megatron.core.packed_seq_params import PackedSeqParams + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + device = torch.device("cuda", rank) + cp_group = parallel_state.get_context_parallel_group() + + gdn, config = _build_gdn(world_size, "chunkwise", torch.float32) + for p in gdn.parameters(): + dist.broadcast(p.data, src=0, group=cp_group) + + seq_lens = [128, 64] + total = sum(seq_lens) + cu = torch.tensor([0, seq_lens[0], total], device=device, dtype=torch.int32) + psp = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + max_seqlen_q=max(seq_lens), + max_seqlen_kv=max(seq_lens), + cp_group=cp_group, + local_cp_size=world_size, + ) + + gen = torch.Generator(device="cpu").manual_seed(29) + hidden_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + grad_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + hidden = _zigzag_shard(hidden_full, cu, world_size, rank) + grad = _zigzag_shard(grad_full, cu, world_size, rank) + + out_eager, in_grad_eager, param_grads_eager = _run_gdn_once(gdn, hidden, psp, grad) + out_recompute, in_grad_recompute, param_grads_recompute = _run_gdn_once( + gdn, + hidden, + psp, + grad, + recompute=True, + ) + + _report_rms(f"[rank{rank}] recompute out", out_recompute, out_eager, KERNEL_RMS_RATIO_FP32) + _report_rms( + f"[rank{rank}] recompute d_hidden", + in_grad_recompute, + in_grad_eager, + KERNEL_RMS_RATIO_FP32, + ) + assert set(param_grads_recompute) == set(param_grads_eager) + for name in param_grads_eager: + _report_rms( + f"[rank{rank}] recompute grad {name}", + param_grads_recompute[name], + param_grads_eager[name], + KERNEL_RMS_RATIO_FP32, + ) + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _worker_tp2_cp2(rank, world_size, _spec, _unused): + """Exercise TP head sharding and CP routing together.""" + assert world_size == 4 + _init_dist(rank, world_size) + import torch.distributed as dist + from megatron.core import parallel_state + from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.process_groups_config import ProcessGroupCollection + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=2, + ) + device = torch.device("cuda", rank) + cp_group = parallel_state.get_context_parallel_group() + cp_rank = cp_group.rank() + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_source = dist.get_process_group_ranks(cp_group)[0] + solo = [dist.new_group([r]) for r in range(world_size)][rank] + + modules = {} + headwise, config = _build_gdn(2, "headwise", torch.float32, tp_size=2) + for p in headwise.parameters(): + dist.broadcast(p.data, src=cp_source, group=cp_group) + modules["headwise"] = headwise + modules["chunkwise"], _ = _build_gdn(2, "chunkwise", torch.float32, tp_size=2) + modules["chunkwise"].load_state_dict(headwise.state_dict()) + sharded_signatures = {} + for mode, module in modules.items(): + sharded_signatures[mode] = { + key: ( + tuple(getattr(value, "global_shape", ())), + tuple(getattr(value, "local_shape", ())), + getattr(value, "axis_fragmentations", None), + ) + for key, value in sorted(module.sharded_state_dict(prefix="mixer.").items()) + } + assert sharded_signatures["headwise"] == sharded_signatures["chunkwise"] + + seq_lens = [128, 64] + total = sum(seq_lens) + cu = torch.tensor([0, seq_lens[0], total], device=device, dtype=torch.int32) + gen = torch.Generator(device="cpu").manual_seed(41) + hidden_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + grad_full = torch.randn(total, 1, config.hidden_size, generator=gen).to(device) + + def _psp(group, cp_size): + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + max_seqlen_q=max(seq_lens), + max_seqlen_kv=max(seq_lens), + cp_group=group, + local_cp_size=cp_size, + ) + + solo_pg = ProcessGroupCollection(tp=tp_group, cp=solo) + out_ref, in_grad_ref, param_grads_ref = _run_gdn_once( + headwise, + hidden_full, + _psp(solo, 1), + grad_full, + pg_collection=solo_pg, + ) + hidden = _zigzag_shard(hidden_full, cu, 2, cp_rank) + grad = _zigzag_shard(grad_full, cu, 2, cp_rank) + out_want = _zigzag_shard(out_ref, cu, 2, cp_rank) + in_grad_want = _zigzag_shard(in_grad_ref, cu, 2, cp_rank) + + for mode, module in modules.items(): + out, in_grad, param_grads = _run_gdn_once(module, hidden, _psp(cp_group, 2), grad) + _report_rms(f"[rank{rank}][{mode}] TP2/CP2 out", out, out_want, RMS_RATIO_FP32) + _report_rms( + f"[rank{rank}][{mode}] TP2/CP2 d_hidden", + in_grad, + in_grad_want, + RMS_RATIO_FP32, + ) + assert set(param_grads) == set(param_grads_ref) + for name, param_grad in param_grads.items(): + summed = param_grad.clone() + dist.all_reduce(summed, group=cp_group) + _report_rms( + f"[rank{rank}][{mode}] TP2/CP2 grad {name}", + summed, + param_grads_ref[name], + RMS_RATIO_FP32, + ) + + dist.barrier() + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + def _worker_layout_round_trip(rank, world_size, _spec, _unused): """zigzag -> contiguous -> zigzag over a real CP group must be token-exact. @@ -571,12 +838,13 @@ def _psp(group, local_cp_size): # 2. ...but a CP=1 micro-batch is legal under any declared mode: it needs no CP # communication at all, so the mode is never consulted. full = torch.randn(total, 1, config.hidden_size, device=device, dtype=torch.bfloat16) - gdn(full, None, packed_seq_params=_psp(solo, 1)) + solo_psp = _psp(solo, 1) + gdn(full, None, packed_seq_params=solo_psp) - # 3. local_cp_size disagreeing with the group it selected means some collective would - # run on the wrong group. + # 3. The final PackedSeqParams must describe the same runtime CP geometry + # through both fields. gdn_hw, _ = _build_gdn(world_size, "headwise", torch.bfloat16) - with pytest.raises(ValueError, match="local_cp_size"): + with pytest.raises(ValueError, match="does not match cp_group.size"): gdn_hw(hidden, None, packed_seq_params=_psp(cp_group, world_size + 1)) # 4. deterministic mode has no CP-context scan. @@ -605,6 +873,8 @@ def _worker_state_dict_invariance(rank, world_size, _spec, _unused): """GDN weights must stay TP-only: same keys and shard dims in every CP mode.""" _init_dist(rank, world_size) + import io + import torch.distributed as dist from megatron.core import parallel_state @@ -613,6 +883,7 @@ def _worker_state_dict_invariance(rank, world_size, _spec, _unused): pipeline_model_parallel_size=1, context_parallel_size=world_size, ) + device = torch.device("cuda", rank) signatures = {} for cp_size, mode in ((1, "headwise"), (world_size, "headwise"), (world_size, "chunkwise")): @@ -630,6 +901,20 @@ def _worker_state_dict_invariance(rank, world_size, _spec, _unused): for k, v in sorted(sharded.items()) }, ) + + # Exercise real serialization and loading, not just key comparison. + checkpoint = io.BytesIO() + torch.save(sd, checkpoint) + checkpoint.seek(0) + loaded = torch.load(checkpoint, map_location=device, weights_only=True) + restored, _ = _build_gdn(cp_size, mode, torch.bfloat16) + restored.load_state_dict(loaded, strict=True) + for name, tensor in sd.items(): + if torch.is_tensor(tensor): + assert torch.equal(restored.state_dict()[name], tensor), ( + f"state_dict round trip changed {name} for {(cp_size, mode)}" + ) + del restored del gdn baseline = signatures[(1, "headwise")] @@ -649,8 +934,17 @@ def _worker_state_dict_invariance(rank, world_size, _spec, _unused): # pytest entry points # --------------------------------------------------------------------------- def _spawn(fn, spec, port, extra=None): + _spawn_world(fn, WORLD_SIZE, spec, port, extra=extra) + + +def _spawn_world(fn, world_size, spec, port, extra=None): os.environ["MASTER_PORT"] = str(port) - mp.spawn(fn, args=(WORLD_SIZE, extra if extra is not None else spec, None), nprocs=WORLD_SIZE, join=True) + mp.spawn( + fn, + args=(world_size, extra if extra is not None else spec, None), + nprocs=world_size, + join=True, + ) @needs_backport @@ -668,6 +962,27 @@ def test_gdn_cp_matches_cp1(dtype_name, port): _spawn(_worker_gdn_module, dtype_name, port) +@needs_backport +def test_deterministic_headwise_cp_matches_cp1(): + """The torch-native rule accepts cp_context=None and stays correct under + headwise CP.""" + _spawn(_worker_deterministic_reference, "n/a", 29547) + + +@needs_backport +def test_chunkwise_recompute_matches_eager(): + """Replaying chunkwise forward during backward preserves every tested + gradient.""" + _spawn(_worker_recompute_parity, "n/a", 29548) + + +@needs_backport +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 CUDA devices for TP2/CP2") +def test_gdn_tp2_cp2_matches_cp1(): + """TP2/CP2 exercises TP head shards and both CP algorithms together.""" + _spawn_world(_worker_tp2_cp2, 4, "n/a", 29549) + + @needs_backport def test_layout_round_trip_over_real_cp_group(): """RFC 5.1 at the collective level: the layout swap is a pure diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py index 49cd4912d..80a1b39a1 100644 --- a/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py @@ -2,13 +2,13 @@ """Unit tests for the GDN chunkwise-CP layout backport (Task 32, phase 1). Covers the pure-tensor half of the backported MCore capability: the two THD CP -partitions (``zigzag`` / ``contiguous``), the route tensors that implement the -all-to-all between them, and the construction-time ``linear_cp_mode`` gate. +partitions (``zigzag`` / ``contiguous``), their agreement with Relax's existing +zigzag sharding, dynamic group resolution, and the construction-time +``linear_cp_mode`` gate. -Everything here runs on CPU with no process group: a CP all-to-all is emulated -locally by decoding every rank's route and delivering the pieces by hand, which -is exactly what NCCL would do and lets us assert token-level identity for CP -sizes we do not have GPUs for. +Everything here runs on CPU with no process group. It validates partition +definitions only; the actual all-to-all round trip is exercised with NCCL in +``test_gdn_chunkwise_cp_gpu.py``. The real-kernel / real-collective half lives in ``test_gdn_chunkwise_cp_gpu.py``. @@ -16,6 +16,8 @@ from __future__ import annotations +import inspect + import pytest import torch @@ -104,14 +106,23 @@ def test_both_layouts_are_permutations_of_each_other(cp_size, lengths_factor): lengths = [2 * cp_size * f for f in lengths_factor] cu = _cu(lengths) total = int(cu[-1]) + zig_by_rank = [] + con_by_rank = [] for rank in range(cp_size): zig = cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, "zigzag") con = cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, "contiguous") + zig_by_rank.append(zig) + con_by_rank.append(con) assert zig.numel() == con.numel() == total // cp_size # contiguous is exactly this rank's span of the flattened buffer assert torch.equal(con, torch.arange(rank * (total // cp_size), (rank + 1) * (total // cp_size))) - # zigzag is two equal chunks per sequence, in local storage order - assert torch.equal(torch.sort(zig).values, torch.sort(zig).values) + + # Across the whole CP group, both layouts are permutations of exactly the + # same global token rows. + assert torch.equal( + torch.cat(zig_by_rank).sort().values, + torch.cat(con_by_rank).sort().values, + ) @pytest.mark.parametrize("cp_size", [2, 4]) @@ -126,6 +137,23 @@ def test_rank_indices_reject_unknown_layout(): cpl.get_thd_context_parallel_rank_indices(_cu([16, 16]), 2, 0, "contiguous_ish") +@pytest.mark.parametrize("layout", ["zigzag", "contiguous"]) +def test_rank_indices_ignore_duplicate_boundaries(layout): + compact = torch.tensor([0, 16, 40], dtype=torch.int64) + padded = torch.tensor([0, 16, 40, 40, 40], dtype=torch.int64) + for rank in range(2): + assert torch.equal( + cpl.get_thd_context_parallel_rank_indices(compact, 2, rank, layout), + cpl.get_thd_context_parallel_rank_indices(padded, 2, rank, layout), + ) + + +@pytest.mark.parametrize("layout", ["zigzag", "contiguous"]) +def test_rank_indices_reject_decreasing_boundaries(layout): + with pytest.raises(ValueError, match="nondecreasing"): + cpl.get_thd_context_parallel_rank_indices(torch.tensor([0, 16, 8]), 2, 0, layout) + + # --------------------------------------------------------------------------- # Dynamic CP group resolution # --------------------------------------------------------------------------- @@ -212,3 +240,9 @@ def test_config_rejects_unresolved_and_unknown_linear_cp_mode(): _gdn_config(context_parallel_size=2, linear_cp_mode=bad) with pytest.raises(AssertionError, match="linear_cp_mode"): _gdn_config(context_parallel_size=4, tensor_model_parallel_size=2, linear_cp_mode=bad) + + +def test_gdn_forward_has_no_per_call_mode_override(): + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + assert "linear_cp_mode" not in inspect.signature(GatedDeltaNet.forward).parameters From 88c6a90dfcf3020d54b667ebd0ebaf626cc2f1e9 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Fri, 7 Aug 2026 20:36:09 +1000 Subject: [PATCH 3/7] fix headwise resolve cu_seqlens check --- .../patch/megatron/20260805-85bced0ae.patch | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/docker/patch/megatron/20260805-85bced0ae.patch b/docker/patch/megatron/20260805-85bced0ae.patch index e24e6852b..4456800a6 100644 --- a/docker/patch/megatron/20260805-85bced0ae.patch +++ b/docker/patch/megatron/20260805-85bced0ae.patch @@ -1582,7 +1582,7 @@ index 8df4df1..f39eff8 100644 - seq_len, + seq_len_global, "cu_seqlens_q", -+ cp_size=cp_size_chunkwise, ++ cp_size=cp_size, ) cu_seqlens_kv = self._resolve_cu_seqlens( packed_seq_params.cu_seqlens_kv_padded, @@ -1590,7 +1590,7 @@ index 8df4df1..f39eff8 100644 - seq_len, + seq_len_global, "cu_seqlens_kv", -+ cp_size=cp_size_chunkwise, ++ cp_size=cp_size, ) assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " @@ -1882,26 +1882,21 @@ index 8df4df1..f39eff8 100644 """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" if cu_seqlens_padded is not None: cu_seqlens = cu_seqlens_padded -@@ -582,6 +761,18 @@ class GatedDeltaNet(MegatronModule): +@@ -582,6 +761,13 @@ class GatedDeltaNet(MegatronModule): f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." ) -+ if cp_size > 1: -+ # Chunkwise CP slices the packed buffer per sequence; a sequence whose padded -+ # length is not a multiple of the CP size cannot be split evenly. (The zigzag -+ # layout needs 2*cp; that stricter check lives in the layout conversion, which -+ # reports the offending lengths.) -+ seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] -+ if (seq_lengths % cp_size != 0).any(): -+ raise ValueError( -+ f"GDN: all per-sequence lengths in {name} must be divisible by " -+ f"cp_size={cp_size}, but got lengths: {seq_lengths.tolist()}" -+ ) ++ seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] ++ if (seq_lengths % cp_size != 0).any(): ++ raise ValueError( ++ f"All per-sequence lengths in cu_seqlens must be divisible by cp_size={cp_size}, " ++ f"but got lengths: {seq_lengths.tolist()}" ++ ) + return cu_seqlens def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): -@@ -780,8 +971,10 @@ def get_parameter_local_cp( +@@ -780,8 +966,10 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ @@ -1914,7 +1909,7 @@ index 8df4df1..f39eff8 100644 # No need to split if CP size is 1. if cp_size == 1: -@@ -800,7 +993,7 @@ def get_parameter_local_cp( +@@ -800,7 +988,7 @@ def get_parameter_local_cp( slices = [slice(None)] * param.dim() dim_size = param.size(dim=dim) slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) @@ -1923,7 +1918,7 @@ index 8df4df1..f39eff8 100644 return param -@@ -935,7 +1128,8 @@ def torch_chunk_gated_delta_rule( +@@ -935,7 +1123,8 @@ def torch_chunk_gated_delta_rule( initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, @@ -1932,7 +1927,7 @@ index 8df4df1..f39eff8 100644 ): # pylint: disable=line-too-long ''' -@@ -948,6 +1142,9 @@ def torch_chunk_gated_delta_rule( +@@ -948,6 +1137,9 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." From 22b1206a19bfed47cfbc1030418d084a6bf4db00 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Mon, 10 Aug 2026 20:16:18 +1000 Subject: [PATCH 4/7] relax cp chunkwise intergration --- relax/backends/megatron/arguments.py | 37 +++ relax/backends/megatron/model.py | 158 +++--------- relax/backends/megatron/model_provider.py | 1 + .../megatron/test_gdn_chunkwise_cp_layout.py | 12 +- .../megatron/test_gdn_cp_mode_stage2.py | 229 ++++++++++++++++++ 5 files changed, 305 insertions(+), 132 deletions(-) create mode 100644 tests/backends/megatron/test_gdn_cp_mode_stage2.py diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index 27aaed382..25c4c92e8 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -126,6 +126,41 @@ def _validate_dynamic_context_parallel(args): args.max_seqlen_per_dp_cp_rank = args.max_tokens_per_gpu +def _validate_linear_cp_mode(args) -> None: + """Fail fast on `--linear-cp-mode` / flag combinations that are invalid for + every model, without needing the HF config. + + Geometry-dependent rejections (e.g. explicit `headwise` on heads not + divisible by `tp*max_cp`) can only be checked once the real GDN head counts + are known, which happens in MCore's `TransformerConfig.__post_init__` gate + -- not here. + """ + mode = getattr(args, "linear_cp_mode", "headwise") + allowed_modes = {"headwise", "chunkwise", "all_gather"} + if mode not in allowed_modes: + raise ValueError( + f"--linear-cp-mode must be one of {sorted(allowed_modes)!r}; got {mode!r}. v1 does not support 'auto'." + ) + + if mode == "chunkwise" and getattr(args, "allgather_cp", False): + raise ValueError( + "--linear-cp-mode=chunkwise is incompatible with --allgather-cp: chunkwise CP requires " + "Megatron's zig-zag THD packing, while --allgather-cp switches the data path to a single " + "contiguous per-rank chunk. Note --allgather-cp is a data/attention packing flag, unrelated " + "to the GDN `all_gather` CP mode." + ) + + cp_may_exceed_one = ( + getattr(args, "dynamic_context_parallel", False) or getattr(args, "context_parallel_size", 1) > 1 + ) + if mode == "chunkwise" and cp_may_exceed_one and getattr(args, "deterministic_mode", False): + raise ValueError( + "--linear-cp-mode=chunkwise does not support --deterministic-mode while CP>1 may occur: " + "the deterministic torch reference path only accepts cp_context=None. Use " + "--linear-cp-mode=headwise or =all_gather for deterministic CP>1 runs." + ) + + def validate_args(args): """Run megatron's own validate_args plus slime-specific megatron validations.""" @@ -172,6 +207,8 @@ class _DeviceProperty: assert args.calculate_per_token_loss, ( "--calculate-per-token-loss must be set when context_parallel_size > 1 or dynamic_context_parallel is enabled (required by Megatron-Bridge)." ) + + _validate_linear_cp_mode(args) return args diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 4196a2d4d..a39d5410c 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -305,14 +305,6 @@ def setup_model_and_optimizer( assert not args.moe_use_upcycling assert args.load is not None or args.pretrained_checkpoint is not None - # Relax the Megatron GDN head-vs-(tp*cp) config gate down to (tp) BEFORE the model - # provider finalizes the TransformerConfig (get_model_provider_func below triggers - # __post_init__), so high-CP GDN configs (e.g. TP2/CP16) validate. The matching - # forward all-gather path is installed by _patch_gdn_for_dynamic_cp after the model - # is built; see both functions for why % tp suffices (GDN weights are TP-only). - if getattr(args, "dynamic_context_parallel", False) or getattr(args, "context_parallel_size", 1) > 1: - _relax_gdn_cp_config_assert() - model = get_model( wrap_model_provider_with_freeze(get_model_provider_func(args, role), args), ModelType.encoder_or_decoder, @@ -338,6 +330,15 @@ def setup_model_and_optimizer( # (dynamic CP, or static context_parallel_size > 1), incl. weight-only # roles that still run forward. _patch_gdn_for_dynamic_cp() + model_config = get_model_config(model[0]) + if getattr(model_config, "experimental_attention_variant", None) == "gated_delta_net" and ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ): + logger.info( + f"[GDN CP] role={role} linear_cp_mode={model_config.linear_cp_mode} " + f"TP={model_config.tensor_model_parallel_size} max_CP={model_config.context_parallel_size} " + f"key_heads={model_config.linear_num_key_heads} value_heads={model_config.linear_num_value_heads}" + ) if args.only_load_weight: return model, None, None @@ -410,84 +411,13 @@ def _gdn_cp_gather_full(qkvzba, cu_seqlens_cpu, cp_size, cp_group): return gdn_cp_gather_full(qkvzba, cu_seqlens_cpu, cp_size, cp_group) -def _relax_gdn_cp_config_assert() -> None: - """Relax Megatron's GDN config gate ``linear_num_{key,value}_heads % - (tp*cp) == 0`` down to ``% tp`` so high-CP GDN configs (e.g. TP2/CP16) - finalize. - - Megatron's ``TransformerConfig.__post_init__`` enforces the *native* cp2hp - (split-sequence -> split-head) divisibility ``heads % (tp * cp)``. - ``_patch_gdn_for_dynamic_cp`` replaces that forward with an all-gather + duplicated - scan whose weights stay **TP-only** (``qk_dim_local_tp = qk_dim // tp``, etc.), so - only ``heads % tp`` is actually required. Without relaxing this config gate, TP2/CP16 - (16 % 32 != 0) aborts at config finalize (``get_model_provider_func`` -> ``finalize`` - -> ``__post_init__``) *before* the forward patch is installed. - - Only intervenes when the native check would reject but the relaxed ``% tp`` check - passes: it temporarily scales the two GDN head counts by ``cp`` (which preserves - ``value % key`` and makes ``heads % (tp*cp)`` hold), runs the original - ``__post_init__``, then restores them. Those head counts are validation-only in - ``__post_init__`` (no stored value is derived from them -- verified against Megatron - core), and ``GatedDeltaNet.__init__`` reads the restored config later, so nothing - downstream sees the temporary values. Idempotent; monkey-patch only (no upstream - edit), matching ``_patch_gdn_for_dynamic_cp``. - """ - try: - from megatron.core.transformer.transformer_config import TransformerConfig - except ImportError: - return - - if getattr(TransformerConfig, "_gdn_cp_relaxed", False): - return - - _orig_post_init = TransformerConfig.__post_init__ - - def _relaxed_post_init(self, *post_init_args, **post_init_kwargs): - if getattr(self, "experimental_attention_variant", None) == "gated_delta_net": - tp = self.tensor_model_parallel_size - cp = self.context_parallel_size - key = self.linear_num_key_heads or 0 - val = self.linear_num_value_heads or 0 - native_bad = cp > 1 and ((key % (tp * cp)) != 0 or (val % (tp * cp)) != 0) - relaxed_ok = tp > 0 and (key % tp) == 0 and (val % tp) == 0 - if native_bad and relaxed_ok: - # key%tp==0 => (key*cp)%(tp*cp)==0, and (val*cp)%(key*cp)==(val%key) so the - # value%key assert is preserved. Restored in `finally` before anything else - # (incl. GatedDeltaNet.__init__) reads the config. - self.linear_num_key_heads = key * cp - self.linear_num_value_heads = val * cp - try: - _orig_post_init(self, *post_init_args, **post_init_kwargs) - finally: - self.linear_num_key_heads = key - self.linear_num_value_heads = val - return - _orig_post_init(self, *post_init_args, **post_init_kwargs) - - TransformerConfig.__post_init__ = _relaxed_post_init - TransformerConfig._gdn_cp_relaxed = True - - def _patch_gdn_for_dynamic_cp() -> None: - """Monkey-patch GatedDeltaNet.forward for CP via all-gather + duplicated - scan. - - Megatron's native GDN forward implements CP by converting "split sequence" - into "split head" (``cp2hp`` all-to-all, ``num_value_heads // tp // cp``), - which forces ``num_heads % (tp * cp) == 0`` and breaks at high CP for - head-light models (e.g. Qwen3.5). This patch keeps that efficient native path - whenever the heads still divide ``tp * cp`` (``native_ok``), and only when - native would break does it fall back to all-gathering the full sequence across - CP, running the recurrent scan duplicated on each rank while keeping relax's - **TP** head-split intact, then re-slicing this rank's shard. The effective - constraint drops to ``num_heads % tp == 0`` (CP16 works), and weight - conversion / DCS sync / checkpoint (all TP-only) are untouched. - - Dynamic CP: size/group are read per micro-batch from ``packed_seq_params`` - (set in get_batch), falling back to the static CP group. The ``cp == 1``, - non-thd, and ``native_ok`` cases keep upstream behavior (swap the dynamic CP - group, call the original forward). Idempotent; avoids editing upstream - Megatron source. + """Patch GDN forward for dynamic CP and Relax's all-gather mode. + + CP=1 and MCore-native headwise/chunkwise modes call the patched MCore + forward directly. Only static ``linear_cp_mode='all_gather'`` with CP>1 + executes Relax's existing fallback. No shared module/config state is + modified. """ try: from megatron.core.ssm.gated_delta_net import GatedDeltaNet @@ -499,23 +429,6 @@ def _patch_gdn_for_dynamic_cp() -> None: _orig_forward = GatedDeltaNet.forward - def _call_orig_with_dynamic_cp( - self, cp_size, cp_group, hidden_states, attention_mask, inference_context, packed_seq_params, *args, **kwargs - ): - # cp == 1 or non-thd: preserve upstream behavior; just point the module at - # the (possibly dynamic) CP group for the original forward. - _orig_cp_size = self.cp_size - _orig_cp_group = self.pg_collection.cp - self.cp_size = cp_size - self.pg_collection.cp = cp_group - try: - return _orig_forward( - self, hidden_states, attention_mask, inference_context, packed_seq_params, *args, **kwargs - ) - finally: - self.cp_size = _orig_cp_size - self.pg_collection.cp = _orig_cp_group - def _dcp_gdn_forward( self, hidden_states, attention_mask, inference_context=None, packed_seq_params=None, *args, **kwargs ): @@ -525,26 +438,16 @@ def _dcp_gdn_forward( from .cp_utils import gdn_cp_slice cp_size, cp_group, cp_rank = _resolve_gdn_cp(self, packed_seq_params) - is_thd = packed_seq_params is not None and getattr(packed_seq_params, "qkv_format", None) == "thd" - # Native cp2hp (head-split) is exact and cheaper (no duplicated scan, GDN - # activation sharded by CP) whenever the heads divide tp*cp. Only fall back - # to the all-gather path when native would break the head split — i.e. when - # num_key_heads is not divisible by tp*cp (covers tp*cp > num_key_heads). - # num_value_heads is a multiple of num_key_heads, so this one check suffices. - native_ok = self.num_key_heads % (self.tp_size * cp_size) == 0 - if cp_size == 1 or not is_thd or native_ok: - return _call_orig_with_dynamic_cp( - self, - cp_size, - cp_group, - hidden_states, - attention_mask, - inference_context, - packed_seq_params, - *args, - **kwargs, + if cp_size == 1 or self.config.linear_cp_mode != "all_gather": + return _orig_forward( + self, hidden_states, attention_mask, inference_context, packed_seq_params, *args, **kwargs ) + is_thd = packed_seq_params is not None and getattr(packed_seq_params, "qkv_format", None) == "thd" + assert is_thd, ( + "GDN linear_cp_mode='all_gather' with cp_size>1 only supports packed (thd) sequences; " + "use linear_cp_mode='headwise' or 'chunkwise' for SBHD/static-batch inputs." + ) assert inference_context is None, "GDN all-gather CP path does not support inference." # Packed (thd) + deterministic is unsupported: a single conv/scan over the # concatenated samples would bleed state across cu_seqlens boundaries, and @@ -604,19 +507,14 @@ def _dcp_gdn_forward( ) # Reuse the module's own prep (split/l2norm/GQA-expand) with CP disabled so - # its internal `// self.cp_size` becomes a no-op. Wrap in the dynamo-disable + # its internal `// cp_size_headwise` becomes a no-op. Wrap in the dynamo-disable # guard added by docker/patch/megatron/20260506-85bced0ae.patch (Qwen3.6 GDN # torch.compile failure); calling _prepare_qkv_for_gated_delta_rule directly # would re-trigger that compile failure. - _saved_cp = self.cp_size - self.cp_size = 1 - try: - with torch._dynamo.config.patch(disable=True): - query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len - ) - finally: - self.cp_size = _saved_cp + with torch._dynamo.config.patch(disable=True): + query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( + qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise=1 + ) # g/beta from the full (un-CP-sliced) A_log / dt_bias. g, beta = self._compute_g_and_beta(self.A_log, self.dt_bias, alpha, beta) diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 7e9805608..213f6ddaf 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -217,6 +217,7 @@ def wrapped_model_provider( "pipeline_model_parallel_size", "virtual_pipeline_model_parallel_size", "context_parallel_size", + "linear_cp_mode", "expert_model_parallel_size", "expert_tensor_parallel_size", "variable_seq_lengths", diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py index 80a1b39a1..a651c8c2a 100644 --- a/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_layout.py @@ -132,6 +132,14 @@ def test_rank_indices_reject_lengths_not_divisible_by_two_cp(cp_size): cpl.get_thd_context_parallel_rank_indices(bad, cp_size, 0, "zigzag") +def test_gdn_rejects_packed_lengths_not_divisible_by_cp(): + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + cu = _cu([8, 6]) + with pytest.raises(ValueError, match="divisible by cp_size=4"): + GatedDeltaNet._resolve_cu_seqlens(None, None, cu, int(cu[-1]), "cu_seqlens_q", cp_size=4) + + def test_rank_indices_reject_unknown_layout(): with pytest.raises(ValueError, match="Unsupported context-parallel layout"): cpl.get_thd_context_parallel_rank_indices(_cu([16, 16]), 2, 0, "contiguous_ish") @@ -223,7 +231,7 @@ def test_chunkwise_config_only_requires_heads_divisible_by_tp(): def test_all_gather_config_uses_the_tp_only_head_rule(): - """`--gdn-cp-mode=all_gather` must be constructible on a non-divisible + """`--linear-cp-mode=all_gather` must be constructible on a non-divisible geometry. Relax's all-gather fallback keeps GDN weights TP-only, so declaring it @@ -234,7 +242,7 @@ def test_all_gather_config_uses_the_tp_only_head_rule(): def test_config_rejects_unresolved_and_unknown_linear_cp_mode(): - """`auto` is resolved before construction; MCore must never see it.""" + """MCore only accepts the three concrete execution modes.""" for bad in ("auto", "allgather", "chunk", ""): with pytest.raises(AssertionError, match="linear_cp_mode"): _gdn_config(context_parallel_size=2, linear_cp_mode=bad) diff --git a/tests/backends/megatron/test_gdn_cp_mode_stage2.py b/tests/backends/megatron/test_gdn_cp_mode_stage2.py new file mode 100644 index 000000000..96034560a --- /dev/null +++ b/tests/backends/megatron/test_gdn_cp_mode_stage2.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""CPU-only tests for the Task 32 Stage 2 Relax-side GDN CP routing. + +Covers the pieces added on top of the Stage 1 FLA/MCore backport +(``test_gdn_chunkwise_cp_layout.py``): the ``--linear-cp-mode`` CLI, invalid +Chunkwise combinations, and the thin runtime dispatcher installed on +``GatedDeltaNet.forward``. + +The dispatcher tests drive ``GatedDeltaNet.forward`` through duck-typed fakes +and hook/counter spies instead of a real distributed process group or FLA +kernel call, per task32-stage2-handoff.md §6.2 ("use hooks/counters, don't +infer routing from numerics"). Real-kernel / real-collective coverage stays in +``test_gdn_chunkwise_cp_gpu.py``. +""" + +from __future__ import annotations + +import argparse +from types import SimpleNamespace + +import pytest + + +pytest.importorskip("megatron.core.context_parallel_layout", reason="requires the patched Megatron-LM") + +from megatron.core.packed_seq_params import PackedSeqParams # noqa: E402 +from megatron.core.ssm.gated_delta_net import GatedDeltaNet # noqa: E402 + +from relax.backends.megatron import arguments as megatron_arguments # noqa: E402 +from relax.backends.megatron import model as gdn_model # noqa: E402 +from relax.backends.megatron.arguments import _validate_linear_cp_mode # noqa: E402 + + +# --------------------------------------------------------------------------- +# Step 1: CLI flag +# --------------------------------------------------------------------------- +def _parse_megatron_args(monkeypatch, *argv): + pytest.importorskip("sglang.srt.server_args") + from relax.utils.arguments import get_slime_extra_args_provider + + monkeypatch.setattr("sys.argv", ["test-linear-cp-mode", *argv]) + return megatron_arguments._megatron_parse_args( + extra_args_provider=get_slime_extra_args_provider(), + ignore_unknown_args=False, + ) + + +def test_linear_cp_mode_flag_defaults_to_headwise(monkeypatch): + args = _parse_megatron_args(monkeypatch) + assert args.linear_cp_mode == "headwise" + + +@pytest.mark.parametrize("mode", ["headwise", "chunkwise", "all_gather"]) +def test_linear_cp_mode_flag_accepts_all_concrete_modes(monkeypatch, mode): + args = _parse_megatron_args(monkeypatch, "--linear-cp-mode", mode) + assert args.linear_cp_mode == mode + + +# --------------------------------------------------------------------------- +# Step 2: argument validation +# --------------------------------------------------------------------------- +def _args(**overrides): + base = dict( + linear_cp_mode="headwise", + allgather_cp=False, + deterministic_mode=False, + dynamic_context_parallel=False, + context_parallel_size=1, + ) + base.update(overrides) + return argparse.Namespace(**base) + + +@pytest.mark.parametrize("bad", ["auto", "allgather"]) +def test_validate_linear_cp_mode_rejects_unsupported_value(bad): + with pytest.raises(ValueError, match="does not support 'auto'|must be one of"): + _validate_linear_cp_mode(_args(linear_cp_mode=bad)) + + +def test_validate_linear_cp_mode_rejects_chunkwise_with_allgather_cp(): + with pytest.raises(ValueError, match="allgather-cp"): + _validate_linear_cp_mode(_args(linear_cp_mode="chunkwise", allgather_cp=True)) + + +@pytest.mark.parametrize("cp_kwargs", [{"context_parallel_size": 2}, {"dynamic_context_parallel": True}]) +def test_validate_linear_cp_mode_rejects_chunkwise_deterministic_when_cp_may_exceed_one(cp_kwargs): + with pytest.raises(ValueError, match="deterministic"): + _validate_linear_cp_mode(_args(linear_cp_mode="chunkwise", deterministic_mode=True, **cp_kwargs)) + + +# --------------------------------------------------------------------------- +# Steps 4-6: runtime dispatcher +# --------------------------------------------------------------------------- +class _FakeGroup: + """Minimal process-group stand-in exposing only .size()/.rank(): the + dispatcher and resolve_cp_group() never issue a real collective.""" + + def __init__(self, size, rank=0): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def rank(self): + return self._rank + + +def _fake_packed_seq_params(qkv_format="thd", cp_group=None, local_cp_size=None): + params = PackedSeqParams(qkv_format=qkv_format) + if cp_group is not None: + params.cp_group = cp_group + if local_cp_size is not None: + params.local_cp_size = local_cp_size + return params + + +def _fake_gdn_module(*, linear_cp_mode, static_cp_size=1, deterministic_mode=False): + """Duck-typed GatedDeltaNet 'self': just enough attributes for the + dispatcher guards to read -- no real nn.Module/CUDA/FLA state.""" + return SimpleNamespace( + pg_collection=SimpleNamespace(cp=_FakeGroup(static_cp_size)), + cp_size=static_cp_size, + config=SimpleNamespace(linear_cp_mode=linear_cp_mode, deterministic_mode=deterministic_mode), + ) + + +@pytest.fixture(autouse=True) +def _isolate_gdn_forward_patch(): + """`_patch_gdn_for_dynamic_cp` idempotently monkey-patches the *shared* + GatedDeltaNet class attribute; save/restore it around every test so it + cannot leak into test_gdn_chunkwise_cp_gpu.py.""" + orig_forward = GatedDeltaNet.forward + orig_patched_flag = getattr(GatedDeltaNet, "_dcp_patched", False) + yield + GatedDeltaNet.forward = orig_forward + GatedDeltaNet._dcp_patched = orig_patched_flag + + +def _install_dispatcher_with_spies(): + """Install the dispatcher over a spy for the original MCore forward.""" + calls = {"orig": 0} + + def spy_orig(self, hidden_states, attention_mask, inference_context=None, packed_seq_params=None, *a, **kw): + calls["orig"] += 1 + return "orig", hidden_states + + GatedDeltaNet.forward = spy_orig + GatedDeltaNet._dcp_patched = False + gdn_model._patch_gdn_for_dynamic_cp() + return calls + + +def test_dispatcher_cp1_goes_to_original_forward_regardless_of_mode(): + for mode in ("headwise", "chunkwise", "all_gather"): + calls = _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode=mode, static_cp_size=1) + out = GatedDeltaNet.forward(m, "hs", None, None, None) + assert calls == {"orig": 1}, mode + assert out == ("orig", "hs") + + +@pytest.mark.parametrize("mode", ["headwise", "chunkwise"]) +def test_dispatcher_headwise_and_chunkwise_cp_gt_1_go_to_original_forward(mode): + calls = _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode=mode, static_cp_size=4) + psp = _fake_packed_seq_params(cp_group=_FakeGroup(4, rank=2), local_cp_size=4) + GatedDeltaNet.forward(m, "hs", None, None, psp) + assert calls == {"orig": 1} + + +def test_dispatcher_all_gather_cp_gt_1_goes_to_relax_fallback(): + calls = _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode="all_gather", static_cp_size=4) + psp = _fake_packed_seq_params(qkv_format="sbhd", cp_group=_FakeGroup(4, rank=3), local_cp_size=4) + with pytest.raises(AssertionError, match=r"packed \(thd\) sequences"): + GatedDeltaNet.forward(m, "hs", None, None, psp) + assert calls == {"orig": 0} + + +def test_dispatcher_prefers_dynamic_group_over_static_group(): + """Runtime CP (from packed_seq_params) must win over the module's static + max-CP group -- e.g. a static CP=8 model running a CP=1 micro-batch must + not take the all_gather branch just because the static group has size 8.""" + calls = _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode="all_gather", static_cp_size=8) + psp = _fake_packed_seq_params(cp_group=_FakeGroup(1, rank=0), local_cp_size=1) + GatedDeltaNet.forward(m, "hs", None, None, psp) + assert calls == {"orig": 1} + + +@pytest.mark.parametrize( + ("mode", "runtime_cp_size"), + [("headwise", 4), ("chunkwise", 4), ("all_gather", 1)], +) +def test_dispatcher_never_mutates_shared_module_or_config_state(mode, runtime_cp_size): + """Covers handoff §6.3: self.cp_size / self.pg_collection.cp / + self.config.linear_cp_mode must be bit-identical before and after, across + every mode and every runtime CP size (static max CP fixed at 8, so a + smaller runtime CP can only come from the dynamic packed_seq_params).""" + _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode=mode, static_cp_size=8) + psp = _fake_packed_seq_params(cp_group=_FakeGroup(runtime_cp_size, rank=0), local_cp_size=runtime_cp_size) + before_cp_size = m.cp_size + before_pg_cp = m.pg_collection.cp + before_mode = m.config.linear_cp_mode + GatedDeltaNet.forward(m, "hs", None, None, psp) + assert m.cp_size == before_cp_size + assert m.pg_collection.cp is before_pg_cp + assert m.config.linear_cp_mode == before_mode + + +# --------------------------------------------------------------------------- +# All-gather guard clauses (checked before any real tensor operation). +# --------------------------------------------------------------------------- +def test_all_gather_fallback_rejects_inference(): + _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode="all_gather", static_cp_size=4) + psp = _fake_packed_seq_params(cp_group=_FakeGroup(4), local_cp_size=4) + with pytest.raises(AssertionError, match="inference"): + GatedDeltaNet.forward(m, "hs", None, object(), psp) + + +def test_all_gather_fallback_rejects_deterministic_mode(): + _install_dispatcher_with_spies() + m = _fake_gdn_module(linear_cp_mode="all_gather", static_cp_size=4, deterministic_mode=True) + psp = _fake_packed_seq_params(cp_group=_FakeGroup(4), local_cp_size=4) + with pytest.raises(AssertionError, match="deterministic mode"): + GatedDeltaNet.forward(m, "hs", None, None, psp) From a3db7db2f4451dfc56bc3e021d9252db2090a9d9 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Thu, 13 Aug 2026 23:18:11 +1000 Subject: [PATCH 5/7] add cp chunkwise route --- .../patch/megatron/20260805-85bced0ae.patch | 719 ++++++++++++++---- .../megatron/test_gdn_chunkwise_cp_route.py | 231 ++++++ 2 files changed, 821 insertions(+), 129 deletions(-) create mode 100644 tests/backends/megatron/test_gdn_chunkwise_cp_route.py diff --git a/docker/patch/megatron/20260805-85bced0ae.patch b/docker/patch/megatron/20260805-85bced0ae.patch index 4456800a6..e99db1fe9 100644 --- a/docker/patch/megatron/20260805-85bced0ae.patch +++ b/docker/patch/megatron/20260805-85bced0ae.patch @@ -526,20 +526,38 @@ index c297a23..587e1de 100644 + return None diff --git a/megatron/core/context_parallel_layout.py b/megatron/core/context_parallel_layout.py new file mode 100644 -index 0000000..e3acdd9 +index 0000000..718b2fc --- /dev/null +++ b/megatron/core/context_parallel_layout.py -@@ -0,0 +1,307 @@ +@@ -0,0 +1,691 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Context parallel tensor layout helpers.""" + -+from typing import List, Optional, Tuple ++from contextlib import contextmanager ++from typing import Any, List, NamedTuple, Optional, Tuple + +import torch + +from megatron.core.tensor_parallel import all_to_all + ++_THD_CP_ROUTE_ATTRS = { ++ ("zigzag", "contiguous"): "cp_partition_route_zigzag_to_contiguous", ++ ("contiguous", "zigzag"): "cp_partition_route_contiguous_to_zigzag", ++} ++ ++ ++@contextmanager ++def _cp_layout_nvtx_range(message: str): ++ active = torch.cuda.is_available() ++ if active: ++ torch.cuda.nvtx.range_push(message) ++ try: ++ yield ++ finally: ++ if active: ++ torch.cuda.nvtx.range_pop() ++ + +def get_thd_context_parallel_rank_indices( + cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, layout: str @@ -614,11 +632,379 @@ index 0000000..e3acdd9 + return rank_positions[torch.argsort(rank_local_pos)] + + ++class ThdCpPartitionRoute(NamedTuple): ++ """Everything one THD zigzag<->contiguous all-to-all needs, precomputed. ++ ++ ``send_rows`` / ``recv_rows`` are ``None`` when the permutation on that side is ++ the identity, which lets the swap skip the gather/scatter entirely. ++ The first six fields are exactly what #5664's ``decode_thd_cp_partition_route()`` ++ returns; see :func:`build_thd_cp_partition_route` for why the encoded tensor form is ++ not used here. ``cu_seqlens``, the layout pair and ``cp_size`` / ``cp_rank`` are ++ Relax-only additions carrying no computation -- they exist so a route cached on a ++ ``PackedSeqParams`` can prove it still describes the conversion being asked for. ++ """ ++ ++ local_source_length: int ++ local_target_length: int ++ send_rows: Optional[torch.Tensor] ++ recv_rows: Optional[torch.Tensor] ++ input_split_sizes: List[int] ++ output_split_sizes: List[int] ++ cu_seqlens: torch.Tensor ++ cp_size: int ++ cp_rank: int ++ source_layout: str ++ target_layout: str ++ ++ ++_ThdLayoutSegment = Tuple[int, int, int] ++ ++ ++def _compact_thd_cu_seqlens_to_list(cu_seqlens: torch.Tensor) -> List[int]: ++ if cu_seqlens.dim() != 1: ++ raise ValueError(f"cu_seqlens must be 1-D, got shape {tuple(cu_seqlens.shape)}.") ++ ++ cu = cu_seqlens.detach().to(device="cpu", dtype=torch.long).tolist() ++ if not cu or cu[0] != 0: ++ raise ValueError(f"cu_seqlens must start at 0, got {cu_seqlens}.") ++ ++ compact_cu: List[int] = [cu[0]] ++ prev = cu[0] ++ for value in cu[1:]: ++ if value < prev: ++ raise ValueError(f"cu_seqlens must be nondecreasing, got {cu_seqlens}.") ++ if value != prev: ++ compact_cu.append(value) ++ prev = value ++ return compact_cu ++ ++ ++def _validate_thd_route_partitioning(cu: List[int], cp_size: int) -> None: ++ total_tokens = cu[-1] ++ if total_tokens % cp_size != 0: ++ raise ValueError( ++ f"Contiguous CP partitioning requires total_tokens={total_tokens} " ++ f"to be divisible by cp_size={cp_size}." ++ ) ++ ++ chunk_divisor = 2 * cp_size ++ bad_seq_lens = [ ++ seq_end - seq_start ++ for seq_start, seq_end in zip(cu[:-1], cu[1:]) ++ if (seq_end - seq_start) % chunk_divisor != 0 ++ ] ++ if bad_seq_lens: ++ raise ValueError( ++ "All packed sequence lengths must be divisible by " ++ f"2 * cp_size ({chunk_divisor}) for zigzag/contiguous CP layout conversion, " ++ f"got {bad_seq_lens}." ++ ) ++ ++ ++def _build_thd_layout_segments( ++ cu: List[int], cp_size: int, cp_rank: int, layout: str ++) -> Tuple[List[_ThdLayoutSegment], int]: ++ """Describe a rank's THD partition as (global_start, length, local_start) runs. ++ ++ Both layouts are unions of contiguous global spans, so the whole route can be ++ derived by intersecting spans instead of materialising per-token index tensors. ++ """ ++ total_tokens = cu[-1] ++ if layout == "contiguous": ++ part_len = total_tokens // cp_size ++ if part_len == 0: ++ return [], 0 ++ return [(cp_rank * part_len, part_len, 0)], part_len ++ ++ if layout != "zigzag": ++ raise ValueError( ++ f"Unsupported context-parallel layout {layout!r} for THD layout segments " ++ f"with cp_size={cp_size}, rank={cp_rank}." ++ ) ++ ++ segments: List[_ThdLayoutSegment] = [] ++ local_start = 0 ++ for seq_start, seq_end in zip(cu[:-1], cu[1:]): ++ seq_len = seq_end - seq_start ++ chunk_len = seq_len // (2 * cp_size) ++ first_chunk = cp_rank ++ second_chunk = 2 * cp_size - cp_rank - 1 ++ segments.append((seq_start + first_chunk * chunk_len, chunk_len, local_start)) ++ segments.append((seq_start + second_chunk * chunk_len, chunk_len, local_start + chunk_len)) ++ local_start += 2 * chunk_len ++ ++ return segments, local_start ++ ++ ++def _intersect_thd_layout_segments( ++ source_segments: List[_ThdLayoutSegment], target_segments: List[_ThdLayoutSegment] ++) -> List[Tuple[int, int, int]]: ++ """Overlap two sorted segment lists into (source_row, target_row, length) runs.""" ++ intersections: List[Tuple[int, int, int]] = [] ++ source_index = 0 ++ target_index = 0 ++ while source_index < len(source_segments) and target_index < len(target_segments): ++ source_global_start, source_len, source_local_start = source_segments[source_index] ++ target_global_start, target_len, target_local_start = target_segments[target_index] ++ source_global_end = source_global_start + source_len ++ target_global_end = target_global_start + target_len ++ ++ overlap_start = max(source_global_start, target_global_start) ++ overlap_end = min(source_global_end, target_global_end) ++ if overlap_start < overlap_end: ++ intersections.append( ++ ( ++ source_local_start + overlap_start - source_global_start, ++ target_local_start + overlap_start - target_global_start, ++ overlap_end - overlap_start, ++ ) ++ ) ++ ++ if source_global_end <= target_global_end: ++ source_index += 1 ++ else: ++ target_index += 1 ++ ++ return intersections ++ ++ ++def _append_range(rows: List[int], start: int, length: int) -> None: ++ rows.extend(range(start, start + length)) ++ ++ ++def _row_list_is_identity(rows: List[int]) -> bool: ++ return all(row == index for index, row in enumerate(rows)) ++ ++ ++def _thd_cp_partition_route_attr_name(source_layout: str, target_layout: str) -> str: ++ try: ++ return _THD_CP_ROUTE_ATTRS[(source_layout, target_layout)] ++ except KeyError as exc: ++ raise ValueError( ++ f"Unsupported CP layout conversion {source_layout!r} -> {target_layout!r} " ++ "for THD route." ++ ) from exc ++ ++ ++def build_thd_cp_partition_route( ++ cu_seqlens: torch.Tensor, ++ cp_size: int, ++ cp_rank: int, ++ source_layout: str, ++ target_layout: str, ++ *, ++ device: Optional[torch.device] = None, ++) -> ThdCpPartitionRoute: ++ """Precompute one THD CP layout conversion route. ++ ++ The route depends only on packed sequence metadata, CP rank/size and the ++ source/target layouts, so it can be reused by every tensor with the same THD ++ sequence axis in the same microbatch. The whole derivation runs on CPU ints ++ after a single ``cu_seqlens`` transfer, which is what keeps the conversion ++ itself free of device-host synchronisation. ++ ++ Deviation from #5664: upstream serialises the result into a single flat ++ ``torch.Tensor`` and decodes it again inside every conversion, so that the route ++ can be a CUDA-graph capture input. Decoding costs three device-to-host copies, ++ which on this path (~200 conversions per microbatch under full recompute) puts ++ back the synchronisation this route exists to remove, and Megatron's full-iteration ++ CUDA graph is rejected for THD layout conversion anyway. We therefore return the ++ decoded form directly -- the fields below are exactly upstream's ++ ``decode_thd_cp_partition_route()`` tuple. Revisit if GDN ever needs to run under ++ graph capture: the route tensors would then have to be stable capture inputs again. ++ """ ++ _thd_cp_partition_route_attr_name(source_layout, target_layout) ++ if device is None: ++ device = cu_seqlens.device ++ ++ with _cp_layout_nvtx_range(f"cp_layout/thd/route/{source_layout}_to_{target_layout}"): ++ cu = _compact_thd_cu_seqlens_to_list(cu_seqlens) ++ _validate_thd_route_partitioning(cu, cp_size) ++ ++ source_segments_by_rank: List[List[_ThdLayoutSegment]] = [] ++ source_lengths: List[int] = [] ++ target_segments_by_rank: List[List[_ThdLayoutSegment]] = [] ++ target_lengths: List[int] = [] ++ for rank in range(cp_size): ++ source_segments, source_length = _build_thd_layout_segments( ++ cu, cp_size, rank, source_layout ++ ) ++ target_segments, target_length = _build_thd_layout_segments( ++ cu, cp_size, rank, target_layout ++ ) ++ source_segments_by_rank.append(source_segments) ++ source_lengths.append(source_length) ++ target_segments_by_rank.append(target_segments) ++ target_lengths.append(target_length) ++ ++ local_source_segments = source_segments_by_rank[cp_rank] ++ local_target_segments = target_segments_by_rank[cp_rank] ++ ++ send_rows_list: List[int] = [] ++ input_split_sizes: List[int] = [] ++ for dst_rank in range(cp_size): ++ intersections = _intersect_thd_layout_segments( ++ local_source_segments, target_segments_by_rank[dst_rank] ++ ) ++ intersections.sort(key=lambda item: item[1]) ++ input_split_size = 0 ++ for source_row, _, length in intersections: ++ _append_range(send_rows_list, source_row, length) ++ input_split_size += length ++ input_split_sizes.append(input_split_size) ++ ++ recv_rows_list: List[int] = [] ++ output_split_sizes: List[int] = [] ++ for src_rank in range(cp_size): ++ intersections = _intersect_thd_layout_segments( ++ source_segments_by_rank[src_rank], local_target_segments ++ ) ++ intersections.sort(key=lambda item: item[1]) ++ output_split_size = 0 ++ for _, target_row, length in intersections: ++ _append_range(recv_rows_list, target_row, length) ++ output_split_size += length ++ output_split_sizes.append(output_split_size) ++ ++ assert len(send_rows_list) == source_lengths[cp_rank] ++ assert len(recv_rows_list) == target_lengths[cp_rank] ++ ++ send_rows = ( ++ None ++ if _row_list_is_identity(send_rows_list) ++ else torch.tensor(send_rows_list, device=device, dtype=torch.long) ++ ) ++ recv_rows = ( ++ None ++ if _row_list_is_identity(recv_rows_list) ++ else torch.tensor(recv_rows_list, device=device, dtype=torch.long) ++ ) ++ return ThdCpPartitionRoute( ++ local_source_length=source_lengths[cp_rank], ++ local_target_length=target_lengths[cp_rank], ++ send_rows=send_rows, ++ recv_rows=recv_rows, ++ input_split_sizes=input_split_sizes, ++ output_split_sizes=output_split_sizes, ++ cu_seqlens=cu_seqlens, ++ cp_size=cp_size, ++ cp_rank=cp_rank, ++ source_layout=source_layout, ++ target_layout=target_layout, ++ ) ++ ++ ++def _thd_cp_partition_route_is_reusable( ++ route: Optional[ThdCpPartitionRoute], ++ cu_seqlens: torch.Tensor, ++ cp_size: int, ++ cp_rank: int, ++ source_layout: str, ++ target_layout: str, ++ device: torch.device, ++) -> bool: ++ """A cached route is only valid for the exact packed boundaries it was built from. ++ ++ ``cu_seqlens`` is compared by object identity, not by value: a value comparison is ++ itself a device-host synchronisation, which is precisely what the route exists to ++ avoid. Identity is also the stronger check -- the same tensor object is the same ++ buffer, whereas equal values could still come from a different microbatch. ++ """ ++ if route is None: ++ return False ++ if route.cu_seqlens is not cu_seqlens: ++ return False ++ if route.cp_size != cp_size or route.cp_rank != cp_rank: ++ return False ++ if route.source_layout != source_layout or route.target_layout != target_layout: ++ return False ++ rows = route.send_rows if route.send_rows is not None else route.recv_rows ++ return rows is None or rows.device == device ++ ++ ++def get_thd_cp_partition_route( ++ packed_seq_params: Optional[Any], ++ cu_seqlens: torch.Tensor, ++ cp_size: int, ++ cp_rank: int, ++ source_layout: str, ++ target_layout: str, ++ *, ++ device: Optional[torch.device] = None, ++) -> ThdCpPartitionRoute: ++ """Return this microbatch's route, building and caching it on first use. ++ ++ Packed boundaries change every microbatch, so the cache is attached to the ++ ``PackedSeqParams`` instance the forward was handed and is only reused while it ++ still describes the very same ``cu_seqlens`` tensor. Every GDN layer in a ++ microbatch, plus its recompute replay, shares one build. ++ ++ Deviation from #5664: upstream expects the routes to have been prebuilt by the data ++ pipeline, so its lookup is a bare ``getattr`` and the build-on-miss path emits a ++ ``FutureWarning``. Here build-on-miss is the intended path -- the object a GDN ++ forward receives is not necessarily the one the data pipeline created, because the ++ Bridge/VLM path can repack ``PackedSeqParams`` after embedding -- so there is ++ nothing to warn about, and the cache instead has to defend itself against reuse ++ across microbatches and against dynamic CP changing ``cp_size``/``cp_rank`` between ++ microbatches on the same module. Callers that do own the final object can still ++ prebuild eagerly via :func:`prebuild_thd_cp_partition_routes`. ++ """ ++ if device is None: ++ device = cu_seqlens.device ++ attr_name = _thd_cp_partition_route_attr_name(source_layout, target_layout) ++ cached = getattr(packed_seq_params, attr_name, None) if packed_seq_params is not None else None ++ if _thd_cp_partition_route_is_reusable( ++ cached, cu_seqlens, cp_size, cp_rank, source_layout, target_layout, device ++ ): ++ return cached ++ ++ route = build_thd_cp_partition_route( ++ cu_seqlens, cp_size, cp_rank, source_layout, target_layout, device=device ++ ) ++ if packed_seq_params is not None: ++ setattr(packed_seq_params, attr_name, route) ++ return route ++ ++ ++def prebuild_thd_cp_partition_routes( ++ packed_seq_params: Optional[Any], ++ cp_group: Optional[torch.distributed.ProcessGroup] = None, ++ cu_seqlens: Optional[torch.Tensor] = None, ++ *, ++ device: Optional[torch.device] = None, ++) -> None: ++ """Eagerly populate both THD CP layout routes for a packed microbatch.""" ++ if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": ++ return ++ if cp_group is None: ++ cp_group = getattr(packed_seq_params, "cp_group", None) ++ if cp_group is None or cp_group.size() <= 1: ++ return ++ if cu_seqlens is None: ++ cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q_padded", None) ++ if cu_seqlens is None: ++ cu_seqlens = getattr(packed_seq_params, "cu_seqlens_q", None) ++ if cu_seqlens is None: ++ return ++ ++ for source_layout, target_layout in _THD_CP_ROUTE_ATTRS: ++ get_thd_cp_partition_route( ++ packed_seq_params, ++ cu_seqlens, ++ cp_group.size(), ++ cp_group.rank(), ++ source_layout, ++ target_layout, ++ device=device, ++ ) ++ ++ +def zigzag_to_contiguous_chunks( + x: torch.Tensor, + cp_group: torch.distributed.ProcessGroup, + seq_dim: int = 0, + cu_seqlens: Optional[torch.Tensor] = None, ++ thd_cp_partition_route: Optional[ThdCpPartitionRoute] = None, +) -> torch.Tensor: + """Permute CP chunks from Megatron zigzag layout to contiguous-time layout. + @@ -628,7 +1014,13 @@ index 0000000..e3acdd9 + """ + if cu_seqlens is not None: + return _zigzag_contiguous_thd_swap( -+ x, cp_group, seq_dim, cu_seqlens, source_layout="zigzag", target_layout="contiguous" ++ x, ++ cp_group, ++ seq_dim, ++ cu_seqlens, ++ source_layout="zigzag", ++ target_layout="contiguous", ++ thd_cp_partition_route=thd_cp_partition_route, + ) + return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=True) + @@ -638,15 +1030,43 @@ index 0000000..e3acdd9 + cp_group: torch.distributed.ProcessGroup, + seq_dim: int = 0, + cu_seqlens: Optional[torch.Tensor] = None, ++ thd_cp_partition_route: Optional[ThdCpPartitionRoute] = None, +) -> torch.Tensor: + """Inverse of :func:`zigzag_to_contiguous_chunks`.""" + if cu_seqlens is not None: + return _zigzag_contiguous_thd_swap( -+ x, cp_group, seq_dim, cu_seqlens, source_layout="contiguous", target_layout="zigzag" ++ x, ++ cp_group, ++ seq_dim, ++ cu_seqlens, ++ source_layout="contiguous", ++ target_layout="zigzag", ++ thd_cp_partition_route=thd_cp_partition_route, + ) + return _zigzag_contiguous_chunk_swap(x, cp_group, seq_dim, to_contiguous=False) + + ++def _pack_thd_cp_route_send_buffer( ++ x: torch.Tensor, local_source_length: int, send_rows: Optional[torch.Tensor] ++) -> torch.Tensor: ++ if local_source_length == 0: ++ return x.narrow(0, 0, 0) ++ if send_rows is None: ++ return x ++ return x.index_select(0, send_rows) ++ ++ ++def _scatter_thd_cp_route_recv_buffer( ++ recv_buf: torch.Tensor, recv_rows: Optional[torch.Tensor], out_shape: Tuple[int, ...] ++) -> torch.Tensor: ++ if recv_rows is None: ++ return recv_buf ++ out = recv_buf.new_empty(out_shape) ++ if recv_rows.numel() > 0: ++ out.index_copy_(0, recv_rows, recv_buf) ++ return out ++ ++ +def _zigzag_contiguous_thd_swap( + x: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], @@ -654,95 +1074,59 @@ index 0000000..e3acdd9 + cu_seqlens: torch.Tensor, + source_layout: str, + target_layout: str, ++ thd_cp_partition_route: Optional[ThdCpPartitionRoute] = None, +) -> torch.Tensor: + """Single-all-to-all THD permutation between zigzag and contiguous layouts. + + The packed THD tensor stays packed: we first group local tokens by their + target CP rank, exchange those groups once, then scatter received tokens -+ back into the target rank-local order. ++ back into the target rank-local order. Which rows go where is described by a ++ :class:`ThdCpPartitionRoute` the caller should have precomputed once for the ++ microbatch; without one this rebuilds it, which is correct but pays the build ++ on every conversion. + """ + cp_size = cp_group.size() if cp_group is not None else 1 + if cp_size == 1: + return x + cp_rank = cp_group.rank() + -+ if seq_dim != 0: -+ x = x.movedim(seq_dim, 0) -+ x = x.contiguous() -+ -+ cu = cu_seqlens.to(device=x.device, dtype=torch.long) -+ # TODO: Let a future CP layout scheduler precompute this routing once per -+ # microbatch from immutable cu_seqlens and pass it through both THD swaps. -+ # Do not cache it across microbatches because packed sequence boundaries change. -+ source_by_rank = [ -+ get_thd_context_parallel_rank_indices(cu, cp_size, rank, source_layout) -+ for rank in range(cp_size) -+ ] -+ target_by_rank = [ -+ get_thd_context_parallel_rank_indices(cu, cp_size, rank, target_layout) -+ for rank in range(cp_size) -+ ] ++ conversion_name = f"{source_layout}_to_{target_layout}" ++ with _cp_layout_nvtx_range(f"cp_layout/thd/swap/{conversion_name}"): ++ if seq_dim != 0: ++ x = x.movedim(seq_dim, 0) ++ x = x.contiguous() ++ ++ route = thd_cp_partition_route ++ if not _thd_cp_partition_route_is_reusable( ++ route, cu_seqlens, cp_size, cp_rank, source_layout, target_layout, x.device ++ ): ++ route = build_thd_cp_partition_route( ++ cu_seqlens, cp_size, cp_rank, source_layout, target_layout, device=x.device ++ ) + -+ local_source_indices = source_by_rank[cp_rank] -+ local_target_indices = target_by_rank[cp_rank] -+ if x.size(0) != local_source_indices.numel(): -+ raise ValueError( -+ f"Local THD tensor length ({x.size(0)}) does not match {source_layout} " -+ f"rank-{cp_rank} partition length ({local_source_indices.numel()})." -+ ) ++ if x.size(0) != route.local_source_length: ++ raise ValueError( ++ f"Local THD tensor length ({x.size(0)}) does not match {source_layout} " ++ f"rank-{cp_rank} partition length ({route.local_source_length})." ++ ) + -+ total_tokens = int(cu[-1].item()) -+ target_owner = torch.empty(total_tokens, device=x.device, dtype=torch.long) -+ target_local_pos = torch.empty(total_tokens, device=x.device, dtype=torch.long) -+ for rank, indices in enumerate(target_by_rank): -+ target_owner[indices] = rank -+ target_local_pos[indices] = torch.arange(indices.numel(), device=x.device) -+ -+ local_target_owner = target_owner[local_source_indices] -+ local_target_pos = target_local_pos[local_source_indices] -+ -+ send_parts: List[torch.Tensor] = [] -+ input_split_sizes: List[int] = [] -+ for dst_rank in range(cp_size): -+ dst_mask = local_target_owner == dst_rank -+ dst_rows = dst_mask.nonzero(as_tuple=False).flatten() -+ if dst_rows.numel() > 0: -+ dst_rows = dst_rows[torch.argsort(local_target_pos[dst_rows])] -+ send_part = x.index_select(0, dst_rows) -+ else: -+ send_part = x.narrow(0, 0, 0) -+ send_parts.append(send_part) -+ input_split_sizes.append(send_part.size(0)) -+ send_buf = torch.cat(send_parts, dim=0).contiguous() -+ -+ output_split_sizes: List[int] = [] -+ recv_target_positions: List[torch.Tensor] = [] -+ for src_rank in range(cp_size): -+ src_indices = source_by_rank[src_rank] -+ src_to_this_rank = target_owner[src_indices] == cp_rank -+ recv_global_indices = src_indices[src_to_this_rank] -+ if recv_global_indices.numel() > 0: -+ recv_positions = target_local_pos[recv_global_indices] -+ recv_positions = recv_positions[torch.argsort(recv_positions)] -+ else: -+ recv_positions = local_target_indices.narrow(0, 0, 0) -+ recv_target_positions.append(recv_positions) -+ output_split_sizes.append(recv_positions.numel()) ++ with _cp_layout_nvtx_range(f"cp_layout/thd/pack/{conversion_name}"): ++ send_buf = _pack_thd_cp_route_send_buffer(x, route.local_source_length, route.send_rows) ++ if not send_buf.is_contiguous(): ++ send_buf = send_buf.contiguous() + -+ recv_buf = all_to_all(cp_group, send_buf, output_split_sizes, input_split_sizes) ++ with _cp_layout_nvtx_range(f"cp_layout/thd/all_to_all/{conversion_name}"): ++ recv_buf = all_to_all( ++ cp_group, send_buf, route.output_split_sizes, route.input_split_sizes ++ ) + -+ out_shape = (local_target_indices.numel(),) + tuple(x.shape[1:]) -+ out = x.new_empty(out_shape) -+ offset = 0 -+ for recv_positions in recv_target_positions: -+ recv_len = recv_positions.numel() -+ if recv_len > 0: -+ out[recv_positions] = recv_buf[offset : offset + recv_len] -+ offset += recv_len ++ with _cp_layout_nvtx_range(f"cp_layout/thd/scatter/{conversion_name}"): ++ out_shape = (route.local_target_length,) + tuple(x.shape[1:]) ++ out = _scatter_thd_cp_route_recv_buffer(recv_buf, route.recv_rows, out_shape) + -+ if seq_dim != 0: -+ out = out.movedim(0, seq_dim) -+ return out.contiguous() ++ if seq_dim != 0: ++ out = out.movedim(0, seq_dim) ++ return out.contiguous() + + +def _zigzag_contiguous_chunk_swap( @@ -1318,7 +1702,7 @@ index 4430a8c..084389b 100644 tensors[key].shape, gbuf_local_start, diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py -index 322f12a..0be9e5e 100644 +index 322f12a..21f43ca 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -64,3 +64,17 @@ class PackedSeqParams: @@ -1438,15 +1822,16 @@ index 465e83f..232caef 100644 tensor_recv_prev = None tensor_recv_next = None diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py -index 8df4df1..f39eff8 100644 +index 8df4df1..66c4d1e 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py -@@ -14,12 +14,16 @@ import torch.nn as nn +@@ -14,12 +14,17 @@ import torch.nn as nn import torch.nn.functional as F from torch import Tensor +from megatron.core.context_parallel_layout import ( + contiguous_to_zigzag_chunks, ++ get_thd_cp_partition_route, + zigzag_to_contiguous_chunks, +) from megatron.core.dist_checkpointing import ShardedTensor @@ -1459,14 +1844,15 @@ index 8df4df1..f39eff8 100644 from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.mamba_context_parallel import ( _all_to_all_cp2hp, -@@ -43,5 +47,6 @@ try: +@@ -42,6 +47,7 @@ from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx + try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True -@@ -233,6 +238,15 @@ class GatedDeltaNet(MegatronModule): +@@ -233,6 +239,15 @@ class GatedDeltaNet(MegatronModule): tp_group=self.pg_collection.tp, ) @@ -1482,7 +1868,7 @@ index 8df4df1..f39eff8 100644 self.reset_parameters() def reset_parameters(self): -@@ -265,6 +279,7 @@ class GatedDeltaNet(MegatronModule): +@@ -265,6 +280,7 @@ class GatedDeltaNet(MegatronModule): packed_seq_params: Optional[PackedSeqParams] = None, sequence_len_offset: Optional[int] = None, *, @@ -1490,14 +1876,10 @@ index 8df4df1..f39eff8 100644 inference_params: Optional[BaseInferenceContext] = None, **kwargs, ): -@@ -286,10 +301,65 @@ class GatedDeltaNet(MegatronModule): - """ - # TODO: Deal with attention_mask -- -+ +@@ -288,8 +304,63 @@ class GatedDeltaNet(MegatronModule): + inference_context = deprecate_inference_params(inference_context, inference_params) -- -+ + + # Bridge repacking may replace PackedSeqParams. Validate the object received + # by every GDN forward; group.size() is local metadata and needs no collective. + if packed_seq_params is not None: @@ -1559,7 +1941,7 @@ index 8df4df1..f39eff8 100644 if inference_context is not None: assert ( -@@ -299,24 +369,30 @@ class GatedDeltaNet(MegatronModule): +@@ -299,53 +370,103 @@ class GatedDeltaNet(MegatronModule): # TODO: support inference raise NotImplementedError("GDN does not support inference for now.") @@ -1572,29 +1954,34 @@ index 8df4df1..f39eff8 100644 ), "Packed sequence does not support deterministic mode." - # Resolve cu_seqlens with alignment padding handling. +- cu_seqlens_q = self._resolve_cu_seqlens( +- packed_seq_params.cu_seqlens_q_padded, +- packed_seq_params.cu_seqlens_q, +- seq_len, +- "cu_seqlens_q", +- ) +- cu_seqlens_kv = self._resolve_cu_seqlens( +- packed_seq_params.cu_seqlens_kv_padded, +- packed_seq_params.cu_seqlens_kv, +- seq_len, +- "cu_seqlens_kv", +- ) +- assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( +- "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " +- f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" +- ) +- num_packed_seqs = cu_seqlens_q.shape[0] - 1 +- assert num_packed_seqs > 0, ( +- "Number of packed sequences must be greater than 0, " +- f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + # Resolve cu_seqlens with alignment padding handling. These are the *global* + # (pre-CP-split) boundaries in both CP modes, taken from the PackedSeqParams + # this forward was handed -- for the Bridge/VLM unsplit path that is the + # post-embedding, post-repack object, so no stale boundary can leak in. - cu_seqlens_q = self._resolve_cu_seqlens( - packed_seq_params.cu_seqlens_q_padded, - packed_seq_params.cu_seqlens_q, -- seq_len, -+ seq_len_global, - "cu_seqlens_q", -+ cp_size=cp_size, ++ cu_seqlens_q, cu_seqlens_kv = self._resolve_thd_cu_seqlens( ++ packed_seq_params, seq_len_global, cp_size ) - cu_seqlens_kv = self._resolve_cu_seqlens( - packed_seq_params.cu_seqlens_kv_padded, - packed_seq_params.cu_seqlens_kv, -- seq_len, -+ seq_len_global, - "cu_seqlens_kv", -+ cp_size=cp_size, - ) - assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( - "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " -@@ -331,21 +407,59 @@ class GatedDeltaNet(MegatronModule): + else: cu_seqlens_q = None cu_seqlens_kv = None @@ -1619,6 +2006,31 @@ index 8df4df1..f39eff8 100644 + cu_seqlens_q, chunkwise_cp_context = self._build_chunkwise_cp_context( + cp_group_chunkwise, cp_size_chunkwise, cu_seqlens_q, seq_len_global, batch + ) ++ ++ # Both zigzag <-> contiguous conversions below are pure functions of the packed ++ # boundaries and the CP group, so resolve them once here and let every GDN layer ++ # of this micro-batch (and its recompute replay) reuse the same routes. ++ route_to_contiguous = None ++ route_to_zigzag = None ++ if cp_size_chunkwise > 1 and is_thd: ++ route_to_contiguous = get_thd_cp_partition_route( ++ packed_seq_params, ++ cu_seqlens_q, ++ cp_size_chunkwise, ++ cp_group_chunkwise.rank(), ++ "zigzag", ++ "contiguous", ++ device=hidden_states.device, ++ ) ++ route_to_zigzag = get_thd_cp_partition_route( ++ packed_seq_params, ++ cu_seqlens_q, ++ cp_size_chunkwise, ++ cp_group_chunkwise.rank(), ++ "contiguous", ++ "zigzag", ++ device=hidden_states.device, ++ ) + # Input projection nvtx_range_push(suffix="in_proj") @@ -1638,6 +2050,7 @@ index 8df4df1..f39eff8 100644 + cp_group_chunkwise, + seq_dim=0, + cu_seqlens=cu_seqlens_q if is_thd else None, ++ thd_cp_partition_route=route_to_contiguous, + ) + nvtx_range_pop(suffix="zigzag_to_contiguous") + @@ -1657,7 +2070,7 @@ index 8df4df1..f39eff8 100644 split_sections=[ self.qk_dim_local_tp, self.qk_dim_local_tp, -@@ -357,12 +471,12 @@ class GatedDeltaNet(MegatronModule): +@@ -357,12 +478,12 @@ class GatedDeltaNet(MegatronModule): ) outputs.append(qkvzba_i) qkvzba = torch.cat(outputs, dim=0) @@ -1672,7 +2085,7 @@ index 8df4df1..f39eff8 100644 split_sections=[ self.qk_dim_local_tp, self.qk_dim_local_tp, -@@ -381,10 +495,10 @@ class GatedDeltaNet(MegatronModule): +@@ -381,10 +502,10 @@ class GatedDeltaNet(MegatronModule): qkv, gate, beta, alpha = torch.split( qkvzba, [ @@ -1687,7 +2100,7 @@ index 8df4df1..f39eff8 100644 ], dim=-1, ) -@@ -403,14 +517,14 @@ class GatedDeltaNet(MegatronModule): +@@ -403,14 +524,14 @@ class GatedDeltaNet(MegatronModule): conv1d_weight = get_parameter_local_cp( self.conv1d.weight, dim=0, @@ -1704,7 +2117,7 @@ index 8df4df1..f39eff8 100644 split_sections=qkv_channels_split_sections, ) if self.conv_bias -@@ -425,7 +539,7 @@ class GatedDeltaNet(MegatronModule): +@@ -425,7 +546,7 @@ class GatedDeltaNet(MegatronModule): stride=self.conv1d.stride, padding=self.conv1d.padding, dilation=self.conv1d.dilation, @@ -1713,7 +2126,7 @@ index 8df4df1..f39eff8 100644 ) qkv = self.act_fn(conv_out[..., :seq_len]) qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d -@@ -439,22 +553,22 @@ class GatedDeltaNet(MegatronModule): +@@ -439,22 +560,22 @@ class GatedDeltaNet(MegatronModule): initial_state=None, output_final_state=False, cu_seqlens=cu_seqlens_q, @@ -1743,7 +2156,7 @@ index 8df4df1..f39eff8 100644 g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") -@@ -469,6 +583,7 @@ class GatedDeltaNet(MegatronModule): +@@ -469,6 +590,7 @@ class GatedDeltaNet(MegatronModule): output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=cu_seqlens_q, @@ -1751,7 +2164,7 @@ index 8df4df1..f39eff8 100644 ) nvtx_range_pop(suffix="gated_delta_rule") -@@ -482,19 +597,31 @@ class GatedDeltaNet(MegatronModule): +@@ -482,19 +604,32 @@ class GatedDeltaNet(MegatronModule): norm_out = norm_out.reshape(batch, seq_len, -1) norm_out = norm_out.transpose(0, 1).contiguous() @@ -1764,6 +2177,7 @@ index 8df4df1..f39eff8 100644 + cp_group=cp_group_chunkwise, + seq_dim=0, + cu_seqlens=cu_seqlens_q if is_thd else None, ++ thd_cp_partition_route=route_to_zigzag, + ) + nvtx_range_pop(suffix="contiguous_to_zigzag") + @@ -1787,7 +2201,7 @@ index 8df4df1..f39eff8 100644 ) # Output projection -@@ -504,6 +631,48 @@ class GatedDeltaNet(MegatronModule): +@@ -504,6 +639,48 @@ class GatedDeltaNet(MegatronModule): return out, out_bias @@ -1836,7 +2250,7 @@ index 8df4df1..f39eff8 100644 @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm -@@ -517,15 +686,23 @@ class GatedDeltaNet(MegatronModule): +@@ -517,15 +694,23 @@ class GatedDeltaNet(MegatronModule): return y @jit_fuser @@ -1862,7 +2276,7 @@ index 8df4df1..f39eff8 100644 dim=-1, ) -@@ -538,7 +715,7 @@ class GatedDeltaNet(MegatronModule): +@@ -538,7 +723,7 @@ class GatedDeltaNet(MegatronModule): query_key = l2norm(query_key.contiguous()) # Split query and key @@ -1871,18 +2285,67 @@ index 8df4df1..f39eff8 100644 query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) -@@ -567,7 +744,9 @@ class GatedDeltaNet(MegatronModule): +@@ -567,7 +752,58 @@ class GatedDeltaNet(MegatronModule): beta = beta.sigmoid() return g, beta - def _resolve_cu_seqlens(self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name): ++ def _resolve_thd_cu_seqlens(self, packed_seq_params, seq_len_global, cp_size): ++ """Resolve and validate this micro-batch's global THD boundaries, once. ++ ++ Every GDN layer is handed the same ``PackedSeqParams``, so the checks below -- ++ each of which reads a device tensor from Python and therefore stalls on the ++ GPU -- only need to run for the first layer that asks. The result is cached on ++ the ``PackedSeqParams`` instance and reused while it still describes the very ++ same boundary tensors, mirroring how the CP layout routes are cached. ++ """ ++ sources = ( ++ packed_seq_params.cu_seqlens_q_padded, ++ packed_seq_params.cu_seqlens_q, ++ packed_seq_params.cu_seqlens_kv_padded, ++ packed_seq_params.cu_seqlens_kv, ++ ) ++ cached = getattr(packed_seq_params, "_gdn_resolved_cu_seqlens", None) ++ if cached is not None: ++ cached_key, cached_value = cached ++ if ( ++ cached_key[0] == seq_len_global ++ and cached_key[1] == cp_size ++ and len(cached_key[2]) == len(sources) ++ and all(a is b for a, b in zip(cached_key[2], sources)) ++ ): ++ return cached_value ++ ++ cu_seqlens_q = self._resolve_cu_seqlens( ++ sources[0], sources[1], seq_len_global, "cu_seqlens_q", cp_size=cp_size ++ ) ++ cu_seqlens_kv = self._resolve_cu_seqlens( ++ sources[2], sources[3], seq_len_global, "cu_seqlens_kv", cp_size=cp_size ++ ) ++ assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( ++ "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " ++ f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" ++ ) ++ num_packed_seqs = cu_seqlens_q.shape[0] - 1 ++ assert num_packed_seqs > 0, ( ++ "Number of packed sequences must be greater than 0, " ++ f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" ++ ) ++ ++ resolved = (cu_seqlens_q, cu_seqlens_kv) ++ packed_seq_params._gdn_resolved_cu_seqlens = ( ++ (seq_len_global, cp_size, sources), ++ resolved, ++ ) ++ return resolved ++ + def _resolve_cu_seqlens( + self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1 + ): """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" if cu_seqlens_padded is not None: cu_seqlens = cu_seqlens_padded -@@ -582,6 +761,13 @@ class GatedDeltaNet(MegatronModule): +@@ -582,6 +818,13 @@ class GatedDeltaNet(MegatronModule): f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." ) @@ -1896,7 +2359,7 @@ index 8df4df1..f39eff8 100644 return cu_seqlens def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): -@@ -780,8 +966,10 @@ def get_parameter_local_cp( +@@ -780,8 +1023,10 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ @@ -1909,7 +2372,7 @@ index 8df4df1..f39eff8 100644 # No need to split if CP size is 1. if cp_size == 1: -@@ -800,7 +988,7 @@ def get_parameter_local_cp( +@@ -800,7 +1045,7 @@ def get_parameter_local_cp( slices = [slice(None)] * param.dim() dim_size = param.size(dim=dim) slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) @@ -1918,8 +2381,7 @@ index 8df4df1..f39eff8 100644 return param -@@ -935,7 +1123,8 @@ def torch_chunk_gated_delta_rule( - initial_state=None, +@@ -945,6 +1190,7 @@ def torch_chunk_gated_delta_rule( output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, @@ -1927,15 +2389,14 @@ index 8df4df1..f39eff8 100644 ): # pylint: disable=line-too-long ''' -@@ -948,6 +1137,9 @@ def torch_chunk_gated_delta_rule( +@@ -957,6 +1203,9 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + assert ( + cp_context is None + ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." -- -+ + initial_dtype = query.dtype if use_qk_l2norm_in_kernel: diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_route.py b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py new file mode 100644 index 000000000..6d426c117 --- /dev/null +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +"""Unit tests for the prebuilt THD CP layout route (Task 32, phase 3). + +Phase 1 derived the zigzag<->contiguous all-to-all plan inside every conversion, +from device tensors, which cost a device-host synchronisation per CP rank per +call. Phase 3 backports NVIDIA/Megatron-LM#5664's idea instead: derive the plan +once per micro-batch on CPU and hand the same route to every GDN layer. + +Two things therefore need proving on CPU, with no process group: + +1. the segment-based route describes *exactly* the permutation the phase-1 + index-based partition described -- otherwise chunkwise CP silently reorders + tokens; +2. a cached route is only ever reused for the micro-batch, CP geometry and + direction it was built for. + +The real all-to-all round trip over NCCL stays in +``test_gdn_chunkwise_cp_gpu.py``. +""" + +from __future__ import annotations + +import pytest +import torch + + +cpl = pytest.importorskip("megatron.core.context_parallel_layout", reason="requires the patched Megatron-LM") + +from megatron.core.packed_seq_params import PackedSeqParams # noqa: E402 + + +DIRECTIONS = [("zigzag", "contiguous"), ("contiguous", "zigzag")] + +# Packed boundary shapes worth covering: single sequence, uneven multi-sequence, +# and a duplicated boundary (an empty padding slot), which the compaction step +# has to drop before the segments line up. +LENGTH_CASES = [ + [1], + [3, 1, 2], + [2, 0, 1, 3], +] + + +def _cu(lengths: list[int], unit: int) -> torch.Tensor: + cu = [0] + for n in lengths: + cu.append(cu[-1] + n * unit) + return torch.tensor(cu, dtype=torch.int64) + + +def _packed_seq_params(cu: torch.Tensor) -> PackedSeqParams: + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=int(cu[-1]), + max_seqlen_kv=int(cu[-1]), + ) + + +def _apply_route_across_ranks( + x: torch.Tensor, cu: torch.Tensor, cp_size: int, source: str, target: str +) -> list[torch.Tensor]: + """Run the route-driven swap for every rank, emulating the all-to-all locally.""" + source_by_rank = [ + cpl.get_thd_context_parallel_rank_indices(cu, cp_size, r, source) for r in range(cp_size) + ] + routes = [ + cpl.build_thd_cp_partition_route(cu, cp_size, r, source, target) for r in range(cp_size) + ] + + send_bufs = [] + for rank, route in enumerate(routes): + local = x[source_by_rank[rank]] + assert local.size(0) == route.local_source_length + send_bufs.append(local if route.send_rows is None else local.index_select(0, route.send_rows)) + + outputs = [] + for dst, route in enumerate(routes): + parts = [] + for src in range(cp_size): + offset = sum(routes[src].input_split_sizes[:dst]) + length = routes[src].input_split_sizes[dst] + assert length == route.output_split_sizes[src], "split sizes disagree between peers" + parts.append(send_bufs[src][offset : offset + length]) + recv = torch.cat(parts, dim=0) + if route.recv_rows is None: + outputs.append(recv) + else: + out = recv.new_empty((route.local_target_length,) + tuple(x.shape[1:])) + out.index_copy_(0, route.recv_rows, recv) + outputs.append(out) + return outputs + + +# --------------------------------------------------------------------------- +# The route is the same permutation phase 1 computed +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("cp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("source,target", DIRECTIONS) +@pytest.mark.parametrize("lengths", LENGTH_CASES) +def test_route_reproduces_index_based_partition(cp_size, source, target, lengths): + cu = _cu(lengths, unit=2 * cp_size) + total = int(cu[-1]) + x = torch.arange(total * 3, dtype=torch.float64).reshape(total, 3) + + got = _apply_route_across_ranks(x, cu, cp_size, source, target) + for rank in range(cp_size): + want = x[cpl.get_thd_context_parallel_rank_indices(cu, cp_size, rank, target)] + assert torch.equal(got[rank], want), f"cp_size={cp_size} rank={rank} {source}->{target}" + + +# --------------------------------------------------------------------------- +# Fail-fast parity with the index-based builder +# --------------------------------------------------------------------------- +def test_route_rejects_lengths_not_divisible_by_two_cp(): + cu = torch.tensor([0, 12], dtype=torch.int64) # 12 % (2 * 4) != 0 + with pytest.raises(ValueError, match="divisible by"): + cpl.get_thd_context_parallel_rank_indices(cu, 4, 0, "zigzag") + with pytest.raises(ValueError, match="divisible by"): + cpl.build_thd_cp_partition_route(cu, 4, 0, "zigzag", "contiguous") + + +def test_route_rejects_malformed_cu_seqlens(): + with pytest.raises(ValueError, match="must start at 0"): + cpl.build_thd_cp_partition_route( + torch.tensor([8, 16], dtype=torch.int64), 2, 0, "zigzag", "contiguous" + ) + with pytest.raises(ValueError, match="nondecreasing"): + cpl.build_thd_cp_partition_route( + torch.tensor([0, 16, 8], dtype=torch.int64), 2, 0, "zigzag", "contiguous" + ) + + +def test_route_rejects_unknown_layout(): + cu = _cu([1], unit=4) + with pytest.raises(ValueError, match="Unsupported CP layout conversion"): + cpl.build_thd_cp_partition_route(cu, 2, 0, "zigzag", "interleaved") + + +# --------------------------------------------------------------------------- +# Caching: reuse only within the micro-batch it was built for +# --------------------------------------------------------------------------- +def test_route_is_cached_per_packed_seq_params(): + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + + first = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + second = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + assert second is first, "a second layer of the same micro-batch must reuse the route" + assert psp.cp_partition_route_zigzag_to_contiguous is first + + +def test_both_directions_are_cached_separately(): + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + + to_contiguous = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + to_zigzag = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "contiguous", "zigzag") + assert to_contiguous is not to_zigzag + assert psp.cp_partition_route_zigzag_to_contiguous is to_contiguous + assert psp.cp_partition_route_contiguous_to_zigzag is to_zigzag + + +def test_route_is_rebuilt_for_new_packed_boundaries(): + """Packed boundaries move every micro-batch; a stale route would corrupt tokens.""" + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + first = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + + next_cu = _cu([2, 2], unit=8) + psp.cu_seqlens_q = next_cu + rebuilt = cpl.get_thd_cp_partition_route(psp, next_cu, 4, 1, "zigzag", "contiguous") + assert rebuilt is not first + assert rebuilt.cu_seqlens is next_cu + + want = cpl.build_thd_cp_partition_route(next_cu, 4, 1, "zigzag", "contiguous") + assert rebuilt.input_split_sizes == want.input_split_sizes + assert rebuilt.output_split_sizes == want.output_split_sizes + for field in ("send_rows", "recv_rows"): + got_rows, want_rows = getattr(rebuilt, field), getattr(want, field) + assert (got_rows is None) == (want_rows is None) + if want_rows is not None: + assert torch.equal(got_rows, want_rows) + + +def test_route_is_rebuilt_when_the_dynamic_cp_geometry_changes(): + """Dynamic CP varies cp_size/cp_rank across micro-batches on one module.""" + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + cp4 = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + + cp2 = cpl.get_thd_cp_partition_route(psp, cu, 2, 1, "zigzag", "contiguous") + assert cp2 is not cp4 + assert (cp2.cp_size, cp2.cp_rank) == (2, 1) + + other_rank = cpl.get_thd_cp_partition_route(psp, cu, 2, 0, "zigzag", "contiguous") + assert other_rank is not cp2 + assert other_rank.cp_rank == 0 + + +def test_prebuild_populates_both_directions(): + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + + class _FakeGroup: + def size(self): + return 4 + + def rank(self): + return 2 + + psp.cp_group = _FakeGroup() + psp.local_cp_size = 4 + cpl.prebuild_thd_cp_partition_routes(psp) + + for attr in ("cp_partition_route_zigzag_to_contiguous", "cp_partition_route_contiguous_to_zigzag"): + route = getattr(psp, attr) + assert route is not None + assert (route.cp_size, route.cp_rank) == (4, 2) + + +def test_prebuild_is_a_noop_without_context_parallelism(): + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + cpl.prebuild_thd_cp_partition_routes(psp) + assert getattr(psp, "cp_partition_route_zigzag_to_contiguous", None) is None + + non_thd = PackedSeqParams(qkv_format="sbhd") + cpl.prebuild_thd_cp_partition_routes(non_thd) + assert getattr(non_thd, "cp_partition_route_zigzag_to_contiguous", None) is None From 6ac74648d2d4ee81d25a1813ccbb574f6c1c70ce Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Thu, 13 Aug 2026 23:26:54 +1000 Subject: [PATCH 6/7] style: apply ruff-format and docformatter to the CP route tests Pre-commit was not run before pushing the previous commit, so CI's ruff-format and docformatter hooks failed on the new test file. Formatting only -- no test logic changed. Co-authored-by: Cursor --- .../megatron/test_gdn_chunkwise_cp_route.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_route.py b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py index 6d426c117..886a88a07 100644 --- a/tests/backends/megatron/test_gdn_chunkwise_cp_route.py +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py @@ -61,13 +61,10 @@ def _packed_seq_params(cu: torch.Tensor) -> PackedSeqParams: def _apply_route_across_ranks( x: torch.Tensor, cu: torch.Tensor, cp_size: int, source: str, target: str ) -> list[torch.Tensor]: - """Run the route-driven swap for every rank, emulating the all-to-all locally.""" - source_by_rank = [ - cpl.get_thd_context_parallel_rank_indices(cu, cp_size, r, source) for r in range(cp_size) - ] - routes = [ - cpl.build_thd_cp_partition_route(cu, cp_size, r, source, target) for r in range(cp_size) - ] + """Run the route-driven swap for every rank, emulating the all-to-all + locally.""" + source_by_rank = [cpl.get_thd_context_parallel_rank_indices(cu, cp_size, r, source) for r in range(cp_size)] + routes = [cpl.build_thd_cp_partition_route(cu, cp_size, r, source, target) for r in range(cp_size)] send_bufs = [] for rank, route in enumerate(routes): @@ -123,13 +120,9 @@ def test_route_rejects_lengths_not_divisible_by_two_cp(): def test_route_rejects_malformed_cu_seqlens(): with pytest.raises(ValueError, match="must start at 0"): - cpl.build_thd_cp_partition_route( - torch.tensor([8, 16], dtype=torch.int64), 2, 0, "zigzag", "contiguous" - ) + cpl.build_thd_cp_partition_route(torch.tensor([8, 16], dtype=torch.int64), 2, 0, "zigzag", "contiguous") with pytest.raises(ValueError, match="nondecreasing"): - cpl.build_thd_cp_partition_route( - torch.tensor([0, 16, 8], dtype=torch.int64), 2, 0, "zigzag", "contiguous" - ) + cpl.build_thd_cp_partition_route(torch.tensor([0, 16, 8], dtype=torch.int64), 2, 0, "zigzag", "contiguous") def test_route_rejects_unknown_layout(): @@ -163,7 +156,8 @@ def test_both_directions_are_cached_separately(): def test_route_is_rebuilt_for_new_packed_boundaries(): - """Packed boundaries move every micro-batch; a stale route would corrupt tokens.""" + """Packed boundaries move every micro-batch; a stale route would corrupt + tokens.""" cu = _cu([3, 1], unit=8) psp = _packed_seq_params(cu) first = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") From de5a26f81bf9dc02531b81a4d05ca5ab727c4122 Mon Sep 17 00:00:00 2001 From: jambow0320 <1193785411@qq.com> Date: Fri, 14 Aug 2026 14:18:03 +1000 Subject: [PATCH 7/7] guard cp route cache against in-place cu_seqlens writes --- .../patch/megatron/20260805-85bced0ae.patch | 55 ++++++++++++------- .../megatron/test_gdn_chunkwise_cp_route.py | 33 +++++++++++ 2 files changed, 69 insertions(+), 19 deletions(-) diff --git a/docker/patch/megatron/20260805-85bced0ae.patch b/docker/patch/megatron/20260805-85bced0ae.patch index e99db1fe9..900fad25d 100644 --- a/docker/patch/megatron/20260805-85bced0ae.patch +++ b/docker/patch/megatron/20260805-85bced0ae.patch @@ -526,10 +526,10 @@ index c297a23..587e1de 100644 + return None diff --git a/megatron/core/context_parallel_layout.py b/megatron/core/context_parallel_layout.py new file mode 100644 -index 0000000..718b2fc +index 0000000..c287ffb --- /dev/null +++ b/megatron/core/context_parallel_layout.py -@@ -0,0 +1,691 @@ +@@ -0,0 +1,699 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Context parallel tensor layout helpers.""" @@ -651,6 +651,7 @@ index 0000000..718b2fc + input_split_sizes: List[int] + output_split_sizes: List[int] + cu_seqlens: torch.Tensor ++ cu_seqlens_version: int + cp_size: int + cp_rank: int + source_layout: str @@ -887,6 +888,7 @@ index 0000000..718b2fc + input_split_sizes=input_split_sizes, + output_split_sizes=output_split_sizes, + cu_seqlens=cu_seqlens, ++ cu_seqlens_version=cu_seqlens._version, + cp_size=cp_size, + cp_rank=cp_rank, + source_layout=source_layout, @@ -905,15 +907,21 @@ index 0000000..718b2fc +) -> bool: + """A cached route is only valid for the exact packed boundaries it was built from. + -+ ``cu_seqlens`` is compared by object identity, not by value: a value comparison is -+ itself a device-host synchronisation, which is precisely what the route exists to -+ avoid. Identity is also the stronger check -- the same tensor object is the same -+ buffer, whereas equal values could still come from a different microbatch. ++ ``cu_seqlens`` is matched on object identity *plus* autograd's version counter, ++ never on value: a value comparison would itself be a device-host synchronisation, ++ which is precisely what the route exists to avoid. Identity alone would not survive ++ a caller that refills a preallocated boundary buffer in place -- a pattern that ++ already exists elsewhere in Megatron -- so the version counter, which every in-place ++ op bumps (views share the counter with their base) and which costs a plain Python ++ attribute read, closes that hole. The one case neither check sees is a write that ++ bypasses the dispatcher entirely, e.g. through ``data_ptr()``. + """ + if route is None: + return False + if route.cu_seqlens is not cu_seqlens: + return False ++ if route.cu_seqlens_version != cu_seqlens._version: ++ return False + if route.cp_size != cp_size or route.cp_rank != cp_rank: + return False + if route.source_layout != source_layout or route.target_layout != target_layout: @@ -1822,7 +1830,7 @@ index 465e83f..232caef 100644 tensor_recv_prev = None tensor_recv_next = None diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py -index 8df4df1..66c4d1e 100644 +index 8df4df1..e41d55f 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -14,12 +14,17 @@ import torch.nn as nn @@ -2285,7 +2293,7 @@ index 8df4df1..66c4d1e 100644 query, key = torch.split(query_key, [split_size, split_size], dim=2) # Expand query and key if needed (grouped query attention) -@@ -567,7 +752,58 @@ class GatedDeltaNet(MegatronModule): +@@ -567,7 +752,67 @@ class GatedDeltaNet(MegatronModule): beta = beta.sigmoid() return g, beta @@ -2298,6 +2306,13 @@ index 8df4df1..66c4d1e 100644 + GPU -- only need to run for the first layer that asks. The result is cached on + the ``PackedSeqParams`` instance and reused while it still describes the very + same boundary tensors, mirroring how the CP layout routes are cached. ++ ++ A cached result is reused only when all four source tensors are still the same ++ objects *and* none of them has been written in place. Identity alone would not ++ notice a caller refilling a preallocated boundary buffer, so autograd's version ++ counter -- which every in-place op bumps, and which costs a plain Python ++ attribute read -- is checked as well. Comparing values instead would reintroduce ++ the device-host sync this cache exists to remove. + """ + sources = ( + packed_seq_params.cu_seqlens_q_padded, @@ -2305,14 +2320,14 @@ index 8df4df1..66c4d1e 100644 + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + ) ++ versions = tuple(None if t is None else t._version for t in sources) + cached = getattr(packed_seq_params, "_gdn_resolved_cu_seqlens", None) + if cached is not None: -+ cached_key, cached_value = cached ++ cached_meta, cached_sources, cached_versions, cached_value = cached + if ( -+ cached_key[0] == seq_len_global -+ and cached_key[1] == cp_size -+ and len(cached_key[2]) == len(sources) -+ and all(a is b for a, b in zip(cached_key[2], sources)) ++ cached_meta == (seq_len_global, cp_size) ++ and all(a is b for a, b in zip(cached_sources, sources)) ++ and cached_versions == versions + ): + return cached_value + @@ -2334,7 +2349,9 @@ index 8df4df1..66c4d1e 100644 + + resolved = (cu_seqlens_q, cu_seqlens_kv) + packed_seq_params._gdn_resolved_cu_seqlens = ( -+ (seq_len_global, cp_size, sources), ++ (seq_len_global, cp_size), ++ sources, ++ versions, + resolved, + ) + return resolved @@ -2345,7 +2362,7 @@ index 8df4df1..66c4d1e 100644 """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" if cu_seqlens_padded is not None: cu_seqlens = cu_seqlens_padded -@@ -582,6 +818,13 @@ class GatedDeltaNet(MegatronModule): +@@ -582,6 +827,13 @@ class GatedDeltaNet(MegatronModule): f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." ) @@ -2359,7 +2376,7 @@ index 8df4df1..66c4d1e 100644 return cu_seqlens def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): -@@ -780,8 +1023,10 @@ def get_parameter_local_cp( +@@ -780,8 +1032,10 @@ def get_parameter_local_cp( torch.Tensor: The local parameter for the current context parallel rank. """ @@ -2372,7 +2389,7 @@ index 8df4df1..66c4d1e 100644 # No need to split if CP size is 1. if cp_size == 1: -@@ -800,7 +1045,7 @@ def get_parameter_local_cp( +@@ -800,7 +1054,7 @@ def get_parameter_local_cp( slices = [slice(None)] * param.dim() dim_size = param.size(dim=dim) slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) @@ -2381,7 +2398,7 @@ index 8df4df1..66c4d1e 100644 return param -@@ -945,6 +1190,7 @@ def torch_chunk_gated_delta_rule( +@@ -945,6 +1199,7 @@ def torch_chunk_gated_delta_rule( output_final_state=False, use_qk_l2norm_in_kernel=False, cu_seqlens=None, @@ -2389,7 +2406,7 @@ index 8df4df1..66c4d1e 100644 ): # pylint: disable=line-too-long ''' -@@ -957,6 +1203,9 @@ def torch_chunk_gated_delta_rule( +@@ -957,6 +1212,9 @@ def torch_chunk_gated_delta_rule( assert ( cu_seqlens is None ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." diff --git a/tests/backends/megatron/test_gdn_chunkwise_cp_route.py b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py index 886a88a07..6276150c6 100644 --- a/tests/backends/megatron/test_gdn_chunkwise_cp_route.py +++ b/tests/backends/megatron/test_gdn_chunkwise_cp_route.py @@ -178,6 +178,39 @@ def test_route_is_rebuilt_for_new_packed_boundaries(): assert torch.equal(got_rows, want_rows) +def test_route_is_rebuilt_when_cu_seqlens_is_mutated_in_place(): + """Identity alone would miss a caller refilling a preallocated boundary + buffer. + + Megatron already has that pattern elsewhere (persistent ``_cu_seqlens_buffer`` + written with ``buf[0] = 0``), so the cache also fingerprints autograd's version + counter, which every in-place write bumps. + """ + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + first = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + + # Same tensor object, refilled with different boundaries. + cu.copy_(_cu([2, 2], unit=8)) + rebuilt = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + assert rebuilt is not first, "an in-place refill must invalidate the cached route" + + want = cpl.build_thd_cp_partition_route(cu, 4, 1, "zigzag", "contiguous") + assert rebuilt.input_split_sizes == want.input_split_sizes + assert rebuilt.output_split_sizes == want.output_split_sizes + + +def test_route_is_rebuilt_when_a_view_of_cu_seqlens_is_mutated(): + """Views share the version counter with their base, so writes through one + count.""" + cu = _cu([3, 1], unit=8) + psp = _packed_seq_params(cu) + first = cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") + + cu[1:] = _cu([2, 2], unit=8)[1:] + assert cpl.get_thd_cp_partition_route(psp, cu, 4, 1, "zigzag", "contiguous") is not first + + def test_route_is_rebuilt_when_the_dynamic_cp_geometry_changes(): """Dynamic CP varies cp_size/cp_rank across micro-batches on one module.""" cu = _cu([3, 1], unit=8)