diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 0c265a508e9..8b3ea9975c7 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -115,6 +115,15 @@ def init( args, role ) + parallel_state = get_parallel_state() + if parallel_state.cp.size > 1: + from miles_plugins.models.cp_utils import detect_and_setup_hybrid_cp + + for model_chunk in self.model: + detect_and_setup_hybrid_cp( + model_chunk, parallel_state.cp.group, parallel_state.cp.rank, parallel_state.cp.size + ) + verify_megatron_parallel_state(self.model) if role == "critic": diff --git a/miles/backends/training_utils/cp_utils.py b/miles/backends/training_utils/cp_utils.py index 0fbba35b02b..e79ccb44698 100644 --- a/miles/backends/training_utils/cp_utils.py +++ b/miles/backends/training_utils/cp_utils.py @@ -1,11 +1,20 @@ +import logging from collections.abc import Callable import torch import torch.distributed as dist +import torch.nn as nn import torch.nn.functional as F from .parallel import get_parallel_state +try: + from fla.ops.cp import build_cp_context as _fla_build_cp_context +except ImportError: + _fla_build_cp_context = None + +logger = logging.getLogger(__name__) + def get_logits_and_tokens_offset_with_cp( total_length: int, @@ -336,3 +345,30 @@ def slice_log_prob_with_cp( return chunk_1 + chunk_2 else: return torch.cat([chunk_1, chunk_2], dim=0) + + +def build_gdn_cp_context(module: nn.Module, cu_seqlens: torch.Tensor, device: torch.device): + """Build fla CP context for a GatedDeltaNet module from packed sequence boundaries. + + Args: + module: GDN module with ``cp_group`` / ``cp_world_size`` / ``conv_kernel_size``. + cu_seqlens: Global packed sequence boundaries (e.g. ``packed_seq_params.cu_seqlens_q``). + device: Target device. + + Returns ``None`` when CP is not configured on the module (``cp_group`` not set). + Raises ``RuntimeError`` if hybrid CP is configured but ``fla.ops.cp`` is missing. + """ + cp_group = getattr(module, "cp_group", None) + if cp_group is None: + return None + if _fla_build_cp_context is None: + raise RuntimeError( + "Hybrid CP requires fla.ops.cp (flash-linear-attention >= 0.4.2) " "but it could not be imported." + ) + if cu_seqlens is None or cu_seqlens.numel() < 2: + raise ValueError(f"Hybrid CP requires valid cu_seqlens (at least 2 elements) but got {cu_seqlens}") + return _fla_build_cp_context( + cu_seqlens=cu_seqlens.to(device=device, dtype=torch.int32), + group=cp_group, + conv1d_kernel_size=module.conv_kernel_size, + ) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index f7726618cbd..375cd6c2c27 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2140,6 +2140,9 @@ def equal(x, y): ), ("rope_theta", "rotary_base", equal), ]: + # FIXME: Qwen3.5 transfomers has bug. + if getattr(hf_config, "model_type", "") == "qwen3_5_moe_text" and hf_config_name == "intermediate_size": + continue if hasattr(hf_config, hf_config_name): if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)): errors.append( diff --git a/miles_plugins/models/cp_utils.py b/miles_plugins/models/cp_utils.py new file mode 100644 index 00000000000..87f4205f343 --- /dev/null +++ b/miles_plugins/models/cp_utils.py @@ -0,0 +1,26 @@ +import logging + +import torch.distributed as dist +import torch.nn as nn + +from miles_plugins.models.hf_attention import HuggingfaceAttention + +logger = logging.getLogger(__name__) + + +def detect_and_setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> int: + """Scan for GatedDeltaNet modules and configure them for native fla CP.""" + count = 0 + for module in model.modules(): + if isinstance(module, HuggingfaceAttention): + linear_attn = getattr(module, "linear_attn", None) + if linear_attn is not None: + linear_attn.cp_group = cp_group + linear_attn.cp_rank = cp_rank + linear_attn.cp_world_size = cp_world_size + module.hybrid_cp = True + count += 1 + + if count > 0: + logger.info(f"Configured hybrid CP on {count} GDN modules (fla native state passing)") + return count diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index 7abe09b0eed..aeacc58ea22 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -38,6 +38,116 @@ def _fix_dtype(d): return ns +def _get_cp_sequence_lengths(cu_seqlens, cp_size, local_total_len=None): + global_seq_lengths = [(cu_seqlens[i + 1] - cu_seqlens[i]).item() for i in range(len(cu_seqlens) - 1)] + local_seq_lengths = [] + for global_seq_len in global_seq_lengths: + if global_seq_len % cp_size != 0: + raise ValueError(f"Expected sequence length {global_seq_len} to be divisible by cp_size={cp_size}") + local_seq_lengths.append(global_seq_len // cp_size) + + if local_total_len is not None and sum(local_seq_lengths) != local_total_len: + raise ValueError(f"Expected local total length {local_total_len}, got {sum(local_seq_lengths)}") + + return global_seq_lengths, local_seq_lengths + + +def _gather_cp_tensors(x, cp_group): + gathered = [torch.empty_like(x) for _ in range(dist.get_world_size(group=cp_group))] + dist.all_gather(gathered, x.contiguous(), group=cp_group) + return gathered + + +def _zigzag_to_packed_shard_impl(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + """Convert zigzag ring-attn layout to the contiguous packed shard expected by fla CP.""" + global_seq_lengths, local_seq_lengths = _get_cp_sequence_lengths(cu_seqlens, cp_size, hidden_states.size(0)) + gathered_by_rank = [ + gathered.split(local_seq_lengths, dim=0) for gathered in _gather_cp_tensors(hidden_states, cp_group) + ] + + full_sequences = [] + for seq_idx, global_seq_len in enumerate(global_seq_lengths): + per_rank = [rank_seqs[seq_idx] for rank_seqs in gathered_by_rank] + if global_seq_len % (2 * cp_size) == 0: + subchunk_len = global_seq_len // (2 * cp_size) + full_seq = torch.cat( + [seq[:subchunk_len] for seq in per_rank] + [seq[subchunk_len:] for seq in per_rank][::-1], + dim=0, + ) + else: + # Final local padding is appended contiguously on each rank, not in zigzag order. + full_seq = torch.cat(per_rank, dim=0) + full_sequences.append(full_seq) + + full_stream = torch.cat(full_sequences, dim=0) if full_sequences else hidden_states[:0] + shard_len = hidden_states.size(0) + return full_stream[cp_rank * shard_len : (cp_rank + 1) * shard_len] + + +def _packed_shard_to_zigzag_impl(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + """Convert contiguous packed shard layout back to zigzag ring-attn layout.""" + global_seq_lengths, local_seq_lengths = _get_cp_sequence_lengths(cu_seqlens, cp_size, hidden_states.size(0)) + full_stream = torch.cat(_gather_cp_tensors(hidden_states, cp_group), dim=0) + full_sequences = full_stream.split(global_seq_lengths, dim=0) + + local_sequences = [] + for full_seq, global_seq_len, local_seq_len in zip( + full_sequences, global_seq_lengths, local_seq_lengths, strict=True + ): + if global_seq_len % (2 * cp_size) == 0: + subchunk_len = global_seq_len // (2 * cp_size) + parts = full_seq.split(subchunk_len, dim=0) + local_sequences.append(torch.cat([parts[cp_rank], parts[2 * cp_size - 1 - cp_rank]], dim=0)) + else: + local_sequences.append(full_seq.split(local_seq_len, dim=0)[cp_rank]) + + return torch.cat(local_sequences, dim=0) if local_sequences else hidden_states[:0] + + +class _ZigzagToPackedShard(torch.autograd.Function): + """Convert zigzag ring-attn layout to contiguous packed shards for native fla CP.""" + + @staticmethod + def forward(ctx, hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + ctx.cp_group = cp_group + ctx.cp_rank = cp_rank + ctx.cp_size = cp_size + ctx.save_for_backward(cu_seqlens) + return _zigzag_to_packed_shard_impl(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size) + + @staticmethod + def backward(ctx, grad_output): + (cu_seqlens,) = ctx.saved_tensors + result = _packed_shard_to_zigzag_impl(grad_output, cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size) + return result, None, None, None, None + + +class _PackedShardToZigzag(torch.autograd.Function): + """Convert contiguous packed shards back to zigzag ring-attn layout.""" + + @staticmethod + def forward(ctx, hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + ctx.cp_group = cp_group + ctx.cp_rank = cp_rank + ctx.cp_size = cp_size + ctx.save_for_backward(cu_seqlens) + return _packed_shard_to_zigzag_impl(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size) + + @staticmethod + def backward(ctx, grad_output): + (cu_seqlens,) = ctx.saved_tensors + result = _zigzag_to_packed_shard_impl(grad_output, cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size) + return result, None, None, None, None + + +def _zigzag_to_packed_shard(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + return _ZigzagToPackedShard.apply(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size) + + +def _packed_shard_to_zigzag(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size): + return _PackedShardToZigzag.apply(hidden_states, cu_seqlens, cp_group, cp_rank, cp_size) + + class _AllGatherForDuplicatedComputation(torch.autograd.Function): """All-gather whose backward just returns the local gradient slice (no reduce). @@ -68,6 +178,10 @@ class HuggingfaceAttention(MegatronModule, ABC): "cross attn" specializations. """ + # Subclasses set this to True when the underlying module handles CP natively + # (e.g. via fla's state-passing CP for DeltaNet), bypassing the all-gather. + hybrid_cp: bool = False + def __init__( self, args, @@ -115,7 +229,22 @@ def forward( group=mpu.get_tensor_model_parallel_group(), ) - if mpu.get_context_parallel_world_size() > 1: + if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp: + cp_size = mpu.get_context_parallel_world_size() + # Native fla CP expects each rank to own a contiguous shard of the + # packed global token stream. In allgather-CP mode the data pipeline + # already provides that layout, so no extra relayout is + # needed here. + if not self.args.allgather_cp: + hidden_states = _zigzag_to_packed_shard( + hidden_states, + cu_seqlens, + mpu.get_context_parallel_group(), + mpu.get_context_parallel_rank(), + cp_size, + ) + + elif mpu.get_context_parallel_world_size() > 1: cp_size = mpu.get_context_parallel_world_size() # Use custom all-gather whose backward returns local gradient # instead of reduce-scatter, since the computation is duplicated. @@ -150,7 +279,17 @@ def forward( output = output.permute(1, 0, 2) # [seq_len, bsz, hidden_dim] - if mpu.get_context_parallel_world_size() > 1: + if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp: + if not self.args.allgather_cp: + output = _packed_shard_to_zigzag( + output, + cu_seqlens, + mpu.get_context_parallel_group(), + mpu.get_context_parallel_rank(), + cp_size, + ) + + elif mpu.get_context_parallel_world_size() > 1: cp_rank = mpu.get_context_parallel_rank() output_list = [] for i in range(len(cu_seqlens) - 1): diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index a796c8c49c5..794cf738081 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -15,6 +15,8 @@ except ImportError: pass +from miles.backends.training_utils.cp_utils import build_gdn_cp_context + from .hf_attention import HuggingfaceAttention, _load_hf_config @@ -88,6 +90,8 @@ def forward( ): batch_size, seq_len, _ = hidden_states.shape + cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) + # Projections (flat layout: [Q_all, K_all, V_all]) mixed_qkv = self.in_proj_qkv(hidden_states) z = self.in_proj_z(hidden_states) @@ -95,10 +99,12 @@ def forward( b = self.in_proj_b(hidden_states) a = self.in_proj_a(hidden_states) - # Convolution on the flat QKV + # Convolution on the flat QKV (pass cp_context for boundary handling) + conv_cu_seqlens = cp_context.cu_seqlens if cp_context is not None else cu_seqlens mixed_qkv, _ = self.conv1d( x=mixed_qkv, - cu_seqlens=cu_seqlens, + cu_seqlens=conv_cu_seqlens, + cp_context=cp_context, ) # Split into Q, K, V (flat split, matching HF layout) @@ -118,17 +124,29 @@ def forward( query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) + if cp_context is not None: + core_attn_out, _ = chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cp_context.cu_seqlens, + cp_context=cp_context, + ) + else: + core_attn_out, _ = chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) z_shape_og = z.shape # reshape input data into 2D tensor diff --git a/miles_plugins/models/qwen3_next.py b/miles_plugins/models/qwen3_next.py index 92e39ff318d..1dbee8acd01 100644 --- a/miles_plugins/models/qwen3_next.py +++ b/miles_plugins/models/qwen3_next.py @@ -18,6 +18,8 @@ except ImportError: pass +from miles.backends.training_utils.cp_utils import build_gdn_cp_context + from .hf_attention import HuggingfaceAttention @@ -108,6 +110,8 @@ def forward( hidden_states: torch.Tensor, cu_seqlens: torch.Tensor = None, ): + cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device) + projected_states_qkvz = self.in_proj_qkvz(hidden_states) projected_states_ba = self.in_proj_ba(hidden_states) query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba) @@ -115,9 +119,11 @@ def forward( mixed_qkv = torch.cat((query, key, value), dim=-1) + conv_cu_seqlens = cp_context.cu_seqlens if cp_context is not None else cu_seqlens mixed_qkv, _ = self.conv1d( x=mixed_qkv, - cu_seqlens=cu_seqlens, + cu_seqlens=conv_cu_seqlens, + cp_context=cp_context, ) query, key, value = torch.split( @@ -140,17 +146,29 @@ def forward( query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) + if cp_context is not None: + core_attn_out, _ = chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cp_context.cu_seqlens, + cp_context=cp_context, + ) + else: + core_attn_out, _ = chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) z_shape_og = z.shape # reshape input data into 2D tensor diff --git a/scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py b/scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py new file mode 100644 index 00000000000..ee70d5f8b05 --- /dev/null +++ b/scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py @@ -0,0 +1,178 @@ +from dataclasses import dataclass +from typing import Literal + +import typer + +import miles.utils.external_utils.command_utils as U + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + mode: Literal["normal", "debug_minimal"] = "normal" + run_id: str = U.create_run_id() + model_name: str = "Qwen3.5-35B-A3B" + megatron_model_type: str = "qwen3.5-35B-A3B" + num_gpus_per_node: int = 8 + hardware: Literal["H200"] = "H200" + enable_eval: bool = True + extra_args: str = "" + data_dir: str = "/root/datasets" + model_dir: str = "/root/models" + megatron_path: str = "/root/Megatron-LM" + + +def prepare(args: ScriptArgs): + U.exec_command(f"mkdir -p {args.model_dir} {args.data_dir}") + U.exec_command("pip install transformers==5.2.0") + U.exec_command(f"hf download Qwen/{args.model_name} --local-dir {args.model_dir}/{args.model_name}") + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + U.hf_download_dataset("zhuzilin/aime-2024", data_dir=args.data_dir) + + U.convert_checkpoint( + model_name=args.model_name, + megatron_model_type=args.megatron_model_type, + num_gpus_per_node=args.num_gpus_per_node, + dir_dst=args.model_dir, + hf_checkpoint=f"{args.model_dir}/{args.model_name}", + megatron_path=args.megatron_path, + ) + + +def execute(args: ScriptArgs): + ref_load_path = f"{args.model_dir}/{args.model_name}_torch_dist" + load_save_path = f"{args.output_dir}/{args.run_id}/checkpoints" + + ckpt_args = ( + f"--hf-checkpoint {args.model_dir}/{args.model_name} " + f"--ref-load {ref_load_path} " + f"--load {load_save_path} " + f"--save {load_save_path} " + f"--save-interval {2 if args.mode == 'debug_minimal' else 20} " + ) + + rollout_args = ( + f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + f"--num-rollout {64 if args.mode == 'debug_minimal' else 3000} " + f"--rollout-batch-size {8 if args.mode == 'debug_minimal' else 32} " + f"--n-samples-per-prompt {2 if args.mode == 'debug_minimal' else 8} " + f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} " + "--rollout-temperature 1 " + f"--global-batch-size {16 if args.mode == 'debug_minimal' else 256} " + "--balance-data " + ) + + eval_args = "" + if (args.mode != "debug_minimal") and args.enable_eval: + eval_args += ( + "--eval-interval 20 " + f"--eval-prompt-data aime {args.data_dir}/aime-2024/aime-2024.jsonl " + "--n-samples-per-eval-prompt 16 " + "--eval-max-response-len 16384 " + "--eval-top-p 1 " + ) + + # CP=2 EP=8: validated on 8x H200 + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 8192 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 8 " + "--sglang-mem-fraction-static 0.7 " + "--sglang-ep-size 8 " + "--sglang-cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128 136 144 152 160 168 176 184 192 200 208 216 224 232 240 248 256 " + # mtp speculative decoding + "--sglang-speculative-algorithm EAGLE " + "--sglang-speculative-num-steps 2 " + "--sglang-speculative-eagle-topk 1 " + "--sglang-speculative-num-draft-tokens 3 " + "--sglang-max-running-requests 512 " + "--sglang-mamba-scheduler-strategy extra_buffer " + ) + + mtp_args = "--enable-mtp-training " "--mtp-num-layers 1 " "--mtp-loss-scaling-factor 0.2 " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--moe-token-dispatcher-type flex " + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--num-gpus-per-node {args.num_gpus_per_node} " + "--colocate " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{mtp_args} " + f"{misc_args} " + f"{args.extra_args} " + ) + + U.execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=args.megatron_model_type, + extra_env_vars={ + "SGLANG_ENABLE_SPEC_V2": "1", + }, + megatron_path=args.megatron_path, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs): + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main) diff --git a/tests/e2e/megatron/test_qwen3_5_35B_A3B_cp.py b/tests/e2e/megatron/test_qwen3_5_35B_A3B_cp.py new file mode 100644 index 00000000000..f951cf4f3a0 --- /dev/null +++ b/tests/e2e/megatron/test_qwen3_5_35B_A3B_cp.py @@ -0,0 +1,153 @@ +"""E2E test for Qwen3.5-35B-A3B with Context Parallel (CP=2 and CP=4). + +Validates that GDN layers use real fla native CP (state passing) instead of +duplicated all-gather computation. See: https://github.com/radixark/miles/issues/878 +""" + +import os + +import miles.utils.external_utils.command_utils as U + +MODEL_NAME = "Qwen3.5-35B-A3B" +MODEL_TYPE = "qwen3.5-35B-A3B" +NUM_GPUS = 8 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.hf_download_dataset("zhuzilin/aime-2024") + U.convert_checkpoint(model_name=MODEL_NAME, megatron_model_type=MODEL_TYPE, num_gpus_per_node=NUM_GPUS) + + +def _execute_with_cp(cp_size: int): + """Run a short training loop with the given context-parallel size.""" + assert NUM_GPUS % cp_size == 0 + ep_size = NUM_GPUS // cp_size + + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} " f"--ref-load /root/{MODEL_NAME}_torch_dist " + + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 8 " + "--rollout-max-response-len 8192 " + "--rollout-temperature 1 " + "--global-batch-size 32 " + "--balance-data " + ) + + eval_args = ( + "--eval-prompt-data aime24 /root/datasets/aime-2024/aime-2024.jsonl " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 16384 " + "--eval-top-k 1 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + f"--context-parallel-size {cp_size} " + f"--expert-model-parallel-size {ep_size} " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 8192 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 8 " + "--sglang-mem-fraction-static 0.7 " + f"--sglang-ep-size {NUM_GPUS} " + "--sglang-max-running-requests 512 " + "--sglang-speculative-algorithm EAGLE " + "--sglang-speculative-num-steps 2 " + "--sglang-speculative-eagle-topk 1 " + "--sglang-speculative-num-draft-tokens 3 " + ) + + mtp_args = "--enable-mtp-training " "--mtp-num-layers 1 " "--mtp-loss-scaling-factor 0.2 " + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--moe-token-dispatcher-type flex " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{mtp_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +def execute_cp2(): + """Qwen3.5-35B-A3B with CP=2.""" + _execute_with_cp(cp_size=2) + + +def execute_cp4(): + """Qwen3.5-35B-A3B with CP=4.""" + _execute_with_cp(cp_size=4) + + +if __name__ == "__main__": + cp_size = int(os.environ.get("CP_SIZE", "2")) + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + _execute_with_cp(cp_size) diff --git a/tests/test_qwen3_5_mtp_bridge_mapping.py b/tests/e2e/megatron/test_qwen3_5_mtp_bridge_mapping.py similarity index 100% rename from tests/test_qwen3_5_mtp_bridge_mapping.py rename to tests/e2e/megatron/test_qwen3_5_mtp_bridge_mapping.py diff --git a/tests/e2e/precision/test_hf_attention_cp_relayout.py b/tests/e2e/precision/test_hf_attention_cp_relayout.py new file mode 100644 index 00000000000..39ab46f9153 --- /dev/null +++ b/tests/e2e/precision/test_hf_attention_cp_relayout.py @@ -0,0 +1,101 @@ +"""Distributed correctness test for zigzag <-> packed-shard hybrid CP relayout. + +Run with: + torchrun --nproc_per_node=2 tests/e2e/precision/test_hf_attention_cp_relayout.py + torchrun --nproc_per_node=4 tests/e2e/precision/test_hf_attention_cp_relayout.py +""" + +import os +import sys + +import torch +import torch.distributed as dist + +from miles_plugins.models.hf_attention import _packed_shard_to_zigzag, _zigzag_to_packed_shard + + +def setup_dist(): + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + return rank, world_size, local_rank + + +def _make_subchunk(sample_id: int, sub_id: int, chunk_len: int, device: torch.device) -> torch.Tensor: + base = sample_id * 1000 + sub_id * 100 + values = torch.arange(base, base + chunk_len, device=device, dtype=torch.float32) + return values.view(-1, 1, 1) + + +def _build_rank_inputs(rank: int, world_size: int, device: torch.device): + chunk_lens = [3, 5] + tail_pad_local_len = 3 + zigzag_chunks = [] + full_sequences = [] + cu = [0] + + for sample_id, chunk_len in enumerate(chunk_lens): + subchunks = [_make_subchunk(sample_id, sub_id, chunk_len, device) for sub_id in range(2 * world_size)] + zigzag_chunks.extend([subchunks[rank], subchunks[2 * world_size - 1 - rank]]) + full_sequences.append(torch.cat(subchunks, dim=0)) + cu.append(cu[-1] + 2 * world_size * chunk_len) + + tail_pad = (rank * 10000 + torch.arange(tail_pad_local_len, device=device, dtype=torch.float32)).view(-1, 1, 1) + zigzag_chunks.append(tail_pad) + full_sequences.append( + torch.cat( + [ + (r * 10000 + torch.arange(tail_pad_local_len, device=device, dtype=torch.float32)).view(-1, 1, 1) + for r in range(world_size) + ], + dim=0, + ) + ) + cu.append(cu[-1] + world_size * tail_pad_local_len) + + zigzag = torch.cat(zigzag_chunks, dim=0).requires_grad_(True) + packed_full = torch.cat(full_sequences, dim=0) + local_len = zigzag.size(0) + packed_shard = packed_full[rank * local_len : (rank + 1) * local_len] + cu_seqlens = torch.tensor(cu, device=device, dtype=torch.int32) + return zigzag, packed_shard, cu_seqlens + + +def test_relayout(rank: int, world_size: int): + device = torch.device(f"cuda:{rank}") + cp_group = dist.group.WORLD + + zigzag, expected_packed_shard, cu_seqlens = _build_rank_inputs(rank, world_size, device) + + packed_shard = _zigzag_to_packed_shard(zigzag, cu_seqlens, cp_group, rank, world_size) + roundtrip = _packed_shard_to_zigzag(packed_shard, cu_seqlens, cp_group, rank, world_size) + + packed_ok = torch.equal(packed_shard, expected_packed_shard) + roundtrip_ok = torch.equal(roundtrip, zigzag) + + loss = roundtrip.sum() + loss.backward() + grad_ok = torch.equal(zigzag.grad, torch.ones_like(zigzag)) + + passed = packed_ok and roundtrip_ok and grad_ok + if rank == 0: + print(f"\n=== HF Attention Hybrid CP Relayout Test CP={world_size} ===") + print(f"zigzag->packed PASS: {packed_ok}") + print(f"roundtrip PASS: {roundtrip_ok}") + print(f"backward PASS: {grad_ok}") + if not passed: + sys.exit(1) + + +def main(): + rank, world_size, _ = setup_dist() + try: + test_relayout(rank, world_size) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/precision/test_qwen3_5_cp_correctness.py b/tests/e2e/precision/test_qwen3_5_cp_correctness.py new file mode 100644 index 00000000000..d0a2f3f32b6 --- /dev/null +++ b/tests/e2e/precision/test_qwen3_5_cp_correctness.py @@ -0,0 +1,147 @@ +"""Correctness test for Qwen3.5 GDN with native fla Context Parallel. + +Run with: + torchrun --nproc_per_node=2 tests/test_qwen3_5_cp_correctness.py # CP=2 + torchrun --nproc_per_node=4 tests/test_qwen3_5_cp_correctness.py # CP=4 + +Validates that GDN forward+backward with native fla CP produces results +consistent with the non-CP (single-rank full-sequence) baseline. +""" + +import os +import sys + +import torch +import torch.distributed as dist + + +def setup_dist(): + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + return rank, world_size, local_rank + + +def build_gdn_module(device, dtype=torch.bfloat16): + """Build a small Qwen3.5 GDN module for testing.""" + + class FakeConfig: + hidden_size = 256 + linear_num_value_heads = 4 + linear_num_key_heads = 2 + linear_key_head_dim = 64 + linear_value_head_dim = 64 + linear_conv_kernel_dim = 4 + hidden_act = "silu" + rms_norm_eps = 1e-6 + + FakeConfig.dtype = dtype + + from miles_plugins.models.qwen3_5 import Qwen3_5GatedDeltaNet + + return Qwen3_5GatedDeltaNet(FakeConfig, layer_idx=0).to(device=device, dtype=dtype) + + +def test_cp_forward_backward(rank, world_size): + device = torch.device(f"cuda:{rank}") + dtype = torch.bfloat16 + + # ---- Reference: full sequence on rank 0 (no CP) ---- + torch.manual_seed(42) + model_ref = build_gdn_module(device, dtype) + + total_seq_len = 128 * world_size # must be divisible by world_size + batch = 1 + + torch.manual_seed(123) + full_hidden = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype, requires_grad=True) + full_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device) + + # Forward without CP + ref_out = model_ref(full_hidden, cu_seqlens=full_cu) + ref_loss = ref_out.sum() + ref_loss.backward() + ref_grad = full_hidden.grad.clone() + + # ---- Test: CP across ranks ---- + torch.manual_seed(42) + model_cp = build_gdn_module(device, dtype) + # Copy weights from ref to ensure identical params + model_cp.load_state_dict(model_ref.state_dict()) + + # Set up CP context on the module + cp_group = dist.group.WORLD + model_cp.cp_group = cp_group + model_cp.cp_rank = rank + model_cp.cp_world_size = world_size + + # Each rank gets its local chunk + local_seq_len = total_seq_len // world_size + start = rank * local_seq_len + end = start + local_seq_len + + torch.manual_seed(123) + full_hidden_cp = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype) + local_hidden = full_hidden_cp[:, start:end, :].clone().contiguous().requires_grad_(True) + + # Global cu_seqlens (build_gdn_cp_context expects global boundaries) + global_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device) + + # Forward with CP + cp_out = model_cp(local_hidden, cu_seqlens=global_cu) + cp_loss = cp_out.sum() + + # Reduce loss across ranks to match reference + dist.all_reduce(cp_loss, op=dist.ReduceOp.SUM) + + cp_loss.backward() + + # ---- Gather outputs for comparison ---- + gathered_out = [torch.zeros_like(cp_out) for _ in range(world_size)] + dist.all_gather(gathered_out, cp_out.contiguous()) + full_cp_out = torch.cat(gathered_out, dim=1) + + gathered_grad = [torch.zeros_like(local_hidden.grad) for _ in range(world_size)] + dist.all_gather(gathered_grad, local_hidden.grad.contiguous()) + full_cp_grad = torch.cat(gathered_grad, dim=1) + + if rank == 0: + # Compare outputs + out_diff = (ref_out.detach().float() - full_cp_out.detach().float()).abs() + out_max_diff = out_diff.max().item() + out_rel_diff = (out_diff / (ref_out.detach().float().abs() + 1e-8)).max().item() + + # Compare gradients + grad_diff = (ref_grad.float() - full_cp_grad.float()).abs() + grad_max_diff = grad_diff.max().item() + grad_rel_diff = (grad_diff / (ref_grad.float().abs() + 1e-8)).max().item() + + print(f"\n=== CP={world_size} Correctness Test ===") + print(f"Forward max abs diff: {out_max_diff:.6e} max rel diff: {out_rel_diff:.6e}") + print(f"Backward max abs diff: {grad_max_diff:.6e} max rel diff: {grad_rel_diff:.6e}") + + # bf16 tolerance: 1e-2 is generous for bf16 accumulated ops + fwd_ok = out_max_diff < 1e-2 + bwd_ok = grad_max_diff < 1e-2 + print(f"Forward PASS: {fwd_ok}") + print(f"Backward PASS: {bwd_ok}") + + if not (fwd_ok and bwd_ok): + print("FAILED!") + sys.exit(1) + else: + print(f"CP={world_size} test PASSED!") + + +def main(): + rank, world_size, _ = setup_dist() + try: + test_cp_forward_backward(rank, world_size) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast/backends/training_utils/__init__.py b/tests/fast/backends/training_utils/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/fast/backends/training_utils/__init__.py @@ -0,0 +1 @@ +