From b827d4204d336237830fc53b64250ec5e1920460 Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Tue, 30 Jun 2026 15:20:33 +0800 Subject: [PATCH 1/6] fsdp rr --- .../rl/grpo/short_math_grpo_routing_replay.py | 334 ++++++++++++++++++ .../rl/grpo/short_math_grpo_routing_replay.sh | 24 ++ src/twinkle/cli/cli.py | 1 + src/twinkle/data_format/input_feature.py | 1 + src/twinkle/data_format/output.py | 1 + src/twinkle/data_format/sampling.py | 1 + .../model/transformers/moe/expert_parallel.py | 85 ++++- .../model/transformers/moe/router_replay.py | 281 +++++++++++++++ .../model/transformers/transformers.py | 83 ++++- src/twinkle/processor/base.py | 38 +- .../sampler/vllm_sampler/vllm_engine.py | 2 + .../sampler/vllm_sampler/vllm_sampler.py | 16 +- 12 files changed, 848 insertions(+), 19 deletions(-) create mode 100644 cookbook/rl/grpo/short_math_grpo_routing_replay.py create mode 100644 cookbook/rl/grpo/short_math_grpo_routing_replay.sh create mode 100644 src/twinkle/model/transformers/moe/router_replay.py diff --git a/cookbook/rl/grpo/short_math_grpo_routing_replay.py b/cookbook/rl/grpo/short_math_grpo_routing_replay.py new file mode 100644 index 000000000..8297239f3 --- /dev/null +++ b/cookbook/rl/grpo/short_math_grpo_routing_replay.py @@ -0,0 +1,334 @@ +"""GRPO training with routing replay for MoE models (Only FSDP backend is supported currently). + +Supports three routing replay modes for MoE expert routing consistency: +- disabled: No routing replay (default behavior) +- R2: Record routing during a forward_only RECORD pass, then replay during training +- R3: vLLM returns routed_experts data (requires vLLM >= 0.14.0), training replays directly +""" +import copy +import re +from typing import List, Tuple, Dict, Any + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import GRPOMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward import GSM8KAccuracyReward +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler +from twinkle.preprocessor.llm import GSM8KProcessor +from twinkle.model.transformers.moe.router_replay import RouterReplayAction + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.6-35B-A3B' +USE_MEGATRON = args.model.strategy != 'native_fsdp' +ROUTER_REPLAY_MODE = args.rl.router_replay_mode or 'R3' + +MODEL_GPUS = args.infra.model_gpus or 4 +MODEL_FSDP = args.infra.fsdp_size or 4 +MODEL_DP = args.infra.dp_size or 1 +MODEL_EP = args.infra.ep_size or 4 +MODEL_TP = args.infra.tp_size or 1 +MODEL_PP = args.infra.pp_size or 1 + +SAMPLER_GPUS = args.infra.sampler_gpus or 2 +SAMPLER_TP = args.sampler.tensor_parallel_size or 2 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +LEARNING_RATE = args.optimizer.learning_rate or 5e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 4 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 1 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 1000 +LORA_RANK = args.lora.lora_r or 16 + +SYSTEM_PROMPT = ('You are a helpful math assistant. Solve the problem with minimal but correct reasoning ' + 'and put your final answer within \\boxed{}.') + +# Validate configuration +if ROUTER_REPLAY_MODE not in ('disabled', 'R2', 'R3'): + raise ValueError(f'Invalid ROUTER_REPLAY_MODE: {ROUTER_REPLAY_MODE}. ' + f"Must be one of 'disabled', 'R2', 'R3'") +if ROUTER_REPLAY_MODE != 'disabled' and USE_MEGATRON: + raise ValueError('Routing replay requires USE_MEGATRON=0 (Only FSDP backend is supported currently)') +if ROUTER_REPLAY_MODE == 'R3': + logger.info('R3 mode: vLLM will return routed_experts data (requires vLLM >= 0.14.0)') +elif ROUTER_REPLAY_MODE == 'R2': + logger.info('R2 mode: Recording routing during forward_only, replaying during training') + +# ========== Reward Functions ========== +class GSM8KBrevityReward(Reward): + """Brevity reward: rewards shorter completions that contain a valid answer. + + Returns 0.0 if no valid answer format (\\boxed{} or ####). + Otherwise returns higher score for shorter completions (1.0 at <=200 chars). + """ + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + + has_answer = bool( + re.search(r'\\boxed\{[^}]+\}', completion) + or re.search(r'####\s*[\-\d,\.]+', completion) + ) + + if not has_answer: + rewards.append(0.0) + else: + length = len(completion) + if length <= 200: + rewards.append(1.0) + else: + rewards.append(max(0.0, 1.0 - (length - 200) / 3000)) + return rewards + + +# ========== Dataset ========== +def create_gsm8k_dataset(): + dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=4096, truncation_strategy='delete', enable_thinking=False) + dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) + dataset.encode(add_generation_prompt=True) + return dataset + + +def compute_rewards( + trajectories: List[Dict[str, Any]], +) -> Tuple[List[float], List[float], List[float]]: + accuracy_reward_fn = GSM8KAccuracyReward() + brevity_reward_fn = GSM8KBrevityReward() + + accuracy_rewards = accuracy_reward_fn(trajectories) + brevity_rewards = brevity_reward_fn(trajectories) + total_rewards = [a + b for a, b in zip(accuracy_rewards, brevity_rewards)] + return total_rewards, brevity_rewards, accuracy_rewards + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU', gpus_per_worker=SAMPLER_TP), + ] + if USE_MEGATRON: + dp_size = MODEL_GPUS // (MODEL_TP * MODEL_PP) + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=dp_size, tp_size=MODEL_TP, pp_size=MODEL_PP, ep_size=MODEL_EP, sequence_parallel=True) + else: + model_mesh = DeviceMesh.from_sizes(fsdp_size=MODEL_FSDP, dp_size=MODEL_DP, ep_size=MODEL_EP) + sampler_dp_size = SAMPLER_GPUS // (SAMPLER_TP) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=sampler_dp_size, tp_size=SAMPLER_TP) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + enable_ep = MODEL_EP > 1 + if enable_ep and not USE_MEGATRON: + lora_config = LoraConfig( + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + target_modules='all-linear', + target_parameters=['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'], + ) + else: + lora_config = LoraConfig( + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + target_modules='all-linear', + ) + + if USE_MEGATRON: + from twinkle.model.megatron import MegatronModel + model = MegatronModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + mixed_precision='bf16', + ) + else: + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + enable_router_replay=(ROUTER_REPLAY_MODE != 'disabled'), + strategy='native_fsdp', + fsdp_config={ + 'expert_parallel': { + 'enabled': enable_ep, + 'router_dtype': 'fp32', + 'keep_router_logits': False, + } + }, + ) + + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + if USE_MEGATRON: + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE) + else: + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + + model.set_loss('GRPOLoss', epsilon=0.2) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + # Configure sampler: R3 mode enables routed_experts return from vLLM + engine_args = { + 'tensor_parallel_size': SAMPLER_TP, + 'gpu_memory_utilization': 0.7, + 'max_model_len': 10000, + 'max_lora_rank': LORA_RANK, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + } + if ROUTER_REPLAY_MODE == 'R3': + engine_args['enable_return_routed_experts'] = True + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args=engine_args, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_gsm8k_dataset, + batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = GRPOMetric() + sampling_params = SamplingParams(max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95) + + optim_step = 0 + logger.info(f'Starting GSM8K GRPO training with router replay mode={ROUTER_REPLAY_MODE}') + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + expand_prompts = [] + for prompt in batch: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + # enable_lora=True used with ckpt_manager.sync_weights(merge_and_sync=False) + # meaning only sync lora weights, if merge_and_sync=True, + # lora will be merged into the base model and sync all weights to vLLM + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + sample_responses = sampler.sample( + expand_prompts, + sampling_params, + ) + if sample_responses and sample_responses[0].sequences: + first_decoded = sample_responses[0].sequences[0].decoded + if isinstance(first_decoded, str): + logger.info('[sample_debug] first_generation=%r', first_decoded[:512]) + + all_input_data: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + + for sample_response in sample_responses: + for sequence in sample_response.sequences: + all_input_data.append(sequence.new_input_feature) + all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + all_completion_lengths.append(len(sequence.tokens)) + + total_rewards, _, _ = compute_rewards(all_input_data) + + advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + + total_completions = len(all_input_data) + + # compute old logps and routed_experts(R2) + # R2: forward_only RECORD pass → get routing data → inject into inputs + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + mb_inputs = all_input_data[mb_start:mb_end] + recompute_output = model.forward_only( + inputs=copy.deepcopy(mb_inputs), + router_replay_action={'R2': RouterReplayAction.RECORD,'R3': RouterReplayAction.REPLAY_FORWARD}.get(ROUTER_REPLAY_MODE), + micro_batch_size=MICRO_BATCH_SIZE, + ) + old_logps = recompute_output.get('logps') + assert old_logps.shape[0] == len(mb_inputs) + for i, mb in enumerate(mb_inputs): + mb['old_logps'] = old_logps[i] + routed_experts = recompute_output.get('routed_experts') + if routed_experts is not None: + assert routed_experts.shape[0] == len(mb_inputs) + for i, mb in enumerate(mb_inputs): + mb['routed_experts'] = routed_experts[i] + + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + mb_inputs = all_input_data[mb_start:mb_end] + mb_old_logps = all_old_logps[mb_start:mb_end] + mb_advantages = advantages[mb_start:mb_end] + recompute_logps = [input.pop('old_logps').unsqueeze(0) for input in mb_inputs] + + mb_output = model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + router_replay_action=RouterReplayAction.REPLAY_FORWARD + ) + model.clip_grad_and_step() + optim_step += 1 + + logps = mb_output.get('logps') + mb_output['logps'] = [logps[i:i+1] for i in range(logps.size(0))] + metrics.accumulate( + mb_inputs, + mb_output, + old_logps=mb_old_logps, + advantages=mb_advantages, + ) + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'math-grpo-checkpoint-{optim_step}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('grpo-gsm8k-checkpoint') + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/cookbook/rl/grpo/short_math_grpo_routing_replay.sh b/cookbook/rl/grpo/short_math_grpo_routing_replay.sh new file mode 100644 index 000000000..23ea801ce --- /dev/null +++ b/cookbook/rl/grpo/short_math_grpo_routing_replay.sh @@ -0,0 +1,24 @@ +#!/bin/sh +set -eu + +python short_math_grpo_routing_replay.py \ + --model-id ms://Qwen/Qwen3.6-35B-A3B \ + --strategy native_fsdp \ + --router_replay_mode R3 \ + --model-gpus 4 \ + --sampler-gpus 2 \ + --fsdp-size 4 \ + --dp-size 1 \ + --ep-size 4 \ + --tensor-parallel-size 2 \ + --num-generations 8 \ + --max-tokens 4096 \ + --batch-size 4 \ + --mini-batch-size 4 \ + --micro-batch-size 1 \ + --max-steps 1000 \ + --lr 5e-5 \ + --lora-r 16 \ + --save-steps 1000 \ + --adapter-name default \ + "$@" diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index 8e9412159..ce720a9f5 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -193,6 +193,7 @@ class RLArgs: gkd_beta: float = 0.5 gkd_temperature: float = 1.0 gkd_topk: int = 64 + router_replay_mode: Literal['disabled', 'R2', 'R3'] = 'disabled' @dataclass diff --git a/src/twinkle/data_format/input_feature.py b/src/twinkle/data_format/input_feature.py index 525f700c2..ebb4d4263 100644 --- a/src/twinkle/data_format/input_feature.py +++ b/src/twinkle/data_format/input_feature.py @@ -36,3 +36,4 @@ class InputFeature(TypedDict, total=False): labels: InputType completion_mask: InputType length: int + routed_experts: InputType diff --git a/src/twinkle/data_format/output.py b/src/twinkle/data_format/output.py index 596252fb6..6a248649b 100644 --- a/src/twinkle/data_format/output.py +++ b/src/twinkle/data_format/output.py @@ -27,6 +27,7 @@ class ModelOutput(TypedDict, total=False): logps: Optional[OutputType] num_tokens: Optional[OutputType] embeddings: Optional[OutputType] + routed_experts: Optional[OutputType] class LossOutput(TypedDict, total=False): diff --git a/src/twinkle/data_format/sampling.py b/src/twinkle/data_format/sampling.py index 01ff0377d..05ecdd641 100644 --- a/src/twinkle/data_format/sampling.py +++ b/src/twinkle/data_format/sampling.py @@ -174,6 +174,7 @@ class SampledSequence: logprobs: Optional[List[List[Tuple[int, float]]]] = None decoded: str = None new_input_feature: InputFeature = None + routed_experts: Optional[Any] = None @dataclass diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 2b6e45ea9..93e447a85 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -63,9 +63,9 @@ def apply_expert_parallel( ep_rank = ep_mesh.get_local_rank() specs = [] - for _, block in find_moe_blocks_with_names(model): + for block_name, block in find_moe_blocks_with_names(model): spec = shard_experts(block, ep_world_size, ep_rank, cfg) - patch_forward(block, ep_group, ep_world_size, cfg) + patch_forward(block, ep_group, ep_world_size, cfg, block_name) specs.append(spec) return specs @@ -160,6 +160,7 @@ def patch_forward( ep_group: dist.ProcessGroup, ep_world_size: int, cfg: ExpertParallelConfig, + block_name: str, ) -> None: """Replace the MoE block forward with EP-aware communication flow. @@ -222,12 +223,17 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): else: raise ValueError(f'Unsupported hidden_states ndim: {hidden_states.ndim}') + # R2 / R3 routing replay: pass block-level replay state + from .router_replay import get_replay_state + replay_state = get_replay_state(block_name) + router_logits, routing_weights, selected_experts = _run_router( gate=gate, hidden_states=hidden_states_2d, top_k=top_k, router_dtype=_get_router_dtype(cfg.router_dtype, hidden_states_2d.dtype), - norm_topk_prob=getattr(block, 'norm_topk_prob', False), + norm_topk_prob=_get_norm_topk_prob(block), + replay_state=replay_state, **kwargs, ) # Keep routing weights in activation dtype before unpermute weighting. @@ -400,6 +406,15 @@ def _get_top_k(block: nn.Module) -> int | None: return int(value) return None +def _get_norm_topk_prob(block: nn.Module) -> bool: + # fix: get norm_topk_prob from gate + gate = _get_gate(block) + if gate is not None and hasattr(gate, 'norm_topk_prob'): + value = getattr(gate, 'norm_topk_prob') + if value is not None: + return bool(value) + # default retrun True + return True def _get_router_dtype(router_dtype: str, default_dtype: torch.dtype) -> torch.dtype: if router_dtype == 'fp32': @@ -419,13 +434,32 @@ def _maybe_run_shared_expert(block: nn.Module, hidden_states_2d: torch.Tensor, c shared = getattr(block, 'shared_experts', None) if shared is None: return None - return _run_module_with_casting(shared, hidden_states_2d) + shared_output = _run_module_with_casting(shared, hidden_states_2d) + # fix: add shared_expert_gate for Qwen3_5Moe + shared_gate = getattr(block, 'shared_expert_gate', None) + if shared_gate is not None: + shared_output = F.sigmoid(shared_gate(hidden_states_2d)) * shared_output + return shared_output def _is_moe_experts(experts: Any) -> bool: - if isinstance(experts, nn.ModuleList): + unwrapped = experts + # Look through PEFT / LoRA wrappers that may wrap the experts module + # or its parameters (e.g. LoraLayer, ParamWrapper). + while True: + previous = unwrapped + if hasattr(unwrapped, 'base_layer'): + unwrapped = unwrapped.base_layer + elif hasattr(unwrapped, 'module'): + unwrapped = unwrapped.module + elif hasattr(unwrapped, '_fsdp_wrapped_module'): + unwrapped = unwrapped._fsdp_wrapped_module + if unwrapped is previous: + break + + if isinstance(unwrapped, nn.ModuleList): return True - if hasattr(experts, 'gate_up_proj') and hasattr(experts, 'down_proj'): + if hasattr(unwrapped, 'gate_up_proj') and hasattr(unwrapped, 'down_proj'): return True return False @@ -507,21 +541,48 @@ def _run_router( top_k: int, router_dtype: torch.dtype, norm_topk_prob: bool, + replay_state: Any = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: gate_kwargs = {} if 'input_ids' in kwargs and _module_forward_accepts_kwarg(gate, 'input_ids'): gate_kwargs['input_ids'] = kwargs['input_ids'] gate_out = gate(hidden_states, **gate_kwargs) + + # Resolve router_logits once — needed by both REPLAY and normal paths. if isinstance(gate_out, tuple) and len(gate_out) >= 3: - router_logits, routing_weights, selected_experts = gate_out[:3] + router_logits = gate_out[0] + else: + router_logits = gate_out + + # Lazy import to avoid circular dependency with router_replay.py + from .router_replay import RouterReplayAction + + # --- REPLAY_FORWARD: use injected selected_experts --- + if (replay_state is not None + and replay_state.action != RouterReplayAction.RECORD + and replay_state.target_indices is not None): + selected_experts = replay_state.target_indices + routing_weights = torch.softmax(router_logits, dim=-1, dtype=router_dtype) + routing_weights = routing_weights.gather(-1, selected_experts) + if norm_topk_prob: + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) return router_logits, routing_weights, selected_experts - router_logits = gate_out - routing_weights = torch.softmax(router_logits, dim=-1, dtype=router_dtype) - routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) - if norm_topk_prob: - routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + # --- Normal path: compute routing_weights and selected_experts --- + if isinstance(gate_out, tuple) and len(gate_out) >= 3: + _, routing_weights, selected_experts = gate_out[:3] + else: + routing_weights = torch.softmax(router_logits, dim=-1, dtype=router_dtype) + routing_weights, selected_experts = torch.topk(routing_weights, top_k, dim=-1) + if norm_topk_prob: + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + + # --- RECORD: save selected_experts after computing them --- + if (replay_state is not None + and replay_state.action == RouterReplayAction.RECORD): + replay_state.recorded_indices = selected_experts.detach().clone() + return router_logits, routing_weights, selected_experts diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py new file mode 100644 index 000000000..06e01ecf6 --- /dev/null +++ b/src/twinkle/model/transformers/moe/router_replay.py @@ -0,0 +1,281 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +""" +FSDP routing replay utilities for HF Transformers MoE models. + +Provides RouterReplayAction, per-MoE-block replay state, and functions for +recording / replaying expert routing decisions during GRPO training. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Optional, Tuple + +import inspect +import torch +import torch.nn as nn + +from twinkle import Platform +from twinkle.utils import get_logger + +logger = get_logger() + + +class RouterReplayAction(Enum): + """A Enum to define the actions for router replay.""" + + RECORD = "record" # Record the topk indices for replay + REPLAY_FORWARD = "replay_forward" # Replay the recorded topk indices for forward pass + REPLAY_BACKWARD = "replay_backward" # Replay topk indices for re-compute during backward pass + +# --------------------------------------------------------------------------- +# Global registry: {block_name: _RouterReplayState} +# --------------------------------------------------------------------------- + +@dataclass +class _RouterReplayState: + """Per-MoE-block replay state, stored in the global ``_registry``.""" + + action: RouterReplayAction = None + recorded_indices: Optional[torch.Tensor] = None # [num_tokens, topk] + target_indices: Optional[torch.Tensor] = None # [num_tokens, topk] + +_registry: Dict[str, _RouterReplayState] = {} + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def set_global_router_replay_action(action: RouterReplayAction) -> None: + """Set *action* on every registered MoE block.""" + for state in _registry.values(): + state.action = action + +def clear_global_router_replay_action() -> None: + """Reset action to None on every registered MoE block.""" + for state in _registry.values(): + state.action = None + +def clear_global_indices() -> None: + """Clear recorded / target indices on every registered MoE block.""" + for state in _registry.values(): + state.recorded_indices = None + state.target_indices = None + +def get_replay_state(block_name: str) -> Optional[_RouterReplayState]: + """Return the replay state for *block_name*, or *None*.""" + return _registry.get(block_name) + +def set_router_replay_data( + routed_experts: torch.Tensor, + model: nn.Module, +) -> None: + """Slice *routed_experts* into per-block + ``target_indices`` and inject them into the registered MoE blocks of *model*. + + Each block receives a ``[num_tokens, topk]`` slice covering the tokens + processed by that layer. + """ + if routed_experts is None: + return + + blocks = _find_moe_blocks_with_names(model) + if not blocks: + return + + if routed_experts.dim() != 4: + raise ValueError( + f'Expected routed_experts with shape [bs, seq_len, layers, topk], ' + f'got {tuple(routed_experts.shape)}' + ) + + # SP: slice full-sequence routed_experts to local SP rank tokens. + from ..strategy.sequence_parallel import sequence_parallel as sp + sp_world_size = getattr(sp, 'sp_world_size', None) or 1 + if sp_world_size > 1: + _, _, _, _, _, _, extra_values = sp.pad_and_split_inputs( + None, + None, + None, + None, + None, + None, + real_position_ids = sp.real_position_ids, + extra_split_values=[(routed_experts, 0, 1)]) + routed_experts = extra_values[0] + + # [bs, seq_len, num_moe_layers, topk] -> [total_seq, num_moe_layers, topk] + routed_experts = routed_experts.flatten(0, 1).to(Platform.get_local_device()) + + num_layers_in_data = routed_experts.shape[1] + + for layer_idx, (name, _) in enumerate(blocks): + state = _registry.get(name) + if state is None: + continue + if layer_idx >= num_layers_in_data: + break + + # Each layer gets [num_tokens, topk] + target = routed_experts[:, layer_idx, :].to(torch.int64) + if target.numel() > 0: + state.target_indices = target + +def get_router_replay_data(model: nn.Module, batch_size=1) -> Optional[torch.Tensor]: + """Collect ``recorded_indices`` from all registered MoE blocks in *model*. + + . note:: + *ep_group* is accepted for signature parity with the Megatron + version but is unused in FSDP — routing runs before the EP + all-to-all so every EP rank records the same indices. + + Returns *None* when no MoE blocks have recorded routing data. + """ + blocks = _find_moe_blocks_with_names(model) + if not blocks: + return None + + layers = [] + for name, _ in blocks: + state = _registry.get(name) + if state is not None and state.recorded_indices is not None: + layers.append(state.recorded_indices) # each: [num_tokens, topk] + + if not layers: + return None + + # Stack: [num_tokens, num_layers, topk] + routed_experts = torch.stack(layers, dim=1) + _, num_layers, topk = routed_experts.shape + # [num_tokens, num_layers, topk] -> [bs, local_seq_len, num_layers, topk] + routed_experts = routed_experts.reshape(batch_size, -1, num_layers, topk) + + # SP: all-gather local routing data across SP ranks along seq_len dim. + from ..strategy.sequence_parallel import sequence_parallel as sp + sp_world_size = getattr(sp, 'sp_world_size', None) or 1 + if sp_world_size > 1: + routed_experts = sp.gather(routed_experts, dim=1, position_ids=sp.real_position_ids) + + # [bs, seq_len, num_layers, topk] + return routed_experts + +def apply_router_replay_patch(model: nn.Module) -> None: + """Register MoE blocks and (for EP=1) wrap their forwards through + ``_run_router()`` so that routing replay works on the HF native path. + + When EP > 1, ``expert_parallel.py`` already patches the forward through + ``_run_router()`` — only registration is needed. + """ + if getattr(model, '_rr_patched', False): + return + + blocks = _find_moe_blocks_with_names(model) + if not blocks: + return + + # Walk *model*, find every MoE block and register it in ``_registry``. + # Safe to call multiple times — already-registered blocks are skipped. + for name, _ in blocks: + if name not in _registry: + _registry[name] = _RouterReplayState() + + # Determine if EP patches are already in place. + # When EP > 1 the block forward has already been replaced by + # patch_forward(); we only need to ensure replay_state is wired in. + # When EP = 1 the original HF forward is intact — wrap it. + + # Check whether the first block's forward has already been EP-patched + _, first_block = blocks[0] + if _is_ep_patched(first_block): + logger.debug( + 'EP patches detected — routing replay piggy-backs on ' + 'patch_forward() replay_state wiring.' + ) + return + + # EP = 1: wrap each MoE block forward through _run_router() + logger.info('Applying FSDP routing replay patch (EP=1 mode).') + _wrap_all_moe_blocks(blocks) + model._rr_patched = True + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _is_ep_patched(block: nn.Module) -> bool: + """Return True if *block* has already been patched by expert_parallel.""" + return getattr(block, '_ep_patched', False) + +def _find_moe_blocks_with_names(model: nn.Module) -> List[tuple[str, nn.Module]]: + from .expert_parallel import find_moe_blocks_with_names + # Strip PEFT wrapper so block names match those used by EP patch_forward + unwrap = model.get_base_model() if hasattr(model, 'get_base_model') else model + blocks = list(find_moe_blocks_with_names(unwrap)) + return blocks + +def _wrap_all_moe_blocks(blocks: List[tuple[str, nn.Module]],) -> None: + """Replace each MoE block's forward with a wrapper that calls + ``_run_router()`` instead of the original gate logic.""" + from .expert_parallel import (_get_gate, _get_norm_topk_prob, _get_top_k, _run_router, + _maybe_run_shared_expert, ExpertParallelConfig) + import types + + for name, block in blocks: + gate = _get_gate(block) + if gate is None: + raise ValueError('MoE block must define gate/router module.') + + top_k = _get_top_k(block) + if top_k is None: + raise ValueError('MoE block must define top_k/num_experts_per_tok.') + + norm_topk_prob = _get_norm_topk_prob(block) + + original_forward = block.forward + return_annotation = inspect.signature(original_forward).return_annotation + returns_router_logits = return_annotation in ( + tuple, + Tuple[torch.Tensor, torch.Tensor | None], + ) + + def _make_patched_forward(_name, _block, _gate, _top_k, _norm_topk_prob, _returns_router_logits): + def patched_forward(self, hidden_states): + orig_shape = hidden_states.shape + if hidden_states.ndim == 3: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_2d = hidden_states.view(-1, hidden_dim) + elif hidden_states.ndim == 2: + batch_size, sequence_length = 1, hidden_states.shape[0] + hidden_dim = hidden_states.shape[1] + hidden_states_2d = hidden_states + else: + raise ValueError(f'Unsupported hidden_states ndim: {hidden_states.ndim}') + + replay_state = _registry.get(_name) + router_logits, routing_weights, selected_experts = _run_router( + gate=_gate, + hidden_states=hidden_states_2d, + top_k=_top_k, + router_dtype=hidden_states_2d.dtype, + norm_topk_prob=_norm_topk_prob, + replay_state=replay_state, + ) + routed_output = self.experts( + hidden_states_2d, selected_experts, routing_weights + ) + + shared_out = _maybe_run_shared_expert(_block, hidden_states_2d, ExpertParallelConfig()) + if shared_out is not None: + routed_output = routed_output + shared_out + + if len(orig_shape) == 3: + routed_output = routed_output.reshape(batch_size, sequence_length, hidden_dim) + if _returns_router_logits: + return routed_output, router_logits + return routed_output + + return patched_forward + + routed_fn = _make_patched_forward(name, block, gate, top_k, norm_topk_prob, returns_router_logits) + block.forward = types.MethodType(routed_fn, block) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 1c312a8e3..33b024351 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -191,6 +191,8 @@ def __init__( self._fsdp_config = dict(fsdp_config or {}) self._ddp_config = ddp_config or {} self._memory_efficient_init = memory_efficient_init + self._router_replay_enabled = bool(kwargs.pop('enable_router_replay', False)) + self._router_replay_applied = False self._decide_strategy(strategy) self.grad_scaler_config = grad_scaler_config if model_id is not None: @@ -367,6 +369,52 @@ def _maybe_apply_expert_parallel(self): ) self._expert_parallel_applied = True + def _maybe_apply_router_replay(self): + """Lazily register MoE blocks and (for EP=1) wrap their forwards.""" + if not self._router_replay_enabled or self._router_replay_applied: + return + model = self.strategy.unwrap_model(self.model) + from .moe.router_replay import apply_router_replay_patch + apply_router_replay_patch(model) + self._router_replay_applied = True + + def _router_replay_setup(self, router_replay_action, routed_experts=None, + batch_size=1, manual_cleanup=False): + """Set up routing replay before a model forward. + + Returns ``cleanup_fn``. + Call *cleanup_fn* after the forward to gather recorded indices(when RECORD) and clear global state. + """ + if not self._router_replay_enabled: + return lambda: None + + self._maybe_apply_router_replay() + + if router_replay_action is None: + return lambda: None + + from .moe.router_replay import ( + set_router_replay_data, set_global_router_replay_action, + clear_global_router_replay_action, clear_global_indices, + RouterReplayAction, get_router_replay_data, + ) + unwrapped = self.strategy.unwrap_model(self.model) + set_global_router_replay_action(router_replay_action) + if router_replay_action == RouterReplayAction.REPLAY_FORWARD: + assert routed_experts is not None, f'routed_experts must be not None' + set_router_replay_data(routed_experts, unwrapped) + + def cleanup(): + recorded = None + if router_replay_action == RouterReplayAction.RECORD: + recorded = get_router_replay_data(unwrapped, batch_size) + if not manual_cleanup: + clear_global_router_replay_action() + clear_global_indices() + return recorded + + return cleanup + def _ensure_optimizer_dp_groups(self): for optimizer_group in self.optimizer_group.values(): if not isinstance(optimizer_group, OptimizerGroup): @@ -427,8 +475,19 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec ) labels: torch.Tensor = inputs.pop('labels', None) optimizer_config.accumulate_metrics(True) + + # Routing replay: respects router_replay_action regardless of caller + router_replay_action = kwargs.pop('router_replay_action', None) + router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) + routed_experts = inputs.pop('routed_experts', None) + rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, + labels.shape[0], router_replay_manual_cleanup) + with _resolve_task_context(self.model, task): outputs = self.model(**inputs) + + recorded_routing = rr_cleanup() + inputs['labels'] = labels if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = (labels != -100).bool() @@ -454,6 +513,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec return_outputs = copy(outputs) if not return_logits: return_outputs['logits'] = None + if recorded_routing is not None: + return_outputs['routed_experts'] = recorded_routing return return_outputs @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) @@ -503,11 +564,22 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T labels = inputs.pop('labels', None) optimizer_config.accumulate_metrics(False) unwrapped_model = self.strategy.unwrap_model(self.model) + + # Routing replay: handle any action generically + router_replay_action = kwargs.pop('router_replay_action', None) + router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) + routed_experts = inputs.pop('routed_experts', None) + rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, + labels.shape[0], router_replay_manual_cleanup) + lora_ctx = ( unwrapped_model.disable_adapter() if disable_lora and isinstance(unwrapped_model, PeftModel) else contextlib.nullcontext()) with _resolve_task_context(self.model, task), lora_ctx: outputs = self.model(**inputs) + + recorded_routing = rr_cleanup() + inputs['labels'] = labels if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = (labels != -100).bool() @@ -533,6 +605,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T return_outputs = copy(outputs) if not return_logits: return_outputs['logits'] = None + if recorded_routing is not None: + return_outputs['routed_experts'] = recorded_routing return return_outputs @remote_function(collect='mean') @@ -612,11 +686,18 @@ def backward(self, **kwargs): if not should_sync and hasattr(self.model, 'no_sync'): no_sync_ctx = self.model.no_sync() + router_replay_action = kwargs.pop('router_replay_action', None) + rr_cleanup = lambda: None + if router_replay_action is not None: + from .moe.router_replay import RouterReplayAction + rr_cleanup = self._router_replay_setup(router_replay_action=RouterReplayAction.REPLAY_BACKWARD) + with no_sync_ctx: if scaler is not None: scaler.scale(loss_value).backward() else: loss_value.backward() + rr_cleanup() optimizer_config.train_status.loss_value = None @@ -634,7 +715,7 @@ def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Tr Returns: The output of the model forward. """ - outputs = self.forward(inputs=inputs, **kwargs) + outputs = self.forward(inputs=inputs, router_replay_manual_cleanup=True, **kwargs) loss = self.calculate_loss(**kwargs) outputs['loss'] = loss self.backward(**kwargs) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index d5e7893f4..e6a061cba 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -43,6 +43,7 @@ class InputProcessor: 'video_grid_thw': 0, 'input_features': 0.0, 'feature_attention_mask': 0, + 'routed_experts': 0, } # VLM fields to concatenate (not pad) in batch @@ -69,6 +70,7 @@ def __init__(self, self.framework = framework self.process_pipeline = [ self.prepare_inputs, + self.align_routed_experts, self.pad_cp, self.collate_fn, self.to_transformers_dict, @@ -114,7 +116,7 @@ def to_tensor(_input): if value is None: continue value = torch.tensor(value) - elif (isinstance(value, list)) and key in ('completion_mask', 'mm_token_type_ids'): + elif (isinstance(value, list)) and key in ('completion_mask', 'mm_token_type_ids', 'routed_experts'): value = torch.tensor(value) elif key in self.VLM_CONCAT_FIELDS: if not isinstance(value[0], torch.Tensor): @@ -447,10 +449,11 @@ def drop_causal_4d_mask(self, inputs: List[InputFeature], **kwargs) -> List[Inpu return inputs @staticmethod - def _pad_sequence(sequences, padding_value, padding_side): + def _pad_sequence(sequences, padding_value, padding_side, concat=None): if padding_side == 'right': from twinkle.utils import pad_and_stack_tensors - return pad_and_stack_tensors(sequences, pad_value=padding_value, concat=sequences[0].dim() >= 2) + return pad_and_stack_tensors(sequences, pad_value=padding_value, + concat=concat if concat is not None else (sequences[0].dim() >= 2)) else: # left padding import torch @@ -650,6 +653,7 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat 'max_length_q', 'max_length_k', 'packed_seq_params', + 'routed_experts', ] + list(InputProcessor.VLM_CONCAT_FIELDS) for key in list(_input.keys()): if key not in _keys: @@ -716,7 +720,8 @@ def is_mm_position_ids(position_ids): num_axes = values[0].shape[0] result[key] = result[key].reshape(len(values), num_axes, -1).permute(1, 0, 2).contiguous() elif isinstance(values[0], torch.Tensor): - result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side) + concat = False if (key == 'routed_experts') else None + result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side, concat) if result[key].dim() == 1: result[key] = result[key].unsqueeze(0) else: @@ -799,3 +804,28 @@ def postprocess_tensor_cp(self, tensor, cu_seqlens=None): from twinkle.utils.torch_utils import gather_cp_load_balanced return gather_cp_load_balanced(tensor, mpu.get_context_parallel_group(), seq_dim=1, cu_seqlens=cu_seqlens) + + def align_routed_experts(self, inputs: Union[List[InputFeature], InputFeature], **kwargs) -> List[InputFeature]: + + def align_to(_input): + routed_experts = _input.get('routed_experts', None) + input_seq_len = _input.get('length', None) + if input_seq_len is None: + input_ids = _input.get('input_ids', None) + input_seq_len = input_ids.shape[1] if input_ids is not None else 0 + if routed_experts is not None: + # The number of experts in the output can be 1 less than (prompt_length + response_token_count) + # This gap of 1 is expected + # For more details, please refer PR https://github.com/vllm-project/vllm/pull/28284 + experts_seq_len = routed_experts.shape[0] + # Padding routed_experts(seq_len, layer_num, topk) seq_len to match the seq_len of the input_ids + padding_routed_experts = routed_experts + padding_len = input_seq_len - experts_seq_len + if padding_len > 0: + padding_routed_experts = torch.nn.functional.pad(routed_experts, (0, 0, 0, 0, 0, padding_len), + 'constant', self.padding_map.get('routed_experts', 0)) + _input['routed_experts'] = padding_routed_experts.unsqueeze(0) + + return _input + + return [align_to(_input) for _input in inputs] \ No newline at end of file diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index 0ed9d715e..64ba9eb1e 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -312,10 +312,12 @@ async def sample(self, stop_reason: StopReason = _map_finish_reason(output.finish_reason) + routed_experts = getattr(output, 'routed_experts', None) sequences.append(SampledSequence( stop_reason=stop_reason, tokens=token_ids, logprobs=seq_logprobs, + routed_experts=routed_experts, )) # Extract prompt logprobs if requested diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 0433ef5a8..cac2e843a 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -251,18 +251,30 @@ async def _sample_single( sequences = [] for seq in response.sequences: if logprobs_only: + new_input_feature = _convert_ndarray_to_list(feat) + routed_experts = None + if seq.routed_experts is not None: + routed_experts = _convert_ndarray_to_list(seq.routed_experts) + new_input_feature['routed_experts'] = routed_experts sampled_seq = SampledSequence( tokens=[], stop_reason=seq.stop_reason, - new_input_feature=_convert_ndarray_to_list(feat), + new_input_feature=new_input_feature, + routed_experts=routed_experts, ) else: + new_input_feature = _convert_ndarray_to_list(self.template.concat_input_feature(feat, seq.tokens)) + routed_experts = None + if seq.routed_experts is not None: + routed_experts = _convert_ndarray_to_list(seq.routed_experts) + new_input_feature['routed_experts'] = routed_experts sampled_seq = SampledSequence( stop_reason=seq.stop_reason, tokens=seq.tokens, logprobs=seq.logprobs, decoded=self.template.decode(seq.tokens), - new_input_feature=_convert_ndarray_to_list(self.template.concat_input_feature(feat, seq.tokens)), + new_input_feature=new_input_feature, + routed_experts=routed_experts, ) sequences.append(sampled_seq) return SampleResponse( From 7fb0210057326a6e3feb0ff39da8af1bf3816a1b Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Tue, 30 Jun 2026 19:01:30 +0800 Subject: [PATCH 2/6] fsdp rr --- src/twinkle/sampler/vllm_sampler/vllm_sampler.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index cac2e843a..3c7b2f686 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -252,29 +252,23 @@ async def _sample_single( for seq in response.sequences: if logprobs_only: new_input_feature = _convert_ndarray_to_list(feat) - routed_experts = None if seq.routed_experts is not None: - routed_experts = _convert_ndarray_to_list(seq.routed_experts) - new_input_feature['routed_experts'] = routed_experts + new_input_feature['routed_experts'] = _convert_ndarray_to_list(seq.routed_experts) sampled_seq = SampledSequence( tokens=[], stop_reason=seq.stop_reason, new_input_feature=new_input_feature, - routed_experts=routed_experts, ) else: new_input_feature = _convert_ndarray_to_list(self.template.concat_input_feature(feat, seq.tokens)) - routed_experts = None if seq.routed_experts is not None: - routed_experts = _convert_ndarray_to_list(seq.routed_experts) - new_input_feature['routed_experts'] = routed_experts + new_input_feature['routed_experts'] = _convert_ndarray_to_list(seq.routed_experts) sampled_seq = SampledSequence( stop_reason=seq.stop_reason, tokens=seq.tokens, logprobs=seq.logprobs, decoded=self.template.decode(seq.tokens), new_input_feature=new_input_feature, - routed_experts=routed_experts, ) sequences.append(sampled_seq) return SampleResponse( From bf32af7b1e7a00de840b352a68169d49ca9143eb Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Mon, 6 Jul 2026 20:31:18 +0800 Subject: [PATCH 3/6] fix --- .../rl/grpo/short_math_grpo_routing_replay.py | 3 ++- .../model/transformers/moe/router_replay.py | 19 ++++++++++++++++--- .../model/transformers/transformers.py | 8 ++++++-- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/cookbook/rl/grpo/short_math_grpo_routing_replay.py b/cookbook/rl/grpo/short_math_grpo_routing_replay.py index 8297239f3..931030396 100644 --- a/cookbook/rl/grpo/short_math_grpo_routing_replay.py +++ b/cookbook/rl/grpo/short_math_grpo_routing_replay.py @@ -297,7 +297,8 @@ def main(): mb_inputs = all_input_data[mb_start:mb_end] mb_old_logps = all_old_logps[mb_start:mb_end] mb_advantages = advantages[mb_start:mb_end] - recompute_logps = [input.pop('old_logps').unsqueeze(0) for input in mb_inputs] + for input in mb_inputs: + input.pop('old_logps', None) mb_output = model.forward_backward( inputs=mb_inputs, diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py index 06e01ecf6..8de0b915b 100644 --- a/src/twinkle/model/transformers/moe/router_replay.py +++ b/src/twinkle/model/transformers/moe/router_replay.py @@ -261,9 +261,22 @@ def patched_forward(self, hidden_states): norm_topk_prob=_norm_topk_prob, replay_state=replay_state, ) - routed_output = self.experts( - hidden_states_2d, selected_experts, routing_weights - ) + + if isinstance(self.experts, nn.ModuleList): + routed_output = torch.zeros_like(hidden_states_2d) + num_experts = len(self.experts) + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) + current_state = hidden_states_2d[top_x] + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + routed_output.index_add_(0, top_x, current_hidden_states.to(hidden_states_2d.dtype)) + else: + routed_output = self.experts( + hidden_states_2d, selected_experts, routing_weights + ) shared_out = _maybe_run_shared_expert(_block, hidden_states_2d, ExpertParallelConfig()) if shared_out is not None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index dbe7b2516..cafe6133e 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -480,8 +480,10 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec router_replay_action = kwargs.pop('router_replay_action', None) router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) routed_experts = inputs.pop('routed_experts', None) + batch_size = labels.shape[0] if labels is not None else ( + inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, - labels.shape[0], router_replay_manual_cleanup) + batch_size, router_replay_manual_cleanup) with _resolve_task_context(self.model, task): outputs = self.model(**inputs) @@ -569,8 +571,10 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T router_replay_action = kwargs.pop('router_replay_action', None) router_replay_manual_cleanup = kwargs.pop('router_replay_manual_cleanup', False) routed_experts = inputs.pop('routed_experts', None) + batch_size = labels.shape[0] if labels is not None else ( + inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, - labels.shape[0], router_replay_manual_cleanup) + batch_size, router_replay_manual_cleanup) lora_ctx = ( unwrapped_model.disable_adapter() From 9b27af7e1110b21babb508edee2791d85b5b632e Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Mon, 6 Jul 2026 20:48:35 +0800 Subject: [PATCH 4/6] fix --- .../rl/grpo/short_math_grpo_routing_replay.py | 2 +- .../model/transformers/moe/expert_parallel.py | 8 +- .../model/transformers/moe/router_replay.py | 73 +++++++++++-------- .../model/transformers/transformers.py | 18 ++--- src/twinkle/processor/base.py | 9 ++- 5 files changed, 59 insertions(+), 51 deletions(-) diff --git a/cookbook/rl/grpo/short_math_grpo_routing_replay.py b/cookbook/rl/grpo/short_math_grpo_routing_replay.py index 931030396..347458e12 100644 --- a/cookbook/rl/grpo/short_math_grpo_routing_replay.py +++ b/cookbook/rl/grpo/short_math_grpo_routing_replay.py @@ -332,4 +332,4 @@ def main(): if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 2a3cfe1a9..4b5010470 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -406,6 +406,7 @@ def _get_top_k(block: nn.Module) -> int | None: return int(value) return None + def _get_norm_topk_prob(block: nn.Module) -> bool: # fix: get norm_topk_prob from gate gate = _get_gate(block) @@ -416,6 +417,7 @@ def _get_norm_topk_prob(block: nn.Module) -> bool: # default retrun True return True + def _get_router_dtype(router_dtype: str, default_dtype: torch.dtype) -> torch.dtype: if router_dtype == 'fp32': return torch.float32 @@ -559,8 +561,7 @@ def _run_router( from .router_replay import RouterReplayAction # --- REPLAY_FORWARD: use injected selected_experts --- - if (replay_state is not None - and replay_state.action != RouterReplayAction.RECORD + if (replay_state is not None and replay_state.action != RouterReplayAction.RECORD and replay_state.target_indices is not None): selected_experts = replay_state.target_indices routing_weights = torch.softmax(router_logits, dim=-1, dtype=router_dtype) @@ -579,8 +580,7 @@ def _run_router( routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) # --- RECORD: save selected_experts after computing them --- - if (replay_state is not None - and replay_state.action == RouterReplayAction.RECORD): + if (replay_state is not None and replay_state.action == RouterReplayAction.RECORD): replay_state.recorded_indices = selected_experts.detach().clone() return router_logits, routing_weights, selected_experts diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py index 8de0b915b..3cfcce124 100644 --- a/src/twinkle/model/transformers/moe/router_replay.py +++ b/src/twinkle/model/transformers/moe/router_replay.py @@ -8,13 +8,12 @@ from __future__ import annotations -from dataclasses import dataclass -from enum import Enum -from typing import Dict, List, Optional, Tuple - import inspect import torch import torch.nn as nn +from dataclasses import dataclass +from enum import Enum +from typing import Dict, List, Optional, Tuple from twinkle import Platform from twinkle.utils import get_logger @@ -25,48 +24,56 @@ class RouterReplayAction(Enum): """A Enum to define the actions for router replay.""" - RECORD = "record" # Record the topk indices for replay - REPLAY_FORWARD = "replay_forward" # Replay the recorded topk indices for forward pass - REPLAY_BACKWARD = "replay_backward" # Replay topk indices for re-compute during backward pass + RECORD = 'record' # Record the topk indices for replay + REPLAY_FORWARD = 'replay_forward' # Replay the recorded topk indices for forward pass + REPLAY_BACKWARD = 'replay_backward' # Replay topk indices for re-compute during backward pass + # --------------------------------------------------------------------------- # Global registry: {block_name: _RouterReplayState} # --------------------------------------------------------------------------- + @dataclass class _RouterReplayState: """Per-MoE-block replay state, stored in the global ``_registry``.""" action: RouterReplayAction = None - recorded_indices: Optional[torch.Tensor] = None # [num_tokens, topk] - target_indices: Optional[torch.Tensor] = None # [num_tokens, topk] + recorded_indices: torch.Tensor | None = None # [num_tokens, topk] + target_indices: torch.Tensor | None = None # [num_tokens, topk] -_registry: Dict[str, _RouterReplayState] = {} + +_registry: dict[str, _RouterReplayState] = {} # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- + def set_global_router_replay_action(action: RouterReplayAction) -> None: """Set *action* on every registered MoE block.""" for state in _registry.values(): state.action = action + def clear_global_router_replay_action() -> None: """Reset action to None on every registered MoE block.""" for state in _registry.values(): state.action = None + def clear_global_indices() -> None: """Clear recorded / target indices on every registered MoE block.""" for state in _registry.values(): state.recorded_indices = None state.target_indices = None -def get_replay_state(block_name: str) -> Optional[_RouterReplayState]: + +def get_replay_state(block_name: str) -> _RouterReplayState | None: """Return the replay state for *block_name*, or *None*.""" return _registry.get(block_name) + def set_router_replay_data( routed_experts: torch.Tensor, model: nn.Module, @@ -85,10 +92,8 @@ def set_router_replay_data( return if routed_experts.dim() != 4: - raise ValueError( - f'Expected routed_experts with shape [bs, seq_len, layers, topk], ' - f'got {tuple(routed_experts.shape)}' - ) + raise ValueError(f'Expected routed_experts with shape [bs, seq_len, layers, topk], ' + f'got {tuple(routed_experts.shape)}') # SP: slice full-sequence routed_experts to local SP rank tokens. from ..strategy.sequence_parallel import sequence_parallel as sp @@ -101,7 +106,7 @@ def set_router_replay_data( None, None, None, - real_position_ids = sp.real_position_ids, + real_position_ids=sp.real_position_ids, extra_split_values=[(routed_experts, 0, 1)]) routed_experts = extra_values[0] @@ -122,7 +127,8 @@ def set_router_replay_data( if target.numel() > 0: state.target_indices = target -def get_router_replay_data(model: nn.Module, batch_size=1) -> Optional[torch.Tensor]: + +def get_router_replay_data(model: nn.Module, batch_size=1) -> torch.Tensor | None: """Collect ``recorded_indices`` from all registered MoE blocks in *model*. . note:: @@ -135,7 +141,7 @@ def get_router_replay_data(model: nn.Module, batch_size=1) -> Optional[torch.Ten blocks = _find_moe_blocks_with_names(model) if not blocks: return None - + layers = [] for name, _ in blocks: state = _registry.get(name) @@ -144,7 +150,7 @@ def get_router_replay_data(model: nn.Module, batch_size=1) -> Optional[torch.Ten if not layers: return None - + # Stack: [num_tokens, num_layers, topk] routed_experts = torch.stack(layers, dim=1) _, num_layers, topk = routed_experts.shape @@ -160,6 +166,7 @@ def get_router_replay_data(model: nn.Module, batch_size=1) -> Optional[torch.Ten # [bs, seq_len, num_layers, topk] return routed_experts + def apply_router_replay_patch(model: nn.Module) -> None: """Register MoE blocks and (for EP=1) wrap their forwards through ``_run_router()`` so that routing replay works on the HF native path. @@ -188,10 +195,8 @@ def apply_router_replay_patch(model: nn.Module) -> None: # Check whether the first block's forward has already been EP-patched _, first_block = blocks[0] if _is_ep_patched(first_block): - logger.debug( - 'EP patches detected — routing replay piggy-backs on ' - 'patch_forward() replay_state wiring.' - ) + logger.debug('EP patches detected — routing replay piggy-backs on ' + 'patch_forward() replay_state wiring.') return # EP = 1: wrap each MoE block forward through _run_router() @@ -199,28 +204,34 @@ def apply_router_replay_patch(model: nn.Module) -> None: _wrap_all_moe_blocks(blocks) model._rr_patched = True + # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- + def _is_ep_patched(block: nn.Module) -> bool: """Return True if *block* has already been patched by expert_parallel.""" return getattr(block, '_ep_patched', False) -def _find_moe_blocks_with_names(model: nn.Module) -> List[tuple[str, nn.Module]]: + +def _find_moe_blocks_with_names(model: nn.Module) -> list[tuple[str, nn.Module]]: from .expert_parallel import find_moe_blocks_with_names + # Strip PEFT wrapper so block names match those used by EP patch_forward unwrap = model.get_base_model() if hasattr(model, 'get_base_model') else model blocks = list(find_moe_blocks_with_names(unwrap)) return blocks -def _wrap_all_moe_blocks(blocks: List[tuple[str, nn.Module]],) -> None: + +def _wrap_all_moe_blocks(blocks: list[tuple[str, nn.Module]], ) -> None: """Replace each MoE block's forward with a wrapper that calls ``_run_router()`` instead of the original gate logic.""" - from .expert_parallel import (_get_gate, _get_norm_topk_prob, _get_top_k, _run_router, - _maybe_run_shared_expert, ExpertParallelConfig) import types + from .expert_parallel import (ExpertParallelConfig, _get_gate, _get_norm_topk_prob, _get_top_k, + _maybe_run_shared_expert, _run_router) + for name, block in blocks: gate = _get_gate(block) if gate is None: @@ -240,6 +251,7 @@ def _wrap_all_moe_blocks(blocks: List[tuple[str, nn.Module]],) -> None: ) def _make_patched_forward(_name, _block, _gate, _top_k, _norm_topk_prob, _returns_router_logits): + def patched_forward(self, hidden_states): orig_shape = hidden_states.shape if hidden_states.ndim == 3: @@ -265,7 +277,8 @@ def patched_forward(self, hidden_states): if isinstance(self.experts, nn.ModuleList): routed_output = torch.zeros_like(hidden_states_2d) num_experts = len(self.experts) - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) + expert_mask = torch.nn.functional.one_hot( + selected_experts, num_classes=num_experts).permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_layer = self.experts[expert_idx] @@ -274,9 +287,7 @@ def patched_forward(self, hidden_states): current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] routed_output.index_add_(0, top_x, current_hidden_states.to(hidden_states_2d.dtype)) else: - routed_output = self.experts( - hidden_states_2d, selected_experts, routing_weights - ) + routed_output = self.experts(hidden_states_2d, selected_experts, routing_weights) shared_out = _maybe_run_shared_expert(_block, hidden_states_2d, ExpertParallelConfig()) if shared_out is not None: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index cafe6133e..7d34230ce 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -378,8 +378,7 @@ def _maybe_apply_router_replay(self): apply_router_replay_patch(model) self._router_replay_applied = True - def _router_replay_setup(self, router_replay_action, routed_experts=None, - batch_size=1, manual_cleanup=False): + def _router_replay_setup(self, router_replay_action, routed_experts=None, batch_size=1, manual_cleanup=False): """Set up routing replay before a model forward. Returns ``cleanup_fn``. @@ -393,11 +392,8 @@ def _router_replay_setup(self, router_replay_action, routed_experts=None, if router_replay_action is None: return lambda: None - from .moe.router_replay import ( - set_router_replay_data, set_global_router_replay_action, - clear_global_router_replay_action, clear_global_indices, - RouterReplayAction, get_router_replay_data, - ) + from .moe.router_replay import (RouterReplayAction, clear_global_indices, clear_global_router_replay_action, + get_router_replay_data, set_global_router_replay_action, set_router_replay_data) unwrapped = self.strategy.unwrap_model(self.model) set_global_router_replay_action(router_replay_action) if router_replay_action == RouterReplayAction.REPLAY_FORWARD: @@ -482,8 +478,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec routed_experts = inputs.pop('routed_experts', None) batch_size = labels.shape[0] if labels is not None else ( inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) - rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, - batch_size, router_replay_manual_cleanup) + rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, batch_size, + router_replay_manual_cleanup) with _resolve_task_context(self.model, task): outputs = self.model(**inputs) @@ -573,8 +569,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T routed_experts = inputs.pop('routed_experts', None) batch_size = labels.shape[0] if labels is not None else ( inputs['input_ids'].shape[0] if 'input_ids' in inputs else 1) - rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, - batch_size, router_replay_manual_cleanup) + rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, batch_size, + router_replay_manual_cleanup) lora_ctx = ( unwrapped_model.disable_adapter() diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index e6a061cba..7b1f9df5d 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -452,8 +452,8 @@ def drop_causal_4d_mask(self, inputs: List[InputFeature], **kwargs) -> List[Inpu def _pad_sequence(sequences, padding_value, padding_side, concat=None): if padding_side == 'right': from twinkle.utils import pad_and_stack_tensors - return pad_and_stack_tensors(sequences, pad_value=padding_value, - concat=concat if concat is not None else (sequences[0].dim() >= 2)) + return pad_and_stack_tensors( + sequences, pad_value=padding_value, concat=concat if concat is not None else (sequences[0].dim() >= 2)) else: # left padding import torch @@ -823,9 +823,10 @@ def align_to(_input): padding_len = input_seq_len - experts_seq_len if padding_len > 0: padding_routed_experts = torch.nn.functional.pad(routed_experts, (0, 0, 0, 0, 0, padding_len), - 'constant', self.padding_map.get('routed_experts', 0)) + 'constant', + self.padding_map.get('routed_experts', 0)) _input['routed_experts'] = padding_routed_experts.unsqueeze(0) return _input - return [align_to(_input) for _input in inputs] \ No newline at end of file + return [align_to(_input) for _input in inputs] From 252630b5e5b7c799d57747683fb70da44cff4775 Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Mon, 6 Jul 2026 20:56:15 +0800 Subject: [PATCH 5/6] fix --- src/twinkle/model/transformers/transformers.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 7d34230ce..c8fec86da 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -397,7 +397,7 @@ def _router_replay_setup(self, router_replay_action, routed_experts=None, batch_ unwrapped = self.strategy.unwrap_model(self.model) set_global_router_replay_action(router_replay_action) if router_replay_action == RouterReplayAction.REPLAY_FORWARD: - assert routed_experts is not None, f'routed_experts must be not None' + assert routed_experts is not None, 'routed_experts must be not None' set_router_replay_data(routed_experts, unwrapped) def cleanup(): @@ -687,7 +687,7 @@ def backward(self, **kwargs): no_sync_ctx = self.model.no_sync() router_replay_action = kwargs.pop('router_replay_action', None) - rr_cleanup = lambda: None + rr_cleanup = None if router_replay_action is not None: from .moe.router_replay import RouterReplayAction rr_cleanup = self._router_replay_setup(router_replay_action=RouterReplayAction.REPLAY_BACKWARD) @@ -697,7 +697,9 @@ def backward(self, **kwargs): scaler.scale(loss_value).backward() else: loss_value.backward() - rr_cleanup() + + if rr_cleanup is not None: + rr_cleanup() optimizer_config.train_status.loss_value = None From 8dd1851f007f91df18822881aca09b78638249cb Mon Sep 17 00:00:00 2001 From: XianlongLi <2286061024@qq.com> Date: Tue, 7 Jul 2026 11:33:42 +0800 Subject: [PATCH 6/6] fix --- .../model/transformers/moe/expert_parallel.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 4b5010470..8fd0e3df4 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -408,14 +408,14 @@ def _get_top_k(block: nn.Module) -> int | None: def _get_norm_topk_prob(block: nn.Module) -> bool: - # fix: get norm_topk_prob from gate - gate = _get_gate(block) - if gate is not None and hasattr(gate, 'norm_topk_prob'): - value = getattr(gate, 'norm_topk_prob') - if value is not None: - return bool(value) - # default retrun True - return True + value = getattr(block, 'norm_topk_prob', None) + if value is None: + # fix: Attempt to fetch from the gate if block does not exist + gate = _get_gate(block) + value = getattr(gate, 'norm_topk_prob', None) + # Default return True. + # For Qwen3-5MoE, the `norm_topk_prob` attribute does not exist, and normalization is performed by default. + return bool(value) if value is not None else True def _get_router_dtype(router_dtype: str, default_dtype: torch.dtype) -> torch.dtype: