From a63a112fd6ce6271ce18953ae8363e6ae00f0d8b Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 05:37:15 +0800 Subject: [PATCH 01/41] feat(lora): add Mixture-of-LoRA routing primitives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### Summary 为 Mixture-of-LoRA 提供训练端和 rollout 端可复用的路由基础逻辑。本次提交只增加公共路由能力和 CPU 测试，不修改现有单 LoRA 路径。 #### Changes 新增 RoutingDecision 和 RoutingStatistics，统一保存 Top-K 前概率、expert 编号、归一化激活权重，以及可跨 rank 聚合的逐 expert 原始统计。实现 FP32 softmax、token-level Top-K 路由、response mask 过滤、归一化熵和 balance loss；expert 选择份额按 valid_tokens * K 归一化，离散选择统计不参与反向传播，router 梯度通过 Top-K 前概率计算。新增 17 个 CPU 测试，覆盖多个 K、temperature、mask、空 response、梯度、entropy 和非法参数。 #### Verification pytest：17 passed。Ruff、docformatter、通用文件检查和冲突标记检查均通过。gitleaks hook 因 H20 访问 proxy.golang.org 时 TLS 超时，未能完成环境安装；失败发生在依赖下载阶段，并非扫描发现问题。 --- relax/utils/mixture_lora.py | 200 +++++++++++++++++++++++ tests/utils/test_mixture_lora_routing.py | 141 ++++++++++++++++ 2 files changed, 341 insertions(+) create mode 100644 relax/utils/mixture_lora.py create mode 100644 tests/utils/test_mixture_lora_routing.py diff --git a/relax/utils/mixture_lora.py b/relax/utils/mixture_lora.py new file mode 100644 index 000000000..489cc5cd9 --- /dev/null +++ b/relax/utils/mixture_lora.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Backend-independent routing primitives for Mixture-of-LoRA.""" + +import math +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class RoutingDecision: + """Token-level Top-K routing results. + + All tensors use ``[token, expert]`` or ``[token, selected_expert]`` + layouts. Probabilities and selected weights are always FP32. + """ + + pre_topk_probs: torch.Tensor + topk_indices: torch.Tensor + post_topk_weights: torch.Tensor + + @property + def num_experts(self) -> int: + return self.pre_topk_probs.shape[-1] + + @property + def top_k(self) -> int: + return self.topk_indices.shape[-1] + + def dense_weights(self) -> torch.Tensor: + """Return Top-K weights scattered into the full expert dimension.""" + + weights = torch.zeros_like(self.pre_topk_probs) + return weights.scatter(-1, self.topk_indices, self.post_topk_weights) + + +@dataclass(frozen=True) +class RoutingStatistics: + """Unnormalized per-expert values that can be reduced across ranks.""" + + pre_topk_prob_sum: torch.Tensor + post_topk_weight_sum: torch.Tensor + selection_count: torch.Tensor + top1_count: torch.Tensor + pre_topk_entropy_sum: torch.Tensor + post_topk_entropy_sum: torch.Tensor + valid_token_count: torch.Tensor + top_k: int + + def _per_token(self, value: torch.Tensor) -> torch.Tensor: + denominator = self.valid_token_count.clamp_min(1) + mean = value / denominator + return mean * (self.valid_token_count > 0).to(mean.dtype) + + @property + def pre_topk_mean_prob(self) -> torch.Tensor: + return self._per_token(self.pre_topk_prob_sum) + + @property + def post_topk_mean_weight(self) -> torch.Tensor: + return self._per_token(self.post_topk_weight_sum) + + @property + def selection_share(self) -> torch.Tensor: + return self._per_token(self.selection_count) / self.top_k + + @property + def top1_fraction(self) -> torch.Tensor: + return self._per_token(self.top1_count) + + @property + def pre_topk_normalized_entropy(self) -> torch.Tensor: + return self._per_token(self.pre_topk_entropy_sum) + + @property + def post_topk_normalized_entropy(self) -> torch.Tensor: + return self._per_token(self.post_topk_entropy_sum) + + @property + def balance_loss(self) -> torch.Tensor: + """Return ``N * sum(F_e * P_e)`` with Top-K-normalized ``F_e``.""" + + # Top-K membership is discrete. Keep F_e detached while P_e remains + # differentiable so the auxiliary loss updates only the router scores. + selection_share = self.selection_share.detach() + loss = self.pre_topk_prob_sum.shape[0] * torch.sum(selection_share * self.pre_topk_mean_prob) + return loss * (self.valid_token_count > 0).to(loss.dtype) + + +def route_topk(router_logits: torch.Tensor, top_k: int, temperature: float) -> RoutingDecision: + """Compute FP32 softmax probabilities and normalized Top-K weights.""" + + if router_logits.ndim != 2: + raise ValueError(f"router_logits must have shape [token, expert], got {tuple(router_logits.shape)}") + if not torch.is_floating_point(router_logits): + raise TypeError(f"router_logits must be floating point, got {router_logits.dtype}") + + num_experts = router_logits.shape[-1] + if not isinstance(top_k, int) or isinstance(top_k, bool) or not 1 <= top_k <= num_experts: + raise ValueError(f"top_k must satisfy 1 <= top_k <= {num_experts}, got {top_k}") + if not isinstance(temperature, (int, float)) or isinstance(temperature, bool): + raise TypeError(f"temperature must be a real number, got {type(temperature).__name__}") + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError(f"temperature must be finite and greater than 0, got {temperature}") + + # Keep router normalization in FP32 even when model parameters and + # activations use a lower-precision dtype. + pre_topk_probs = torch.softmax(router_logits.float() / temperature, dim=-1) + selected_probs, topk_indices = torch.topk(pre_topk_probs, k=top_k, dim=-1) + post_topk_weights = selected_probs / selected_probs.sum(dim=-1, keepdim=True) + return RoutingDecision( + pre_topk_probs=pre_topk_probs, + topk_indices=topk_indices, + post_topk_weights=post_topk_weights, + ) + + +def _normalized_entropy(probs: torch.Tensor, num_choices: int) -> torch.Tensor: + if num_choices <= 1: + return torch.zeros(probs.shape[0], dtype=probs.dtype, device=probs.device) + + log_probs = probs.clamp_min(torch.finfo(probs.dtype).tiny).log() + return -(probs * log_probs).sum(dim=-1) / math.log(num_choices) + + +def compute_routing_statistics( + decision: RoutingDecision, + response_mask: torch.Tensor, +) -> RoutingStatistics: + """Collect response-token routing values without converting tensors to + Python.""" + + _validate_routing_decision(decision) + if response_mask.ndim != 1 or response_mask.shape[0] != decision.pre_topk_probs.shape[0]: + raise ValueError( + f"response_mask must have shape [token] matching the routing decision, got {tuple(response_mask.shape)}" + ) + if response_mask.device != decision.pre_topk_probs.device: + raise ValueError("response_mask and routing tensors must be on the same device") + + probs = decision.pre_topk_probs + indices = decision.topk_indices + weights = decision.post_topk_weights + num_experts = decision.num_experts + top_k = decision.top_k + + valid_mask = response_mask.to(dtype=torch.bool) + token_weights = valid_mask.to(dtype=probs.dtype) + selected_token_weights = token_weights.unsqueeze(-1).expand(-1, top_k) + + pre_topk_prob_sum = torch.sum(probs * token_weights.unsqueeze(-1), dim=0) + + # Accumulate selected experts directly instead of materializing another + # token-by-expert tensor for every metric. + post_topk_weight_sum = torch.zeros(num_experts, dtype=probs.dtype, device=probs.device) + post_topk_weight_sum.scatter_add_( + 0, + indices.reshape(-1), + (weights * selected_token_weights).reshape(-1), + ) + + selection_count = torch.zeros_like(post_topk_weight_sum) + selection_count.scatter_add_(0, indices.reshape(-1), selected_token_weights.reshape(-1)) + + top1_count = torch.zeros_like(post_topk_weight_sum) + top1_count.scatter_add_(0, indices[:, 0], token_weights) + + pre_topk_entropy_sum = torch.sum(_normalized_entropy(probs, num_experts) * token_weights) + post_topk_entropy_sum = torch.sum(_normalized_entropy(weights, top_k) * token_weights) + + return RoutingStatistics( + pre_topk_prob_sum=pre_topk_prob_sum, + post_topk_weight_sum=post_topk_weight_sum, + selection_count=selection_count, + top1_count=top1_count, + pre_topk_entropy_sum=pre_topk_entropy_sum, + post_topk_entropy_sum=post_topk_entropy_sum, + valid_token_count=token_weights.sum(), + top_k=top_k, + ) + + +def _validate_routing_decision(decision: RoutingDecision) -> None: + probs = decision.pre_topk_probs + indices = decision.topk_indices + weights = decision.post_topk_weights + + if probs.ndim != 2 or indices.ndim != 2 or weights.ndim != 2: + raise ValueError("routing decision tensors must all be two-dimensional") + if indices.shape != weights.shape or indices.shape[0] != probs.shape[0]: + raise ValueError("routing decision tensor shapes do not agree") + if indices.shape[1] < 1 or indices.shape[1] > probs.shape[1]: + raise ValueError("routing decision has an invalid Top-K dimension") + if probs.device != indices.device or probs.device != weights.device: + raise ValueError("routing decision tensors must be on the same device") + if probs.dtype != torch.float32 or weights.dtype != torch.float32: + raise TypeError("routing probabilities and weights must be FP32") + if indices.dtype != torch.long: + raise TypeError("topk_indices must use torch.long") diff --git a/tests/utils/test_mixture_lora_routing.py b/tests/utils/test_mixture_lora_routing.py new file mode 100644 index 000000000..52343629d --- /dev/null +++ b/tests/utils/test_mixture_lora_routing.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import math + +import pytest +import torch + +from relax.utils.mixture_lora import RoutingDecision, compute_routing_statistics, route_topk + + +def _normalized_entropy(values): + if len(values) <= 1: + return 0.0 + return -sum(value * math.log(value) for value in values if value > 0) / math.log(len(values)) + + +def test_route_topk_returns_fp32_normalized_weights(): + logits = torch.tensor([[4.0, 1.0, 3.0, 2.0], [0.0, 5.0, 2.0, 1.0]], dtype=torch.bfloat16) + + decision = route_topk(logits, top_k=2, temperature=1.0) + + assert decision.pre_topk_probs.shape == (2, 4) + assert decision.topk_indices.shape == (2, 2) + assert decision.post_topk_weights.shape == (2, 2) + assert decision.pre_topk_probs.dtype == torch.float32 + assert decision.post_topk_weights.dtype == torch.float32 + assert decision.topk_indices.dtype == torch.long + assert torch.equal(decision.topk_indices, torch.tensor([[0, 2], [1, 2]])) + assert torch.allclose(decision.post_topk_weights.sum(dim=-1), torch.ones(2)) + + dense_weights = decision.dense_weights() + assert torch.allclose(dense_weights.sum(dim=-1), torch.ones(2)) + assert torch.equal(dense_weights == 0, torch.tensor([[False, True, False, True], [True, False, False, True]])) + + +def test_route_topk_temperature_changes_distribution(): + logits = torch.tensor([[3.0, 2.0, 1.0]]) + + cold = route_topk(logits, top_k=2, temperature=0.5) + warm = route_topk(logits, top_k=2, temperature=2.0) + + assert cold.pre_topk_probs.max() > warm.pre_topk_probs.max() + assert cold.post_topk_weights.max() > warm.post_topk_weights.max() + + +@pytest.mark.parametrize("top_k", [1, 2, 4]) +def test_balance_loss_uniform_baseline_does_not_depend_on_top_k(top_k): + decision = route_topk(torch.zeros(8, 4), top_k=top_k, temperature=1.0) + stats = compute_routing_statistics(decision, torch.ones(8, dtype=torch.bool)) + + assert stats.balance_loss == pytest.approx(1.0) + assert stats.selection_share.sum() == pytest.approx(1.0) + + +def test_routing_statistics_use_only_response_tokens(): + pre_topk_probs = torch.tensor( + [ + [0.6, 0.3, 0.1], + [0.2, 0.5, 0.3], + [0.1, 0.2, 0.7], + ] + ) + topk_indices = torch.tensor([[0, 1], [1, 2], [2, 1]]) + post_topk_weights = torch.tensor([[2 / 3, 1 / 3], [5 / 8, 3 / 8], [7 / 9, 2 / 9]]) + decision = RoutingDecision(pre_topk_probs, topk_indices, post_topk_weights) + + stats = compute_routing_statistics(decision, torch.tensor([1, 0, 1], dtype=torch.bool)) + + assert stats.valid_token_count == 2 + assert torch.allclose(stats.pre_topk_mean_prob, torch.tensor([0.35, 0.25, 0.40])) + assert torch.allclose(stats.post_topk_mean_weight, torch.tensor([1 / 3, 5 / 18, 7 / 18])) + assert torch.allclose(stats.selection_share, torch.tensor([0.25, 0.50, 0.25])) + assert torch.allclose(stats.top1_fraction, torch.tensor([0.50, 0.0, 0.50])) + assert stats.balance_loss == pytest.approx(0.9375) + + expected_pre_entropy = (_normalized_entropy([0.6, 0.3, 0.1]) + _normalized_entropy([0.1, 0.2, 0.7])) / 2 + expected_post_entropy = (_normalized_entropy([2 / 3, 1 / 3]) + _normalized_entropy([7 / 9, 2 / 9])) / 2 + assert stats.pre_topk_normalized_entropy == pytest.approx(expected_pre_entropy) + assert stats.post_topk_normalized_entropy == pytest.approx(expected_post_entropy) + + +def test_balance_loss_backpropagates_through_pre_topk_probabilities(): + logits = torch.tensor([[3.0, 2.0, 1.0], [2.0, 0.0, 1.0]], requires_grad=True) + decision = route_topk(logits, top_k=2, temperature=1.0) + + compute_routing_statistics(decision, torch.ones(2, dtype=torch.bool)).balance_loss.backward() + + assert logits.grad is not None + assert torch.isfinite(logits.grad).all() + assert torch.count_nonzero(logits.grad) > 0 + + +def test_k_one_has_zero_post_topk_entropy(): + decision = route_topk(torch.tensor([[2.0, 1.0], [0.0, 3.0]]), top_k=1, temperature=1.0) + stats = compute_routing_statistics(decision, torch.ones(2, dtype=torch.bool)) + + assert stats.post_topk_normalized_entropy == 0 + + +def test_no_response_tokens_returns_zero_statistics_and_loss(): + logits = torch.tensor([[2.0, 1.0], [0.0, 3.0]], requires_grad=True) + decision = route_topk(logits, top_k=1, temperature=1.0) + stats = compute_routing_statistics(decision, torch.zeros(2, dtype=torch.bool)) + + assert stats.valid_token_count == 0 + assert torch.count_nonzero(stats.pre_topk_mean_prob) == 0 + assert torch.count_nonzero(stats.post_topk_mean_weight) == 0 + assert stats.balance_loss == 0 + + stats.balance_loss.backward() + assert torch.count_nonzero(logits.grad) == 0 + + +@pytest.mark.parametrize( + ("top_k", "temperature", "error"), + [ + (0, 1.0, ValueError), + (4, 1.0, ValueError), + (1, 0.0, ValueError), + (1, float("inf"), ValueError), + (True, 1.0, ValueError), + (1, True, TypeError), + ], +) +def test_route_topk_rejects_invalid_configuration(top_k, temperature, error): + with pytest.raises(error): + route_topk(torch.ones(2, 3), top_k=top_k, temperature=temperature) + + +def test_route_topk_rejects_invalid_logits(): + with pytest.raises(ValueError, match="shape"): + route_topk(torch.ones(2, 3, 4), top_k=2, temperature=1.0) + with pytest.raises(TypeError, match="floating point"): + route_topk(torch.ones(2, 3, dtype=torch.long), top_k=2, temperature=1.0) + + +def test_statistics_reject_mask_with_wrong_shape(): + decision = route_topk(torch.ones(2, 3), top_k=2, temperature=1.0) + + with pytest.raises(ValueError, match="response_mask"): + compute_routing_statistics(decision, torch.ones(2, 1)) From 9deae67c3dc428bf6458b512a0dc05118616771c Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 06:23:13 +0800 Subject: [PATCH 02/41] feat(lora): add Mixture-of-LoRA execution schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### Summary 完成 Mixture-of-LoRA 阶段 A 的公共配置、参数 schema 和无参数执行器接口，为 Megatron 与 SGLang 后续实现提供同一组结构和数学定义。 #### Changes 新增 MixtureLoraConfig 及 N、R、K、temperature、aux loss、alpha 和目标层校验；固定 expert A/B 与 router 的参数名称和全局 shape；新增 checkpoint/transport 共用的状态与 TP 分片描述；新增显式并行上下文、RoutedLoRAExecutor 协议和纯 PyTorch dense executor；增加逐 site balance loss 平均函数，并保证执行器不持有参数、不修改路由结果。 #### Verification Mixture-of-LoRA L0 测试 38 passed，覆盖 FP32、FP16、BF16、输出、梯度、mask、schema、可重复性和异常输入。现有单 LoRA 回归 43 passed、5 skipped。全仓 pre-commit 全部通过，包含 Ruff、docformatter 和 gitleaks。 --- relax/utils/mixture_lora.py | 309 +++++++++++++++++++++++ tests/utils/test_mixture_lora_routing.py | 201 ++++++++++++++- 2 files changed, 509 insertions(+), 1 deletion(-) diff --git a/relax/utils/mixture_lora.py b/relax/utils/mixture_lora.py index 489cc5cd9..9c54e232d 100644 --- a/relax/utils/mixture_lora.py +++ b/relax/utils/mixture_lora.py @@ -4,8 +4,240 @@ import math from dataclasses import dataclass +from typing import Literal, Protocol, Sequence, runtime_checkable import torch +from torch import nn + + +MIXTURE_LORA_SCHEMA_VERSION = 1 +MixtureLoraParameterKind = Literal["experts.lora_A", "experts.lora_B", "router.weight"] +_PARAMETER_KINDS = {"experts.lora_A", "experts.lora_B", "router.weight"} + + +@dataclass(frozen=True) +class MixtureLoraConfig: + """Backend-independent Mixture-of-LoRA configuration.""" + + num_experts: int + rank: int + top_k: int + temperature: float + aux_loss_coef: float + alpha: float + target_modules: tuple[str, ...] + schema_version: int = MIXTURE_LORA_SCHEMA_VERSION + + def __post_init__(self) -> None: + _validate_positive_int("num_experts", self.num_experts) + if self.num_experts <= 1: + raise ValueError(f"num_experts must be greater than 1, got {self.num_experts}") + _validate_positive_int("rank", self.rank) + _validate_positive_int("top_k", self.top_k) + if self.top_k > self.num_experts: + raise ValueError(f"top_k must be no greater than num_experts, got {self.top_k} and {self.num_experts}") + _validate_finite_number("temperature", self.temperature, minimum=0.0, minimum_inclusive=False) + _validate_finite_number("aux_loss_coef", self.aux_loss_coef, minimum=0.0, minimum_inclusive=True) + _validate_finite_number("alpha", self.alpha, minimum=0.0, minimum_inclusive=False) + _validate_positive_int("schema_version", self.schema_version) + + if isinstance(self.target_modules, str): + raise TypeError("target_modules must be a sequence of module names, not a string") + target_modules = tuple(self.target_modules) + if not target_modules or any(not isinstance(target, str) or not target.strip() for target in target_modules): + raise ValueError("target_modules must contain non-empty module names") + if len(set(target_modules)) != len(target_modules): + raise ValueError("target_modules must not contain duplicates") + object.__setattr__(self, "target_modules", target_modules) + + @property + def scale(self) -> float: + return self.alpha / self.rank + + +@dataclass(frozen=True) +class MixtureLoraStateSpec: + """Stable identity and global layout for one Mixture-of-LoRA tensor.""" + + schema_version: int + site_id: str + parameter_kind: MixtureLoraParameterKind + global_shape: tuple[int, ...] + dtype: torch.dtype + + def __post_init__(self) -> None: + _validate_positive_int("schema_version", self.schema_version) + if not isinstance(self.site_id, str) or not self.site_id.strip(): + raise ValueError("site_id must be a non-empty string") + if self.parameter_kind not in _PARAMETER_KINDS: + raise ValueError(f"unsupported Mixture-of-LoRA parameter kind: {self.parameter_kind}") + + global_shape = tuple(self.global_shape) + expected_ndim = 2 if self.parameter_kind == "router.weight" else 3 + if len(global_shape) != expected_ndim or any( + not isinstance(size, int) or isinstance(size, bool) or size <= 0 for size in global_shape + ): + raise ValueError( + f"global_shape for {self.parameter_kind} must contain {expected_ndim} positive dimensions, " + f"got {global_shape}" + ) + if not isinstance(self.dtype, torch.dtype) or not self.dtype.is_floating_point: + raise TypeError(f"dtype must be a floating-point torch dtype, got {self.dtype}") + object.__setattr__(self, "global_shape", global_shape) + + @property + def parameter_name(self) -> str: + return f"{self.site_id}.mixture_lora.{self.parameter_kind}" + + +@dataclass(frozen=True) +class TransportTensorSpec: + """Tensor identity plus the TP shard that is being transported.""" + + state: MixtureLoraStateSpec + tp_shard_dim: int | None + tp_rank: int + tp_world_size: int + + def __post_init__(self) -> None: + if not isinstance(self.state, MixtureLoraStateSpec): + raise TypeError(f"state must be a MixtureLoraStateSpec, got {type(self.state).__name__}") + _validate_positive_int("tp_world_size", self.tp_world_size) + if not isinstance(self.tp_rank, int) or isinstance(self.tp_rank, bool): + raise TypeError(f"tp_rank must be an integer, got {type(self.tp_rank).__name__}") + if not 0 <= self.tp_rank < self.tp_world_size: + raise ValueError(f"tp_rank must satisfy 0 <= tp_rank < tp_world_size, got {self.tp_rank}") + if self.tp_shard_dim is not None: + if not isinstance(self.tp_shard_dim, int) or isinstance(self.tp_shard_dim, bool): + raise TypeError(f"tp_shard_dim must be an integer or None, got {type(self.tp_shard_dim).__name__}") + if not 0 <= self.tp_shard_dim < len(self.state.global_shape): + raise ValueError( + f"tp_shard_dim must index global_shape {self.state.global_shape}, got {self.tp_shard_dim}" + ) + + @property + def parameter_name(self) -> str: + return self.state.parameter_name + + @property + def schema_version(self) -> int: + return self.state.schema_version + + @property + def site_id(self) -> str: + return self.state.site_id + + @property + def parameter_kind(self) -> MixtureLoraParameterKind: + return self.state.parameter_kind + + @property + def global_shape(self) -> tuple[int, ...]: + return self.state.global_shape + + @property + def dtype(self) -> torch.dtype: + return self.state.dtype + + +@dataclass(frozen=True) +class RoutedLoRAParallelContext: + """Parallel layout passed explicitly to a routed LoRA executor.""" + + target_module: str + sequence_parallel: bool = False + tensor_parallel_group: object | None = None + + def __post_init__(self) -> None: + if not isinstance(self.target_module, str) or not self.target_module.strip(): + raise ValueError("target_module must be a non-empty string") + if not isinstance(self.sequence_parallel, bool): + raise TypeError(f"sequence_parallel must be a bool, got {type(self.sequence_parallel).__name__}") + + +@runtime_checkable +class RoutedLoRAExecutor(Protocol): + """Parameter-free expert execution interface shared by model backends.""" + + def execute( + self, + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + routing_decision: "RoutingDecision", + scale: float, + parallel_context: RoutedLoRAParallelContext | None, + ) -> torch.Tensor: ... + + +class DenseRoutedLoRAExecutor(nn.Module): + """Compute every expert with batched einsums and apply sparse weights.""" + + def forward( + self, + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + routing_decision: "RoutingDecision", + scale: float, + parallel_context: RoutedLoRAParallelContext | None = None, + ) -> torch.Tensor: + return self.execute(x, lora_a, lora_b, routing_decision, scale, parallel_context) + + def execute( + self, + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + routing_decision: "RoutingDecision", + scale: float, + parallel_context: RoutedLoRAParallelContext | None = None, + ) -> torch.Tensor: + _validate_dense_executor_inputs(x, lora_a, lora_b, routing_decision, scale, parallel_context) + + input_shape = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + expert_hidden = torch.einsum("ti,nri->tnr", x_flat, lora_a) + expert_outputs = torch.einsum("tnr,nor->tno", expert_hidden, lora_b) + routing_weights = routing_decision.dense_weights().to(dtype=expert_outputs.dtype) + delta = torch.sum(expert_outputs * routing_weights.unsqueeze(-1), dim=1) + return (delta * scale).reshape(*input_shape, lora_b.shape[1]) + + +def build_mixture_lora_state_specs( + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, + dtype: torch.dtype, +) -> tuple[MixtureLoraStateSpec, ...]: + """Build the executor-independent global parameter schema for one site.""" + + _validate_positive_int("input_size", input_size) + _validate_positive_int("output_size", output_size) + return ( + MixtureLoraStateSpec( + config.schema_version, + site_id, + "experts.lora_A", + (config.num_experts, config.rank, input_size), + dtype, + ), + MixtureLoraStateSpec( + config.schema_version, + site_id, + "experts.lora_B", + (config.num_experts, output_size, config.rank), + dtype, + ), + MixtureLoraStateSpec( + config.schema_version, + site_id, + "router.weight", + (config.num_experts, input_size), + dtype, + ), + ) @dataclass(frozen=True) @@ -181,6 +413,14 @@ def compute_routing_statistics( ) +def mean_routing_balance_loss(statistics: Sequence[RoutingStatistics]) -> torch.Tensor: + """Average independently computed balance losses across routed sites.""" + + if not statistics: + raise ValueError("statistics must contain at least one routed site") + return torch.stack([site_statistics.balance_loss for site_statistics in statistics]).mean() + + def _validate_routing_decision(decision: RoutingDecision) -> None: probs = decision.pre_topk_probs indices = decision.topk_indices @@ -198,3 +438,72 @@ def _validate_routing_decision(decision: RoutingDecision) -> None: raise TypeError("routing probabilities and weights must be FP32") if indices.dtype != torch.long: raise TypeError("topk_indices must use torch.long") + + +def _validate_dense_executor_inputs( + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + decision: RoutingDecision, + scale: float, + parallel_context: RoutedLoRAParallelContext | None, +) -> None: + _validate_routing_decision(decision) + if x.ndim < 2: + raise ValueError(f"x must have at least two dimensions, got {tuple(x.shape)}") + if lora_a.ndim != 3 or lora_b.ndim != 3: + raise ValueError("lora_a and lora_b must have shapes [expert, rank, input] and [expert, output, rank]") + if not torch.is_floating_point(x) or not torch.is_floating_point(lora_a) or not torch.is_floating_point(lora_b): + raise TypeError("x, lora_a, and lora_b must be floating-point tensors") + if x.dtype != lora_a.dtype or x.dtype != lora_b.dtype: + raise TypeError( + f"x, lora_a, and lora_b must use the same dtype, got {x.dtype}, {lora_a.dtype}, {lora_b.dtype}" + ) + if x.device != lora_a.device or x.device != lora_b.device or x.device != decision.pre_topk_probs.device: + raise ValueError("x, expert parameters, and routing tensors must be on the same device") + + num_experts, rank, input_size = lora_a.shape + if lora_b.shape[0] != num_experts or lora_b.shape[2] != rank: + raise ValueError(f"lora_a and lora_b expert/rank dimensions do not agree: {lora_a.shape} and {lora_b.shape}") + if x.shape[-1] != input_size: + raise ValueError(f"x hidden size {x.shape[-1]} does not match lora_a input size {input_size}") + if decision.num_experts != num_experts: + raise ValueError( + f"routing decision has {decision.num_experts} experts but adapter parameters have {num_experts}" + ) + token_count = math.prod(x.shape[:-1]) + if decision.pre_topk_probs.shape[0] != token_count: + raise ValueError( + f"routing decision has {decision.pre_topk_probs.shape[0]} tokens but x contains {token_count}" + ) + _validate_finite_number("scale", scale) + + # TP collectives depend on the target layer contract and are implemented + # by backend executors. The shared dense executor handles local tensors. + if parallel_context is not None and parallel_context.tensor_parallel_group is not None: + raise NotImplementedError("DenseRoutedLoRAExecutor does not perform tensor-parallel collectives") + + +def _validate_positive_int(name: str, value: int) -> None: + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be an integer, got {type(value).__name__}") + if value <= 0: + raise ValueError(f"{name} must be greater than 0, got {value}") + + +def _validate_finite_number( + name: str, + value: float, + minimum: float | None = None, + minimum_inclusive: bool = True, +) -> None: + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError(f"{name} must be a real number, got {type(value).__name__}") + if not math.isfinite(value): + raise ValueError(f"{name} must be finite, got {value}") + if minimum is None: + return + if minimum_inclusive and value < minimum: + raise ValueError(f"{name} must be at least {minimum}, got {value}") + if not minimum_inclusive and value <= minimum: + raise ValueError(f"{name} must be greater than {minimum}, got {value}") diff --git a/tests/utils/test_mixture_lora_routing.py b/tests/utils/test_mixture_lora_routing.py index 52343629d..177033939 100644 --- a/tests/utils/test_mixture_lora_routing.py +++ b/tests/utils/test_mixture_lora_routing.py @@ -5,7 +5,20 @@ import pytest import torch -from relax.utils.mixture_lora import RoutingDecision, compute_routing_statistics, route_topk +from relax.utils.mixture_lora import ( + MIXTURE_LORA_SCHEMA_VERSION, + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + MixtureLoraStateSpec, + RoutedLoRAExecutor, + RoutedLoRAParallelContext, + RoutingDecision, + TransportTensorSpec, + build_mixture_lora_state_specs, + compute_routing_statistics, + mean_routing_balance_loss, + route_topk, +) def _normalized_entropy(values): @@ -14,6 +27,94 @@ def _normalized_entropy(values): return -sum(value * math.log(value) for value in values if value > 0) / math.log(len(values)) +def _config(**overrides): + values = { + "num_experts": 4, + "rank": 2, + "top_k": 2, + "temperature": 1.0, + "aux_loss_coef": 0.01, + "alpha": 4.0, + "target_modules": ("linear_qkv", "linear_proj"), + } + values.update(overrides) + return MixtureLoraConfig(**values) + + +def _reference_routed_lora(x, lora_a, lora_b, decision, scale): + outputs = [] + for token, indices, weights in zip(x.reshape(-1, x.shape[-1]), decision.topk_indices, decision.post_topk_weights): + delta = torch.zeros(lora_b.shape[1], dtype=x.dtype) + for expert_index, weight in zip(indices, weights): + expert_output = lora_b[expert_index] @ (lora_a[expert_index] @ token) + delta = delta + weight.to(x.dtype) * expert_output + outputs.append(delta * scale) + return torch.stack(outputs).reshape(*x.shape[:-1], lora_b.shape[1]) + + +def test_mixture_lora_config_exposes_scale_and_normalizes_targets(): + config = _config(target_modules=["linear_qkv", "linear_proj"]) + + assert config.schema_version == MIXTURE_LORA_SCHEMA_VERSION + assert config.target_modules == ("linear_qkv", "linear_proj") + assert config.scale == 2.0 + + +@pytest.mark.parametrize( + ("overrides", "error"), + [ + ({"num_experts": 1}, ValueError), + ({"rank": 0}, ValueError), + ({"top_k": 0}, ValueError), + ({"top_k": 5}, ValueError), + ({"temperature": 0.0}, ValueError), + ({"aux_loss_coef": -0.1}, ValueError), + ({"alpha": float("inf")}, ValueError), + ({"target_modules": ()}, ValueError), + ({"target_modules": "linear_qkv"}, TypeError), + ({"target_modules": ("linear_qkv", "linear_qkv")}, ValueError), + ], +) +def test_mixture_lora_config_rejects_invalid_values(overrides, error): + with pytest.raises(error): + _config(**overrides) + + +def test_state_specs_fix_parameter_names_and_global_shapes(): + site_id = "decoder.layers.0.self_attention.linear_qkv" + + specs = build_mixture_lora_state_specs(_config(), site_id, input_size=8, output_size=12, dtype=torch.bfloat16) + + assert [spec.parameter_name for spec in specs] == [ + f"{site_id}.mixture_lora.experts.lora_A", + f"{site_id}.mixture_lora.experts.lora_B", + f"{site_id}.mixture_lora.router.weight", + ] + assert [spec.global_shape for spec in specs] == [(4, 2, 8), (4, 12, 2), (4, 8)] + assert all(spec.dtype == torch.bfloat16 for spec in specs) + + transport = TransportTensorSpec(specs[1], tp_shard_dim=1, tp_rank=1, tp_world_size=2) + assert transport.parameter_name == specs[1].parameter_name + assert transport.schema_version == MIXTURE_LORA_SCHEMA_VERSION + assert transport.site_id == site_id + assert transport.parameter_kind == "experts.lora_B" + assert transport.global_shape == (4, 12, 2) + assert transport.dtype == torch.bfloat16 + + +def test_state_and_transport_specs_reject_invalid_layouts(): + with pytest.raises(ValueError, match="global_shape"): + MixtureLoraStateSpec(1, "site", "router.weight", (4, 8, 2), torch.float32) + with pytest.raises(TypeError, match="floating-point"): + MixtureLoraStateSpec(1, "site", "router.weight", (4, 8), torch.int64) + + state = MixtureLoraStateSpec(1, "site", "router.weight", (4, 8), torch.float32) + with pytest.raises(ValueError, match="tp_rank"): + TransportTensorSpec(state, tp_shard_dim=None, tp_rank=2, tp_world_size=2) + with pytest.raises(ValueError, match="tp_shard_dim"): + TransportTensorSpec(state, tp_shard_dim=2, tp_rank=0, tp_world_size=2) + + def test_route_topk_returns_fp32_normalized_weights(): logits = torch.tensor([[4.0, 1.0, 3.0, 2.0], [0.0, 5.0, 2.0, 1.0]], dtype=torch.bfloat16) @@ -90,6 +191,26 @@ def test_balance_loss_backpropagates_through_pre_topk_probabilities(): assert torch.count_nonzero(logits.grad) > 0 +def test_balance_loss_is_averaged_across_sites_after_site_local_statistics(): + first_logits = torch.tensor([[3.0, 2.0, 1.0], [2.0, 0.0, 1.0]], requires_grad=True) + second_logits = torch.tensor([[1.0, 3.0, 2.0], [0.0, 2.0, 3.0]], requires_grad=True) + response_mask = torch.ones(2, dtype=torch.bool) + first = compute_routing_statistics(route_topk(first_logits, 2, 1.0), response_mask) + second = compute_routing_statistics(route_topk(second_logits, 2, 1.0), response_mask) + + loss = mean_routing_balance_loss((first, second)) + + torch.testing.assert_close(loss, (first.balance_loss + second.balance_loss) / 2) + loss.backward() + assert torch.count_nonzero(first_logits.grad) > 0 + assert torch.count_nonzero(second_logits.grad) > 0 + + +def test_balance_loss_requires_at_least_one_site(): + with pytest.raises(ValueError, match="at least one"): + mean_routing_balance_loss(()) + + def test_k_one_has_zero_post_topk_entropy(): decision = route_topk(torch.tensor([[2.0, 1.0], [0.0, 3.0]]), top_k=1, temperature=1.0) stats = compute_routing_statistics(decision, torch.ones(2, dtype=torch.bool)) @@ -111,6 +232,84 @@ def test_no_response_tokens_returns_zero_statistics_and_loss(): assert torch.count_nonzero(logits.grad) == 0 +@pytest.mark.parametrize( + ("dtype", "atol"), + [(torch.float32, 1e-6), (torch.float16, 1e-2), (torch.bfloat16, 2e-2)], +) +def test_dense_executor_matches_independent_expert_loop(dtype, atol): + torch.manual_seed(7) + x = torch.randn(2, 3, 5, dtype=dtype) + lora_a = torch.randn(4, 2, 5, dtype=dtype) + lora_b = torch.randn(4, 7, 2, dtype=dtype) + decision = route_topk(torch.randn(6, 4), top_k=2, temperature=0.7) + original_indices = decision.topk_indices.clone() + original_weights = decision.post_topk_weights.clone() + executor = DenseRoutedLoRAExecutor() + context = RoutedLoRAParallelContext(target_module="linear_qkv", sequence_parallel=False) + + actual = executor.execute(x, lora_a, lora_b, decision, scale=0.5, parallel_context=context) + expected = _reference_routed_lora(x, lora_a, lora_b, decision, scale=0.5) + + assert isinstance(executor, RoutedLoRAExecutor) + assert executor.state_dict() == {} + assert list(executor.parameters()) == [] + assert actual.shape == (2, 3, 7) + assert actual.dtype == dtype + assert torch.allclose(actual.float(), expected.float(), atol=atol) + assert torch.equal(decision.topk_indices, original_indices) + assert torch.equal(decision.post_topk_weights, original_weights) + + +def test_dense_executor_backpropagates_to_experts_input_and_router(): + torch.manual_seed(11) + x = torch.randn(4, 5, requires_grad=True) + lora_a = torch.randn(3, 2, 5, requires_grad=True) + lora_b = torch.randn(3, 7, 2, requires_grad=True) + logits = torch.tensor([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0], [2.0, 1.0, 3.0], [3.0, 1.0, 2.0]], requires_grad=True) + decision = route_topk(logits, top_k=2, temperature=1.0) + + DenseRoutedLoRAExecutor()(x, lora_a, lora_b, decision, scale=1.0).square().mean().backward() + + for tensor in (x, lora_a, lora_b, logits): + assert tensor.grad is not None + assert torch.isfinite(tensor.grad).all() + assert torch.count_nonzero(tensor.grad) > 0 + + +def test_dense_executor_is_repeatable_with_fixed_seed(): + def run_once(): + torch.manual_seed(19) + x = torch.randn(3, 4) + lora_a = torch.randn(4, 2, 4) + lora_b = torch.randn(4, 6, 2) + decision = route_topk(torch.randn(3, 4), top_k=2, temperature=1.0) + return DenseRoutedLoRAExecutor()(x, lora_a, lora_b, decision, scale=0.5) + + assert torch.equal(run_once(), run_once()) + + +def test_dense_executor_rejects_inconsistent_inputs_and_implicit_collectives(): + x = torch.randn(2, 5) + lora_a = torch.randn(4, 2, 5) + lora_b = torch.randn(4, 7, 2) + decision = route_topk(torch.randn(2, 4), top_k=2, temperature=1.0) + executor = DenseRoutedLoRAExecutor() + + with pytest.raises(ValueError, match="hidden size"): + executor(x[:, :4], lora_a, lora_b, decision, scale=1.0) + with pytest.raises(ValueError, match="tokens"): + executor(x.repeat(2, 1), lora_a, lora_b, decision, scale=1.0) + with pytest.raises(NotImplementedError, match="tensor-parallel"): + executor( + x, + lora_a, + lora_b, + decision, + scale=1.0, + parallel_context=RoutedLoRAParallelContext("linear_proj", tensor_parallel_group=object()), + ) + + @pytest.mark.parametrize( ("top_k", "temperature", "error"), [ From bd4a2026971c50c6f8ebcd8010cc7cfd3a61f75d Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 06:37:45 +0800 Subject: [PATCH 03/41] feat(lora): validate Mixture-of-LoRA arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### Summary 接入 Mixture-of-LoRA 命令行参数和启动阶段校验，并保持单 expert 配置继续使用现有 LoRA rollout 路径。 #### Changes 新增 expert 数量、router Top-K、temperature 和 balance loss coefficient 参数；N 大于 1 时要求完整 router 配置、colocate、SGLang TP/DP 为 1，并拒绝 fully-async、merge mode、adapter mode 和不支持的目标层；新增 Mixture 启用判断和共享配置构造；将旧 LoRA rollout mode 校验集中到独立函数，N=1 仍自动选择现有 merge 路径。 #### Verification 参数、PEFT 和单 LoRA 回归共 101 passed、5 skipped。全仓 pre-commit 全部通过，包含 Ruff、docformatter 和 gitleaks。 --- relax/utils/arguments.py | 124 ++++++++++-- relax/utils/megatron_peft_utils.py | 24 +++ tests/utils/test_arguments_mixture_lora.py | 212 +++++++++++++++++++++ 3 files changed, 344 insertions(+), 16 deletions(-) create mode 100644 tests/utils/test_arguments_mixture_lora.py diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 22430e883..e5d90906f 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -14,6 +14,7 @@ from relax.utils import device as device_utils from relax.utils.env import Envs from relax.utils.logging_utils import get_logger +from relax.utils.megatron_peft_utils import build_mixture_lora_config from relax.utils.opd.opd_utils import ( add_opd_arguments, is_managed_opd_teacher_enabled, @@ -1515,6 +1516,30 @@ def add_algo_arguments(parser): default=0, help="LoRA rank for parameter-efficient fine-tuning (0=disabled).", ) + parser.add_argument( + "--lora-num-experts", + type=int, + default=1, + help="Number of LoRA experts. Values greater than 1 enable token-level routing.", + ) + parser.add_argument( + "--lora-router-top-k", + type=int, + default=None, + help="Number of experts selected for each token in Mixture-of-LoRA.", + ) + parser.add_argument( + "--lora-router-temperature", + type=float, + default=None, + help="Softmax temperature for Mixture-of-LoRA router logits.", + ) + parser.add_argument( + "--lora-router-aux-loss-coef", + type=float, + default=None, + help="Coefficient applied to the Mixture-of-LoRA balance loss.", + ) parser.add_argument( "--lora-alpha", type=int, @@ -2762,6 +2787,88 @@ def _normalize_sync_ppo_kl_args(args) -> bool: return True +def _validate_lora_args(args) -> None: + """Validate Mixture settings before applying the legacy LoRA rollout + mode.""" + + num_experts = getattr(args, "lora_num_experts", 1) + if not isinstance(num_experts, int) or isinstance(num_experts, bool): + raise TypeError(f"--lora-num-experts must be an integer, got {type(num_experts).__name__}.") + if num_experts < 1: + raise ValueError(f"--lora-num-experts must be >= 1, got {num_experts}.") + + mixture_fields = ( + ("lora_router_top_k", "--lora-router-top-k"), + ("lora_router_temperature", "--lora-router-temperature"), + ("lora_router_aux_loss_coef", "--lora-router-aux-loss-coef"), + ) + if num_experts == 1: + unexpected = [option for attribute, option in mixture_fields if getattr(args, attribute, None) is not None] + if unexpected: + raise ValueError(f"{', '.join(unexpected)} require --lora-num-experts greater than 1.") + else: + if getattr(args, "lora_rank", 0) <= 0: + raise ValueError("--lora-num-experts greater than 1 requires --lora-rank greater than 0.") + missing = [option for attribute, option in mixture_fields if getattr(args, attribute, None) is None] + if missing: + raise ValueError( + "--lora-num-experts greater than 1 requires explicit values for " + ", ".join(missing) + "." + ) + if getattr(args, "lora_merge_mode", False) or getattr(args, "lora_adapter_mode", False): + raise ValueError("Mixture-of-LoRA does not support --lora-merge-mode or --lora-adapter-mode.") + if getattr(args, "fully_async", False): + raise ValueError("Mixture-of-LoRA does not support --fully-async.") + if not getattr(args, "colocate", False): + raise ValueError("Mixture-of-LoRA requires --colocate.") + + dp_size = getattr(args, "sglang_dp_size", getattr(args, "sglang_data_parallel_size", 1)) + if dp_size != 1: + raise ValueError(f"Mixture-of-LoRA requires SGLang DP size 1, got {dp_size}.") + + tp_size = getattr(args, "sglang_tp_size", None) + if tp_size is None: + pp_size = getattr(args, "sglang_pipeline_parallel_size", getattr(args, "sglang_pp_size", 1)) + gpus_per_engine = getattr(args, "rollout_num_gpus_per_engine", 1) + if pp_size < 1 or gpus_per_engine % pp_size != 0: + raise ValueError( + "Mixture-of-LoRA requires --rollout-num-gpus-per-engine to be divisible by the SGLang PP size." + ) + tp_size = gpus_per_engine // pp_size + if tp_size != 1: + raise ValueError(f"Mixture-of-LoRA requires SGLang TP size 1, got {tp_size}.") + + target_modules = getattr(args, "lora_target_modules", ()) + unsupported_targets = sorted(set(target_modules) - {"linear_qkv", "linear_proj"}) + if unsupported_targets: + raise ValueError( + "Mixture-of-LoRA currently supports only linear_qkv and linear_proj; " + f"unsupported targets: {', '.join(unsupported_targets)}." + ) + + # Reuse the shared config validation for K, temperature, coefficient, + # alpha, rank, and duplicate target modules. + build_mixture_lora_config(args) + return + + if getattr(args, "lora_rank", 0) <= 0: + return + if getattr(args, "lora_merge_mode", False) and getattr(args, "lora_adapter_mode", False): + raise ValueError( + "--lora-merge-mode and --lora-adapter-mode are mutually exclusive; pick one LoRA rollout path." + ) + if getattr(args, "lora_adapter_mode", False) and getattr(args, "sglang_dp_size", 1) != 1: + raise ValueError( + "--lora-adapter-mode requires --sglang-dp-size 1 (SGLang dynamic LoRA loading does not " + "support dp_size > 1)." + ) + if not getattr(args, "lora_merge_mode", False) and not getattr(args, "lora_adapter_mode", False): + logger.info( + "LoRA enabled (lora_rank=%d): forcing --lora-merge-mode (default supported LoRA rollout path).", + args.lora_rank, + ) + args.lora_merge_mode = True + + def slime_validate_args(args): # Backward compatibility: old scripts may pass --enable-gloo-process-groups if not hasattr(args, "use_gloo_process_groups"): @@ -2800,22 +2907,7 @@ def slime_validate_args(args): if args.max_staleness < 0: raise ValueError("--max-staleness must be >= 0.") - if getattr(args, "lora_rank", 0) > 0: - if getattr(args, "lora_merge_mode", False) and getattr(args, "lora_adapter_mode", False): - raise ValueError( - "--lora-merge-mode and --lora-adapter-mode are mutually exclusive; pick one LoRA rollout path." - ) - if getattr(args, "lora_adapter_mode", False) and getattr(args, "sglang_dp_size", 1) != 1: - raise ValueError( - "--lora-adapter-mode requires --sglang-dp-size 1 (SGLang dynamic LoRA loading does not " - "support dp_size > 1)." - ) - if not getattr(args, "lora_merge_mode", False) and not getattr(args, "lora_adapter_mode", False): - logger.info( - "LoRA enabled (lora_rank=%d): forcing --lora-merge-mode (default supported LoRA rollout path).", - args.lora_rank, - ) - args.lora_merge_mode = True + _validate_lora_args(args) # Refuse SGLANG_ENABLE_SPEC_V2=1 with speculative decoding on SGLang <= 0.5.9. # There, spec_v2 routes requests through EAGLEWorkerV2.verify(), which does diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 7fd78a9ab..6cd3752da 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/relax/utils/megatron_peft_utils.py @@ -7,6 +7,8 @@ import torch +from relax.utils.mixture_lora import MixtureLoraConfig + # Fixed name under which the trained policy LoRA adapter is registered on the rollout # engines in adapter mode. Generation requests must pass ``lora_path=LORA_ADAPTER_NAME`` @@ -187,6 +189,26 @@ def is_lora_enabled(args) -> bool: return hasattr(args, "lora_rank") and args.lora_rank > 0 +def is_mixture_lora_enabled(args) -> bool: + """Return whether training uses more than one routed LoRA expert.""" + return is_lora_enabled(args) and getattr(args, "lora_num_experts", 1) > 1 + + +def build_mixture_lora_config(args) -> MixtureLoraConfig | None: + """Build the shared Mixture-of-LoRA config from validated arguments.""" + if not is_mixture_lora_enabled(args): + return None + return MixtureLoraConfig( + num_experts=args.lora_num_experts, + rank=args.lora_rank, + top_k=args.lora_router_top_k, + temperature=args.lora_router_temperature, + aux_loss_coef=args.lora_router_aux_loss_coef, + alpha=args.lora_alpha, + target_modules=tuple(args.lora_target_modules), + ) + + def is_lora_merge_mode(args) -> bool: """Check if LoRA merge mode is enabled. @@ -305,7 +327,9 @@ def build_lora_peft(args): "write_hf_peft_adapter", "extract_lora_delta", "is_lora_enabled", + "is_mixture_lora_enabled", "is_lora_merge_mode", "is_lora_adapter_mode", + "build_mixture_lora_config", "build_lora_peft", ] diff --git a/tests/utils/test_arguments_mixture_lora.py b/tests/utils/test_arguments_mixture_lora.py new file mode 100644 index 000000000..cb658434c --- /dev/null +++ b/tests/utils/test_arguments_mixture_lora.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import argparse +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from relax.utils.megatron_peft_utils import build_mixture_lora_config, is_mixture_lora_enabled + + +@pytest.fixture() +def arguments_module(monkeypatch): + router_pkg = ModuleType("sglang_router") + launch_router = ModuleType("sglang_router.launch_router") + launch_router.RouterArgs = object + monkeypatch.setitem(sys.modules, "sglang_router", router_pkg) + monkeypatch.setitem(sys.modules, "sglang_router.launch_router", launch_router) + + sglang_arguments = ModuleType("relax.backends.sglang.arguments") + sglang_arguments.sglang_parse_args = lambda: None + sglang_arguments.validate_args = lambda args: args + monkeypatch.setitem(sys.modules, "relax.backends.sglang.arguments", sglang_arguments) + + device = ModuleType("relax.utils.device") + device.get_dist_backend = lambda: "gloo" + monkeypatch.setitem(sys.modules, "relax.utils.device", device) + + eval_config = ModuleType("relax.utils.training.eval_config") + eval_config.EvalDatasetConfig = dict + eval_config.build_eval_dataset_configs = lambda args, datasets_config, defaults: [] + eval_config.build_named_prompt_data_configs = lambda values: [] + eval_config.ensure_dataset_list = lambda values: values or [] + monkeypatch.setitem(sys.modules, "relax.utils.training.eval_config", eval_config) + + sys.modules.pop("relax.utils.arguments", None) + module = importlib.import_module("relax.utils.arguments") + yield module + sys.modules.pop("relax.utils.arguments", None) + + +def _args(**overrides): + defaults = dict( + lora_rank=16, + lora_num_experts=4, + lora_router_top_k=2, + lora_router_temperature=1.0, + lora_router_aux_loss_coef=0.01, + lora_alpha=32, + lora_target_modules=["linear_qkv", "linear_proj"], + lora_merge_mode=False, + lora_adapter_mode=False, + fully_async=False, + colocate=True, + sglang_dp_size=1, + sglang_tp_size=1, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def test_parser_defaults_preserve_single_lora_path(arguments_module): + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + args = parser.parse_args([]) + + assert args.lora_num_experts == 1 + assert args.lora_router_top_k is None + assert args.lora_router_temperature is None + assert args.lora_router_aux_loss_coef is None + + +def test_parser_accepts_explicit_mixture_lora_configuration(arguments_module): + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + args = parser.parse_args( + [ + "--lora-rank", + "16", + "--lora-num-experts", + "4", + "--lora-router-top-k", + "2", + "--lora-router-temperature", + "0.8", + "--lora-router-aux-loss-coef", + "0.01", + ] + ) + + assert args.lora_rank == 16 + assert args.lora_num_experts == 4 + assert args.lora_router_top_k == 2 + assert args.lora_router_temperature == 0.8 + assert args.lora_router_aux_loss_coef == 0.01 + + +def test_valid_mixture_configuration_uses_shared_config(arguments_module): + args = _args() + + arguments_module._validate_lora_args(args) + config = build_mixture_lora_config(args) + + assert is_mixture_lora_enabled(args) + assert config.num_experts == 4 + assert config.rank == 16 + assert config.top_k == 2 + assert config.temperature == 1.0 + assert config.aux_loss_coef == 0.01 + assert config.scale == 2.0 + assert args.lora_merge_mode is False + assert args.lora_adapter_mode is False + + +def test_missing_mixture_values_are_reported_together(arguments_module): + args = _args( + lora_router_top_k=None, + lora_router_temperature=None, + lora_router_aux_loss_coef=None, + ) + + with pytest.raises(ValueError) as error: + arguments_module._validate_lora_args(args) + + message = str(error.value) + assert "--lora-router-top-k" in message + assert "--lora-router-temperature" in message + assert "--lora-router-aux-loss-coef" in message + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"lora_num_experts": 0}, "lora-num-experts"), + ({"lora_rank": 0}, "lora-rank"), + ({"lora_router_top_k": 0}, "top_k"), + ({"lora_router_top_k": 5}, "top_k"), + ({"lora_router_temperature": 0.0}, "temperature"), + ({"lora_router_temperature": float("inf")}, "temperature"), + ({"lora_router_aux_loss_coef": -0.1}, "aux_loss_coef"), + ({"lora_router_aux_loss_coef": float("inf")}, "aux_loss_coef"), + ({"lora_target_modules": ["linear_qkv", "linear_qkv"]}, "duplicates"), + ], +) +def test_mixture_configuration_rejects_invalid_values(arguments_module, overrides, message): + with pytest.raises((TypeError, ValueError), match=message): + arguments_module._validate_lora_args(_args(**overrides)) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"lora_merge_mode": True}, "merge-mode"), + ({"lora_adapter_mode": True}, "adapter-mode"), + ({"fully_async": True}, "fully-async"), + ({"colocate": False}, "colocate"), + ({"sglang_dp_size": 2}, "DP size 1"), + ({"sglang_tp_size": 2}, "TP size 1"), + ({"lora_target_modules": ["linear_fc1"]}, "linear_fc1"), + ], +) +def test_mixture_configuration_rejects_unsupported_modes(arguments_module, overrides, message): + with pytest.raises(ValueError, match=message): + arguments_module._validate_lora_args(_args(**overrides)) + + +def test_mixture_tp_size_can_be_derived_from_rollout_and_pp(arguments_module): + args = _args(rollout_num_gpus_per_engine=2, sglang_pipeline_parallel_size=2) + del args.sglang_tp_size + + arguments_module._validate_lora_args(args) + + +def test_single_lora_keeps_legacy_auto_merge_behavior(arguments_module): + args = _args( + lora_num_experts=1, + lora_router_top_k=None, + lora_router_temperature=None, + lora_router_aux_loss_coef=None, + ) + + arguments_module._validate_lora_args(args) + + assert not is_mixture_lora_enabled(args) + assert build_mixture_lora_config(args) is None + assert args.lora_merge_mode is True + assert args.lora_adapter_mode is False + + +def test_legacy_namespace_without_mixture_fields_is_unchanged(arguments_module): + args = SimpleNamespace(lora_rank=8, lora_merge_mode=False, lora_adapter_mode=False, sglang_dp_size=1) + + arguments_module._validate_lora_args(args) + + assert args.lora_merge_mode is True + + +def test_single_lora_rejects_mixture_only_values(arguments_module): + args = _args( + lora_num_experts=1, + lora_router_top_k=1, + lora_router_temperature=None, + lora_router_aux_loss_coef=None, + ) + + with pytest.raises(ValueError, match="lora-router-top-k"): + arguments_module._validate_lora_args(args) From dea7f5d04dbfcae4369e91dafd1addc36fd2fddc Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 07:12:13 +0800 Subject: [PATCH 04/41] feat(lora): integrate Mixture-of-LoRA training modules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: - 接入 Megatron 训练端的 Mixture-of-LoRA 核心模块和 Bridge PEFT 注入路径。 - 保持单 expert 配置继续使用原有 LoRA 实现。 Changes: - 新增打包的 expert A/B 参数、FP32 token router、dense Top-K 执行和线性层包装器。 - 在模型 provider 中按 expert 数选择实现，并在 optimizer 建立前检查 base 冻结及 expert/router 可训练状态。 - 扩展参数分类与分项统计，保持 base state key 和 Megatron 线性层返回协议。 - 增加 CPU 数值、梯度、初始化、state key、Bridge 匹配、provider 路径和 CUDA profiler 测试。 Verification: - 相关 pytest：137 passed, 5 skipped。 - H20 CUDA forward/backward 与 profiler 测试通过。 - pre-commit run --all-files 全部通过。 --- relax/backends/megatron/mixture_lora.py | 256 +++++++++++++++++ relax/backends/megatron/model_provider.py | 47 ++- relax/utils/megatron_peft_utils.py | 85 +++++- tests/backends/megatron/test_mixture_lora.py | 267 ++++++++++++++++++ .../megatron/test_model_provider_vpp.py | 42 +++ tests/utils/test_megatron_peft_utils.py | 59 ++++ 6 files changed, 746 insertions(+), 10 deletions(-) create mode 100644 relax/backends/megatron/mixture_lora.py create mode 100644 tests/backends/megatron/test_mixture_lora.py diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py new file mode 100644 index 000000000..3dc70de89 --- /dev/null +++ b/relax/backends/megatron/mixture_lora.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Megatron model modules and Bridge injection for Mixture-of-LoRA.""" + +import math +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from relax.utils.mixture_lora import DenseRoutedLoRAExecutor, MixtureLoraConfig, RoutingDecision, route_topk + + +class MixtureLoRAExperts(nn.Module): + """LoRA expert parameters stored in one stable logical layout.""" + + def __init__( + self, + config: MixtureLoraConfig, + input_size: int, + output_size: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.lora_A = nn.Parameter( + torch.empty(config.num_experts, config.rank, input_size, device=device, dtype=dtype) + ) + self.lora_B = nn.Parameter( + torch.empty(config.num_experts, output_size, config.rank, device=device, dtype=dtype) + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + # Match the current Bridge LoRA initialization independently for each expert. + for expert_weight in self.lora_A: + nn.init.xavier_uniform_(expert_weight) + nn.init.zeros_(self.lora_B) + + +class MixtureLoRARouter(nn.Module): + """Per-token linear router with FP32 logits.""" + + def __init__( + self, + num_experts: int, + input_size: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(num_experts, input_size, device=device, dtype=dtype)) + nn.init.normal_(self.weight, mean=0.0, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x.float(), self.weight.float()) + + +class MixtureLoRAAdapter(nn.Module): + """Route tokens across LoRA experts and combine their outputs.""" + + def __init__( + self, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, + *, + dropout: float, + device: torch.device, + dtype: torch.dtype, + ) -> None: + super().__init__() + if not isinstance(site_id, str) or not site_id.strip(): + raise ValueError("site_id must be a non-empty string") + if not math.isfinite(dropout) or not 0.0 <= dropout < 1.0: + raise ValueError(f"dropout must satisfy 0 <= dropout < 1, got {dropout}") + if not dtype.is_floating_point: + raise TypeError(f"dtype must be floating point, got {dtype}") + + self.config = config + self.site_id = site_id + self.experts = MixtureLoRAExperts(config, input_size, output_size, device=device, dtype=dtype) + self.router = MixtureLoRARouter(config.num_experts, input_size, device=device, dtype=dtype) + self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() + self.executor = DenseRoutedLoRAExecutor() + + def route(self, x: torch.Tensor) -> RoutingDecision: + logits = self.router(x.reshape(-1, x.shape[-1])) + return route_topk(logits, self.config.top_k, self.config.temperature) + + def forward_with_routing(self, x: torch.Tensor) -> tuple[torch.Tensor, RoutingDecision]: + decision = self.route(x) + delta = self.executor( + self.dropout(x), + self.experts.lora_A, + self.experts.lora_B, + decision, + self.config.scale, + ) + return delta, decision + + def forward(self, x: torch.Tensor) -> torch.Tensor: + delta, _ = self.forward_with_routing(x) + return delta + + +class MixtureParallelLinearAdapter(nn.Module): + """Add a routed LoRA delta while preserving Megatron's linear protocol.""" + + def __init__( + self, + to_wrap: nn.Module, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, + *, + dropout: float, + ) -> None: + super().__init__() + try: + first_parameter = next(to_wrap.parameters()) + except StopIteration as error: + raise ValueError(f"Mixture-of-LoRA target {site_id} has no parameters") from error + + self.to_wrap = to_wrap + self.mixture_lora = MixtureLoRAAdapter( + config, + site_id, + input_size, + output_size, + dropout=dropout, + device=first_parameter.device, + dtype=first_parameter.dtype, + ) + self._adapter_enabled = True + + def enable_adapter_layers(self) -> None: + self._adapter_enabled = True + + def disable_adapter_layers(self) -> None: + self._adapter_enabled = False + + def base_linear_forward( + self, x: torch.Tensor, *args: Any, **kwargs: Any + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """Normalize the return forms used by Megatron parallel linear layers.""" + + result = self.to_wrap(x, *args, **kwargs) + if not isinstance(result, tuple): + raise TypeError(f"{type(self.to_wrap).__name__} must return a tuple, got {type(result).__name__}") + + bias = None + adapter_input = x + if len(result) == 2: + output, bias = result + if isinstance(output, tuple) and len(output) == 2: + output, adapter_input = output + elif len(result) == 3: + output, bias, adapter_input = result + else: + raise ValueError(f"{type(self.to_wrap).__name__} returned an unsupported tuple of length {len(result)}") + return output, bias, adapter_input + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor | None]: + output, bias, adapter_input = self.base_linear_forward(x, *args, **kwargs) + if not self._adapter_enabled: + return output, bias + adapter_output = self.mixture_lora(adapter_input.contiguous()).reshape(output.shape) + return output + adapter_output, bias + + def state_dict( + self, + destination: dict[str, Any] | None = None, + prefix: str = "", + keep_vars: bool = False, + ) -> dict[str, Any]: + """Keep base keys unchanged and store routed parameters under mixture_lora.""" + + if destination is None: + destination = {} + self.to_wrap.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) + self.mixture_lora.state_dict( + destination=destination, + prefix=f"{prefix}mixture_lora.", + keep_vars=keep_vars, + ) + return destination + + +def build_mixture_lora_peft(config: MixtureLoraConfig, dropout: float): + """Build a Bridge PEFT object that injects routed adapters at matched sites.""" + + try: + from megatron.bridge.peft.base import PEFT + from megatron.bridge.peft.module_matcher import ModuleMatcher + from megatron.bridge.peft.utils import get_adapter_attributes_from_linear + from megatron.core import parallel_state + except ImportError as error: + raise RuntimeError( + "Mixture-of-LoRA training requires a Megatron-Bridge image with PEFT support. " + "Please upgrade the training image." + ) from error + + @dataclass + class MixtureLoRAPEFT(PEFT, ModuleMatcher): + target_modules: list[str] = field(default_factory=list) + mixture_config: MixtureLoraConfig | None = None + dropout: float = 0.0 + + def transform( + self, + module: nn.Module, + name: str | None = None, + prefix: str | None = None, + ) -> nn.Module: + if isinstance(module, MixtureParallelLinearAdapter): + return module + match = self.match(module, name, prefix) + if match is None: + return module + if self.mixture_config is None: + raise RuntimeError("Mixture-of-LoRA PEFT is missing its configuration") + if parallel_state.get_tensor_model_parallel_world_size() != 1: + raise NotImplementedError("Mixture-of-LoRA tensor parallel execution is not implemented yet") + + _, full_name = match + attributes = get_adapter_attributes_from_linear(module) + return MixtureParallelLinearAdapter( + module, + self.mixture_config, + full_name, + attributes.in_features, + attributes.out_features, + dropout=self.dropout, + ) + + return MixtureLoRAPEFT( + target_modules=list(config.target_modules), + mixture_config=config, + dropout=dropout, + ) + + +__all__ = [ + "MixtureLoRAAdapter", + "MixtureLoRAExperts", + "MixtureLoRARouter", + "MixtureParallelLinearAdapter", + "build_mixture_lora_peft", +] diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 7e9805608..e514c5909 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -25,7 +25,14 @@ from relax.utils.device import is_npu_available from relax.utils.logging_utils import get_logger -from relax.utils.megatron_peft_utils import build_lora_peft, count_adapter_parameters, is_lora_enabled +from relax.utils.megatron_peft_utils import ( + build_lora_peft, + build_mixture_lora_config, + count_adapter_parameters, + is_lora_enabled, + is_mixture_lora_enabled, + validate_and_count_mixture_lora_parameters, +) from relax.utils.misc import load_function from relax.utils.training.ppo_utils import install_critic_value_head_in_provider @@ -472,17 +479,39 @@ def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwarg model = original_provider(pre_process=pre_process, post_process=post_process) try: - peft = build_lora_peft(args) + if is_mixture_lora_enabled(args): + from relax.backends.megatron.mixture_lora import build_mixture_lora_peft + + mixture_config = build_mixture_lora_config(args) + if mixture_config is None: + raise RuntimeError("Mixture-of-LoRA is enabled but its validated configuration is missing") + peft = build_mixture_lora_peft(mixture_config, args.lora_dropout) + else: + peft = build_lora_peft(args) model = peft(model, training=True) + mixture_parameter_counts = None + if is_mixture_lora_enabled(args): + mixture_parameter_counts = validate_and_count_mixture_lora_parameters(model) if dist.is_initialized() and dist.get_rank() == 0: - adapter_params, total_params, percentage = count_adapter_parameters(model) - logger.info( - f"LoRA enabled: rank={args.lora_rank}, alpha={args.lora_alpha}, " - f"adapter_params={adapter_params:,} ({percentage:.2f}% of {total_params:,} total)" - ) + if is_mixture_lora_enabled(args): + if mixture_parameter_counts is None: + raise RuntimeError("Mixture-of-LoRA parameter validation did not run") + base_params, expert_params, router_params, total_params = mixture_parameter_counts + trainable_percentage = 100 * (expert_params + router_params) / total_params + logger.info( + f"Mixture-of-LoRA enabled: experts={args.lora_num_experts}, rank={args.lora_rank}, " + f"top_k={args.lora_router_top_k}, base_params={base_params:,}, " + f"expert_params={expert_params:,}, router_params={router_params:,}, " + f"trainable={trainable_percentage:.2f}% of {total_params:,} total" + ) + else: + adapter_params, total_params, percentage = count_adapter_parameters(model) + logger.info( + f"LoRA enabled: rank={args.lora_rank}, alpha={args.lora_alpha}, " + f"adapter_params={adapter_params:,} ({percentage:.2f}% of {total_params:,} total)" + ) except RuntimeError: - # build_lora_peft already raises a clear upgrade-hint message when the - # Megatron-Bridge image lacks PEFT support; don't shadow it. + # PEFT builders provide actionable configuration and dependency errors. raise except Exception as e: raise RuntimeError( diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 6cd3752da..ad6a0516c 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/relax/utils/megatron_peft_utils.py @@ -14,6 +14,11 @@ # engines in adapter mode. Generation requests must pass ``lora_path=LORA_ADAPTER_NAME`` # for the adapter to take effect (see sglang_rollout.generate and _push_lora_adapter). LORA_ADAPTER_NAME = "relax_policy_lora" +_MIXTURE_LORA_EXPERT_SUFFIXES = ( + "mixture_lora.experts.lora_A", + "mixture_lora.experts.lora_B", +) +_MIXTURE_LORA_ROUTER_SUFFIX = "mixture_lora.router.weight" def count_adapter_parameters(model) -> Tuple[int, int, float]: @@ -46,6 +51,52 @@ def count_adapter_parameters(model) -> Tuple[int, int, float]: return adapter_params, total_params, percentage +def validate_and_count_mixture_lora_parameters(model) -> tuple[int, int, int, int]: + """Validate trainable parameters and return base/expert/router/total + counts.""" + + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + if not isinstance(unwrapped, list): + unwrapped = [unwrapped] + + base_params = 0 + expert_params = 0 + router_params = 0 + total_params = 0 + frozen_mixture_params = [] + unexpected_trainable_params = [] + + for chunk in unwrapped: + for name, param in chunk.named_parameters(): + param_count = param.numel() + total_params += param_count + if is_mixture_lora_expert_param(name): + expert_params += param_count + if not param.requires_grad: + frozen_mixture_params.append(name) + elif is_mixture_lora_router_param(name): + router_params += param_count + if not param.requires_grad: + frozen_mixture_params.append(name) + else: + base_params += param_count + if param.requires_grad: + unexpected_trainable_params.append(name) + + if expert_params == 0 or router_params == 0: + raise RuntimeError("Mixture-of-LoRA injection produced no expert or router parameters") + if frozen_mixture_params: + names = ", ".join(frozen_mixture_params[:5]) + raise RuntimeError(f"Mixture-of-LoRA parameters must remain trainable: {names}") + if unexpected_trainable_params: + names = ", ".join(unexpected_trainable_params[:5]) + raise RuntimeError(f"Base parameters must be frozen during Mixture-of-LoRA training: {names}") + + return base_params, expert_params, router_params, total_params + + # Fused Megatron module name -> the HF-style projections it expands to. Megatron is the # canonical form used everywhere internally (CLI, injection, weight sync); the mapping is # one-to-many because a single fused Megatron linear covers several HF projections @@ -106,7 +157,35 @@ def is_lora_adapter_param(name: str) -> bool: - Megatron-Bridge: ``...adapter.linear_in.weight`` / ``...adapter.linear_out.weight`` - Standard PEFT: ``...lora_A.weight`` / ``...lora_B.weight`` """ - return ".lora_A." in name or ".lora_B." in name or ".adapter.linear_in." in name or ".adapter.linear_out." in name + return ( + ".lora_A." in name + or ".lora_B." in name + or ".adapter.linear_in." in name + or ".adapter.linear_out." in name + or is_mixture_lora_param(name) + ) + + +def is_mixture_lora_param(name: str) -> bool: + """Return whether ``name`` identifies an expert or router parameter.""" + + return is_mixture_lora_expert_param(name) or is_mixture_lora_router_param(name) + + +def is_mixture_lora_expert_param(name: str) -> bool: + """Return whether ``name`` identifies a Mixture LoRA A or B tensor.""" + + return _has_parameter_suffix(name, _MIXTURE_LORA_EXPERT_SUFFIXES) + + +def is_mixture_lora_router_param(name: str) -> bool: + """Return whether ``name`` identifies a Mixture LoRA router tensor.""" + + return _has_parameter_suffix(name, (_MIXTURE_LORA_ROUTER_SUFFIX,)) + + +def _has_parameter_suffix(name: str, suffixes: tuple[str, ...]) -> bool: + return any(name == suffix or name.endswith(f".{suffix}") for suffix in suffixes) def build_hf_peft_config_dict( @@ -322,11 +401,15 @@ def build_lora_peft(args): __all__ = [ "LORA_ADAPTER_NAME", "count_adapter_parameters", + "validate_and_count_mixture_lora_parameters", "convert_megatron_to_hf_target_modules", "MEGATRON_TO_HF_MODULES", "write_hf_peft_adapter", "extract_lora_delta", "is_lora_enabled", + "is_mixture_lora_param", + "is_mixture_lora_expert_param", + "is_mixture_lora_router_param", "is_mixture_lora_enabled", "is_lora_merge_mode", "is_lora_adapter_mode", diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py new file mode 100644 index 000000000..5cf0625b4 --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import sys +import types +from dataclasses import dataclass, field +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from relax.backends.megatron.mixture_lora import ( + MixtureLoRAAdapter, + MixtureParallelLinearAdapter, + build_mixture_lora_peft, +) +from relax.utils.mixture_lora import MixtureLoraConfig + + +def _config(*, num_experts=3, top_k=2, rank=2, alpha=4.0): + return MixtureLoraConfig( + num_experts=num_experts, + rank=rank, + top_k=top_k, + temperature=0.7, + aux_loss_coef=0.01, + alpha=alpha, + target_modules=("linear_qkv", "linear_proj"), + ) + + +class _TupleLinear(torch.nn.Module): + def __init__(self, input_size, output_size, *, return_mode="standard"): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(output_size, input_size)) + self.bias = torch.nn.Parameter(torch.randn(output_size)) + self.return_mode = return_mode + + def forward(self, x): + output = F.linear(x, self.weight) + if self.return_mode == "standard": + return output, self.bias + if self.return_mode == "layernorm": + return (output, x + 1.0), self.bias + if self.return_mode == "three": + return output, self.bias, x + 1.0 + raise AssertionError(f"unknown return mode: {self.return_mode}") + + +@pytest.mark.parametrize(("num_experts", "top_k"), [(3, 2), (3, 3)]) +def test_mixture_lora_forward_matches_expert_reference(num_experts, top_k): + torch.manual_seed(7) + config = _config(num_experts=num_experts, top_k=top_k) + adapter = MixtureLoRAAdapter( + config, + "decoder.layers.0.self_attention.linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + adapter.experts.lora_A.copy_(torch.randn_like(adapter.experts.lora_A)) + adapter.experts.lora_B.copy_(torch.randn_like(adapter.experts.lora_B)) + adapter.router.weight.copy_(torch.randn_like(adapter.router.weight)) + + x = torch.randn(2, 3, 4) + actual, decision = adapter.forward_with_routing(x) + x_flat = x.reshape(-1, x.shape[-1]) + dense_weights = decision.dense_weights() + expected_tokens = [] + for token, token_weights in zip(x_flat, dense_weights, strict=True): + expert_sum = torch.zeros(5) + for expert_index, expert_weight in enumerate(token_weights): + hidden = adapter.experts.lora_A[expert_index] @ token + expert_output = adapter.experts.lora_B[expert_index] @ hidden + expert_sum += expert_weight * expert_output + expected_tokens.append(expert_sum * config.scale) + expected = torch.stack(expected_tokens).reshape(2, 3, 5) + + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_mixture_lora_router_normalizes_in_fp32(dtype): + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=dtype, + ) + + decision = adapter.route(torch.randn(2, 3, 4, dtype=dtype)) + + assert decision.pre_topk_probs.dtype == torch.float32 + assert decision.post_topk_weights.dtype == torch.float32 + + +def test_mixture_lora_parameter_layout_and_initialization(): + config = _config(num_experts=4, rank=3) + adapter = MixtureLoRAAdapter( + config, + "linear_proj", + 6, + 7, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert adapter.experts.lora_A.shape == (4, 3, 6) + assert adapter.experts.lora_B.shape == (4, 7, 3) + assert adapter.router.weight.shape == (4, 6) + assert torch.count_nonzero(adapter.experts.lora_A) > 0 + assert torch.count_nonzero(adapter.router.weight) > 0 + assert torch.count_nonzero(adapter.experts.lora_B) == 0 + assert list(adapter.executor.parameters()) == [] + assert adapter.executor.state_dict() == {} + + +def test_mixture_lora_wrapper_freezes_base_and_routes_gradients(): + torch.manual_seed(11) + base = _TupleLinear(4, 5) + for parameter in base.parameters(): + parameter.requires_grad = False + wrapper = MixtureParallelLinearAdapter(base, _config(), "linear_qkv", 4, 5, dropout=0.0) + with torch.no_grad(): + wrapper.mixture_lora.experts.lora_B.normal_(mean=0.0, std=0.2) + + output, _ = wrapper(torch.randn(3, 2, 4)) + output.square().mean().backward() + + assert all(parameter.grad is None for parameter in base.parameters()) + assert wrapper.mixture_lora.experts.lora_A.grad is not None + assert wrapper.mixture_lora.experts.lora_B.grad is not None + assert wrapper.mixture_lora.router.weight.grad is not None + assert torch.count_nonzero(wrapper.mixture_lora.router.weight.grad) > 0 + + +@pytest.mark.parametrize("return_mode", ["standard", "layernorm", "three"]) +def test_mixture_lora_wrapper_preserves_linear_return_protocol(return_mode): + base = _TupleLinear(4, 5, return_mode=return_mode) + wrapper = MixtureParallelLinearAdapter(base, _config(), "linear_qkv", 4, 5, dropout=0.0) + x = torch.randn(2, 3, 4) + + base_output, base_bias, _ = wrapper.base_linear_forward(x) + output, bias = wrapper(x) + + assert output.shape == base_output.shape == (2, 3, 5) + assert bias is base_bias + torch.testing.assert_close(output, base_output) + + wrapper.disable_adapter_layers() + disabled_output, disabled_bias = wrapper(x) + torch.testing.assert_close(disabled_output, base_output) + assert disabled_bias is base_bias + + +def test_mixture_lora_wrapper_keeps_base_state_keys_stable(): + wrapper = MixtureParallelLinearAdapter(_TupleLinear(4, 5), _config(), "linear_qkv", 4, 5, dropout=0.0) + + state = wrapper.state_dict() + + assert set(state) == { + "weight", + "bias", + "mixture_lora.experts.lora_A", + "mixture_lora.experts.lora_B", + "mixture_lora.router.weight", + } + + +def _install_fake_bridge(monkeypatch): + megatron = types.ModuleType("megatron") + bridge = types.ModuleType("megatron.bridge") + peft_package = types.ModuleType("megatron.bridge.peft") + base_module = types.ModuleType("megatron.bridge.peft.base") + matcher_module = types.ModuleType("megatron.bridge.peft.module_matcher") + utils_module = types.ModuleType("megatron.bridge.peft.utils") + core_module = types.ModuleType("megatron.core") + + @dataclass + class FakePEFT: + params_to_save: set[str] = field(default_factory=set, init=False) + + def __call__(self, model, training=True): + for parameter in model.parameters(): + parameter.requires_grad = False + model.linear_qkv = self.transform(model.linear_qkv, "linear_qkv", "decoder.layers.0.self_attention") + return model + + @dataclass + class FakeModuleMatcher: + target_modules: list[str] = field(default_factory=list) + + def match(self, module, name=None, prefix=None): + if name not in self.target_modules: + return None + return name, f"{prefix}.{name}" if prefix else name + + base_module.PEFT = FakePEFT + matcher_module.ModuleMatcher = FakeModuleMatcher + utils_module.get_adapter_attributes_from_linear = lambda module: SimpleNamespace( + in_features=module.weight.shape[1], out_features=module.weight.shape[0] + ) + core_module.parallel_state = SimpleNamespace(get_tensor_model_parallel_world_size=lambda: 1) + + modules = { + "megatron": megatron, + "megatron.bridge": bridge, + "megatron.bridge.peft": peft_package, + "megatron.bridge.peft.base": base_module, + "megatron.bridge.peft.module_matcher": matcher_module, + "megatron.bridge.peft.utils": utils_module, + "megatron.core": core_module, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + +def test_mixture_lora_peft_uses_bridge_matcher_and_freezes_base(monkeypatch): + _install_fake_bridge(monkeypatch) + config = _config() + peft = build_mixture_lora_peft(config, dropout=0.0) + model = torch.nn.Module() + model.linear_qkv = _TupleLinear(4, 5) + + transformed = peft(model, training=True) + + assert isinstance(transformed.linear_qkv, MixtureParallelLinearAdapter) + assert transformed.linear_qkv.mixture_lora.site_id == "decoder.layers.0.self_attention.linear_qkv" + assert all(not parameter.requires_grad for parameter in transformed.linear_qkv.to_wrap.parameters()) + assert all(parameter.requires_grad for parameter in transformed.linear_qkv.mixture_lora.parameters()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the forward profiler check") +def test_mixture_lora_cuda_forward_backward_has_no_tensor_to_host_sync(): + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 16, + 24, + dropout=0.0, + device=torch.device("cuda"), + dtype=torch.float16, + ) + with torch.no_grad(): + adapter.experts.lora_B.normal_(mean=0.0, std=0.02) + x = torch.randn(8, 4, 16, device="cuda", dtype=torch.float16) + + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA] + ) as profile: + output = adapter(x) + output.float().square().mean().backward() + torch.cuda.synchronize() + + event_names = {event.key for event in profile.key_averages()} + assert "aten::_local_scalar_dense" not in event_names + assert output.shape == (8, 4, 24) + assert adapter.experts.lora_A.grad is not None + assert adapter.experts.lora_B.grad is not None + assert adapter.router.weight.grad is not None diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 82c20510b..92a5447d6 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -251,6 +251,48 @@ def __init__(self): assert not hasattr(model, ppo_utils._RELAX_HF_OUTPUT_LAYER_ATTR) +def test_model_provider_uses_mixture_lora_for_multiple_experts(monkeypatch): + module, _ = _load_model_provider(monkeypatch) + model = object() + config = object() + calls = [] + mixture_module = types.ModuleType("relax.backends.megatron.mixture_lora") + + def build_mixture_lora_peft(received_config, dropout): + calls.append((received_config, dropout)) + return lambda received_model, training: received_model + + mixture_module.build_mixture_lora_peft = build_mixture_lora_peft + monkeypatch.setitem(sys.modules, "relax.backends.megatron.mixture_lora", mixture_module) + monkeypatch.setattr(module, "build_mixture_lora_config", lambda args: config) + monkeypatch.setattr(module, "build_lora_peft", lambda args: pytest.fail("single LoRA factory was called")) + monkeypatch.setattr(module, "validate_and_count_mixture_lora_parameters", lambda model: (100, 20, 4, 124)) + args = SimpleNamespace(lora_rank=16, lora_num_experts=4, lora_dropout=0.1) + + provider = module.wrap_model_provider_with_lora(lambda **kwargs: model, args) + + assert provider() is model + assert calls == [(config, 0.1)] + + +def test_model_provider_keeps_single_expert_on_existing_lora_path(monkeypatch): + module, _ = _load_model_provider(monkeypatch) + model = object() + calls = [] + + def build_lora_peft(args): + calls.append(args.lora_num_experts) + return lambda received_model, training: received_model + + monkeypatch.setattr(module, "build_lora_peft", build_lora_peft) + args = SimpleNamespace(lora_rank=16, lora_num_experts=1) + + provider = module.wrap_model_provider_with_lora(lambda **kwargs: model, args) + + assert provider() is model + assert calls == [1] + + def test_critic_value_head_validation_accepts_ddp_and_optimizer_ownership(monkeypatch): class _FakeBridgeModel(torch.nn.Module): def __init__(self): diff --git a/tests/utils/test_megatron_peft_utils.py b/tests/utils/test_megatron_peft_utils.py index f16219e21..1a3ddc9e9 100644 --- a/tests/utils/test_megatron_peft_utils.py +++ b/tests/utils/test_megatron_peft_utils.py @@ -29,6 +29,8 @@ is_lora_adapter_param, is_lora_enabled, is_lora_merge_mode, + is_mixture_lora_param, + validate_and_count_mixture_lora_parameters, write_hf_peft_adapter, ) @@ -94,6 +96,18 @@ def test_is_lora_adapter_param_guards_against_substring_false_positive(self): assert is_lora_adapter_param("decoder.layers.0.mlp.lora_gate.weight") is False assert is_lora_adapter_param("decoder.layers.0.mlp.lora_Attention.weight") is False + @pytest.mark.parametrize( + "name", + [ + "decoder.layers.0.self_attention.linear_qkv.mixture_lora.experts.lora_A", + "decoder.layers.0.self_attention.linear_qkv.mixture_lora.experts.lora_B", + "decoder.layers.0.self_attention.linear_qkv.mixture_lora.router.weight", + ], + ) + def test_is_mixture_lora_param_matches_stable_parameter_names(self, name): + assert is_mixture_lora_param(name) is True + assert is_lora_adapter_param(name) is True + class TestWriteHfPeftAdapter: def test_write_hf_peft_adapter_round_trip(self, tmp_path): @@ -213,6 +227,51 @@ def named_parameters(self, *a, **k): assert pct == 0 +class TestValidateAndCountMixtureLoraParameters: + @staticmethod + def _stub_unwrap_model(monkeypatch): + fake_utils = types.ModuleType("megatron.core.utils") + fake_utils.unwrap_model = lambda model: model + monkeypatch.setitem(sys.modules, "megatron", types.ModuleType("megatron")) + monkeypatch.setitem(sys.modules, "megatron.core", types.ModuleType("megatron.core")) + monkeypatch.setitem(sys.modules, "megatron.core.utils", fake_utils) + + @staticmethod + def _model(): + model = torch.nn.Module() + model.base_weight = torch.nn.Parameter(torch.zeros(5), requires_grad=False) + model.mixture_lora = torch.nn.Module() + model.mixture_lora.experts = torch.nn.Module() + model.mixture_lora.experts.lora_A = torch.nn.Parameter(torch.zeros(2, 3, 4)) + model.mixture_lora.experts.lora_B = torch.nn.Parameter(torch.zeros(2, 5, 3)) + model.mixture_lora.router = torch.nn.Module() + model.mixture_lora.router.weight = torch.nn.Parameter(torch.zeros(2, 4)) + return model + + def test_validate_and_count_mixture_lora_parameters(self, monkeypatch): + self._stub_unwrap_model(monkeypatch) + + counts = validate_and_count_mixture_lora_parameters(self._model()) + + assert counts == (5, 54, 8, 67) + + def test_validate_rejects_trainable_base_parameter(self, monkeypatch): + self._stub_unwrap_model(monkeypatch) + model = self._model() + model.base_weight.requires_grad = True + + with pytest.raises(RuntimeError, match="Base parameters must be frozen"): + validate_and_count_mixture_lora_parameters(model) + + def test_validate_rejects_frozen_router_parameter(self, monkeypatch): + self._stub_unwrap_model(monkeypatch) + model = self._model() + model.mixture_lora.router.weight.requires_grad = False + + with pytest.raises(RuntimeError, match="must remain trainable"): + validate_and_count_mixture_lora_parameters(model) + + class TestBridgeParamPrefixes: """Base<->adapter name prefixes must round-trip so merge-mode can pair them. From 8c166f2e76296580a1e84cde321230195234b093 Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 07:51:52 +0800 Subject: [PATCH 05/41] feat(lora): add routed balance loss context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: - 为每个训练 microbatch 建立 Mixture-of-LoRA 路由上下文。 - 将逐 site balance loss 直接附着到模型激活，保持 policy loss 标量不变。 Changes: - 使用 full_loss_masks 过滤 prompt、padding 和 dummy token，并处理 batch-first mask 与 Megatron 激活布局。 - 普通 loss、per-token loss 和动态 batch 共用同一个 microbatch 缩放 helper。 - 按全模型 site 数归一化 aux loss，分离可导 aux tensor 与无梯度路由记录。 - 捕获 activation checkpoint 的路由上下文，并补齐冻结 base 时的 recompute input-grad 路径。 - 增加 mask、缩放、dummy、router 梯度、checkpoint 恢复、真实 Bridge 和真实 ColumnParallelLinear 测试。 Verification: - 相关 pytest：152 passed, 8 skipped。 - 真实 Megatron/Bridge 定向测试：2 passed。 - H20 CUDA profiler 测试通过。 - pre-commit run --all-files 全部通过。 --- relax/backends/megatron/loss.py | 44 ++- relax/backends/megatron/mixture_lora.py | 257 +++++++++++++++- relax/backends/megatron/model.py | 126 +++++++- relax/backends/megatron/model_provider.py | 8 +- tests/backends/megatron/test_mixture_lora.py | 281 +++++++++++++++++- .../megatron/test_model_provider_vpp.py | 2 + 6 files changed, 672 insertions(+), 46 deletions(-) diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 506bd6221..05754ec3b 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -42,6 +42,7 @@ maybe_padded_total_lengths, slice_log_prob_with_cp, ) +from .mixture_lora import get_microbatch_objective_scale def get_responses( @@ -1406,34 +1407,25 @@ def loss_function( is_dummy = batch.get("__is_dummy__", False) explicit_loss_scale = batch.get("__loss_scale__", None) - # Rescale the loss for Megatron's gradient accumulation. The non-per-token - # branch folds in the DP(+CP) world size (cancelled by DDP's 1/dp_cp grad - # scaling); the per-token branch does NO CP scaling (normalization is the - # all-reduced CP-local token count in finalize_model_grads). + # This is the final scale after Megatron's schedule divides non-per-token + # losses by num_microbatches. Routed aux losses use the same helper before + # they are attached to intermediate activations. global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size) + microbatch_objective_scale = get_microbatch_objective_scale( + calculate_per_token_loss=args.calculate_per_token_loss, + is_dummy=is_dummy, + explicit_loss_scale=explicit_loss_scale, + num_microbatches=num_microbatches, + global_batch_size=global_batch_size, + data_parallel_world_size_with_cp=mpu.get_data_parallel_world_size(with_context_parallel=True), + ) if not args.calculate_per_token_loss: - if is_dummy: - # Zero-out gradient contribution but keep the autograd graph - # connected so PP/CP backward collectives still complete. - loss = 0.0 * loss - elif explicit_loss_scale is not None: - loss = loss * explicit_loss_scale - else: - loss = ( - loss - * num_microbatches - / global_batch_size - * mpu.get_data_parallel_world_size(with_context_parallel=True) - ) - else: - if is_dummy: - loss = 0.0 * loss - # Non-dummy per-token path: do NOT scale by cp_size. `loss` is the - # CP-local token-sum; finalize_model_grads normalizes the summed gradient - # by the all-reduced CP-local `num_tokens`. A `* cp_size` here would weight - # each sample by its CP degree — wrong when CP differs across micro-batches - # (dynamic CP). Under static CP the removed factor exactly cancels the old - # full-count denominator, leaving the final loss/grad unchanged. + loss = loss * microbatch_objective_scale * num_microbatches + elif is_dummy: + # Keep the graph connected so PP/CP backward collectives still complete. + loss = 0.0 * loss + # The non-dummy per-token path is normalized by the all-reduced CP-local + # token count in finalize_model_grads, so it needs no scale here. effective_num_tokens = torch.zeros_like(num_tokens) if is_dummy else num_tokens log_values = torch.tensor( diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 3dc70de89..667391b57 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -3,14 +3,248 @@ """Megatron model modules and Bridge injection for Mixture-of-LoRA.""" import math +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass, field -from typing import Any +from functools import wraps +from typing import Any, Iterator import torch import torch.nn.functional as F from torch import nn -from relax.utils.mixture_lora import DenseRoutedLoRAExecutor, MixtureLoraConfig, RoutingDecision, route_topk +from relax.utils.mixture_lora import ( + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + RoutingDecision, + RoutingStatistics, + compute_routing_statistics, + route_topk, +) + + +@dataclass(frozen=True) +class MixtureLoRARoutingRecord: + """Detached routing values retained for one microbatch and site.""" + + key: tuple[int, int, str] + statistics: RoutingStatistics + balance_loss: torch.Tensor + aux_loss: torch.Tensor + + +class _AttachAuxLoss(torch.autograd.Function): + """Attach an auxiliary loss to an activation with an explicit scale.""" + + @staticmethod + def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor, backward_scale: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(aux_loss, backward_scale) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]: + aux_loss, backward_scale = ctx.saved_tensors + aux_loss_grad = torch.ones_like(aux_loss) * backward_scale.reshape(()) + return grad_output, aux_loss_grad, None + + +@dataclass +class MixtureLoRARoutingContext: + """Routing state and aux-loss scaling for one training microbatch.""" + + optimizer_step: int + microbatch_id: int + response_mask: torch.Tensor + num_microbatches: int + num_sites: int + num_samples: int + calculate_per_token_loss: bool + objective_scale: float + main_loss_backward_scale: torch.Tensor + is_dummy: bool = False + records: dict[str, MixtureLoRARoutingRecord] = field(default_factory=dict, init=False) + + def __post_init__(self) -> None: + if self.optimizer_step < 0 or self.microbatch_id < 0: + raise ValueError("optimizer_step and microbatch_id must be non-negative") + if self.num_microbatches <= 0 or self.num_sites <= 0 or self.num_samples <= 0: + raise ValueError("num_microbatches, num_sites, and num_samples must be positive") + if not torch.is_tensor(self.response_mask) or self.response_mask.ndim not in (1, 2): + raise ValueError("response_mask must be a one- or two-dimensional tensor") + if not torch.is_tensor(self.main_loss_backward_scale) or self.main_loss_backward_scale.numel() != 1: + raise ValueError("main_loss_backward_scale must be a one-element tensor") + if not math.isfinite(self.objective_scale) or self.objective_scale < 0: + raise ValueError("objective_scale must be finite and non-negative") + + def response_mask_for(self, x: torch.Tensor) -> torch.Tensor: + """Align a batch-first mask with Megatron's activation layout.""" + + activation_shape = tuple(x.shape[:-1]) + mask = self.response_mask + if tuple(mask.shape) == activation_shape: + aligned = mask + elif x.ndim == 3 and mask.ndim == 2 and tuple(mask.shape) == (x.shape[1], x.shape[0]): + aligned = mask.transpose(0, 1) + else: + raise ValueError( + f"response_mask shape {tuple(mask.shape)} does not match activation token layout {activation_shape}" + ) + if aligned.device != x.device: + raise ValueError("response_mask and routed activation must be on the same device") + if self.is_dummy: + aligned = torch.zeros_like(aligned) + return aligned.reshape(-1) + + def attach_aux_loss( + self, + output: torch.Tensor, + x: torch.Tensor, + site_id: str, + config: MixtureLoraConfig, + decision: RoutingDecision, + ) -> torch.Tensor: + """Attach one site's balance loss and replace its detached record.""" + + response_mask = self.response_mask_for(x) + statistics = compute_routing_statistics(decision, response_mask) + # This is a microbatch-level F_e * P_e objective. It is not an + # average of independent per-token losses. + balance_loss = statistics.balance_loss + site_aux_loss = balance_loss * (config.aux_loss_coef / self.num_sites) + if self.calculate_per_token_loss: + aux_loss_payload = site_aux_loss * statistics.valid_token_count + else: + aux_loss_payload = site_aux_loss * self.num_samples + + objective_aux_loss = aux_loss_payload * self.objective_scale + key = (self.optimizer_step, self.microbatch_id, site_id) + self.records[site_id] = MixtureLoRARoutingRecord( + key=key, + statistics=_detach_routing_statistics(statistics), + balance_loss=balance_loss.detach(), + aux_loss=objective_aux_loss.detach(), + ) + + backward_scale = self.main_loss_backward_scale.to(device=output.device) * self.objective_scale + return _AttachAuxLoss.apply(output, aux_loss_payload, backward_scale) + + +_ACTIVE_ROUTING_CONTEXT: ContextVar[MixtureLoRARoutingContext | None] = ContextVar( + "mixture_lora_routing_context", default=None +) + + +@contextmanager +def activate_mixture_lora_routing_context(context: MixtureLoRARoutingContext) -> Iterator[None]: + """Make a microbatch routing context visible to routed adapters.""" + + token = _ACTIVE_ROUTING_CONTEXT.set(context) + try: + yield + finally: + _ACTIVE_ROUTING_CONTEXT.reset(token) + + +def get_mixture_lora_routing_context() -> MixtureLoRARoutingContext | None: + """Return the active microbatch routing context, if one exists.""" + + return _ACTIVE_ROUTING_CONTEXT.get() + + +def install_mixture_lora_checkpoint_context() -> None: + """Restore the active routing context during Megatron recomputation.""" + + from megatron.core import tensor_parallel + + checkpoint = tensor_parallel.checkpoint + if getattr(checkpoint, "_relax_mixture_lora_context", False): + return + + @wraps(checkpoint) + def checkpoint_with_routing_context(function, distribute_saved_activations, *args): + routing_context = get_mixture_lora_routing_context() + if routing_context is None: + return checkpoint(function, distribute_saved_activations, *args) + + @wraps(function) + def run_with_routing_context(*function_args): + with activate_mixture_lora_routing_context(routing_context): + return function(*function_args) + + return checkpoint(run_with_routing_context, distribute_saved_activations, *args) + + checkpoint_with_routing_context._relax_mixture_lora_context = True + tensor_parallel.checkpoint = checkpoint_with_routing_context + + +def ensure_mixture_lora_recompute_inputs_grad(model) -> None: + """Keep full activation recompute reachable when the base is frozen.""" + + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + model_chunks = unwrapped if isinstance(unwrapped, list) else [unwrapped] + for model_chunk in model_chunks: + config = getattr(model_chunk, "config", None) + if config is None or getattr(config, "recompute_method", None) is None: + continue + for module in model_chunk.modules(): + if not isinstance(module, TransformerBlock) or getattr( + module, "_relax_mixture_lora_input_grad_patched", False + ): + continue + original_forward = module.forward + + @wraps(original_forward) + def forward_with_input_grad(hidden_states, *args, _original_forward=original_forward, **kwargs): + if ( + torch.is_tensor(hidden_states) + and hidden_states.is_floating_point() + and not hidden_states.requires_grad + ): + hidden_states = hidden_states.detach().requires_grad_(True) + return _original_forward(hidden_states, *args, **kwargs) + + module.forward = forward_with_input_grad + module._relax_mixture_lora_input_grad_patched = True + + +def get_microbatch_objective_scale( + *, + calculate_per_token_loss: bool, + is_dummy: bool, + explicit_loss_scale: float | None, + num_microbatches: int, + global_batch_size: int, + data_parallel_world_size_with_cp: int, +) -> float: + """Return the final per-microbatch scale applied before parameter grads.""" + + if num_microbatches <= 0 or global_batch_size <= 0 or data_parallel_world_size_with_cp <= 0: + raise ValueError("microbatch, global batch, and data-parallel sizes must be positive") + if is_dummy: + return 0.0 + if calculate_per_token_loss: + return 1.0 + if explicit_loss_scale is not None: + if not math.isfinite(explicit_loss_scale) or explicit_loss_scale < 0: + raise ValueError("explicit_loss_scale must be finite and non-negative") + return explicit_loss_scale / num_microbatches + return data_parallel_world_size_with_cp / global_batch_size + + +def _detach_routing_statistics(statistics: RoutingStatistics) -> RoutingStatistics: + return RoutingStatistics( + pre_topk_prob_sum=statistics.pre_topk_prob_sum.detach(), + post_topk_weight_sum=statistics.post_topk_weight_sum.detach(), + selection_count=statistics.selection_count.detach(), + top1_count=statistics.top1_count.detach(), + pre_topk_entropy_sum=statistics.pre_topk_entropy_sum.detach(), + post_topk_entropy_sum=statistics.post_topk_entropy_sum.detach(), + valid_token_count=statistics.valid_token_count.detach(), + top_k=statistics.top_k, + ) class MixtureLoRAExperts(nn.Module): @@ -102,6 +336,9 @@ def forward_with_routing(self, x: torch.Tensor) -> tuple[torch.Tensor, RoutingDe decision, self.config.scale, ) + routing_context = get_mixture_lora_routing_context() + if routing_context is not None: + delta = routing_context.attach_aux_loss(delta, x, self.site_id, self.config, decision) return delta, decision def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -149,7 +386,8 @@ def disable_adapter_layers(self) -> None: def base_linear_forward( self, x: torch.Tensor, *args: Any, **kwargs: Any ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: - """Normalize the return forms used by Megatron parallel linear layers.""" + """Normalize the return forms used by Megatron parallel linear + layers.""" result = self.to_wrap(x, *args, **kwargs) if not isinstance(result, tuple): @@ -180,7 +418,8 @@ def state_dict( prefix: str = "", keep_vars: bool = False, ) -> dict[str, Any]: - """Keep base keys unchanged and store routed parameters under mixture_lora.""" + """Keep base keys unchanged and store routed parameters under + mixture_lora.""" if destination is None: destination = {} @@ -194,7 +433,8 @@ def state_dict( def build_mixture_lora_peft(config: MixtureLoraConfig, dropout: float): - """Build a Bridge PEFT object that injects routed adapters at matched sites.""" + """Build a Bridge PEFT object that injects routed adapters at matched + sites.""" try: from megatron.bridge.peft.base import PEFT @@ -248,9 +488,16 @@ def transform( __all__ = [ + "MixtureLoRARoutingContext", + "MixtureLoRARoutingRecord", "MixtureLoRAAdapter", "MixtureLoRAExperts", "MixtureLoRARouter", "MixtureParallelLinearAdapter", + "activate_mixture_lora_routing_context", "build_mixture_lora_peft", + "ensure_mixture_lora_recompute_inputs_grad", + "get_microbatch_objective_scale", + "get_mixture_lora_routing_context", + "install_mixture_lora_checkpoint_context", ] diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 4196a2d4d..15b891dbd 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -7,7 +7,7 @@ import uuid from argparse import Namespace from collections.abc import Callable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from functools import partial from pathlib import Path @@ -31,7 +31,7 @@ from relax.utils.data.stream_dataloader import StreamingTQIterator from relax.utils.env import Envs from relax.utils.logging_utils import get_logger -from relax.utils.megatron_peft_utils import is_lora_enabled +from relax.utils.megatron_peft_utils import is_lora_enabled, is_mixture_lora_enabled from relax.utils.memory_utils import clear_memory from relax.utils.opd.opd_utils import consume_opd_train_data from relax.utils.timer import timer @@ -45,12 +45,82 @@ from .checkpoint import load_checkpoint, save_checkpoint from .data import DataIterator, get_batch from .loss import loss_function +from .mixture_lora import ( + MixtureLoRARoutingContext, + MixtureParallelLinearAdapter, + activate_mixture_lora_routing_context, + get_microbatch_objective_scale, +) from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze logger = get_logger(__name__) +def _get_global_mixture_lora_site_count(args: Namespace, model: Sequence[DDP]) -> int: + cached_count = getattr(args, "_mixture_lora_global_site_count", None) + if cached_count is not None: + return cached_count + + local_site_ids = { + module.mixture_lora.site_id + for model_chunk in model + for module in model_chunk.modules() + if isinstance(module, MixtureParallelLinearAdapter) + } + site_count = torch.tensor( + len(local_site_ids), + dtype=torch.int64, + device=next(model[0].parameters()).device, + ) + if mpu.get_pipeline_model_parallel_world_size() > 1: + torch.distributed.all_reduce(site_count, group=mpu.get_pipeline_model_parallel_group()) + global_site_count = int(site_count.item()) + if global_site_count <= 0: + raise RuntimeError("Mixture-of-LoRA is enabled but the model contains no routed sites") + args._mixture_lora_global_site_count = global_site_count + return global_site_count + + +def _build_mixture_lora_routing_context( + args: Namespace, + batch: dict, + model: GPTModel, + *, + optimizer_step: int, + microbatch_id: int, + num_microbatches: int, + num_sites: int, +) -> MixtureLoRARoutingContext: + model_config = get_model_config(model) + loss_scale_input = torch.ones(1, device=batch["full_loss_masks"].device) + main_loss_backward_scale = ( + model_config.grad_scale_func(loss_scale_input) + if model_config.grad_scale_func is not None + else loss_scale_input + ) + objective_scale = get_microbatch_objective_scale( + calculate_per_token_loss=args.calculate_per_token_loss, + is_dummy=batch.get("__is_dummy__", False), + explicit_loss_scale=batch.get("__loss_scale__", None), + num_microbatches=num_microbatches, + global_batch_size=batch.get("dynamic_global_batch_size", args.global_batch_size), + data_parallel_world_size_with_cp=mpu.get_data_parallel_world_size(with_context_parallel=True), + ) + return MixtureLoRARoutingContext( + optimizer_step=optimizer_step, + microbatch_id=microbatch_id, + response_mask=batch["full_loss_masks"], + num_microbatches=num_microbatches, + num_sites=num_sites, + num_samples=len(batch["response_lengths"]), + calculate_per_token_loss=args.calculate_per_token_loss, + objective_scale=objective_scale, + main_loss_backward_scale=main_loss_backward_scale.detach().clone(), + is_dummy=batch.get("__is_dummy__", False), + ) + + def _find_lm_output_layer(model: torch.nn.Module) -> torch.nn.Module | None: """Walk DDP / bridge-VL wrappers to the lm_head; None on non-last PP stages. @@ -984,6 +1054,11 @@ def train_one_step( custom_before_train_step_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler) main_loss_has_tokens = False + mixture_lora_enabled = is_mixture_lora_enabled(args) + mixture_lora_num_sites = _get_global_mixture_lora_site_count(args, model) if mixture_lora_enabled else 0 + # Checkpoint closures retain these contexts until their microbatch backward completes. + routing_contexts: list[MixtureLoRARoutingContext] = [] + next_routing_microbatch_id = 0 def forward_step( data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False @@ -1003,7 +1078,7 @@ def forward_step( (loss, num_elems, {"keys": list[str], "values": torch.Tensor}). """ - nonlocal main_loss_has_tokens + nonlocal main_loss_has_tokens, next_routing_microbatch_id is_vl_model = getattr(args, "is_vl_model", False) sft_chunked = _should_use_sft_chunked(args) # Get the batch. @@ -1037,6 +1112,25 @@ def forward_step( if args.ci_test and args.enable_mtp_training: main_loss_has_tokens = main_loss_has_tokens or _main_loss_has_tokens(batch) + routing_context = None + if mixture_lora_enabled: + routing_context = _build_mixture_lora_routing_context( + args, + batch, + model, + optimizer_step=step_id, + microbatch_id=next_routing_microbatch_id, + num_microbatches=num_microbatches, + num_sites=mixture_lora_num_sites, + ) + routing_contexts.append(routing_context) + next_routing_microbatch_id += 1 + + def routing_scope(): + if routing_context is None: + return nullcontext() + return activate_mixture_lora_routing_context(routing_context) + if Envs.ENABLE_ROUTING_REPLAY: old_stage = os.environ["ROUTING_REPLAY_STAGE"] os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" @@ -1051,14 +1145,15 @@ def forward_step( # chunked-logits incompatibility is enforced as a hard assert in # arguments.py.slime_validate_args, so sft_chunked is guaranteed # False here — no runtime fallback or advisory needed. - output_tensor = model.build_schedule_plan( - input_ids=batch["tokens"], - position_ids=None, - attention_mask=None, - labels=None, - packed_seq_params=batch["packed_seq_params"], - loss_mask=batch["full_loss_masks"], - ) + with routing_scope(): + output_tensor = model.build_schedule_plan( + input_ids=batch["tokens"], + position_ids=None, + attention_mask=None, + labels=None, + packed_seq_params=batch["packed_seq_params"], + loss_mask=batch["full_loss_masks"], + ) else: has_mm_inputs = batch.get("multimodal_train_inputs", None) is not None needs_unsplit = is_vl_model or has_mm_inputs or getattr(args, "uses_unsplit_forward", False) @@ -1111,9 +1206,11 @@ def forward_step( model, mtp_output_layer_calls=mtp_output_layer_calls, ) as lm_head_forward: - output_tensor = model(**forward_kwargs) + with routing_scope(): + output_tensor = model(**forward_kwargs) else: - output_tensor = model(**forward_kwargs) + with routing_scope(): + output_tensor = model(**forward_kwargs) if Envs.ENABLE_ROUTING_REPLAY: os.environ["ROUTING_REPLAY_STAGE"] = old_stage @@ -1170,6 +1267,9 @@ def forward_step( if _dcp_orig_cp_group is not None: inner.pg_collection.cp = _dcp_orig_cp_group + # All checkpoint recomputation and aux backward work is complete here. + routing_contexts.clear() + # CI check: verify only MTP parameters have non-zero gradients when truncation happens # This check must happen before optimizer.step() as gradients may be modified during step if args.ci_test and args.enable_mtp_training: diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index e514c5909..10a9a1a2b 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -480,7 +480,11 @@ def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwarg try: if is_mixture_lora_enabled(args): - from relax.backends.megatron.mixture_lora import build_mixture_lora_peft + from relax.backends.megatron.mixture_lora import ( + build_mixture_lora_peft, + ensure_mixture_lora_recompute_inputs_grad, + install_mixture_lora_checkpoint_context, + ) mixture_config = build_mixture_lora_config(args) if mixture_config is None: @@ -491,6 +495,8 @@ def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwarg model = peft(model, training=True) mixture_parameter_counts = None if is_mixture_lora_enabled(args): + ensure_mixture_lora_recompute_inputs_grad(model) + install_mixture_lora_checkpoint_context() mixture_parameter_counts = validate_and_count_mixture_lora_parameters(model) if dist.is_initialized() and dist.get_rank() == 0: if is_mixture_lora_enabled(args): diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 5cf0625b4..360e5df90 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -11,10 +11,16 @@ from relax.backends.megatron.mixture_lora import ( MixtureLoRAAdapter, + MixtureLoRARoutingContext, MixtureParallelLinearAdapter, + activate_mixture_lora_routing_context, build_mixture_lora_peft, + ensure_mixture_lora_recompute_inputs_grad, + get_microbatch_objective_scale, + get_mixture_lora_routing_context, + install_mixture_lora_checkpoint_context, ) -from relax.utils.mixture_lora import MixtureLoraConfig +from relax.utils.mixture_lora import MixtureLoraConfig, compute_routing_statistics def _config(*, num_experts=3, top_k=2, rank=2, alpha=4.0): @@ -174,6 +180,209 @@ def test_mixture_lora_wrapper_keeps_base_state_keys_stable(): } +@pytest.mark.parametrize( + ("calculate_per_token_loss", "is_dummy", "explicit_loss_scale", "expected"), + [ + (False, False, None, 0.125), + (False, False, 0.5, 0.125), + (True, False, None, 1.0), + (False, True, None, 0.0), + (True, True, None, 0.0), + ], +) +def test_microbatch_objective_scale_matches_training_modes( + calculate_per_token_loss, is_dummy, explicit_loss_scale, expected +): + scale = get_microbatch_objective_scale( + calculate_per_token_loss=calculate_per_token_loss, + is_dummy=is_dummy, + explicit_loss_scale=explicit_loss_scale, + num_microbatches=4, + global_batch_size=32, + data_parallel_world_size_with_cp=4, + ) + + assert scale == expected + + +def test_routing_context_aligns_batch_first_response_mask_and_records_site(): + adapter = MixtureLoRAAdapter( + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + response_mask = torch.tensor([[0, 1, 1], [1, 0, 0]], dtype=torch.bool) + context = MixtureLoRARoutingContext( + optimizer_step=7, + microbatch_id=2, + response_mask=response_mask, + num_microbatches=4, + num_sites=2, + num_samples=2, + calculate_per_token_loss=False, + objective_scale=0.25, + main_loss_backward_scale=torch.ones(1), + ) + x = torch.randn(3, 2, 4) + + with activate_mixture_lora_routing_context(context): + _, decision = adapter.forward_with_routing(x) + + record = context.records[adapter.site_id] + expected_statistics = compute_routing_statistics(decision, response_mask.transpose(0, 1).reshape(-1)) + assert record.key == (7, 2, adapter.site_id) + torch.testing.assert_close(record.statistics.valid_token_count, torch.tensor(3.0)) + torch.testing.assert_close(record.balance_loss, expected_statistics.balance_loss) + torch.testing.assert_close( + record.aux_loss, + expected_statistics.balance_loss * adapter.config.aux_loss_coef / 2 * 2 * 0.25, + ) + assert get_mixture_lora_routing_context() is None + + +@pytest.mark.parametrize("calculate_per_token_loss", [False, True]) +def test_routing_context_attaches_expected_router_gradient(calculate_per_token_loss): + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + response_mask = torch.tensor([[0, 1, 1], [1, 1, 0]], dtype=torch.bool) + objective_scale = 1.0 if calculate_per_token_loss else 0.25 + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=response_mask, + num_microbatches=2, + num_sites=1, + num_samples=2, + calculate_per_token_loss=calculate_per_token_loss, + objective_scale=objective_scale, + main_loss_backward_scale=torch.ones(1), + ) + x = torch.randn(3, 2, 4) + + with activate_mixture_lora_routing_context(context): + output, decision = adapter.forward_with_routing(x) + statistics = compute_routing_statistics(decision, response_mask.transpose(0, 1).reshape(-1)) + count_multiplier = statistics.valid_token_count if calculate_per_token_loss else context.num_samples + expected_objective = statistics.balance_loss * adapter.config.aux_loss_coef * count_multiplier * objective_scale + expected_router_grad = torch.autograd.grad(expected_objective, adapter.router.weight, retain_graph=True)[0] + + output.sum().backward() + + torch.testing.assert_close(adapter.router.weight.grad, expected_router_grad) + + +def test_dummy_routing_context_records_zero_aux_loss(): + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=1, + response_mask=torch.ones(1, 3), + num_microbatches=2, + num_sites=1, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=0.0, + main_loss_backward_scale=torch.ones(1), + is_dummy=True, + ) + + with activate_mixture_lora_routing_context(context): + adapter(torch.randn(3, 1, 4)) + + record = context.records[adapter.site_id] + torch.testing.assert_close(record.statistics.valid_token_count, torch.tensor(0.0)) + torch.testing.assert_close(record.balance_loss, torch.tensor(0.0)) + torch.testing.assert_close(record.aux_loss, torch.tensor(0.0)) + + +def test_checkpoint_wrapper_restores_captured_routing_context(monkeypatch): + megatron_module = types.ModuleType("megatron") + core_module = types.ModuleType("megatron.core") + tensor_parallel_module = types.ModuleType("megatron.core.tensor_parallel") + captured_functions = [] + + def checkpoint(function, distribute_saved_activations, *args): + captured_functions.append(function) + return function(*args) + + tensor_parallel_module.checkpoint = checkpoint + core_module.tensor_parallel = tensor_parallel_module + monkeypatch.setitem(sys.modules, "megatron", megatron_module) + monkeypatch.setitem(sys.modules, "megatron.core", core_module) + monkeypatch.setitem(sys.modules, "megatron.core.tensor_parallel", tensor_parallel_module) + install_mixture_lora_checkpoint_context() + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=torch.ones(1, 1), + num_microbatches=1, + num_sites=1, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + ) + observed_contexts = [] + + def run(x): + observed_contexts.append(get_mixture_lora_routing_context()) + return x + + with activate_mixture_lora_routing_context(context): + tensor_parallel_module.checkpoint(run, False, torch.ones(1)) + captured_function = captured_functions[0] + observed_contexts.clear() + + captured_function(torch.ones(1)) + + assert observed_contexts == [context] + + +def test_recompute_input_grad_patch_recognizes_mixture_adapter(monkeypatch): + transformer_block_module = types.ModuleType("megatron.core.transformer.transformer_block") + utils_module = types.ModuleType("megatron.core.utils") + + class FakeTransformerBlock(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states + + transformer_block_module.TransformerBlock = FakeTransformerBlock + utils_module.unwrap_model = lambda model: model + monkeypatch.setitem(sys.modules, "megatron.core.transformer", types.ModuleType("megatron.core.transformer")) + monkeypatch.setitem(sys.modules, "megatron.core.transformer.transformer_block", transformer_block_module) + monkeypatch.setitem(sys.modules, "megatron.core.utils", utils_module) + model = torch.nn.Module() + model.config = SimpleNamespace(recompute_method="uniform") + model.block = FakeTransformerBlock() + + ensure_mixture_lora_recompute_inputs_grad(model) + patched_forward = model.block.forward + output = model.block(torch.ones(2, 3)) + ensure_mixture_lora_recompute_inputs_grad(model) + + assert output.requires_grad + assert model.block.forward is patched_forward + + def _install_fake_bridge(monkeypatch): megatron = types.ModuleType("megatron") bridge = types.ModuleType("megatron.bridge") @@ -237,6 +446,76 @@ def test_mixture_lora_peft_uses_bridge_matcher_and_freezes_base(monkeypatch): assert all(parameter.requires_grad for parameter in transformed.linear_qkv.mixture_lora.parameters()) +def test_mixture_lora_peft_instantiates_with_real_bridge_when_available(): + pytest.importorskip("megatron.bridge.peft.base") + + peft = build_mixture_lora_peft(_config(), dropout=0.0) + + assert peft.target_modules == ["linear_qkv", "linear_proj"] + assert peft.mixture_config == _config() + + +def test_mixture_lora_peft_wraps_real_column_parallel_linear_when_available(tmp_path): + pytest.importorskip("megatron.bridge.peft.base") + from megatron.core import parallel_state + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + from megatron.core.transformer.transformer_config import TransformerConfig + + if torch.distributed.is_initialized(): + pytest.skip("test requires ownership of the single-process distributed state") + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{tmp_path / 'distributed_init'}", + rank=0, + world_size=1, + ) + parallel_state.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + try: + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=4, + num_attention_heads=1, + use_cpu_initialization=True, + ) + model = torch.nn.Module() + model.linear_qkv = ColumnParallelLinear( + 4, + 5, + config=transformer_config, + init_method=lambda weight: torch.nn.init.normal_(weight, mean=0.0, std=0.02), + bias=False, + gather_output=False, + skip_bias_add=True, + ) + + transformed = build_mixture_lora_peft(_config(), dropout=0.0)(model, training=True) + routing_context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=torch.tensor([[0, 1, 1]], dtype=torch.bool), + num_microbatches=1, + num_sites=1, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + ) + with activate_mixture_lora_routing_context(routing_context): + output, bias = transformed.linear_qkv(torch.randn(3, 1, 4)) + output.sum().backward() + + assert isinstance(transformed.linear_qkv, MixtureParallelLinearAdapter) + assert output.shape == (3, 1, 5) + assert bias is None + assert all(not parameter.requires_grad for parameter in transformed.linear_qkv.to_wrap.parameters()) + assert all(parameter.requires_grad for parameter in transformed.linear_qkv.mixture_lora.parameters()) + assert all(parameter.grad is None for parameter in transformed.linear_qkv.to_wrap.parameters()) + assert torch.count_nonzero(transformed.linear_qkv.mixture_lora.router.weight.grad) > 0 + finally: + parallel_state.destroy_model_parallel() + torch.distributed.destroy_process_group() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the forward profiler check") def test_mixture_lora_cuda_forward_backward_has_no_tensor_to_host_sync(): adapter = MixtureLoRAAdapter( diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 92a5447d6..992988a8d 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -263,6 +263,8 @@ def build_mixture_lora_peft(received_config, dropout): return lambda received_model, training: received_model mixture_module.build_mixture_lora_peft = build_mixture_lora_peft + mixture_module.ensure_mixture_lora_recompute_inputs_grad = lambda model: None + mixture_module.install_mixture_lora_checkpoint_context = lambda: None monkeypatch.setitem(sys.modules, "relax.backends.megatron.mixture_lora", mixture_module) monkeypatch.setattr(module, "build_mixture_lora_config", lambda args: config) monkeypatch.setattr(module, "build_lora_peft", lambda args: pytest.fail("single LoRA factory was called")) From e825c080618c56097cd3fe5285964019f5875a24 Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 08:13:22 +0800 Subject: [PATCH 06/41] feat(lora): aggregate routed expert metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: - 在 optimizer step 结束后汇总每个 Mixture-of-LoRA site 的路由统计并接入现有训练日志。 - 保留实际训练目标中的 aux loss，同时输出用于判断 expert 塌缩的逐 site 与全局指标。 Changes: - 按固定 site_id 表打包 Top-K 前概率、Top-K 后权重、选择次数、Top-1 次数、熵和有效 token 数。 - 先在 DP/CP group 汇总原始统计，再在 pipeline group 合并各 stage 持有的 site。 - 输出 expert 平均概率、平均激活权重、选择份额、Top-1 比例、归一化熵、balance loss 和实际 aux loss。 - 使用训练目标权重汇总 balance loss，支持普通 loss、per-token loss 和 dummy microbatch。 - activation recompute 使用相同 key 覆盖记录，step 聚合完成后清理上下文。 - 增加手算指标、空记录、recompute 去重及两进程 DP/PP 聚合测试。 Verification: - 相关 pytest：108 passed, 2 skipped。 - 真实 Megatron/Bridge 定向测试：2 passed。 - 真实 Megatron/Bridge 环境导入训练模块成功。 - pre-commit run --all-files 全部通过。 --- relax/backends/megatron/mixture_lora.py | 121 ++++++++++++++- relax/backends/megatron/model.py | 95 +++++++++--- tests/backends/megatron/test_mixture_lora.py | 131 ++++++++++++++++ .../megatron/test_mixture_lora_distributed.py | 142 ++++++++++++++++++ 4 files changed, 464 insertions(+), 25 deletions(-) create mode 100644 tests/backends/megatron/test_mixture_lora_distributed.py diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 667391b57..31360183c 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -31,6 +31,7 @@ class MixtureLoRARoutingRecord: statistics: RoutingStatistics balance_loss: torch.Tensor aux_loss: torch.Tensor + objective_weight: torch.Tensor class _AttachAuxLoss(torch.autograd.Function): @@ -111,18 +112,21 @@ def attach_aux_loss( # average of independent per-token losses. balance_loss = statistics.balance_loss site_aux_loss = balance_loss * (config.aux_loss_coef / self.num_sites) - if self.calculate_per_token_loss: - aux_loss_payload = site_aux_loss * statistics.valid_token_count - else: - aux_loss_payload = site_aux_loss * self.num_samples - - objective_aux_loss = aux_loss_payload * self.objective_scale + sample_or_token_count = ( + statistics.valid_token_count + if self.calculate_per_token_loss + else statistics.valid_token_count.new_tensor(self.num_samples) + ) + objective_weight = sample_or_token_count * self.objective_scale + aux_loss_payload = site_aux_loss * sample_or_token_count + objective_aux_loss = site_aux_loss * objective_weight key = (self.optimizer_step, self.microbatch_id, site_id) self.records[site_id] = MixtureLoRARoutingRecord( key=key, statistics=_detach_routing_statistics(statistics), balance_loss=balance_loss.detach(), aux_loss=objective_aux_loss.detach(), + objective_weight=objective_weight.detach(), ) backward_scale = self.main_loss_backward_scale.to(device=output.device) * self.objective_scale @@ -247,6 +251,109 @@ def _detach_routing_statistics(statistics: RoutingStatistics) -> RoutingStatisti ) +def pack_mixture_lora_routing_records( + contexts: list[MixtureLoRARoutingContext], + site_ids: tuple[str, ...], + *, + num_experts: int, + top_k: int, + device: torch.device, +) -> torch.Tensor: + """Pack local step statistics into a fixed site-major tensor.""" + + if not site_ids or len(set(site_ids)) != len(site_ids): + raise ValueError("site_ids must be non-empty and unique") + if num_experts <= 0 or not 1 <= top_k <= num_experts: + raise ValueError("num_experts and top_k do not describe a valid router") + + scalar_fields = 6 + row_width = 4 * num_experts + scalar_fields + packed = torch.zeros(len(site_ids), row_width, dtype=torch.float64, device=device) + site_indices = {site_id: index for index, site_id in enumerate(site_ids)} + scalar_offset = 4 * num_experts + + for context in contexts: + for site_id, record in context.records.items(): + if site_id not in site_indices: + raise ValueError(f"routing record references unknown site_id {site_id!r}") + statistics = record.statistics + if statistics.pre_topk_prob_sum.numel() != num_experts or statistics.top_k != top_k: + raise ValueError(f"routing record for {site_id!r} does not match the configured router") + + row = packed[site_indices[site_id]] + row[0:num_experts] += statistics.pre_topk_prob_sum.to(device=device, dtype=packed.dtype) + row[num_experts : 2 * num_experts] += statistics.post_topk_weight_sum.to(device=device, dtype=packed.dtype) + row[2 * num_experts : 3 * num_experts] += statistics.selection_count.to(device=device, dtype=packed.dtype) + row[3 * num_experts : 4 * num_experts] += statistics.top1_count.to(device=device, dtype=packed.dtype) + row[scalar_offset] += statistics.pre_topk_entropy_sum.to(device=device, dtype=packed.dtype) + row[scalar_offset + 1] += statistics.post_topk_entropy_sum.to(device=device, dtype=packed.dtype) + row[scalar_offset + 2] += statistics.valid_token_count.to(device=device, dtype=packed.dtype) + row[scalar_offset + 3] += (record.balance_loss * record.objective_weight).to( + device=device, dtype=packed.dtype + ) + row[scalar_offset + 4] += record.objective_weight.to(device=device, dtype=packed.dtype) + row[scalar_offset + 5] += record.aux_loss.to(device=device, dtype=packed.dtype) + + return packed + + +def mixture_lora_metrics_from_packed_records( + packed: torch.Tensor, + site_ids: tuple[str, ...], + *, + num_experts: int, + top_k: int, + calculate_per_token_loss: bool, + data_parallel_world_size_with_cp: int, +) -> dict[str, torch.Tensor]: + """Compute routed-site and global metrics after distributed reduction.""" + + expected_shape = (len(site_ids), 4 * num_experts + 6) + if tuple(packed.shape) != expected_shape: + raise ValueError(f"packed routing statistics must have shape {expected_shape}, got {tuple(packed.shape)}") + if data_parallel_world_size_with_cp <= 0: + raise ValueError("data_parallel_world_size_with_cp must be positive") + + scalar_offset = 4 * num_experts + metrics: dict[str, torch.Tensor] = {} + + def divide_or_zero(numerator: torch.Tensor, denominator: torch.Tensor) -> torch.Tensor: + nonzero = denominator > 0 + safe_denominator = torch.where(nonzero, denominator, torch.ones_like(denominator)) + return numerator / safe_denominator * nonzero.to(numerator.dtype) + + def add_statistics(prefix: str, row: torch.Tensor) -> None: + valid_token_count = row[scalar_offset + 2] + pre_topk_mean_prob = divide_or_zero(row[0:num_experts], valid_token_count) + post_topk_mean_weight = divide_or_zero(row[num_experts : 2 * num_experts], valid_token_count) + selection_share = divide_or_zero(row[2 * num_experts : 3 * num_experts], valid_token_count * top_k) + top1_fraction = divide_or_zero(row[3 * num_experts : 4 * num_experts], valid_token_count) + for expert_id in range(num_experts): + metrics[f"{prefix}/expert_{expert_id}_pre_topk_mean_prob"] = pre_topk_mean_prob[expert_id] + metrics[f"{prefix}/expert_{expert_id}_post_topk_mean_weight"] = post_topk_mean_weight[expert_id] + metrics[f"{prefix}/expert_{expert_id}_selection_share"] = selection_share[expert_id] + metrics[f"{prefix}/expert_{expert_id}_top1_fraction"] = top1_fraction[expert_id] + metrics[f"{prefix}/pre_topk_normalized_entropy"] = divide_or_zero(row[scalar_offset], valid_token_count) + metrics[f"{prefix}/post_topk_normalized_entropy"] = divide_or_zero(row[scalar_offset + 1], valid_token_count) + + for site_id, row in zip(site_ids, packed, strict=True): + prefix = f"molora/{site_id}" + add_statistics(prefix, row) + metrics[f"{prefix}/balance_loss"] = divide_or_zero(row[scalar_offset + 3], row[scalar_offset + 4]) + + global_row = packed.sum(dim=0) + add_statistics("molora/global", global_row) + aux_loss_sum = global_row[scalar_offset + 5] + if calculate_per_token_loss: + # Every configured site routes the same response-token stream. Use one + # site's count so the site sum in the numerator is not divided twice. + aux_denominator = packed[0, scalar_offset + 2] + else: + aux_denominator = packed.new_tensor(data_parallel_world_size_with_cp) + metrics["molora/aux_loss"] = divide_or_zero(aux_loss_sum, aux_denominator) + return metrics + + class MixtureLoRAExperts(nn.Module): """LoRA expert parameters stored in one stable logical layout.""" @@ -500,4 +607,6 @@ def transform( "get_microbatch_objective_scale", "get_mixture_lora_routing_context", "install_mixture_lora_checkpoint_context", + "mixture_lora_metrics_from_packed_records", + "pack_mixture_lora_routing_records", ] diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 15b891dbd..490e22b1b 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -50,6 +50,8 @@ MixtureParallelLinearAdapter, activate_mixture_lora_routing_context, get_microbatch_objective_scale, + mixture_lora_metrics_from_packed_records, + pack_mixture_lora_routing_records, ) from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze @@ -57,29 +59,74 @@ logger = get_logger(__name__) -def _get_global_mixture_lora_site_count(args: Namespace, model: Sequence[DDP]) -> int: - cached_count = getattr(args, "_mixture_lora_global_site_count", None) - if cached_count is not None: - return cached_count +def _get_global_mixture_lora_metadata(args: Namespace, model: Sequence[DDP]) -> tuple[tuple[str, ...], int, int]: + cached_metadata = getattr(args, "_mixture_lora_global_metadata", None) + if cached_metadata is not None: + return cached_metadata - local_site_ids = { - module.mixture_lora.site_id + local_adapters = [ + module for model_chunk in model for module in model_chunk.modules() if isinstance(module, MixtureParallelLinearAdapter) - } - site_count = torch.tensor( - len(local_site_ids), - dtype=torch.int64, - device=next(model[0].parameters()).device, + ] + local_site_ids = sorted({module.mixture_lora.site_id for module in local_adapters}) + if len(local_site_ids) != len(local_adapters): + raise RuntimeError("Mixture-of-LoRA site_id values must be unique within a pipeline stage") + + pipeline_world_size = mpu.get_pipeline_model_parallel_world_size() + if pipeline_world_size > 1: + gathered_site_ids: list[list[str] | None] = [None] * pipeline_world_size + torch.distributed.all_gather_object( + gathered_site_ids, + local_site_ids, + group=mpu.get_pipeline_model_parallel_group(), + ) + global_site_ids = tuple(sorted(site_id for stage_ids in gathered_site_ids for site_id in stage_ids or [])) + else: + global_site_ids = tuple(local_site_ids) + if not global_site_ids: + raise RuntimeError("Mixture-of-LoRA is enabled but the model contains no routed sites") + if len(global_site_ids) != len(set(global_site_ids)): + raise RuntimeError("Mixture-of-LoRA site_id values must be unique across pipeline stages") + + metadata = (global_site_ids, args.lora_num_experts, args.lora_router_top_k) + args._mixture_lora_global_metadata = metadata + return metadata + + +def _reduce_mixture_lora_routing_metrics( + args: Namespace, + contexts: list[MixtureLoRARoutingContext], + metadata: tuple[tuple[str, ...], int, int], + device: torch.device, +) -> dict[str, torch.Tensor]: + """Reduce detached routing records after the pipeline schedule finishes.""" + + site_ids, num_experts, top_k = metadata + packed = pack_mixture_lora_routing_records( + contexts, + site_ids, + num_experts=num_experts, + top_k=top_k, + device=device, ) + data_parallel_world_size_with_cp = mpu.get_data_parallel_world_size(with_context_parallel=True) + if data_parallel_world_size_with_cp > 1: + torch.distributed.all_reduce( + packed, + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) if mpu.get_pipeline_model_parallel_world_size() > 1: - torch.distributed.all_reduce(site_count, group=mpu.get_pipeline_model_parallel_group()) - global_site_count = int(site_count.item()) - if global_site_count <= 0: - raise RuntimeError("Mixture-of-LoRA is enabled but the model contains no routed sites") - args._mixture_lora_global_site_count = global_site_count - return global_site_count + torch.distributed.all_reduce(packed, group=mpu.get_pipeline_model_parallel_group()) + return mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=num_experts, + top_k=top_k, + calculate_per_token_loss=args.calculate_per_token_loss, + data_parallel_world_size_with_cp=data_parallel_world_size_with_cp, + ) def _build_mixture_lora_routing_context( @@ -1055,7 +1102,8 @@ def train_one_step( main_loss_has_tokens = False mixture_lora_enabled = is_mixture_lora_enabled(args) - mixture_lora_num_sites = _get_global_mixture_lora_site_count(args, model) if mixture_lora_enabled else 0 + mixture_lora_metadata = _get_global_mixture_lora_metadata(args, model) if mixture_lora_enabled else None + mixture_lora_num_sites = len(mixture_lora_metadata[0]) if mixture_lora_metadata is not None else 0 # Checkpoint closures retain these contexts until their microbatch backward completes. routing_contexts: list[MixtureLoRARoutingContext] = [] next_routing_microbatch_id = 0 @@ -1267,7 +1315,15 @@ def routing_scope(): if _dcp_orig_cp_group is not None: inner.pg_collection.cp = _dcp_orig_cp_group - # All checkpoint recomputation and aux backward work is complete here. + mixture_lora_metrics = {} + if mixture_lora_metadata is not None: + mixture_lora_metrics = _reduce_mixture_lora_routing_metrics( + args, + routing_contexts, + mixture_lora_metadata, + next(model[0].parameters()).device, + ) + # All checkpoint recomputation, aux backward, and metric reduction work is complete here. routing_contexts.clear() # CI check: verify only MTP parameters have non-zero gradients when truncation happens @@ -1354,6 +1410,7 @@ def routing_scope(): # CP degree under dynamic CP (and is a no-op under static CP, where the # count previously carried the cancelling cp factor). loss_reduced[key] = value / num_samples_or_tokens + loss_reduced.update(mixture_lora_metrics) return loss_reduced, grad_norm return {}, grad_norm diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 360e5df90..31098253e 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -19,6 +19,8 @@ get_microbatch_objective_scale, get_mixture_lora_routing_context, install_mixture_lora_checkpoint_context, + mixture_lora_metrics_from_packed_records, + pack_mixture_lora_routing_records, ) from relax.utils.mixture_lora import MixtureLoraConfig, compute_routing_statistics @@ -312,6 +314,135 @@ def test_dummy_routing_context_records_zero_aux_loss(): torch.testing.assert_close(record.statistics.valid_token_count, torch.tensor(0.0)) torch.testing.assert_close(record.balance_loss, torch.tensor(0.0)) torch.testing.assert_close(record.aux_loss, torch.tensor(0.0)) + torch.testing.assert_close(record.objective_weight, torch.tensor(0.0)) + + +@pytest.mark.parametrize("calculate_per_token_loss", [False, True]) +def test_routing_records_pack_and_report_step_metrics(calculate_per_token_loss): + config = _config() + adapters = [ + MixtureLoRAAdapter( + config, + site_id, + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + for site_id in ("layers.0.linear_qkv", "layers.1.linear_qkv") + ] + context = MixtureLoRARoutingContext( + optimizer_step=3, + microbatch_id=1, + response_mask=torch.tensor([[0, 1, 1], [1, 1, 0]], dtype=torch.bool), + num_microbatches=2, + num_sites=2, + num_samples=2, + calculate_per_token_loss=calculate_per_token_loss, + objective_scale=1.0 if calculate_per_token_loss else 0.25, + main_loss_backward_scale=torch.ones(1), + ) + x = torch.randn(3, 2, 4) + with activate_mixture_lora_routing_context(context): + for adapter in adapters: + adapter(x) + + site_ids = tuple(adapter.site_id for adapter in adapters) + packed = pack_mixture_lora_routing_records( + [context], + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + device=torch.device("cpu"), + ) + metrics = mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + calculate_per_token_loss=calculate_per_token_loss, + data_parallel_world_size_with_cp=2, + ) + + for site_id in site_ids: + record = context.records[site_id] + prefix = f"molora/{site_id}" + for expert_id in range(config.num_experts): + torch.testing.assert_close( + metrics[f"{prefix}/expert_{expert_id}_pre_topk_mean_prob"], + record.statistics.pre_topk_mean_prob[expert_id].double(), + ) + torch.testing.assert_close( + metrics[f"{prefix}/expert_{expert_id}_post_topk_mean_weight"], + record.statistics.post_topk_mean_weight[expert_id].double(), + ) + torch.testing.assert_close(metrics[f"{prefix}/balance_loss"], record.balance_loss.double()) + torch.testing.assert_close( + sum(metrics[f"molora/global/expert_{expert_id}_selection_share"] for expert_id in range(config.num_experts)), + torch.tensor(1.0, dtype=torch.float64), + ) + aux_loss_sum = sum(record.aux_loss for record in context.records.values()).double() + denominator = ( + next(iter(context.records.values())).statistics.valid_token_count.double() + if calculate_per_token_loss + else torch.tensor(2.0, dtype=torch.float64) + ) + torch.testing.assert_close(metrics["molora/aux_loss"], aux_loss_sum / denominator) + + +def test_recompute_replaces_routing_record_instead_of_counting_twice(): + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=torch.ones(1, 2), + num_microbatches=1, + num_sites=1, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + ) + + with activate_mixture_lora_routing_context(context): + adapter(torch.ones(2, 1, 4)) + _, recomputed_decision = adapter.forward_with_routing(torch.full((2, 1, 4), 2.0)) + + assert len(context.records) == 1 + expected = compute_routing_statistics(recomputed_decision, torch.ones(2)) + torch.testing.assert_close( + context.records[adapter.site_id].statistics.pre_topk_prob_sum, expected.pre_topk_prob_sum + ) + + +def test_empty_routing_records_produce_zero_metrics(): + site_ids = ("layers.0.linear_qkv",) + packed = pack_mixture_lora_routing_records( + [], + site_ids, + num_experts=3, + top_k=2, + device=torch.device("cpu"), + ) + metrics = mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=3, + top_k=2, + calculate_per_token_loss=True, + data_parallel_world_size_with_cp=1, + ) + + assert all(metric.item() == 0.0 for metric in metrics.values()) def test_checkpoint_wrapper_restores_captured_routing_context(monkeypatch): diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py new file mode 100644 index 000000000..582828a4e --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -0,0 +1,142 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from datetime import timedelta + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from relax.backends.megatron.mixture_lora import ( + MixtureLoRAAdapter, + MixtureLoRARoutingContext, + activate_mixture_lora_routing_context, + mixture_lora_metrics_from_packed_records, + pack_mixture_lora_routing_records, +) +from relax.utils.mixture_lora import MixtureLoraConfig + + +def _config() -> MixtureLoraConfig: + return MixtureLoraConfig( + num_experts=3, + rank=2, + top_k=2, + temperature=1.0, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_qkv",), + ) + + +def _routing_context(site_id: str, mask: torch.Tensor, *, num_sites: int, objective_scale: float): + config = _config() + adapter = MixtureLoRAAdapter( + config, + site_id, + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + adapter.router.weight.zero_() + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=mask, + num_microbatches=1, + num_sites=num_sites, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=objective_scale, + main_loss_backward_scale=torch.ones(1), + ) + with activate_mixture_lora_routing_context(context): + adapter(torch.ones(mask.shape[1], mask.shape[0], 4)) + return context + + +def _assert_uniform_metrics(metrics: dict[str, torch.Tensor], prefix: str) -> None: + pre_topk_sum = sum(metrics[f"{prefix}/expert_{expert_id}_pre_topk_mean_prob"] for expert_id in range(3)) + post_topk_sum = sum(metrics[f"{prefix}/expert_{expert_id}_post_topk_mean_weight"] for expert_id in range(3)) + selection_sum = sum(metrics[f"{prefix}/expert_{expert_id}_selection_share"] for expert_id in range(3)) + top1_sum = sum(metrics[f"{prefix}/expert_{expert_id}_top1_fraction"] for expert_id in range(3)) + torch.testing.assert_close(pre_topk_sum, torch.tensor(1.0, dtype=torch.float64)) + torch.testing.assert_close(post_topk_sum, torch.tensor(1.0, dtype=torch.float64)) + torch.testing.assert_close(selection_sum, torch.tensor(1.0, dtype=torch.float64)) + torch.testing.assert_close(top1_sum, torch.tensor(1.0, dtype=torch.float64)) + torch.testing.assert_close( + metrics[f"{prefix}/pre_topk_normalized_entropy"], torch.tensor(1.0, dtype=torch.float64) + ) + + +def _distributed_routing_metrics_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + # DP/CP-style reduction: both ranks contribute tokens for the same site. + local_mask = torch.tensor([[1, 1]], dtype=torch.bool) if rank == 0 else torch.tensor([[1, 0]], dtype=torch.bool) + context = _routing_context("layers.0.linear_qkv", local_mask, num_sites=1, objective_scale=0.5) + packed = pack_mixture_lora_routing_records( + [context], + ("layers.0.linear_qkv",), + num_experts=3, + top_k=2, + device=torch.device("cpu"), + ) + dist.all_reduce(packed) + metrics = mixture_lora_metrics_from_packed_records( + packed, + ("layers.0.linear_qkv",), + num_experts=3, + top_k=2, + calculate_per_token_loss=False, + data_parallel_world_size_with_cp=world_size, + ) + _assert_uniform_metrics(metrics, "molora/layers.0.linear_qkv") + torch.testing.assert_close( + metrics["molora/layers.0.linear_qkv/balance_loss"], torch.tensor(1.0, dtype=torch.float64) + ) + torch.testing.assert_close(metrics["molora/aux_loss"], torch.tensor(0.005, dtype=torch.float64)) + + # PP-style reduction: each rank owns a different site and leaves the + # other row at zero before the step-end collective. + site_ids = ("layers.0.linear_qkv", "layers.1.linear_qkv") + local_site_id = site_ids[rank] + context = _routing_context(local_site_id, torch.ones(1, 2, dtype=torch.bool), num_sites=2, objective_scale=1.0) + packed = pack_mixture_lora_routing_records( + [context], + site_ids, + num_experts=3, + top_k=2, + device=torch.device("cpu"), + ) + dist.all_reduce(packed) + metrics = mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=3, + top_k=2, + calculate_per_token_loss=False, + data_parallel_world_size_with_cp=1, + ) + for site_id in site_ids: + _assert_uniform_metrics(metrics, f"molora/{site_id}") + torch.testing.assert_close( + metrics[f"molora/{site_id}/balance_loss"], torch.tensor(1.0, dtype=torch.float64) + ) + _assert_uniform_metrics(metrics, "molora/global") + torch.testing.assert_close(metrics["molora/aux_loss"], torch.tensor(0.01, dtype=torch.float64)) + finally: + dist.destroy_process_group() + + +def test_routing_metrics_reduce_across_two_real_processes(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-gloo-init'}" + mp.spawn(_distributed_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) From 69f751d5fb8615786eec2e3cfe59c21715dafebd Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 08:39:54 +0800 Subject: [PATCH 07/41] feat(lora): support tensor and sequence parallel routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: - 为 Mixture-of-LoRA 增加 qkv 与 proj 的 Tensor Parallel、Sequence Parallel 执行路径。 - 保持 TP 分片后的 forward、expert/router 梯度和 aux loss 与单卡数学参考一致。 Changes: - qkv 的 A 沿 rank 维分片，B 沿输出维分片，router 在 TP rank 间复制并同步梯度。 - proj 的 A/router 沿输入维分片，B 沿输出维分片，显式汇总低秩中间结果和 router logits。 - 增加 sequence gather/scatter 及其反向通信，所有 collective 显式使用构造时保存的 TP group。 - qkv aux loss 按 TP size 缩放后再同步 replicated router 梯度；路由指标仅由 TP rank 0 记录。 - 保留 TP=1 参数形状与原有执行结果，并校验 rank、输入和输出维度的可分片性。 - 增加两进程 qkv/proj、SP 开关、policy gradient、aux gradient、mask 对齐和统计去重测试。 Verification: - 相关 pytest：109 passed, 2 skipped。 - 两进程 TP/SP gloo 对照测试通过。 - 真实 Megatron/Bridge 定向测试：2 passed。 - H20 CUDA 单卡 profiler 测试通过。 - pre-commit run --all-files 全部通过。 Limitations: - 当前 H20 节点仅暴露 1 张 GPU，尚未执行 NCCL TP=2 测试。 --- relax/backends/megatron/mixture_lora.py | 284 ++++++++++++++++-- tests/backends/megatron/test_mixture_lora.py | 10 +- .../megatron/test_mixture_lora_distributed.py | 248 ++++++++++++++- 3 files changed, 517 insertions(+), 25 deletions(-) diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 31360183c..7bed95339 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -14,7 +14,6 @@ from torch import nn from relax.utils.mixture_lora import ( - DenseRoutedLoRAExecutor, MixtureLoraConfig, RoutingDecision, RoutingStatistics, @@ -103,9 +102,15 @@ def attach_aux_loss( site_id: str, config: MixtureLoraConfig, decision: RoutingDecision, + *, + record_statistics: bool = True, + backward_divisor: int = 1, ) -> torch.Tensor: """Attach one site's balance loss and replace its detached record.""" + if backward_divisor <= 0: + raise ValueError("backward_divisor must be positive") + response_mask = self.response_mask_for(x) statistics = compute_routing_statistics(decision, response_mask) # This is a microbatch-level F_e * P_e objective. It is not an @@ -120,16 +125,19 @@ def attach_aux_loss( objective_weight = sample_or_token_count * self.objective_scale aux_loss_payload = site_aux_loss * sample_or_token_count objective_aux_loss = site_aux_loss * objective_weight - key = (self.optimizer_step, self.microbatch_id, site_id) - self.records[site_id] = MixtureLoRARoutingRecord( - key=key, - statistics=_detach_routing_statistics(statistics), - balance_loss=balance_loss.detach(), - aux_loss=objective_aux_loss.detach(), - objective_weight=objective_weight.detach(), - ) + if record_statistics: + key = (self.optimizer_step, self.microbatch_id, site_id) + self.records[site_id] = MixtureLoRARoutingRecord( + key=key, + statistics=_detach_routing_statistics(statistics), + balance_loss=balance_loss.detach(), + aux_loss=objective_aux_loss.detach(), + objective_weight=objective_weight.detach(), + ) - backward_scale = self.main_loss_backward_scale.to(device=output.device) * self.objective_scale + backward_scale = ( + self.main_loss_backward_scale.to(device=output.device) * self.objective_scale / backward_divisor + ) return _AttachAuxLoss.apply(output, aux_loss_payload, backward_scale) @@ -354,6 +362,139 @@ def add_statistics(prefix: str, row: torch.Tensor) -> None: return metrics +class _CopyToTensorParallelRegion(torch.autograd.Function): + """Keep the forward value and sum tensor-parallel input gradients.""" + + @staticmethod + def forward(ctx, value: torch.Tensor, group: Any) -> torch.Tensor: + ctx.group = group + return value + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + grad_input = grad_output.contiguous().clone() + torch.distributed.all_reduce(grad_input, group=ctx.group) + return grad_input, None + + +class _ReduceFromTensorParallelRegion(torch.autograd.Function): + """Sum tensor-parallel partials while leaving backward rank-local.""" + + @staticmethod + def forward(ctx, value: torch.Tensor, group: Any) -> torch.Tensor: + del ctx + output = value.contiguous().clone() + torch.distributed.all_reduce(output, group=group) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + del ctx + return grad_output, None + + +class _GatherLastDimFromTensorParallelRegion(torch.autograd.Function): + """Gather hidden shards and return the local shard in backward.""" + + @staticmethod + def forward(ctx, value: torch.Tensor, group: Any, rank: int, world_size: int) -> torch.Tensor: + ctx.rank = rank + ctx.world_size = world_size + gathered = [torch.empty_like(value) for _ in range(world_size)] + torch.distributed.all_gather(gathered, value.contiguous(), group=group) + return torch.cat(gathered, dim=-1) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + return grad_output.chunk(ctx.world_size, dim=-1)[ctx.rank].contiguous(), None, None, None + + +class _GatherFirstDimFromTensorParallelRegion(torch.autograd.Function): + """Gather sequence shards and reduce-scatter their input gradients.""" + + @staticmethod + def forward(ctx, value: torch.Tensor, group: Any, rank: int, world_size: int) -> torch.Tensor: + ctx.group = group + ctx.rank = rank + ctx.world_size = world_size + gathered = [torch.empty_like(value) for _ in range(world_size)] + torch.distributed.all_gather(gathered, value.contiguous(), group=group) + return torch.cat(gathered, dim=0) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + grad_input = grad_output.contiguous().clone() + torch.distributed.all_reduce(grad_input, group=ctx.group) + return grad_input.chunk(ctx.world_size, dim=0)[ctx.rank].contiguous(), None, None, None + + +class _ScatterFirstDimToTensorParallelRegion(torch.autograd.Function): + """Scatter sequence tokens and gather their output gradients.""" + + @staticmethod + def forward(ctx, value: torch.Tensor, group: Any, rank: int, world_size: int) -> torch.Tensor: + ctx.group = group + ctx.world_size = world_size + if value.shape[0] % world_size != 0: + raise ValueError("sequence dimension must be divisible by tensor parallel size") + return value.chunk(world_size, dim=0)[rank].contiguous() + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + gathered = [torch.empty_like(grad_output) for _ in range(ctx.world_size)] + torch.distributed.all_gather(gathered, grad_output.contiguous(), group=ctx.group) + return torch.cat(gathered, dim=0), None, None, None + + +class MegatronDenseRoutedLoRAExecutor(nn.Module): + """Dense expert execution with explicit Megatron tensor-parallel + collectives.""" + + def __init__(self, *, input_is_parallel: bool, tp_group: Any, tp_rank: int, tp_world_size: int) -> None: + super().__init__() + self.input_is_parallel = input_is_parallel + self.tp_group = tp_group + self.tp_rank = tp_rank + self.tp_world_size = tp_world_size + + def forward( + self, + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + routing_decision: RoutingDecision, + scale: float, + ) -> torch.Tensor: + input_shape = x.shape[:-1] + x_flat = x.reshape(-1, x.shape[-1]) + expert_hidden = torch.einsum("ti,nri->tnr", x_flat, lora_a) + if self.tp_world_size > 1: + if self.input_is_parallel: + expert_hidden = _ReduceFromTensorParallelRegion.apply(expert_hidden, self.tp_group) + else: + expert_hidden = _GatherLastDimFromTensorParallelRegion.apply( + expert_hidden, + self.tp_group, + self.tp_rank, + self.tp_world_size, + ) + # B is output-sharded, so its input gradient must include every + # output shard before it reaches A. + expert_hidden = _CopyToTensorParallelRegion.apply(expert_hidden, self.tp_group) + + expert_outputs = torch.einsum("tnr,nor->tno", expert_hidden, lora_b) + if self.input_is_parallel and self.tp_world_size > 1: + expert_outputs = _GatherLastDimFromTensorParallelRegion.apply( + expert_outputs, + self.tp_group, + self.tp_rank, + self.tp_world_size, + ) + routing_weights = routing_decision.dense_weights().to(dtype=expert_outputs.dtype) + delta = torch.sum(expert_outputs * routing_weights.unsqueeze(-1), dim=1) + return (delta * scale).reshape(*input_shape, expert_outputs.shape[-1]) + + class MixtureLoRAExperts(nn.Module): """LoRA expert parameters stored in one stable logical layout.""" @@ -363,15 +504,21 @@ def __init__( input_size: int, output_size: int, *, + local_rank: int | None = None, + local_input_size: int | None = None, + local_output_size: int | None = None, device: torch.device, dtype: torch.dtype, ) -> None: super().__init__() + local_rank = config.rank if local_rank is None else local_rank + local_input_size = input_size if local_input_size is None else local_input_size + local_output_size = output_size if local_output_size is None else local_output_size self.lora_A = nn.Parameter( - torch.empty(config.num_experts, config.rank, input_size, device=device, dtype=dtype) + torch.empty(config.num_experts, local_rank, local_input_size, device=device, dtype=dtype) ) self.lora_B = nn.Parameter( - torch.empty(config.num_experts, output_size, config.rank, device=device, dtype=dtype) + torch.empty(config.num_experts, local_output_size, config.rank, device=device, dtype=dtype) ) self.reset_parameters() @@ -414,6 +561,11 @@ def __init__( dropout: float, device: torch.device, dtype: torch.dtype, + input_is_parallel: bool = False, + sequence_parallel: bool = False, + tp_group: Any = None, + tp_rank: int = 0, + tp_world_size: int = 1, ) -> None: super().__init__() if not isinstance(site_id, str) or not site_id.strip(): @@ -422,30 +574,102 @@ def __init__( raise ValueError(f"dropout must satisfy 0 <= dropout < 1, got {dropout}") if not dtype.is_floating_point: raise TypeError(f"dtype must be floating point, got {dtype}") + if tp_world_size <= 0 or not 0 <= tp_rank < tp_world_size: + raise ValueError("tp_rank and tp_world_size do not describe a valid tensor-parallel group") + if tp_world_size > 1 and tp_group is None: + raise ValueError("tp_group is required when tensor parallelism is enabled") + if input_size % tp_world_size != 0 or output_size % tp_world_size != 0: + raise ValueError("Mixture-of-LoRA input and output sizes must be divisible by tensor parallel size") + if not input_is_parallel and config.rank % tp_world_size != 0: + raise ValueError("Mixture-of-LoRA rank must be divisible by tensor parallel size for column layers") self.config = config self.site_id = site_id - self.experts = MixtureLoRAExperts(config, input_size, output_size, device=device, dtype=dtype) - self.router = MixtureLoRARouter(config.num_experts, input_size, device=device, dtype=dtype) + self.input_is_parallel = input_is_parallel + self.sequence_parallel = sequence_parallel + self.tp_group = tp_group + self.tp_rank = tp_rank + self.tp_world_size = tp_world_size + local_rank = config.rank if input_is_parallel else config.rank // tp_world_size + local_input_size = input_size // tp_world_size if input_is_parallel else input_size + local_output_size = output_size // tp_world_size + self.experts = MixtureLoRAExperts( + config, + input_size, + output_size, + local_rank=local_rank, + local_input_size=local_input_size, + local_output_size=local_output_size, + device=device, + dtype=dtype, + ) + self.router = MixtureLoRARouter(config.num_experts, local_input_size, device=device, dtype=dtype) self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() - self.executor = DenseRoutedLoRAExecutor() + self.executor = MegatronDenseRoutedLoRAExecutor( + input_is_parallel=input_is_parallel, + tp_group=tp_group, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + if tp_world_size > 1 and not input_is_parallel: + self._synchronize_replicated_router() + self.router.weight.register_hook(self._reduce_replicated_router_gradient) + + def _synchronize_replicated_router(self) -> None: + source_rank = torch.distributed.get_global_rank(self.tp_group, 0) + with torch.no_grad(): + torch.distributed.broadcast(self.router.weight, src=source_rank, group=self.tp_group) + + def _reduce_replicated_router_gradient(self, gradient: torch.Tensor) -> torch.Tensor: + reduced_gradient = gradient.contiguous() + torch.distributed.all_reduce(reduced_gradient, group=self.tp_group) + return reduced_gradient def route(self, x: torch.Tensor) -> RoutingDecision: logits = self.router(x.reshape(-1, x.shape[-1])) + if self.input_is_parallel and self.tp_world_size > 1: + logits = _ReduceFromTensorParallelRegion.apply(logits, self.tp_group) return route_topk(logits, self.config.top_k, self.config.temperature) def forward_with_routing(self, x: torch.Tensor) -> tuple[torch.Tensor, RoutingDecision]: - decision = self.route(x) + routed_input = x + if not self.input_is_parallel and self.tp_world_size > 1: + if self.sequence_parallel: + routed_input = _GatherFirstDimFromTensorParallelRegion.apply( + routed_input, + self.tp_group, + self.tp_rank, + self.tp_world_size, + ) + else: + routed_input = _CopyToTensorParallelRegion.apply(routed_input, self.tp_group) + + decision = self.route(routed_input) delta = self.executor( - self.dropout(x), + self.dropout(routed_input), self.experts.lora_A, self.experts.lora_B, decision, self.config.scale, ) + if self.input_is_parallel and self.sequence_parallel and self.tp_world_size > 1: + delta = _ScatterFirstDimToTensorParallelRegion.apply( + delta, + self.tp_group, + self.tp_rank, + self.tp_world_size, + ) routing_context = get_mixture_lora_routing_context() if routing_context is not None: - delta = routing_context.attach_aux_loss(delta, x, self.site_id, self.config, decision) + delta = routing_context.attach_aux_loss( + delta, + routed_input, + self.site_id, + self.config, + decision, + record_statistics=self.tp_rank == 0, + backward_divisor=self.tp_world_size if not self.input_is_parallel else 1, + ) return delta, decision def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -465,6 +689,11 @@ def __init__( output_size: int, *, dropout: float, + input_is_parallel: bool = False, + sequence_parallel: bool = False, + tp_group: Any = None, + tp_rank: int = 0, + tp_world_size: int = 1, ) -> None: super().__init__() try: @@ -481,6 +710,11 @@ def __init__( dropout=dropout, device=first_parameter.device, dtype=first_parameter.dtype, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + tp_group=tp_group, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ) self._adapter_enabled = True @@ -573,11 +807,12 @@ def transform( return module if self.mixture_config is None: raise RuntimeError("Mixture-of-LoRA PEFT is missing its configuration") - if parallel_state.get_tensor_model_parallel_world_size() != 1: - raise NotImplementedError("Mixture-of-LoRA tensor parallel execution is not implemented yet") - _, full_name = match attributes = get_adapter_attributes_from_linear(module) + tp_world_size = parallel_state.get_tensor_model_parallel_world_size() + tp_group = getattr(module, "tp_group", None) + if tp_world_size > 1 and tp_group is None: + tp_group = parallel_state.get_tensor_model_parallel_group() return MixtureParallelLinearAdapter( module, self.mixture_config, @@ -585,6 +820,12 @@ def transform( attributes.in_features, attributes.out_features, dropout=self.dropout, + input_is_parallel=attributes.input_is_parallel, + sequence_parallel=getattr(getattr(module, "config", None), "sequence_parallel", False) + and not attributes.disable_sequence_parallel_comm, + tp_group=tp_group, + tp_rank=parallel_state.get_tensor_model_parallel_rank(), + tp_world_size=tp_world_size, ) return MixtureLoRAPEFT( @@ -601,6 +842,7 @@ def transform( "MixtureLoRAExperts", "MixtureLoRARouter", "MixtureParallelLinearAdapter", + "MegatronDenseRoutedLoRAExecutor", "activate_mixture_lora_routing_context", "build_mixture_lora_peft", "ensure_mixture_lora_recompute_inputs_grad", diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 31098253e..14f098b3e 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -545,9 +545,15 @@ def match(self, module, name=None, prefix=None): base_module.PEFT = FakePEFT matcher_module.ModuleMatcher = FakeModuleMatcher utils_module.get_adapter_attributes_from_linear = lambda module: SimpleNamespace( - in_features=module.weight.shape[1], out_features=module.weight.shape[0] + in_features=module.weight.shape[1], + out_features=module.weight.shape[0], + input_is_parallel=False, + disable_sequence_parallel_comm=True, + ) + core_module.parallel_state = SimpleNamespace( + get_tensor_model_parallel_world_size=lambda: 1, + get_tensor_model_parallel_rank=lambda: 0, ) - core_module.parallel_state = SimpleNamespace(get_tensor_model_parallel_world_size=lambda: 1) modules = { "megatron": megatron, diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index 582828a4e..7c0861986 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -5,6 +5,7 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp +import torch.nn.functional as F from relax.backends.megatron.mixture_lora import ( MixtureLoRAAdapter, @@ -13,7 +14,7 @@ mixture_lora_metrics_from_packed_records, pack_mixture_lora_routing_records, ) -from relax.utils.mixture_lora import MixtureLoraConfig +from relax.utils.mixture_lora import DenseRoutedLoRAExecutor, MixtureLoraConfig, route_topk def _config() -> MixtureLoraConfig: @@ -81,7 +82,9 @@ def _distributed_routing_metrics_worker(rank: int, world_size: int, init_method: ) try: # DP/CP-style reduction: both ranks contribute tokens for the same site. - local_mask = torch.tensor([[1, 1]], dtype=torch.bool) if rank == 0 else torch.tensor([[1, 0]], dtype=torch.bool) + local_mask = ( + torch.tensor([[1, 1]], dtype=torch.bool) if rank == 0 else torch.tensor([[1, 0]], dtype=torch.bool) + ) context = _routing_context("layers.0.linear_qkv", local_mask, num_sites=1, objective_scale=0.5) packed = pack_mixture_lora_routing_records( [context], @@ -140,3 +143,244 @@ def _distributed_routing_metrics_worker(rank: int, world_size: int, init_method: def test_routing_metrics_reduce_across_two_real_processes(tmp_path): init_method = f"file://{tmp_path / 'mixture-lora-gloo-init'}" mp.spawn(_distributed_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) + + +def _reference_forward_and_backward( + config: MixtureLoraConfig, + x: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + router_weight: torch.Tensor, +): + reference_x = x.detach().clone().requires_grad_(True) + reference_a = lora_a.detach().clone().requires_grad_(True) + reference_b = lora_b.detach().clone().requires_grad_(True) + reference_router = router_weight.detach().clone().requires_grad_(True) + logits = F.linear(reference_x.reshape(-1, reference_x.shape[-1]).float(), reference_router.float()) + decision = route_topk(logits, config.top_k, config.temperature) + output = DenseRoutedLoRAExecutor()(reference_x, reference_a, reference_b, decision, config.scale) + output.square().sum().backward() + return reference_x, reference_a, reference_b, reference_router, output.detach() + + +def _assert_tp_adapter_matches_reference(rank: int, *, input_is_parallel: bool, sequence_parallel: bool) -> None: + config = MixtureLoraConfig( + num_experts=4, + rank=4, + top_k=2, + temperature=0.8, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_proj" if input_is_parallel else "linear_qkv",), + ) + torch.manual_seed(1234 + int(input_is_parallel) * 10 + int(sequence_parallel)) + device = torch.device("cpu") + sequence_length, batch_size, input_size, output_size = 4, 2, 6, 8 + full_x = torch.randn(sequence_length, batch_size, input_size, device=device) + full_a = torch.randn(config.num_experts, config.rank, input_size, device=device) + full_b = torch.randn(config.num_experts, output_size, config.rank, device=device) + full_router = torch.randn(config.num_experts, input_size, device=device) + reference_x, reference_a, reference_b, reference_router, reference_output = _reference_forward_and_backward( + config, + full_x, + full_a, + full_b, + full_router, + ) + + adapter = MixtureLoRAAdapter( + config, + "linear_proj" if input_is_parallel else "linear_qkv", + input_size, + output_size, + dropout=0.0, + device=device, + dtype=torch.float32, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + tp_group=dist.group.WORLD, + tp_rank=rank, + tp_world_size=2, + ) + input_slice = slice(rank * input_size // 2, (rank + 1) * input_size // 2) + output_slice = slice(rank * output_size // 2, (rank + 1) * output_size // 2) + rank_slice = slice(rank * config.rank // 2, (rank + 1) * config.rank // 2) + with torch.no_grad(): + if input_is_parallel: + adapter.experts.lora_A.copy_(full_a[:, :, input_slice]) + adapter.router.weight.copy_(full_router[:, input_slice]) + else: + adapter.experts.lora_A.copy_(full_a[:, rank_slice, :]) + adapter.router.weight.copy_(full_router) + adapter.experts.lora_B.copy_(full_b[:, output_slice, :]) + + if input_is_parallel: + local_x = full_x[:, :, input_slice].detach().clone().requires_grad_(True) + elif sequence_parallel: + sequence_slice = slice(rank * sequence_length // 2, (rank + 1) * sequence_length // 2) + local_x = full_x[sequence_slice].detach().clone().requires_grad_(True) + else: + local_x = full_x.detach().clone().requires_grad_(True) + + output = adapter(local_x) + if input_is_parallel: + if sequence_parallel: + sequence_slice = slice(rank * sequence_length // 2, (rank + 1) * sequence_length // 2) + expected_output = reference_output[sequence_slice] + else: + expected_output = reference_output + else: + expected_output = reference_output[:, :, output_slice] + torch.testing.assert_close(output, expected_output, atol=2e-5, rtol=2e-5) + output.square().sum().backward() + + if input_is_parallel: + expected_input_grad = reference_x.grad[:, :, input_slice] + expected_a_grad = reference_a.grad[:, :, input_slice] + expected_router_grad = reference_router.grad[:, input_slice] + else: + expected_input_grad = reference_x.grad[sequence_slice] if sequence_parallel else reference_x.grad + expected_a_grad = reference_a.grad[:, rank_slice, :] + expected_router_grad = reference_router.grad + torch.testing.assert_close(local_x.grad, expected_input_grad, atol=3e-5, rtol=3e-5) + torch.testing.assert_close(adapter.experts.lora_A.grad, expected_a_grad, atol=3e-5, rtol=3e-5) + torch.testing.assert_close( + adapter.experts.lora_B.grad, + reference_b.grad[:, output_slice, :], + atol=3e-5, + rtol=3e-5, + ) + torch.testing.assert_close(adapter.router.weight.grad, expected_router_grad, atol=3e-5, rtol=3e-5) + dist.barrier() + + +def _assert_tp_aux_loss_matches_reference(rank: int, *, input_is_parallel: bool, sequence_parallel: bool) -> None: + config = MixtureLoraConfig( + num_experts=4, + rank=4, + top_k=2, + temperature=0.8, + aux_loss_coef=0.03, + alpha=4.0, + target_modules=("linear_proj" if input_is_parallel else "linear_qkv",), + ) + torch.manual_seed(4321 + int(input_is_parallel)) + input_size, output_size = 6, 8 + full_x = torch.randn(4, 2, input_size) + full_a = torch.randn(config.num_experts, config.rank, input_size) + full_b = torch.randn(config.num_experts, output_size, config.rank) + full_router = torch.randn(config.num_experts, input_size) + response_mask = torch.tensor([[0, 1, 1, 1], [1, 1, 0, 0]], dtype=torch.bool) + + reference = MixtureLoRAAdapter( + config, + "linear_proj" if input_is_parallel else "linear_qkv", + input_size, + output_size, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + input_is_parallel=input_is_parallel, + ) + with torch.no_grad(): + reference.experts.lora_A.copy_(full_a) + reference.experts.lora_B.copy_(full_b) + reference.router.weight.copy_(full_router) + reference_context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=response_mask, + num_microbatches=1, + num_sites=1, + num_samples=2, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + ) + with activate_mixture_lora_routing_context(reference_context): + reference_output = reference(full_x) + (reference_output.sum() * 0.0).backward() + + adapter = MixtureLoRAAdapter( + config, + "linear_proj" if input_is_parallel else "linear_qkv", + input_size, + output_size, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + tp_group=dist.group.WORLD, + tp_rank=rank, + tp_world_size=2, + ) + input_slice = slice(rank * input_size // 2, (rank + 1) * input_size // 2) + output_slice = slice(rank * output_size // 2, (rank + 1) * output_size // 2) + rank_slice = slice(rank * config.rank // 2, (rank + 1) * config.rank // 2) + with torch.no_grad(): + if input_is_parallel: + adapter.experts.lora_A.copy_(full_a[:, :, input_slice]) + adapter.router.weight.copy_(full_router[:, input_slice]) + else: + adapter.experts.lora_A.copy_(full_a[:, rank_slice, :]) + adapter.router.weight.copy_(full_router) + adapter.experts.lora_B.copy_(full_b[:, output_slice, :]) + if input_is_parallel: + local_x = full_x[:, :, input_slice] + elif sequence_parallel: + sequence_slice = slice(rank * full_x.shape[0] // 2, (rank + 1) * full_x.shape[0] // 2) + local_x = full_x[sequence_slice] + else: + local_x = full_x + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=response_mask, + num_microbatches=1, + num_sites=1, + num_samples=2, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + ) + with activate_mixture_lora_routing_context(context): + output = adapter(local_x) + (output.sum() * 0.0).backward() + + expected_router_grad = ( + reference.router.weight.grad[:, input_slice] if input_is_parallel else reference.router.weight.grad + ) + torch.testing.assert_close(adapter.router.weight.grad, expected_router_grad, atol=3e-5, rtol=3e-5) + assert (adapter.site_id in context.records) == (rank == 0) + dist.barrier() + + +def _tensor_parallel_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + for input_is_parallel in (False, True): + for sequence_parallel in (False, True): + _assert_tp_adapter_matches_reference( + rank, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + ) + _assert_tp_aux_loss_matches_reference( + rank, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + ) + finally: + dist.destroy_process_group() + + +def test_tensor_parallel_forward_and_backward_match_single_rank_reference(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-tp-gloo-init'}" + mp.spawn(_tensor_parallel_worker, args=(2, init_method), nprocs=2, join=True) From 007b981ca80b6c76929dfc7021d2b41735f6de81 Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 08:56:56 +0800 Subject: [PATCH 08/41] feat(lora): checkpoint routed adapters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: - 将 Mixture-of-LoRA expert、router 和配置接入 Megatron 原生 distributed checkpoint。 - 恢复时校验 Mixture 配置，并保持现有单 LoRA 的 HF PEFT 导出行为。 Changes: - qkv/proj 分别声明 A、B 和 router 的全局 shape 与 TP 分片轴。 - 以模块 extra state 保存 schema version、site_id、N/R/K/T/C、alpha、target、输入输出维度和 dtype。 - 加载配置不一致时列出具体字段并停止恢复。 - Mixture 模式跳过仅适用于标准单 LoRA 的 _save_lora_to_checkpoint 导出；N=1 路径不变。 - checkpoint 不记录 dense executor 类型，后续执行器可以共享同一参数结构。 - 增加配置冲突、真实 torch_dist 写盘/恢复、参数、路由、输出和 loss 一致性测试。 Verification: - 相关 pytest：110 passed, 2 skipped。 - 真实 Megatron distributed checkpoint 保存/恢复测试通过。 - 真实 Megatron/Bridge 定向测试通过。 - pre-commit run --all-files 全部通过。 --- relax/backends/megatron/mixture_lora.py | 75 +++++++++++++++++ relax/backends/megatron/model.py | 4 +- tests/backends/megatron/test_mixture_lora.py | 85 ++++++++++++++++++++ 3 files changed, 162 insertions(+), 2 deletions(-) diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 7bed95339..2ce749607 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -585,6 +585,8 @@ def __init__( self.config = config self.site_id = site_id + self.input_size = input_size + self.output_size = output_size self.input_is_parallel = input_is_parallel self.sequence_parallel = sequence_parallel self.tp_group = tp_group @@ -625,6 +627,42 @@ def _reduce_replicated_router_gradient(self, gradient: torch.Tensor) -> torch.Te torch.distributed.all_reduce(reduced_gradient, group=self.tp_group) return reduced_gradient + def get_extra_state(self) -> dict[str, Any]: + """Return the configuration required to validate checkpoint restore.""" + + return { + "schema_version": self.config.schema_version, + "site_id": self.site_id, + "num_experts": self.config.num_experts, + "rank": self.config.rank, + "top_k": self.config.top_k, + "temperature": self.config.temperature, + "aux_loss_coef": self.config.aux_loss_coef, + "alpha": self.config.alpha, + "target_modules": list(self.config.target_modules), + "input_size": self.input_size, + "output_size": self.output_size, + "dtype": str(self.router.weight.dtype), + } + + def set_extra_state(self, state: dict[str, Any]) -> None: + """Validate saved Mixture-of-LoRA metadata before loading tensors.""" + + if not isinstance(state, dict): + raise RuntimeError(f"Mixture-of-LoRA checkpoint metadata must be a dict, got {type(state).__name__}") + expected = self.get_extra_state() + mismatches = { + key: (expected_value, state.get(key)) + for key, expected_value in expected.items() + if state.get(key) != expected_value + } + if mismatches: + details = ", ".join( + f"{key}: expected {expected_value!r}, checkpoint has {saved_value!r}" + for key, (expected_value, saved_value) in mismatches.items() + ) + raise RuntimeError(f"Mixture-of-LoRA checkpoint metadata mismatch for {self.site_id}: {details}") + def route(self, x: torch.Tensor) -> RoutingDecision: logits = self.router(x.reshape(-1, x.shape[-1])) if self.input_is_parallel and self.tp_world_size > 1: @@ -772,6 +810,43 @@ def state_dict( ) return destination + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int], ...] = (), + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Combine the base layer and routed TP shards for native + checkpointing.""" + + from megatron.core import parallel_state + from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + + sharded_state = self.to_wrap.sharded_state_dict(prefix, sharded_offsets, metadata) + adapter_state = self.mixture_lora.state_dict(prefix="", keep_vars=True) + axis_map = { + "experts.lora_A": 2 if self.mixture_lora.input_is_parallel else 1, + "experts.lora_B": 1, + } + if self.mixture_lora.input_is_parallel: + axis_map["router.weight"] = 1 + dp_cp_group = ( + metadata["dp_cp_group"] + if metadata is not None and metadata.get("dp_cp_group") is not None + else parallel_state.get_data_parallel_group(with_context_parallel=True) + ) + sharded_state.update( + make_sharded_tensors_for_checkpoint( + adapter_state, + f"{prefix}mixture_lora.", + axis_map, + sharded_offsets, + tp_group=self.mixture_lora.tp_group, + dp_cp_group=dp_cp_group, + ) + ) + return sharded_state + def build_mixture_lora_peft(config: MixtureLoraConfig, dropout: float): """Build a Bridge PEFT object that injects routed adapters at matched diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 490e22b1b..5c6a1c4fc 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1682,7 +1682,7 @@ def save( train_data_iterator=None, preprocess_common_state_dict_fn=None, ) - if is_lora_enabled(args): + if is_lora_enabled(args) and not is_mixture_lora_enabled(args): checkpoint_dir = Path(args.save) / f"iter_{iteration:07d}" _save_lora_to_checkpoint(model, str(checkpoint_dir), args) if should_disable_forward_pre_hook(args): @@ -1748,7 +1748,7 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: if allow_missing_mtp_keys and is_export_writer: reconcile_hf_export_index(str(path), reference_hf_dir=args.hf_checkpoint, supplement_mtp=True) - if is_lora_enabled(args): + if is_lora_enabled(args) and not is_mixture_lora_enabled(args): _save_lora_to_checkpoint(model, str(path), args, bridge=bridge) if should_log: logger.info(f"Successfully saved HuggingFace model to {path}") diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 14f098b3e..e98b78dfc 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. +import copy import sys import types from dataclasses import dataclass, field @@ -176,12 +177,47 @@ def test_mixture_lora_wrapper_keeps_base_state_keys_stable(): assert set(state) == { "weight", "bias", + "mixture_lora._extra_state", "mixture_lora.experts.lora_A", "mixture_lora.experts.lora_B", "mixture_lora.router.weight", } +def test_mixture_lora_state_restores_output_and_validates_metadata(): + source = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + target = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + source.experts.lora_B.normal_(mean=0.0, std=0.2) + x = torch.randn(3, 2, 4) + expected = source(x) + + state = copy.deepcopy(source.state_dict()) + target.load_state_dict(state) + + torch.testing.assert_close(target(x), expected) + mismatched_state = copy.deepcopy(state) + mismatched_state["_extra_state"]["top_k"] = 1 + with pytest.raises(RuntimeError, match=r"top_k: expected 2, checkpoint has 1"): + target.load_state_dict(mismatched_state) + + @pytest.mark.parametrize( ("calculate_per_token_loss", "is_dummy", "explicit_loss_scale", "expected"), [ @@ -648,6 +684,55 @@ def test_mixture_lora_peft_wraps_real_column_parallel_linear_when_available(tmp_ assert all(parameter.requires_grad for parameter in transformed.linear_qkv.mixture_lora.parameters()) assert all(parameter.grad is None for parameter in transformed.linear_qkv.to_wrap.parameters()) assert torch.count_nonzero(transformed.linear_qkv.mixture_lora.router.weight.grad) > 0 + sharded_state = transformed.linear_qkv.sharded_state_dict( + metadata={"dp_cp_group": parallel_state.get_data_parallel_group(with_context_parallel=True)} + ) + assert { + "mixture_lora._extra_state", + "mixture_lora.experts.lora_A", + "mixture_lora.experts.lora_B", + "mixture_lora.router.weight", + }.issubset(sharded_state) + assert sharded_state["mixture_lora.experts.lora_A"].global_shape == (3, 2, 4) + assert sharded_state["mixture_lora.experts.lora_B"].global_shape == (3, 5, 2) + assert sharded_state["mixture_lora.router.weight"].global_shape == (3, 4) + + from megatron.core import dist_checkpointing + + checkpoint_dir = tmp_path / "mixture_lora_dist_checkpoint" + checkpoint_dir.mkdir() + fixed_input = torch.randn(3, 1, 4) + expected_output = transformed.linear_qkv(fixed_input)[0].detach() + expected_loss = expected_output.square().mean() + expected_routing = transformed.linear_qkv.mixture_lora.route(fixed_input) + expected_parameters = { + name: parameter.detach().clone() + for name, parameter in transformed.linear_qkv.mixture_lora.named_parameters() + } + dist_checkpointing.save(sharded_state, str(checkpoint_dir)) + with torch.no_grad(): + transformed.linear_qkv.mixture_lora.experts.lora_A.zero_() + transformed.linear_qkv.mixture_lora.experts.lora_B.zero_() + transformed.linear_qkv.mixture_lora.router.weight.zero_() + load_template = transformed.linear_qkv.sharded_state_dict( + metadata={"dp_cp_group": parallel_state.get_data_parallel_group(with_context_parallel=True)} + ) + loaded_state = dist_checkpointing.load(load_template, str(checkpoint_dir)) + mixture_state = { + key.removeprefix("mixture_lora."): value + for key, value in loaded_state.items() + if key.startswith("mixture_lora.") + } + transformed.linear_qkv.mixture_lora.load_state_dict(mixture_state) + restored_output = transformed.linear_qkv(fixed_input)[0] + restored_routing = transformed.linear_qkv.mixture_lora.route(fixed_input) + torch.testing.assert_close(restored_output, expected_output) + torch.testing.assert_close(restored_output.square().mean(), expected_loss) + torch.testing.assert_close(restored_routing.pre_topk_probs, expected_routing.pre_topk_probs) + torch.testing.assert_close(restored_routing.post_topk_weights, expected_routing.post_topk_weights) + assert torch.equal(restored_routing.topk_indices, expected_routing.topk_indices) + for name, parameter in transformed.linear_qkv.mixture_lora.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) finally: parallel_state.destroy_model_parallel() torch.distributed.destroy_process_group() From 07d7905b627c49b1900270ff406f936d3c3488e2 Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 12:17:59 +0800 Subject: [PATCH 09/41] feat(lora): add SGLang rollout model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 Qwen3 Mixture-of-LoRA SGLang 外部模型，在 qkv_proj 和 o_proj 上复用训练侧 Top-K 路由与 dense executor。 通过启动前 JSON 环境变量传递路由配置，并校验外部模型包冲突、参数名称、shape 和 dtype。补充单 LoRA 兼容、训练与 rollout 数值一致、BF16 router 及权重加载测试。 --- relax/backends/sglang/sglang_engine.py | 12 +- relax/models/qwen3_mixture_lora/__init__.py | 3 + .../qwen3_mixture_lora/sglang/__init__.py | 3 + .../models/qwen3_mixture_lora/sglang/model.py | 211 ++++++++++++++++++ relax/utils/env.py | 1 + relax/utils/mixture_lora.py | 108 +++++++++ tests/backends/sglang/test_mixture_lora.py | 55 +++++ .../qwen3_mixture_lora/test_sglang_model.py | 159 +++++++++++++ tests/utils/test_mixture_lora_routing.py | 60 +++++ 9 files changed, 610 insertions(+), 2 deletions(-) create mode 100644 relax/models/qwen3_mixture_lora/__init__.py create mode 100644 relax/models/qwen3_mixture_lora/sglang/__init__.py create mode 100644 relax/models/qwen3_mixture_lora/sglang/model.py create mode 100644 tests/backends/sglang/test_mixture_lora.py create mode 100644 tests/models/qwen3_mixture_lora/test_sglang_model.py diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index a927e667b..7e80f49bd 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -24,7 +24,12 @@ from relax.utils.env import Envs from relax.utils.http_utils import get_host_info, router_worker_base_url from relax.utils.logging_utils import get_logger -from relax.utils.megatron_peft_utils import convert_megatron_to_hf_target_modules, is_lora_enabled +from relax.utils.megatron_peft_utils import ( + build_mixture_lora_config, + convert_megatron_to_hf_target_modules, + is_lora_enabled, +) +from relax.utils.mixture_lora import configure_mixture_lora_external_model logger = get_logger(__name__) @@ -466,7 +471,10 @@ def _init_normal(self, server_args_dict): # Must be set before launch_server_process() spawns child process # (multiprocessing start_method='spawn'), because the child inherits # the parent's os.environ at spawn time. - external_pkg = getattr(self.args, "sglang_external_model_package", None) + external_pkg = configure_mixture_lora_external_model( + build_mixture_lora_config(self.args), + getattr(self.args, "sglang_external_model_package", None), + ) if external_pkg: os.environ["SGLANG_EXTERNAL_MODEL_PACKAGE"] = external_pkg os.environ["SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE"] = external_pkg diff --git a/relax/models/qwen3_mixture_lora/__init__.py b/relax/models/qwen3_mixture_lora/__init__.py new file mode 100644 index 000000000..a5335cdd9 --- /dev/null +++ b/relax/models/qwen3_mixture_lora/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Qwen3 Mixture-of-LoRA model integrations.""" diff --git a/relax/models/qwen3_mixture_lora/sglang/__init__.py b/relax/models/qwen3_mixture_lora/sglang/__init__.py new file mode 100644 index 000000000..0014fed8f --- /dev/null +++ b/relax/models/qwen3_mixture_lora/sglang/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Qwen3 Mixture-of-LoRA SGLang external model package.""" diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py new file mode 100644 index 000000000..fb893176c --- /dev/null +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -0,0 +1,211 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""SGLang Qwen3 external model with token-routed LoRA experts.""" + +from types import MethodType +from typing import Any, Iterable + +import torch +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.models.qwen3 import Qwen3ForCausalLM +from torch import nn +from torch.nn import functional as F + +from relax.utils.env import Envs +from relax.utils.mixture_lora import ( + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + deserialize_mixture_lora_config, + megatron_mixture_lora_name_to_sglang, + route_topk, +) + + +class SGLangMixtureLoRA(nn.Module): + """Parameter-compatible dense Mixture-of-LoRA execution for rollout.""" + + def __init__( + self, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.config = config + self.site_id = site_id + self.experts = nn.Module() + self.experts.register_parameter( + "lora_A", + nn.Parameter(torch.empty(config.num_experts, config.rank, input_size, device=device, dtype=dtype)), + ) + self.experts.register_parameter( + "lora_B", + nn.Parameter(torch.empty(config.num_experts, output_size, config.rank, device=device, dtype=dtype)), + ) + self.router = nn.Linear(input_size, config.num_experts, bias=False, device=device, dtype=dtype) + self.executor = DenseRoutedLoRAExecutor() + self.reset_parameters() + + def reset_parameters(self) -> None: + for expert_weight in self.experts.lora_A: + nn.init.xavier_uniform_(expert_weight) + nn.init.zeros_(self.experts.lora_B) + nn.init.normal_(self.router.weight, mean=0.0, std=0.02) + + def route(self, x: torch.Tensor): + logits = F.linear(x.reshape(-1, x.shape[-1]).float(), self.router.weight.float()) + return route_topk(logits, self.config.top_k, self.config.temperature) + + def forward_with_routing(self, x: torch.Tensor): + decision = self.route(x) + output = self.executor( + x, + self.experts.lora_A, + self.experts.lora_B, + decision, + self.config.scale, + ) + return output, decision + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output, _ = self.forward_with_routing(x) + return output + + +def _routed_linear_forward(linear, x: torch.Tensor, *args: Any, **kwargs: Any): + base_result = linear._relax_mixture_lora_base_forward(x, *args, **kwargs) + if not isinstance(base_result, tuple) or len(base_result) != 2: + raise TypeError(f"{type(linear).__name__} must return an (output, bias) tuple") + output, bias = base_result + delta = linear.mixture_lora(x).reshape(output.shape) + return output + delta, bias + + +def attach_sglang_mixture_lora( + linear: nn.Module, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, +) -> None: + """Add routed parameters while preserving SGLang base parameter names.""" + + if hasattr(linear, "mixture_lora"): + raise RuntimeError(f"SGLang linear {site_id} already has a Mixture-of-LoRA adapter") + try: + base_parameter = next(linear.parameters()) + except StopIteration as error: + raise ValueError(f"SGLang linear {site_id} has no parameters") from error + linear.add_module( + "mixture_lora", + SGLangMixtureLoRA( + config, + site_id, + input_size, + output_size, + device=base_parameter.device, + dtype=base_parameter.dtype, + ), + ) + linear._relax_mixture_lora_base_forward = linear.forward + linear.forward = MethodType(_routed_linear_forward, linear) + + +def load_sglang_mixture_lora_weights( + model: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Validate and copy one chunk of routed weights into an SGLang model.""" + + parameters = dict(model.named_parameters()) + loaded_names = set() + for source_name, loaded_weight in weights: + target_name = ( + source_name + if source_name.startswith("model.layers.") + else megatron_mixture_lora_name_to_sglang(source_name) + ) + if target_name in loaded_names: + raise ValueError(f"Duplicate Mixture-of-LoRA weight: {target_name}") + if target_name not in parameters: + raise ValueError(f"Unknown Mixture-of-LoRA weight for SGLang: {target_name}") + parameter = parameters[target_name] + if tuple(loaded_weight.shape) != tuple(parameter.shape): + raise ValueError( + f"Mixture-of-LoRA weight {target_name} has shape {tuple(loaded_weight.shape)}, " + f"expected {tuple(parameter.shape)}" + ) + if loaded_weight.dtype != parameter.dtype: + raise TypeError( + f"Mixture-of-LoRA weight {target_name} has dtype {loaded_weight.dtype}, expected {parameter.dtype}" + ) + with torch.no_grad(): + parameter.copy_(loaded_weight.to(device=parameter.device)) + loaded_names.add(target_name) + return loaded_names + + +class Qwen3MixtureLoRAForCausalLM(Qwen3ForCausalLM): + """Qwen3 external model that routes LoRA experts at attention + projections.""" + + def __init__(self, config, quant_config=None, prefix: str = "") -> None: + runtime_config = Envs.RELAX_MIXTURE_LORA_CONFIG + if runtime_config is None: + raise RuntimeError("RELAX_MIXTURE_LORA_CONFIG must be set before constructing the external model") + self.mixture_lora_config = deserialize_mixture_lora_config(runtime_config) + if get_tensor_model_parallel_world_size() != 1: + raise ValueError("Qwen3 Mixture-of-LoRA rollout currently requires SGLang TP=1") + super().__init__(config, quant_config=quant_config, prefix=prefix) + self._install_mixture_lora() + + def _install_mixture_lora(self) -> None: + targets = set(self.mixture_lora_config.target_modules) + supported_targets = {"linear_qkv", "linear_proj"} + if not targets.issubset(supported_targets): + raise ValueError(f"Unsupported SGLang Mixture-of-LoRA targets: {sorted(targets - supported_targets)}") + start_layer = getattr(self.model, "start_layer", 0) + for local_layer_id, layer in enumerate(self.model.layers): + layer_id = start_layer + local_layer_id + attention = layer.self_attn + if "linear_qkv" in targets: + site_id = f"decoder.layers.{layer_id}.self_attention.linear_qkv" + attach_sglang_mixture_lora( + attention.qkv_proj, + self.mixture_lora_config, + site_id, + attention.qkv_proj.input_size, + attention.qkv_proj.output_size, + ) + if "linear_proj" in targets: + site_id = f"decoder.layers.{layer_id}.self_attention.linear_proj" + attach_sglang_mixture_lora( + attention.o_proj, + self.mixture_lora_config, + site_id, + attention.o_proj.input_size, + attention.o_proj.output_size, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + mixture_weights = [] + + def base_weight_iterator(): + for name, weight in weights: + if ".mixture_lora." in name: + mixture_weights.append((name, weight)) + else: + yield name, weight + + super().load_weights(base_weight_iterator()) + if not mixture_weights: + return + + load_sglang_mixture_lora_weights(self, mixture_weights) + + +EntryClass = Qwen3MixtureLoRAForCausalLM diff --git a/relax/utils/env.py b/relax/utils/env.py index d875e8fe5..49349a3cb 100644 --- a/relax/utils/env.py +++ b/relax/utils/env.py @@ -203,6 +203,7 @@ class Envs(metaclass=_EnvsMeta): # ------------- LoRA ------------- RELAX_LORA_LIVE_DIR = EnvProperty("RELAX_LORA_LIVE_DIR", str, None) + RELAX_MIXTURE_LORA_CONFIG = EnvProperty("RELAX_MIXTURE_LORA_CONFIG", str, None) # ------------- Extra Modules ------------- RELAX_EXTRA_MODULES = EnvProperty("RELAX_EXTRA_MODULES", str, "") diff --git a/relax/utils/mixture_lora.py b/relax/utils/mixture_lora.py index 9c54e232d..de2f9c336 100644 --- a/relax/utils/mixture_lora.py +++ b/relax/utils/mixture_lora.py @@ -2,7 +2,9 @@ """Backend-independent routing primitives for Mixture-of-LoRA.""" +import json import math +import os from dataclasses import dataclass from typing import Literal, Protocol, Sequence, runtime_checkable @@ -13,6 +15,16 @@ MIXTURE_LORA_SCHEMA_VERSION = 1 MixtureLoraParameterKind = Literal["experts.lora_A", "experts.lora_B", "router.weight"] _PARAMETER_KINDS = {"experts.lora_A", "experts.lora_B", "router.weight"} +_CONFIG_JSON_FIELDS = { + "schema_version", + "num_experts", + "rank", + "top_k", + "temperature", + "aux_loss_coef", + "alpha", + "target_modules", +} @dataclass(frozen=True) @@ -55,6 +67,102 @@ def scale(self) -> float: return self.alpha / self.rank +def serialize_mixture_lora_config(config: MixtureLoraConfig) -> str: + """Serialize validated runtime configuration for spawned rollout + workers.""" + + if not isinstance(config, MixtureLoraConfig): + raise TypeError(f"config must be MixtureLoraConfig, got {type(config).__name__}") + payload = { + "schema_version": config.schema_version, + "num_experts": config.num_experts, + "rank": config.rank, + "top_k": config.top_k, + "temperature": config.temperature, + "aux_loss_coef": config.aux_loss_coef, + "alpha": config.alpha, + "target_modules": list(config.target_modules), + } + return json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def deserialize_mixture_lora_config(raw: str) -> MixtureLoraConfig: + """Parse and validate runtime configuration inherited by rollout + workers.""" + + if not isinstance(raw, str) or not raw.strip(): + raise ValueError("Mixture-of-LoRA runtime configuration must be a non-empty JSON string") + try: + payload = json.loads(raw) + except json.JSONDecodeError as error: + raise ValueError(f"Invalid Mixture-of-LoRA runtime configuration JSON: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError("Mixture-of-LoRA runtime configuration JSON must contain an object") + fields = set(payload) + if fields != _CONFIG_JSON_FIELDS: + missing = sorted(_CONFIG_JSON_FIELDS - fields) + unexpected = sorted(fields - _CONFIG_JSON_FIELDS) + raise ValueError( + f"Mixture-of-LoRA runtime configuration fields do not match schema: " + f"missing={missing}, unexpected={unexpected}" + ) + return MixtureLoraConfig( + schema_version=payload["schema_version"], + num_experts=payload["num_experts"], + rank=payload["rank"], + top_k=payload["top_k"], + temperature=payload["temperature"], + aux_loss_coef=payload["aux_loss_coef"], + alpha=payload["alpha"], + target_modules=tuple(payload["target_modules"]), + ) + + +def configure_mixture_lora_external_model( + config: MixtureLoraConfig | None, + external_package: str | None, +) -> str | None: + """Set the validated rollout configuration before spawning SGLang.""" + + if config is None: + return external_package + mixture_external_package = "relax.models.qwen3_mixture_lora.sglang" + if external_package not in (None, mixture_external_package): + raise ValueError( + "Mixture-of-LoRA requires the Qwen3 external model package, " + f"but sglang_external_model_package is {external_package!r}" + ) + os.environ["RELAX_MIXTURE_LORA_CONFIG"] = serialize_mixture_lora_config(config) + return mixture_external_package + + +def megatron_mixture_lora_name_to_sglang(parameter_name: str) -> str: + """Map one stable Megatron Mixture parameter name to Qwen3 SGLang.""" + + if not isinstance(parameter_name, str) or ".mixture_lora." not in parameter_name: + raise ValueError(f"Invalid Mixture-of-LoRA parameter name: {parameter_name!r}") + site_id, parameter_kind = parameter_name.split(".mixture_lora.", maxsplit=1) + if parameter_kind not in _PARAMETER_KINDS: + raise ValueError(f"Unsupported Mixture-of-LoRA parameter kind: {parameter_kind!r}") + site_parts = site_id.split(".") + if len(site_parts) != 5 or site_parts[:2] != ["decoder", "layers"] or site_parts[3] != "self_attention": + raise ValueError(f"Unsupported Mixture-of-LoRA site_id: {site_id!r}") + try: + layer_id = int(site_parts[2]) + except ValueError as error: + raise ValueError(f"Mixture-of-LoRA layer id must be an integer: {site_parts[2]!r}") from error + if layer_id < 0: + raise ValueError(f"Mixture-of-LoRA layer id must be non-negative: {layer_id}") + target_mapping = { + "linear_qkv": "qkv_proj", + "linear_proj": "o_proj", + } + target = site_parts[4] + if target not in target_mapping: + raise ValueError(f"Unsupported Mixture-of-LoRA target module: {target!r}") + return f"model.layers.{layer_id}.self_attn.{target_mapping[target]}.mixture_lora.{parameter_kind}" + + @dataclass(frozen=True) class MixtureLoraStateSpec: """Stable identity and global layout for one Mixture-of-LoRA tensor.""" diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py new file mode 100644 index 000000000..d658c9bbe --- /dev/null +++ b/tests/backends/sglang/test_mixture_lora.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import os + +import pytest + +from relax.utils.mixture_lora import ( + MixtureLoraConfig, + configure_mixture_lora_external_model, + deserialize_mixture_lora_config, +) + + +def _config(**overrides): + values = { + "rank": 16, + "num_experts": 4, + "top_k": 2, + "temperature": 0.8, + "aux_loss_coef": 0.01, + "alpha": 32.0, + "target_modules": ("linear_qkv", "linear_proj"), + } + values.update(overrides) + return MixtureLoraConfig(**values) + + +def test_mixture_lora_configures_external_package_before_spawn(monkeypatch): + monkeypatch.delenv("RELAX_MIXTURE_LORA_CONFIG", raising=False) + + package = configure_mixture_lora_external_model(_config(), None) + + assert package == "relax.models.qwen3_mixture_lora.sglang" + config = deserialize_mixture_lora_config(os.environ["RELAX_MIXTURE_LORA_CONFIG"]) + assert config.num_experts == 4 + assert config.rank == 16 + assert config.top_k == 2 + assert config.target_modules == ("linear_qkv", "linear_proj") + + +def test_single_lora_does_not_enable_external_mixture_model(monkeypatch): + monkeypatch.delenv("RELAX_MIXTURE_LORA_CONFIG", raising=False) + + package = configure_mixture_lora_external_model( + None, + "custom.single_lora.package", + ) + + assert package == "custom.single_lora.package" + assert "RELAX_MIXTURE_LORA_CONFIG" not in os.environ + + +def test_mixture_lora_rejects_conflicting_external_package(): + with pytest.raises(ValueError, match="requires the Qwen3 external model package"): + configure_mixture_lora_external_model(_config(), "custom.other.package") diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py new file mode 100644 index 000000000..4a231fb7c --- /dev/null +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from relax.backends.megatron.mixture_lora import MixtureLoRAAdapter +from relax.models.qwen3_mixture_lora.sglang.model import ( + SGLangMixtureLoRA, + attach_sglang_mixture_lora, + load_sglang_mixture_lora_weights, +) +from relax.utils.mixture_lora import MixtureLoraConfig + + +def _config(): + return MixtureLoraConfig( + num_experts=4, + rank=2, + top_k=2, + temperature=0.8, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_qkv", "linear_proj"), + ) + + +class _TupleLinear(nn.Module): + def __init__(self, input_size: int, output_size: int): + super().__init__() + self.weight = nn.Parameter(torch.randn(output_size, input_size)) + + def forward(self, x): + return F.linear(x, self.weight), None + + +def _fake_qwen_model_with_routed_qkv(): + model = nn.Module() + model.model = nn.Module() + layer = nn.Module() + layer.self_attn = nn.Module() + layer.self_attn.qkv_proj = _TupleLinear(4, 6) + model.model.layers = nn.ModuleList([layer]) + attach_sglang_mixture_lora( + layer.self_attn.qkv_proj, + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + ) + return model + + +def test_sglang_dense_adapter_matches_training_adapter(): + config = _config() + training_adapter = MixtureLoRAAdapter( + config, + "linear_qkv", + 4, + 6, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + rollout_adapter = SGLangMixtureLoRA( + config, + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + training_adapter.experts.lora_B.normal_(mean=0.0, std=0.2) + rollout_adapter.experts.lora_A.copy_(training_adapter.experts.lora_A) + rollout_adapter.experts.lora_B.copy_(training_adapter.experts.lora_B) + rollout_adapter.router.weight.copy_(training_adapter.router.weight) + x = torch.randn(5, 4) + + training_output, training_decision = training_adapter.forward_with_routing(x) + rollout_output, rollout_decision = rollout_adapter.forward_with_routing(x) + + torch.testing.assert_close(rollout_output, training_output) + torch.testing.assert_close(rollout_decision.pre_topk_probs, training_decision.pre_topk_probs) + torch.testing.assert_close(rollout_decision.post_topk_weights, training_decision.post_topk_weights) + assert torch.equal(rollout_decision.topk_indices, training_decision.topk_indices) + + +def test_sglang_router_uses_fp32_logits_with_bfloat16_parameters(): + adapter = SGLangMixtureLoRA( + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + decision = adapter.route(torch.randn(5, 4, dtype=torch.bfloat16)) + + assert decision.pre_topk_probs.dtype == torch.float32 + assert decision.post_topk_weights.dtype == torch.float32 + + +def test_attached_adapter_preserves_base_parameter_name_and_adds_delta(): + model = _fake_qwen_model_with_routed_qkv() + linear = model.model.layers[0].self_attn.qkv_proj + x = torch.randn(3, 4) + base_output = F.linear(x, linear.weight) + with torch.no_grad(): + linear.mixture_lora.experts.lora_B.normal_(mean=0.0, std=0.2) + + output, bias = linear(x) + + assert bias is None + assert output.shape == base_output.shape + assert not torch.equal(output, base_output) + parameter_names = set(dict(model.named_parameters())) + assert "model.layers.0.self_attn.qkv_proj.weight" in parameter_names + assert "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight" in parameter_names + + +def test_sglang_weight_loader_maps_training_names_and_validates_tensors(): + model = _fake_qwen_model_with_routed_qkv() + prefix = "decoder.layers.0.self_attention.linear_qkv.mixture_lora" + weights = { + f"{prefix}.experts.lora_A": torch.randn(4, 2, 4), + f"{prefix}.experts.lora_B": torch.randn(4, 6, 2), + f"{prefix}.router.weight": torch.randn(4, 4), + } + + loaded_names = load_sglang_mixture_lora_weights(model, weights.items()) + + assert loaded_names == { + "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_A", + "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_B", + "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight", + } + parameters = dict(model.named_parameters()) + for source_name, loaded_weight in weights.items(): + target_name = source_name.replace( + "decoder.layers.0.self_attention.linear_qkv", + "model.layers.0.self_attn.qkv_proj", + ) + torch.testing.assert_close(parameters[target_name], loaded_weight) + + with pytest.raises(ValueError, match="has shape"): + load_sglang_mixture_lora_weights(model, [(f"{prefix}.router.weight", torch.randn(4, 5))]) + with pytest.raises(TypeError, match="has dtype"): + load_sglang_mixture_lora_weights( + model, + [(f"{prefix}.router.weight", torch.randn(4, 4, dtype=torch.float64))], + ) + with pytest.raises(ValueError, match="Unknown Mixture-of-LoRA weight"): + load_sglang_mixture_lora_weights( + model, + [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], + ) diff --git a/tests/utils/test_mixture_lora_routing.py b/tests/utils/test_mixture_lora_routing.py index 177033939..571c383b8 100644 --- a/tests/utils/test_mixture_lora_routing.py +++ b/tests/utils/test_mixture_lora_routing.py @@ -16,8 +16,11 @@ TransportTensorSpec, build_mixture_lora_state_specs, compute_routing_statistics, + deserialize_mixture_lora_config, mean_routing_balance_loss, + megatron_mixture_lora_name_to_sglang, route_topk, + serialize_mixture_lora_config, ) @@ -41,6 +44,63 @@ def _config(**overrides): return MixtureLoraConfig(**values) +def test_runtime_config_json_round_trip_is_canonical(): + config = _config() + + serialized = serialize_mixture_lora_config(config) + + assert deserialize_mixture_lora_config(serialized) == config + assert serialize_mixture_lora_config(deserialize_mixture_lora_config(serialized)) == serialized + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + ("", "non-empty JSON string"), + ("[]", "must contain an object"), + ('{"num_experts":4}', "fields do not match schema"), + ("not-json", "Invalid Mixture-of-LoRA runtime configuration JSON"), + ], +) +def test_runtime_config_json_rejects_incomplete_or_invalid_input(raw, message): + with pytest.raises(ValueError, match=message): + deserialize_mixture_lora_config(raw) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ( + "decoder.layers.0.self_attention.linear_qkv.mixture_lora.experts.lora_A", + "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_A", + ), + ( + "decoder.layers.17.self_attention.linear_qkv.mixture_lora.experts.lora_B", + "model.layers.17.self_attn.qkv_proj.mixture_lora.experts.lora_B", + ), + ( + "decoder.layers.2.self_attention.linear_proj.mixture_lora.router.weight", + "model.layers.2.self_attn.o_proj.mixture_lora.router.weight", + ), + ], +) +def test_megatron_mixture_parameter_name_maps_to_sglang(source, expected): + assert megatron_mixture_lora_name_to_sglang(source) == expected + + +@pytest.mark.parametrize( + "name", + [ + "decoder.layers.x.self_attention.linear_qkv.mixture_lora.router.weight", + "decoder.layers.0.mlp.linear_fc1.mixture_lora.router.weight", + "decoder.layers.0.self_attention.linear_qkv.mixture_lora.unknown", + ], +) +def test_megatron_mixture_parameter_name_rejects_unsupported_schema(name): + with pytest.raises(ValueError): + megatron_mixture_lora_name_to_sglang(name) + + def _reference_routed_lora(x, lora_a, lora_b, decision, scale): outputs = [] for token, indices, weights in zip(x.reshape(-1, x.shape[-1]), decision.topk_indices, decision.post_topk_weights): From 755dc4c36ee3e14102958656347b722d9f82890f Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 12:25:55 +0800 Subject: [PATCH 10/41] fix(lora): validate SGLang rollout execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLang PP worker 加载 Mixture 参数时跳过不属于本 stage 的 layer，同时保留本地未知参数的严格报错。 新增 H20 BF16 CUDA Graph capture 测试和 PP layer 过滤测试，确认 routed adapter 可以沿用现有 graph 执行路径。 --- .../models/qwen3_mixture_lora/sglang/model.py | 5 +++ .../qwen3_mixture_lora/test_sglang_model.py | 40 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index fb893176c..1564340a7 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -132,6 +132,11 @@ def load_sglang_mixture_lora_weights( if target_name in loaded_names: raise ValueError(f"Duplicate Mixture-of-LoRA weight: {target_name}") if target_name not in parameters: + layer_id = int(target_name.split(".", maxsplit=3)[2]) + start_layer = getattr(getattr(model, "model", None), "start_layer", None) + end_layer = getattr(getattr(model, "model", None), "end_layer", None) + if start_layer is not None and end_layer is not None and not start_layer <= layer_id < end_layer: + continue raise ValueError(f"Unknown Mixture-of-LoRA weight for SGLang: {target_name}") parameter = parameters[target_name] if tuple(loaded_weight.shape) != tuple(parameter.shape): diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 4a231fb7c..6325c07db 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -157,3 +157,43 @@ def test_sglang_weight_loader_maps_training_names_and_validates_tensors(): model, [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], ) + + +def test_sglang_weight_loader_skips_layers_owned_by_another_pp_stage(): + model = _fake_qwen_model_with_routed_qkv() + model.model.start_layer = 0 + model.model.end_layer = 1 + + loaded_names = load_sglang_mixture_lora_weights( + model, + [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], + ) + + assert loaded_names == set() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_sglang_adapter_can_be_captured_by_cuda_graph(): + adapter = SGLangMixtureLoRA( + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 16, + 24, + device=torch.device("cuda"), + dtype=torch.bfloat16, + ) + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + adapter(x) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = adapter(x) + graph.replay() + + assert captured_output.shape == (8, 24) + assert torch.isfinite(captured_output).all() From a48c17fc4b33597dda20570cc04b2d95e707852a Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 12:43:52 +0800 Subject: [PATCH 11/41] feat(lora): sync routed adapters to rollout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 Mixture-of-LoRA 权重收集器，按固定 schema 汇总 PP/TP 分片，并将 Megatron 的 GQA 分组 QKV 排布转换为 SGLang 的 Q/K/V 连续排布。 首轮同步发送 base、全部 expert 和 router，后续只发送 expert 和 router；只有最后一个 routed chunk 携带新 weight version。复用现有 pause、flush、IPC update 和 continue 流程，并从 raw/Bridge HF iterator 中排除 Mixture 参数。 补充 QKV 转换、HF 过滤、首轮与后续批次、CPU mirror、生命周期顺序和单 LoRA 回归测试。 --- .../hf_weight_iterator_bridge.py | 3 + .../hf_weight_iterator_direct.py | 3 + .../weight_update/mixture_lora_sync.py | 295 ++++++++++++++++++ .../update_weight_from_tensor.py | 115 ++++++- .../test_mixture_lora_weight_sync.py | 288 +++++++++++++++++ 5 files changed, 699 insertions(+), 5 deletions(-) create mode 100644 relax/backends/megatron/weight_update/mixture_lora_sync.py create mode 100644 tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py index ad05ae770..7d2e37e03 100644 --- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py +++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py @@ -17,6 +17,7 @@ is_lora_adapter_param, is_lora_enabled, is_lora_merge_mode, + is_mixture_lora_param, ) from relax.utils.types import ParamInfo @@ -243,6 +244,8 @@ def _build_param_info_buckets(args, model, collect_adapters=False): vanilla_key_map = {} adapter_map: dict[str, dict[str, str]] = {} for (v_name, v_param), (g_name, _g_param) in zip(vanilla_iter, global_iter, strict=True): + if is_mixture_lora_param(g_name): + continue if collect_adapters and is_lora_adapter_param(g_name): # LoRA adapter param: keep it out of the conversion buckets (no standalone # bridge mapping) and record its vanilla key for load-time merging. diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_direct.py b/relax/backends/megatron/weight_update/hf_weight_iterator_direct.py index 83976cb4a..bb88e3d6d 100644 --- a/relax/backends/megatron/weight_update/hf_weight_iterator_direct.py +++ b/relax/backends/megatron/weight_update/hf_weight_iterator_direct.py @@ -9,6 +9,7 @@ from relax.utils import device as device_utils from relax.utils.distributed_utils import get_gloo_group +from relax.utils.megatron_peft_utils import is_mixture_lora_param from relax.utils.types import ParamInfo from ..sglang import monkey_patch_torch_reductions @@ -152,6 +153,8 @@ def _get_megatron_local_param_infos(args: Namespace, model: Sequence[torch.nn.Mo param_infos = {} rank = dist.get_rank() for name, param in named_params_and_buffers(args, model): + if is_mixture_lora_param(name): + continue param_infos[name] = ParamInfo( name=name, dtype=param.dtype, diff --git a/relax/backends/megatron/weight_update/mixture_lora_sync.py b/relax/backends/megatron/weight_update/mixture_lora_sync.py new file mode 100644 index 000000000..33d91e3bc --- /dev/null +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Collect and convert Mixture-of-LoRA parameters for rollout sync.""" + +from argparse import Namespace +from dataclasses import dataclass +from typing import Iterator, Mapping, Sequence + +import torch +import torch.distributed as dist +from megatron.core import mpu + +from relax.backends.megatron.misc_utils import strip_param_name_prefix +from relax.utils import device as device_utils +from relax.utils.megatron_peft_utils import build_mixture_lora_config, is_mixture_lora_param +from relax.utils.mixture_lora import ( + MixtureLoraStateSpec, + build_mixture_lora_state_specs, +) + +from .common import named_params_and_buffers + + +@dataclass(frozen=True) +class MixtureLoraParamInfo: + """One routed tensor and the training rank that owns its PP stage.""" + + state: MixtureLoraStateSpec + local_shape: tuple[int, ...] + tp_shard_dim: int | None + src_rank: int + weight_key: str + + +def _qwen3_attention_dimensions(args: Namespace) -> tuple[int, int, int]: + head_dim = getattr(args, "kv_channels", None) + if head_dim is None: + head_dim = args.hidden_size // args.num_attention_heads + if args.num_attention_heads % args.num_query_groups != 0: + raise ValueError("num_attention_heads must be divisible by num_query_groups") + return head_dim, args.num_attention_heads, args.num_query_groups + + +def _site_dimensions(args: Namespace, site_id: str) -> tuple[int, int]: + head_dim, num_attention_heads, num_query_groups = _qwen3_attention_dimensions(args) + target = site_id.rsplit(".", maxsplit=1)[-1] + if target == "linear_qkv": + return args.hidden_size, (num_attention_heads + 2 * num_query_groups) * head_dim + if target == "linear_proj": + return num_attention_heads * head_dim, args.hidden_size + raise ValueError(f"Unsupported Mixture-of-LoRA site: {site_id!r}") + + +def _parameter_kind(parameter_name: str) -> str: + for kind in ("experts.lora_A", "experts.lora_B", "router.weight"): + if parameter_name.endswith(f".mixture_lora.{kind}"): + return kind + raise ValueError(f"Unsupported Mixture-of-LoRA parameter name: {parameter_name!r}") + + +def _tp_shard_dim(site_id: str, parameter_kind: str) -> int | None: + target = site_id.rsplit(".", maxsplit=1)[-1] + if target == "linear_qkv": + return { + "experts.lora_A": 1, + "experts.lora_B": 1, + "router.weight": None, + }[parameter_kind] + if target == "linear_proj": + return { + "experts.lora_A": 2, + "experts.lora_B": 1, + "router.weight": 1, + }[parameter_kind] + raise ValueError(f"Unsupported Mixture-of-LoRA site: {site_id!r}") + + +def _qkv_lora_b_to_sglang( + tensor: torch.Tensor, + *, + num_attention_heads: int, + num_query_groups: int, + head_dim: int, +) -> torch.Tensor: + """Convert Megatron's group-interleaved QKV rows to Q/K/V blocks.""" + + queries_per_group = num_attention_heads // num_query_groups + expected_rows = (num_attention_heads + 2 * num_query_groups) * head_dim + if tensor.shape[1] != expected_rows: + raise ValueError(f"QKV LoRA B has {tensor.shape[1]} rows, expected {expected_rows}") + grouped = tensor.reshape( + tensor.shape[0], + num_query_groups, + queries_per_group + 2, + head_dim, + tensor.shape[-1], + ) + query, key, value = torch.split(grouped, [queries_per_group, 1, 1], dim=2) + return torch.cat( + [ + query.reshape(tensor.shape[0], -1, tensor.shape[-1]), + key.reshape(tensor.shape[0], -1, tensor.shape[-1]), + value.reshape(tensor.shape[0], -1, tensor.shape[-1]), + ], + dim=1, + ) + + +def merge_mixture_lora_tp_shards( + info: MixtureLoraParamInfo, + shards: Sequence[torch.Tensor], + *, + num_attention_heads: int, + num_query_groups: int, + head_dim: int, +) -> torch.Tensor: + """Reconstruct one global tensor and convert backend-specific layout.""" + + if not shards: + raise ValueError(f"No TP shards supplied for {info.state.parameter_name}") + if any(tuple(shard.shape) != info.local_shape for shard in shards): + raise ValueError(f"TP shard shape mismatch for {info.state.parameter_name}") + if info.tp_shard_dim is None: + if len(shards) != 1: + raise ValueError(f"Replicated parameter {info.state.parameter_name} expects one tensor") + merged = shards[0] + else: + merged = torch.cat(tuple(shards), dim=info.tp_shard_dim) + + if info.state.site_id.endswith(".linear_qkv") and info.state.parameter_kind == "experts.lora_B": + merged = _qkv_lora_b_to_sglang( + merged, + num_attention_heads=num_attention_heads, + num_query_groups=num_query_groups, + head_dim=head_dim, + ) + if tuple(merged.shape) != info.state.global_shape: + raise ValueError( + f"Reconstructed {info.state.parameter_name} has shape {tuple(merged.shape)}, " + f"expected {info.state.global_shape}" + ) + if merged.dtype != info.state.dtype: + raise TypeError( + f"Reconstructed {info.state.parameter_name} has dtype {merged.dtype}, expected {info.state.dtype}" + ) + return merged.contiguous() + + +class MixtureLoraSync: + """Gather every PP/TP shard into SGLang-ready routed tensors.""" + + def __init__(self, args: Namespace, model: Sequence[torch.nn.Module]) -> None: + self.args = args + self.model = model + self.config = build_mixture_lora_config(args) + if self.config is None: + raise ValueError("MixtureLoraSync requires an enabled Mixture-of-LoRA configuration") + self.base_sync_done = False + self.param_infos = self._build_param_infos() + if not self.param_infos: + raise ValueError("No Mixture-of-LoRA parameters were found in the training model") + self._validate_param_infos() + + def _validate_param_infos(self) -> None: + expected_kinds = {"experts.lora_A", "experts.lora_B", "router.weight"} + kinds_by_site: dict[str, set[str]] = {} + for info in self.param_infos: + kinds_by_site.setdefault(info.state.site_id, set()).add(info.state.parameter_kind) + incomplete = { + site_id: sorted(expected_kinds - parameter_kinds) + for site_id, parameter_kinds in kinds_by_site.items() + if parameter_kinds != expected_kinds + } + if incomplete: + raise ValueError(f"Incomplete Mixture-of-LoRA parameter schema: {incomplete}") + + def _build_param_infos(self) -> tuple[MixtureLoraParamInfo, ...]: + rank = dist.get_rank() + tp_world_size = mpu.get_tensor_model_parallel_world_size() + local_infos: dict[str, MixtureLoraParamInfo] = {} + vanilla = named_params_and_buffers(self.args, self.model, convert_to_global_name=False) + global_names = named_params_and_buffers(self.args, self.model, convert_to_global_name=True) + for (vanilla_name, parameter), (global_name, _) in zip(vanilla, global_names, strict=True): + parameter_name = strip_param_name_prefix(global_name) + if not is_mixture_lora_param(parameter_name): + continue + kind = _parameter_kind(parameter_name) + site_id = parameter_name.removesuffix(f".mixture_lora.{kind}") + input_size, output_size = _site_dimensions(self.args, site_id) + state = next( + spec + for spec in build_mixture_lora_state_specs( + self.config, + site_id, + input_size, + output_size, + parameter.dtype, + ) + if spec.parameter_kind == kind + ) + shard_dim = _tp_shard_dim(site_id, kind) + expected_local_shape = list(state.global_shape) + if shard_dim is not None: + if expected_local_shape[shard_dim] % tp_world_size != 0: + raise ValueError(f"{state.parameter_name} cannot be sharded over TP={tp_world_size}") + expected_local_shape[shard_dim] //= tp_world_size + if tuple(parameter.shape) != tuple(expected_local_shape): + raise ValueError( + f"Training parameter {state.parameter_name} has local shape {tuple(parameter.shape)}, " + f"expected {tuple(expected_local_shape)}" + ) + weight_key = global_name if self.args.megatron_to_hf_mode == "raw" else vanilla_name + local_infos[state.parameter_name] = MixtureLoraParamInfo( + state=state, + local_shape=tuple(parameter.shape), + tp_shard_dim=shard_dim, + src_rank=rank, + weight_key=weight_key, + ) + + if mpu.get_pipeline_model_parallel_world_size() > 1: + gathered_infos = [None] * mpu.get_pipeline_model_parallel_world_size() + dist.all_gather_object( + (rank, local_infos), + object_list=gathered_infos, + group=mpu.get_pipeline_model_parallel_group(), + ) + for _, stage_infos in gathered_infos: + for name, info in stage_infos.items(): + previous = local_infos.get(name) + if previous is not None and previous != info: + raise ValueError(f"Conflicting Mixture-of-LoRA metadata for {name}") + local_infos[name] = info + return tuple(local_infos[name] for name in sorted(local_infos)) + + def get_weight_chunks( + self, + local_weights: Mapping[str, torch.Tensor], + ) -> Iterator[list[tuple[str, torch.Tensor]]]: + """Yield full routed tensors in deterministic, size-bounded chunks.""" + + rank = dist.get_rank() + device = device_utils.make_current_torch_device() + tp_world_size = mpu.get_tensor_model_parallel_world_size() + tp_group = mpu.get_tensor_model_parallel_group() + pp_world_size = mpu.get_pipeline_model_parallel_world_size() + pp_group = mpu.get_pipeline_model_parallel_group() if pp_world_size > 1 else None + pp_ranks = set(dist.get_process_group_ranks(pp_group)) if pp_group is not None else {rank} + head_dim, num_attention_heads, num_query_groups = _qwen3_attention_dimensions(self.args) + chunk: list[tuple[str, torch.Tensor]] = [] + chunk_size = 0 + + for info in self.param_infos: + if rank == info.src_rank: + if info.weight_key not in local_weights: + raise KeyError(f"Missing Mixture-of-LoRA weight {info.weight_key!r}") + source_tensor = local_weights[info.weight_key] + if tuple(source_tensor.shape) != info.local_shape: + raise ValueError( + f"Mixture-of-LoRA weight {info.weight_key!r} has shape {tuple(source_tensor.shape)}, " + f"expected {info.local_shape}" + ) + if source_tensor.dtype != info.state.dtype: + raise TypeError( + f"Mixture-of-LoRA weight {info.weight_key!r} has dtype {source_tensor.dtype}, " + f"expected {info.state.dtype}" + ) + local_tensor = source_tensor.to(device=device) + else: + local_tensor = torch.empty(info.local_shape, dtype=info.state.dtype, device=device) + if pp_group is not None and info.src_rank in pp_ranks: + dist.broadcast(local_tensor, src=info.src_rank, group=pp_group) + + if info.tp_shard_dim is None or tp_world_size == 1: + shards = [local_tensor] + else: + shards = [torch.empty_like(local_tensor) for _ in range(tp_world_size)] + dist.all_gather(shards, local_tensor.contiguous(), group=tp_group) + full_tensor = merge_mixture_lora_tp_shards( + info, + shards, + num_attention_heads=num_attention_heads, + num_query_groups=num_query_groups, + head_dim=head_dim, + ) + tensor_size = full_tensor.numel() * full_tensor.element_size() + if chunk and chunk_size + tensor_size > self.args.update_weight_buffer_size: + yield chunk + chunk = [] + chunk_size = 0 + chunk.append((info.state.parameter_name, full_tensor)) + chunk_size += tensor_size + + if chunk: + yield chunk diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index e38c85ae5..b2570cd8d 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -25,11 +25,13 @@ is_lora_adapter_param, is_lora_enabled, is_lora_merge_mode, + is_mixture_lora_enabled, ) from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer from .hf_weight_iterator_base import HfWeightIteratorBase from .lora_adapter_sync import LoraAdapterSync +from .mixture_lora_sync import MixtureLoraSync from .update_weight_from_distributed import ( connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, @@ -39,6 +41,26 @@ logger = get_logger(__name__) +_CURRENT_WEIGHT_VERSION = object() + + +def iter_mixture_weight_updates(base_chunks, mixture_chunks, *, include_base: bool, weight_version: int): + """Yield base first and attach the version only to the final routed + chunk.""" + + if include_base: + for chunk in base_chunks: + yield chunk, None + + iterator = iter(mixture_chunks) + try: + pending = next(iterator) + except StopIteration as error: + raise ValueError("Mixture-of-LoRA weight sync produced no routed tensors") from error + for chunk in iterator: + yield pending, None + pending = chunk + yield pending, weight_version class UpdateWeightFromTensor: @@ -73,10 +95,12 @@ def __init__( self.lora_enabled = is_lora_enabled(args) self.lora_merge_mode = is_lora_merge_mode(args) if self.lora_enabled else False self.lora_adapter_mode = is_lora_adapter_mode(args) if self.lora_enabled else False + self.mixture_lora_enabled = is_mixture_lora_enabled(args) # Adapter-mode incremental-sync state (base-once + adapter-delta protocol) lives in the # shared LoraAdapterSync helper; the backend keeps only its in-memory transport below. self._lora_sync = LoraAdapterSync(args, model) if self.lora_adapter_mode else None + self._mixture_lora_sync = MixtureLoraSync(args, model) if self.mixture_lora_enabled else None self._hf_weight_iterator = HfWeightIteratorBase.create( args=args, model=model, model_name=model_name, quantization_config=quantization_config @@ -131,6 +155,9 @@ def connect_rollout_engines( self.use_distribute = len(rollout_engines) > colocate_engine_nums + if self.mixture_lora_enabled and self.use_distribute: + raise ValueError("Mixture-of-LoRA weight sync currently requires colocated rollout engines") + if self.use_distribute: self.rollout_engines = rollout_engines[:colocate_engine_nums] self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:] @@ -191,6 +218,9 @@ def update_weights(self) -> None: if self.lora_enabled and self.lora_adapter_mode: self._update_weights_adapter_mode() return + if self.mixture_lora_enabled: + self._update_weights_mixture_lora() + return self.weight_version += 1 @@ -271,6 +301,75 @@ def update_weights(self) -> None: ray.get([engine.continue_generation.remote() for engine in all_engines]) dist.barrier(group=get_gloo_group()) + def _update_weights_mixture_lora(self) -> None: + """Sync the frozen base once and all routed parameters every step.""" + + next_weight_version = self.weight_version + 1 + include_base = not self._mixture_lora_sync.base_sync_done + all_engines = list(self.rollout_engines) + rank = dist.get_rank() + + if rank == 0: + ray.get([engine.pause_generation.remote() for engine in all_engines]) + ray.get([engine.flush_cache.remote() for engine in all_engines]) + if ( + include_base + and self.quantization_config + and self.quantization_config["quant_method"] in ["compressed-tensors"] + ): + post_process_weights( + restore_weights_before_load=True, + post_process_quantization=False, + rollout_engines=all_engines, + ) + dist.barrier(group=get_gloo_group()) + + local_weights = self.weights_getter() + base_chunks = self._hf_weight_iterator.get_hf_weight_chunks(local_weights) + mixture_chunks = self._mixture_lora_sync.get_weight_chunks(local_weights) + updates = iter_mixture_weight_updates( + base_chunks, + mixture_chunks, + include_base=include_base, + weight_version=next_weight_version, + ) + self._send_weight_update_stream(updates) + + dist.barrier(group=get_gloo_group()) + if rank == 0: + if ( + include_base + and self.quantization_config + and self.quantization_config["quant_method"] in ["compressed-tensors"] + ): + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=all_engines, + ) + ray.get([engine.continue_generation.remote() for engine in all_engines]) + dist.barrier(group=get_gloo_group()) + + self.weight_version = next_weight_version + self._mixture_lora_sync.base_sync_done = True + + def _send_weight_update_stream(self, updates) -> None: + """Pipeline conversion collectives with the preceding IPC request.""" + + previous_refs: list[ObjectRef] = [] + previous_tensors = None + for named_tensors, weight_version in updates: + refs, long_lived_tensors = self._send_hf_params(named_tensors, weight_version=weight_version) + if previous_refs: + ray.get(previous_refs) + del previous_tensors + previous_refs = refs + previous_tensors = long_lived_tensors + device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) + if previous_refs: + ray.get(previous_refs) + del previous_tensors + def _update_weights_adapter_mode(self) -> None: """LoRA adapter mode: sync base once, then push only the adapter each step. @@ -476,8 +575,14 @@ def _push_lora_adapter(self, all_params: Mapping[str, torch.Tensor], *, first_sy finally: set_sharing_strategy(prev_strategy) - def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: + def _send_hf_params( + self, + hf_named_tensors, + *, + weight_version: int | None | object = _CURRENT_WEIGHT_VERSION, + ) -> tuple[list[ObjectRef], Any]: all_refs = [] + resolved_weight_version = self.weight_version if weight_version is _CURRENT_WEIGHT_VERSION else weight_version long_lived_tensors = None if self._ipc_engine is not None: @@ -486,7 +591,7 @@ def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: ipc_engine=self._ipc_engine, ipc_gather_src=self._ipc_gather_src, ipc_gather_group=self._ipc_gather_group, - weight_version=self.weight_version, + weight_version=resolved_weight_version, ) all_refs.extend(refs_colocated) @@ -494,7 +599,7 @@ def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: refs_distributed = update_weights_from_distributed( self._group_name, self._model_update_groups, - self.weight_version, + resolved_weight_version, self.distributed_rollout_engines, hf_named_tensors, ) @@ -510,7 +615,7 @@ def _send_to_colocated_engine( ipc_engine, ipc_gather_src, ipc_gather_group, - weight_version, + weight_version: int | None, ) -> tuple[list[ObjectRef], Any]: # Placeholder ranks (GPU slots reserved but no engine) have no gather group. # gather_object is only collective among group members, so we skip entirely. @@ -571,7 +676,7 @@ def _send_to_colocated_engine( ipc_engine.update_weights_from_tensor.remote( serialized_named_tensors=[tensors[i] for tensors in serialized_named_tensors], load_format="flattened_bucket", - weight_version=str(weight_version), + weight_version=None if weight_version is None else str(weight_version), ) ) diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py new file mode 100644 index 000000000..1a7fcf922 --- /dev/null +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -0,0 +1,288 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from relax.backends.megatron.weight_update.mixture_lora_sync import ( + MixtureLoraParamInfo, + merge_mixture_lora_tp_shards, +) +from relax.backends.megatron.weight_update.update_weight_from_tensor import iter_mixture_weight_updates +from relax.utils.mixture_lora import MixtureLoraStateSpec + + +def _info(site, kind, global_shape, local_shape, shard_dim): + return MixtureLoraParamInfo( + state=MixtureLoraStateSpec( + schema_version=1, + site_id=f"decoder.layers.0.self_attention.{site}", + parameter_kind=kind, + global_shape=global_shape, + dtype=torch.float32, + ), + local_shape=local_shape, + tp_shard_dim=shard_dim, + src_rank=0, + weight_key="weight", + ) + + +def test_qkv_lora_b_tp_shards_are_converted_from_group_layout_to_qkv_blocks(): + # Two query groups, two query heads per group, then one K and one V head. + grouped = torch.tensor([[[[10.0], [11.0], [20.0], [30.0], [12.0], [13.0], [21.0], [31.0]]]]).reshape(1, 8, 1) + info = _info("linear_qkv", "experts.lora_B", (1, 8, 1), (1, 4, 1), 1) + + merged = merge_mixture_lora_tp_shards( + info, + grouped.chunk(2, dim=1), + num_attention_heads=4, + num_query_groups=2, + head_dim=1, + ) + + expected = torch.tensor([10.0, 11.0, 12.0, 13.0, 20.0, 21.0, 30.0, 31.0]).reshape(1, 8, 1) + torch.testing.assert_close(merged, expected) + + +@pytest.mark.parametrize( + ("site", "kind", "global_shape", "local_shape", "shard_dim"), + [ + ("linear_qkv", "experts.lora_A", (2, 4, 6), (2, 2, 6), 1), + ("linear_proj", "experts.lora_A", (2, 4, 6), (2, 4, 3), 2), + ("linear_proj", "experts.lora_B", (2, 6, 4), (2, 3, 4), 1), + ("linear_proj", "router.weight", (2, 6), (2, 3), 1), + ], +) +def test_non_qkv_output_shards_concatenate_on_the_schema_axis( + site, + kind, + global_shape, + local_shape, + shard_dim, +): + info = _info(site, kind, global_shape, local_shape, shard_dim) + shards = [torch.zeros(local_shape), torch.ones(local_shape)] + + merged = merge_mixture_lora_tp_shards( + info, + shards, + num_attention_heads=4, + num_query_groups=2, + head_dim=1, + ) + + assert merged.shape == global_shape + first, second = merged.chunk(2, dim=shard_dim) + assert torch.equal(first, shards[0]) + assert torch.equal(second, shards[1]) + + +def test_replicated_router_rejects_multiple_tp_copies(): + info = _info("linear_qkv", "router.weight", (2, 6), (2, 6), None) + + with pytest.raises(ValueError, match="expects one tensor"): + merge_mixture_lora_tp_shards( + info, + [torch.zeros(2, 6), torch.zeros(2, 6)], + num_attention_heads=4, + num_query_groups=2, + head_dim=1, + ) + + +def test_reconstructed_tensor_validates_shape_and_dtype(): + info = _info("linear_proj", "router.weight", (2, 6), (2, 3), 1) + + with pytest.raises(ValueError, match="TP shard shape mismatch"): + merge_mixture_lora_tp_shards( + info, + [torch.zeros(2, 4), torch.zeros(2, 4)], + num_attention_heads=4, + num_query_groups=2, + head_dim=1, + ) + with pytest.raises(TypeError, match="has dtype"): + merge_mixture_lora_tp_shards( + info, + [torch.zeros(2, 3, dtype=torch.float64), torch.zeros(2, 3, dtype=torch.float64)], + num_attention_heads=4, + num_query_groups=2, + head_dim=1, + ) + + +def _named_parameters(): + base = torch.nn.Parameter(torch.zeros(4, 4)) + mixture = torch.nn.Parameter(torch.zeros(2, 2, 4)) + return [ + ("module.module.decoder.layers.0.self_attention.linear_qkv.weight", base), + ( + "module.module.decoder.layers.0.self_attention.linear_qkv.mixture_lora.experts.lora_A", + mixture, + ), + ] + + +def test_direct_hf_iterator_excludes_mixture_parameters(): + from relax.backends.megatron.weight_update.hf_weight_iterator_direct import _get_megatron_local_param_infos + + args = SimpleNamespace(update_weight_buffer_size=1024, mtp_num_layers=None) + + def gather_single_process(obj, object_list, group=None): + del group + object_list[0] = obj + + with ( + patch( + "relax.backends.megatron.weight_update.hf_weight_iterator_direct.named_params_and_buffers", + return_value=iter(_named_parameters()), + ), + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.get_world_size", return_value=1), + patch("torch.distributed.all_gather_object", side_effect=gather_single_process), + patch( + "relax.backends.megatron.weight_update.hf_weight_iterator_direct.get_gloo_group", + return_value=None, + ), + patch("megatron.core.mpu.get_pipeline_model_parallel_world_size", return_value=1), + patch("megatron.core.mpu.get_expert_model_parallel_world_size", return_value=1), + ): + infos = _get_megatron_local_param_infos(args, model=[]) + + assert [info.name for info in infos] == ["module.module.decoder.layers.0.self_attention.linear_qkv.weight"] + + +def test_bridge_hf_iterator_excludes_mixture_parameters(): + from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import _build_param_info_buckets + + vanilla = [(f"vp_stages.0.{name}", parameter) for name, parameter in _named_parameters()] + args = SimpleNamespace(update_weight_buffer_size=1024, num_experts=None) + with ( + patch( + "relax.backends.megatron.weight_update.hf_weight_iterator_bridge.named_params_and_buffers", + side_effect=[iter(vanilla), iter(_named_parameters())], + ), + patch("torch.distributed.get_rank", return_value=0), + patch("megatron.core.mpu.get_pipeline_model_parallel_world_size", return_value=1), + patch("megatron.core.mpu.get_expert_model_parallel_world_size", return_value=1), + patch("megatron.core.mpu.get_tensor_model_parallel_world_size", return_value=1), + ): + expert_buckets, base_buckets, _, _ = _build_param_info_buckets(args, model=[]) + + assert expert_buckets == [] + assert [[info.name for info in bucket] for bucket in base_buckets] == [ + ["module.module.decoder.layers.0.self_attention.linear_qkv.weight"] + ] + + +def test_first_sync_sends_base_then_routes_and_versions_only_the_final_chunk(): + updates = list( + iter_mixture_weight_updates( + base_chunks=[["base-0"], ["base-1"]], + mixture_chunks=[["mixture-0"], ["mixture-1"]], + include_base=True, + weight_version=7, + ) + ) + + assert updates == [ + (["base-0"], None), + (["base-1"], None), + (["mixture-0"], None), + (["mixture-1"], 7), + ] + + +def test_subsequent_sync_skips_base_and_sends_all_routed_parameters(): + def base_chunks_must_not_be_read(): + raise AssertionError("base chunks were read after the first sync") + yield + + updates = list( + iter_mixture_weight_updates( + base_chunks=base_chunks_must_not_be_read(), + mixture_chunks=[["mixture-0"], ["mixture-1"]], + include_base=False, + weight_version=8, + ) + ) + + assert updates == [(["mixture-0"], None), (["mixture-1"], 8)] + + +def test_sync_rejects_an_empty_mixture_parameter_set(): + with pytest.raises(ValueError, match="produced no routed tensors"): + list( + iter_mixture_weight_updates( + base_chunks=[["base"]], + mixture_chunks=[], + include_base=True, + weight_version=1, + ) + ) + + +def test_update_lifecycle_keeps_call_order_and_skips_base_after_first_sync(): + from relax.backends.megatron.weight_update.update_weight_from_tensor import UpdateWeightFromTensor + + events = [] + + class RemoteMethod: + def __init__(self, name): + self.name = name + + def remote(self): + events.append(self.name) + return self.name + + engine = SimpleNamespace( + pause_generation=RemoteMethod("pause"), + flush_cache=RemoteMethod("flush"), + continue_generation=RemoteMethod("continue"), + ) + local_weights = {"cpu-mirror": torch.ones(1)} + updater = UpdateWeightFromTensor.__new__(UpdateWeightFromTensor) + updater.weight_version = 0 + updater.rollout_engines = [engine] + updater.quantization_config = None + updater.weights_getter = MagicMock(return_value=local_weights) + updater._hf_weight_iterator = MagicMock() + updater._hf_weight_iterator.get_hf_weight_chunks.return_value = iter([["base"]]) + updater._mixture_lora_sync = SimpleNamespace(base_sync_done=False) + updater._mixture_lora_sync.get_weight_chunks = MagicMock( + side_effect=[iter([["mixture-1"]]), iter([["mixture-2"]])] + ) + sent_updates = [] + + def record_updates(updates): + events.append("update") + sent_updates.append(list(updates)) + + updater._send_weight_update_stream = record_updates + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.barrier"), + patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + patch("ray.get", side_effect=lambda refs: refs), + ): + updater._update_weights_mixture_lora() + updater._hf_weight_iterator.get_hf_weight_chunks.return_value = iter([["must-not-be-read"]]) + updater._update_weights_mixture_lora() + + assert events == ["pause", "flush", "update", "continue", "pause", "flush", "update", "continue"] + assert sent_updates == [ + [(["base"], None), (["mixture-1"], 1)], + [(["mixture-2"], 2)], + ] + assert updater.weights_getter.call_count == 2 + assert len(updater._mixture_lora_sync.get_weight_chunks.call_args_list) == 2 + assert all( + recorded_call.args[0] is local_weights + for recorded_call in updater._mixture_lora_sync.get_weight_chunks.call_args_list + ) + assert updater.weight_version == 2 + assert updater._mixture_lora_sync.base_sync_done is True From 76c014f6d8243e0003dbf0e2790ecf34d4bc402b Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 13:09:52 +0800 Subject: [PATCH 12/41] fix(lora): activate SGLang external model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLang 按 Hugging Face checkpoint 的 architecture 名称注册外部模型。本次将入口类名改为 Qwen3ForCausalLM，确保 Mixture-of-LoRA 实现覆盖内置 Qwen3，而不是启动后继续使用原模型。 Mixture-of-LoRA 是纯文本模型，启动子进程前会清理多模态外部处理器环境变量，并补充入口注册和环境隔离测试。H20 实机已验证预填充、CUDA Graph 解码及连续三次 router/expert 在线更新。 --- relax/backends/sglang/sglang_engine.py | 29 ++++++++++++++----- .../models/qwen3_mixture_lora/sglang/model.py | 8 +++-- tests/backends/sglang/test_mixture_lora.py | 16 ++++++++++ .../qwen3_mixture_lora/test_sglang_model.py | 5 ++++ 4 files changed, 48 insertions(+), 10 deletions(-) diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 7e80f49bd..e4a2a6dea 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -147,6 +147,20 @@ def _resolve_external_model_arch(package_name): return None +def _configure_external_model_environment(external_pkg: str, *, text_only: bool) -> str | None: + os.environ["SGLANG_EXTERNAL_MODEL_PACKAGE"] = external_pkg + if text_only: + os.environ.pop("SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE", None) + os.environ.pop("SGLANG_EXTERNAL_MM_MODEL_ARCH", None) + return None + + os.environ["SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE"] = external_pkg + arch = _resolve_external_model_arch(external_pkg) + if arch: + os.environ["SGLANG_EXTERNAL_MM_MODEL_ARCH"] = arch + return arch + + def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: multiprocessing.set_start_method("spawn", force=True) server_args.host = server_args.host.strip("[]") @@ -471,17 +485,18 @@ def _init_normal(self, server_args_dict): # Must be set before launch_server_process() spawns child process # (multiprocessing start_method='spawn'), because the child inherits # the parent's os.environ at spawn time. + mixture_lora_config = build_mixture_lora_config(self.args) external_pkg = configure_mixture_lora_external_model( - build_mixture_lora_config(self.args), + mixture_lora_config, getattr(self.args, "sglang_external_model_package", None), ) if external_pkg: - os.environ["SGLANG_EXTERNAL_MODEL_PACKAGE"] = external_pkg - os.environ["SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE"] = external_pkg - arch = _resolve_external_model_arch(external_pkg) - if arch: - os.environ["SGLANG_EXTERNAL_MM_MODEL_ARCH"] = arch - logger.info(f"Set SGLANG_EXTERNAL_MODEL_PACKAGE={external_pkg}, SGLANG_EXTERNAL_MM_MODEL_ARCH={arch}") + is_mixture_lora = mixture_lora_config is not None + arch = _configure_external_model_environment(external_pkg, text_only=is_mixture_lora) + if not is_mixture_lora: + logger.info(f"Set SGLANG_EXTERNAL_MODEL_PACKAGE={external_pkg}, SGLANG_EXTERNAL_MM_MODEL_ARCH={arch}") + else: + logger.info(f"Set SGLANG_EXTERNAL_MODEL_PACKAGE={external_pkg} for Mixture-of-LoRA") # Warm the OS page cache for this engine's HF checkpoint before the SGLang # subprocess mmaps the safetensors. Applies uniformly to rollout / genrm / diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index 1564340a7..5b9d0ce58 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -7,7 +7,7 @@ import torch from sglang.srt.distributed import get_tensor_model_parallel_world_size -from sglang.srt.models.qwen3 import Qwen3ForCausalLM +from sglang.srt.models.qwen3 import Qwen3ForCausalLM as SGLangQwen3ForCausalLM from torch import nn from torch.nn import functional as F @@ -154,7 +154,9 @@ def load_sglang_mixture_lora_weights( return loaded_names -class Qwen3MixtureLoRAForCausalLM(Qwen3ForCausalLM): +# SGLang registers external models by the checkpoint architecture name. Keeping +# this name replaces its built-in Qwen3 entry while preserving checkpoint metadata. +class Qwen3ForCausalLM(SGLangQwen3ForCausalLM): """Qwen3 external model that routes LoRA experts at attention projections.""" @@ -213,4 +215,4 @@ def base_weight_iterator(): load_sglang_mixture_lora_weights(self, mixture_weights) -EntryClass = Qwen3MixtureLoRAForCausalLM +EntryClass = Qwen3ForCausalLM diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py index d658c9bbe..88be4db97 100644 --- a/tests/backends/sglang/test_mixture_lora.py +++ b/tests/backends/sglang/test_mixture_lora.py @@ -4,6 +4,7 @@ import pytest +from relax.backends.sglang.sglang_engine import _configure_external_model_environment from relax.utils.mixture_lora import ( MixtureLoraConfig, configure_mixture_lora_external_model, @@ -53,3 +54,18 @@ def test_single_lora_does_not_enable_external_mixture_model(monkeypatch): def test_mixture_lora_rejects_conflicting_external_package(): with pytest.raises(ValueError, match="requires the Qwen3 external model package"): configure_mixture_lora_external_model(_config(), "custom.other.package") + + +def test_text_external_model_clears_multimodal_environment(monkeypatch): + monkeypatch.setenv("SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE", "stale.processor") + monkeypatch.setenv("SGLANG_EXTERNAL_MM_MODEL_ARCH", "StaleArchitecture") + + arch = _configure_external_model_environment( + "relax.models.qwen3_mixture_lora.sglang", + text_only=True, + ) + + assert arch is None + assert os.environ["SGLANG_EXTERNAL_MODEL_PACKAGE"] == "relax.models.qwen3_mixture_lora.sglang" + assert "SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE" not in os.environ + assert "SGLANG_EXTERNAL_MM_MODEL_ARCH" not in os.environ diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 6325c07db..5cd340543 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -7,6 +7,7 @@ from relax.backends.megatron.mixture_lora import MixtureLoRAAdapter from relax.models.qwen3_mixture_lora.sglang.model import ( + EntryClass, SGLangMixtureLoRA, attach_sglang_mixture_lora, load_sglang_mixture_lora_weights, @@ -26,6 +27,10 @@ def _config(): ) +def test_external_entry_class_overrides_the_checkpoint_architecture(): + assert EntryClass.__name__ == "Qwen3ForCausalLM" + + class _TupleLinear(nn.Module): def __init__(self, input_size: int, output_size: int): super().__init__() From db7a685465aff4d2211a4aaf36759f11a02a0139 Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 13:23:04 +0800 Subject: [PATCH 13/41] fix(lora): recover failed rollout updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mixture-of-LoRA 权重更新按暂停、发送和恢复三个阶段同步各训练 rank 的失败状态。任一阶段失败时恢复 rollout generation，并保持原 weight version 与 base 同步状态，避免服务停在暂停状态或发布半完成版本。 IPC 流水等待上一 chunk 失败时仍先完成既有 chunk barrier，防止其他 rank 永久等待。新增 engine 更新失败和 chunk barrier 失败注入测试；Task 25 回归 173 passed，全量 pre-commit 通过。 --- .../update_weight_from_tensor.py | 87 +++++++++++++++---- .../test_mixture_lora_weight_sync.py | 72 ++++++++++++++- 2 files changed, 140 insertions(+), 19 deletions(-) diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index b2570cd8d..ad055faf9 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -308,8 +308,12 @@ def _update_weights_mixture_lora(self) -> None: include_base = not self._mixture_lora_sync.base_sync_done all_engines = list(self.rollout_engines) rank = dist.get_rank() + quantization_restored = False - if rank == 0: + def pause_and_flush() -> None: + nonlocal quantization_restored + if rank != 0: + return ray.get([engine.pause_generation.remote() for engine in all_engines]) ray.get([engine.flush_cache.remote() for engine in all_engines]) if ( @@ -322,23 +326,27 @@ def _update_weights_mixture_lora(self) -> None: post_process_quantization=False, rollout_engines=all_engines, ) - dist.barrier(group=get_gloo_group()) - - local_weights = self.weights_getter() - base_chunks = self._hf_weight_iterator.get_hf_weight_chunks(local_weights) - mixture_chunks = self._mixture_lora_sync.get_weight_chunks(local_weights) - updates = iter_mixture_weight_updates( - base_chunks, - mixture_chunks, - include_base=include_base, - weight_version=next_weight_version, - ) - self._send_weight_update_stream(updates) + quantization_restored = True + + def send_weights() -> None: + local_weights = self.weights_getter() + base_chunks = self._hf_weight_iterator.get_hf_weight_chunks(local_weights) + mixture_chunks = self._mixture_lora_sync.get_weight_chunks(local_weights) + updates = iter_mixture_weight_updates( + base_chunks, + mixture_chunks, + include_base=include_base, + weight_version=next_weight_version, + ) + self._send_weight_update_stream(updates) - dist.barrier(group=get_gloo_group()) - if rank == 0: + def resume_generation(*, finish_quantization: bool) -> None: + if rank != 0: + return if ( - include_base + finish_quantization + and quantization_restored + and include_base and self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"] ): @@ -348,11 +356,48 @@ def _update_weights_mixture_lora(self) -> None: rollout_engines=all_engines, ) ray.get([engine.continue_generation.remote() for engine in all_engines]) - dist.barrier(group=get_gloo_group()) + + for phase_name, operation in (("pause and flush", pause_and_flush), ("send weights", send_weights)): + local_error, phase_failed = self._run_synchronized_weight_update_phase(operation) + if not phase_failed: + continue + + cleanup_error, _ = self._run_synchronized_weight_update_phase( + lambda: resume_generation(finish_quantization=True) + ) + if cleanup_error is not None: + logger.error( + "Failed to resume rollout generation after Mixture-of-LoRA update failure", + exc_info=(type(cleanup_error), cleanup_error, cleanup_error.__traceback__), + ) + if local_error is not None: + raise local_error + raise RuntimeError(f"Mixture-of-LoRA weight update phase {phase_name!r} failed on another rank") + + local_error, phase_failed = self._run_synchronized_weight_update_phase( + lambda: resume_generation(finish_quantization=True) + ) + if phase_failed: + if local_error is not None: + raise local_error + raise RuntimeError("Mixture-of-LoRA weight update phase 'resume generation' failed on another rank") self.weight_version = next_weight_version self._mixture_lora_sync.base_sync_done = True + @staticmethod + def _run_synchronized_weight_update_phase(operation): + """Run one update phase and share its failure state across ranks.""" + + local_error = None + try: + operation() + except Exception as error: + local_error = error + failed = torch.tensor([local_error is not None], dtype=torch.int32) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=get_gloo_group()) + return local_error, bool(failed.item()) + def _send_weight_update_stream(self, updates) -> None: """Pipeline conversion collectives with the preceding IPC request.""" @@ -360,12 +405,18 @@ def _send_weight_update_stream(self, updates) -> None: previous_tensors = None for named_tensors, weight_version in updates: refs, long_lived_tensors = self._send_hf_params(named_tensors, weight_version=weight_version) + previous_error = None if previous_refs: - ray.get(previous_refs) + try: + ray.get(previous_refs) + except Exception as error: + previous_error = error del previous_tensors previous_refs = refs previous_tensors = long_lived_tensors device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) + if previous_error is not None: + raise previous_error if previous_refs: ray.get(previous_refs) del previous_tensors diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 1a7fcf922..6cfc88480 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -265,7 +265,7 @@ def record_updates(updates): with ( patch("torch.distributed.get_rank", return_value=0), - patch("torch.distributed.barrier"), + patch("torch.distributed.all_reduce"), patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), patch("ray.get", side_effect=lambda refs: refs), ): @@ -286,3 +286,73 @@ def record_updates(updates): ) assert updater.weight_version == 2 assert updater._mixture_lora_sync.base_sync_done is True + + +def test_update_failure_resumes_generation_and_keeps_previous_version(): + from relax.backends.megatron.weight_update.update_weight_from_tensor import UpdateWeightFromTensor + + events = [] + + class RemoteMethod: + def __init__(self, name): + self.name = name + + def remote(self): + events.append(self.name) + return self.name + + engine = SimpleNamespace( + pause_generation=RemoteMethod("pause"), + flush_cache=RemoteMethod("flush"), + continue_generation=RemoteMethod("continue"), + ) + updater = UpdateWeightFromTensor.__new__(UpdateWeightFromTensor) + updater.weight_version = 4 + updater.rollout_engines = [engine] + updater.quantization_config = None + updater.weights_getter = MagicMock(return_value={"weight": torch.ones(1)}) + updater._hf_weight_iterator = MagicMock() + updater._hf_weight_iterator.get_hf_weight_chunks.return_value = iter([[("base", torch.ones(1))]]) + updater._mixture_lora_sync = SimpleNamespace(base_sync_done=False) + updater._mixture_lora_sync.get_weight_chunks = MagicMock(return_value=iter([[("router", torch.ones(1))]])) + + def fail_update(_updates): + events.append("update") + raise RuntimeError("engine rejected routed weights") + + updater._send_weight_update_stream = fail_update + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_reduce"), + patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + patch("ray.get", side_effect=lambda refs: refs), + pytest.raises(RuntimeError, match="engine rejected routed weights"), + ): + updater._update_weights_mixture_lora() + + assert events == ["pause", "flush", "update", "continue"] + assert updater.weight_version == 4 + assert updater._mixture_lora_sync.base_sync_done is False + + +def test_stream_failure_reaches_chunk_barrier_before_raising(): + from relax.backends.megatron.weight_update.update_weight_from_tensor import UpdateWeightFromTensor + + updater = UpdateWeightFromTensor.__new__(UpdateWeightFromTensor) + updater._send_hf_params = MagicMock( + side_effect=[(["first-ref"], ["first-tensor"]), (["second-ref"], ["second-tensor"])] + ) + + with ( + patch("ray.get", side_effect=RuntimeError("engine update failed")), + patch( + "relax.backends.megatron.weight_update.update_weight_from_tensor." + "device_utils.maybe_backend_barrier_on_weight_chunk" + ) as chunk_barrier, + patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + pytest.raises(RuntimeError, match="engine update failed"), + ): + updater._send_weight_update_stream([([("first", torch.ones(1))], None), ([("second", torch.ones(1))], 5)]) + + assert chunk_barrier.call_count == 2 From fad97d12cbce18fadc89b7affca84338a3b96e1b Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 13:38:47 +0800 Subject: [PATCH 14/41] feat(lora): support static context parallel routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 静态 CP 下按 site 在 context-parallel group 汇总 Top-K 入选次数、路由概率和有效 response token 数。每个 rank 只保留本地概率和的可导贡献，随后由现有 DP/CP 梯度归约合成完整序列的 balance loss。 普通 loss 与 per-token loss 分别使用正确的样本或全局 token 缩放，路由日志中的 aux loss 不随 CP 数放大。新增 CP=2、空 token rank、梯度和指标参考测试，并在启动阶段明确拒绝首版不支持的动态 CP。Task 25 完整回归通过，全量 pre-commit 通过。 --- relax/backends/megatron/mixture_lora.py | 59 ++++++++-- relax/backends/megatron/model.py | 4 + relax/utils/arguments.py | 2 + .../megatron/test_mixture_lora_distributed.py | 106 +++++++++++++++++- tests/utils/test_arguments_mixture_lora.py | 1 + 5 files changed, 164 insertions(+), 8 deletions(-) diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 2ce749607..4c748498d 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -61,6 +61,8 @@ class MixtureLoRARoutingContext: calculate_per_token_loss: bool objective_scale: float main_loss_backward_scale: torch.Tensor + context_parallel_group: Any = None + context_parallel_world_size: int = 1 is_dummy: bool = False records: dict[str, MixtureLoRARoutingRecord] = field(default_factory=dict, init=False) @@ -75,6 +77,10 @@ def __post_init__(self) -> None: raise ValueError("main_loss_backward_scale must be a one-element tensor") if not math.isfinite(self.objective_scale) or self.objective_scale < 0: raise ValueError("objective_scale must be finite and non-negative") + if self.context_parallel_world_size <= 0: + raise ValueError("context_parallel_world_size must be positive") + if self.context_parallel_world_size > 1 and self.context_parallel_group is None: + raise ValueError("context_parallel_group is required when context parallelism is enabled") def response_mask_for(self, x: torch.Tensor) -> torch.Tensor: """Align a batch-first mask with Megatron's activation layout.""" @@ -113,24 +119,27 @@ def attach_aux_loss( response_mask = self.response_mask_for(x) statistics = compute_routing_statistics(decision, response_mask) - # This is a microbatch-level F_e * P_e objective. It is not an - # average of independent per-token losses. - balance_loss = statistics.balance_loss - site_aux_loss = balance_loss * (config.aux_loss_coef / self.num_sites) + gradient_balance_loss, recorded_balance_loss, global_valid_token_count = _context_parallel_balance_losses( + statistics, + self.context_parallel_group, + self.context_parallel_world_size, + ) + site_aux_loss = gradient_balance_loss * (config.aux_loss_coef / self.num_sites) + recorded_site_aux_loss = recorded_balance_loss * (config.aux_loss_coef / self.num_sites) sample_or_token_count = ( - statistics.valid_token_count + global_valid_token_count if self.calculate_per_token_loss else statistics.valid_token_count.new_tensor(self.num_samples) ) objective_weight = sample_or_token_count * self.objective_scale aux_loss_payload = site_aux_loss * sample_or_token_count - objective_aux_loss = site_aux_loss * objective_weight + objective_aux_loss = recorded_site_aux_loss * objective_weight / self.context_parallel_world_size if record_statistics: key = (self.optimizer_step, self.microbatch_id, site_id) self.records[site_id] = MixtureLoRARoutingRecord( key=key, statistics=_detach_routing_statistics(statistics), - balance_loss=balance_loss.detach(), + balance_loss=recorded_balance_loss, aux_loss=objective_aux_loss.detach(), objective_weight=objective_weight.detach(), ) @@ -259,6 +268,42 @@ def _detach_routing_statistics(statistics: RoutingStatistics) -> RoutingStatisti ) +def _context_parallel_balance_losses( + statistics: RoutingStatistics, + context_parallel_group: Any, + context_parallel_world_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the local gradient contribution and global CP balance loss.""" + + if context_parallel_world_size == 1: + balance_loss = statistics.balance_loss + return balance_loss, balance_loss.detach(), statistics.valid_token_count + + num_experts = statistics.pre_topk_prob_sum.numel() + reduced = torch.cat( + ( + statistics.pre_topk_prob_sum.detach(), + statistics.selection_count.detach(), + statistics.valid_token_count.detach().reshape(1), + ) + ).clone() + torch.distributed.all_reduce(reduced, group=context_parallel_group) + global_pre_topk_prob_sum = reduced[:num_experts] + global_selection_count = reduced[num_experts : 2 * num_experts] + global_valid_token_count = reduced[-1] + denominator = global_valid_token_count.clamp_min(1) + has_valid_tokens = (global_valid_token_count > 0).to(reduced.dtype) + selection_share = global_selection_count / (denominator * statistics.top_k) + + # Keep only this rank's probability sum differentiable. CP/DDP gradient + # reduction adds the rank-local contributions into the full-sequence loss. + local_mean_prob = statistics.pre_topk_prob_sum / denominator + local_balance_loss = num_experts * torch.sum(selection_share * local_mean_prob) * has_valid_tokens + global_mean_prob = global_pre_topk_prob_sum / denominator + global_balance_loss = num_experts * torch.sum(selection_share * global_mean_prob) * has_valid_tokens + return local_balance_loss, global_balance_loss, global_valid_token_count + + def pack_mixture_lora_routing_records( contexts: list[MixtureLoRARoutingContext], site_ids: tuple[str, ...], diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 5c6a1c4fc..2d68fa1de 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -164,6 +164,10 @@ def _build_mixture_lora_routing_context( calculate_per_token_loss=args.calculate_per_token_loss, objective_scale=objective_scale, main_loss_backward_scale=main_loss_backward_scale.detach().clone(), + context_parallel_group=( + mpu.get_context_parallel_group() if mpu.get_context_parallel_world_size() > 1 else None + ), + context_parallel_world_size=mpu.get_context_parallel_world_size(), is_dummy=batch.get("__is_dummy__", False), ) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index e5d90906f..0b3699e6f 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -2818,6 +2818,8 @@ def _validate_lora_args(args) -> None: raise ValueError("Mixture-of-LoRA does not support --lora-merge-mode or --lora-adapter-mode.") if getattr(args, "fully_async", False): raise ValueError("Mixture-of-LoRA does not support --fully-async.") + if getattr(args, "dynamic_context_parallel", False): + raise ValueError("Mixture-of-LoRA currently supports static context parallelism only.") if not getattr(args, "colocate", False): raise ValueError("Mixture-of-LoRA requires --colocate.") diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index 7c0861986..e5008abd7 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -14,7 +14,12 @@ mixture_lora_metrics_from_packed_records, pack_mixture_lora_routing_records, ) -from relax.utils.mixture_lora import DenseRoutedLoRAExecutor, MixtureLoraConfig, route_topk +from relax.utils.mixture_lora import ( + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + compute_routing_statistics, + route_topk, +) def _config() -> MixtureLoraConfig: @@ -145,6 +150,105 @@ def test_routing_metrics_reduce_across_two_real_processes(tmp_path): mp.spawn(_distributed_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) +def _context_parallel_balance_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + config = _config() + full_logits = torch.tensor( + [ + [2.0, 0.5, -1.0], + [0.2, 1.7, -0.4], + [-0.3, 0.1, 2.1], + [1.2, -0.7, 0.4], + ] + ) + for calculate_per_token_loss in (False, True): + for empty_second_rank in (False, True): + full_mask = torch.tensor( + [1, 1, 0, 0] if empty_second_rank else [1, 1, 1, 0], + dtype=torch.bool, + ) + local_logits = full_logits.chunk(world_size)[rank].clone().requires_grad_(True) + local_mask = full_mask.chunk(world_size)[rank] + decision = route_topk(local_logits, config.top_k, config.temperature) + objective_scale = 1.0 if calculate_per_token_loss else float(world_size) + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=2 * int(calculate_per_token_loss) + int(empty_second_rank), + response_mask=local_mask.unsqueeze(0), + num_microbatches=1, + num_sites=1, + num_samples=1, + calculate_per_token_loss=calculate_per_token_loss, + objective_scale=objective_scale, + main_loss_backward_scale=torch.ones(1), + context_parallel_group=dist.group.WORLD, + context_parallel_world_size=world_size, + ) + output = local_logits.sum() * 0.0 + torch.zeros(2, 1, 1) + attached = context.attach_aux_loss( + output, + torch.zeros(2, 1, 1), + "layers.0.linear_qkv", + config, + decision, + ) + attached.sum().backward() + + reference_logits = full_logits.clone().requires_grad_(True) + reference_decision = route_topk(reference_logits, config.top_k, config.temperature) + reference_balance_loss = compute_routing_statistics(reference_decision, full_mask).balance_loss + objective_count = full_mask.sum() if calculate_per_token_loss else 1 + (reference_balance_loss * config.aux_loss_coef * objective_count * objective_scale).backward() + expected_gradient = reference_logits.grad.chunk(world_size)[rank] + torch.testing.assert_close(local_logits.grad, expected_gradient, atol=1e-6, rtol=1e-6) + + record = context.records["layers.0.linear_qkv"] + torch.testing.assert_close(record.balance_loss, reference_balance_loss.detach()) + torch.testing.assert_close( + record.aux_loss, + reference_balance_loss.detach() + * config.aux_loss_coef + * objective_count + * objective_scale + / world_size, + ) + packed = pack_mixture_lora_routing_records( + [context], + ("layers.0.linear_qkv",), + num_experts=config.num_experts, + top_k=config.top_k, + device=torch.device("cpu"), + ) + dist.all_reduce(packed) + metrics = mixture_lora_metrics_from_packed_records( + packed, + ("layers.0.linear_qkv",), + num_experts=config.num_experts, + top_k=config.top_k, + calculate_per_token_loss=calculate_per_token_loss, + data_parallel_world_size_with_cp=world_size, + ) + torch.testing.assert_close( + metrics["molora/aux_loss"], + (reference_balance_loss.detach() * config.aux_loss_coef).to(torch.float64), + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_context_parallel_balance_loss_matches_full_sequence_reference(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-cp-gloo-init'}" + mp.spawn(_context_parallel_balance_worker, args=(2, init_method), nprocs=2, join=True) + + def _reference_forward_and_backward( config: MixtureLoraConfig, x: torch.Tensor, diff --git a/tests/utils/test_arguments_mixture_lora.py b/tests/utils/test_arguments_mixture_lora.py index cb658434c..821695934 100644 --- a/tests/utils/test_arguments_mixture_lora.py +++ b/tests/utils/test_arguments_mixture_lora.py @@ -158,6 +158,7 @@ def test_mixture_configuration_rejects_invalid_values(arguments_module, override ({"lora_merge_mode": True}, "merge-mode"), ({"lora_adapter_mode": True}, "adapter-mode"), ({"fully_async": True}, "fully-async"), + ({"dynamic_context_parallel": True}, "static context parallelism"), ({"colocate": False}, "colocate"), ({"sglang_dp_size": 2}, "DP size 1"), ({"sglang_tp_size": 2}, "TP size 1"), From f607b5e8a00e6f7b7c1aea790e8d3ce1619f84de Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 13:44:05 +0800 Subject: [PATCH 15/41] test(lora): verify resumed optimizer step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 补充 Mixture-of-LoRA checkpoint 恢复后的下一步一致性测试。第一个 optimizer step 后保存 expert/router、AdamW、LR scheduler、iteration 和 Torch RNG，再比较连续训练与恢复训练的第二步结果。 测试启用 dropout 并确认 router 在第二步发生更新；恢复后的 loss、全部参数、优化器动量、scheduler 和 iteration 与不中断路径一致。Task 25 回归 176 passed，全量 pre-commit 通过。 --- tests/backends/megatron/test_mixture_lora.py | 72 ++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index e98b78dfc..c52e179c5 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -218,6 +218,78 @@ def test_mixture_lora_state_restores_output_and_validates_metadata(): target.load_state_dict(mismatched_state) +def test_checkpoint_restore_matches_uninterrupted_next_optimizer_step(): + torch.manual_seed(1234) + adapter = MixtureLoRAAdapter( + _config(), + "linear_qkv", + 4, + 5, + dropout=0.25, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + adapter.experts.lora_B.normal_(mean=0.0, std=0.2) + optimizer = torch.optim.AdamW(adapter.parameters(), lr=0.03, weight_decay=0.01) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.8) + batches = ( + torch.linspace(-1.0, 1.0, steps=24).reshape(3, 2, 4), + torch.linspace(0.8, -0.6, steps=24).reshape(3, 2, 4), + ) + + def train_step(x): + optimizer.zero_grad() + loss = adapter(x).square().mean() + loss.backward() + optimizer.step() + scheduler.step() + return loss.detach() + + train_step(batches[0]) + checkpoint = { + "model": copy.deepcopy(adapter.state_dict()), + "optimizer": copy.deepcopy(optimizer.state_dict()), + "scheduler": copy.deepcopy(scheduler.state_dict()), + "iteration": 1, + "rng": torch.get_rng_state().clone(), + } + + expected_loss = train_step(batches[1]) + expected_parameters = {name: parameter.detach().clone() for name, parameter in adapter.named_parameters()} + expected_optimizer_state = { + name: { + key: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in optimizer.state[parameter].items() + } + for name, parameter in adapter.named_parameters() + } + expected_scheduler_state = copy.deepcopy(scheduler.state_dict()) + expected_iteration = checkpoint["iteration"] + 1 + assert not torch.equal(expected_parameters["router.weight"], checkpoint["model"]["router.weight"]) + + torch.rand(17) + adapter.load_state_dict(checkpoint["model"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + torch.set_rng_state(checkpoint["rng"]) + restored_iteration = checkpoint["iteration"] + restored_loss = train_step(batches[1]) + restored_iteration += 1 + + torch.testing.assert_close(restored_loss, expected_loss) + assert restored_iteration == expected_iteration + assert scheduler.state_dict() == expected_scheduler_state + for name, parameter in adapter.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) + for key, expected_value in expected_optimizer_state[name].items(): + restored_value = optimizer.state[parameter][key] + if torch.is_tensor(expected_value): + torch.testing.assert_close(restored_value, expected_value) + else: + assert restored_value == expected_value + + @pytest.mark.parametrize( ("calculate_per_token_loss", "is_dummy", "explicit_loss_scale", "expected"), [ From 707c41d3ab0d03ea8d6bdc10e32ccd0b22dcc17c Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Sun, 9 Aug 2026 13:54:01 +0800 Subject: [PATCH 16/41] feat(lora): add Mixture training recipe and guides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 Qwen3-4B、DAPO math、GRPO、8 卡 colocate 的 Mixture-of-LoRA recipe，显式配置 N=4、R=16、Top-K、temperature 和 balance coefficient。训练侧使用 TP=2/SP，rollout 启动八个独立的 TP=1/DP=1 SGLang engine，并保留命令行覆盖参数。 新增中英文使用文档和 VitePress 导航，说明启动条件、支持范围、路由公式、逐 site 指标、权重同步和 checkpoint 恢复。Bash 语法检查及全量 pre-commit 通过。 --- docs/.vitepress/config.mts | 6 +- docs/en/guide/mixture-lora.md | 94 +++++++++++ docs/zh/guide/mixture-lora.md | 94 +++++++++++ .../text/run-qwen3-4B-mixture-lora-8xgpu.sh | 149 ++++++++++++++++++ 4 files changed, 341 insertions(+), 2 deletions(-) create mode 100644 docs/en/guide/mixture-lora.md create mode 100644 docs/zh/guide/mixture-lora.md create mode 100755 scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index 33e7188e5..2bb3db537 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -271,7 +271,8 @@ export default defineConfig({ { text: 'Metrics Service', link: '/en/guide/metrics-service-detailed' }, { text: 'Notification System', link: '/en/guide/notification-system' }, { text: 'Update Weights Pipeline', link: '/en/guide/update-weights-pipeline' }, - { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' } + { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' }, + { text: 'Mixture-of-LoRA RL Training', link: '/en/guide/mixture-lora' } ] }, { @@ -382,7 +383,8 @@ export default defineConfig({ { text: 'Metrics 服务', link: '/zh/guide/metrics-service-detailed' }, { text: '通知系统', link: '/zh/guide/notification-system' }, { text: '权重更新流水线优化', link: '/zh/guide/update-weights-pipeline' }, - { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' } + { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' }, + { text: 'Mixture-of-LoRA RL 训练', link: '/zh/guide/mixture-lora' } ] }, { diff --git a/docs/en/guide/mixture-lora.md b/docs/en/guide/mixture-lora.md new file mode 100644 index 000000000..e69d019fa --- /dev/null +++ b/docs/en/guide/mixture-lora.md @@ -0,0 +1,94 @@ +# Mixture-of-LoRA RL Training + +Mixture-of-LoRA keeps the base model frozen and trains multiple LoRA experts at each target projection. A learned token-level router selects `K` experts and combines their outputs with normalized Top-K weights. + +Use this path when one LoRA adapter does not provide enough capacity. Existing single LoRA behavior remains active when `--lora-num-experts 1` is used or omitted. + +## Configuration + +The Mixture path is enabled when both `--lora-rank` is positive and `--lora-num-experts` is greater than one. + +```bash +--lora-rank 16 +--lora-alpha 32 +--lora-target-modules linear_qkv linear_proj +--lora-dropout 0.0 +--lora-num-experts 4 +--lora-router-top-k 2 +--lora-router-temperature 1.0 +--lora-router-aux-loss-coef 0.01 +``` + +| Option | Meaning | +| --- | --- | +| `--lora-num-experts` | Number of LoRA experts at each routed projection. Values greater than one enable Mixture-of-LoRA. | +| `--lora-router-top-k` | Number of experts selected for each token. It must satisfy `1 <= K <= N`. | +| `--lora-router-temperature` | Temperature applied before the router softmax. | +| `--lora-router-aux-loss-coef` | Coefficient for the per-site router balance loss. Use `0` to disable its gradient while retaining routing metrics. | + +The three router options must be provided explicitly for `N > 1`. Mixture mode does not use `--lora-merge-mode` or `--lora-adapter-mode`. + +## Supported Setup + +The first implementation supports: + +- dense Qwen3 models; +- `linear_qkv` and `linear_proj` attention targets; +- Megatron data, tensor, sequence, pipeline, and static context parallelism; +- synchronous colocate training; +- multiple independent SGLang engines where each engine uses TP=1 and DP=1; +- native Megatron distributed checkpoints. + +Fully async training, dynamic context parallelism, VLMs, MoE base models, MLP targets, and SGLang engines with internal TP or DP greater than one are rejected during startup. + +## Routing and Balance Loss + +For each token, the router computes FP32 probabilities over `N` experts, selects the configured Top-K entries, and renormalizes the selected probabilities to sum to one. Training and SGLang rollout call the same routing function and use equivalent dense expert equations; the Megatron executor additionally handles its TP/SP collectives. + +The balance objective is calculated independently for each routed site and averaged over all sites: + +```text +L_balance = N * sum_e(F_e * P_e) +F_e = selection_count_e / (valid_response_tokens * K) +P_e = mean pre-Top-K router probability for expert e +``` + +Prompt, padding, and dummy tokens do not enter the balance loss or routing metrics. Under static context parallelism, the counts and probability sums are combined over the CP group before the objective is formed. + +## Routing Metrics + +Metrics are emitted under `molora//...` and `molora/global/...`: + +- `expert__pre_topk_mean_prob`; +- `expert__post_topk_mean_weight`; +- `expert__selection_share`; +- `expert__top1_fraction`; +- `pre_topk_normalized_entropy` and `post_topk_normalized_entropy`; +- `balance_loss` per site and `molora/aux_loss` globally. + +Use the per-site post-Top-K weights, selection shares, and entropy to detect collapse. Global metrics are useful summaries but can hide a collapsed layer. + +## Rollout Weight Updates + +Relax starts the Qwen3 SGLang external model automatically when Mixture mode is enabled. The first colocate update sends the frozen base plus all expert and router tensors. Later updates send all current expert and router tensors without resending the base. A weight version is published only after the final routed tensor is accepted. + +If an update fails, Relax resumes generation, keeps the previous weight version, and reports the failed update instead of silently serving a partial policy. + +## Checkpoints + +Expert and router tensors are ordinary model parameters in the native Megatron distributed checkpoint. The same checkpoint also restores optimizer, scheduler, iteration, and RNG state. Mixture mode does not create a separate HF PEFT adapter export. + +Resume by launching the same recipe with `--load` and `--save` pointing to the existing output directory. The saved Mixture metadata is checked against the current expert count, rank, Top-K, temperature, coefficient, alpha, target modules, dtype, and site dimensions before tensors are loaded. + +## Qwen3-4B DAPO Recipe + +The reference script runs Qwen3-4B GRPO for 200 rollouts on eight colocated GPUs: + +```bash +MODEL_PATH=/path/to/Qwen3-4B \ +PROMPT_DATA=/path/to/dapo-math-17k.jsonl \ +OUTPUT_DIR=/path/to/qwen3-4b-mixture-lora \ +bash scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh +``` + +The actor uses TP=2 with sequence parallelism. The rollout allocation creates eight independent one-GPU SGLang engines. Environment variables such as `NUM_ROLLOUT`, `LORA_NUM_EXPERTS`, `LORA_RANK`, and `LORA_ROUTER_TOP_K` can override the recipe values. Additional Relax arguments can be appended to the command. diff --git a/docs/zh/guide/mixture-lora.md b/docs/zh/guide/mixture-lora.md new file mode 100644 index 000000000..6d14503a1 --- /dev/null +++ b/docs/zh/guide/mixture-lora.md @@ -0,0 +1,94 @@ +# Mixture-of-LoRA RL 训练 + +Mixture-of-LoRA 冻结基座模型，在每个目标投影层训练多个 LoRA expert。可训练的 token-level router 为每个 token 选择 `K` 个 expert，再用归一化后的 Top-K 权重组合这些 expert 的输出。 + +当单个 LoRA adapter 容量不足时可以使用这条路径。`--lora-num-experts 1` 或不传该参数时，Relax 继续使用现有单 LoRA 实现。 + +## 参数配置 + +`--lora-rank` 大于零且 `--lora-num-experts` 大于一时启用 Mixture 路径。 + +```bash +--lora-rank 16 +--lora-alpha 32 +--lora-target-modules linear_qkv linear_proj +--lora-dropout 0.0 +--lora-num-experts 4 +--lora-router-top-k 2 +--lora-router-temperature 1.0 +--lora-router-aux-loss-coef 0.01 +``` + +| 参数 | 含义 | +| --- | --- | +| `--lora-num-experts` | 每个目标投影层的 LoRA expert 数量。值大于一时启用 Mixture-of-LoRA。 | +| `--lora-router-top-k` | 每个 token 选中的 expert 数量，必须满足 `1 <= K <= N`。 | +| `--lora-router-temperature` | router softmax 使用的温度。 | +| `--lora-router-aux-loss-coef` | 逐 site balance loss 的系数。设为 `0` 时不产生这部分梯度，但仍输出路由指标。 | + +`N > 1` 时必须明确提供三个 router 参数。Mixture 模式不使用 `--lora-merge-mode` 或 `--lora-adapter-mode`。 + +## 支持范围 + +首版支持： + +- Qwen3 dense 模型； +- attention 的 `linear_qkv` 和 `linear_proj`； +- Megatron data、tensor、sequence、pipeline 和静态 context parallel； +- 同步 colocate 训练； +- 多个独立 SGLang engine，每个 engine 内部使用 TP=1、DP=1； +- Megatron 原生 distributed checkpoint。 + +启动时会拒绝 fully async、动态 context parallel、VLM、MoE 基座、MLP target，以及内部 TP 或 DP 大于一的 SGLang engine。 + +## 路由与 Balance Loss + +router 为每个 token 使用 FP32 计算 `N` 个 expert 的概率，选出 Top-K 后重新归一化，使选中权重之和为一。训练端和 SGLang rollout 端调用同一套路由函数并使用等价的 dense expert 计算；Megatron 执行器还负责 TP/SP collective。 + +每个 routed site 单独计算 balance objective，再对所有 site 取平均： + +```text +L_balance = N * sum_e(F_e * P_e) +F_e = selection_count_e / (valid_response_tokens * K) +P_e = expert e 在 Top-K 前的平均 router 概率 +``` + +Prompt、padding 和 dummy token 不参与 balance loss 与路由指标。静态 context parallel 下会先在 CP group 中汇总入选次数、概率和与有效 token 数，再计算完整序列的 objective。 + +## 路由指标 + +指标名称使用 `molora//...` 和 `molora/global/...`： + +- `expert__pre_topk_mean_prob`； +- `expert__post_topk_mean_weight`； +- `expert__selection_share`； +- `expert__top1_fraction`； +- `pre_topk_normalized_entropy` 和 `post_topk_normalized_entropy`； +- 每个 site 的 `balance_loss` 与全局 `molora/aux_loss`。 + +判断路由塌缩时应查看每个 site 的 Top-K 后平均权重、选择份额和熵。全局指标只用于汇总，可能掩盖个别层的塌缩。 + +## Rollout 权重更新 + +启用 Mixture 后，Relax 会自动启动 Qwen3 SGLang external model。第一次 colocate 更新发送冻结的基座参数以及全部 expert/router 参数；后续更新只发送当前全部 expert/router 参数。最后一组 routed tensor 加载成功后才发布新的 weight version。 + +更新失败时会恢复 generation、保留原 weight version，并报告错误，不会静默使用只加载了一部分的新策略。 + +## Checkpoint + +Expert 和 router 是 Megatron 原生 distributed checkpoint 中的普通模型参数。Optimizer、scheduler、iteration 和 RNG 状态沿用同一个 checkpoint 恢复。Mixture 模式不会额外导出一份 HF PEFT adapter。 + +恢复训练时使用同一 recipe，并让 `--load` 和 `--save` 指向已有输出目录。加载 tensor 前会检查 checkpoint 中的 expert 数、rank、Top-K、temperature、coefficient、alpha、target module、dtype 和 site 维度是否与当前配置一致。 + +## Qwen3-4B DAPO Recipe + +参考脚本使用八张 colocate GPU，让 Qwen3-4B 在 DAPO math 上运行 200 个 rollout 的 GRPO： + +```bash +MODEL_PATH=/path/to/Qwen3-4B \ +PROMPT_DATA=/path/to/dapo-math-17k.jsonl \ +OUTPUT_DIR=/path/to/qwen3-4b-mixture-lora \ +bash scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh +``` + +Actor 使用 TP=2 和 sequence parallel。Rollout 资源会建立八个独立的单卡 SGLang engine。可以通过 `NUM_ROLLOUT`、`LORA_NUM_EXPERTS`、`LORA_RANK`、`LORA_ROUTER_TOP_K` 等环境变量覆盖 recipe 中的值，也可以在命令末尾继续追加 Relax 参数。 diff --git a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh new file mode 100755 index 000000000..3e68787d8 --- /dev/null +++ b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh @@ -0,0 +1,149 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3-4B Mixture-of-LoRA GRPO on DAPO math with 8 colocated GPUs. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +if [[ -z "${RELAX_ENTRYPOINT_MODE:-}" ]]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-4B.sh" + +now="$(date '+%Y-%m-%d-%H-%M-%S')" +PROJECT_NAME="${PROJECT_NAME:-Relax/dev/qwen3-4b-mixture-lora}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" +MODEL_PATH="${MODEL_PATH:-${MODEL_DIR}/Qwen3-4B}" +PROMPT_DATA="${PROMPT_DATA:-${DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl}" +OUTPUT_DIR="${OUTPUT_DIR:-${EXP_DIR}/Qwen3-4B_mixture_lora_8xgpu}" +NUM_ROLLOUT="${NUM_ROLLOUT:-200}" + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_PATH}" + --ref-load "${MODEL_PATH}" + --megatron-to-hf-mode bridge + --warm-hf-checkpoint-page-cache + --load "${OUTPUT_DIR}" + --save "${OUTPUT_DIR}" + --save-interval "${SAVE_INTERVAL:-50}" +) + +LORA_ARGS=( + --lora-rank "${LORA_RANK:-16}" + --lora-alpha "${LORA_ALPHA:-32}" + --lora-target-modules linear_qkv linear_proj + --lora-dropout 0.0 + --lora-num-experts "${LORA_NUM_EXPERTS:-4}" + --lora-router-top-k "${LORA_ROUTER_TOP_K:-2}" + --lora-router-temperature "${LORA_ROUTER_TEMPERATURE:-1.0}" + --lora-router-aux-loss-coef "${LORA_ROUTER_AUX_LOSS_COEF:-0.01}" +) + +ROLLOUT_ARGS=( + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type dapo + --reward-key score + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-16}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-8}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-8192}" + --rollout-temperature 1.0 + --global-batch-size "${GLOBAL_BATCH_SIZE:-128}" + --balance-data + --use-fault-tolerance +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.0 + --kl-loss-type low_var_kl + --entropy-coef 0.0 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr "${LR:-1e-5}" + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --initial-loss-scale 32768 + --min-loss-scale 1 + --use-precision-aware-optimizer + --no-store-param-remainders +) + +PARALLEL_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --calculate-per-token-loss + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-9216}" + --log-probs-max-tokens-per-gpu "${LOG_PROBS_MAX_TOKENS_PER_GPU:-30720}" +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static "${SGLANG_MEM_FRACTION_STATIC:-0.7}" +) + +LOGGING_ARGS=( + --use-metrics-service + --tb-project-name "${PROJECT_NAME}" + --tb-experiment-name "qwen3-4b-mixture-lora-${now}" +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --skip-eval-before-train +) + +mkdir -p "${OUTPUT_DIR}" "${OUTPUT_DIR}/logs" +if [[ -z "${RUNTIME_ENV_JSON:-}" ]]; then + RUNTIME_ENV_JSON='{}' +fi + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="${RAY_ADDRESS:-http://127.0.0.1:8265}" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 8], "rollout": [1, 8]}' \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --fp16 \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${LOGGING_ARGS[@]}" \ + "${PARALLEL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${LORA_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "$@" 2>&1 | tee "${OUTPUT_DIR}/logs/qwen3-4b-mixture-lora-${now}.log" From 95834995f5470cdefa932246360da65fefd1323b Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Sun, 9 Aug 2026 19:45:40 +0800 Subject: [PATCH 17/41] fix(mixture-lora): support pipeline-parallel synchronization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary 修复 Mixture-of-LoRA 在 PP/VPP 训练中的 site 标识、权重转换和 rollout 同步问题，使各 pipeline stage 的 expert 与 router 能稳定聚合并更新到 SGLang。 Changes - 使用全局层偏移生成稳定 site_id，并传递 virtual pipeline stage。 - base 权重仅由参数所属 stage 执行 Bridge 转换，转换后在 PP group 广播；修正转换耗时统计。 - 汇总各 PP stage 的 Mixture 参数元数据，补充 selective CPU offload 和 SGLang base CPU backup。 - 增加 rank/TP 整除、VLM/MoE 拒绝和 Bridge 可选依赖保护，并增强同步失败日志。 - 补充 PP site、权重同步、offload、参数校验和 rollout 回归测试。 Verification - Task 25、两卡分布式及单 LoRA 兼容回归：192 passed。 - Bridge 与权重同步定向回归：33 passed。 - pre-commit：除独立执行的 gitleaks hook 外，其余 hook 全部通过。 - Gitleaks v8.24.2 staged scan：no leaks found。 - A800 8 卡 PP=2、TP=2、SP、colocate 两步训练完成，包含两次训练后权重同步和 iteration 0/1 checkpoint。 --- relax/backends/megatron/arguments.py | 6 + relax/backends/megatron/mixture_lora.py | 25 +++- relax/backends/megatron/model_provider.py | 2 +- .../backends/megatron/weight_update/common.py | 2 +- .../hf_weight_iterator_bridge.py | 31 ++--- .../weight_update/mixture_lora_sync.py | 42 +++++-- .../update_weight_from_tensor.py | 5 + relax/backends/sglang/sglang_engine.py | 6 + relax/utils/arguments.py | 6 + relax/utils/megatron_bridge_utils.py | 34 ++++++ .../text/run-qwen3-4B-mixture-lora-8xgpu.sh | 2 +- tests/backends/megatron/test_mixture_lora.py | 56 +++++++++ .../megatron/test_mixture_lora_arguments.py | 23 ++++ .../test_mixture_lora_weight_sync.py | 108 ++++++++++++++++++ tests/backends/sglang/test_mixture_lora.py | 34 +++++- tests/utils/test_arguments_mixture_lora.py | 2 + 16 files changed, 352 insertions(+), 32 deletions(-) create mode 100644 tests/backends/megatron/test_mixture_lora_arguments.py diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index 27aaed382..bcb2c7fda 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -215,6 +215,9 @@ def equal(x, y): # RoPE kernel cannot handle this and produces numerically different results from the # unfused HF/SGLang implementation, causing training-inference log-prob mismatch. is_multimodal = hasattr(hf_config, "text_config") or hasattr(hf_config, "thinker_config") + mixture_lora_enabled = getattr(args, "lora_num_experts", 1) > 1 + if mixture_lora_enabled and is_multimodal: + errors.append("Mixture-of-LoRA currently supports text-only base models; multimodal models are unsupported.") if is_multimodal and getattr(args, "apply_rope_fusion", False): errors.append( "Multimodal models use multi-axis RoPE (list of tensors) which is incompatible " @@ -230,6 +233,9 @@ def equal(x, y): if hasattr(hf_config, "text_config"): hf_config = hf_config.text_config + if mixture_lora_enabled and _is_moe_config(hf_config): + errors.append("Mixture-of-LoRA currently supports dense base models; MoE base models are unsupported.") + validate_dense_ffn = not _is_moe_config(hf_config) or _has_dense_moe_layers(args) for hf_config_name, megatron_config_name, compare_fn in ( diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 4c748498d..830f2d952 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -3,6 +3,7 @@ """Megatron model modules and Bridge injection for Mixture-of-LoRA.""" import math +import re from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field @@ -893,7 +894,7 @@ def sharded_state_dict( return sharded_state -def build_mixture_lora_peft(config: MixtureLoraConfig, dropout: float): +def build_mixture_lora_peft(config: MixtureLoraConfig, dropout: float, vp_stage: int | None = None): """Build a Bridge PEFT object that injects routed adapters at matched sites.""" @@ -929,6 +930,11 @@ def transform( raise RuntimeError("Mixture-of-LoRA PEFT is missing its configuration") _, full_name = match attributes = get_adapter_attributes_from_linear(module) + site_id = _global_mixture_lora_site_id( + full_name, + getattr(module, "config", None), + vp_stage, + ) tp_world_size = parallel_state.get_tensor_model_parallel_world_size() tp_group = getattr(module, "tp_group", None) if tp_world_size > 1 and tp_group is None: @@ -936,7 +942,7 @@ def transform( return MixtureParallelLinearAdapter( module, self.mixture_config, - full_name, + site_id, attributes.in_features, attributes.out_features, dropout=self.dropout, @@ -972,3 +978,18 @@ def transform( "mixture_lora_metrics_from_packed_records", "pack_mixture_lora_routing_records", ] +_DECODER_LAYER_SITE_PATTERN = re.compile(r"^(?P.*decoder\.layers\.)(?P\d+)(?P\..+)$") + + +def _global_mixture_lora_site_id(site_id: str, transformer_config: Any, vp_stage: int | None) -> str: + """Convert a pipeline-local decoder layer name into its global name.""" + + match = _DECODER_LAYER_SITE_PATTERN.fullmatch(site_id) + if match is None or transformer_config is None: + return site_id + + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + layer_offset = get_transformer_layer_offset(transformer_config, vp_stage=vp_stage) + global_layer = int(match.group("layer")) + layer_offset + return f"{match.group('prefix')}{global_layer}{match.group('suffix')}" diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 10a9a1a2b..19479e1a7 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -489,7 +489,7 @@ def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwarg mixture_config = build_mixture_lora_config(args) if mixture_config is None: raise RuntimeError("Mixture-of-LoRA is enabled but its validated configuration is missing") - peft = build_mixture_lora_peft(mixture_config, args.lora_dropout) + peft = build_mixture_lora_peft(mixture_config, args.lora_dropout, vp_stage=vp_stage) else: peft = build_lora_peft(args) model = peft(model, training=True) diff --git a/relax/backends/megatron/weight_update/common.py b/relax/backends/megatron/weight_update/common.py index f39a87a82..4ce3e16d2 100644 --- a/relax/backends/megatron/weight_update/common.py +++ b/relax/backends/megatron/weight_update/common.py @@ -213,7 +213,7 @@ def _maybe_get_cpu_backup(x: torch.Tensor) -> torch.Tensor: # frees a param's GPU storage via storage().resize_(0), it stashes the CPU copy # on the tensor as ``_relax_cpu_offload_data``. If the GPU storage is empty, read # from that CPU copy instead of touching the now-invalid CUDA storage. - if getattr(x, "_relax_cpu_offload_data", None) is not None and x.storage().size() == 0: + if getattr(x, "_relax_cpu_offload_data", None) is not None and x.untyped_storage().nbytes() == 0: return x._relax_cpu_offload_data # torch_memory_saver path: only usable when its LD_PRELOAD hook is active; diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py index 7d2e37e03..43c133652 100644 --- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py +++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py @@ -110,30 +110,33 @@ def _iter_hf_params(self, megatron_local_weights): yield from results del all_converted, results - # --- Non-expert weights: original path --- + # --- Non-expert weights --- for bucket_infos in self._non_expert_buckets: - t_b0 = time.monotonic() - params = _load_and_broadcast( + params = _load_to_gpu( bucket_infos, megatron_local_weights, self._vanilla_key_map, device, rank, merge_fn=merge_fn ) - t_b1 = time.monotonic() - t_bcast_total += t_b1 - t_b0 - + all_converted = [] for info, param in zip(bucket_infos, params, strict=True): t_g0 = time.monotonic() gathered = all_gather_param(self.args, info.name, param) t_g1 = time.monotonic() t_gather_total += t_g1 - t_g0 - - converted = self._bridge_converter.convert(info.name, gathered) - t_convert_total += time.monotonic() - t_g1 - - param_count += len(converted) - yield from converted - del gathered, converted - + if rank == info.src_rank: + t_c0 = time.monotonic() + all_converted.append(self._bridge_converter.convert(info.name, gathered)) + t_convert_total += time.monotonic() - t_c0 + else: + all_converted.append(None) + del gathered del params + t_b0 = time.monotonic() + results = _broadcast_converted_bucket(bucket_infos, all_converted, device) + t_bcast_total += time.monotonic() - t_b0 + param_count += len(results) + yield from results + del all_converted, results + if rank == 0: logger.info( "[Bridge Fast] params=%d | bcast=%.1fs | tp_gather=%.1fs | convert=%.1fs | total=%.1fs", diff --git a/relax/backends/megatron/weight_update/mixture_lora_sync.py b/relax/backends/megatron/weight_update/mixture_lora_sync.py index 33d91e3bc..9a6392f61 100644 --- a/relax/backends/megatron/weight_update/mixture_lora_sync.py +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -32,6 +32,30 @@ class MixtureLoraParamInfo: weight_key: str +def _gather_pipeline_param_infos( + local_infos: dict[str, MixtureLoraParamInfo], + *, + pipeline_group, + pipeline_world_size: int, +) -> dict[str, MixtureLoraParamInfo]: + """Collect parameter metadata from every stage in one PP group.""" + + gathered_infos = [None] * pipeline_world_size + dist.all_gather_object( + gathered_infos, + (dist.get_rank(), local_infos), + group=pipeline_group, + ) + merged_infos: dict[str, MixtureLoraParamInfo] = {} + for _, stage_infos in gathered_infos: + for name, info in stage_infos.items(): + previous = merged_infos.get(name) + if previous is not None and previous != info: + raise ValueError(f"Conflicting Mixture-of-LoRA metadata for {name}") + merged_infos[name] = info + return merged_infos + + def _qwen3_attention_dimensions(args: Namespace) -> tuple[int, int, int]: head_dim = getattr(args, "kv_channels", None) if head_dim is None: @@ -218,19 +242,13 @@ def _build_param_infos(self) -> tuple[MixtureLoraParamInfo, ...]: weight_key=weight_key, ) - if mpu.get_pipeline_model_parallel_world_size() > 1: - gathered_infos = [None] * mpu.get_pipeline_model_parallel_world_size() - dist.all_gather_object( - (rank, local_infos), - object_list=gathered_infos, - group=mpu.get_pipeline_model_parallel_group(), + pipeline_world_size = mpu.get_pipeline_model_parallel_world_size() + if pipeline_world_size > 1: + local_infos = _gather_pipeline_param_infos( + local_infos, + pipeline_group=mpu.get_pipeline_model_parallel_group(), + pipeline_world_size=pipeline_world_size, ) - for _, stage_infos in gathered_infos: - for name, info in stage_infos.items(): - previous = local_infos.get(name) - if previous is not None and previous != info: - raise ValueError(f"Conflicting Mixture-of-LoRA metadata for {name}") - local_infos[name] = info return tuple(local_infos[name] for name in sorted(local_infos)) def get_weight_chunks( diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index ad055faf9..680e10d4c 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -394,6 +394,11 @@ def _run_synchronized_weight_update_phase(operation): operation() except Exception as error: local_error = error + logger.error( + "Weight update phase failed on rank %d before failure synchronization", + dist.get_rank(), + exc_info=(type(error), error, error.__traceback__), + ) failed = torch.tensor([local_error is not None], dtype=torch.int32) dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=get_gloo_group()) return local_error, bool(failed.item()) diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index e4a2a6dea..2962d4075 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -28,6 +28,7 @@ build_mixture_lora_config, convert_megatron_to_hf_target_modules, is_lora_enabled, + is_mixture_lora_enabled, ) from relax.utils.mixture_lora import configure_mixture_lora_external_model @@ -1377,6 +1378,11 @@ def _compute_server_args( # the first wake (other modes re-push full base every step and never notice). kwargs["enable_weights_cpu_backup"] = True + # Mixture sync also sends the frozen base only once. Keep its CPU copy alive + # while colocate sleep releases the GPU weight pages between rollout phases. + if is_mixture_lora_enabled(args): + kwargs["enable_weights_cpu_backup"] = True + if worker_type == "prefill": kwargs["disaggregation_mode"] = "prefill" kwargs["load_balance_method"] = "round_robin" diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 0b3699e6f..833004d68 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -2809,6 +2809,12 @@ def _validate_lora_args(args) -> None: else: if getattr(args, "lora_rank", 0) <= 0: raise ValueError("--lora-num-experts greater than 1 requires --lora-rank greater than 0.") + training_tp_size = getattr(args, "tensor_model_parallel_size", 1) + if args.lora_rank % training_tp_size != 0: + raise ValueError( + f"Mixture-of-LoRA requires --lora-rank to be divisible by the Megatron TP size; " + f"got rank {args.lora_rank} and TP size {training_tp_size}." + ) missing = [option for attribute, option in mixture_fields if getattr(args, attribute, None) is None] if missing: raise ValueError( diff --git a/relax/utils/megatron_bridge_utils.py b/relax/utils/megatron_bridge_utils.py index f1d96d477..0a0de7e6b 100644 --- a/relax/utils/megatron_bridge_utils.py +++ b/relax/utils/megatron_bridge_utils.py @@ -3,6 +3,7 @@ from contextlib import contextmanager from relax.utils.logging_utils import get_logger +from relax.utils.megatron_peft_utils import is_mixture_lora_param logger = get_logger(__name__) @@ -57,8 +58,40 @@ def _with_progress_tracking(self, tasks, description: str, show_progress: bool = _patch_progress_tracking_show_elapsed() +_bridge_mixture_filter_active: contextvars.ContextVar[bool] = contextvars.ContextVar( + "relax_bridge_mixture_filter_active", default=False +) +_bridge_mixture_filter_patched = False + + +def _ensure_bridge_mixture_filter_patched(): + """Teach Bridge base conversion to treat routed parameters as adapters.""" + + global _bridge_mixture_filter_patched + if _bridge_mixture_filter_patched: + return + + try: + from megatron.bridge.models.conversion.peft_bridge import MegatronPeftBridge + except ImportError: + logger.warning("Megatron-Bridge PEFT helpers are unavailable; skipping the Mixture-of-LoRA parameter filter.") + return + + original = MegatronPeftBridge._is_adapter_param_name + + def _is_adapter_param_name(self, param_name: str) -> bool: + return original(self, param_name) or ( + _bridge_mixture_filter_active.get() and is_mixture_lora_param(param_name) + ) + + MegatronPeftBridge._is_adapter_param_name = _is_adapter_param_name + _bridge_mixture_filter_patched = True + + @contextmanager def patch_megatron_model(model): + _ensure_bridge_mixture_filter_patched() + mixture_filter_token = _bridge_mixture_filter_active.set(True) unwrapped_model = unwrap_model(model)[0] model_config = unwrapped_model.config attribute_was_added = False @@ -69,6 +102,7 @@ def patch_megatron_model(model): try: yield finally: + _bridge_mixture_filter_active.reset(mixture_filter_token) if attribute_was_added: delattr(model_config, "share_embeddings_and_output_weights") diff --git a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh index 3e68787d8..cb4358670 100755 --- a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh +++ b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh @@ -4,7 +4,7 @@ # # Qwen3-4B Mixture-of-LoRA GRPO on DAPO math with 8 colocated GPUs. -set -euo pipefail +set -ex SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" if [[ -z "${RELAX_ENTRYPOINT_MODE:-}" ]]; then diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index c52e179c5..456a2b78d 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -23,6 +23,7 @@ mixture_lora_metrics_from_packed_records, pack_mixture_lora_routing_records, ) +from relax.utils import megatron_bridge_utils from relax.utils.mixture_lora import MixtureLoraConfig, compute_routing_statistics @@ -38,6 +39,37 @@ def _config(*, num_experts=3, top_k=2, rank=2, alpha=4.0): ) +def test_bridge_base_load_filters_mixture_parameters_only_inside_patch_context(): + peft_bridge = pytest.importorskip("megatron.bridge.models.conversion.peft_bridge") + bridge = peft_bridge.MegatronPeftBridge() + parameter_name = "decoder.layers.0.self_attention.linear_qkv.mixture_lora.router.weight" + model = torch.nn.Module() + model.config = SimpleNamespace() + model.share_embeddings_and_output_weights = False + + assert bridge._is_adapter_param_name(parameter_name) is False + with megatron_bridge_utils.patch_megatron_model([model]): + assert bridge._is_adapter_param_name(parameter_name) is True + assert bridge._is_adapter_param_name("decoder.layers.0.self_attention.linear_qkv.weight") is False + assert bridge._is_adapter_param_name(parameter_name) is False + + +def test_bridge_mixture_filter_tolerates_missing_bridge_helpers(monkeypatch): + original_import = __import__ + + def import_without_peft_bridge(name, *args, **kwargs): + if name == "megatron.bridge.models.conversion.peft_bridge": + raise ImportError("bridge PEFT helpers unavailable") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(megatron_bridge_utils, "_bridge_mixture_filter_patched", False) + monkeypatch.setattr("builtins.__import__", import_without_peft_bridge) + + megatron_bridge_utils._ensure_bridge_mixture_filter_patched() + + assert megatron_bridge_utils._bridge_mixture_filter_patched is False + + class _TupleLinear(torch.nn.Module): def __init__(self, input_size, output_size, *, return_mode="standard"): super().__init__() @@ -691,6 +723,30 @@ def test_mixture_lora_peft_uses_bridge_matcher_and_freezes_base(monkeypatch): assert all(parameter.requires_grad for parameter in transformed.linear_qkv.mixture_lora.parameters()) +def test_mixture_lora_peft_uses_global_layer_id_with_pipeline_offset(monkeypatch): + _install_fake_bridge(monkeypatch) + transformer_package = types.ModuleType("megatron.core.transformer") + transformer_layer_module = types.ModuleType("megatron.core.transformer.transformer_layer") + observed = {} + + def get_transformer_layer_offset(config, vp_stage=None): + observed["config"] = config + observed["vp_stage"] = vp_stage + return 18 + + transformer_layer_module.get_transformer_layer_offset = get_transformer_layer_offset + monkeypatch.setitem(sys.modules, "megatron.core.transformer", transformer_package) + monkeypatch.setitem(sys.modules, "megatron.core.transformer.transformer_layer", transformer_layer_module) + model = torch.nn.Module() + model.linear_qkv = _TupleLinear(4, 5) + model.linear_qkv.config = SimpleNamespace(num_layers=36) + + transformed = build_mixture_lora_peft(_config(), dropout=0.0, vp_stage=1)(model, training=True) + + assert transformed.linear_qkv.mixture_lora.site_id == "decoder.layers.18.self_attention.linear_qkv" + assert observed == {"config": model.linear_qkv.to_wrap.config, "vp_stage": 1} + + def test_mixture_lora_peft_instantiates_with_real_bridge_when_available(): pytest.importorskip("megatron.bridge.peft.base") diff --git a/tests/backends/megatron/test_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py new file mode 100644 index 000000000..b554922fc --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from types import SimpleNamespace + +import pytest + +from relax.backends.megatron.arguments import _hf_validate_args + + +def test_mixture_lora_rejects_multimodal_hf_config(): + args = SimpleNamespace(lora_num_experts=4, apply_rope_fusion=False) + hf_config = SimpleNamespace(text_config=SimpleNamespace()) + + with pytest.raises(AssertionError, match="text-only base models"): + _hf_validate_args(args, hf_config) + + +def test_mixture_lora_rejects_moe_hf_config(): + args = SimpleNamespace(lora_num_experts=4) + hf_config = SimpleNamespace(num_experts=8) + + with pytest.raises(AssertionError, match="dense base models"): + _hf_validate_args(args, hf_config) diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 6cfc88480..2d3fb4654 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -6,12 +6,83 @@ import pytest import torch +from relax.backends.megatron.weight_update.common import _maybe_get_cpu_backup +from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge from relax.backends.megatron.weight_update.mixture_lora_sync import ( MixtureLoraParamInfo, + _gather_pipeline_param_infos, merge_mixture_lora_tp_shards, ) from relax.backends.megatron.weight_update.update_weight_from_tensor import iter_mixture_weight_updates from relax.utils.mixture_lora import MixtureLoraStateSpec +from relax.utils.types import ParamInfo + + +def test_selective_offload_uses_cpu_copy_only_after_live_storage_is_released(): + cpu_copy = torch.arange(4, dtype=torch.float32) + released = torch.empty(0) + released._relax_cpu_offload_data = cpu_copy + + assert _maybe_get_cpu_backup(released) is cpu_copy + + resident = torch.ones(1) + resident._relax_cpu_offload_data = cpu_copy + + assert _maybe_get_cpu_backup(resident) is resident + + +def _run_non_expert_bridge_iterator(*, current_rank: int, src_rank: int): + name = "decoder.layers.0.self_attention.linear_proj.to_wrap.weight" + info = ParamInfo( + name=name, + dtype=torch.float32, + shape=torch.Size((4, 8)), + attrs={}, + size=4 * 8 * 4, + src_rank=src_rank, + ) + parameter = torch.nn.Parameter(torch.ones(info.shape), requires_grad=False) + converted = [("model.layers.0.self_attn.o_proj.weight", torch.full(info.shape, 2.0))] + iterator = HfWeightIteratorBridge.__new__(HfWeightIteratorBridge) + iterator.args = MagicMock() + iterator._bridge_converter = MagicMock() + iterator._bridge_converter.convert.return_value = converted + iterator.lora_merge_mode = False + iterator._expert_buckets = [] + iterator._non_expert_buckets = [[info]] + iterator._vanilla_key_map = {info.name: info.name} + + def broadcast_owner_result(bucket_infos, all_converted, device): + assert bucket_infos == [info] + assert device == "cpu" + assert all_converted == ([converted] if current_rank == src_rank else [None]) + return converted + + module = "relax.backends.megatron.weight_update.hf_weight_iterator_bridge" + with ( + patch(f"{module}.dist.get_rank", return_value=current_rank), + patch(f"{module}.device_utils.make_current_torch_device", return_value="cpu"), + patch(f"{module}._load_to_gpu", return_value=[parameter]), + patch(f"{module}.all_gather_param", return_value=parameter), + patch(f"{module}._broadcast_converted_bucket", side_effect=broadcast_owner_result), + ): + result = list(iterator._iter_hf_params({info.name: parameter})) + + return iterator, result, converted + + +def test_non_expert_owner_stage_runs_bridge_conversion(): + iterator, result, converted = _run_non_expert_bridge_iterator(current_rank=0, src_rank=0) + + assert result == converted + iterator._bridge_converter.convert.assert_called_once() + + +def test_non_expert_remote_stage_only_receives_converted_result(): + iterator, result, converted = _run_non_expert_bridge_iterator(current_rank=1, src_rank=0) + + assert result == converted + iterator._bridge_converter.convert.assert_not_called() def _info(site, kind, global_shape, local_shape, shard_dim): @@ -30,6 +101,43 @@ def _info(site, kind, global_shape, local_shape, shard_dim): ) +def test_pipeline_metadata_gather_uses_output_list_first_and_merges_stages(): + stage_0 = _info("linear_qkv", "router.weight", (2, 6), (2, 6), None) + stage_1 = MixtureLoraParamInfo( + state=MixtureLoraStateSpec( + schema_version=1, + site_id="decoder.layers.1.self_attention.linear_qkv", + parameter_kind="router.weight", + global_shape=(2, 6), + dtype=torch.float32, + ), + local_shape=(2, 6), + tp_shard_dim=None, + src_rank=1, + weight_key="stage-1-weight", + ) + + def gather(output_list, input_object, group=None): + assert input_object == (0, {stage_0.state.parameter_name: stage_0}) + assert group == "pp-group" + output_list[:] = [input_object, (1, {stage_1.state.parameter_name: stage_1})] + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_gather_object", side_effect=gather), + ): + merged = _gather_pipeline_param_infos( + {stage_0.state.parameter_name: stage_0}, + pipeline_group="pp-group", + pipeline_world_size=2, + ) + + assert merged == { + stage_0.state.parameter_name: stage_0, + stage_1.state.parameter_name: stage_1, + } + + def test_qkv_lora_b_tp_shards_are_converted_from_group_layout_to_qkv_blocks(): # Two query groups, two query heads per group, then one K and one V head. grouped = torch.tensor([[[[10.0], [11.0], [20.0], [30.0], [12.0], [13.0], [21.0], [31.0]]]]).reshape(1, 8, 1) diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py index 88be4db97..0dbf9eb2b 100644 --- a/tests/backends/sglang/test_mixture_lora.py +++ b/tests/backends/sglang/test_mixture_lora.py @@ -1,10 +1,11 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. import os +from types import SimpleNamespace import pytest -from relax.backends.sglang.sglang_engine import _configure_external_model_environment +from relax.backends.sglang.sglang_engine import _compute_server_args, _configure_external_model_environment from relax.utils.mixture_lora import ( MixtureLoraConfig, configure_mixture_lora_external_model, @@ -69,3 +70,34 @@ def test_text_external_model_clears_multimodal_environment(monkeypatch): assert os.environ["SGLANG_EXTERNAL_MODEL_PACKAGE"] == "relax.models.qwen3_mixture_lora.sglang" assert "SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE" not in os.environ assert "SGLANG_EXTERNAL_MM_MODEL_ARCH" not in os.environ + + +def test_mixture_lora_keeps_base_cpu_backup_for_colocate_sleep(): + args = SimpleNamespace( + rollout_num_gpus_per_engine=1, + num_gpus_per_node=8, + colocate=True, + hf_checkpoint="/models/Qwen3-4B", + seed=42, + offload_rollout=True, + sglang_pp_size=1, + sglang_dp_size=1, + sglang_ep_size=1, + use_rollout_routing_replay=False, + fp16=True, + lora_rank=16, + lora_num_experts=4, + lora_adapter_mode=False, + ) + + kwargs, _ = _compute_server_args( + args, + rank=0, + dist_init_addr="127.0.0.1:1234", + nccl_port=1235, + host="127.0.0.1", + port=30000, + base_gpu_id=0, + ) + + assert kwargs["enable_weights_cpu_backup"] is True diff --git a/tests/utils/test_arguments_mixture_lora.py b/tests/utils/test_arguments_mixture_lora.py index 821695934..1b69d456a 100644 --- a/tests/utils/test_arguments_mixture_lora.py +++ b/tests/utils/test_arguments_mixture_lora.py @@ -55,6 +55,7 @@ def _args(**overrides): colocate=True, sglang_dp_size=1, sglang_tp_size=1, + tensor_model_parallel_size=2, ) defaults.update(overrides) return SimpleNamespace(**defaults) @@ -138,6 +139,7 @@ def test_missing_mixture_values_are_reported_together(arguments_module): [ ({"lora_num_experts": 0}, "lora-num-experts"), ({"lora_rank": 0}, "lora-rank"), + ({"lora_rank": 15}, "divisible"), ({"lora_router_top_k": 0}, "top_k"), ({"lora_router_top_k": 5}, "top_k"), ({"lora_router_temperature": 0.0}, "temperature"), From 6d1d751768fe8cffe962894bd53ab01353c7a5db Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Sun, 9 Aug 2026 23:53:13 +0800 Subject: [PATCH 18/41] fix(mixture-lora): align reshardable checkpoint resume MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary 对齐 Megatron 原生分布式 checkpoint 设计，使 Mixture-of-LoRA 在保持 TP/PP 不变时支持调整 DP size 后恢复 optimizer、expert、router 和训练进度。 Changes - recipe 启用 fully reshardable distributed optimizer，并在中英文文档说明跨 DP 恢复条件。 - 补齐 TP 路由指标同步，使日志 rank 在任意 TP rank 上都能获得完整指标。 - 增加 DP 梯度与指标、PP stage 梯度与指标、PP checkpoint 归属及 DP 缩容恢复测试。 - 更新 VPP Mixture factory 回归测试，并让纯 CPU checkpoint 测试不依赖 GPU 空闲显存。 Verification - Task 25、单 LoRA 及相关 VPP 定向回归：232 passed。 - ruff、冲突标记和 whitespace 检查通过。 - Qwen3-4B、DAPO math、colocate 从 4 进程 DP=2 checkpoint 恢复到 2 进程 DP=1，完成 step 2 并保存 iteration 2。 - 恢复后四个 expert 的全局平均激活权重约为 24.5%、26.2%、22.6%、26.7%，未出现单 expert 塌缩。 --- docs/en/guide/mixture-lora.md | 2 +- docs/zh/guide/mixture-lora.md | 2 +- relax/backends/megatron/model.py | 4 + .../text/run-qwen3-4B-mixture-lora-8xgpu.sh | 1 + ...est_mixture_lora_checkpoint_distributed.py | 264 ++++++++++++++++ .../megatron/test_mixture_lora_distributed.py | 287 ++++++++++++++++++ .../megatron/test_model_provider_vpp.py | 6 +- 7 files changed, 561 insertions(+), 5 deletions(-) create mode 100644 tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py diff --git a/docs/en/guide/mixture-lora.md b/docs/en/guide/mixture-lora.md index e69d019fa..46fe65f36 100644 --- a/docs/en/guide/mixture-lora.md +++ b/docs/en/guide/mixture-lora.md @@ -76,7 +76,7 @@ If an update fails, Relax resumes generation, keeps the previous weight version, ## Checkpoints -Expert and router tensors are ordinary model parameters in the native Megatron distributed checkpoint. The same checkpoint also restores optimizer, scheduler, iteration, and RNG state. Mixture mode does not create a separate HF PEFT adapter export. +Expert and router tensors are ordinary model parameters in the native Megatron distributed checkpoint. The same checkpoint also restores optimizer, scheduler, iteration, and RNG state. Mixture mode does not create a separate HF PEFT adapter export. The provided recipe enables Megatron's fully reshardable distributed-optimizer format so a saved run can be resumed with a different data-parallel size while keeping TP and PP unchanged. Resume by launching the same recipe with `--load` and `--save` pointing to the existing output directory. The saved Mixture metadata is checked against the current expert count, rank, Top-K, temperature, coefficient, alpha, target modules, dtype, and site dimensions before tensors are loaded. diff --git a/docs/zh/guide/mixture-lora.md b/docs/zh/guide/mixture-lora.md index 6d14503a1..01f04e417 100644 --- a/docs/zh/guide/mixture-lora.md +++ b/docs/zh/guide/mixture-lora.md @@ -76,7 +76,7 @@ Prompt、padding 和 dummy token 不参与 balance loss 与路由指标。静态 ## Checkpoint -Expert 和 router 是 Megatron 原生 distributed checkpoint 中的普通模型参数。Optimizer、scheduler、iteration 和 RNG 状态沿用同一个 checkpoint 恢复。Mixture 模式不会额外导出一份 HF PEFT adapter。 +Expert 和 router 是 Megatron 原生 distributed checkpoint 中的普通模型参数。Optimizer、scheduler、iteration 和 RNG 状态沿用同一个 checkpoint 恢复。Mixture 模式不会额外导出一份 HF PEFT adapter。随附的 recipe 启用了 Megatron 的 fully reshardable distributed optimizer 格式，因此保持 TP 和 PP 不变时，可以修改 DP size 后继续训练。 恢复训练时使用同一 recipe，并让 `--load` 和 `--save` 指向已有输出目录。加载 tensor 前会检查 checkpoint 中的 expert 数、rank、Top-K、temperature、coefficient、alpha、target module、dtype 和 site 维度是否与当前配置一致。 diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 2d68fa1de..45c2ff41a 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -119,6 +119,10 @@ def _reduce_mixture_lora_routing_metrics( ) if mpu.get_pipeline_model_parallel_world_size() > 1: torch.distributed.all_reduce(packed, group=mpu.get_pipeline_model_parallel_group()) + # Routed sites record statistics on TP rank 0 only. Replicate the reduced + # step metrics so logging remains correct regardless of the selected rank. + if mpu.get_tensor_model_parallel_world_size() > 1: + torch.distributed.all_reduce(packed, group=mpu.get_tensor_model_parallel_group()) return mixture_lora_metrics_from_packed_records( packed, site_ids, diff --git a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh index cb4358670..05ae83a7f 100755 --- a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh +++ b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh @@ -30,6 +30,7 @@ CKPT_ARGS=( --load "${OUTPUT_DIR}" --save "${OUTPUT_DIR}" --save-interval "${SAVE_INTERVAL:-50}" + --dist-ckpt-optim-fully-reshardable ) LORA_ARGS=( diff --git a/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py b/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py new file mode 100644 index 000000000..0967f6818 --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py @@ -0,0 +1,264 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from datetime import timedelta +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from relax.backends.megatron.mixture_lora import MixtureParallelLinearAdapter +from relax.utils.mixture_lora import MixtureLoraConfig + + +def _config() -> MixtureLoraConfig: + return MixtureLoraConfig( + num_experts=3, + rank=2, + top_k=2, + temperature=0.7, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_qkv",), + ) + + +def _save_cpu_checkpoint(dist_checkpointing, sharded_state: dict, checkpoint_dir: str) -> None: + """Avoid MCore's unconditional CUDA sync when every saved tensor is on + CPU.""" + + with patch.object(torch.cuda, "synchronize"): + dist_checkpointing.save(sharded_state, checkpoint_dir) + + +def _pipeline_checkpoint_worker( + rank: int, + world_size: int, + init_method: str, + checkpoint_dir: str, +) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=90), + ) + from megatron.core import dist_checkpointing, parallel_state + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + from megatron.core.transformer.transformer_config import TransformerConfig + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=world_size, + ) + try: + transformer_config = TransformerConfig( + num_layers=world_size, + hidden_size=4, + num_attention_heads=1, + pipeline_model_parallel_size=world_size, + pipeline_dtype=torch.float32, + use_cpu_initialization=True, + ) + model = torch.nn.Module() + model.linear_qkv = ColumnParallelLinear( + 4, + 5, + config=transformer_config, + init_method=lambda weight: torch.nn.init.normal_(weight, mean=0.0, std=0.02), + bias=False, + gather_output=False, + skip_bias_add=True, + ) + expected_site_id = f"decoder.layers.{rank}.self_attention.linear_qkv" + model.linear_qkv = MixtureParallelLinearAdapter( + model.linear_qkv, + _config(), + expected_site_id, + 4, + 5, + dropout=0.0, + tp_group=parallel_state.get_tensor_model_parallel_group(), + tp_rank=0, + tp_world_size=1, + ) + mixture = model.linear_qkv.mixture_lora + assert mixture.site_id == expected_site_id + + generator = torch.Generator().manual_seed(7100 + rank) + with torch.no_grad(): + for parameter in mixture.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator)) + expected_parameters = {name: parameter.detach().clone() for name, parameter in mixture.named_parameters()} + + prefix = f"{expected_site_id}." + sharded_state = model.linear_qkv.sharded_state_dict( + prefix=prefix, + metadata={"dp_cp_group": parallel_state.get_data_parallel_group(with_context_parallel=True)}, + ) + local_mixture_keys = tuple(sorted(key for key in sharded_state if ".mixture_lora." in key)) + gathered_keys = [None] * world_size + dist.all_gather_object(gathered_keys, local_mixture_keys) + expected_keys = { + f"decoder.layers.{stage}.self_attention.linear_qkv.mixture_lora.{suffix}" + for stage in range(world_size) + for suffix in ("_extra_state", "experts.lora_A", "experts.lora_B", "router.weight") + } + assert set().union(*map(set, gathered_keys)) == expected_keys + assert set(gathered_keys[0]).isdisjoint(gathered_keys[1]) + + if rank == 0: + Path(checkpoint_dir).mkdir(parents=True, exist_ok=True) + dist.barrier() + _save_cpu_checkpoint(dist_checkpointing, sharded_state, checkpoint_dir) + + with torch.no_grad(): + for parameter in mixture.parameters(): + parameter.zero_() + load_template = model.linear_qkv.sharded_state_dict( + prefix=prefix, + metadata={"dp_cp_group": parallel_state.get_data_parallel_group(with_context_parallel=True)}, + ) + loaded_state = dist_checkpointing.load(load_template, checkpoint_dir) + mixture_prefix = f"{prefix}mixture_lora." + local_mixture_state = { + key.removeprefix(mixture_prefix): value + for key, value in loaded_state.items() + if key.startswith(mixture_prefix) + } + mixture.load_state_dict(local_mixture_state) + for name, parameter in mixture.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) + assert mixture.site_id == expected_site_id + dist.barrier() + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def test_pipeline_checkpoint_restores_each_stage_parameters(tmp_path): + pytest.importorskip("megatron.bridge.peft.base") + checkpoint_dir = tmp_path / "mixture_lora_pp_checkpoint" + init_method = f"file://{tmp_path / 'mixture-lora-pp-checkpoint-gloo-init'}" + mp.spawn( + _pipeline_checkpoint_worker, + args=(2, init_method, str(checkpoint_dir)), + nprocs=2, + join=True, + ) + + +def _data_parallel_checkpoint_worker( + rank: int, + world_size: int, + init_method: str, + checkpoint_dir: str, + save_checkpoint: bool, +) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=90), + ) + from megatron.core import dist_checkpointing, parallel_state + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + from megatron.core.transformer.transformer_config import TransformerConfig + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + try: + torch.manual_seed(8100) + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=4, + num_attention_heads=1, + use_cpu_initialization=True, + ) + base = ColumnParallelLinear( + 4, + 6, + config=transformer_config, + init_method=lambda weight: torch.nn.init.normal_(weight, mean=0.0, std=0.02), + bias=False, + gather_output=False, + skip_bias_add=True, + ) + adapter = MixtureParallelLinearAdapter( + base, + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + dropout=0.0, + tp_group=parallel_state.get_tensor_model_parallel_group(), + tp_rank=0, + tp_world_size=1, + ) + generator = torch.Generator().manual_seed(8200) + with torch.no_grad(): + for parameter in adapter.mixture_lora.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator)) + expected_parameters = { + name: parameter.detach().clone() for name, parameter in adapter.mixture_lora.named_parameters() + } + prefix = "decoder.layers.0.self_attention.linear_qkv." + metadata = {"dp_cp_group": parallel_state.get_data_parallel_group(with_context_parallel=True)} + + if save_checkpoint: + if rank == 0: + Path(checkpoint_dir).mkdir(parents=True, exist_ok=True) + dist.barrier() + _save_cpu_checkpoint( + dist_checkpointing, + adapter.sharded_state_dict(prefix=prefix, metadata=metadata), + checkpoint_dir, + ) + else: + with torch.no_grad(): + for parameter in adapter.mixture_lora.parameters(): + parameter.zero_() + loaded_state = dist_checkpointing.load( + adapter.sharded_state_dict(prefix=prefix, metadata=metadata), + checkpoint_dir, + ) + mixture_prefix = f"{prefix}mixture_lora." + adapter.mixture_lora.load_state_dict( + { + key.removeprefix(mixture_prefix): value + for key, value in loaded_state.items() + if key.startswith(mixture_prefix) + } + ) + for name, parameter in adapter.mixture_lora.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) + dist.barrier() + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def test_checkpoint_reshards_mixture_parameters_when_data_parallel_size_shrinks(tmp_path): + pytest.importorskip("megatron.bridge.peft.base") + checkpoint_dir = tmp_path / "mixture_lora_dp_checkpoint" + save_init_method = f"file://{tmp_path / 'mixture-lora-dp-save-gloo-init'}" + load_init_method = f"file://{tmp_path / 'mixture-lora-dp-load-gloo-init'}" + + mp.spawn( + _data_parallel_checkpoint_worker, + args=(2, save_init_method, str(checkpoint_dir), True), + nprocs=2, + join=True, + ) + mp.spawn( + _data_parallel_checkpoint_worker, + args=(1, load_init_method, str(checkpoint_dir), False), + nprocs=1, + join=True, + ) diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index e5008abd7..75361a294 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -1,11 +1,13 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. from datetime import timedelta +from types import SimpleNamespace import torch import torch.distributed as dist import torch.multiprocessing as mp import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel from relax.backends.megatron.mixture_lora import ( MixtureLoRAAdapter, @@ -150,6 +152,291 @@ def test_routing_metrics_reduce_across_two_real_processes(tmp_path): mp.spawn(_distributed_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) +def _tensor_parallel_routing_metrics_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + from megatron.core import parallel_state + + from relax.backends.megatron.model import _reduce_mixture_lora_routing_metrics + + parallel_state.initialize_model_parallel(tensor_model_parallel_size=world_size) + try: + site_ids = ("layers.0.linear_qkv",) + contexts = ( + [_routing_context(site_ids[0], torch.ones(1, 2, dtype=torch.bool), num_sites=1, objective_scale=1.0)] + if rank == 0 + else [] + ) + metrics = _reduce_mixture_lora_routing_metrics( + SimpleNamespace(calculate_per_token_loss=False), + contexts, + (site_ids, 3, 2), + torch.device("cpu"), + ) + _assert_uniform_metrics(metrics, f"molora/{site_ids[0]}") + torch.testing.assert_close(metrics["molora/aux_loss"], torch.tensor(0.01, dtype=torch.float64)) + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def test_routing_metrics_are_available_on_every_tensor_parallel_rank(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-tp-metrics-gloo-init'}" + mp.spawn(_tensor_parallel_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) + + +def _new_distributed_adapter(config: MixtureLoraConfig, site_id: str) -> MixtureLoRAAdapter: + return MixtureLoRAAdapter( + config, + site_id, + 4, + 4, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + +def _set_distributed_adapter_weights(adapter: MixtureLoRAAdapter, seed: int) -> None: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + adapter.experts.lora_A.copy_(torch.randn(adapter.experts.lora_A.shape, generator=generator)) + adapter.experts.lora_B.copy_(torch.randn(adapter.experts.lora_B.shape, generator=generator)) + adapter.router.weight.copy_(torch.randn(adapter.router.weight.shape, generator=generator)) + + +def _distributed_microbatch(rank: int) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator().manual_seed(6100 + rank) + x = torch.randn(3, 1, 4, generator=generator) + mask = torch.tensor([[1, 1, rank == 0]], dtype=torch.bool) + return x, mask + + +def _new_routing_context( + site_ids: tuple[str, ...], + mask: torch.Tensor, + *, + microbatch_id: int, + main_loss_backward_scale: float, +) -> MixtureLoRARoutingContext: + return MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=microbatch_id, + response_mask=mask, + num_microbatches=1, + num_sites=len(site_ids), + num_samples=mask.shape[0], + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.tensor([main_loss_backward_scale]), + ) + + +def _routing_metrics_from_contexts( + contexts: list[MixtureLoRARoutingContext], + site_ids: tuple[str, ...], + config: MixtureLoraConfig, + *, + data_parallel_world_size: int, +) -> dict[str, torch.Tensor]: + packed = pack_mixture_lora_routing_records( + contexts, + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + device=torch.device("cpu"), + ) + return mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + calculate_per_token_loss=False, + data_parallel_world_size_with_cp=data_parallel_world_size, + ) + + +def _assert_metric_dict_close(actual: dict[str, torch.Tensor], expected: dict[str, torch.Tensor]) -> None: + assert actual.keys() == expected.keys() + for name in actual: + torch.testing.assert_close(actual[name], expected[name], atol=1e-8, rtol=1e-8, msg=name) + + +def _data_parallel_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + config = _config() + site_ids = ("layers.0.linear_qkv",) + adapter = _new_distributed_adapter(config, site_ids[0]) + _set_distributed_adapter_weights(adapter, 6200) + distributed_adapter = DistributedDataParallel(adapter) + + local_x, local_mask = _distributed_microbatch(rank) + local_context = _new_routing_context( + site_ids, + local_mask, + microbatch_id=rank, + main_loss_backward_scale=1.0, + ) + with activate_mixture_lora_routing_context(local_context): + local_output = distributed_adapter(local_x) + local_output.square().mean().backward() + + reference = _new_distributed_adapter(config, site_ids[0]) + _set_distributed_adapter_weights(reference, 6200) + reference_contexts = [] + for microbatch_rank in range(world_size): + reference_x, reference_mask = _distributed_microbatch(microbatch_rank) + reference_context = _new_routing_context( + site_ids, + reference_mask, + microbatch_id=microbatch_rank, + main_loss_backward_scale=1.0 / world_size, + ) + with activate_mixture_lora_routing_context(reference_context): + reference_output = reference(reference_x) + (reference_output.square().mean() / world_size).backward() + reference_contexts.append(reference_context) + + for distributed_param, reference_param in zip( + distributed_adapter.module.parameters(), reference.parameters(), strict=True + ): + torch.testing.assert_close(distributed_param.grad, reference_param.grad, atol=2e-6, rtol=2e-6) + + packed = pack_mixture_lora_routing_records( + [local_context], + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + device=torch.device("cpu"), + ) + dist.all_reduce(packed) + actual_metrics = mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + calculate_per_token_loss=False, + data_parallel_world_size_with_cp=world_size, + ) + expected_metrics = _routing_metrics_from_contexts( + reference_contexts, + site_ids, + config, + data_parallel_world_size=world_size, + ) + _assert_metric_dict_close(actual_metrics, expected_metrics) + finally: + dist.destroy_process_group() + + +def test_data_parallel_gradients_and_metrics_match_microbatch_reference(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-dp-gloo-init'}" + mp.spawn(_data_parallel_worker, args=(2, init_method), nprocs=2, join=True) + + +def _pipeline_parallel_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + config = _config() + site_ids = ("layers.0.linear_qkv", "layers.1.linear_qkv") + x, mask = _distributed_microbatch(0) + + reference_adapters = [_new_distributed_adapter(config, site_id) for site_id in site_ids] + for stage, reference_adapter in enumerate(reference_adapters): + _set_distributed_adapter_weights(reference_adapter, 6300 + stage) + reference_context = _new_routing_context( + site_ids, + mask, + microbatch_id=0, + main_loss_backward_scale=1.0, + ) + with activate_mixture_lora_routing_context(reference_context): + reference_hidden = reference_adapters[0](x) + reference_output = reference_adapters[1](reference_hidden) + reference_output.square().mean().backward() + + local_adapter = _new_distributed_adapter(config, site_ids[rank]) + _set_distributed_adapter_weights(local_adapter, 6300 + rank) + local_context = _new_routing_context( + site_ids, + mask, + microbatch_id=0, + main_loss_backward_scale=1.0, + ) + if rank == 0: + with activate_mixture_lora_routing_context(local_context): + local_hidden = local_adapter(x) + torch.testing.assert_close(local_hidden, reference_hidden.detach()) + dist.send(local_hidden.detach(), dst=1) + hidden_gradient = torch.empty_like(local_hidden) + dist.recv(hidden_gradient, src=1) + local_hidden.backward(hidden_gradient) + else: + local_hidden = torch.empty_like(reference_hidden) + dist.recv(local_hidden, src=0) + local_hidden.requires_grad_(True) + with activate_mixture_lora_routing_context(local_context): + local_output = local_adapter(local_hidden) + torch.testing.assert_close(local_output, reference_output.detach()) + local_output.square().mean().backward() + dist.send(local_hidden.grad, dst=0) + + for local_param, reference_param in zip( + local_adapter.parameters(), reference_adapters[rank].parameters(), strict=True + ): + torch.testing.assert_close(local_param.grad, reference_param.grad, atol=2e-6, rtol=2e-6) + assert local_adapter.router.weight.grad.norm() > 0 + + packed = pack_mixture_lora_routing_records( + [local_context], + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + device=torch.device("cpu"), + ) + dist.all_reduce(packed) + actual_metrics = mixture_lora_metrics_from_packed_records( + packed, + site_ids, + num_experts=config.num_experts, + top_k=config.top_k, + calculate_per_token_loss=False, + data_parallel_world_size_with_cp=1, + ) + expected_metrics = _routing_metrics_from_contexts( + [reference_context], + site_ids, + config, + data_parallel_world_size=1, + ) + _assert_metric_dict_close(actual_metrics, expected_metrics) + finally: + dist.destroy_process_group() + + +def test_pipeline_parallel_stage_gradients_and_metrics_match_reference(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-pp-gloo-init'}" + mp.spawn(_pipeline_parallel_worker, args=(2, init_method), nprocs=2, join=True) + + def _context_parallel_balance_worker(rank: int, world_size: int, init_method: str) -> None: dist.init_process_group( backend="gloo", diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 992988a8d..e59ad1c7e 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -258,8 +258,8 @@ def test_model_provider_uses_mixture_lora_for_multiple_experts(monkeypatch): calls = [] mixture_module = types.ModuleType("relax.backends.megatron.mixture_lora") - def build_mixture_lora_peft(received_config, dropout): - calls.append((received_config, dropout)) + def build_mixture_lora_peft(received_config, dropout, vp_stage=None): + calls.append((received_config, dropout, vp_stage)) return lambda received_model, training: received_model mixture_module.build_mixture_lora_peft = build_mixture_lora_peft @@ -274,7 +274,7 @@ def build_mixture_lora_peft(received_config, dropout): provider = module.wrap_model_provider_with_lora(lambda **kwargs: model, args) assert provider() is model - assert calls == [(config, 0.1)] + assert calls == [(config, 0.1, None)] def test_model_provider_keeps_single_expert_on_existing_lora_path(monkeypatch): From 317cc4413545580fc8811a417f4a080ad53d8e52 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Tue, 11 Aug 2026 00:22:15 +0800 Subject: [PATCH 19/41] =?UTF-8?q?fix(recipe):=20Mixture-of-LoRA=20?= =?UTF-8?q?=E9=AA=8C=E6=94=B6=E9=85=8D=E7=BD=AE=E6=94=B9=E7=94=A8=20BF16?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将 Qwen3-4B Mixture-of-LoRA recipe 从 FP16 切换为 BF16，并移除仅适用于 FP16 动态缩放的 loss scale 参数。同步更新中英文使用文档，使公开 recipe 与已完成的 200-step 验收实验配置一致。 --- docs/en/guide/mixture-lora.md | 2 +- docs/zh/guide/mixture-lora.md | 2 +- scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh | 4 +--- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/docs/en/guide/mixture-lora.md b/docs/en/guide/mixture-lora.md index 46fe65f36..08a508422 100644 --- a/docs/en/guide/mixture-lora.md +++ b/docs/en/guide/mixture-lora.md @@ -91,4 +91,4 @@ OUTPUT_DIR=/path/to/qwen3-4b-mixture-lora \ bash scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh ``` -The actor uses TP=2 with sequence parallelism. The rollout allocation creates eight independent one-GPU SGLang engines. Environment variables such as `NUM_ROLLOUT`, `LORA_NUM_EXPERTS`, `LORA_RANK`, and `LORA_ROUTER_TOP_K` can override the recipe values. Additional Relax arguments can be appended to the command. +The actor uses TP=2 with sequence parallelism. The rollout allocation creates eight independent one-GPU SGLang engines. Environment variables such as `NUM_ROLLOUT`, `LORA_NUM_EXPERTS`, `LORA_RANK`, and `LORA_ROUTER_TOP_K` can override the recipe values. The recipe uses BF16. Additional Relax arguments can be appended to the command. diff --git a/docs/zh/guide/mixture-lora.md b/docs/zh/guide/mixture-lora.md index 01f04e417..c00bdedf7 100644 --- a/docs/zh/guide/mixture-lora.md +++ b/docs/zh/guide/mixture-lora.md @@ -91,4 +91,4 @@ OUTPUT_DIR=/path/to/qwen3-4b-mixture-lora \ bash scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh ``` -Actor 使用 TP=2 和 sequence parallel。Rollout 资源会建立八个独立的单卡 SGLang engine。可以通过 `NUM_ROLLOUT`、`LORA_NUM_EXPERTS`、`LORA_RANK`、`LORA_ROUTER_TOP_K` 等环境变量覆盖 recipe 中的值，也可以在命令末尾继续追加 Relax 参数。 +Actor 使用 TP=2 和 sequence parallel。Rollout 资源会建立八个独立的单卡 SGLang engine。可以通过 `NUM_ROLLOUT`、`LORA_NUM_EXPERTS`、`LORA_RANK`、`LORA_ROUTER_TOP_K` 等环境变量覆盖 recipe 中的值。这份 recipe 使用 BF16，也可以在命令末尾继续追加 Relax 参数。 diff --git a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh index 05ae83a7f..6ba09df44 100755 --- a/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh +++ b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh @@ -80,8 +80,6 @@ OPTIMIZER_ARGS=( --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 - --initial-loss-scale 32768 - --min-loss-scale 1 --use-precision-aware-optimizer --no-store-param-remainders ) @@ -135,7 +133,7 @@ ray job submit ${RAY_NO_WAIT:+--no-wait} --address="${RAY_ADDRESS:-http://127.0. --max-staleness 0 \ --num-data-storage-units 1 \ --colocate \ - --fp16 \ + --bf16 \ --use-health-check \ "${MODEL_ARGS[@]}" \ "${CKPT_ARGS[@]}" \ From 87dfe2f21601e99bffa4289f706b319cc5174815 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Tue, 11 Aug 2026 01:14:50 +0800 Subject: [PATCH 20/41] =?UTF-8?q?test(mixture-lora):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E5=85=A8=E9=87=8F=E5=9B=9E=E5=BD=92=E6=B5=8B=E8=AF=95=E4=BE=9D?= =?UTF-8?q?=E8=B5=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为 SGLang router registration fixture 补充 Mixture 配置与 external model stub，并为 SFT loss 测试提供有效的 DP world size。修复全量 pytest 中由新增 Mixture import 和严格 loss 缩放校验触发的测试错误。 --- tests/backends/megatron/test_sft_chunked_ce.py | 1 + tests/backends/sglang/test_router_registration.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/tests/backends/megatron/test_sft_chunked_ce.py b/tests/backends/megatron/test_sft_chunked_ce.py index 48cc12a14..db5c87a43 100644 --- a/tests/backends/megatron/test_sft_chunked_ce.py +++ b/tests/backends/megatron/test_sft_chunked_ce.py @@ -531,6 +531,7 @@ def spy_checkpoint(_function, *args, **kwargs): monkeypatch.setattr(_loss_mod, "checkpoint", spy_checkpoint) monkeypatch.setattr(_loss_mod, "get_cp_local_num_tokens", lambda *args, **kwargs: torch.tensor(1)) monkeypatch.setattr(_loss_mod, "get_sum_of_sample_mean", lambda *args, **kwargs: object()) + monkeypatch.setattr(_loss_mod.mpu, "get_data_parallel_world_size", lambda **_kwargs: 1) args = Namespace( loss_type="sft", diff --git a/tests/backends/sglang/test_router_registration.py b/tests/backends/sglang/test_router_registration.py index 871a0216f..48c1fd887 100644 --- a/tests/backends/sglang/test_router_registration.py +++ b/tests/backends/sglang/test_router_registration.py @@ -61,10 +61,16 @@ def sglang_engine_module(monkeypatch): monkeypatch.setitem(sys.modules, "relax.utils.logging_utils", logging_utils) megatron_peft_utils = ModuleType("relax.utils.megatron_peft_utils") + megatron_peft_utils.build_mixture_lora_config = lambda _args: None megatron_peft_utils.convert_megatron_to_hf_target_modules = lambda value: value megatron_peft_utils.is_lora_enabled = lambda _args: False + megatron_peft_utils.is_mixture_lora_enabled = lambda _args: False monkeypatch.setitem(sys.modules, "relax.utils.megatron_peft_utils", megatron_peft_utils) + mixture_lora = ModuleType("relax.utils.mixture_lora") + mixture_lora.configure_mixture_lora_external_model = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "relax.utils.mixture_lora", mixture_lora) + sys.modules.pop("relax.backends.sglang.sglang_engine", None) module = importlib.import_module("relax.backends.sglang.sglang_engine") yield module From f9ce3fbf4ecadb660a5b47442a49402747d416c0 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Tue, 11 Aug 2026 02:30:20 +0800 Subject: [PATCH 21/41] =?UTF-8?q?test(mixture-lora):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E5=8F=AF=E9=80=89=E5=90=8E=E7=AB=AF=E4=BE=9D=E8=B5=96=E4=BF=9D?= =?UTF-8?q?=E6=8A=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 默认 CPU CI 未安装 Megatron 或 SGLang 时，按仓库现有测试约定跳过依赖真实后端的新增用例，避免在测试收集阶段失败。 安装完整后端的环境仍会执行参数校验、权重同步、SGLang 配置和模型一致性测试。 --- .../megatron/test_mixture_lora_arguments.py | 7 ++++++- .../test_mixture_lora_weight_sync.py | 19 +++++++++++++------ tests/backends/sglang/test_mixture_lora.py | 14 ++++++++++++-- .../qwen3_mixture_lora/test_sglang_model.py | 12 +++++++++--- 4 files changed, 40 insertions(+), 12 deletions(-) diff --git a/tests/backends/megatron/test_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py index b554922fc..6f78658f4 100644 --- a/tests/backends/megatron/test_mixture_lora_arguments.py +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -4,7 +4,12 @@ import pytest -from relax.backends.megatron.arguments import _hf_validate_args + +# The default CPU CI does not install Megatron. The validation tests run in +# the official Relax image with the real Megatron argument definitions. +pytest.importorskip("megatron.training.arguments") + +from relax.backends.megatron.arguments import _hf_validate_args # noqa: E402 def test_mixture_lora_rejects_multimodal_hf_config(): diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 2d3fb4654..378003627 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -6,16 +6,23 @@ import pytest import torch -from relax.backends.megatron.weight_update.common import _maybe_get_cpu_backup -from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge -from relax.backends.megatron.weight_update.mixture_lora_sync import ( + +# These tests exercise the real Megatron Bridge weight conversion and +# distributed synchronization modules, which are not installed in CPU CI. +pytest.importorskip("megatron.bridge.peft.lora") + +from relax.backends.megatron.weight_update.common import _maybe_get_cpu_backup # noqa: E402 +from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge # noqa: E402 +from relax.backends.megatron.weight_update.mixture_lora_sync import ( # noqa: E402 MixtureLoraParamInfo, _gather_pipeline_param_infos, merge_mixture_lora_tp_shards, ) -from relax.backends.megatron.weight_update.update_weight_from_tensor import iter_mixture_weight_updates -from relax.utils.mixture_lora import MixtureLoraStateSpec -from relax.utils.types import ParamInfo +from relax.backends.megatron.weight_update.update_weight_from_tensor import ( # noqa: E402 + iter_mixture_weight_updates, +) +from relax.utils.mixture_lora import MixtureLoraStateSpec # noqa: E402 +from relax.utils.types import ParamInfo # noqa: E402 def test_selective_offload_uses_cpu_copy_only_after_live_storage_is_released(): diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py index 0dbf9eb2b..539c9a470 100644 --- a/tests/backends/sglang/test_mixture_lora.py +++ b/tests/backends/sglang/test_mixture_lora.py @@ -5,8 +5,18 @@ import pytest -from relax.backends.sglang.sglang_engine import _compute_server_args, _configure_external_model_environment -from relax.utils.mixture_lora import ( + +# The engine imports the real SGLang router and Megatron-backed checkpoint +# client. CPU CI omits these backend dependencies. +pytest.importorskip("megatron.core") +pytest.importorskip("sglang.srt.server_args") +pytest.importorskip("sglang_router") + +from relax.backends.sglang.sglang_engine import ( # noqa: E402 + _compute_server_args, + _configure_external_model_environment, +) +from relax.utils.mixture_lora import ( # noqa: E402 MixtureLoraConfig, configure_mixture_lora_external_model, deserialize_mixture_lora_config, diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 5cd340543..4b78a6c6f 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -5,14 +5,20 @@ import torch.nn.functional as F from torch import nn -from relax.backends.megatron.mixture_lora import MixtureLoRAAdapter -from relax.models.qwen3_mixture_lora.sglang.model import ( + +# This parity suite needs both model backends. They are available in the +# official Relax image but intentionally absent from the default CPU CI. +pytest.importorskip("megatron.bridge") +pytest.importorskip("sglang.srt.models.qwen3") + +from relax.backends.megatron.mixture_lora import MixtureLoRAAdapter # noqa: E402 +from relax.models.qwen3_mixture_lora.sglang.model import ( # noqa: E402 EntryClass, SGLangMixtureLoRA, attach_sglang_mixture_lora, load_sglang_mixture_lora_weights, ) -from relax.utils.mixture_lora import MixtureLoraConfig +from relax.utils.mixture_lora import MixtureLoraConfig # noqa: E402 def _config(): From a82eb6262826313acacec0502de8efd246fcccd3 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Tue, 11 Aug 2026 11:42:14 +0800 Subject: [PATCH 22/41] =?UTF-8?q?test(mixture-lora):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E5=BC=A0=E9=87=8F=E5=B9=B6=E8=A1=8C=E6=8C=87=E6=A0=87=E7=94=A8?= =?UTF-8?q?=E4=BE=8B=E5=AE=88=E5=8D=AB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 默认 CPU CI 未安装 Megatron 时，张量并行路由指标用例仍会失败：mp.spawn 拉起的子进程是全新解释器，不加载 pytest 与 conftest，worker 内的 from megatron.core import parallel_state 直接抛 ModuleNotFoundError，并以 ProcessRaisedException 冒泡，使三个 Python 版本的测试 job 全部中断。 在 mp.spawn 之前补上 megatron.core 与 megatron.training 的 importorskip，并为该文件补充缺失的 pytest 导入。守卫模块按 worker 的实际依赖链选取：parallel_state 来自 megatron.core，_reduce_mixture_lora_routing_metrics 所在的 relax.backends.megatron.model 在模块级还需要 megatron.training。 这是 f9ce3fb 守卫补齐的延续，该 commit 覆盖了模块级可见后端导入的四个文件，遗漏了本文件这种藏在子进程里的导入。写法与 test_mixture_lora_checkpoint_distributed.py 现有约定一致，安装完整后端的官方镜像仍会完整执行该用例。 Co-Authored-By: Claude Opus 5 (1M context) --- tests/backends/megatron/test_mixture_lora_distributed.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index 75361a294..e719b3b6d 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -3,6 +3,7 @@ from datetime import timedelta from types import SimpleNamespace +import pytest import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -186,6 +187,10 @@ def _tensor_parallel_routing_metrics_worker(rank: int, world_size: int, init_met def test_routing_metrics_are_available_on_every_tensor_parallel_rank(tmp_path): + # The spawned workers need the real Megatron tensor-parallel state, which the + # default CPU CI does not install. This case runs in the official Relax image. + pytest.importorskip("megatron.core") + pytest.importorskip("megatron.training") init_method = f"file://{tmp_path / 'mixture-lora-tp-metrics-gloo-init'}" mp.spawn(_tensor_parallel_routing_metrics_worker, args=(2, init_method), nprocs=2, join=True) From 5d86614dfabce4ff99e83896b82ff9310945ebcd Mon Sep 17 00:00:00 2001 From: PopHirasawa <1831651457@qq.com> Date: Wed, 12 Aug 2026 23:17:06 +0800 Subject: [PATCH 23/41] =?UTF-8?q?test(sglang):=20=E8=A1=A5=E9=BD=90=20S3?= =?UTF-8?q?=20loader=20=E7=9A=84=20Mixture=20=E6=B5=8B=E8=AF=95=E6=A1=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 最新 main 新增的 S3 loader 测试会构造精简版 megatron_peft_utils 模块，再单独导入 sglang_engine。Task 25 为 sglang_engine 增加 Mixture 配置与模式判断导入后，原 fixture 缺少对应符号，导致测试在收集阶段失败。 为 fixture 补充 build_mixture_lora_config 和 is_mixture_lora_enabled 两个 stub，并固定返回未启用 Mixture，使测试继续覆盖原有 S3 模型加载路径。该修改只影响测试替身，不改变 S3、SGLang 或 Mixture 的运行时行为。 验证：tests/test_s3_model_loader.py 共 51 个用例全部通过。 --- tests/test_s3_model_loader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_s3_model_loader.py b/tests/test_s3_model_loader.py index ef30c8683..693c1c0f4 100644 --- a/tests/test_s3_model_loader.py +++ b/tests/test_s3_model_loader.py @@ -124,8 +124,10 @@ class ServerArgs: monkeypatch.setitem(sys.modules, "relax.utils.logging_utils", logging_utils) megatron_peft_utils = ModuleType("relax.utils.megatron_peft_utils") + megatron_peft_utils.build_mixture_lora_config = lambda _args: None megatron_peft_utils.convert_megatron_to_hf_target_modules = lambda value: value megatron_peft_utils.is_lora_enabled = lambda _args: False + megatron_peft_utils.is_mixture_lora_enabled = lambda _args: False monkeypatch.setitem(sys.modules, "relax.utils.megatron_peft_utils", megatron_peft_utils) module_name = "relax.backends.sglang.sglang_engine" From f2972fe15b28d366235f5308a36947522e6590c0 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 22:58:24 +0800 Subject: [PATCH 24/41] =?UTF-8?q?fix(sglang):=20=E4=BF=AE=E5=A4=8D=20Mixtu?= =?UTF-8?q?re=20=E5=9C=A8=20PP=20stage=20=E7=9A=84=E5=B1=82=E6=B3=A8?= =?UTF-8?q?=E5=85=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLang 的 PP 模型使用全局 layer 表，并以占位模块表示其他 stage 的层。Mixture 安装逻辑改为只遍历当前 stage 的 [start_layer, end_layer) 区间，直接使用全局 layer id，避免访问 PPMissingLayer 或重复偏移层编号。 新增 PP stage 回归测试，使用无 self_attn 的占位层验证只为本 stage 真实层安装 qkv/proj adapter，并检查同步 schema 使用正确的全局 site id。 验证：tests/models/qwen3_mixture_lora/test_sglang_model.py，8 passed。 --- .../models/qwen3_mixture_lora/sglang/model.py | 5 ++-- .../qwen3_mixture_lora/test_sglang_model.py | 27 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index 5b9d0ce58..bde415b1c 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -176,8 +176,9 @@ def _install_mixture_lora(self) -> None: if not targets.issubset(supported_targets): raise ValueError(f"Unsupported SGLang Mixture-of-LoRA targets: {sorted(targets - supported_targets)}") start_layer = getattr(self.model, "start_layer", 0) - for local_layer_id, layer in enumerate(self.model.layers): - layer_id = start_layer + local_layer_id + end_layer = getattr(self.model, "end_layer", len(self.model.layers)) + for layer_id in range(start_layer, end_layer): + layer = self.model.layers[layer_id] attention = layer.self_attn if "linear_qkv" in targets: site_id = f"decoder.layers.{layer_id}.self_attention.linear_qkv" diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 4b78a6c6f..a8511558d 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -183,6 +183,33 @@ def test_sglang_weight_loader_skips_layers_owned_by_another_pp_stage(): assert loaded_names == set() +def test_sglang_installs_adapters_only_on_layers_owned_by_pp_stage(): + model = nn.Module() + model.mixture_lora_config = _config() + model.model = nn.Module() + model.model.start_layer = 1 + model.model.end_layer = 2 + + missing_layer = nn.Module() + local_layer = nn.Module() + local_layer.self_attn = nn.Module() + local_layer.self_attn.qkv_proj = _TupleLinear(4, 6) + local_layer.self_attn.qkv_proj.input_size = 4 + local_layer.self_attn.qkv_proj.output_size = 6 + local_layer.self_attn.o_proj = _TupleLinear(6, 4) + local_layer.self_attn.o_proj.input_size = 6 + local_layer.self_attn.o_proj.output_size = 4 + model.model.layers = nn.ModuleList([missing_layer, local_layer]) + + EntryClass._install_mixture_lora(model) + + assert not hasattr(missing_layer, "self_attn") + assert local_layer.self_attn.qkv_proj.mixture_lora.site_id == ( + "decoder.layers.1.self_attention.linear_qkv" + ) + assert local_layer.self_attn.o_proj.mixture_lora.site_id == "decoder.layers.1.self_attention.linear_proj" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_sglang_adapter_can_be_captured_by_cuda_graph(): adapter = SGLangMixtureLoRA( From 84550ac81f74085dfdcfb227a4e5885af1075b70 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:02:30 +0800 Subject: [PATCH 25/41] =?UTF-8?q?fix(sglang):=20=E4=BD=BF=E7=94=A8?= =?UTF-8?q?=E6=B5=AE=E7=82=B9=E8=AE=A1=E7=AE=97=E7=B1=BB=E5=9E=8B=E5=88=9B?= =?UTF-8?q?=E5=BB=BA=20Mixture=20adapter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 量化 SGLang linear 的首个参数可能是 int32 或 int8 的 packed 权重，不能作为 LoRA expert 和 router 的参数类型。安装 adapter 时继续从 base parameter 获取设备，但优先使用 linear.params_dtype 作为浮点计算类型。 当 linear 无法提供浮点 params_dtype，且首个参数也是整数类型时，在模型构造阶段给出明确错误，避免创建需要梯度的整数 Parameter。 新增量化 linear 测试，覆盖 int32 packed 权重配合 BF16 params_dtype，以及缺少浮点 dtype 时的失败路径。验证：SGLang 模型测试 10 passed。 --- .../models/qwen3_mixture_lora/sglang/model.py | 8 ++++++- .../qwen3_mixture_lora/test_sglang_model.py | 23 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index bde415b1c..c4b77656d 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -100,6 +100,12 @@ def attach_sglang_mixture_lora( base_parameter = next(linear.parameters()) except StopIteration as error: raise ValueError(f"SGLang linear {site_id} has no parameters") from error + adapter_dtype = getattr(linear, "params_dtype", base_parameter.dtype) + if not isinstance(adapter_dtype, torch.dtype) or not adapter_dtype.is_floating_point: + raise TypeError( + f"SGLang linear {site_id} must expose a floating-point params_dtype for Mixture-of-LoRA, " + f"got {adapter_dtype}" + ) linear.add_module( "mixture_lora", SGLangMixtureLoRA( @@ -108,7 +114,7 @@ def attach_sglang_mixture_lora( input_size, output_size, device=base_parameter.device, - dtype=base_parameter.dtype, + dtype=adapter_dtype, ), ) linear._relax_mixture_lora_base_forward = linear.forward diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index a8511558d..102487523 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -46,6 +46,13 @@ def forward(self, x): return F.linear(x, self.weight), None +class _QuantizedTupleLinear(nn.Module): + def __init__(self, params_dtype): + super().__init__() + self.weight_packed = nn.Parameter(torch.ones(4, 4, dtype=torch.int32), requires_grad=False) + self.params_dtype = params_dtype + + def _fake_qwen_model_with_routed_qkv(): model = nn.Module() model.model = nn.Module() @@ -132,6 +139,22 @@ def test_attached_adapter_preserves_base_parameter_name_and_adds_delta(): assert "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight" in parameter_names +def test_attached_adapter_uses_quantized_linear_params_dtype(): + linear = _QuantizedTupleLinear(torch.bfloat16) + + attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) + + assert {parameter.dtype for parameter in linear.mixture_lora.parameters()} == {torch.bfloat16} + assert linear.mixture_lora.experts.lora_A.device == linear.weight_packed.device + + +def test_attached_adapter_rejects_missing_floating_point_params_dtype(): + linear = _QuantizedTupleLinear(torch.int32) + + with pytest.raises(TypeError, match="floating-point params_dtype"): + attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) + + def test_sglang_weight_loader_maps_training_names_and_validates_tensors(): model = _fake_qwen_model_with_routed_qkv() prefix = "decoder.layers.0.self_attention.linear_qkv.mixture_lora" From c65f640fb33f93c16b1529dbb29b8343084aa762 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:18:32 +0800 Subject: [PATCH 26/41] =?UTF-8?q?fix(weight-sync):=20=E9=98=BB=E6=AD=A2?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E7=9A=84=20Mixture=20=E6=9B=B4=E6=96=B0?= =?UTF-8?q?=E6=81=A2=E5=A4=8D=E6=9C=8D=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 带版本号的最终分块会在所有 rank 确认前序请求成功后再发送，避免失败更新被发布为完整策略。 同步失败后保持 rollout 暂停并保留原版本，补充故障注入测试及中英文恢复说明。 --- docs/en/guide/mixture-lora.md | 2 +- docs/zh/guide/mixture-lora.md | 2 +- .../update_weight_from_tensor.py | 27 ++++++++++++++----- .../test_mixture_lora_weight_sync.py | 12 ++++++--- 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/docs/en/guide/mixture-lora.md b/docs/en/guide/mixture-lora.md index 08a508422..4a2bc8854 100644 --- a/docs/en/guide/mixture-lora.md +++ b/docs/en/guide/mixture-lora.md @@ -72,7 +72,7 @@ Use the per-site post-Top-K weights, selection shares, and entropy to detect col Relax starts the Qwen3 SGLang external model automatically when Mixture mode is enabled. The first colocate update sends the frozen base plus all expert and router tensors. Later updates send all current expert and router tensors without resending the base. A weight version is published only after the final routed tensor is accepted. -If an update fails, Relax resumes generation, keeps the previous weight version, and reports the failed update instead of silently serving a partial policy. +If an update fails, Relax keeps generation paused, preserves the previous weight version, and reports the failure. Restart or fully resynchronize the rollout engines before serving requests again, because some unversioned chunks may already have been transferred. The final versioned chunk is sent only after every rank confirms that the preceding chunks succeeded. ## Checkpoints diff --git a/docs/zh/guide/mixture-lora.md b/docs/zh/guide/mixture-lora.md index c00bdedf7..2fd91c667 100644 --- a/docs/zh/guide/mixture-lora.md +++ b/docs/zh/guide/mixture-lora.md @@ -72,7 +72,7 @@ Prompt、padding 和 dummy token 不参与 balance loss 与路由指标。静态 启用 Mixture 后，Relax 会自动启动 Qwen3 SGLang external model。第一次 colocate 更新发送冻结的基座参数以及全部 expert/router 参数；后续更新只发送当前全部 expert/router 参数。最后一组 routed tensor 加载成功后才发布新的 weight version。 -更新失败时会恢复 generation、保留原 weight version，并报告错误，不会静默使用只加载了一部分的新策略。 +更新失败时会保持 generation 暂停、保留原 weight version，并报告错误。由于部分未带版本号的分块可能已经传输，需要重启或完整重同步 rollout engine 后才能继续提供服务。只有所有 rank 都确认前序分块成功后，才会发送带新版本号的最后一个分块。 ## Checkpoint diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index 70a9c3da6..202b00ecb 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -362,14 +362,11 @@ def resume_generation(*, finish_quantization: bool) -> None: if not phase_failed: continue - cleanup_error, _ = self._run_synchronized_weight_update_phase( - lambda: resume_generation(finish_quantization=True) + logger.error( + "Mixture-of-LoRA weight update failed during %s; rollout generation remains paused " + "until the engines are restarted or fully resynchronized", + phase_name, ) - if cleanup_error is not None: - logger.error( - "Failed to resume rollout generation after Mixture-of-LoRA update failure", - exc_info=(type(cleanup_error), cleanup_error, cleanup_error.__traceback__), - ) if local_error is not None: raise local_error raise RuntimeError(f"Mixture-of-LoRA weight update phase {phase_name!r} failed on another rank") @@ -408,7 +405,22 @@ def _send_weight_update_stream(self, updates) -> None: previous_refs: list[ObjectRef] = [] previous_tensors = None + has_previous_chunk = False for named_tensors, weight_version in updates: + # The versioned final chunk makes the whole update visible. Confirm + # the preceding chunk on every rank before publishing that version. + if weight_version is not None and has_previous_chunk: + previous_error, previous_failed = self._run_synchronized_weight_update_phase( + lambda: ray.get(previous_refs) if previous_refs else None + ) + del previous_tensors + previous_refs = [] + previous_tensors = None + if previous_failed: + if previous_error is not None: + raise previous_error + raise RuntimeError("A preceding Mixture-of-LoRA weight chunk failed on another rank") + refs, long_lived_tensors = self._send_hf_params(named_tensors, weight_version=weight_version) previous_error = None if previous_refs: @@ -419,6 +431,7 @@ def _send_weight_update_stream(self, updates) -> None: del previous_tensors previous_refs = refs previous_tensors = long_lived_tensors + has_previous_chunk = True device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) if previous_error is not None: raise previous_error diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 378003627..3c5c65aae 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -403,7 +403,7 @@ def record_updates(updates): assert updater._mixture_lora_sync.base_sync_done is True -def test_update_failure_resumes_generation_and_keeps_previous_version(): +def test_update_failure_keeps_generation_paused_and_preserves_version(): from relax.backends.megatron.weight_update.update_weight_from_tensor import UpdateWeightFromTensor events = [] @@ -446,12 +446,12 @@ def fail_update(_updates): ): updater._update_weights_mixture_lora() - assert events == ["pause", "flush", "update", "continue"] + assert events == ["pause", "flush", "update"] assert updater.weight_version == 4 assert updater._mixture_lora_sync.base_sync_done is False -def test_stream_failure_reaches_chunk_barrier_before_raising(): +def test_stream_failure_does_not_publish_the_final_versioned_chunk(): from relax.backends.megatron.weight_update.update_weight_from_tensor import UpdateWeightFromTensor updater = UpdateWeightFromTensor.__new__(UpdateWeightFromTensor) @@ -465,9 +465,13 @@ def test_stream_failure_reaches_chunk_barrier_before_raising(): "relax.backends.megatron.weight_update.update_weight_from_tensor." "device_utils.maybe_backend_barrier_on_weight_chunk" ) as chunk_barrier, + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_reduce"), patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), pytest.raises(RuntimeError, match="engine update failed"), ): updater._send_weight_update_stream([([("first", torch.ones(1))], None), ([("second", torch.ones(1))], 5)]) - assert chunk_barrier.call_count == 2 + updater._send_hf_params.assert_called_once() + assert updater._send_hf_params.call_args.kwargs["weight_version"] is None + assert chunk_barrier.call_count == 1 From 43fcc978dc4a9c3004e05f367d4c7d2fbbb8e3b2 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:28:46 +0800 Subject: [PATCH 27/41] =?UTF-8?q?fix(bridge):=20=E7=BB=9F=E4=B8=80=20EP=20?= =?UTF-8?q?=E5=A4=8D=E5=88=B6=E5=8F=82=E6=95=B0=E7=9A=84=E8=BD=AC=E6=8D=A2?= =?UTF-8?q?=E6=9D=A5=E6=BA=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EP metadata 合并时为 replicated non-expert 参数选择最小全局 rank，保证每个 converted slot 只有一个 owner。 补充 metadata 单测和真实 EP=2 Gloo collective 回归，同时验证不同 expert shard 仍完整保留。 --- .../hf_weight_iterator_bridge.py | 8 +- .../test_mixture_lora_weight_sync.py | 108 ++++++++++++++++++ 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py index 43c133652..25f487958 100644 --- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py +++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py @@ -291,9 +291,9 @@ def _build_param_info_buckets(args, model, collect_adapters=False): else: local_infos[name] = info - # Exchange across EP so every rank has all expert indices. - # Only expert params need src_rank update — non-expert params are - # replicated across EP and already have the correct PP-local src_rank. + # Exchange across EP so every rank has all expert indices. Replicated + # non-expert params also need one canonical source: the converted-tensor + # broadcast protocol requires at most one owner for each parameter slot. if ep_size > 1: ep_infos_list: list[None | tuple[int, dict]] = [None] * ep_size dist.all_gather_object( @@ -305,7 +305,7 @@ def _build_param_info_buckets(args, model, collect_adapters=False): for name, info in infos.items(): if name not in local_infos: local_infos[name] = dataclasses.replace(info, src_rank=src_rank) - elif ".experts." in name and info.src_rank < local_infos[name].src_rank: + elif info.src_rank < local_infos[name].src_rank: local_infos[name] = dataclasses.replace(local_infos[name], src_rank=info.src_rank) # Sort deterministically and split expert / non-expert. diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 3c5c65aae..a14723df2 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -1,10 +1,14 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. +from dataclasses import replace +from datetime import timedelta from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp # These tests exercise the real Megatron Bridge weight conversion and @@ -293,6 +297,110 @@ def test_bridge_hf_iterator_excludes_mixture_parameters(): ] +def test_bridge_ep_metadata_selects_one_owner_for_replicated_non_expert_params(): + from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import _build_param_info_buckets + + non_expert_name = "module.module.decoder.layers.0.self_attention.linear_qkv.weight" + local_expert_name = "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight1" + remote_expert_name = "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight0" + local_params = [ + (non_expert_name, torch.nn.Parameter(torch.zeros(4, 4))), + (local_expert_name, torch.nn.Parameter(torch.zeros(4, 4))), + ] + args = SimpleNamespace(update_weight_buffer_size=1024, num_experts=2) + + def gather_ep(obj, object_list, group=None): + rank, local_infos = obj + assert rank == 3 + assert group == "ep-group" + remote_infos = { + non_expert_name: replace(local_infos[non_expert_name], src_rank=2), + remote_expert_name: replace(local_infos[local_expert_name], name=remote_expert_name, src_rank=2), + } + object_list[:] = [obj, (2, remote_infos)] + + with ( + patch( + "relax.backends.megatron.weight_update.hf_weight_iterator_bridge.named_params_and_buffers", + side_effect=[iter(local_params), iter(local_params)], + ), + patch("torch.distributed.get_rank", return_value=3), + patch("torch.distributed.all_gather_object", side_effect=gather_ep), + patch("megatron.core.mpu.get_pipeline_model_parallel_world_size", return_value=1), + patch("megatron.core.mpu.get_expert_model_parallel_world_size", return_value=2), + patch("megatron.core.mpu.get_expert_model_parallel_group", return_value="ep-group"), + patch("megatron.core.mpu.get_tensor_model_parallel_world_size", return_value=1), + patch("megatron.core.mpu.get_expert_tensor_parallel_world_size", return_value=1), + ): + expert_buckets, non_expert_buckets, _, _ = _build_param_info_buckets(args, model=[]) + + expert_infos = {info.name: info for bucket in expert_buckets for info in bucket} + non_expert_infos = {info.name: info for bucket in non_expert_buckets for info in bucket} + assert non_expert_infos[non_expert_name].src_rank == 2 + assert expert_infos[local_expert_name].src_rank == 3 + assert expert_infos[remote_expert_name].src_rank == 2 + + +def _bridge_ep_owner_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + from megatron.core import parallel_state + + from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import ( + _broadcast_converted_phase, + _build_param_info_buckets, + ) + + parallel_state.initialize_model_parallel(expert_model_parallel_size=world_size) + try: + non_expert_name = "module.module.decoder.layers.0.self_attention.linear_qkv.weight" + expert_name = f"module.module.decoder.layers.0.mlp.experts.linear_fc1.weight{rank}" + local_params = [ + (non_expert_name, torch.nn.Parameter(torch.zeros(4, 4))), + (expert_name, torch.nn.Parameter(torch.zeros(4, 4))), + ] + args = SimpleNamespace(update_weight_buffer_size=1024, num_experts=world_size) + with patch( + "relax.backends.megatron.weight_update.hf_weight_iterator_bridge.named_params_and_buffers", + side_effect=[iter(local_params), iter(local_params)], + ): + expert_buckets, non_expert_buckets, _, _ = _build_param_info_buckets(args, model=[]) + + expert_infos = {info.name: info for bucket in expert_buckets for info in bucket} + non_expert_infos = {info.name: info for bucket in non_expert_buckets for info in bucket} + assert set(expert_infos) == { + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight0", + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight1", + } + info = non_expert_infos[non_expert_name] + assert info.src_rank == 0 + + converted = [("model.layers.0.self_attn.qkv_proj.weight", torch.full((2, 2), 7.0))] + local_converted = [converted if rank == info.src_rank else None] + result = _broadcast_converted_phase( + [info], + local_converted, + torch.device("cpu"), + rank, + parallel_state.get_expert_model_parallel_group(), + ) + assert result[0][0][0] == converted[0][0] + torch.testing.assert_close(result[0][0][1], converted[0][1]) + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def test_bridge_ep_replicated_non_expert_has_one_real_collective_owner(tmp_path): + init_method = f"file://{tmp_path / 'bridge-ep-owner-gloo-init'}" + mp.spawn(_bridge_ep_owner_worker, args=(2, init_method), nprocs=2, join=True) + + def test_first_sync_sends_base_then_routes_and_versions_only_the_final_chunk(): updates = list( iter_mixture_weight_updates( From acee9fa413124073915af88c8912f6c854dbc04d Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:45:17 +0800 Subject: [PATCH 28/41] =?UTF-8?q?fix(megatron):=20=E6=8C=89=20TP=20?= =?UTF-8?q?=E8=AF=AD=E4=B9=89=E5=88=9D=E5=A7=8B=E5=8C=96=20LoRA=20A=20?= =?UTF-8?q?=E5=88=86=E7=89=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mixture expert 复用 MCore 的 CPU master-weight 切分和 CUDA model-parallel RNG tracker，避免 qkv rank-axis shard 重复。 新增 CPU Gloo TP2 与 CUDA NCCL TP2 初始化回归，并保留单卡初始化及 LoRA B 零初始化行为。 --- relax/backends/megatron/mixture_lora.py | 52 ++++++++- .../megatron/test_mixture_lora_distributed.py | 100 ++++++++++++++++++ 2 files changed, 149 insertions(+), 3 deletions(-) diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 830f2d952..46e8779c1 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -555,6 +555,10 @@ def __init__( local_output_size: int | None = None, device: torch.device, dtype: torch.dtype, + input_is_parallel: bool = False, + tp_rank: int = 0, + tp_world_size: int = 1, + use_cpu_initialization: bool | None = None, ) -> None: super().__init__() local_rank = config.rank if local_rank is None else local_rank @@ -566,12 +570,48 @@ def __init__( self.lora_B = nn.Parameter( torch.empty(config.num_experts, local_output_size, config.rank, device=device, dtype=dtype) ) + self._full_rank = config.rank + self._full_input_size = input_size + self._input_is_parallel = input_is_parallel + self._tp_rank = tp_rank + self._tp_world_size = tp_world_size + self._use_cpu_initialization = ( + device.type == "cpu" if use_cpu_initialization is None else use_cpu_initialization + ) self.reset_parameters() def reset_parameters(self) -> None: - # Match the current Bridge LoRA initialization independently for each expert. - for expert_weight in self.lora_A: - nn.init.xavier_uniform_(expert_weight) + if self._tp_world_size == 1: + for expert_weight in self.lora_A: + nn.init.xavier_uniform_(expert_weight) + else: + from megatron.core.tensor_parallel.layers import ( + _initialize_affine_weight_cpu, + _initialize_affine_weight_gpu, + ) + + partition_dim = 1 if self._input_is_parallel else 0 + per_partition_size = self.lora_A.shape[partition_dim + 1] + for expert_weight in self.lora_A: + if self._use_cpu_initialization: + _initialize_affine_weight_cpu( + expert_weight, + self._full_rank, + self._full_input_size, + per_partition_size, + partition_dim, + nn.init.xavier_uniform_, + params_dtype=expert_weight.dtype, + rank=self._tp_rank, + world_size=self._tp_world_size, + skip_set_tensor_parallel_attributes=True, + ) + else: + _initialize_affine_weight_gpu( + expert_weight, + nn.init.xavier_uniform_, + partition_dim=partition_dim, + ) nn.init.zeros_(self.lora_B) @@ -612,6 +652,7 @@ def __init__( tp_group: Any = None, tp_rank: int = 0, tp_world_size: int = 1, + use_cpu_initialization: bool | None = None, ) -> None: super().__init__() if not isinstance(site_id, str) or not site_id.strip(): @@ -650,6 +691,10 @@ def __init__( local_output_size=local_output_size, device=device, dtype=dtype, + input_is_parallel=input_is_parallel, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + use_cpu_initialization=use_cpu_initialization, ) self.router = MixtureLoRARouter(config.num_experts, local_input_size, device=device, dtype=dtype) self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() @@ -799,6 +844,7 @@ def __init__( tp_group=tp_group, tp_rank=tp_rank, tp_world_size=tp_world_size, + use_cpu_initialization=getattr(getattr(to_wrap, "config", None), "use_cpu_initialization", None), ) self._adapter_enabled = True diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index e719b3b6d..3ae6f3ad2 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -780,3 +780,103 @@ def _tensor_parallel_worker(rank: int, world_size: int, init_method: str) -> Non def test_tensor_parallel_forward_and_backward_match_single_rank_reference(tmp_path): init_method = f"file://{tmp_path / 'mixture-lora-tp-gloo-init'}" mp.spawn(_tensor_parallel_worker, args=(2, init_method), nprocs=2, join=True) + + +def _tensor_parallel_initialization_worker(rank: int, world_size: int, init_method: str) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + torch.manual_seed(1234) + config = MixtureLoraConfig( + num_experts=2, + rank=4, + top_k=1, + temperature=1.0, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_qkv",), + ) + adapter = MixtureLoRAAdapter( + config, + "linear_qkv", + input_size=8, + output_size=8, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + tp_group=dist.group.WORLD, + tp_rank=rank, + tp_world_size=world_size, + use_cpu_initialization=True, + ) + shards = [torch.empty_like(adapter.experts.lora_A) for _ in range(world_size)] + dist.all_gather(shards, adapter.experts.lora_A) + full_a = torch.cat(shards, dim=1) + assert full_a.shape == (config.num_experts, config.rank, 8) + assert not torch.equal(shards[0], shards[1]) + finally: + dist.destroy_process_group() + + +def test_tensor_parallel_lora_a_initialization_uses_distinct_rank_blocks(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-tp-init-gloo'}" + mp.spawn(_tensor_parallel_initialization_worker, args=(2, init_method), nprocs=2, join=True) + + +def _tensor_parallel_cuda_initialization_worker(rank: int, world_size: int, init_method: str) -> None: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + try: + model_parallel_cuda_manual_seed( + 1234, + tp_rank=rank, + ep_rank=0, + etp_rank=0, + force_reset_rng=True, + ) + config = MixtureLoraConfig( + num_experts=2, + rank=4, + top_k=1, + temperature=1.0, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=("linear_qkv",), + ) + adapter = MixtureLoRAAdapter( + config, + "linear_qkv", + input_size=8, + output_size=8, + dropout=0.0, + device=torch.device("cuda", rank), + dtype=torch.float32, + tp_group=dist.group.WORLD, + tp_rank=rank, + tp_world_size=world_size, + use_cpu_initialization=False, + ) + shards = [torch.empty_like(adapter.experts.lora_A) for _ in range(world_size)] + dist.all_gather(shards, adapter.experts.lora_A) + assert not torch.equal(shards[0], shards[1]) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_tensor_parallel_cuda_lora_a_initialization_uses_model_parallel_rng(tmp_path): + init_method = f"file://{tmp_path / 'mixture-lora-tp-init-nccl'}" + mp.spawn(_tensor_parallel_cuda_initialization_worker, args=(2, init_method), nprocs=2, join=True) From b7b77f0900afc718bba329334369252b9dc62ffb Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:52:43 +0800 Subject: [PATCH 29/41] =?UTF-8?q?fix(sglang):=20=E9=9A=94=E7=A6=BB=20polic?= =?UTF-8?q?y=20=E7=9A=84=20Mixture=20external=20model?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLangEngine 默认关闭 Mixture 自动注入，仅常规 policy rollout actor 显式开启，避免 GenRM 和 teacher 继承 policy adapter。 无 Mixture 配置时清理残留的 router 环境变量，并补充 policy 与辅助 engine 角色隔离测试。 --- relax/backends/sglang/sglang_engine.py | 4 ++- relax/distributed/ray/rollout.py | 1 + relax/utils/mixture_lora.py | 1 + tests/backends/sglang/test_mixture_lora.py | 36 +++++++++++++++++++++- 4 files changed, 40 insertions(+), 2 deletions(-) diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 2840a6ad8..a87663f01 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -352,6 +352,7 @@ def __init__( sglang_overrides: dict | None = None, num_gpus_per_engine: int | None = None, register_sigterm_handler: bool = False, + enable_mixture_lora_external_model: bool = False, ): self.args = args self.rank = rank @@ -359,6 +360,7 @@ def __init__( self.base_gpu_id = base_gpu_id self.sglang_overrides = sglang_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine + self.enable_mixture_lora_external_model = enable_mixture_lora_external_model self._evicted = threading.Event() self._is_weight_updating: bool = False self._router_worker_id: str | None = None @@ -577,7 +579,7 @@ def _init_normal(self, server_args_dict, *, apply_policy_load_plan: bool = True) # Must be set before launch_server_process() spawns child process # (multiprocessing start_method='spawn'), because the child inherits # the parent's os.environ at spawn time. - mixture_lora_config = build_mixture_lora_config(self.args) + mixture_lora_config = build_mixture_lora_config(self.args) if self.enable_mixture_lora_external_model else None external_pkg = configure_mixture_lora_external_model( mixture_lora_config, getattr(self.args, "sglang_external_model_package", None), diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index ea59f4f63..348f3473a 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -534,6 +534,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis sglang_overrides=self.sglang_overrides, num_gpus_per_engine=self.num_gpus_per_engine, register_sigterm_handler=self.is_scaled_out, + enable_mixture_lora_external_model=True, ) rollout_engines.append((global_rank, rollout_engine)) diff --git a/relax/utils/mixture_lora.py b/relax/utils/mixture_lora.py index de2f9c336..7dd37f36d 100644 --- a/relax/utils/mixture_lora.py +++ b/relax/utils/mixture_lora.py @@ -125,6 +125,7 @@ def configure_mixture_lora_external_model( """Set the validated rollout configuration before spawning SGLang.""" if config is None: + os.environ.pop("RELAX_MIXTURE_LORA_CONFIG", None) return external_package mixture_external_package = "relax.models.qwen3_mixture_lora.sglang" if external_package not in (None, mixture_external_package): diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py index 539c9a470..cbc6c824e 100644 --- a/tests/backends/sglang/test_mixture_lora.py +++ b/tests/backends/sglang/test_mixture_lora.py @@ -51,7 +51,7 @@ def test_mixture_lora_configures_external_package_before_spawn(monkeypatch): def test_single_lora_does_not_enable_external_mixture_model(monkeypatch): - monkeypatch.delenv("RELAX_MIXTURE_LORA_CONFIG", raising=False) + monkeypatch.setenv("RELAX_MIXTURE_LORA_CONFIG", "stale-policy-config") package = configure_mixture_lora_external_model( None, @@ -62,6 +62,40 @@ def test_single_lora_does_not_enable_external_mixture_model(monkeypatch): assert "RELAX_MIXTURE_LORA_CONFIG" not in os.environ +def test_sglang_engine_enables_mixture_external_model_only_for_policy_rollout(monkeypatch): + from unittest.mock import MagicMock + + from relax.backends.sglang import sglang_engine + from relax.backends.sglang.sglang_engine import SGLangEngine + + args = SimpleNamespace(optimize_routing_replay=False, warm_hf_checkpoint_page_cache=False) + policy = SGLangEngine(args, rank=0, enable_mixture_lora_external_model=True) + auxiliary = SGLangEngine(args, rank=0) + for engine in (policy, auxiliary): + engine.server_host = "127.0.0.1" + engine.server_port = 30000 + engine.worker_type = "regular" + engine._skip_router_registration = True + + config = _config() + build_config = MagicMock(return_value=config) + configure_external = MagicMock(return_value=None) + monkeypatch.setattr(sglang_engine, "build_mixture_lora_config", build_config) + monkeypatch.setattr(sglang_engine, "configure_mixture_lora_external_model", configure_external) + monkeypatch.setattr(sglang_engine, "_apply_sglang_policy_load_plan", lambda server_args, _args: server_args) + monkeypatch.setattr(sglang_engine, "ServerArgs", lambda **kwargs: kwargs) + process = MagicMock() + process.is_alive.return_value = False + monkeypatch.setattr(sglang_engine, "launch_server_process", lambda _server_args: process) + + policy._init_normal({}) + auxiliary._init_normal({}) + + build_config.assert_called_once_with(args) + assert configure_external.call_args_list[0].args == (config, None) + assert configure_external.call_args_list[1].args == (None, None) + + def test_mixture_lora_rejects_conflicting_external_package(): with pytest.raises(ValueError, match="requires the Qwen3 external model package"): configure_mixture_lora_external_model(_config(), "custom.other.package") From e9abbafc8d1a08bae0d5b1355998bf8e2171a3aa Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:55:08 +0800 Subject: [PATCH 30/41] =?UTF-8?q?fix(arguments):=20=E6=8B=92=E7=BB=9D=20Mi?= =?UTF-8?q?xture=20=E4=B8=8E=20MTP=20=E5=90=8C=E6=97=B6=E5=90=AF=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mixture rollout mapper 尚未定义 MTP adapter site，启动校验在 mtp_num_layers 大于零时直接报告不支持。 补充启用 MTP 的失败测试和 mtp_num_layers=None 的兼容测试。 --- relax/backends/megatron/arguments.py | 2 ++ .../megatron/test_mixture_lora_arguments.py | 13 +++++++++++++ 2 files changed, 15 insertions(+) diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index bcb2c7fda..2736cc0bf 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -218,6 +218,8 @@ def equal(x, y): mixture_lora_enabled = getattr(args, "lora_num_experts", 1) > 1 if mixture_lora_enabled and is_multimodal: errors.append("Mixture-of-LoRA currently supports text-only base models; multimodal models are unsupported.") + if mixture_lora_enabled and int(getattr(args, "mtp_num_layers", 0) or 0) > 0: + errors.append("Mixture-of-LoRA does not currently support MTP layers.") if is_multimodal and getattr(args, "apply_rope_fusion", False): errors.append( "Multimodal models use multi-axis RoPE (list of tensors) which is incompatible " diff --git a/tests/backends/megatron/test_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py index 6f78658f4..7543e0da0 100644 --- a/tests/backends/megatron/test_mixture_lora_arguments.py +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -26,3 +26,16 @@ def test_mixture_lora_rejects_moe_hf_config(): with pytest.raises(AssertionError, match="dense base models"): _hf_validate_args(args, hf_config) + + +def test_mixture_lora_rejects_mtp_layers(): + args = SimpleNamespace(lora_num_experts=4, mtp_num_layers=1) + + with pytest.raises(AssertionError, match="does not currently support MTP"): + _hf_validate_args(args, SimpleNamespace()) + + +def test_mixture_lora_accepts_disabled_mtp(): + args = SimpleNamespace(lora_num_experts=4, mtp_num_layers=None) + + _hf_validate_args(args, SimpleNamespace()) From 06894c78b34d1f1af617700830ba198ad495e7e3 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Thu, 13 Aug 2026 23:57:55 +0800 Subject: [PATCH 31/41] =?UTF-8?q?fix(arguments):=20=E6=8B=92=E7=BB=9D?= =?UTF-8?q?=E4=BD=8E=E7=B2=BE=E5=BA=A6=20Mixture=20full=20recompute?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 当前 routing context 只包装 MCore tensor-parallel checkpoint，TE FP8/FP4 full recompute 会绕过 aux-loss 重计算路径。 仅拒绝 FP8/FP4 与 full recompute 的组合，并覆盖 selective 与 BF16 full recompute 的兼容边界。 --- relax/backends/megatron/arguments.py | 3 +++ .../megatron/test_mixture_lora_arguments.py | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index 2736cc0bf..4e0814086 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -220,6 +220,9 @@ def equal(x, y): errors.append("Mixture-of-LoRA currently supports text-only base models; multimodal models are unsupported.") if mixture_lora_enabled and int(getattr(args, "mtp_num_layers", 0) or 0) > 0: errors.append("Mixture-of-LoRA does not currently support MTP layers.") + low_precision_training = bool(getattr(args, "fp8", None) or getattr(args, "fp4", None)) + if mixture_lora_enabled and low_precision_training and getattr(args, "recompute_granularity", None) == "full": + errors.append("Mixture-of-LoRA does not currently support FP8/FP4 training with full recompute.") if is_multimodal and getattr(args, "apply_rope_fusion", False): errors.append( "Multimodal models use multi-axis RoPE (list of tensors) which is incompatible " diff --git a/tests/backends/megatron/test_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py index 7543e0da0..df8e6cb2f 100644 --- a/tests/backends/megatron/test_mixture_lora_arguments.py +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -39,3 +39,28 @@ def test_mixture_lora_accepts_disabled_mtp(): args = SimpleNamespace(lora_num_experts=4, mtp_num_layers=None) _hf_validate_args(args, SimpleNamespace()) + + +@pytest.mark.parametrize("precision", ["fp8", "fp4"]) +def test_mixture_lora_rejects_low_precision_full_recompute(precision): + args = SimpleNamespace( + lora_num_experts=4, + recompute_granularity="full", + fp8="e4m3" if precision == "fp8" else None, + fp4="nvfp4" if precision == "fp4" else None, + ) + + with pytest.raises(AssertionError, match="FP8/FP4 training with full recompute"): + _hf_validate_args(args, SimpleNamespace()) + + +def test_mixture_lora_accepts_fp8_without_full_recompute(): + args = SimpleNamespace(lora_num_experts=4, recompute_granularity="selective", fp8="e4m3", fp4=None) + + _hf_validate_args(args, SimpleNamespace()) + + +def test_mixture_lora_accepts_bf16_full_recompute(): + args = SimpleNamespace(lora_num_experts=4, recompute_granularity="full", fp8=None, fp4=None) + + _hf_validate_args(args, SimpleNamespace()) From 39974f820102fece672ca4ff07d7e743c46ca09a Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Fri, 14 Aug 2026 00:00:09 +0800 Subject: [PATCH 32/41] =?UTF-8?q?fix(arguments):=20=E9=99=90=E5=88=B6=20Mi?= =?UTF-8?q?xture=20=E4=BD=BF=E7=94=A8=20dense=20Qwen3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 训练注入和 SGLang rollout mapper 当前只实现 Qwen3 布局，HF 参数校验要求文本 config 的 model_type 为 qwen3。 补充 dense Llama 拒绝与 dense Qwen3 接受测试，并明确既有正向 fixture 的模型类型。 --- relax/backends/megatron/arguments.py | 2 ++ .../megatron/test_mixture_lora_arguments.py | 23 +++++++++++++++---- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index 4e0814086..5a6ddd5de 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -238,6 +238,8 @@ def equal(x, y): if hasattr(hf_config, "text_config"): hf_config = hf_config.text_config + if mixture_lora_enabled and getattr(hf_config, "model_type", None) != "qwen3": + errors.append("Mixture-of-LoRA currently supports Qwen3 base models only.") if mixture_lora_enabled and _is_moe_config(hf_config): errors.append("Mixture-of-LoRA currently supports dense base models; MoE base models are unsupported.") diff --git a/tests/backends/megatron/test_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py index df8e6cb2f..804ba72f4 100644 --- a/tests/backends/megatron/test_mixture_lora_arguments.py +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -22,7 +22,7 @@ def test_mixture_lora_rejects_multimodal_hf_config(): def test_mixture_lora_rejects_moe_hf_config(): args = SimpleNamespace(lora_num_experts=4) - hf_config = SimpleNamespace(num_experts=8) + hf_config = SimpleNamespace(model_type="qwen3", num_experts=8) with pytest.raises(AssertionError, match="dense base models"): _hf_validate_args(args, hf_config) @@ -38,7 +38,7 @@ def test_mixture_lora_rejects_mtp_layers(): def test_mixture_lora_accepts_disabled_mtp(): args = SimpleNamespace(lora_num_experts=4, mtp_num_layers=None) - _hf_validate_args(args, SimpleNamespace()) + _hf_validate_args(args, SimpleNamespace(model_type="qwen3")) @pytest.mark.parametrize("precision", ["fp8", "fp4"]) @@ -51,16 +51,29 @@ def test_mixture_lora_rejects_low_precision_full_recompute(precision): ) with pytest.raises(AssertionError, match="FP8/FP4 training with full recompute"): - _hf_validate_args(args, SimpleNamespace()) + _hf_validate_args(args, SimpleNamespace(model_type="qwen3")) def test_mixture_lora_accepts_fp8_without_full_recompute(): args = SimpleNamespace(lora_num_experts=4, recompute_granularity="selective", fp8="e4m3", fp4=None) - _hf_validate_args(args, SimpleNamespace()) + _hf_validate_args(args, SimpleNamespace(model_type="qwen3")) def test_mixture_lora_accepts_bf16_full_recompute(): args = SimpleNamespace(lora_num_experts=4, recompute_granularity="full", fp8=None, fp4=None) - _hf_validate_args(args, SimpleNamespace()) + _hf_validate_args(args, SimpleNamespace(model_type="qwen3")) + + +def test_mixture_lora_rejects_non_qwen3_dense_model(): + args = SimpleNamespace(lora_num_experts=4) + + with pytest.raises(AssertionError, match="supports Qwen3 base models only"): + _hf_validate_args(args, SimpleNamespace(model_type="llama")) + + +def test_mixture_lora_accepts_dense_qwen3_model(): + args = SimpleNamespace(lora_num_experts=4) + + _hf_validate_args(args, SimpleNamespace(model_type="qwen3")) From af1c40478959ae279f3336a9c24de91057a6dc00 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Fri, 14 Aug 2026 00:10:02 +0800 Subject: [PATCH 33/41] =?UTF-8?q?fix(megatron):=20=E6=98=BE=E5=BC=8F?= =?UTF-8?q?=E4=BC=A0=E9=80=92=20Mixture=20activation=20layout?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Routing context 使用 qkv_format 区分 bshd 与 thd；bshd 始终将 batch-first mask 转为 sequence-first，避免 B 等于 S 时跳过转置。 补充方阵非对称 mask、packed THD 及既有 DP/PP/CP/TP 路径回归。 --- relax/backends/megatron/mixture_lora.py | 28 ++++++++--- relax/backends/megatron/model.py | 1 + tests/backends/megatron/test_mixture_lora.py | 47 +++++++++++++++++++ .../megatron/test_mixture_lora_distributed.py | 5 ++ .../qwen3_mixture_lora/test_sglang_model.py | 4 +- 5 files changed, 76 insertions(+), 9 deletions(-) diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora.py index 46e8779c1..e0f10c02f 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora.py @@ -62,6 +62,7 @@ class MixtureLoRARoutingContext: calculate_per_token_loss: bool objective_scale: float main_loss_backward_scale: torch.Tensor + activation_layout: str context_parallel_group: Any = None context_parallel_world_size: int = 1 is_dummy: bool = False @@ -82,20 +83,35 @@ def __post_init__(self) -> None: raise ValueError("context_parallel_world_size must be positive") if self.context_parallel_world_size > 1 and self.context_parallel_group is None: raise ValueError("context_parallel_group is required when context parallelism is enabled") + if self.activation_layout not in ("bshd", "thd"): + raise ValueError(f"unsupported activation_layout: {self.activation_layout!r}") def response_mask_for(self, x: torch.Tensor) -> torch.Tensor: """Align a batch-first mask with Megatron's activation layout.""" activation_shape = tuple(x.shape[:-1]) mask = self.response_mask - if tuple(mask.shape) == activation_shape: - aligned = mask - elif x.ndim == 3 and mask.ndim == 2 and tuple(mask.shape) == (x.shape[1], x.shape[0]): - aligned = mask.transpose(0, 1) - else: + if x.ndim != 3 or mask.ndim != 2: raise ValueError( - f"response_mask shape {tuple(mask.shape)} does not match activation token layout {activation_shape}" + f"{self.activation_layout} response_mask shape {tuple(mask.shape)} does not match " + f"activation token layout {activation_shape}" ) + if self.activation_layout == "bshd": + expected_mask_shape = (x.shape[1], x.shape[0]) + if tuple(mask.shape) != expected_mask_shape: + raise ValueError( + f"bshd response_mask shape {tuple(mask.shape)} does not match " + f"sequence-first activation shape {activation_shape}" + ) + aligned = mask.transpose(0, 1) + else: + expected_mask_shape = (1, x.shape[0]) + if x.shape[1] != 1 or tuple(mask.shape) != expected_mask_shape: + raise ValueError( + f"thd response_mask shape {tuple(mask.shape)} does not match " + f"packed activation shape {activation_shape}" + ) + aligned = mask.transpose(0, 1) if aligned.device != x.device: raise ValueError("response_mask and routed activation must be on the same device") if self.is_dummy: diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 2ca7fc3ef..0eb6d6405 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -170,6 +170,7 @@ def _build_mixture_lora_routing_context( calculate_per_token_loss=args.calculate_per_token_loss, objective_scale=objective_scale, main_loss_backward_scale=main_loss_backward_scale.detach().clone(), + activation_layout=args.qkv_format, context_parallel_group=( mpu.get_context_parallel_group() if mpu.get_context_parallel_world_size() > 1 else None ), diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 456a2b78d..2b957e2c6 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -368,6 +368,7 @@ def test_routing_context_aligns_batch_first_response_mask_and_records_site(): calculate_per_token_loss=False, objective_scale=0.25, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) x = torch.randn(3, 2, 4) @@ -386,6 +387,46 @@ def test_routing_context_aligns_batch_first_response_mask_and_records_site(): assert get_mixture_lora_routing_context() is None +def test_bshd_square_response_mask_is_always_transposed(): + response_mask = torch.tensor([[1, 0], [1, 1]], dtype=torch.bool) + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=response_mask, + num_microbatches=1, + num_sites=1, + num_samples=2, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", + ) + + aligned = context.response_mask_for(torch.zeros(2, 2, 4)) + + torch.testing.assert_close(aligned, response_mask.transpose(0, 1).reshape(-1)) + + +def test_thd_response_mask_aligns_with_packed_activation(): + response_mask = torch.tensor([[1, 0, 1]], dtype=torch.bool) + context = MixtureLoRARoutingContext( + optimizer_step=0, + microbatch_id=0, + response_mask=response_mask, + num_microbatches=1, + num_sites=1, + num_samples=1, + calculate_per_token_loss=False, + objective_scale=1.0, + main_loss_backward_scale=torch.ones(1), + activation_layout="thd", + ) + + aligned = context.response_mask_for(torch.zeros(3, 1, 4)) + + torch.testing.assert_close(aligned, response_mask.reshape(-1)) + + @pytest.mark.parametrize("calculate_per_token_loss", [False, True]) def test_routing_context_attaches_expected_router_gradient(calculate_per_token_loss): adapter = MixtureLoRAAdapter( @@ -409,6 +450,7 @@ def test_routing_context_attaches_expected_router_gradient(calculate_per_token_l calculate_per_token_loss=calculate_per_token_loss, objective_scale=objective_scale, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) x = torch.randn(3, 2, 4) @@ -444,6 +486,7 @@ def test_dummy_routing_context_records_zero_aux_loss(): calculate_per_token_loss=False, objective_scale=0.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", is_dummy=True, ) @@ -482,6 +525,7 @@ def test_routing_records_pack_and_report_step_metrics(calculate_per_token_loss): calculate_per_token_loss=calculate_per_token_loss, objective_scale=1.0 if calculate_per_token_loss else 0.25, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) x = torch.randn(3, 2, 4) with activate_mixture_lora_routing_context(context): @@ -551,6 +595,7 @@ def test_recompute_replaces_routing_record_instead_of_counting_twice(): calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) with activate_mixture_lora_routing_context(context): @@ -611,6 +656,7 @@ def checkpoint(function, distribute_saved_activations, *args): calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) observed_contexts = [] @@ -800,6 +846,7 @@ def test_mixture_lora_peft_wraps_real_column_parallel_linear_when_available(tmp_ calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) with activate_mixture_lora_routing_context(routing_context): output, bias = transformed.linear_qkv(torch.randn(3, 1, 4)) diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index 3ae6f3ad2..d6caa2da0 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -60,6 +60,7 @@ def _routing_context(site_id: str, mask: torch.Tensor, *, num_sites: int, object calculate_per_token_loss=False, objective_scale=objective_scale, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) with activate_mixture_lora_routing_context(context): adapter(torch.ones(mask.shape[1], mask.shape[0], 4)) @@ -239,6 +240,7 @@ def _new_routing_context( calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.tensor([main_loss_backward_scale]), + activation_layout="bshd", ) @@ -480,6 +482,7 @@ def _context_parallel_balance_worker(rank: int, world_size: int, init_method: st calculate_per_token_loss=calculate_per_token_loss, objective_scale=objective_scale, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", context_parallel_group=dist.group.WORLD, context_parallel_world_size=world_size, ) @@ -692,6 +695,7 @@ def _assert_tp_aux_loss_matches_reference(rank: int, *, input_is_parallel: bool, calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) with activate_mixture_lora_routing_context(reference_context): reference_output = reference(full_x) @@ -739,6 +743,7 @@ def _assert_tp_aux_loss_matches_reference(rank: int, *, input_is_parallel: bool, calculate_per_token_loss=False, objective_scale=1.0, main_loss_backward_scale=torch.ones(1), + activation_layout="bshd", ) with activate_mixture_lora_routing_context(context): output = adapter(local_x) diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 102487523..0db3769f4 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -227,9 +227,7 @@ def test_sglang_installs_adapters_only_on_layers_owned_by_pp_stage(): EntryClass._install_mixture_lora(model) assert not hasattr(missing_layer, "self_attn") - assert local_layer.self_attn.qkv_proj.mixture_lora.site_id == ( - "decoder.layers.1.self_attention.linear_qkv" - ) + assert local_layer.self_attn.qkv_proj.mixture_lora.site_id == ("decoder.layers.1.self_attention.linear_qkv") assert local_layer.self_attn.o_proj.mixture_lora.site_id == "decoder.layers.1.self_attention.linear_proj" From bcf61bdfbf237f55c7ceb05fe5b8c5b4c204f322 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Fri, 14 Aug 2026 02:53:05 +0800 Subject: [PATCH 34/41] =?UTF-8?q?fix(sglang):=20=E4=BF=AE=E5=A4=8D=20Mixtu?= =?UTF-8?q?re=20=E5=9C=A8=20rollout=20PP=20=E4=B8=8B=E7=9A=84=E6=9D=83?= =?UTF-8?q?=E9=87=8D=E5=90=8C=E6=AD=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLang 的 tensor 更新接口按 TP rank 选择序列化 bucket。rollout 使用 TP=1、PP=2 时，两个 pipeline stage 都会读取索引 0，原有 CUDA IPC 句柄只能由首个 stage 所在 GPU 正确打开。 仅对 Mixture-of-LoRA 且 SGLang PP 大于 1 的 colocate 同步改用 host flattened bucket，并临时切换 file_system 共享策略后恢复原值。全参数、单 LoRA 和 PP=1 继续使用原有 CUDA IPC 路径。 增加 host/device 传输与共享策略恢复单测；完成 Qwen3-4B BF16、训练 TP=2、SGLang TP=1/PP=2 的完整 GRPO step，覆盖初始同步、rollout、反向、checkpoint 和训练后同步。 --- .../update_weight_from_tensor.py | 36 +++++++++++++++-- .../test_mixture_lora_weight_sync.py | 40 ++++++++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index 202b00ecb..9e9bfbb6e 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -661,6 +661,7 @@ def _send_hf_params( ipc_gather_src=self._ipc_gather_src, ipc_gather_group=self._ipc_gather_group, weight_version=resolved_weight_version, + use_host_tensors=(self.mixture_lora_enabled and getattr(self.args, "sglang_pp_size", 1) > 1), ) all_refs.extend(refs_colocated) @@ -685,6 +686,7 @@ def _send_to_colocated_engine( ipc_gather_src, ipc_gather_group, weight_version: int | None, + use_host_tensors: bool = False, ) -> tuple[list[ObjectRef], Any]: # Placeholder ranks (GPU slots reserved but no engine) have no gather group. # gather_object is only collective among group members, so we skip entirely. @@ -700,7 +702,10 @@ def _send_to_colocated_engine( # devices. Synchronous copy: this runs on the weight-update path (not the # rollout/train hot path) and FlattenedTensorBucket may flatten on a # different stream — correctness over a few µs. - cur_device = make_current_torch_device() + # SGLang indexes serialized buckets by TP rank. Pipeline stages therefore + # share one entry when TP=1, so device IPC handles would point every stage + # at the first stage's GPU. Host tensors are safe for all PP stages to read. + cur_device = torch.device("cpu") if use_host_tensors else make_current_torch_device() hf_named_tensors = [ (name, tensor.to(cur_device) if tensor.device != cur_device else tensor) for name, tensor in hf_named_tensors ] @@ -726,7 +731,12 @@ def _send_to_colocated_engine( "metadata": metadata, } long_live_tensors.append(flattened_tensor_data) - serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True)) + serialized_tensors.append( + _serialize_flattened_tensor_data( + flattened_tensor_data, + use_file_system_sharing=use_host_tensors, + ) + ) serialized_named_tensors = ( [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None @@ -755,7 +765,10 @@ def _send_to_colocated_engine( if empty_serialized_tensor is None: empty_tensor_data = _empty_flattened_tensor_data(cur_device) long_live_tensors.append(empty_tensor_data) - empty_serialized_tensor = MultiprocessingSerializer.serialize(empty_tensor_data, output_str=True) + empty_serialized_tensor = _serialize_flattened_tensor_data( + empty_tensor_data, + use_file_system_sharing=use_host_tensors, + ) serialized_tensors_for_bucket.append(empty_serialized_tensor) kwargs = { "serialized_named_tensors": serialized_tensors_for_bucket, @@ -767,6 +780,23 @@ def _send_to_colocated_engine( return refs, long_live_tensors +def _serialize_flattened_tensor_data(flattened_tensor_data, *, use_file_system_sharing: bool): + if not use_file_system_sharing: + return MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True) + + # CPU storage serialized with the default file-descriptor strategy is tied + # to the producer process's multiprocessing auth key. SGLang PP workers run + # in separate processes, so use named shared-memory files for this path. + from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy + + previous_strategy = get_sharing_strategy() + try: + set_sharing_strategy("file_system") + return MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True) + finally: + set_sharing_strategy(previous_strategy) + + def _empty_flattened_tensor_data(device): return { "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=device), diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index a14723df2..472439f1d 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -3,7 +3,7 @@ from dataclasses import replace from datetime import timedelta from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest import torch @@ -23,12 +23,50 @@ merge_mixture_lora_tp_shards, ) from relax.backends.megatron.weight_update.update_weight_from_tensor import ( # noqa: E402 + _send_to_colocated_engine, iter_mixture_weight_updates, ) from relax.utils.mixture_lora import MixtureLoraStateSpec # noqa: E402 from relax.utils.types import ParamInfo # noqa: E402 +@pytest.mark.parametrize(("use_host_tensors", "device_lookup_calls"), [(False, 1), (True, 0)]) +def test_colocated_mixture_transfer_selects_device_by_sglang_pipeline_mode( + use_host_tensors, + device_lookup_calls, +): + tensor = torch.ones(2) + engine = MagicMock() + engine.update_weights_from_tensor.remote.return_value = "ref" + + module = "relax.backends.megatron.weight_update.update_weight_from_tensor" + with ( + patch(f"{module}.dist.get_rank", return_value=0), + patch(f"{module}.dist.get_world_size", return_value=1), + patch( + f"{module}.dist.gather_object", + side_effect=lambda value, object_gather_list, **_: object_gather_list.__setitem__(0, value), + ), + patch(f"{module}.make_current_torch_device", return_value=torch.device("cpu")) as current_device, + patch(f"{module}.torch.multiprocessing.get_sharing_strategy", return_value="file_descriptor"), + patch(f"{module}.torch.multiprocessing.set_sharing_strategy") as set_sharing_strategy, + ): + refs, long_lived = _send_to_colocated_engine( + [("weight", tensor)], + ipc_engine=engine, + ipc_gather_src=0, + ipc_gather_group="group", + weight_version=1, + use_host_tensors=use_host_tensors, + ) + + assert refs == ["ref"] + assert long_lived + assert current_device.call_count == device_lookup_calls + expected_sharing_calls = [call("file_system"), call("file_descriptor")] if use_host_tensors else [] + assert set_sharing_strategy.call_args_list == expected_sharing_calls + + def test_selective_offload_uses_cpu_copy_only_after_live_storage_is_released(): cpu_copy = torch.arange(4, dtype=torch.float32) released = torch.empty(0) From ad3e1450e5fdfa7fbb55adcce7391a21ff7ba078 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Fri, 14 Aug 2026 03:36:29 +0800 Subject: [PATCH 35/41] =?UTF-8?q?test(megatron):=20=E8=B7=B3=E8=BF=87?= =?UTF-8?q?=E7=BC=BA=E5=B0=91=20TP=20=E4=BE=9D=E8=B5=96=E7=9A=84=20Mixture?= =?UTF-8?q?=20=E7=94=A8=E4=BE=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TP 初始化回归依赖 Megatron Core 的并行线性层实现，但默认 CPU CI 不安装该可选后端。 在三个 TP 初始化测试入口增加函数级 importorskip：无 Megatron 环境只跳过相关用例，完整训练环境仍执行真实前向、反向和初始化检查。 --- tests/backends/megatron/test_mixture_lora_distributed.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index d6caa2da0..d302c38a7 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -783,6 +783,7 @@ def _tensor_parallel_worker(rank: int, world_size: int, init_method: str) -> Non def test_tensor_parallel_forward_and_backward_match_single_rank_reference(tmp_path): + pytest.importorskip("megatron.core.tensor_parallel.layers") init_method = f"file://{tmp_path / 'mixture-lora-tp-gloo-init'}" mp.spawn(_tensor_parallel_worker, args=(2, init_method), nprocs=2, join=True) @@ -829,6 +830,7 @@ def _tensor_parallel_initialization_worker(rank: int, world_size: int, init_meth def test_tensor_parallel_lora_a_initialization_uses_distinct_rank_blocks(tmp_path): + pytest.importorskip("megatron.core.tensor_parallel.layers") init_method = f"file://{tmp_path / 'mixture-lora-tp-init-gloo'}" mp.spawn(_tensor_parallel_initialization_worker, args=(2, init_method), nprocs=2, join=True) @@ -883,5 +885,6 @@ def _tensor_parallel_cuda_initialization_worker(rank: int, world_size: int, init @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") def test_tensor_parallel_cuda_lora_a_initialization_uses_model_parallel_rng(tmp_path): + pytest.importorskip("megatron.core.tensor_parallel.layers") init_method = f"file://{tmp_path / 'mixture-lora-tp-init-nccl'}" mp.spawn(_tensor_parallel_cuda_initialization_worker, args=(2, init_method), nprocs=2, join=True) From 1deb023d4a91306e0c5f974a65dcaf324de25011 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:36:54 +0800 Subject: [PATCH 36/41] =?UTF-8?q?test(megatron):=20=E4=BF=AE=E6=AD=A3=20Mi?= =?UTF-8?q?xture=20=E6=9D=83=E9=87=8D=E5=90=8C=E6=AD=A5=E7=94=A8=E4=BE=8B?= =?UTF-8?q?=E7=9A=84=E6=89=93=E6=A1=A9=E7=9B=AE=E6=A0=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit hf_weight_iterator_bridge 在模块导入时执行 from megatron.core import mpu，绑定的是导入那一刻 sys.modules 里的对象。同目录的 test_dtype_codes.py 与 test_broadcast_converted.py 会在各自导入阶段装入 megatron.core 桩模块，所以按目录运行 tests/backends/megatron/weight_update/ 时，bridge 绑定到的是桩 mpu，而用例改写的是 megatron.core.mpu，两者并非同一个对象，两个 bridge 用例因此报 AttributeError；单独运行该文件反而正常。 改为用 patch.object 直接替换 bridge 模块上的 mpu 属性，无论它绑定的是真实模块还是桩模块都能命中。打桩返回值与原先完全一致，只是换了作用目标，未改动任何被测代码。 按目录运行 tests/backends/megatron/weight_update/ 由 73 passed、2 failed 变为 75 passed；单文件运行 23 passed。 --- .../test_mixture_lora_weight_sync.py | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 472439f1d..1cabc08bf 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -312,6 +312,21 @@ def gather_single_process(obj, object_list, group=None): assert [info.name for info in infos] == ["module.module.decoder.layers.0.self_attention.linear_qkv.weight"] +def _patch_bridge_mpu(**parallel_state): + """Patch the ``mpu`` handle the bridge bound at import time. + + ``hf_weight_iterator_bridge`` does ``from megatron.core import mpu``, and + sibling suites import it while a stub ``megatron.core`` sits in + ``sys.modules``, so patching ``megatron.core.mpu`` can leave the object the + bridge actually calls untouched. + """ + + from relax.backends.megatron.weight_update import hf_weight_iterator_bridge + + stub = SimpleNamespace(**{name: (lambda value=value: value) for name, value in parallel_state.items()}) + return patch.object(hf_weight_iterator_bridge, "mpu", stub) + + def test_bridge_hf_iterator_excludes_mixture_parameters(): from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import _build_param_info_buckets @@ -323,9 +338,11 @@ def test_bridge_hf_iterator_excludes_mixture_parameters(): side_effect=[iter(vanilla), iter(_named_parameters())], ), patch("torch.distributed.get_rank", return_value=0), - patch("megatron.core.mpu.get_pipeline_model_parallel_world_size", return_value=1), - patch("megatron.core.mpu.get_expert_model_parallel_world_size", return_value=1), - patch("megatron.core.mpu.get_tensor_model_parallel_world_size", return_value=1), + _patch_bridge_mpu( + get_pipeline_model_parallel_world_size=1, + get_expert_model_parallel_world_size=1, + get_tensor_model_parallel_world_size=1, + ), ): expert_buckets, base_buckets, _, _ = _build_param_info_buckets(args, model=[]) @@ -364,11 +381,13 @@ def gather_ep(obj, object_list, group=None): ), patch("torch.distributed.get_rank", return_value=3), patch("torch.distributed.all_gather_object", side_effect=gather_ep), - patch("megatron.core.mpu.get_pipeline_model_parallel_world_size", return_value=1), - patch("megatron.core.mpu.get_expert_model_parallel_world_size", return_value=2), - patch("megatron.core.mpu.get_expert_model_parallel_group", return_value="ep-group"), - patch("megatron.core.mpu.get_tensor_model_parallel_world_size", return_value=1), - patch("megatron.core.mpu.get_expert_tensor_parallel_world_size", return_value=1), + _patch_bridge_mpu( + get_pipeline_model_parallel_world_size=1, + get_expert_model_parallel_world_size=2, + get_expert_model_parallel_group="ep-group", + get_tensor_model_parallel_world_size=1, + get_expert_tensor_parallel_world_size=1, + ), ): expert_buckets, non_expert_buckets, _, _ = _build_param_info_buckets(args, model=[]) From 948d38ffcca104ac1a81706697501327ffc85085 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:36:54 +0800 Subject: [PATCH 37/41] =?UTF-8?q?refactor(mixture-lora):=20=E6=8C=89?= =?UTF-8?q?=E5=90=8E=E7=AB=AF=E5=BD=92=E5=B1=9E=E9=87=8D=E5=91=BD=E5=90=8D?= =?UTF-8?q?=E4=B8=A4=E4=B8=AA=E6=A8=A1=E5=9D=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit relax/utils/mixture_lora.py 与 relax/backends/megatron/mixture_lora.py 同名不同包，导入语句和调用栈里难以一眼分辨，而 SGLang 侧还要再加一个同主题模块，混淆只会更严重。 前者改名为 relax/utils/mixture_lora_common.py，只放后端无关的配置、路由计算、参数命名与传输描述，不允许导入任何训练或推理后端；后者改名为 relax/backends/megatron/mixture_lora_modules.py，只放依赖 Megatron 并行状态的模块与 Bridge 注入。两个文件的模块说明同步写清这条边界，其余改动全部是导入路径同步，没有行为变化。 pytest tests/utils/test_mixture_lora_routing.py tests/backends/megatron/test_mixture_lora.py tests/backends/sglang/test_router_registration.py 为 97 passed。 --- relax/backends/megatron/loss.py | 2 +- .../{mixture_lora.py => mixture_lora_modules.py} | 9 +++++++-- relax/backends/megatron/model.py | 2 +- relax/backends/megatron/model_provider.py | 2 +- .../backends/megatron/weight_update/mixture_lora_sync.py | 2 +- relax/backends/sglang/sglang_engine.py | 2 +- relax/models/qwen3_mixture_lora/sglang/model.py | 2 +- relax/utils/megatron_peft_utils.py | 2 +- relax/utils/{mixture_lora.py => mixture_lora_common.py} | 8 +++++++- tests/backends/megatron/test_mixture_lora.py | 4 ++-- .../megatron/test_mixture_lora_checkpoint_distributed.py | 4 ++-- tests/backends/megatron/test_mixture_lora_distributed.py | 4 ++-- tests/backends/megatron/test_model_provider_vpp.py | 4 ++-- .../weight_update/test_mixture_lora_weight_sync.py | 2 +- tests/backends/sglang/test_mixture_lora.py | 2 +- tests/backends/sglang/test_router_registration.py | 4 ++-- tests/models/qwen3_mixture_lora/test_sglang_model.py | 4 ++-- tests/utils/test_mixture_lora_routing.py | 2 +- 18 files changed, 36 insertions(+), 25 deletions(-) rename relax/backends/megatron/{mixture_lora.py => mixture_lora_modules.py} (99%) rename relax/utils/{mixture_lora.py => mixture_lora_common.py} (98%) diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 5a6855046..d0c9d08f5 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -42,7 +42,7 @@ maybe_padded_total_lengths, slice_log_prob_with_cp, ) -from .mixture_lora import get_microbatch_objective_scale +from .mixture_lora_modules import get_microbatch_objective_scale def get_responses( diff --git a/relax/backends/megatron/mixture_lora.py b/relax/backends/megatron/mixture_lora_modules.py similarity index 99% rename from relax/backends/megatron/mixture_lora.py rename to relax/backends/megatron/mixture_lora_modules.py index e0f10c02f..57c517644 100644 --- a/relax/backends/megatron/mixture_lora.py +++ b/relax/backends/megatron/mixture_lora_modules.py @@ -1,6 +1,11 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Megatron model modules and Bridge injection for Mixture-of-LoRA.""" +"""Megatron-side Mixture-of-LoRA modules, routing context, and Bridge +injection. + +Everything here depends on Megatron parallel state; the backend-independent +configuration and routing math lives in ``relax.utils.mixture_lora_common``. +""" import math import re @@ -14,7 +19,7 @@ import torch.nn.functional as F from torch import nn -from relax.utils.mixture_lora import ( +from relax.utils.mixture_lora_common import ( MixtureLoraConfig, RoutingDecision, RoutingStatistics, diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 0eb6d6405..7fa58f0fb 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -47,7 +47,7 @@ from .checkpoint import load_checkpoint, save_checkpoint from .data import DataIterator, get_batch from .loss import loss_function -from .mixture_lora import ( +from .mixture_lora_modules import ( MixtureLoRARoutingContext, MixtureParallelLinearAdapter, activate_mixture_lora_routing_context, diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 19479e1a7..60406d84d 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -480,7 +480,7 @@ def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwarg try: if is_mixture_lora_enabled(args): - from relax.backends.megatron.mixture_lora import ( + from relax.backends.megatron.mixture_lora_modules import ( build_mixture_lora_peft, ensure_mixture_lora_recompute_inputs_grad, install_mixture_lora_checkpoint_context, diff --git a/relax/backends/megatron/weight_update/mixture_lora_sync.py b/relax/backends/megatron/weight_update/mixture_lora_sync.py index 9a6392f61..dbbbf7414 100644 --- a/relax/backends/megatron/weight_update/mixture_lora_sync.py +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -13,7 +13,7 @@ from relax.backends.megatron.misc_utils import strip_param_name_prefix from relax.utils import device as device_utils from relax.utils.megatron_peft_utils import build_mixture_lora_config, is_mixture_lora_param -from relax.utils.mixture_lora import ( +from relax.utils.mixture_lora_common import ( MixtureLoraStateSpec, build_mixture_lora_state_specs, ) diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index a87663f01..f3b33fa04 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -38,7 +38,7 @@ is_lora_enabled, is_mixture_lora_enabled, ) -from relax.utils.mixture_lora import configure_mixture_lora_external_model +from relax.utils.mixture_lora_common import configure_mixture_lora_external_model from relax.utils.model_source import ModelSource, SGLangLoadPlan from relax.utils.s3_model_loader import ( get_s3_model_cached_path, diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index c4b77656d..c57c13702 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -12,7 +12,7 @@ from torch.nn import functional as F from relax.utils.env import Envs -from relax.utils.mixture_lora import ( +from relax.utils.mixture_lora_common import ( DenseRoutedLoRAExecutor, MixtureLoraConfig, deserialize_mixture_lora_config, diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index ad6a0516c..3d85c3ba8 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/relax/utils/megatron_peft_utils.py @@ -7,7 +7,7 @@ import torch -from relax.utils.mixture_lora import MixtureLoraConfig +from relax.utils.mixture_lora_common import MixtureLoraConfig # Fixed name under which the trained policy LoRA adapter is registered on the rollout diff --git a/relax/utils/mixture_lora.py b/relax/utils/mixture_lora_common.py similarity index 98% rename from relax/utils/mixture_lora.py rename to relax/utils/mixture_lora_common.py index 7dd37f36d..5b80d1866 100644 --- a/relax/utils/mixture_lora.py +++ b/relax/utils/mixture_lora_common.py @@ -1,6 +1,12 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Backend-independent routing primitives for Mixture-of-LoRA.""" +"""Backend-independent Mixture-of-LoRA contracts. + +Configuration, routing math, parameter naming, and transport specs shared by +the Megatron training modules (``mixture_lora_modules``) and the SGLang rollout +modules (``relax.models.mixture_lora_sglang``). Nothing here may import a +training or inference backend. +""" import json import math diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 2b957e2c6..7574ad09e 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -10,7 +10,7 @@ import torch import torch.nn.functional as F -from relax.backends.megatron.mixture_lora import ( +from relax.backends.megatron.mixture_lora_modules import ( MixtureLoRAAdapter, MixtureLoRARoutingContext, MixtureParallelLinearAdapter, @@ -24,7 +24,7 @@ pack_mixture_lora_routing_records, ) from relax.utils import megatron_bridge_utils -from relax.utils.mixture_lora import MixtureLoraConfig, compute_routing_statistics +from relax.utils.mixture_lora_common import MixtureLoraConfig, compute_routing_statistics def _config(*, num_experts=3, top_k=2, rank=2, alpha=4.0): diff --git a/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py b/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py index 0967f6818..5f362d0c4 100644 --- a/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_checkpoint_distributed.py @@ -9,8 +9,8 @@ import torch.distributed as dist import torch.multiprocessing as mp -from relax.backends.megatron.mixture_lora import MixtureParallelLinearAdapter -from relax.utils.mixture_lora import MixtureLoraConfig +from relax.backends.megatron.mixture_lora_modules import MixtureParallelLinearAdapter +from relax.utils.mixture_lora_common import MixtureLoraConfig def _config() -> MixtureLoraConfig: diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index d302c38a7..d2ed90bee 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -10,14 +10,14 @@ import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel -from relax.backends.megatron.mixture_lora import ( +from relax.backends.megatron.mixture_lora_modules import ( MixtureLoRAAdapter, MixtureLoRARoutingContext, activate_mixture_lora_routing_context, mixture_lora_metrics_from_packed_records, pack_mixture_lora_routing_records, ) -from relax.utils.mixture_lora import ( +from relax.utils.mixture_lora_common import ( DenseRoutedLoRAExecutor, MixtureLoraConfig, compute_routing_statistics, diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index e59ad1c7e..940215987 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -256,7 +256,7 @@ def test_model_provider_uses_mixture_lora_for_multiple_experts(monkeypatch): model = object() config = object() calls = [] - mixture_module = types.ModuleType("relax.backends.megatron.mixture_lora") + mixture_module = types.ModuleType("relax.backends.megatron.mixture_lora_modules") def build_mixture_lora_peft(received_config, dropout, vp_stage=None): calls.append((received_config, dropout, vp_stage)) @@ -265,7 +265,7 @@ def build_mixture_lora_peft(received_config, dropout, vp_stage=None): mixture_module.build_mixture_lora_peft = build_mixture_lora_peft mixture_module.ensure_mixture_lora_recompute_inputs_grad = lambda model: None mixture_module.install_mixture_lora_checkpoint_context = lambda: None - monkeypatch.setitem(sys.modules, "relax.backends.megatron.mixture_lora", mixture_module) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.mixture_lora_modules", mixture_module) monkeypatch.setattr(module, "build_mixture_lora_config", lambda args: config) monkeypatch.setattr(module, "build_lora_peft", lambda args: pytest.fail("single LoRA factory was called")) monkeypatch.setattr(module, "validate_and_count_mixture_lora_parameters", lambda model: (100, 20, 4, 124)) diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 1cabc08bf..64c762127 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -26,7 +26,7 @@ _send_to_colocated_engine, iter_mixture_weight_updates, ) -from relax.utils.mixture_lora import MixtureLoraStateSpec # noqa: E402 +from relax.utils.mixture_lora_common import MixtureLoraStateSpec # noqa: E402 from relax.utils.types import ParamInfo # noqa: E402 diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py index cbc6c824e..287227663 100644 --- a/tests/backends/sglang/test_mixture_lora.py +++ b/tests/backends/sglang/test_mixture_lora.py @@ -16,7 +16,7 @@ _compute_server_args, _configure_external_model_environment, ) -from relax.utils.mixture_lora import ( # noqa: E402 +from relax.utils.mixture_lora_common import ( # noqa: E402 MixtureLoraConfig, configure_mixture_lora_external_model, deserialize_mixture_lora_config, diff --git a/tests/backends/sglang/test_router_registration.py b/tests/backends/sglang/test_router_registration.py index 721f7c197..956987fd1 100644 --- a/tests/backends/sglang/test_router_registration.py +++ b/tests/backends/sglang/test_router_registration.py @@ -67,9 +67,9 @@ def sglang_engine_module(monkeypatch): megatron_peft_utils.is_mixture_lora_enabled = lambda _args: False monkeypatch.setitem(sys.modules, "relax.utils.megatron_peft_utils", megatron_peft_utils) - mixture_lora = ModuleType("relax.utils.mixture_lora") + mixture_lora = ModuleType("relax.utils.mixture_lora_common") mixture_lora.configure_mixture_lora_external_model = lambda *_args, **_kwargs: None - monkeypatch.setitem(sys.modules, "relax.utils.mixture_lora", mixture_lora) + monkeypatch.setitem(sys.modules, "relax.utils.mixture_lora_common", mixture_lora) sys.modules.pop("relax.backends.sglang.sglang_engine", None) module = importlib.import_module("relax.backends.sglang.sglang_engine") diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 0db3769f4..4159fdd54 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -11,14 +11,14 @@ pytest.importorskip("megatron.bridge") pytest.importorskip("sglang.srt.models.qwen3") -from relax.backends.megatron.mixture_lora import MixtureLoRAAdapter # noqa: E402 +from relax.backends.megatron.mixture_lora_modules import MixtureLoRAAdapter # noqa: E402 from relax.models.qwen3_mixture_lora.sglang.model import ( # noqa: E402 EntryClass, SGLangMixtureLoRA, attach_sglang_mixture_lora, load_sglang_mixture_lora_weights, ) -from relax.utils.mixture_lora import MixtureLoraConfig # noqa: E402 +from relax.utils.mixture_lora_common import MixtureLoraConfig # noqa: E402 def _config(): diff --git a/tests/utils/test_mixture_lora_routing.py b/tests/utils/test_mixture_lora_routing.py index 571c383b8..5b375d756 100644 --- a/tests/utils/test_mixture_lora_routing.py +++ b/tests/utils/test_mixture_lora_routing.py @@ -5,7 +5,7 @@ import pytest import torch -from relax.utils.mixture_lora import ( +from relax.utils.mixture_lora_common import ( MIXTURE_LORA_SCHEMA_VERSION, DenseRoutedLoRAExecutor, MixtureLoraConfig, From e1560eab76adb459077d467c63775f850e0c4af5 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:40:32 +0800 Subject: [PATCH 38/41] =?UTF-8?q?fix(mixture-lora):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E8=B7=AF=E7=94=B1=E9=80=82=E9=85=8D=E5=99=A8=E7=9A=84=E5=BC=A0?= =?UTF-8?q?=E9=87=8F=E5=B9=B6=E8=A1=8C=E6=A0=87=E8=AE=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 张量并行开启时，专家权重通过按专家切出的视图逐个初始化。视图与 Parameter 共享存储，数值能写回去，但 Megatron 在视图上设置的 tensor_model_parallel 属性会随视图一起消失。结果是 lora_A、lora_B 以及行并行站点上的 router.weight 都保持默认的未标记状态，梯度裁剪把它们当成各 rank 的重复参数，全局梯度范数只统计其中一份，裁剪系数偏大；行并行与列并行的切分轴不同，这个偏差还会随站点类型变化。 初始化完成后显式给 lora_A、lora_B 打标记，router 在构造时按站点类型打标记：列并行站点的 router 是复制参数，标记为重复；行并行站点的 router 沿输入轴切分，每一份都要计入全局范数。已有标记的参数不重复覆盖。 切分轴此前在适配器初始化、sharded checkpoint 的 axis_map、rollout 权重同步三处各写了一份，容易改一处漏两处。统一收敛到 relax/utils/mixture_lora_common.py 的 mixture_lora_tp_partition_dims，三处都从这张表读取，None 表示复制参数。 新增两个单测覆盖行并行与列并行下三类参数的标记结果；分布式用例补充 param_is_not_tensor_parallel_duplicate 断言，确认切分参数确实参与全局梯度范数。8 卡环境下 tests/backends/megatron 三个 Mixture 分布式用例集 24 passed；CPU 侧 tests/backends/megatron/test_mixture_lora.py 与 tests/utils/test_mixture_lora_routing.py 88 passed，权重同步与 SGLang 用例 29 passed。 --- .../backends/megatron/mixture_lora_modules.py | 63 +++++++++++++++++-- .../weight_update/mixture_lora_sync.py | 20 +++--- relax/utils/mixture_lora_common.py | 21 +++++++ tests/backends/megatron/test_mixture_lora.py | 60 ++++++++++++++++++ .../megatron/test_mixture_lora_distributed.py | 9 +++ 5 files changed, 154 insertions(+), 19 deletions(-) diff --git a/relax/backends/megatron/mixture_lora_modules.py b/relax/backends/megatron/mixture_lora_modules.py index 57c517644..3c9b4067f 100644 --- a/relax/backends/megatron/mixture_lora_modules.py +++ b/relax/backends/megatron/mixture_lora_modules.py @@ -24,6 +24,7 @@ RoutingDecision, RoutingStatistics, compute_routing_statistics, + mixture_lora_tp_partition_dims, route_topk, ) @@ -562,6 +563,29 @@ def forward( return (delta * scale).reshape(*input_shape, expert_outputs.shape[-1]) +def mark_routed_tensor_parallel_shard(parameter: nn.Parameter, *, partition_dim: int | None) -> None: + """Publish Megatron's tensor-parallel metadata on a routed parameter. + + Grad-norm clipping and sharded checkpointing read ``tensor_model_parallel`` + off the Parameter itself. Initializing through per-expert slice views + leaves the Parameter unmarked, so Megatron falls back to its + ``tensor_model_parallel=False`` default, treats every sharded adapter + tensor as a tensor-parallel duplicate, and clips against a partial global + norm. + """ + + from megatron.core.tensor_parallel.layers import set_tensor_model_parallel_attributes + + if hasattr(parameter, "tensor_model_parallel"): + return + set_tensor_model_parallel_attributes( + parameter, + partition_dim is not None, + -1 if partition_dim is None else partition_dim, + 1, + ) + + class MixtureLoRAExperts(nn.Module): """LoRA expert parameters stored in one stable logical layout.""" @@ -602,6 +626,7 @@ def __init__( self.reset_parameters() def reset_parameters(self) -> None: + partition_dims = mixture_lora_tp_partition_dims(self._input_is_parallel) if self._tp_world_size == 1: for expert_weight in self.lora_A: nn.init.xavier_uniform_(expert_weight) @@ -611,7 +636,11 @@ def reset_parameters(self) -> None: _initialize_affine_weight_gpu, ) - partition_dim = 1 if self._input_is_parallel else 0 + # The loop below initializes per-expert slice views: they share + # storage with lora_A so the values land in the Parameter, but any + # attribute Megatron sets on a view dies with it. lora_A and lora_B + # are therefore marked explicitly once the values are in place. + partition_dim = partition_dims["experts.lora_A"] - 1 per_partition_size = self.lora_A.shape[partition_dim + 1] for expert_weight in self.lora_A: if self._use_cpu_initialization: @@ -634,6 +663,9 @@ def reset_parameters(self) -> None: partition_dim=partition_dim, ) nn.init.zeros_(self.lora_B) + if self._tp_world_size > 1: + mark_routed_tensor_parallel_shard(self.lora_A, partition_dim=partition_dims["experts.lora_A"]) + mark_routed_tensor_parallel_shard(self.lora_B, partition_dim=partition_dims["experts.lora_B"]) class MixtureLoRARouter(nn.Module): @@ -646,10 +678,20 @@ def __init__( *, device: torch.device, dtype: torch.dtype, + input_is_parallel: bool = False, + tp_world_size: int = 1, ) -> None: super().__init__() self.weight = nn.Parameter(torch.empty(num_experts, input_size, device=device, dtype=dtype)) nn.init.normal_(self.weight, mean=0.0, std=0.02) + if tp_world_size > 1: + # Column-parallel sites keep a replicated router (marked as a + # tensor-parallel duplicate); row-parallel sites shard it along the + # input axis and must contribute every shard to the global norm. + mark_routed_tensor_parallel_shard( + self.weight, + partition_dim=mixture_lora_tp_partition_dims(input_is_parallel)["router.weight"], + ) def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear(x.float(), self.weight.float()) @@ -717,7 +759,14 @@ def __init__( tp_world_size=tp_world_size, use_cpu_initialization=use_cpu_initialization, ) - self.router = MixtureLoRARouter(config.num_experts, local_input_size, device=device, dtype=dtype) + self.router = MixtureLoRARouter( + config.num_experts, + local_input_size, + device=device, + dtype=dtype, + input_is_parallel=input_is_parallel, + tp_world_size=tp_world_size, + ) self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() self.executor = MegatronDenseRoutedLoRAExecutor( input_is_parallel=input_is_parallel, @@ -938,11 +987,12 @@ def sharded_state_dict( sharded_state = self.to_wrap.sharded_state_dict(prefix, sharded_offsets, metadata) adapter_state = self.mixture_lora.state_dict(prefix="", keep_vars=True) axis_map = { - "experts.lora_A": 2 if self.mixture_lora.input_is_parallel else 1, - "experts.lora_B": 1, + parameter_kind: partition_dim + for parameter_kind, partition_dim in mixture_lora_tp_partition_dims( + self.mixture_lora.input_is_parallel + ).items() + if partition_dim is not None } - if self.mixture_lora.input_is_parallel: - axis_map["router.weight"] = 1 dp_cp_group = ( metadata["dp_cp_group"] if metadata is not None and metadata.get("dp_cp_group") is not None @@ -1042,6 +1092,7 @@ def transform( "get_microbatch_objective_scale", "get_mixture_lora_routing_context", "install_mixture_lora_checkpoint_context", + "mark_routed_tensor_parallel_shard", "mixture_lora_metrics_from_packed_records", "pack_mixture_lora_routing_records", ] diff --git a/relax/backends/megatron/weight_update/mixture_lora_sync.py b/relax/backends/megatron/weight_update/mixture_lora_sync.py index dbbbf7414..e6224efb7 100644 --- a/relax/backends/megatron/weight_update/mixture_lora_sync.py +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -16,6 +16,7 @@ from relax.utils.mixture_lora_common import ( MixtureLoraStateSpec, build_mixture_lora_state_specs, + mixture_lora_tp_partition_dims, ) from .common import named_params_and_buffers @@ -84,19 +85,12 @@ def _parameter_kind(parameter_name: str) -> str: def _tp_shard_dim(site_id: str, parameter_kind: str) -> int | None: target = site_id.rsplit(".", maxsplit=1)[-1] - if target == "linear_qkv": - return { - "experts.lora_A": 1, - "experts.lora_B": 1, - "router.weight": None, - }[parameter_kind] - if target == "linear_proj": - return { - "experts.lora_A": 2, - "experts.lora_B": 1, - "router.weight": 1, - }[parameter_kind] - raise ValueError(f"Unsupported Mixture-of-LoRA site: {site_id!r}") + if target not in ("linear_qkv", "linear_proj"): + raise ValueError(f"Unsupported Mixture-of-LoRA site: {site_id!r}") + # linear_proj consumes a tensor-parallel-sharded input; linear_qkv produces + # a sharded output. The shard axes themselves come from the shared table so + # training, checkpointing and rollout sync cannot drift apart. + return mixture_lora_tp_partition_dims(input_is_parallel=target == "linear_proj")[parameter_kind] def _qkv_lora_b_to_sglang( diff --git a/relax/utils/mixture_lora_common.py b/relax/utils/mixture_lora_common.py index 5b80d1866..a84bb2ca1 100644 --- a/relax/utils/mixture_lora_common.py +++ b/relax/utils/mixture_lora_common.py @@ -355,6 +355,27 @@ def build_mixture_lora_state_specs( ) +def mixture_lora_tp_partition_dims(input_is_parallel: bool) -> dict[MixtureLoraParameterKind, int | None]: + """Map each routed parameter to its tensor-parallel shard axis. + + The axes refer to the stable stored layout described by + ``build_mixture_lora_state_specs``; ``None`` marks a replicated parameter. + Row-parallel sites (``input_is_parallel``) shard ``lora_A`` and the router + on their input axis, column-parallel sites shard ``lora_A`` on its rank + axis and replicate the router. ``lora_B`` is always output-sharded. + + Every consumer that has to reason about sharding — parameter marking, + sharded checkpoints, and rollout weight sync — reads its axes from here so + the three cannot drift apart. + """ + + return { + "experts.lora_A": 2 if input_is_parallel else 1, + "experts.lora_B": 1, + "router.weight": 1 if input_is_parallel else None, + } + + @dataclass(frozen=True) class RoutingDecision: """Token-level Top-K routing results. diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py index 7574ad09e..dfcc518e2 100644 --- a/tests/backends/megatron/test_mixture_lora.py +++ b/tests/backends/megatron/test_mixture_lora.py @@ -12,6 +12,8 @@ from relax.backends.megatron.mixture_lora_modules import ( MixtureLoRAAdapter, + MixtureLoRAExperts, + MixtureLoRARouter, MixtureLoRARoutingContext, MixtureParallelLinearAdapter, activate_mixture_lora_routing_context, @@ -163,6 +165,64 @@ def test_mixture_lora_parameter_layout_and_initialization(): assert adapter.executor.state_dict() == {} +@pytest.mark.parametrize( + ("input_is_parallel", "lora_a_partition_dim", "router_is_parallel"), + [(False, 1, False), (True, 2, True)], +) +def test_mixture_lora_parameters_publish_tensor_parallel_metadata( + input_is_parallel, lora_a_partition_dim, router_is_parallel +): + pytest.importorskip("megatron.core.tensor_parallel.layers") + config = _config(num_experts=4, rank=4) + experts = MixtureLoRAExperts( + config, + 8, + 6, + local_rank=config.rank if input_is_parallel else config.rank // 2, + local_input_size=4 if input_is_parallel else 8, + local_output_size=6 if input_is_parallel else 3, + device=torch.device("cpu"), + dtype=torch.float32, + input_is_parallel=input_is_parallel, + tp_rank=1, + tp_world_size=2, + use_cpu_initialization=True, + ) + router = MixtureLoRARouter( + config.num_experts, + 4 if input_is_parallel else 8, + device=torch.device("cpu"), + dtype=torch.float32, + input_is_parallel=input_is_parallel, + tp_world_size=2, + ) + + # Grad-norm clipping reads these off the Parameter, not off the per-expert + # slice views the initializers run on. + assert experts.lora_A.tensor_model_parallel is True + assert experts.lora_A.partition_dim == lora_a_partition_dim + assert experts.lora_B.tensor_model_parallel is True + assert experts.lora_B.partition_dim == 1 + assert router.weight.tensor_model_parallel is router_is_parallel + assert torch.count_nonzero(experts.lora_A) > 0 + + +def test_mixture_lora_parameters_stay_unmarked_without_tensor_parallelism(): + adapter = MixtureLoRAAdapter( + _config(), + "linear_proj", + 6, + 7, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert not hasattr(adapter.experts.lora_A, "tensor_model_parallel") + assert not hasattr(adapter.experts.lora_B, "tensor_model_parallel") + assert not hasattr(adapter.router.weight, "tensor_model_parallel") + + def test_mixture_lora_wrapper_freezes_base_and_routes_gradients(): torch.manual_seed(11) base = _TupleLinear(4, 5) diff --git a/tests/backends/megatron/test_mixture_lora_distributed.py b/tests/backends/megatron/test_mixture_lora_distributed.py index d2ed90bee..fbab86c71 100644 --- a/tests/backends/megatron/test_mixture_lora_distributed.py +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -825,6 +825,15 @@ def _tensor_parallel_initialization_worker(rank: int, world_size: int, init_meth full_a = torch.cat(shards, dim=1) assert full_a.shape == (config.num_experts, config.rank, 8) assert not torch.equal(shards[0], shards[1]) + + # Every shard of a sharded adapter parameter has to reach the global + # grad-norm reduction; only the replicated router is a duplicate that + # rank 0 alone contributes. + from megatron.core.tensor_parallel.layers import param_is_not_tensor_parallel_duplicate + + assert param_is_not_tensor_parallel_duplicate(adapter.experts.lora_A, tp_group=dist.group.WORLD) + assert param_is_not_tensor_parallel_duplicate(adapter.experts.lora_B, tp_group=dist.group.WORLD) + assert param_is_not_tensor_parallel_duplicate(adapter.router.weight, tp_group=dist.group.WORLD) == (rank == 0) finally: dist.destroy_process_group() From a4fdc5d52f68b7ea3288a9420321ef24b544e245 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:42:05 +0800 Subject: [PATCH 39/41] =?UTF-8?q?fix(weight-update):=20=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E5=88=86=E5=9D=97=E5=8F=91=E9=80=81=E7=9A=84=E8=B7=A8=20rank?= =?UTF-8?q?=20=E5=A4=B1=E8=B4=A5=E5=90=8C=E6=AD=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 权重同步有三处结构相同的流水发送循环：先把上一块的引用交给引擎确认，再发下一块，最后统一等待。等待这一步并不是各 rank 对称的，只有 IPC gather 源那个 rank 持有 object ref，引擎侧出错也只会在它上面抛出。base 与 adapter 两条路径的循环直接裸调 ray.get，完全没有把这个失败告诉别的 rank；Mixture 那条只在发布版本号的最后一块之前做了一次失败标志 all-reduce，中间块的异常仍然是就地抛出。结果是某一块传输失败时，其余 rank 察觉不到，会继续走进下一块的集合通信，一直阻塞到分布式超时，日志里先看到的是超时而不是真正出错的那一块。 新增 relax/backends/megatron/weight_update/synchronized_send.py，提供三个原语：run_synchronized_phase 负责阶段级操作的失败同步，raise_on_any_rank_failure 把任意 rank 的本地异常经 gloo all-reduce 变成全体 rank 一起抛出，send_chunks_pipelined 负责流水发送与确认。update_weight_from_tensor 的三处循环统一改为调用 send_chunks_pipelined，中间块、发布版本号前的确认和收尾的最后一块都经过失败同步；阶段级操作改用 run_synchronized_phase，删除原有的 _run_synchronized_weight_update_phase 与随之不再使用的导入。 成功路径的行为没有变化：发送顺序、张量存活范围和每块的 barrier 都与原来一致。变化的是失败路径，从只有 gather 源 rank 抛出、其余 rank 等待超时，变成任意 rank 出错时全体一起抛出同一个异常。 新增 tests/backends/megatron/weight_update/test_synchronized_send.py 覆盖正常发送顺序、单 rank 失败时全体抛出、以及失败后不再发布后续分块；同步用例中的打桩目标同步指向新模块。 --- .../weight_update/synchronized_send.py | 127 ++++++++++++++++++ .../update_weight_from_tensor.py | 121 +++++------------ .../test_mixture_lora_weight_sync.py | 5 +- .../weight_update/test_synchronized_send.py | 127 ++++++++++++++++++ 4 files changed, 289 insertions(+), 91 deletions(-) create mode 100644 relax/backends/megatron/weight_update/synchronized_send.py create mode 100644 tests/backends/megatron/weight_update/test_synchronized_send.py diff --git a/relax/backends/megatron/weight_update/synchronized_send.py b/relax/backends/megatron/weight_update/synchronized_send.py new file mode 100644 index 000000000..b1d6d6199 --- /dev/null +++ b/relax/backends/megatron/weight_update/synchronized_send.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pipelined weight transport with cross-rank failure synchronization. + +Every weight-sync path sends chunk ``N`` while chunk ``N+1`` is still being +converted, then waits on chunk ``N``'s in-flight request before releasing its +tensors. That wait is *not* rank-symmetric: only the IPC gather-source rank +holds the object refs, so an engine-side failure raises there and nowhere else. +Without a shared failure flag the remaining ranks walk into the next chunk's +collectives and block until the distributed timeout. + +The helpers here turn any such rank-local error into an all-rank abort by +all-reducing the failure flag over Gloo, and expose one pipelined send loop that +the base, adapter-mode and Mixture-of-LoRA paths all drive. +""" + +from collections.abc import Callable, Iterable +from typing import Any + +import ray +import torch +import torch.distributed as dist +from ray import ObjectRef + +from relax.utils import device as device_utils +from relax.utils.distributed_utils import get_gloo_group +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + + +def run_synchronized_phase( + operation: Callable[[], Any], + *, + description: str = "weight update phase", +) -> tuple[Exception | None, bool]: + """Run one phase and share its failure state across every rank. + + Returns the local exception (``None`` when this rank succeeded) and whether + *any* rank failed. All ranks must call this in the same order. + """ + + local_error: Exception | None = None + try: + operation() + except Exception as error: # noqa: BLE001 - re-raised by the caller after the collective + local_error = error + logger.error( + "%s failed on rank %d before failure synchronization", + description, + dist.get_rank(), + exc_info=(type(error), error, error.__traceback__), + ) + failed = torch.tensor([local_error is not None], dtype=torch.int32) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=get_gloo_group()) + return local_error, bool(failed.item()) + + +def raise_on_any_rank_failure( + operation: Callable[[], Any], + *, + description: str = "weight update phase", +) -> None: + """Run ``operation`` and abort on every rank if it failed anywhere.""" + + local_error, failed = run_synchronized_phase(operation, description=description) + if not failed: + return + if local_error is not None: + raise local_error + raise RuntimeError(f"{description} failed on another rank") + + +def send_chunks_pipelined( + chunks: Iterable[Any], + send_chunk: Callable[[Any], tuple[list[ObjectRef], Any]], + *, + description: str = "weight chunk transfer", + confirm_before: Callable[[Any], bool] | None = None, +) -> None: + """Send weight chunks, overlapping each transfer with the next conversion. + + ``send_chunk`` converts and ships one chunk and returns its in-flight refs + plus the tensors that must stay alive until those refs resolve. The wait on + the previous chunk is deferred to this iteration so transfer and conversion + overlap, and it is failure-synchronized so a gather-source-only error aborts + all ranks instead of stranding them in the next collective. + + ``confirm_before`` marks chunks that must not be sent until every preceding + chunk is confirmed on every rank — the Mixture-of-LoRA path uses it for the + versioned chunk that publishes the update. + + Collective contract: all ranks must iterate the same number of chunks and + agree on ``confirm_before``, which the shared chunk iterators guarantee. + """ + + pending_refs: list[ObjectRef] = [] + pending_tensors: Any = None + + def drain(refs: list[ObjectRef]) -> Callable[[], None]: + def wait() -> None: + if refs: + ray.get(refs) + + return wait + + for chunk in chunks: + if confirm_before is not None and confirm_before(chunk): + raise_on_any_rank_failure(drain(pending_refs), description=f"preceding {description}") + del pending_tensors + pending_refs = [] + pending_tensors = None + + refs, long_lived_tensors = send_chunk(chunk) + # Confirm the previous chunk on every rank before dropping the tensors + # that back its shared-memory payload. + raise_on_any_rank_failure(drain(pending_refs), description=description) + del pending_tensors + pending_refs = refs + pending_tensors = long_lived_tensors + # Backend-specific per-chunk synchronization lives in device utils so + # this path stays hardware-agnostic. + device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) + + raise_on_any_rank_failure(drain(pending_refs), description=f"final {description}") + del pending_tensors diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index 9e9bfbb6e..a31d49607 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py @@ -14,7 +14,6 @@ from ray.actor import ActorHandle from relax.backends.megatron.misc_utils import strip_param_name_prefix -from relax.utils import device as device_utils from relax.utils import megatron_bridge_utils from relax.utils.device import make_current_torch_device from relax.utils.distributed_utils import get_gloo_group @@ -32,6 +31,7 @@ from .hf_weight_iterator_base import HfWeightIteratorBase from .lora_adapter_sync import LoraAdapterSync from .mixture_lora_sync import MixtureLoraSync +from .synchronized_send import run_synchronized_phase, send_chunks_pipelined from .update_weight_from_distributed import ( connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, @@ -263,27 +263,14 @@ def update_weights(self) -> None: # Pipeline: when chunk N's IPC refs are in-flight on the engine, # chunk N+1's HF conversion + serialize + gather can proceed in - # parallel. We defer ``ray.get`` to the *next* iteration so the - # two stages overlap. - prev_refs: list[ObjectRef] = [] - prev_long_lived_tensors = None + # parallel. The shared primitive defers each ``ray.get`` to the + # *next* iteration and synchronizes its failures across ranks. with export_ctx: - for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): - refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) - # Wait for the *previous* chunk's IPC to finish before - # releasing its GPU tensors. - if prev_refs: - ray.get(prev_refs) - del prev_long_lived_tensors - prev_refs = refs - prev_long_lived_tensors = long_lived_tensors - # Backend-specific per-chunk synchronization is handled in device - # utils so this path stays hardware-agnostic. - device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) - # Drain the last chunk. - if prev_refs: - ray.get(prev_refs) - del prev_long_lived_tensors + send_chunks_pipelined( + self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights), + self._send_hf_params, + description="base weight chunk transfer", + ) # All ranks must finish sending before rank 0 triggers Marlin repack, # otherwise engines in slower gather groups may still be processing @@ -358,7 +345,9 @@ def resume_generation(*, finish_quantization: bool) -> None: ray.get([engine.continue_generation.remote() for engine in all_engines]) for phase_name, operation in (("pause and flush", pause_and_flush), ("send weights", send_weights)): - local_error, phase_failed = self._run_synchronized_weight_update_phase(operation) + local_error, phase_failed = run_synchronized_phase( + operation, description=f"Mixture-of-LoRA weight update phase {phase_name!r}" + ) if not phase_failed: continue @@ -371,8 +360,9 @@ def resume_generation(*, finish_quantization: bool) -> None: raise local_error raise RuntimeError(f"Mixture-of-LoRA weight update phase {phase_name!r} failed on another rank") - local_error, phase_failed = self._run_synchronized_weight_update_phase( - lambda: resume_generation(finish_quantization=True) + local_error, phase_failed = run_synchronized_phase( + lambda: resume_generation(finish_quantization=True), + description="Mixture-of-LoRA weight update phase 'resume generation'", ) if phase_failed: if local_error is not None: @@ -382,62 +372,20 @@ def resume_generation(*, finish_quantization: bool) -> None: self.weight_version = next_weight_version self._mixture_lora_sync.base_sync_done = True - @staticmethod - def _run_synchronized_weight_update_phase(operation): - """Run one update phase and share its failure state across ranks.""" + def _send_weight_update_stream(self, updates) -> None: + """Pipeline conversion collectives with the preceding IPC request. - local_error = None - try: - operation() - except Exception as error: - local_error = error - logger.error( - "Weight update phase failed on rank %d before failure synchronization", - dist.get_rank(), - exc_info=(type(error), error, error.__traceback__), - ) - failed = torch.tensor([local_error is not None], dtype=torch.int32) - dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=get_gloo_group()) - return local_error, bool(failed.item()) + ``updates`` yields ``(named_tensors, weight_version)``; the versioned + final chunk makes the whole update visible, so every preceding chunk is + confirmed on every rank before it is sent. + """ - def _send_weight_update_stream(self, updates) -> None: - """Pipeline conversion collectives with the preceding IPC request.""" - - previous_refs: list[ObjectRef] = [] - previous_tensors = None - has_previous_chunk = False - for named_tensors, weight_version in updates: - # The versioned final chunk makes the whole update visible. Confirm - # the preceding chunk on every rank before publishing that version. - if weight_version is not None and has_previous_chunk: - previous_error, previous_failed = self._run_synchronized_weight_update_phase( - lambda: ray.get(previous_refs) if previous_refs else None - ) - del previous_tensors - previous_refs = [] - previous_tensors = None - if previous_failed: - if previous_error is not None: - raise previous_error - raise RuntimeError("A preceding Mixture-of-LoRA weight chunk failed on another rank") - - refs, long_lived_tensors = self._send_hf_params(named_tensors, weight_version=weight_version) - previous_error = None - if previous_refs: - try: - ray.get(previous_refs) - except Exception as error: - previous_error = error - del previous_tensors - previous_refs = refs - previous_tensors = long_lived_tensors - has_previous_chunk = True - device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group()) - if previous_error is not None: - raise previous_error - if previous_refs: - ray.get(previous_refs) - del previous_tensors + send_chunks_pipelined( + updates, + lambda update: self._send_hf_params(update[0], weight_version=update[1]), + description="Mixture-of-LoRA weight chunk transfer", + confirm_before=lambda update: update[1] is not None, + ) def _update_weights_adapter_mode(self) -> None: """LoRA adapter mode: sync base once, then push only the adapter each @@ -505,18 +453,11 @@ def _sync_base_and_lora(self) -> None: # 1) Send BASE weights only. In adapter mode the HF iterator pulls adapter params OUT of # the conversion buckets (collect_adapters=True) without merging, so only base weights # flow through SGLang's base-model load_weights (which has no notion of lora_A/lora_B). - prev_refs: list[ObjectRef] = [] - prev_long_lived_tensors = None - for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): - refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) - if prev_refs: - ray.get(prev_refs) - del prev_long_lived_tensors - prev_refs = refs - prev_long_lived_tensors = long_lived_tensors - if prev_refs: - ray.get(prev_refs) - del prev_long_lived_tensors + send_chunks_pipelined( + self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights), + self._send_hf_params, + description="adapter-mode base weight chunk transfer", + ) dist.barrier(group=get_gloo_group()) diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index 64c762127..dd5b8c8f9 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -547,6 +547,7 @@ def record_updates(updates): patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.all_reduce"), patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), patch("ray.get", side_effect=lambda refs: refs), ): updater._update_weights_mixture_lora() @@ -606,6 +607,7 @@ def fail_update(_updates): patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.all_reduce"), patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), patch("ray.get", side_effect=lambda refs: refs), pytest.raises(RuntimeError, match="engine rejected routed weights"), ): @@ -627,12 +629,13 @@ def test_stream_failure_does_not_publish_the_final_versioned_chunk(): with ( patch("ray.get", side_effect=RuntimeError("engine update failed")), patch( - "relax.backends.megatron.weight_update.update_weight_from_tensor." + "relax.backends.megatron.weight_update.synchronized_send." "device_utils.maybe_backend_barrier_on_weight_chunk" ) as chunk_barrier, patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.all_reduce"), patch("relax.backends.megatron.weight_update.update_weight_from_tensor.get_gloo_group", return_value=None), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), pytest.raises(RuntimeError, match="engine update failed"), ): updater._send_weight_update_stream([([("first", torch.ones(1))], None), ([("second", torch.ones(1))], 5)]) diff --git a/tests/backends/megatron/weight_update/test_synchronized_send.py b/tests/backends/megatron/weight_update/test_synchronized_send.py new file mode 100644 index 000000000..e2da7b740 --- /dev/null +++ b/tests/backends/megatron/weight_update/test_synchronized_send.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from relax.backends.megatron.weight_update.synchronized_send import ( + raise_on_any_rank_failure, + run_synchronized_phase, + send_chunks_pipelined, +) + + +@contextmanager +def _single_rank_collectives(*, peer_failed: bool = False): + """Stand in for the Gloo collectives on a one-rank world. + + ``peer_failed`` simulates another rank reporting a failure, which is what + the all-reduce exists to propagate. + """ + + def fake_all_reduce(tensor, *args, **kwargs): + if peer_failed: + tensor.fill_(1) + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), + ): + yield + + +def test_pipelined_send_defers_each_wait_until_the_next_chunk_is_in_flight(): + events = [] + + def send_chunk(chunk): + events.append(f"send:{chunk}") + return [f"ref:{chunk}"], f"tensors:{chunk}" + + with ( + _single_rank_collectives(), + patch("ray.get", side_effect=lambda refs: events.append(f"wait:{refs}")), + ): + send_chunks_pipelined(["a", "b"], send_chunk) + + assert events == ["send:a", "send:b", "wait:['ref:a']", "wait:['ref:b']"] + + +def test_pipelined_send_aborts_when_an_intermediate_chunk_fails_on_another_rank(): + """The IPC wait only raises on the gather-source rank. + + Every other rank has to learn about it from the all-reduce, otherwise it + blocks in the next chunk's collectives. + """ + + send_chunk = MagicMock(side_effect=[(["ref:a"], "tensors:a"), (["ref:b"], "tensors:b"), (["ref:c"], "tensors:c")]) + + with ( + _single_rank_collectives(peer_failed=True), + patch("ray.get") as ray_get, + pytest.raises(RuntimeError, match="failed on another rank"), + ): + send_chunks_pipelined(["a", "b", "c"], send_chunk) + + # The failure surfaces at the first chunk boundary, long before the final + # drain that used to be the only synchronization point. + assert send_chunk.call_count == 1 + ray_get.assert_not_called() + + +def test_pipelined_send_confirms_predecessors_before_a_marked_chunk(): + events = [] + + def send_chunk(chunk): + events.append(f"send:{chunk[0]}") + return [f"ref:{chunk[0]}"], None + + with ( + _single_rank_collectives(), + patch("ray.get", side_effect=lambda refs: events.append(f"wait:{refs}")), + ): + send_chunks_pipelined( + [("a", None), ("b", 7)], + send_chunk, + confirm_before=lambda chunk: chunk[1] is not None, + ) + + assert events == ["send:a", "wait:['ref:a']", "send:b", "wait:['ref:b']"] + + +def test_synchronized_phase_reports_the_local_error_and_the_shared_flag(): + failure = RuntimeError("engine rejected the update") + + with _single_rank_collectives(): + local_error, failed = run_synchronized_phase(MagicMock(side_effect=failure)) + clean_error, clean_failed = run_synchronized_phase(MagicMock()) + + assert local_error is failure + assert failed is True + assert clean_error is None + assert clean_failed is False + + +def test_raise_on_any_rank_failure_reraises_the_local_error_unchanged(): + failure = ValueError("bad chunk") + + with _single_rank_collectives(), pytest.raises(ValueError) as excinfo: + raise_on_any_rank_failure(MagicMock(side_effect=failure)) + + assert excinfo.value is failure + + +def test_failure_flag_is_reduced_as_an_integer_tensor(): + reduced = [] + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_reduce", side_effect=lambda tensor, *a, **k: reduced.append(tensor.clone())), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), + ): + run_synchronized_phase(MagicMock()) + + assert reduced[0].dtype == torch.int32 + assert reduced[0].tolist() == [0] From 7c544ce6600d41ca504881158b8346a03522e410 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:42:38 +0800 Subject: [PATCH 40/41] =?UTF-8?q?fix(mixture-lora):=20=E8=B7=AF=E7=94=B1?= =?UTF-8?q?=E6=9D=83=E9=87=8D=E6=A0=A1=E9=AA=8C=E6=94=B9=E4=B8=BA=E5=85=A8?= =?UTF-8?q?=E4=BD=93=20rank=20=E5=90=8C=E6=AD=A5=E5=A4=B1=E8=B4=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 路由权重的形状、dtype 和缺键检查原本写在生成器循环里，与 PP 广播、TP all-gather 交替执行。某个 rank 的分片有问题时，它在广播之前就地抛出异常直接退出，而其余 rank 已经进入同一轮集合通信，只能等到分布式超时才失败，日志里最先看到的是超时而不是那个真正出错的张量，排查成本很高。 改为在任何集合通信之前先完成本 rank 全部分片的校验，并借助上一提交的 raise_on_any_rank_failure 把校验结果广播给所有 rank：只要有一个 rank 校验失败，全体 rank 立即一起抛错，不会有人卡在集合通信上。校验通过后再进入原有的广播与 all-gather 流程，检查逻辑与报错信息保持不变，只是拆成 _collect_local_tensors 与 _iter_weight_chunks 两步。 新增两个单测：一个确认缺键和形状不符会在任何广播、all-gather 发出之前抛错，另一个确认本 rank 分片正常但其他 rank 失败时同样会抛错退出。 --- .../weight_update/mixture_lora_sync.py | 69 ++++++++++++++----- .../test_mixture_lora_weight_sync.py | 59 ++++++++++++++++ 2 files changed, 112 insertions(+), 16 deletions(-) diff --git a/relax/backends/megatron/weight_update/mixture_lora_sync.py b/relax/backends/megatron/weight_update/mixture_lora_sync.py index e6224efb7..a91d48dfa 100644 --- a/relax/backends/megatron/weight_update/mixture_lora_sync.py +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -20,6 +20,7 @@ ) from .common import named_params_and_buffers +from .synchronized_send import raise_on_any_rank_failure @dataclass(frozen=True) @@ -249,10 +250,59 @@ def get_weight_chunks( self, local_weights: Mapping[str, torch.Tensor], ) -> Iterator[list[tuple[str, torch.Tensor]]]: - """Yield full routed tensors in deterministic, size-bounded chunks.""" + """Yield full routed tensors in deterministic, size-bounded chunks. + + Each rank checks the shards it owns *before* the first collective and + shares the outcome with every other rank: the PP broadcast and the TP + all-gather below are collective, so a rank-local raise in between would + leave the peers of the failing rank blocked until the distributed + timeout instead of reporting the bad tensor. + """ - rank = dist.get_rank() device = device_utils.make_current_torch_device() + local_tensors: dict[str, torch.Tensor] = {} + raise_on_any_rank_failure( + lambda: local_tensors.update(self._collect_local_tensors(local_weights, device=device)), + description="Mixture-of-LoRA routed weight validation", + ) + return self._iter_weight_chunks(local_tensors, device=device) + + def _collect_local_tensors( + self, + local_weights: Mapping[str, torch.Tensor], + *, + device: torch.device, + ) -> dict[str, torch.Tensor]: + """Validate and stage every routed tensor this rank owns.""" + + rank = dist.get_rank() + staged: dict[str, torch.Tensor] = {} + for info in self.param_infos: + if rank != info.src_rank: + continue + if info.weight_key not in local_weights: + raise KeyError(f"Missing Mixture-of-LoRA weight {info.weight_key!r}") + source_tensor = local_weights[info.weight_key] + if tuple(source_tensor.shape) != info.local_shape: + raise ValueError( + f"Mixture-of-LoRA weight {info.weight_key!r} has shape {tuple(source_tensor.shape)}, " + f"expected {info.local_shape}" + ) + if source_tensor.dtype != info.state.dtype: + raise TypeError( + f"Mixture-of-LoRA weight {info.weight_key!r} has dtype {source_tensor.dtype}, " + f"expected {info.state.dtype}" + ) + staged[info.state.parameter_name] = source_tensor.to(device=device) + return staged + + def _iter_weight_chunks( + self, + local_tensors: Mapping[str, torch.Tensor], + *, + device: torch.device, + ) -> Iterator[list[tuple[str, torch.Tensor]]]: + rank = dist.get_rank() tp_world_size = mpu.get_tensor_model_parallel_world_size() tp_group = mpu.get_tensor_model_parallel_group() pp_world_size = mpu.get_pipeline_model_parallel_world_size() @@ -264,20 +314,7 @@ def get_weight_chunks( for info in self.param_infos: if rank == info.src_rank: - if info.weight_key not in local_weights: - raise KeyError(f"Missing Mixture-of-LoRA weight {info.weight_key!r}") - source_tensor = local_weights[info.weight_key] - if tuple(source_tensor.shape) != info.local_shape: - raise ValueError( - f"Mixture-of-LoRA weight {info.weight_key!r} has shape {tuple(source_tensor.shape)}, " - f"expected {info.local_shape}" - ) - if source_tensor.dtype != info.state.dtype: - raise TypeError( - f"Mixture-of-LoRA weight {info.weight_key!r} has dtype {source_tensor.dtype}, " - f"expected {info.state.dtype}" - ) - local_tensor = source_tensor.to(device=device) + local_tensor = local_tensors[info.state.parameter_name] else: local_tensor = torch.empty(info.local_shape, dtype=info.state.dtype, device=device) if pp_group is not None and info.src_rank in pp_ranks: diff --git a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py index dd5b8c8f9..d5284fcd8 100644 --- a/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. +from contextlib import contextmanager from dataclasses import replace from datetime import timedelta from types import SimpleNamespace @@ -187,6 +188,64 @@ def gather(output_list, input_object, group=None): } +@contextmanager +def _single_rank_routed_sync(*, peer_failed: bool = False): + """Run routed-weight validation on a one-rank stand-in for the world.""" + + def fake_all_reduce(tensor, *args, **kwargs): + if peer_failed: + tensor.fill_(1) + + with ( + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.all_reduce", side_effect=fake_all_reduce), + patch("relax.backends.megatron.weight_update.synchronized_send.get_gloo_group", return_value=None), + patch( + "relax.backends.megatron.weight_update.mixture_lora_sync.device_utils.make_current_torch_device", + return_value=torch.device("cpu"), + ), + patch("torch.distributed.broadcast") as broadcast, + patch("torch.distributed.all_gather") as all_gather, + ): + yield broadcast, all_gather + + +def test_routed_weight_validation_rejects_bad_tensors_before_any_collective(): + from relax.backends.megatron.weight_update.mixture_lora_sync import MixtureLoraSync + + sync = MixtureLoraSync.__new__(MixtureLoraSync) + sync.param_infos = (_info("linear_qkv", "router.weight", (2, 6), (2, 6), None),) + + with _single_rank_routed_sync() as (broadcast, all_gather): + with pytest.raises(KeyError, match="Missing Mixture-of-LoRA weight"): + sync.get_weight_chunks({}) + with pytest.raises(ValueError, match="has shape"): + sync.get_weight_chunks({"weight": torch.zeros(2, 5)}) + + broadcast.assert_not_called() + all_gather.assert_not_called() + + +def test_routed_weight_validation_aborts_every_rank_when_one_rank_fails(): + """A rank-local raise before the PP/TP collectives strands its peers. + + The validation result is therefore shared before the first collective, so a + rank whose own shards are fine still aborts instead of blocking. + """ + + from relax.backends.megatron.weight_update.mixture_lora_sync import MixtureLoraSync + + sync = MixtureLoraSync.__new__(MixtureLoraSync) + sync.param_infos = (_info("linear_qkv", "router.weight", (2, 6), (2, 6), None),) + + with _single_rank_routed_sync(peer_failed=True) as (broadcast, all_gather): + with pytest.raises(RuntimeError, match="failed on another rank"): + sync.get_weight_chunks({"weight": torch.zeros(2, 6)}) + + broadcast.assert_not_called() + all_gather.assert_not_called() + + def test_qkv_lora_b_tp_shards_are_converted_from_group_layout_to_qkv_blocks(): # Two query groups, two query heads per group, then one K and one V head. grouped = torch.tensor([[[[10.0], [11.0], [20.0], [30.0], [12.0], [13.0], [21.0], [31.0]]]]).reshape(1, 8, 1) From 2a0910004aea0ad26289da9ab454341bc6d67359 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Mon, 17 Aug 2026 22:43:16 +0800 Subject: [PATCH 41/41] =?UTF-8?q?refactor(sglang):=20=E6=8A=BD=E5=87=BA?= =?UTF-8?q?=E4=B8=8E=E6=9E=B6=E6=9E=84=E6=97=A0=E5=85=B3=E7=9A=84=20Mixtur?= =?UTF-8?q?e=20=E6=8E=A8=E7=90=86=E4=BE=A7=E5=AE=9E=E7=8E=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit relax/models/qwen3_mixture_lora/sglang/model.py 里同时放着两类代码：路由注入、前向改写、权重加载这些对任何架构都一样的逻辑，以及 Qwen3 特有的注入点声明。再接入一个模型架构就得整体复制一遍，两份实现随后各自演进，很快就会出现只修了一边的情况。 架构无关的部分移到 relax/models/mixture_lora_sglang.py：运行时配置读取、路由线性层前向、适配器挂载、路由权重加载，以及供各架构复用的 MixtureLoraSGLangModelMixin（包含 TP 必须为 1 的检查，以及把 mixture_lora 权重从常规权重里分流后再交给父类加载）。qwen3_mixture_lora/sglang/model.py 只剩注入点声明与 EntryClass，代码从 213 行降到 22 行。行为和对外类名保持不变。 测试同步拆分：通用逻辑的用例移到 tests/models/test_mixture_lora_sglang.py，tests/models/qwen3_mixture_lora/test_sglang_model.py 只保留 Qwen3 注入点与注册相关断言。 --- relax/models/mixture_lora_sglang.py | 265 +++++++++++++++ .../models/qwen3_mixture_lora/sglang/model.py | 213 +----------- .../qwen3_mixture_lora/test_sglang_model.py | 241 ++------------ tests/models/test_mixture_lora_sglang.py | 303 ++++++++++++++++++ 4 files changed, 601 insertions(+), 421 deletions(-) create mode 100644 relax/models/mixture_lora_sglang.py create mode 100644 tests/models/test_mixture_lora_sglang.py diff --git a/relax/models/mixture_lora_sglang.py b/relax/models/mixture_lora_sglang.py new file mode 100644 index 000000000..986413edc --- /dev/null +++ b/relax/models/mixture_lora_sglang.py @@ -0,0 +1,265 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Architecture-independent Mixture-of-LoRA support for SGLang models. + +Rollout injection, routed execution and weight loading are identical for every +SGLang architecture: only the mapping from a decoder layer to the linears that +carry routed experts differs. Concrete external models therefore mix in +:class:`MixtureLoraSGLangModelMixin` and implement +:meth:`MixtureLoraSGLangModelMixin.mixture_lora_site_modules`; everything else +lives here so a second architecture costs one small subclass. +""" + +from types import MethodType +from typing import Any, Iterable + +import torch +from torch import nn +from torch.nn import functional as F + +from relax.utils.env import Envs +from relax.utils.mixture_lora_common import ( + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + deserialize_mixture_lora_config, + megatron_mixture_lora_name_to_sglang, + route_topk, +) + + +__all__ = [ + "MixtureLoraSGLangModelMixin", + "SGLangMixtureLoRA", + "attach_sglang_mixture_lora", + "load_sglang_mixture_lora_weights", + "mixture_lora_site_id", + "read_runtime_mixture_lora_config", +] + +# Site ids are the Megatron parameter names the training backend publishes, so +# rollout and training agree on one vocabulary for every architecture. +SUPPORTED_MIXTURE_LORA_TARGETS = frozenset({"linear_qkv", "linear_proj"}) + + +def mixture_lora_site_id(layer_id: int, target: str) -> str: + """Build the Megatron site id for one routed linear.""" + + return f"decoder.layers.{layer_id}.self_attention.{target}" + + +def read_runtime_mixture_lora_config() -> MixtureLoraConfig: + """Read the routing configuration the trainer handed to this engine.""" + + runtime_config = Envs.RELAX_MIXTURE_LORA_CONFIG + if runtime_config is None: + raise RuntimeError("RELAX_MIXTURE_LORA_CONFIG must be set before constructing the external model") + return deserialize_mixture_lora_config(runtime_config) + + +class SGLangMixtureLoRA(nn.Module): + """Parameter-compatible dense Mixture-of-LoRA execution for rollout.""" + + def __init__( + self, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.config = config + self.site_id = site_id + self.experts = nn.Module() + self.experts.register_parameter( + "lora_A", + nn.Parameter(torch.empty(config.num_experts, config.rank, input_size, device=device, dtype=dtype)), + ) + self.experts.register_parameter( + "lora_B", + nn.Parameter(torch.empty(config.num_experts, output_size, config.rank, device=device, dtype=dtype)), + ) + self.router = nn.Linear(input_size, config.num_experts, bias=False, device=device, dtype=dtype) + self.executor = DenseRoutedLoRAExecutor() + self.reset_parameters() + + def reset_parameters(self) -> None: + for expert_weight in self.experts.lora_A: + nn.init.xavier_uniform_(expert_weight) + nn.init.zeros_(self.experts.lora_B) + nn.init.normal_(self.router.weight, mean=0.0, std=0.02) + + def route(self, x: torch.Tensor): + logits = F.linear(x.reshape(-1, x.shape[-1]).float(), self.router.weight.float()) + return route_topk(logits, self.config.top_k, self.config.temperature) + + def forward_with_routing(self, x: torch.Tensor): + decision = self.route(x) + output = self.executor( + x, + self.experts.lora_A, + self.experts.lora_B, + decision, + self.config.scale, + ) + return output, decision + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output, _ = self.forward_with_routing(x) + return output + + +def _routed_linear_forward(linear, x: torch.Tensor, *args: Any, **kwargs: Any): + base_result = linear._relax_mixture_lora_base_forward(x, *args, **kwargs) + if not isinstance(base_result, tuple) or len(base_result) != 2: + raise TypeError(f"{type(linear).__name__} must return an (output, bias) tuple") + output, bias = base_result + delta = linear.mixture_lora(x).reshape(output.shape) + return output + delta, bias + + +def attach_sglang_mixture_lora( + linear: nn.Module, + config: MixtureLoraConfig, + site_id: str, + input_size: int, + output_size: int, +) -> None: + """Add routed parameters while preserving SGLang base parameter names.""" + + if hasattr(linear, "mixture_lora"): + raise RuntimeError(f"SGLang linear {site_id} already has a Mixture-of-LoRA adapter") + try: + base_parameter = next(linear.parameters()) + except StopIteration as error: + raise ValueError(f"SGLang linear {site_id} has no parameters") from error + adapter_dtype = getattr(linear, "params_dtype", base_parameter.dtype) + if not isinstance(adapter_dtype, torch.dtype) or not adapter_dtype.is_floating_point: + raise TypeError( + f"SGLang linear {site_id} must expose a floating-point params_dtype for Mixture-of-LoRA, " + f"got {adapter_dtype}" + ) + linear.add_module( + "mixture_lora", + SGLangMixtureLoRA( + config, + site_id, + input_size, + output_size, + device=base_parameter.device, + dtype=adapter_dtype, + ), + ) + linear._relax_mixture_lora_base_forward = linear.forward + linear.forward = MethodType(_routed_linear_forward, linear) + + +def load_sglang_mixture_lora_weights( + model: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Validate and copy one chunk of routed weights into an SGLang model.""" + + parameters = dict(model.named_parameters()) + loaded_names = set() + for source_name, loaded_weight in weights: + target_name = ( + source_name + if source_name.startswith("model.layers.") + else megatron_mixture_lora_name_to_sglang(source_name) + ) + if target_name in loaded_names: + raise ValueError(f"Duplicate Mixture-of-LoRA weight: {target_name}") + if target_name not in parameters: + layer_id = int(target_name.split(".", maxsplit=3)[2]) + start_layer = getattr(getattr(model, "model", None), "start_layer", None) + end_layer = getattr(getattr(model, "model", None), "end_layer", None) + if start_layer is not None and end_layer is not None and not start_layer <= layer_id < end_layer: + continue + raise ValueError(f"Unknown Mixture-of-LoRA weight for SGLang: {target_name}") + parameter = parameters[target_name] + if tuple(loaded_weight.shape) != tuple(parameter.shape): + raise ValueError( + f"Mixture-of-LoRA weight {target_name} has shape {tuple(loaded_weight.shape)}, " + f"expected {tuple(parameter.shape)}" + ) + if loaded_weight.dtype != parameter.dtype: + raise TypeError( + f"Mixture-of-LoRA weight {target_name} has dtype {loaded_weight.dtype}, expected {parameter.dtype}" + ) + with torch.no_grad(): + parameter.copy_(loaded_weight.to(device=parameter.device)) + loaded_names.add(target_name) + return loaded_names + + +class MixtureLoraSGLangModelMixin: + """Add token-routed LoRA experts to an SGLang causal-LM model. + + Mix in *before* the SGLang base class so the routed weights are split off + before its loader sees them:: + + class Qwen3ForCausalLM(MixtureLoraSGLangModelMixin, SGLangQwen3ForCausalLM): + def mixture_lora_site_modules(self, layer_id): ... + """ + + supported_mixture_lora_targets: frozenset[str] = SUPPORTED_MIXTURE_LORA_TARGETS + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.mixture_lora_config = read_runtime_mixture_lora_config() + # Routed parameters are replicated whole on every rollout worker, so a + # sharded engine would serve mismatched experts. + from sglang.srt.distributed import get_tensor_model_parallel_world_size + + if get_tensor_model_parallel_world_size() != 1: + raise ValueError(f"{type(self).__name__} Mixture-of-LoRA rollout currently requires SGLang TP=1") + super().__init__(*args, **kwargs) + self.install_mixture_lora() + + def mixture_lora_site_modules(self, layer_id: int) -> dict[str, nn.Module]: + """Map each supported target name to the linear it wraps in a layer.""" + + raise NotImplementedError + + def install_mixture_lora(self) -> None: + """Wrap every routed linear this pipeline stage owns.""" + + targets = set(self.mixture_lora_config.target_modules) + unsupported = sorted(targets - set(self.supported_mixture_lora_targets)) + if unsupported: + raise ValueError(f"Unsupported SGLang Mixture-of-LoRA targets: {unsupported}") + start_layer = getattr(self.model, "start_layer", 0) + end_layer = getattr(self.model, "end_layer", len(self.model.layers)) + for layer_id in range(start_layer, end_layer): + site_modules = self.mixture_lora_site_modules(layer_id) + missing = sorted(targets - set(site_modules)) + if missing: + raise ValueError(f"{type(self).__name__} exposes no Mixture-of-LoRA site for {missing}") + for target in sorted(targets): + linear = site_modules[target] + attach_sglang_mixture_lora( + linear, + self.mixture_lora_config, + mixture_lora_site_id(layer_id, target), + linear.input_size, + linear.output_size, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Route the checkpoint stream: base to SGLang, routed tensors here.""" + + mixture_weights = [] + + def base_weight_iterator(): + for name, weight in weights: + if ".mixture_lora." in name: + mixture_weights.append((name, weight)) + else: + yield name, weight + + result = super().load_weights(base_weight_iterator()) + if mixture_weights: + load_sglang_mixture_lora_weights(self, mixture_weights) + return result diff --git a/relax/models/qwen3_mixture_lora/sglang/model.py b/relax/models/qwen3_mixture_lora/sglang/model.py index c57c13702..e755e1ef9 100644 --- a/relax/models/qwen3_mixture_lora/sglang/model.py +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -2,224 +2,21 @@ """SGLang Qwen3 external model with token-routed LoRA experts.""" -from types import MethodType -from typing import Any, Iterable - -import torch -from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.models.qwen3 import Qwen3ForCausalLM as SGLangQwen3ForCausalLM from torch import nn -from torch.nn import functional as F - -from relax.utils.env import Envs -from relax.utils.mixture_lora_common import ( - DenseRoutedLoRAExecutor, - MixtureLoraConfig, - deserialize_mixture_lora_config, - megatron_mixture_lora_name_to_sglang, - route_topk, -) - - -class SGLangMixtureLoRA(nn.Module): - """Parameter-compatible dense Mixture-of-LoRA execution for rollout.""" - - def __init__( - self, - config: MixtureLoraConfig, - site_id: str, - input_size: int, - output_size: int, - *, - device: torch.device, - dtype: torch.dtype, - ) -> None: - super().__init__() - self.config = config - self.site_id = site_id - self.experts = nn.Module() - self.experts.register_parameter( - "lora_A", - nn.Parameter(torch.empty(config.num_experts, config.rank, input_size, device=device, dtype=dtype)), - ) - self.experts.register_parameter( - "lora_B", - nn.Parameter(torch.empty(config.num_experts, output_size, config.rank, device=device, dtype=dtype)), - ) - self.router = nn.Linear(input_size, config.num_experts, bias=False, device=device, dtype=dtype) - self.executor = DenseRoutedLoRAExecutor() - self.reset_parameters() - - def reset_parameters(self) -> None: - for expert_weight in self.experts.lora_A: - nn.init.xavier_uniform_(expert_weight) - nn.init.zeros_(self.experts.lora_B) - nn.init.normal_(self.router.weight, mean=0.0, std=0.02) - - def route(self, x: torch.Tensor): - logits = F.linear(x.reshape(-1, x.shape[-1]).float(), self.router.weight.float()) - return route_topk(logits, self.config.top_k, self.config.temperature) - - def forward_with_routing(self, x: torch.Tensor): - decision = self.route(x) - output = self.executor( - x, - self.experts.lora_A, - self.experts.lora_B, - decision, - self.config.scale, - ) - return output, decision - - def forward(self, x: torch.Tensor) -> torch.Tensor: - output, _ = self.forward_with_routing(x) - return output - - -def _routed_linear_forward(linear, x: torch.Tensor, *args: Any, **kwargs: Any): - base_result = linear._relax_mixture_lora_base_forward(x, *args, **kwargs) - if not isinstance(base_result, tuple) or len(base_result) != 2: - raise TypeError(f"{type(linear).__name__} must return an (output, bias) tuple") - output, bias = base_result - delta = linear.mixture_lora(x).reshape(output.shape) - return output + delta, bias - -def attach_sglang_mixture_lora( - linear: nn.Module, - config: MixtureLoraConfig, - site_id: str, - input_size: int, - output_size: int, -) -> None: - """Add routed parameters while preserving SGLang base parameter names.""" - - if hasattr(linear, "mixture_lora"): - raise RuntimeError(f"SGLang linear {site_id} already has a Mixture-of-LoRA adapter") - try: - base_parameter = next(linear.parameters()) - except StopIteration as error: - raise ValueError(f"SGLang linear {site_id} has no parameters") from error - adapter_dtype = getattr(linear, "params_dtype", base_parameter.dtype) - if not isinstance(adapter_dtype, torch.dtype) or not adapter_dtype.is_floating_point: - raise TypeError( - f"SGLang linear {site_id} must expose a floating-point params_dtype for Mixture-of-LoRA, " - f"got {adapter_dtype}" - ) - linear.add_module( - "mixture_lora", - SGLangMixtureLoRA( - config, - site_id, - input_size, - output_size, - device=base_parameter.device, - dtype=adapter_dtype, - ), - ) - linear._relax_mixture_lora_base_forward = linear.forward - linear.forward = MethodType(_routed_linear_forward, linear) - - -def load_sglang_mixture_lora_weights( - model: nn.Module, - weights: Iterable[tuple[str, torch.Tensor]], -) -> set[str]: - """Validate and copy one chunk of routed weights into an SGLang model.""" - - parameters = dict(model.named_parameters()) - loaded_names = set() - for source_name, loaded_weight in weights: - target_name = ( - source_name - if source_name.startswith("model.layers.") - else megatron_mixture_lora_name_to_sglang(source_name) - ) - if target_name in loaded_names: - raise ValueError(f"Duplicate Mixture-of-LoRA weight: {target_name}") - if target_name not in parameters: - layer_id = int(target_name.split(".", maxsplit=3)[2]) - start_layer = getattr(getattr(model, "model", None), "start_layer", None) - end_layer = getattr(getattr(model, "model", None), "end_layer", None) - if start_layer is not None and end_layer is not None and not start_layer <= layer_id < end_layer: - continue - raise ValueError(f"Unknown Mixture-of-LoRA weight for SGLang: {target_name}") - parameter = parameters[target_name] - if tuple(loaded_weight.shape) != tuple(parameter.shape): - raise ValueError( - f"Mixture-of-LoRA weight {target_name} has shape {tuple(loaded_weight.shape)}, " - f"expected {tuple(parameter.shape)}" - ) - if loaded_weight.dtype != parameter.dtype: - raise TypeError( - f"Mixture-of-LoRA weight {target_name} has dtype {loaded_weight.dtype}, expected {parameter.dtype}" - ) - with torch.no_grad(): - parameter.copy_(loaded_weight.to(device=parameter.device)) - loaded_names.add(target_name) - return loaded_names +from relax.models.mixture_lora_sglang import MixtureLoraSGLangModelMixin # SGLang registers external models by the checkpoint architecture name. Keeping # this name replaces its built-in Qwen3 entry while preserving checkpoint metadata. -class Qwen3ForCausalLM(SGLangQwen3ForCausalLM): +class Qwen3ForCausalLM(MixtureLoraSGLangModelMixin, SGLangQwen3ForCausalLM): """Qwen3 external model that routes LoRA experts at attention projections.""" - def __init__(self, config, quant_config=None, prefix: str = "") -> None: - runtime_config = Envs.RELAX_MIXTURE_LORA_CONFIG - if runtime_config is None: - raise RuntimeError("RELAX_MIXTURE_LORA_CONFIG must be set before constructing the external model") - self.mixture_lora_config = deserialize_mixture_lora_config(runtime_config) - if get_tensor_model_parallel_world_size() != 1: - raise ValueError("Qwen3 Mixture-of-LoRA rollout currently requires SGLang TP=1") - super().__init__(config, quant_config=quant_config, prefix=prefix) - self._install_mixture_lora() - - def _install_mixture_lora(self) -> None: - targets = set(self.mixture_lora_config.target_modules) - supported_targets = {"linear_qkv", "linear_proj"} - if not targets.issubset(supported_targets): - raise ValueError(f"Unsupported SGLang Mixture-of-LoRA targets: {sorted(targets - supported_targets)}") - start_layer = getattr(self.model, "start_layer", 0) - end_layer = getattr(self.model, "end_layer", len(self.model.layers)) - for layer_id in range(start_layer, end_layer): - layer = self.model.layers[layer_id] - attention = layer.self_attn - if "linear_qkv" in targets: - site_id = f"decoder.layers.{layer_id}.self_attention.linear_qkv" - attach_sglang_mixture_lora( - attention.qkv_proj, - self.mixture_lora_config, - site_id, - attention.qkv_proj.input_size, - attention.qkv_proj.output_size, - ) - if "linear_proj" in targets: - site_id = f"decoder.layers.{layer_id}.self_attention.linear_proj" - attach_sglang_mixture_lora( - attention.o_proj, - self.mixture_lora_config, - site_id, - attention.o_proj.input_size, - attention.o_proj.output_size, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - mixture_weights = [] - - def base_weight_iterator(): - for name, weight in weights: - if ".mixture_lora." in name: - mixture_weights.append((name, weight)) - else: - yield name, weight - - super().load_weights(base_weight_iterator()) - if not mixture_weights: - return - - load_sglang_mixture_lora_weights(self, mixture_weights) + def mixture_lora_site_modules(self, layer_id: int) -> dict[str, nn.Module]: + attention = self.model.layers[layer_id].self_attn + return {"linear_qkv": attention.qkv_proj, "linear_proj": attention.o_proj} EntryClass = Qwen3ForCausalLM diff --git a/tests/models/qwen3_mixture_lora/test_sglang_model.py b/tests/models/qwen3_mixture_lora/test_sglang_model.py index 4159fdd54..988d68e5c 100644 --- a/tests/models/qwen3_mixture_lora/test_sglang_model.py +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -6,18 +6,12 @@ from torch import nn -# This parity suite needs both model backends. They are available in the +# This suite needs the SGLang Qwen3 model definition. It is available in the # official Relax image but intentionally absent from the default CPU CI. -pytest.importorskip("megatron.bridge") pytest.importorskip("sglang.srt.models.qwen3") -from relax.backends.megatron.mixture_lora_modules import MixtureLoRAAdapter # noqa: E402 -from relax.models.qwen3_mixture_lora.sglang.model import ( # noqa: E402 - EntryClass, - SGLangMixtureLoRA, - attach_sglang_mixture_lora, - load_sglang_mixture_lora_weights, -) +from relax.models.mixture_lora_sglang import MixtureLoraSGLangModelMixin # noqa: E402 +from relax.models.qwen3_mixture_lora.sglang.model import EntryClass # noqa: E402 from relax.utils.mixture_lora_common import MixtureLoraConfig # noqa: E402 @@ -33,226 +27,47 @@ def _config(): ) -def test_external_entry_class_overrides_the_checkpoint_architecture(): - assert EntryClass.__name__ == "Qwen3ForCausalLM" - - class _TupleLinear(nn.Module): def __init__(self, input_size: int, output_size: int): super().__init__() self.weight = nn.Parameter(torch.randn(output_size, input_size)) + self.input_size = input_size + self.output_size = output_size def forward(self, x): return F.linear(x, self.weight), None -class _QuantizedTupleLinear(nn.Module): - def __init__(self, params_dtype): - super().__init__() - self.weight_packed = nn.Parameter(torch.ones(4, 4, dtype=torch.int32), requires_grad=False) - self.params_dtype = params_dtype +def test_external_entry_class_overrides_the_checkpoint_architecture(): + assert EntryClass.__name__ == "Qwen3ForCausalLM" -def _fake_qwen_model_with_routed_qkv(): - model = nn.Module() +def test_qwen3_model_reuses_the_shared_mixture_lora_machinery(): + # Everything except the site mapping below is architecture-independent, so + # a second architecture only has to supply its own linears. + assert issubclass(EntryClass, MixtureLoraSGLangModelMixin) + assert EntryClass.__mro__.index(MixtureLoraSGLangModelMixin) < EntryClass.__mro__.index(nn.Module) + assert "install_mixture_lora" not in vars(EntryClass) + assert "load_weights" not in vars(EntryClass) + + +def test_qwen3_sites_map_to_the_attention_projections(): + model = EntryClass.__new__(EntryClass) + nn.Module.__init__(model) + model.mixture_lora_config = _config() model.model = nn.Module() layer = nn.Module() layer.self_attn = nn.Module() layer.self_attn.qkv_proj = _TupleLinear(4, 6) + layer.self_attn.o_proj = _TupleLinear(6, 4) model.model.layers = nn.ModuleList([layer]) - attach_sglang_mixture_lora( - layer.self_attn.qkv_proj, - _config(), - "decoder.layers.0.self_attention.linear_qkv", - 4, - 6, - ) - return model - - -def test_sglang_dense_adapter_matches_training_adapter(): - config = _config() - training_adapter = MixtureLoRAAdapter( - config, - "linear_qkv", - 4, - 6, - dropout=0.0, - device=torch.device("cpu"), - dtype=torch.float32, - ) - rollout_adapter = SGLangMixtureLoRA( - config, - "decoder.layers.0.self_attention.linear_qkv", - 4, - 6, - device=torch.device("cpu"), - dtype=torch.float32, - ) - with torch.no_grad(): - training_adapter.experts.lora_B.normal_(mean=0.0, std=0.2) - rollout_adapter.experts.lora_A.copy_(training_adapter.experts.lora_A) - rollout_adapter.experts.lora_B.copy_(training_adapter.experts.lora_B) - rollout_adapter.router.weight.copy_(training_adapter.router.weight) - x = torch.randn(5, 4) - - training_output, training_decision = training_adapter.forward_with_routing(x) - rollout_output, rollout_decision = rollout_adapter.forward_with_routing(x) - - torch.testing.assert_close(rollout_output, training_output) - torch.testing.assert_close(rollout_decision.pre_topk_probs, training_decision.pre_topk_probs) - torch.testing.assert_close(rollout_decision.post_topk_weights, training_decision.post_topk_weights) - assert torch.equal(rollout_decision.topk_indices, training_decision.topk_indices) - - -def test_sglang_router_uses_fp32_logits_with_bfloat16_parameters(): - adapter = SGLangMixtureLoRA( - _config(), - "decoder.layers.0.self_attention.linear_qkv", - 4, - 6, - device=torch.device("cpu"), - dtype=torch.bfloat16, - ) - - decision = adapter.route(torch.randn(5, 4, dtype=torch.bfloat16)) - - assert decision.pre_topk_probs.dtype == torch.float32 - assert decision.post_topk_weights.dtype == torch.float32 - - -def test_attached_adapter_preserves_base_parameter_name_and_adds_delta(): - model = _fake_qwen_model_with_routed_qkv() - linear = model.model.layers[0].self_attn.qkv_proj - x = torch.randn(3, 4) - base_output = F.linear(x, linear.weight) - with torch.no_grad(): - linear.mixture_lora.experts.lora_B.normal_(mean=0.0, std=0.2) - - output, bias = linear(x) - - assert bias is None - assert output.shape == base_output.shape - assert not torch.equal(output, base_output) - parameter_names = set(dict(model.named_parameters())) - assert "model.layers.0.self_attn.qkv_proj.weight" in parameter_names - assert "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight" in parameter_names - -def test_attached_adapter_uses_quantized_linear_params_dtype(): - linear = _QuantizedTupleLinear(torch.bfloat16) + site_modules = model.mixture_lora_site_modules(0) + model.install_mixture_lora() - attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) - - assert {parameter.dtype for parameter in linear.mixture_lora.parameters()} == {torch.bfloat16} - assert linear.mixture_lora.experts.lora_A.device == linear.weight_packed.device - - -def test_attached_adapter_rejects_missing_floating_point_params_dtype(): - linear = _QuantizedTupleLinear(torch.int32) - - with pytest.raises(TypeError, match="floating-point params_dtype"): - attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) - - -def test_sglang_weight_loader_maps_training_names_and_validates_tensors(): - model = _fake_qwen_model_with_routed_qkv() - prefix = "decoder.layers.0.self_attention.linear_qkv.mixture_lora" - weights = { - f"{prefix}.experts.lora_A": torch.randn(4, 2, 4), - f"{prefix}.experts.lora_B": torch.randn(4, 6, 2), - f"{prefix}.router.weight": torch.randn(4, 4), - } - - loaded_names = load_sglang_mixture_lora_weights(model, weights.items()) - - assert loaded_names == { - "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_A", - "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_B", - "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight", + assert site_modules == { + "linear_qkv": layer.self_attn.qkv_proj, + "linear_proj": layer.self_attn.o_proj, } - parameters = dict(model.named_parameters()) - for source_name, loaded_weight in weights.items(): - target_name = source_name.replace( - "decoder.layers.0.self_attention.linear_qkv", - "model.layers.0.self_attn.qkv_proj", - ) - torch.testing.assert_close(parameters[target_name], loaded_weight) - - with pytest.raises(ValueError, match="has shape"): - load_sglang_mixture_lora_weights(model, [(f"{prefix}.router.weight", torch.randn(4, 5))]) - with pytest.raises(TypeError, match="has dtype"): - load_sglang_mixture_lora_weights( - model, - [(f"{prefix}.router.weight", torch.randn(4, 4, dtype=torch.float64))], - ) - with pytest.raises(ValueError, match="Unknown Mixture-of-LoRA weight"): - load_sglang_mixture_lora_weights( - model, - [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], - ) - - -def test_sglang_weight_loader_skips_layers_owned_by_another_pp_stage(): - model = _fake_qwen_model_with_routed_qkv() - model.model.start_layer = 0 - model.model.end_layer = 1 - - loaded_names = load_sglang_mixture_lora_weights( - model, - [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], - ) - - assert loaded_names == set() - - -def test_sglang_installs_adapters_only_on_layers_owned_by_pp_stage(): - model = nn.Module() - model.mixture_lora_config = _config() - model.model = nn.Module() - model.model.start_layer = 1 - model.model.end_layer = 2 - - missing_layer = nn.Module() - local_layer = nn.Module() - local_layer.self_attn = nn.Module() - local_layer.self_attn.qkv_proj = _TupleLinear(4, 6) - local_layer.self_attn.qkv_proj.input_size = 4 - local_layer.self_attn.qkv_proj.output_size = 6 - local_layer.self_attn.o_proj = _TupleLinear(6, 4) - local_layer.self_attn.o_proj.input_size = 6 - local_layer.self_attn.o_proj.output_size = 4 - model.model.layers = nn.ModuleList([missing_layer, local_layer]) - - EntryClass._install_mixture_lora(model) - - assert not hasattr(missing_layer, "self_attn") - assert local_layer.self_attn.qkv_proj.mixture_lora.site_id == ("decoder.layers.1.self_attention.linear_qkv") - assert local_layer.self_attn.o_proj.mixture_lora.site_id == "decoder.layers.1.self_attention.linear_proj" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_sglang_adapter_can_be_captured_by_cuda_graph(): - adapter = SGLangMixtureLoRA( - _config(), - "decoder.layers.0.self_attention.linear_qkv", - 16, - 24, - device=torch.device("cuda"), - dtype=torch.bfloat16, - ) - x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - adapter(x) - torch.cuda.current_stream().wait_stream(stream) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_output = adapter(x) - graph.replay() - - assert captured_output.shape == (8, 24) - assert torch.isfinite(captured_output).all() + assert layer.self_attn.qkv_proj.mixture_lora.site_id == "decoder.layers.0.self_attention.linear_qkv" + assert layer.self_attn.o_proj.mixture_lora.site_id == "decoder.layers.0.self_attention.linear_proj" diff --git a/tests/models/test_mixture_lora_sglang.py b/tests/models/test_mixture_lora_sglang.py new file mode 100644 index 000000000..3dc01167b --- /dev/null +++ b/tests/models/test_mixture_lora_sglang.py @@ -0,0 +1,303 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + + +# The parity assertions need the training backend; the mixin needs SGLang. +pytest.importorskip("megatron.bridge") +pytest.importorskip("sglang.srt.distributed") + +from relax.backends.megatron.mixture_lora_modules import MixtureLoRAAdapter # noqa: E402 +from relax.models.mixture_lora_sglang import ( # noqa: E402 + MixtureLoraSGLangModelMixin, + SGLangMixtureLoRA, + attach_sglang_mixture_lora, + load_sglang_mixture_lora_weights, +) +from relax.utils.mixture_lora_common import MixtureLoraConfig # noqa: E402 + + +def _config(target_modules=("linear_qkv", "linear_proj")): + return MixtureLoraConfig( + num_experts=4, + rank=2, + top_k=2, + temperature=0.8, + aux_loss_coef=0.01, + alpha=4.0, + target_modules=target_modules, + ) + + +class _TupleLinear(nn.Module): + def __init__(self, input_size: int, output_size: int): + super().__init__() + self.weight = nn.Parameter(torch.randn(output_size, input_size)) + self.input_size = input_size + self.output_size = output_size + + def forward(self, x): + return F.linear(x, self.weight), None + + +class _QuantizedTupleLinear(nn.Module): + def __init__(self, params_dtype): + super().__init__() + self.weight_packed = nn.Parameter(torch.ones(4, 4, dtype=torch.int32), requires_grad=False) + self.params_dtype = params_dtype + + +class _RoutedAttentionModel(MixtureLoraSGLangModelMixin, nn.Module): + """Minimal stand-in for an SGLang causal LM with routed attention.""" + + def __init__(self, num_layers: int, *, start_layer: int = 0, end_layer: int | None = None): + nn.Module.__init__(self) + self.mixture_lora_config = _config() + self.model = nn.Module() + self.model.start_layer = start_layer + self.model.end_layer = num_layers if end_layer is None else end_layer + layers = [] + for layer_id in range(num_layers): + layer = nn.Module() + if start_layer <= layer_id < self.model.end_layer: + layer.self_attn = nn.Module() + layer.self_attn.qkv_proj = _TupleLinear(4, 6) + layer.self_attn.o_proj = _TupleLinear(6, 4) + layers.append(layer) + self.model.layers = nn.ModuleList(layers) + + def mixture_lora_site_modules(self, layer_id: int) -> dict[str, nn.Module]: + attention = self.model.layers[layer_id].self_attn + return {"linear_qkv": attention.qkv_proj, "linear_proj": attention.o_proj} + + +def _fake_model_with_routed_qkv(): + model = nn.Module() + model.model = nn.Module() + layer = nn.Module() + layer.self_attn = nn.Module() + layer.self_attn.qkv_proj = _TupleLinear(4, 6) + model.model.layers = nn.ModuleList([layer]) + attach_sglang_mixture_lora( + layer.self_attn.qkv_proj, + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + ) + return model + + +def test_sglang_dense_adapter_matches_training_adapter(): + config = _config() + training_adapter = MixtureLoRAAdapter( + config, + "linear_qkv", + 4, + 6, + dropout=0.0, + device=torch.device("cpu"), + dtype=torch.float32, + ) + rollout_adapter = SGLangMixtureLoRA( + config, + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + device=torch.device("cpu"), + dtype=torch.float32, + ) + with torch.no_grad(): + training_adapter.experts.lora_B.normal_(mean=0.0, std=0.2) + rollout_adapter.experts.lora_A.copy_(training_adapter.experts.lora_A) + rollout_adapter.experts.lora_B.copy_(training_adapter.experts.lora_B) + rollout_adapter.router.weight.copy_(training_adapter.router.weight) + x = torch.randn(5, 4) + + training_output, training_decision = training_adapter.forward_with_routing(x) + rollout_output, rollout_decision = rollout_adapter.forward_with_routing(x) + + torch.testing.assert_close(rollout_output, training_output) + torch.testing.assert_close(rollout_decision.pre_topk_probs, training_decision.pre_topk_probs) + torch.testing.assert_close(rollout_decision.post_topk_weights, training_decision.post_topk_weights) + assert torch.equal(rollout_decision.topk_indices, training_decision.topk_indices) + + +def test_sglang_router_uses_fp32_logits_with_bfloat16_parameters(): + adapter = SGLangMixtureLoRA( + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 4, + 6, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + decision = adapter.route(torch.randn(5, 4, dtype=torch.bfloat16)) + + assert decision.pre_topk_probs.dtype == torch.float32 + assert decision.post_topk_weights.dtype == torch.float32 + + +def test_attached_adapter_preserves_base_parameter_name_and_adds_delta(): + model = _fake_model_with_routed_qkv() + linear = model.model.layers[0].self_attn.qkv_proj + x = torch.randn(3, 4) + base_output = F.linear(x, linear.weight) + with torch.no_grad(): + linear.mixture_lora.experts.lora_B.normal_(mean=0.0, std=0.2) + + output, bias = linear(x) + + assert bias is None + assert output.shape == base_output.shape + assert not torch.equal(output, base_output) + parameter_names = set(dict(model.named_parameters())) + assert "model.layers.0.self_attn.qkv_proj.weight" in parameter_names + assert "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight" in parameter_names + + +def test_attached_adapter_uses_quantized_linear_params_dtype(): + linear = _QuantizedTupleLinear(torch.bfloat16) + + attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) + + assert {parameter.dtype for parameter in linear.mixture_lora.parameters()} == {torch.bfloat16} + assert linear.mixture_lora.experts.lora_A.device == linear.weight_packed.device + + +def test_attached_adapter_rejects_missing_floating_point_params_dtype(): + linear = _QuantizedTupleLinear(torch.int32) + + with pytest.raises(TypeError, match="floating-point params_dtype"): + attach_sglang_mixture_lora(linear, _config(), "decoder.layers.0.self_attention.linear_qkv", 4, 6) + + +def test_sglang_weight_loader_maps_training_names_and_validates_tensors(): + model = _fake_model_with_routed_qkv() + prefix = "decoder.layers.0.self_attention.linear_qkv.mixture_lora" + weights = { + f"{prefix}.experts.lora_A": torch.randn(4, 2, 4), + f"{prefix}.experts.lora_B": torch.randn(4, 6, 2), + f"{prefix}.router.weight": torch.randn(4, 4), + } + + loaded_names = load_sglang_mixture_lora_weights(model, weights.items()) + + assert loaded_names == { + "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_A", + "model.layers.0.self_attn.qkv_proj.mixture_lora.experts.lora_B", + "model.layers.0.self_attn.qkv_proj.mixture_lora.router.weight", + } + parameters = dict(model.named_parameters()) + for source_name, loaded_weight in weights.items(): + target_name = source_name.replace( + "decoder.layers.0.self_attention.linear_qkv", + "model.layers.0.self_attn.qkv_proj", + ) + torch.testing.assert_close(parameters[target_name], loaded_weight) + + with pytest.raises(ValueError, match="has shape"): + load_sglang_mixture_lora_weights(model, [(f"{prefix}.router.weight", torch.randn(4, 5))]) + with pytest.raises(TypeError, match="has dtype"): + load_sglang_mixture_lora_weights( + model, + [(f"{prefix}.router.weight", torch.randn(4, 4, dtype=torch.float64))], + ) + with pytest.raises(ValueError, match="Unknown Mixture-of-LoRA weight"): + load_sglang_mixture_lora_weights( + model, + [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], + ) + + +def test_sglang_weight_loader_skips_layers_owned_by_another_pp_stage(): + model = _fake_model_with_routed_qkv() + model.model.start_layer = 0 + model.model.end_layer = 1 + + loaded_names = load_sglang_mixture_lora_weights( + model, + [("decoder.layers.1.self_attention.linear_qkv.mixture_lora.router.weight", torch.randn(4, 4))], + ) + + assert loaded_names == set() + + +def test_mixin_installs_adapters_only_on_layers_owned_by_pp_stage(): + model = _RoutedAttentionModel(2, start_layer=1, end_layer=2) + + model.install_mixture_lora() + + assert not hasattr(model.model.layers[0], "self_attn") + local_attention = model.model.layers[1].self_attn + assert local_attention.qkv_proj.mixture_lora.site_id == "decoder.layers.1.self_attention.linear_qkv" + assert local_attention.o_proj.mixture_lora.site_id == "decoder.layers.1.self_attention.linear_proj" + + +def test_mixin_rejects_targets_the_architecture_does_not_expose(): + model = _RoutedAttentionModel(1) + model.mixture_lora_config = _config(target_modules=("linear_qkv", "linear_fc1")) + + with pytest.raises(ValueError, match="Unsupported SGLang Mixture-of-LoRA targets"): + model.install_mixture_lora() + + +def test_mixin_load_weights_splits_routed_tensors_from_base_weights(): + model = _RoutedAttentionModel(1) + model.install_mixture_lora() + base_weights = [] + + class _Base: + def load_weights(self, weights): + base_weights.extend(weights) + return "base-done" + + # Stand in for the SGLang base class the mixin cooperates with. + model.__class__ = type("_Patched", (MixtureLoraSGLangModelMixin, _Base, nn.Module), {}) + prefix = "decoder.layers.0.self_attention.linear_qkv.mixture_lora" + routed = ("router.weight", torch.randn(4, 4)) + + result = model.load_weights( + [ + ("model.layers.0.self_attn.qkv_proj.weight", torch.randn(6, 4)), + (f"{prefix}.{routed[0]}", routed[1]), + ] + ) + + assert result == "base-done" + assert [name for name, _ in base_weights] == ["model.layers.0.self_attn.qkv_proj.weight"] + torch.testing.assert_close( + model.model.layers[0].self_attn.qkv_proj.mixture_lora.router.weight, + routed[1], + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_sglang_adapter_can_be_captured_by_cuda_graph(): + adapter = SGLangMixtureLoRA( + _config(), + "decoder.layers.0.self_attention.linear_qkv", + 16, + 24, + device=torch.device("cuda"), + dtype=torch.bfloat16, + ) + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + adapter(x) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = adapter(x) + graph.replay() + + assert captured_output.shape == (8, 24) + assert torch.isfinite(captured_output).all()