perf: accelerate Qwen3.5 OpenR1MM training

docker/patch/megatron/20260915-4a4de8657.patch

@@ -112,7 +112,7 @@ index ba632f3..99e8aa9 100644
      freeze_language_model: bool = False
      freeze_vision_model: bool = False
 diff --git a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
-index 4ef12d7..01c5180 100644
+index 4ef12d7..7bf9ad0 100644
 --- a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
 +++ b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
 @@ -16,7 +16,9 @@ import logging
@@ -275,6 +275,15 @@ index 4ef12d7..01c5180 100644
 
              if self.config.sequence_parallel:
                  tp_group = self.config._pg_collection.tp if self.config._pg_collection is not None else None
+@@ -532,7 +622,7 @@ class KimiK25VLModel(MegatronModule):
+             runtime_gather_output=runtime_gather_output,
+             packed_seq_params=packed_seq_params,
+         )
+-        if cp_size > 1:
++        if cp_size > 1 and loss_mask is not None:
+             return outputs, loss_mask
+         return outputs
+
 diff --git a/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/megatron/bridge/models/qwen/qwen3_moe_bridge.py
 index a6476f4..cab29da 100644
 --- a/megatron/bridge/models/qwen/qwen3_moe_bridge.py

docker/patch/sglang/v0.5.17.patch

@@ -1154,7 +1154,7 @@ index 6eb0ad1a..d27815f3 100644
          """Load a new LoRA adapter without re-launching the engine."""
 
 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
-index ac38b466..ea54bac1 100644
+index ac38b46..abdc04f 100644
 --- a/python/sglang/srt/entrypoints/http_server.py
 +++ b/python/sglang/srt/entrypoints/http_server.py
 @@ -132,6 +132,7 @@ from sglang.srt.managers.io_struct import (
@@ -1208,6 +1208,15 @@ index ac38b466..ea54bac1 100644
 
      server_args = _global_state.tokenizer_manager.server_args
 
+@@ -946,7 +955,7 @@ async def classify_request(obj: EmbeddingReqInput, request: Request):
+ @app.api_route("/flush_cache", methods=["GET", "POST"])
+ @auth_level(AuthLevel.ADMIN_OPTIONAL)
+ async def flush_cache(timeout: float = Query(0.0, ge=0.0)):
+-    """Flush the radix cache."""
++    """Flush runtime caches, including multimodal embeddings."""
+     ret = await _global_state.tokenizer_manager.flush_cache(timeout_s=timeout)
+     if ret.success:
+         content = (
 @@ -1409,6 +1418,22 @@ async def update_weights_from_ipc(
          return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
 
@@ -1546,18 +1555,26 @@ index a45d122f..0334e353 100755
          ):
              # Even the last remaining request cannot fit in memory.
 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
-index a9615a14c5..aca7bd7061 100644
+index a9615a1..b2dfb92 100644
 --- a/python/sglang/srt/managers/scheduler.py
 +++ b/python/sglang/srt/managers/scheduler.py
-@@ -141,6 +141,7 @@ from sglang.srt.managers.io_struct import (
+@@ -99,6 +99,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
+ from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config
+ from sglang.srt.lora.lora_drainer import LoRADrainer
+ from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader
++from sglang.srt.managers import mm_utils
+ from sglang.srt.managers.disagg_service import maybe_create_ascend_config_store
+ from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+ from sglang.srt.managers.io_struct import (
+@@ -141,6 +142,7 @@ from sglang.srt.managers.io_struct import (
      LoadLoRAAdapterReqOutput,
      OpenSessionReqInput,
      PauseGenerationReqInput,
 +    PostProcessWeightsReqInput,
      ProfileReq,
      ReleaseMemoryOccupationReqInput,
      RemoveExternalCorpusReqInput,
-@@ -161,6 +162,8 @@ from sglang.srt.managers.io_struct import (
+@@ -161,6 +163,8 @@ from sglang.srt.managers.io_struct import (
      TokenizedGenerateReqInput,
      UnloadLoRAAdapterReqInput,
      UnloadLoRAAdapterReqOutput,
@@ -1566,7 +1583,7 @@ index a9615a14c5..aca7bd7061 100644
      UpdateWeightFromDiskReqInput,
      UpdateWeightsFromDistributedReqInput,
      UpdateWeightsFromIPCReqInput,
-@@ -280,6 +283,8 @@ from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+@@ -280,6 +284,8 @@ from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
  from sglang.srt.speculative.dflash_utils import validate_dflash_request
  from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
  from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -1575,7 +1592,7 @@ index a9615a14c5..aca7bd7061 100644
  from sglang.srt.utils import (
      DynamicGradMode,
      configure_gc_logger,
-@@ -1539,6 +1544,10 @@ class Scheduler(
+@@ -1539,6 +1545,10 @@ class Scheduler(
                      UpdateWeightsFromIPCReqInput,
                      self.weight_updater.update_weights_from_ipc,
                  ),
@@ -1586,7 +1603,7 @@ index a9615a14c5..aca7bd7061 100644
                  (
                      GetWeightsByNameReqInput,
                      self.weight_updater.get_weights_by_name,
-@@ -1571,6 +1580,10 @@ class Scheduler(
+@@ -1571,6 +1581,10 @@ class Scheduler(
                      LoadLoRAAdapterFromTensorsReqInput,
                      self.load_lora_adapter_from_tensors,
                  ),
@@ -1597,7 +1614,7 @@ index a9615a14c5..aca7bd7061 100644
                  (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
                  (PauseGenerationReqInput, self.pause_generation),
                  (ContinueGenerationReqInput, self.continue_generation),
-@@ -1627,6 +1640,8 @@ class Scheduler(
+@@ -1627,6 +1641,8 @@ class Scheduler(
          self.tree_cache.release_host_resources()
          if self.decode_offload_manager is not None:
              self.decode_offload_manager.release_host_resources()
@@ -1606,7 +1623,29 @@ index a9615a14c5..aca7bd7061 100644
 
      def run_event_loop(self) -> None:
          """Run the scheduler's event loop.
-@@ -4457,6 +4472,12 @@ class Scheduler(
+@@ -4133,7 +4149,7 @@ class Scheduler(
+         return DetachHiCacheStorageReqOutput(success=False, message=msg)
+
+     def flush_cache(self, empty_cache: bool = True):
+-        """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache."""
++        """Flush runtime caches and optionally empty the device allocator cache."""
+         if self.is_fully_idle():
+             self.cur_batch_for_debug = None
+             self.last_batch = None
+@@ -4146,6 +4162,12 @@ class Scheduler(
+             if self.draft_worker:
+                 self.draft_worker.clear_cache_pool()
+
++            # Drop cached embeddings so weight updates cannot reuse stale encoder
++            # outputs. Do this before trimming the allocator so colocated workloads
++            # can reclaim the memory.
++            if mm_utils.embedding_cache is not None:
++                mm_utils.embedding_cache.clear()
++
+             if empty_cache:
+                 current_platform.empty_cache()
+             # Per-DP-group leader logs once: ranks within a DP group are
+@@ -4457,6 +4479,12 @@ class Scheduler(
                  # The request will still run one decode forward pass.
                  # Then we reuse all existing code to clean up the KV cache allocation.
                  logger.debug(f"Abort running request. {req.rid=}")
@@ -1619,7 +1658,7 @@ index a9615a14c5..aca7bd7061 100644
                  req.to_finish = FINISH_ABORT()
 
      def _pause_engine(self) -> Tuple[List[Req], int]:
-@@ -4669,6 +4690,14 @@ class Scheduler(
+@@ -4669,6 +4697,14 @@ class Scheduler(
          result = self.tp_worker.load_lora_adapter_from_tensors(recv_req)
          return result
 

relax/backends/megatron/actor.py

@@ -91,6 +91,7 @@
 from relax.utils.tracking_utils import init_tracking
 from relax.utils.training import train_dump_utils
 from relax.utils.training.data_fields import build_data_fields
+from relax.utils.training.inline_log_probs import can_reuse_colocate_log_probs
 from relax.utils.training.routing_replay import RoutingReplay
 from relax.utils.types import RolloutBatch
 from relax.utils.utils import (
@@ -1971,6 +1972,12 @@ def train_actor(
             data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
         else:
             data_iterator, num_microbatches = prepared_data_iterator, prepared_num_microbatches
+        inline_old_log_probs = hasattr(self, "weights_backuper") and can_reuse_colocate_log_probs(
+            self.args,
+            num_microbatches,
+            backup_tags=self.weights_backuper.backup_tags,
+            rollout_data_postprocess=self.rollout_data_postprocess,
+        )
         # Create a separate iterator with a larger token budget for ref/teacher log-probs
         if (
             not is_offline_mode(self.args)
@@ -2037,7 +2044,11 @@ def train_actor(
                     )
 
                 self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
-                if not standard_dpo and (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics):
+                if (
+                    not standard_dpo
+                    and not inline_old_log_probs
+                    and (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics)
+                ):
                     if self.args.use_routing_replay:
                         if self.args.use_rollout_routing_replay:
                             os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
@@ -2069,7 +2080,7 @@ def train_actor(
                 # because we may need normalize the whole rollout.
                 capture_hooks.begin_rollout_for(self.args, rollout_id)
                 try:
-                    compute_advantages_and_returns(self.args, rollout_data)
+                    compute_advantages_and_returns(self.args, rollout_data, inline_log_probs=inline_old_log_probs)
                     capture_hooks.capture_rollout_advantage(rollout_data=rollout_data, args=self.args)
                 finally:
                     capture_hooks.end_rollout_for()

relax/backends/megatron/loss.py

@@ -618,7 +618,9 @@ def get_values(
     return torch.empty((0,), device=logits.device), res
 
 
-def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) -> None:
+def compute_advantages_and_returns(
+    args: Namespace, rollout_data: RolloutBatch, *, inline_log_probs: bool = False
+) -> None:
     """Compute advantages and returns in-place based on
     `args.advantage_estimator`.
 
@@ -665,6 +667,11 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
     if args.kl_coef == 0 or not log_probs:
         # when kl_coef is 0, we won't compute ref_log_prob
         xs = log_probs if log_probs is not None else values
+        if xs is None and inline_log_probs and args.kl_coef == 0:
+            # Only shape/device are used. Sampling probabilities are never PPO old log-probs.
+            xs = rollout_data.get("rollout_log_probs")
+            if xs is None:
+                raise ValueError("Colocate true-on-policy GRPO requires rollout log-probs for zero-KL shapes.")
         kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
     else:
         kl = [

relax/backends/megatron/weight_update/bridge_converter.py

@@ -428,6 +428,30 @@ def _lookup_task(self, global_name: str) -> tuple[str, Any]:
 
         return global_name, None
 
+    def supports_packed_experts(self, names: Collection[str]) -> bool:
+        """Whether mappings export canonical gated experts accepted by
+        FusedMoE."""
+        for name in names:
+            global_name = _normalize_global_name(name)
+            task = self._bridge_task_map.get(global_name)
+            mapping = task.mapping if task is not None else self._lookup_mapping(global_name)[1]
+            hf_param = getattr(mapping, "hf_param", None)
+            outputs = list(hf_param.values()) if isinstance(hf_param, dict) else [hf_param]
+            projections = set()
+            for output in outputs:
+                match = re.fullmatch(
+                    r"model\.(?:language_model\.)?layers\.\d+\.mlp\.experts\.(?:\d+\.)?"
+                    r"(gate|up|down|gate_up)_proj(?:\.weight)?",
+                    str(output),
+                )
+                if match is None:
+                    return False
+                projections.add(match[1])
+            expected = ({"gate", "up"}, {"gate_up"}) if ".linear_fc1." in name else ({"down"},)
+            if projections not in expected:
+                return False
+        return True
+
     def _lookup_mapping(self, global_name: str) -> tuple[str, Any]:
         """Resolve ``global_name`` straight from the mapping registry.