# KTMoELayerWrapper — nn.Module replacing HF MoE layers for SFT # SPDX-License-Identifier: Apache-2.0 """ KTMoELayerWrapper: drop-in nn.Module replacement for HuggingFace MoE layers. Delegates expert computation to the C++ KTMoEWrapper backend, with support for gradient checkpointing, PEFT LoRA on experts, LoRA Experts (separate small MLPs on GPU), shared experts, and multi-GPU rank-0-only execution. """ from __future__ import annotations import logging import math import os from contextlib import nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch import torch.nn as nn import torch.nn.functional as F if TYPE_CHECKING: from .lora import LoRAExperts from .arch import MOEArchConfig from .autograd import ( KTMoEFunction, _poison_checkpoint_cache, _sync_any_rank_exception, _sync_rank0_exception, ) from .config import KTActivationPolicy from .dist_utils import ( _all_gather_checkpoint_state, _all_gather_qlens, _checkpoint_hook_mode, _dist_gather_varlen_to_rank0, _dist_scatter_varlen_from_rank0, _qlen_offsets, ) logger = logging.getLogger(__name__) _KT_SFT_DEBUG = os.environ.get("KT_SFT_DEBUG", "0") == "1" class _DeepseekRouterLinear(nn.Linear): """PEFT-compatible proxy preserving DeepSeek's FP32 gate projection.""" def forward(self, input: torch.Tensor) -> torch.Tensor: weight = self.weight.float() bias = self.bias.float() if self.bias is not None else None return F.linear(input.float(), weight, bias) @dataclass(frozen=True) class DeepSeekRoutingSpec: """Validated routing semantics shared by remote and native DeepSeek MoE.""" expert_num: int n_group: int topk_group: int top_k: int norm_topk_prob: bool routed_scaling_factor: float @classmethod def from_modules( cls, *, moe: nn.Module, router: nn.Module, expert_num: int, ) -> "DeepSeekRoutingSpec": sources = [ ("moe", moe), ("router", router), ("moe.config", getattr(moe, "config", None)), ("router.config", getattr(router, "config", None)), ] def resolve(name: str, aliases: tuple[str, ...]) -> Any: observed: list[tuple[str, Any]] = [] for source_name, source in sources: if source is None: continue for alias in aliases: if hasattr(source, alias): value = getattr(source, alias) if value is not None: observed.append((f"{source_name}.{alias}", value)) break if not observed: raise ValueError( f"DeepSeek routing metadata {name!r} is missing; checked " + ", ".join(source_name for source_name, _ in sources) ) first_source, first_value = observed[0] for source_name, value in observed[1:]: if value != first_value: raise ValueError( f"Conflicting DeepSeek routing metadata for {name}: " f"{first_source}={first_value!r}, {source_name}={value!r}" ) return first_value def positive_int(name: str, aliases: tuple[str, ...]) -> int: value = resolve(name, aliases) if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise ValueError( f"DeepSeek routing metadata {name} must be a positive integer, " f"got {value!r}" ) return int(value) n_group = positive_int("n_group", ("n_group",)) topk_group = positive_int("topk_group", ("topk_group",)) routed_expert_num = positive_int( "n_routed_experts", ("n_routed_experts", "num_local_experts"), ) top_k = positive_int( "top_k", ("top_k", "num_experts_per_tok"), ) norm_topk_prob = resolve("norm_topk_prob", ("norm_topk_prob",)) if not isinstance(norm_topk_prob, bool): raise ValueError( "DeepSeek routing metadata norm_topk_prob must be bool, " f"got {norm_topk_prob!r}" ) routed_scaling_factor = float( resolve("routed_scaling_factor", ("routed_scaling_factor",)) ) if not math.isfinite(routed_scaling_factor) or routed_scaling_factor <= 0: raise ValueError( "DeepSeek routed_scaling_factor must be finite and positive, " f"got {routed_scaling_factor!r}" ) expert_num = int(expert_num) if routed_expert_num != expert_num: raise ValueError( "DeepSeek routed expert count does not match the KT architecture: " f"routing={routed_expert_num}, kt={expert_num}" ) if expert_num <= 0 or expert_num % n_group: raise ValueError( "DeepSeek expert_num must be positive and divisible by n_group, " f"got expert_num={expert_num}, n_group={n_group}" ) experts_per_group = expert_num // n_group if experts_per_group < 2: raise ValueError( "DeepSeek group-top2 routing requires at least two experts per group" ) if topk_group > n_group: raise ValueError( f"DeepSeek topk_group={topk_group} exceeds n_group={n_group}" ) if top_k > topk_group * experts_per_group: raise ValueError( f"DeepSeek top_k={top_k} exceeds the {topk_group * experts_per_group} " "experts available in selected groups" ) return cls( expert_num=expert_num, n_group=n_group, topk_group=topk_group, top_k=top_k, norm_topk_prob=norm_topk_prob, routed_scaling_factor=routed_scaling_factor, ) def _activation_checkpoint_action(policy: KTActivationPolicy, phase: str) -> str: """Map public activation policy and checkpoint phase to collective control.""" if policy.gpu == "recompute" and policy.cpu == "retain": if phase == "first_forward": return "cache_first_forward" if phase == "recompute": return "reuse_recompute" return "normal" def _validate_activation_checkpoint_phase( policy: KTActivationPolicy, phase: str, *, requires_backward: bool = False, ) -> None: if phase not in ("none", "first_forward", "recompute"): raise RuntimeError(f"Cannot determine activation checkpoint phase: {phase}") if policy.gpu == "retain" and phase in ("first_forward", "recompute"): raise RuntimeError( "activation_policy.gpu=retain conflicts with active gradient checkpointing" ) if requires_backward and policy.gpu == "recompute" and phase == "none": raise RuntimeError( "activation_policy.gpu=recompute requires non-reentrant gradient " "checkpointing with the KT context_fn" ) def _strip_kt_zero_storage_from_state_dict(module, state_dict, prefix, local_metadata) -> None: """Never serialize expert placeholders that do not contain real weights.""" del local_metadata for name, param in module.named_parameters(): if getattr(param, "_kt_zero_storage", False): state_dict.pop(f"{prefix}{name}", None) def _supply_kt_zero_storage_for_state_dict_load( module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ) -> None: """Keep placeholder keys out of checkpoints without reporting them missing.""" del local_metadata, strict, missing_keys, unexpected_keys, error_msgs for name, param in module.named_parameters(): if getattr(param, "_kt_zero_storage", False): state_dict[f"{prefix}{name}"] = param class KTMoELayerWrapper(nn.Module): """Wrapper for MoE layer using KTMoEWrapper.""" def __init__( self, original_moe: nn.Module, wrapper: Any, lora_params: dict[str, nn.Parameter] | None, # Kept for backward compatibility, but ignored moe_config: MOEArchConfig, hidden_size: int, layer_idx: int, lora_experts: "LoRAExperts | None" = None, full_weight_grad: bool | None = None, uses_authoritative_optimizer_grads: bool | None = None, activation_policy: KTActivationPolicy | None = None, ): super().__init__() self._is_kt_moe_wrapper = True self.wrapper = wrapper self.moe_config = moe_config self.hidden_size = hidden_size self.layer_idx = layer_idx self.router_type = moe_config.router_type # IMPORTANT: Register submodules in the SAME ORDER as original MoE module # so that PEFT's named_modules() traversal order matches baseline. # This ensures kaiming_uniform_ calls happen in the same sequence. # Qwen3MoeSparseMoeBlock order: gate FIRST, then experts. # 1. gate/router FIRST - keep original attribute name for PEFT compatibility router_attr = moe_config.router_attr # "gate" for Qwen3/DeepSeek original_router = getattr(original_moe, router_attr, None) self._original_router = None # Set when router is not nn.Linear (e.g. TopKRouter) if original_router is not None and isinstance(original_router, nn.Linear): # transformers <=4.x / some models: gate is nn.Linear - register directly. setattr(self, router_attr, original_router) elif original_router is not None and hasattr(original_router, "weight") and isinstance( getattr(original_router, "weight"), nn.Parameter ): # transformers v5+: gate is a TopKRouter with nn.Parameter weight. # Wrap it in nn.Linear so PEFT can discover and inject LoRA. # The nn.Linear shares the same weight tensor - LoRA applied to it # is equivalent to LoRA on the original gate. router_weight = original_router.weight router_linear_cls = _DeepseekRouterLinear if self.router_type == "deepseek_gate" else nn.Linear router_linear = router_linear_cls( router_weight.shape[1], router_weight.shape[0], bias=False, ) router_linear.weight = router_weight # share the same parameter setattr(self, router_attr, router_linear) # Keep the original router for forward (top-k selection logic) self._original_router = original_router else: setattr(self, router_attr, original_router) self._router_attr = router_attr self._deepseek_routing_spec: DeepSeekRoutingSpec | None = None if self.router_type == "deepseek_gate" and self._original_router is not None: self._deepseek_routing_spec = DeepSeekRoutingSpec.from_modules( moe=original_moe, router=self._original_router, expert_num=moe_config.expert_num, ) correction_bias = getattr( self._original_router, "e_score_correction_bias", None, ) if not isinstance(correction_bias, torch.Tensor) or tuple( correction_bias.shape ) != (int(moe_config.expert_num),): raise ValueError( f"Layer {layer_idx}: DeepSeek correction bias must have shape " f"({moe_config.expert_num},), got " f"{getattr(correction_bias, 'shape', None)}" ) # 2. experts SECOND (this is what PEFT targets for LoRA) experts_attr = moe_config.experts_attr # typically "experts" setattr(self, experts_attr, getattr(original_moe, experts_attr, None)) self._experts_attr = experts_attr # 3. Shared expert (if any). DeepSeek/GLM use ``shared_experts`` while # Qwen2-MoE/Qwen3.5 use ``shared_expert`` plus a sigmoid gate. Preserve # the original attribute names so state_dict/checkpoint keys stay # compatible with the Hugging Face model. self._shared_expert_attr: str | None = None self._shared_expert_gate_attr: str | None = None if moe_config.has_shared_experts: for shared_expert_attr in ("shared_experts", "shared_expert"): shared_expert = getattr(original_moe, shared_expert_attr, None) if shared_expert is not None: setattr(self, shared_expert_attr, shared_expert) self._shared_expert_attr = shared_expert_attr break if self._shared_expert_attr is None: raise ValueError( f"Layer {layer_idx}: architecture declares shared experts, " "but the MoE module has neither 'shared_experts' nor 'shared_expert'" ) gate_candidates = ( ("shared_expert_gate", "shared_experts_gate") if self._shared_expert_attr == "shared_expert" else ("shared_experts_gate", "shared_expert_gate") ) for shared_expert_gate_attr in gate_candidates: shared_expert_gate = getattr(original_moe, shared_expert_gate_attr, None) if shared_expert_gate is not None: setattr(self, shared_expert_gate_attr, shared_expert_gate) self._shared_expert_gate_attr = shared_expert_gate_attr break if self._shared_expert_attr == "shared_expert" and self._shared_expert_gate_attr is None: raise ValueError( f"Layer {layer_idx}: singular 'shared_expert' requires " "'shared_expert_gate' for Qwen-style gated output" ) # 4. lora_experts (separate LoRA expert MLPs, different from PEFT LoRA on experts) self.lora_experts = lora_experts # PEFT LoRA tracking (set by kt_adapt_peft_lora) # _peft_lora_modules: {expert_idx: {proj_name: (lora_A, lora_B)}} self._peft_lora_modules: dict[int, dict[str, tuple[nn.Module, nn.Module]]] | None = None self._lora_pointers_dirty = False self._kt_managed_lora_enabled = False # Training-mode flags must be identical on every distributed rank even # though only rank 0 owns the backend object. if full_weight_grad is None: full_weight_grad = getattr(wrapper, "_full_weight_grad", False) if wrapper is not None else False if uses_authoritative_optimizer_grads is None: uses_authoritative_optimizer_grads = bool( wrapper is not None and getattr(wrapper, "_uses_authoritative_optimizer_grads", False) ) self._full_weight_grad = bool(full_weight_grad) self._uses_authoritative_optimizer_grads = bool(uses_authoritative_optimizer_grads) if activation_policy is None: # Direct callers without the new policy retain the old behavior: # normal training keeps both activation tiers, while the legacy # reuse flag keeps CPU activations across GPU recomputation. if wrapper is not None and getattr(wrapper, "reuse_checkpoint_forward", False): activation_policy = KTActivationPolicy(cpu="retain", gpu="recompute") else: activation_policy = KTActivationPolicy(cpu="retain", gpu="retain") self._kt_activation_policy = activation_policy self.register_state_dict_post_hook(_strip_kt_zero_storage_from_state_dict) self.register_load_state_dict_pre_hook(_supply_kt_zero_storage_for_state_dict_load) def _compute_shared_expert(self, hidden_states: torch.Tensor) -> torch.Tensor | None: if self._shared_expert_attr is None: return None shared_expert = getattr(self, self._shared_expert_attr) shared_output = shared_expert(hidden_states) if self._shared_expert_gate_attr is not None: shared_expert_gate = getattr(self, self._shared_expert_gate_attr) shared_output = torch.sigmoid(shared_expert_gate(hidden_states)) * shared_output return shared_output def _refresh_backend_weights(self) -> None: if self._lora_pointers_dirty: self.update_lora_pointers() self._lora_pointers_dirty = False if self._full_weight_grad and getattr(self.wrapper, "_kt_full_checkpoint_load_failed", False): raise RuntimeError( f"Layer {self.layer_idx}: a previous KT Full checkpoint load failed; " "reload a valid checkpoint before running forward" ) if self._full_weight_grad and getattr(self.wrapper, "_base_weights_dirty", False): with torch.profiler.record_function("kt.sft.base_weight_reload"): self.wrapper.update_base_weights() self.wrapper._base_weights_dirty = False def _apply(self, fn, recurse=True): # Protect experts from device transfer (PEFT LoRA should stay on CPU for KT) saved_experts = None experts_attr = getattr(self, "_experts_attr", None) if experts_attr is not None and getattr(self, experts_attr, None) is not None: saved_experts = getattr(self, experts_attr) self._modules.pop(experts_attr, None) result = super()._apply(fn, recurse) if saved_experts is not None: self._modules[experts_attr] = saved_experts return result def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: import torch.distributed as dist dist_on = dist.is_initialized() and dist.get_world_size() > 1 with torch.profiler.record_function("kt.sft.routing"): topk_ids, topk_weights = self._compute_routing(hidden_states) train_lora = bool( self._kt_managed_lora_enabled or (self._peft_lora_modules is not None and len(self._peft_lora_modules) > 0) or getattr(self, "_fused_expert_lora_params", None) ) full_weight_grad = self._full_weight_grad save_for_backward = ( self.training and torch.is_grad_enabled() and (hidden_states.requires_grad or topk_weights.requires_grad or train_lora or full_weight_grad) ) use_autograd_path = save_for_backward checkpoint_mode = _checkpoint_hook_mode() checkpoint_action = _activation_checkpoint_action(self._kt_activation_policy, checkpoint_mode) if not dist_on: _validate_activation_checkpoint_phase( self._kt_activation_policy, checkpoint_mode, requires_backward=use_autograd_path, ) cache_checkpoint_forward = checkpoint_action == "cache_first_forward" reuse_cached_forward = checkpoint_action == "reuse_recompute" save_for_backward_submit = use_autograd_path if checkpoint_mode == "first_forward" and not cache_checkpoint_forward: save_for_backward_submit = False with torch.profiler.record_function("kt.sft.submit_and_gpu_experts"): gpu_output, all_qlens = self._submit_and_compute_gpu( hidden_states, topk_ids, topk_weights, save_for_backward_submit, checkpoint_mode, checkpoint_action, ) # Use KTMoEFunction whenever backward is needed so KT backward and LoRA # gradient paths remain connected. if use_autograd_path: # A requires-grad sentinel keeps the custom autograd node alive on # non-rank-0 fused/full ranks that intentionally own no KT params. lora_ref = hidden_states.new_empty((), requires_grad=(train_lora or full_weight_grad)) if train_lora and self._peft_lora_modules: found_lora_ref = False for expert_loras in self._peft_lora_modules.values(): for lora_A, lora_B in expert_loras.values(): if hasattr(lora_A, "weight") and lora_A.weight.requires_grad: lora_ref = lora_A.weight found_lora_ref = True break if found_lora_ref: break elif train_lora and getattr(self, "_fused_expert_lora_params", None): lora_ref = self._fused_expert_lora_params[0] elif full_weight_grad and self.wrapper is not None: # In full mode, use base weight param as autograd sentinel if self.wrapper.gate_proj_buf is not None: lora_ref = self.wrapper.gate_proj_buf with torch.profiler.record_function("kt.sft.autograd_apply_and_cpu_sync"): moe_output = KTMoEFunction.apply( hidden_states, topk_ids, topk_weights, self.wrapper, lora_ref, self.hidden_size, self.moe_config.num_experts_per_tok, self.layer_idx, save_for_backward, train_lora, all_qlens, cache_checkpoint_forward, reuse_cached_forward, # Base weight params for full mode gradient flow self.wrapper.gate_proj_buf if full_weight_grad and self.wrapper is not None else None, self.wrapper.up_proj_buf if full_weight_grad and self.wrapper is not None else None, self.wrapper.down_proj_buf if full_weight_grad and self.wrapper is not None else None, ) else: moe_output = self._sync_forward_output_no_autograd( hidden_states=hidden_states, all_qlens=all_qlens, ) if gpu_output is not None: moe_output = moe_output + gpu_output return moe_output def _sync_forward_output_no_autograd( self, hidden_states: torch.Tensor, all_qlens: list[int] | tuple[int, ...] | None, ) -> torch.Tensor: """Sync CPU expert output without creating KTMoEFunction autograd nodes.""" import torch.distributed as dist original_device = hidden_states.device original_dtype = hidden_states.dtype batch_size, seq_len, _ = hidden_states.shape qlen = batch_size * seq_len dist_on = dist.is_initialized() and dist.get_world_size() > 1 rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist_on else 1 if dist_on: if all_qlens is None: all_qlens_list = _all_gather_qlens(qlen, original_device, world_size) else: all_qlens_list = [int(q) for q in all_qlens] if len(all_qlens_list) != world_size: raise RuntimeError(f"all_qlens length mismatch: got {len(all_qlens_list)}, expected {world_size}") if int(all_qlens_list[rank]) != qlen: raise RuntimeError(f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens_list[rank]}") total_qlen = sum(all_qlens_list) sync_error = None if rank == 0: try: if self.wrapper is None: raise RuntimeError("Rank0 wrapper is required in distributed KT overlap path.") cpu_output = self.wrapper.sync_forward(output_device=original_device) cpu_output = cpu_output.to(dtype=original_dtype).view(total_qlen, self.hidden_size) offsets = _qlen_offsets(all_qlens_list) scatter_list = [ cpu_output[offsets[i] : offsets[i + 1]].contiguous() for i in range(world_size) ] except Exception as exc: sync_error = exc else: scatter_list = None _sync_rank0_exception( sync_error, device=original_device, context=f"Layer {self.layer_idx} CPU inference forward failed", ) output_flat = _dist_scatter_varlen_from_rank0( rank0_chunks=scatter_list, all_qlens=all_qlens_list, rank=rank, world_size=world_size, feature_shape=(self.hidden_size,), device=original_device, dtype=original_dtype, ) output = output_flat.view(batch_size, seq_len, self.hidden_size) del output_flat return output if self.wrapper is not None: cpu_output = self.wrapper.sync_forward(output_device=original_device) output = cpu_output.view(batch_size, seq_len, self.hidden_size).to(dtype=original_dtype) return output return torch.empty(batch_size, seq_len, self.hidden_size, device=original_device, dtype=original_dtype) def _compute_routing(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: router = getattr(self, self._router_attr) router_grad_enabled = self.training and torch.is_grad_enabled() and any( parameter.requires_grad for parameter in router.parameters() ) routing_context = nullcontext() if router_grad_enabled else torch.no_grad() def finish( topk_ids: torch.Tensor, topk_weights: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if topk_weights.is_floating_point(): topk_weights = topk_weights.to(torch.bfloat16) if router_grad_enabled and not topk_weights.requires_grad: raise RuntimeError( f"Layer {self.layer_idx}: trainable router produced detached routing weights" ) return topk_ids, topk_weights with routing_context: if self.router_type == "deepseek_gate" and self._original_router is None: # Native fallback; weighted gates use the proxy below so router LoRA stays active. was_training = router.training if was_training: router.eval() router_output = router(hidden_states) if was_training: router.train() if not isinstance(router_output, (tuple, list)) or len(router_output) != 2: raise RuntimeError( f"Layer {self.layer_idx}: DeepSeek gate must return " "(topk_ids, topk_weights)" ) topk_ids, topk_weights = router_output return finish(topk_ids, topk_weights) # When _original_router is set, self.gate is an nn.Linear wrapper # around the TopKRouter's weight. Use it (with PEFT LoRA if # applied) for the linear projection, then replicate top-k logic. if self._original_router is not None: orig_router = self._original_router router_logits = router(hidden_states.view(-1, self.hidden_size)) if self.router_type == "deepseek_gate": spec = self._deepseek_routing_spec if spec is None: raise RuntimeError( f"Layer {self.layer_idx}: DeepSeek routing spec was not initialized" ) router_probs = torch.sigmoid(router_logits.float()) correction_bias = getattr(orig_router, "e_score_correction_bias", None) if correction_bias is None: raise RuntimeError( f"Layer {self.layer_idx}: DeepSeek noaux_tc gate is missing " "e_score_correction_bias" ) scores_for_choice = router_probs + correction_bias.to( device=router_probs.device, dtype=router_probs.dtype, ) group_scores = ( scores_for_choice.view( -1, spec.n_group, spec.expert_num // spec.n_group, ) .topk(2, dim=-1)[0] .sum(dim=-1) ) group_idx = torch.topk( group_scores, k=spec.topk_group, dim=-1, sorted=False, )[1] group_mask = torch.zeros_like(group_scores, dtype=torch.bool) group_mask.scatter_(1, group_idx, True) score_mask = ( group_mask.unsqueeze(-1) .expand( -1, spec.n_group, spec.expert_num // spec.n_group, ) .reshape(-1, spec.expert_num) ) scores_for_choice = scores_for_choice.masked_fill( ~score_mask, 0.0, ) topk_ids = torch.topk( scores_for_choice, k=spec.top_k, dim=-1, sorted=False, )[1] topk_weights = router_probs.gather(1, topk_ids) if spec.norm_topk_prob: topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) topk_weights = topk_weights * spec.routed_scaling_factor return finish(topk_ids, topk_weights) if self.router_type == "glm4_moe_gate": router_probs = torch.sigmoid(router_logits.float()) correction_bias = getattr(orig_router, "e_score_correction_bias", None) if correction_bias is None: router_logits_for_choice = router_probs else: router_logits_for_choice = router_probs + correction_bias.to( device=router_probs.device, dtype=router_probs.dtype, ) n_group = getattr(orig_router, "n_group", 1) topk_group = getattr(orig_router, "topk_group", n_group) expert_num = self.moe_config.expert_num group_scores = ( router_logits_for_choice.view(-1, n_group, expert_num // n_group) .topk(2, dim=-1)[0] .sum(dim=-1) ) group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] group_mask = torch.zeros_like(group_scores) group_mask.scatter_(1, group_idx, 1) score_mask = ( group_mask.unsqueeze(-1) .expand(-1, n_group, expert_num // n_group) .reshape(-1, expert_num) ) scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0) top_k = getattr(orig_router, "top_k", self.moe_config.num_experts_per_tok) topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1] topk_weights = router_probs.gather(1, topk_ids) if getattr(orig_router, "norm_topk_prob", True): topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) topk_weights = topk_weights * getattr(orig_router, "routed_scaling_factor", 1.0) return finish(topk_ids, topk_weights) router_probs = F.softmax(router_logits, dtype=torch.float, dim=-1) top_k = getattr(orig_router, "top_k", self.moe_config.num_experts_per_tok) norm_topk_prob = getattr(orig_router, "norm_topk_prob", True) topk_weights, topk_ids = torch.topk(router_probs, top_k, dim=-1) if norm_topk_prob: topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) topk_weights = topk_weights.to(router_logits.dtype) return finish(topk_ids, topk_weights) router_output = router(hidden_states.view(-1, self.hidden_size)) # transformers v5 TopKRouter returns (router_logits, router_scores, router_indices) # directly — scores/indices are already topk-normalized. if isinstance(router_output, tuple): if len(router_output) >= 3: _logits, topk_weights, topk_ids = router_output[0], router_output[1], router_output[2] return finish(topk_ids, topk_weights) router_output = router_output[0] router_logits = router_output routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(routing_weights, self.moe_config.num_experts_per_tok, dim=-1) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) return finish(topk_ids, topk_weights) def _submit_and_compute_gpu( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, save_for_backward: bool, checkpoint_phase: str = "none", checkpoint_action: str = "normal", ) -> tuple[torch.Tensor | None, list[int] | None]: import torch.distributed as dist batch_size, seq_len, _ = hidden_states.shape original_device = hidden_states.device original_dtype = hidden_states.dtype dist_on = dist.is_initialized() and dist.get_world_size() > 1 rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist_on else 1 qlen = batch_size * seq_len reuse_cached_forward = checkpoint_action == "reuse_recompute" if dist_on: ownership_error = None try: wrapped_world_size = int(getattr(self, "_kt_world_size_at_wrap", world_size)) if wrapped_world_size != world_size: raise RuntimeError( f"KT wrapper was created for world_size={wrapped_world_size}, " f"but the active process group has world_size={world_size}" ) if rank == 0 and self.wrapper is None: raise RuntimeError("rank 0 does not own the KT backend") if rank != 0 and self.wrapper is not None: raise RuntimeError(f"rank {rank} unexpectedly owns a KT backend") except Exception as exc: ownership_error = exc all_qlens = _all_gather_checkpoint_state( qlen, layer_idx=self.layer_idx, phase=checkpoint_phase, action=checkpoint_action, cpu_policy=self._kt_activation_policy.cpu, gpu_policy=self._kt_activation_policy.gpu, owner_valid=ownership_error is None, device=original_device, world_size=world_size, ) if int(all_qlens[rank]) != qlen: raise RuntimeError(f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens[rank]}") _validate_activation_checkpoint_phase( self._kt_activation_policy, checkpoint_phase, requires_backward=save_for_backward, ) if ownership_error is not None: raise RuntimeError( f"Layer {self.layer_idx} KT ownership validation failed: {ownership_error}" ) from ownership_error total_qlen = sum(all_qlens) refresh_error = None if rank == 0: try: self._refresh_backend_weights() if checkpoint_action == "cache_first_forward": validate_empty = getattr(self.wrapper, "validate_checkpoint_cache_empty", None) if validate_empty is not None: validate_empty() elif getattr(self.wrapper, "_kt_has_cached_forward", False): raise RuntimeError( "Checkpoint forward cache is still live before a new first forward" ) except Exception as exc: refresh_error = exc if reuse_cached_forward: cache_error = refresh_error if rank == 0 and cache_error is None: try: if self.wrapper is None: raise RuntimeError("Rank 0 does not own the KT backend") validate_cache = getattr(self.wrapper, "validate_checkpoint_output", None) if validate_cache is not None: validate_cache(total_qlen) elif not getattr(self.wrapper, "_kt_has_cached_forward", False): raise RuntimeError("No cached checkpoint forward output is available") except Exception as exc: cache_error = exc if rank == 0 and cache_error is not None: _poison_checkpoint_cache(self.wrapper, cache_error) _sync_rank0_exception( cache_error, device=original_device, context=f"Layer {self.layer_idx} checkpoint cache validation failed", ) else: hs_flat = hidden_states.view(qlen, self.hidden_size).contiguous() expert_ids = topk_ids.view(qlen, self.moe_config.num_experts_per_tok).contiguous() weights = topk_weights.view(qlen, self.moe_config.num_experts_per_tok).contiguous() gathered_hs = _dist_gather_varlen_to_rank0( hs_flat.detach(), all_qlens=all_qlens, rank=rank, world_size=world_size, ) gathered_ids = _dist_gather_varlen_to_rank0( expert_ids.detach(), all_qlens=all_qlens, rank=rank, world_size=world_size, ) gathered_wts = _dist_gather_varlen_to_rank0( weights.detach(), all_qlens=all_qlens, rank=rank, world_size=world_size, ) submit_error = refresh_error if rank == 0 and submit_error is None: try: if self.wrapper is None: raise RuntimeError("Rank 0 does not own the KT backend") self.wrapper.submit_forward( torch.cat(gathered_hs, dim=0), torch.cat(gathered_ids, dim=0), torch.cat(gathered_wts, dim=0), save_for_backward=save_for_backward, ) except Exception as exc: submit_error = exc if rank == 0 and submit_error is not None and checkpoint_action == "cache_first_forward": _poison_checkpoint_cache(self.wrapper, submit_error) _sync_rank0_exception( submit_error, device=original_device, context=f"Layer {self.layer_idx} CPU forward submission failed", ) gpu_output = None if self._shared_expert_attr is not None or self.lora_experts is not None: gpu_error = None try: gpu_output = self._compute_shared_expert(hidden_states) if gpu_output is not None: gpu_output = gpu_output.to(dtype=original_dtype) if self.lora_experts is not None: lora_out = self.lora_experts(hidden_states) gpu_output = lora_out if gpu_output is None else gpu_output + lora_out except Exception as exc: gpu_error = exc try: _sync_any_rank_exception( gpu_error, device=original_device, context=f"Layer {self.layer_idx} local GPU expert forward failed", ) except Exception as exc: if checkpoint_action in ("cache_first_forward", "reuse_recompute"): _poison_checkpoint_cache(self.wrapper, exc) raise return gpu_output, all_qlens else: # ---- Single-GPU path: submit + GPU compute ---- self._refresh_backend_weights() if checkpoint_action == "cache_first_forward": try: validate_empty = getattr(self.wrapper, "validate_checkpoint_cache_empty", None) if validate_empty is not None: validate_empty() elif getattr(self.wrapper, "_kt_has_cached_forward", False): raise RuntimeError( "Checkpoint forward cache is still live before a new first forward" ) except Exception as exc: _poison_checkpoint_cache(self.wrapper, exc) raise input_flat = hidden_states.view(qlen, self.hidden_size) expert_ids = topk_ids.view(qlen, self.moe_config.num_experts_per_tok) weights = topk_weights.view(qlen, self.moe_config.num_experts_per_tok) # Avoid passing graph-attached tensors into C++ cache. submit_hs = input_flat.detach() submit_ids = expert_ids.detach() submit_wts = weights.detach() if reuse_cached_forward: validate_cache = getattr(self.wrapper, "validate_checkpoint_output", None) if validate_cache is not None: validate_cache(qlen) elif not getattr(self.wrapper, "_kt_has_cached_forward", False): raise RuntimeError("No cached checkpoint forward output is available") else: try: self.wrapper.submit_forward( submit_hs, submit_ids, submit_wts, save_for_backward=save_for_backward, ) except Exception as exc: if checkpoint_action == "cache_first_forward": _poison_checkpoint_cache(self.wrapper, exc) raise # GPU compute: shared_experts + lora_experts try: gpu_output = self._compute_shared_expert(hidden_states) if self.lora_experts is not None: lora_out = self.lora_experts(hidden_states) gpu_output = lora_out if gpu_output is None else gpu_output + lora_out except Exception as exc: if checkpoint_action in ("cache_first_forward", "reuse_recompute"): _poison_checkpoint_cache(self.wrapper, exc) raise return gpu_output, None def update_lora_pointers(self): """Sync PEFT LoRA weights to C++ kernel after optimizer update.""" # Skip if wrapper is None (non-rank-0 processes) if self.wrapper is None: return # Skip if wrapper is not properly initialized if not getattr(self.wrapper, "_weights_loaded", False): logger.warning(f"Layer {self.layer_idx}: Skipping update_lora_pointers - weights not loaded") return if not getattr(self.wrapper, "_lora_initialized", False): logger.warning(f"Layer {self.layer_idx}: Skipping update_lora_pointers - LoRA not initialized") return # PEFT weights are views into wrapper's contiguous buffers — # optimizer.step() already updated them in-place, just re-sync to C++. self.wrapper.update_lora_weights()