diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c2157f054..087ddd4d4 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' } ] }, { @@ -383,7 +384,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..4a2bc8854 --- /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 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 + +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. + +## 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. 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 new file mode 100644 index 000000000..2fd91c667 --- /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，并报告错误。由于部分未带版本号的分块可能已经传输，需要重启或完整重同步 rollout engine 后才能继续提供服务。只有所有 rank 都确认前序分块成功后，才会发送带新版本号的最后一个分块。 + +## Checkpoint + +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 维度是否与当前配置一致。 + +## 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 中的值。这份 recipe 使用 BF16，也可以在命令末尾继续追加 Relax 参数。 diff --git a/relax/backends/megatron/arguments.py b/relax/backends/megatron/arguments.py index 27aaed382..5a6ddd5de 100644 --- a/relax/backends/megatron/arguments.py +++ b/relax/backends/megatron/arguments.py @@ -215,6 +215,14 @@ 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 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 " @@ -230,6 +238,11 @@ 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.") + 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/loss.py b/relax/backends/megatron/loss.py index 372a139cf..d0c9d08f5 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_modules import get_microbatch_objective_scale def get_responses( @@ -1415,34 +1416,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_modules.py b/relax/backends/megatron/mixture_lora_modules.py new file mode 100644 index 000000000..3c9b4067f --- /dev/null +++ b/relax/backends/megatron/mixture_lora_modules.py @@ -0,0 +1,1113 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""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 +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field +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_common import ( + MixtureLoraConfig, + RoutingDecision, + RoutingStatistics, + compute_routing_statistics, + mixture_lora_tp_partition_dims, + 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 + objective_weight: 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 + activation_layout: str + 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) + + 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") + 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") + 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 x.ndim != 3 or mask.ndim != 2: + raise ValueError( + 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: + 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, + *, + 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) + 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 = ( + 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 = 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=recorded_balance_loss, + 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_divisor + ) + 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, + ) + + +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, ...], + *, + 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 _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]) + + +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.""" + + def __init__( + self, + config: MixtureLoraConfig, + 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, + 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 + 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, local_rank, local_input_size, device=device, dtype=dtype) + ) + 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: + 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) + else: + from megatron.core.tensor_parallel.layers import ( + _initialize_affine_weight_cpu, + _initialize_affine_weight_gpu, + ) + + # 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: + _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) + 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): + """Per-token linear router with FP32 logits.""" + + def __init__( + self, + num_experts: int, + input_size: int, + *, + 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()) + + +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, + input_is_parallel: bool = False, + sequence_parallel: bool = False, + 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(): + 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}") + 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.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 + 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, + 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, + 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, + 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 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: + 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]: + 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(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, + 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: + 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, + 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: + 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, + input_is_parallel=input_is_parallel, + sequence_parallel=sequence_parallel, + 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 + + 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 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 = { + 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 + } + 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, vp_stage: int | None = None): + """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") + _, 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: + tp_group = parallel_state.get_tensor_model_parallel_group() + return MixtureParallelLinearAdapter( + module, + self.mixture_config, + site_id, + 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( + target_modules=list(config.target_modules), + mixture_config=config, + dropout=dropout, + ) + + +__all__ = [ + "MixtureLoRARoutingContext", + "MixtureLoRARoutingRecord", + "MixtureLoRAAdapter", + "MixtureLoRAExperts", + "MixtureLoRARouter", + "MixtureParallelLinearAdapter", + "MegatronDenseRoutedLoRAExecutor", + "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", + "mark_routed_tensor_parallel_shard", + "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.py b/relax/backends/megatron/model.py index 077e461d9..7fa58f0fb 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -8,7 +8,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 @@ -33,7 +33,7 @@ from relax.utils.env import Envs from relax.utils.logging_utils import get_logger from relax.utils.megatron_bridge_utils import patch_megatron_model -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 @@ -47,12 +47,138 @@ from .checkpoint import load_checkpoint, save_checkpoint from .data import DataIterator, get_batch from .loss import loss_function +from .mixture_lora_modules import ( + MixtureLoRARoutingContext, + 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 logger = get_logger(__name__) +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_adapters = [ + module + for model_chunk in model + for module in model_chunk.modules() + if isinstance(module, MixtureParallelLinearAdapter) + ] + 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(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, + 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( + 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(), + activation_layout=args.qkv_format, + 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), + ) + + 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. @@ -986,6 +1112,12 @@ 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_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 def forward_step( data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False @@ -1005,7 +1137,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. @@ -1039,6 +1171,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" @@ -1053,14 +1204,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) @@ -1113,9 +1265,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 @@ -1172,6 +1326,17 @@ def forward_step( if _dcp_orig_cp_group is not None: inner.pg_collection.cp = _dcp_orig_cp_group + 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 # This check must happen before optimizer.step() as gradients may be modified during step if args.ci_test and args.enable_mtp_training: @@ -1256,6 +1421,7 @@ def forward_step( # 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 @@ -1527,7 +1693,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): @@ -1718,7 +1884,7 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP], *, force_sync: bo fp8_writer.result.modules_to_not_convert, ) - 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/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 7e9805608..60406d84d 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,45 @@ 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_modules 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: + raise RuntimeError("Mixture-of-LoRA is enabled but its validated configuration is missing") + 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) + 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: - 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/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 ad05ae770..25f487958 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 @@ -109,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", @@ -243,6 +247,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. @@ -285,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( @@ -299,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/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..a91d48dfa --- /dev/null +++ b/relax/backends/megatron/weight_update/mixture_lora_sync.py @@ -0,0 +1,344 @@ +# 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_common import ( + MixtureLoraStateSpec, + build_mixture_lora_state_specs, + mixture_lora_tp_partition_dims, +) + +from .common import named_params_and_buffers +from .synchronized_send import raise_on_any_rank_failure + + +@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 _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: + 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 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( + 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, + ) + + 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, + ) + 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. + + 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. + """ + + 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() + 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: + 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: + 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/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 344ac9436..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 @@ -25,11 +24,14 @@ 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 .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, @@ -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 @@ -233,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 @@ -271,6 +288,105 @@ 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() + quantization_restored = False + + 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 ( + 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, + ) + 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) + + def resume_generation(*, finish_quantization: bool) -> None: + if rank != 0: + return + if ( + finish_quantization + and quantization_restored + and 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]) + + for phase_name, operation in (("pause and flush", pause_and_flush), ("send weights", send_weights)): + local_error, phase_failed = run_synchronized_phase( + operation, description=f"Mixture-of-LoRA weight update phase {phase_name!r}" + ) + if not phase_failed: + continue + + 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 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 = 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: + 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 + + def _send_weight_update_stream(self, updates) -> None: + """Pipeline conversion collectives with the preceding IPC request. + + ``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. + """ + + 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 step. @@ -337,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()) @@ -476,8 +585,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 +601,8 @@ 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, + use_host_tensors=(self.mixture_lora_enabled and getattr(self.args, "sglang_pp_size", 1) > 1), ) all_refs.extend(refs_colocated) @@ -494,7 +610,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 +626,8 @@ def _send_to_colocated_engine( ipc_engine, ipc_gather_src, ipc_gather_group, - weight_version, + 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. @@ -526,7 +643,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 ] @@ -552,7 +672,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 @@ -581,18 +706,38 @@ 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, "load_format": "flattened_bucket", - "weight_version": str(weight_version), + "weight_version": None if weight_version is None else str(weight_version), } refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs)) 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/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 15ca1d172..f3b33fa04 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -32,7 +32,13 @@ 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, + is_mixture_lora_enabled, +) +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, @@ -230,6 +236,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("[]") @@ -332,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 @@ -339,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 @@ -557,14 +579,18 @@ 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. - external_pkg = getattr(self.args, "sglang_external_model_package", None) + 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), + ) 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 / @@ -1450,6 +1476,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/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/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/__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..e755e1ef9 --- /dev/null +++ b/relax/models/qwen3_mixture_lora/sglang/model.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""SGLang Qwen3 external model with token-routed LoRA experts.""" + +from sglang.srt.models.qwen3 import Qwen3ForCausalLM as SGLangQwen3ForCausalLM +from torch import nn + +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(MixtureLoraSGLangModelMixin, SGLangQwen3ForCausalLM): + """Qwen3 external model that routes LoRA experts at attention + projections.""" + + 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/relax/utils/arguments.py b/relax/utils/arguments.py index 3bf40c090..695e1f14e 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -15,6 +15,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, @@ -1561,6 +1562,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, @@ -2961,6 +2986,96 @@ 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.") + 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( + "--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 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.") + + 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"): @@ -2999,22 +3114,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/env.py b/relax/utils/env.py index b0a9397a3..2eb8769f6 100644 --- a/relax/utils/env.py +++ b/relax/utils/env.py @@ -209,6 +209,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/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/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 7fd78a9ab..3d85c3ba8 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/relax/utils/megatron_peft_utils.py @@ -7,11 +7,18 @@ import torch +from relax.utils.mixture_lora_common 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`` # 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]: @@ -44,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 @@ -104,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( @@ -187,6 +268,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. @@ -300,12 +401,18 @@ 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", + "build_mixture_lora_config", "build_lora_peft", ] diff --git a/relax/utils/mixture_lora_common.py b/relax/utils/mixture_lora_common.py new file mode 100644 index 000000000..a84bb2ca1 --- /dev/null +++ b/relax/utils/mixture_lora_common.py @@ -0,0 +1,645 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""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 +import os +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"} +_CONFIG_JSON_FIELDS = { + "schema_version", + "num_experts", + "rank", + "top_k", + "temperature", + "aux_loss_coef", + "alpha", + "target_modules", +} + + +@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 + + +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: + 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): + 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.""" + + 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, + ), + ) + + +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. + + 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 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 + 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") + + +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/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..6ba09df44 --- /dev/null +++ b/scripts/training/text/run-qwen3-4B-mixture-lora-8xgpu.sh @@ -0,0 +1,148 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3-4B Mixture-of-LoRA GRPO on DAPO math with 8 colocated GPUs. + +set -ex + +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}" + --dist-ckpt-optim-fully-reshardable +) + +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 + --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 \ + --bf16 \ + --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" diff --git a/tests/backends/megatron/test_mixture_lora.py b/tests/backends/megatron/test_mixture_lora.py new file mode 100644 index 000000000..dfcc518e2 --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora.py @@ -0,0 +1,1003 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import copy +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_modules import ( + MixtureLoRAAdapter, + MixtureLoRAExperts, + MixtureLoRARouter, + 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, + mixture_lora_metrics_from_packed_records, + pack_mixture_lora_routing_records, +) +from relax.utils import megatron_bridge_utils +from relax.utils.mixture_lora_common import MixtureLoraConfig, compute_routing_statistics + + +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"), + ) + + +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__() + 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() == {} + + +@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) + 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._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) + + +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"), + [ + (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), + activation_layout="bshd", + ) + 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 + + +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( + _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), + activation_layout="bshd", + ) + 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), + activation_layout="bshd", + 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)) + 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), + activation_layout="bshd", + ) + 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), + activation_layout="bshd", + ) + + 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): + 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), + activation_layout="bshd", + ) + 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") + 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], + 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, + ) + + 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()) + + +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") + + 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), + activation_layout="bshd", + ) + 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 + 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() + + +@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_mixture_lora_arguments.py b/tests/backends/megatron/test_mixture_lora_arguments.py new file mode 100644 index 000000000..804ba72f4 --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora_arguments.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from types import SimpleNamespace + +import pytest + + +# 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(): + 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(model_type="qwen3", num_experts=8) + + 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(model_type="qwen3")) + + +@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(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(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(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")) 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..5f362d0c4 --- /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_modules import MixtureParallelLinearAdapter +from relax.utils.mixture_lora_common 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 new file mode 100644 index 000000000..fbab86c71 --- /dev/null +++ b/tests/backends/megatron/test_mixture_lora_distributed.py @@ -0,0 +1,899 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +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_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_common import ( + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + compute_routing_statistics, + route_topk, +) + + +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), + activation_layout="bshd", + ) + 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) + + +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): + # 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) + + +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]), + activation_layout="bshd", + ) + + +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", + 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), + activation_layout="bshd", + 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, + 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), + activation_layout="bshd", + ) + 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), + activation_layout="bshd", + ) + 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): + 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) + + +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]) + + # 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() + + +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) + + +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): + 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) diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 82c20510b..940215987 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -251,6 +251,50 @@ 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_modules") + + 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 + 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_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)) + 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, None)] + + +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/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/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..d5284fcd8 --- /dev/null +++ b/tests/backends/megatron/weight_update/test_mixture_lora_weight_sync.py @@ -0,0 +1,704 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from contextlib import contextmanager +from dataclasses import replace +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock, call, 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 +# 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 ( # noqa: E402 + _send_to_colocated_engine, + iter_mixture_weight_updates, +) +from relax.utils.mixture_lora_common 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) + 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): + 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_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, + } + + +@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) + 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 _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 + + 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_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=[]) + + 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_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_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=[]) + + 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( + 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.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() + 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 + + +def test_update_failure_keeps_generation_paused_and_preserves_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("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"), + ): + updater._update_weights_mixture_lora() + + assert events == ["pause", "flush", "update"] + assert updater.weight_version == 4 + assert updater._mixture_lora_sync.base_sync_done is False + + +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) + 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.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)]) + + 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 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] diff --git a/tests/backends/sglang/test_mixture_lora.py b/tests/backends/sglang/test_mixture_lora.py new file mode 100644 index 000000000..287227663 --- /dev/null +++ b/tests/backends/sglang/test_mixture_lora.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import os +from types import SimpleNamespace + +import pytest + + +# 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_common import ( # noqa: E402 + 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.setenv("RELAX_MIXTURE_LORA_CONFIG", "stale-policy-config") + + 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_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") + + +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 + + +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/backends/sglang/test_router_registration.py b/tests/backends/sglang/test_router_registration.py index 84cc8f3e0..956987fd1 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_common") + mixture_lora.configure_mixture_lora_external_model = lambda *_args, **_kwargs: None + 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") yield module 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..988d68e5c --- /dev/null +++ b/tests/models/qwen3_mixture_lora/test_sglang_model.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + + +# 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("sglang.srt.models.qwen3") + +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 + + +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)) + self.input_size = input_size + self.output_size = output_size + + def forward(self, x): + return F.linear(x, self.weight), None + + +def test_external_entry_class_overrides_the_checkpoint_architecture(): + assert EntryClass.__name__ == "Qwen3ForCausalLM" + + +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]) + + site_modules = model.mixture_lora_site_modules(0) + model.install_mixture_lora() + + assert site_modules == { + "linear_qkv": layer.self_attn.qkv_proj, + "linear_proj": layer.self_attn.o_proj, + } + 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() 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" diff --git a/tests/utils/test_arguments_mixture_lora.py b/tests/utils/test_arguments_mixture_lora.py new file mode 100644 index 000000000..1b69d456a --- /dev/null +++ b/tests/utils/test_arguments_mixture_lora.py @@ -0,0 +1,215 @@ +# 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, + tensor_model_parallel_size=2, + ) + 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_rank": 15}, "divisible"), + ({"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"), + ({"dynamic_context_parallel": True}, "static context parallelism"), + ({"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) 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. diff --git a/tests/utils/test_mixture_lora_routing.py b/tests/utils/test_mixture_lora_routing.py new file mode 100644 index 000000000..5b375d756 --- /dev/null +++ b/tests/utils/test_mixture_lora_routing.py @@ -0,0 +1,400 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import math + +import pytest +import torch + +from relax.utils.mixture_lora_common import ( + MIXTURE_LORA_SCHEMA_VERSION, + DenseRoutedLoRAExecutor, + MixtureLoraConfig, + MixtureLoraStateSpec, + RoutedLoRAExecutor, + RoutedLoRAParallelContext, + RoutingDecision, + 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, +) + + +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 _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 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): + 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) + + 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_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)) + + 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( + ("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"), + [ + (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))