Feat/task25 mixture of lora

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' }
               ]
             },
             {

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/<site_id>/...` and `molora/global/...`:
+
+- `expert_<id>_pre_topk_mean_prob`;
+- `expert_<id>_post_topk_mean_weight`;
+- `expert_<id>_selection_share`;
+- `expert_<id>_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.

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/<site_id>/...` 和 `molora/global/...`：
+
+- `expert_<id>_pre_topk_mean_prob`；
+- `expert_<id>_post_topk_mean_weight`；
+- `expert_<id>_selection_share`；
+- `expert_<id>_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 参数。

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 (

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(