feat: MoVA value expert routing replay

miles/backends/megatron_utils/actor.py

@@ -21,7 +21,7 @@
 from miles.utils.processing_utils import load_tokenizer
 from miles.utils.ray_utils import Box
 from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups
-from miles.utils.replay_base import all_replay_managers
+from miles.utils.replay_base import Replay, all_replay_managers
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils import init_tracking
 from miles.utils.types import RolloutBatch
@@ -39,7 +39,7 @@
 from .lora_utils import is_lora_enabled
 from .model import forward_only, initialize_model_and_optimizer, save, train
 from .parallel import verify_megatron_parallel_state
-from .replay_utils import get_register_replay_list_func
+from .replay_utils import get_replay_columns
 from .update_weight.common import named_params_and_buffers
 from .update_weight.update_weight_from_distributed.broadcast import UpdateWeightFromDistributed
 from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P
@@ -749,8 +749,7 @@ def _fill_replay_data(
         num_microbatches,
         rollout_data,
         data_key: str,
-        replay_list: list,
-        register_replay_list_func,
+        replay_columns: list[tuple[Replay, int]],
         if_sp_region=True,
     ):
         if data_key not in rollout_data:
@@ -807,7 +806,8 @@ def pad_func(data, pad):
                 start, end = seqlen // tp_size * tp_rank, seqlen // tp_size * (tp_rank + 1)
                 replay_data = replay_data[start:end]
 
-            register_replay_list_func(replay_list, replay_data, self.model)
+            for replay, column in replay_columns:
+                replay.record(replay_data[:, column])
 
         del rollout_data[data_key]
 
@@ -892,14 +892,17 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
         data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
 
         for m in all_replay_managers:
-            if self._use_rollout_replay(m):
+            if not self._use_rollout_replay(m):
+                continue
+            for kind, data_key in m.data_keys.items():
+                if data_key not in rollout_data:
+                    continue
                 self._fill_replay_data(
                     data_iterator,
                     num_microbatches,
                     rollout_data,
-                    data_key=m.data_key,
-                    replay_list=m.replays,
-                    register_replay_list_func=get_register_replay_list_func(m),
+                    data_key=data_key,
+                    replay_columns=get_replay_columns(kind, self.model),
                     if_sp_region=m.if_sp_region,
                 )
 

miles/backends/megatron_utils/replay_utils.py

@@ -1,33 +1,17 @@
-from megatron.core.transformer.transformer_block import get_num_layers_to_build
-from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+from megatron.core.utils import unwrap_model
 
-from miles.utils.replay_base import BaseReplayManager, RoutingReplayManager
+from miles.utils.replay_base import Replay
 
 
-def _register_replay_list_moe(replay_list, replay_data, models):
-    layer_indices = []
-    replay_idx = 0
-    for vp_stage, model in enumerate(models):
-        config = model.module.config
-        num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
-        offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
-        for layer_id in range(offset, offset + num_layers_to_build):
-            if isinstance(config.moe_layer_freq, int):
-                if layer_id % config.moe_layer_freq != 0:
-                    continue
-            elif isinstance(config.moe_layer_freq, list):
-                assert len(config.moe_layer_freq) == config.num_layers
-                if config.moe_layer_freq[layer_id] == 0:
-                    continue
-            layer_indices.append(layer_id)
+def get_replay_columns(kind: str, models) -> list[tuple[Replay, int]]:
+    is_value = kind == "value"
+    dense_prefix = unwrap_model(models[0]).config.mova_num_dense_layers if is_value else 0
 
-    for replay_idx, layer_idx in enumerate(layer_indices):
-        layer_data = replay_data[:, layer_idx]
-        replay_list[replay_idx].record(layer_data)
-
-
-def get_register_replay_list_func(manager: BaseReplayManager):
-    if isinstance(manager, RoutingReplayManager):
-        return _register_replay_list_moe
-    else:
-        raise ValueError(f"Unsupported manager type: {type(manager)}")
+    columns = []
+    for model in unwrap_model(models):
+        for layer in model.decoder.layers:
+            owner = getattr(layer.self_attention, "value_projection", None) if is_value else layer.mlp
+            replay = getattr(getattr(owner, "router", None), "routing_replay", None)
+            if replay is not None:
+                columns.append((replay, layer.layer_number - 1 - dense_prefix))
+    return columns

miles/backends/training_utils/data.py

@@ -79,8 +79,9 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch:
                 )
             )
         ]
-    if "rollout_routed_experts" in rollout_data:
-        rollout_data["rollout_routed_experts"] = [torch.from_numpy(r) for r in rollout_data["rollout_routed_experts"]]
+    for key in ("rollout_routed_experts", "rollout_routed_value_experts"):
+        if key in rollout_data:
+            rollout_data[key] = [torch.from_numpy(r) for r in rollout_data[key]]
     return rollout_data
 
 

miles/backends/training_utils/log_utils.py

@@ -200,6 +200,7 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc
                 "loss_masks",
                 "sample_indices",
                 "rollout_routed_experts",
+                "rollout_routed_value_experts",
                 "max_seq_lens",
                 "dynamic_global_batch_size",
                 "weight_versions",

miles/ray/rollout.py

@@ -833,6 +833,8 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
 
         if samples[0].rollout_routed_experts is not None:
             train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples]
+        if samples[0].rollout_routed_value_experts is not None:
+            train_data["rollout_routed_value_experts"] = [sample.rollout_routed_value_experts for sample in samples]
 
         if samples[0].train_metadata is not None:
             train_data["metadata"] = [sample.train_metadata for sample in samples]
@@ -963,6 +965,7 @@ def _stat(xs):
                 "rollout_log_probs",
                 "rollout_sampling_masks",
                 "rollout_routed_experts",
+                "rollout_routed_value_experts",
                 "prompt",
                 "teacher_log_probs",
                 "weight_versions",

miles/rollout/generate_utils/generate_endpoint_utils.py

@@ -97,6 +97,7 @@ async def update_sample_from_response(
 
     # TODO handle multi-turn cases (may need concat instead of assignment)
     sample.rollout_routed_experts = get_rollout_topk_from_response(args, output, sample, "routed_experts")
+    sample.rollout_routed_value_experts = get_rollout_topk_from_response(args, output, sample, "routed_value_experts")
 
     # TODO may unify (currently there are both methods inside Sample and separate functions)
     sample.update_from_meta_info(args, output["meta_info"])
@@ -107,4 +108,9 @@ def get_rollout_topk_from_response(args, output, sample, key):
     if info is None:
         return None
     x = np.frombuffer(pybase64.b64decode(info.encode("ascii")), dtype=np.int32)
-    return x.reshape(len(sample.tokens) - 1, args.num_layers, args.moe_router_topk)
+    num_layers = args.num_layers
+    router_topk = args.moe_router_topk
+    if key == "routed_value_experts":
+        num_layers -= args.mova_num_dense_layers
+        router_topk = args.mova_router_topk
+    return x.reshape(len(sample.tokens) - 1, num_layers, router_topk)

miles/rollout/generate_utils/openai_endpoint_utils.py

@@ -644,6 +644,11 @@ def apply_merged_session_sample(
     if merged.rollout_routed_experts is not None:
         choice = {"meta_info": {"routed_experts": merged.rollout_routed_experts}}
         sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts")
+    if merged.rollout_routed_value_experts is not None:
+        choice = {"meta_info": {"routed_value_experts": merged.rollout_routed_value_experts}}
+        sample.rollout_routed_value_experts = get_rollout_topk_from_response(
+            args, choice, sample, "routed_value_experts"
+        )
 
     routed_experts_shape = None
     routed_experts_decoded_bytes = 0
@@ -770,6 +775,7 @@ def _compute_sample_from_openai_record(
     sample.response_length = len(output_token_ids)
     sample.loss_mask = [1] * len(output_token_ids)
     sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts")
+    sample.rollout_routed_value_experts = get_rollout_topk_from_response(args, choice, sample, "routed_value_experts")
 
     if not hasattr(sample, "metadata") or sample.metadata is None:
         sample.metadata = {}
@@ -817,6 +823,8 @@ def _strip_last_output_tokens_without_decode(sample: Sample, trim_count: int) ->
         sample.loss_mask = sample.loss_mask[:keep_tokens]
     if sample.rollout_routed_experts is not None:
         sample.rollout_routed_experts = sample.rollout_routed_experts[: len(sample.tokens) - 1]
+    if sample.rollout_routed_value_experts is not None:
+        sample.rollout_routed_value_experts = sample.rollout_routed_value_experts[: len(sample.tokens) - 1]
 
 
 def truncate_samples_by_total_tokens(
@@ -868,4 +876,6 @@ def _truncate_sample_output(sample: Sample, keep_tokens: int, tokenizer) -> None
         sample.loss_mask = sample.loss_mask[:keep_tokens]
     if sample.rollout_routed_experts is not None:
         sample.rollout_routed_experts = sample.rollout_routed_experts[: len(sample.tokens) - 1]
+    if sample.rollout_routed_value_experts is not None:
+        sample.rollout_routed_value_experts = sample.rollout_routed_value_experts[: len(sample.tokens) - 1]
     sample.status = Sample.Status.TRUNCATED

miles/rollout/generate_utils/sample_utils.py

@@ -72,6 +72,8 @@ def _fill_defaults(sample: Sample):
         assert obs_len > 0, f"obs_len must be > 0, got {obs_len}"
         if a.rollout_routed_experts is not None:
             assert a.rollout_routed_experts.shape[0] <= b.rollout_routed_experts.shape[0]
+        if a.rollout_routed_value_experts is not None:
+            assert a.rollout_routed_value_experts.shape[0] <= b.rollout_routed_value_experts.shape[0]
         assert a.status == Sample.Status.COMPLETED, f"a.status must be COMPLETED, got {a.status}"
 
         return _create_with_all_fields(
@@ -91,6 +93,7 @@ def _fill_defaults(sample: Sample):
             rollout_log_probs=a.rollout_log_probs + [0.0] * obs_len + b.rollout_log_probs,
             rollout_sampling_mask=None,
             rollout_routed_experts=b.rollout_routed_experts,
+            rollout_routed_value_experts=b.rollout_routed_value_experts,
             remove_sample=_merge_equal_value("remove_sample"),
             status=b.status,
             metadata=metadata,

miles/rollout/session/session_types.py

@@ -30,6 +30,7 @@ class MergedSessionSample(BaseModel):
     weight_versions: list[str] = Field(default_factory=list)
     prefix_cache_meta_infos: list[dict] = Field(default_factory=list)
     rollout_routed_experts: str | None = None
+    rollout_routed_value_experts: str | None = None
 
 
 class GetMergedSessionResponse(BaseModel):

miles/rollout/session/sessions.py

@@ -164,18 +164,20 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord:
                 compact_meta_info[key] = meta_info[key]
 
         # Preserve routed experts if present. get_rollout_topk_from_response()
-        # may read this from choice or meta_info depending on implementation.
-        if "routed_experts" in meta_info:
-            compact_meta_info["routed_experts"] = meta_info["routed_experts"]
+        # may read these from choice or meta_info depending on implementation.
+        for key in ("routed_experts", "routed_value_experts"):
+            if key in meta_info:
+                compact_meta_info[key] = meta_info[key]
 
         compact_choice = {
             "prompt_token_ids": choice.get("prompt_token_ids", []),
             "finish_reason": choice.get("finish_reason"),
             "meta_info": compact_meta_info,
         }
 
-        if "routed_experts" in choice:
-            compact_choice["routed_experts"] = choice["routed_experts"]
+        for key in ("routed_experts", "routed_value_experts"):
+            if key in choice:
+                compact_choice[key] = choice[key]
 
         return SessionRecord(
             timestamp=record.timestamp,
@@ -262,6 +264,7 @@ def _merge_session_records_to_sample(
         weight_versions: list[str] = []
         prefix_cache_meta_infos: list[dict] = []
         routed_experts = None
+        routed_value_experts = None
         prev_checkpoint: list[int] | None = None
         final_cursor = 0
         final_status = "completed"
@@ -345,6 +348,7 @@ def _merge_session_records_to_sample(
             )
 
             routed_experts = meta_info.get("routed_experts", choice.get("routed_experts"))
+            routed_value_experts = meta_info.get("routed_value_experts", choice.get("routed_value_experts"))
             final_status = _status_from_finish_reason(choice.get("finish_reason"))
             if not is_last:
                 assert (
@@ -401,6 +405,7 @@ def _merge_session_records_to_sample(
             weight_versions=weight_versions,
             prefix_cache_meta_infos=prefix_cache_meta_infos,
             rollout_routed_experts=routed_experts,
+            rollout_routed_value_experts=routed_value_experts,
         )
 
     @app.get("/sessions/{session_id}")

miles/rollout/sglang_rollout.py

@@ -8,14 +8,14 @@
 from typing import Any
 
 import numpy as np
-import pybase64
 import sglang_router
 from packaging.version import parse
 from tqdm import tqdm
 
 from miles.backends.megatron_utils.lora_utils import LORA_ADAPTER_NAME, is_lora_enabled
 from miles.rollout.base_types import GenerateFnInput, RolloutFnEvalOutput, RolloutFnTrainOutput
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
+from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response
 from miles.rollout.inference_rollout.compatibility import load_generate_function
 from miles.utils import dumper_utils
 from miles.utils.async_utils import run
@@ -216,15 +216,12 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
             sample.rollout_log_probs = []
         sample.rollout_log_probs += new_response_log_probs
 
-    if "routed_experts" in output["meta_info"]:
-        sample.rollout_routed_experts = np.frombuffer(
-            pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")),
-            dtype=np.int32,
-        ).reshape(
-            len(sample.tokens) - 1,
-            args.num_layers,
-            args.moe_router_topk,
-        )
+    for key, attr in (
+        ("routed_experts", "rollout_routed_experts"),
+        ("routed_value_experts", "rollout_routed_value_experts"),
+    ):
+        if key in output["meta_info"]:
+            setattr(sample, attr, get_rollout_topk_from_response(args, output, sample, key))
 
     sample.update_from_meta_info(args, output["meta_info"])
 

miles/utils/replay_base.py

@@ -52,6 +52,7 @@ def clear_forward(self):
 class BaseReplayManager:
     name: str = ""
     filename: str = ""
+    data_keys: dict[str, str] = {}
 
     def __init__(self):
         self.replays: list[Replay] = []
@@ -220,7 +221,7 @@ def check_replay_result(self, old_topk_fn, scores, topk, top_indices, *args, **k
 class RoutingReplayManager(BaseReplayManager):
     name = "routing"
     filename = "routing_replay.pt"
-    data_key = "rollout_routed_experts"
+    data_keys = {"ffn": "rollout_routed_experts", "value": "rollout_routed_value_experts"}
     if_sp_region = True
     enable_check_replay_result = False
     replay_check_threshold = 1e-2

miles/utils/types.py

@@ -92,6 +92,9 @@ class Sample:
     rollout_routed_experts: numpy.ndarray | None = (
         None  # Routed experts from rollout engine. shape: (num_tokens-1, num_layers, moe_router_topk), dtype=int32
     )
+    rollout_routed_value_experts: numpy.ndarray | None = (
+        None  # Value-routed experts. shape: (num_tokens-1, num_layers - mova_num_dense_layers, mova_router_topk), dtype=int32
+    )
     remove_sample: bool = False
 
     class Status(Enum):
@@ -239,6 +242,10 @@ def validate(self):
             actual = len(self.rollout_routed_experts)
             expect = len(self.tokens) - 1
             assert actual == expect, f"rollout_routed_experts length ({actual}) != len(tokens) - 1 ({expect})"
+        if self.rollout_routed_value_experts is not None:
+            actual = len(self.rollout_routed_value_experts)
+            expect = len(self.tokens) - 1
+            assert actual == expect, f"rollout_routed_value_experts length ({actual}) != len(tokens) - 1 ({expect})"
 
     def strip_last_output_tokens(self, n: int, tokenizer) -> None:
         """Remove the last *n* output tokens and all associated per-token info."""
@@ -258,6 +265,8 @@ def strip_last_output_tokens(self, n: int, tokenizer) -> None:
         self.response = tokenizer.decode(self.tokens[-self.response_length :]) if self.response_length > 0 else ""
         if self.rollout_routed_experts is not None:
             self.rollout_routed_experts = self.rollout_routed_experts[:-n]
+        if self.rollout_routed_value_experts is not None:
+            self.rollout_routed_value_experts = self.rollout_routed_value_experts[:-n]
 
     def reset_for_retry(self) -> None:
         """Reset generated outputs so the original prompt can be re-sampled.
@@ -276,6 +285,7 @@ def reset_for_retry(self) -> None:
         self.rollout_log_probs = None
         self.rollout_sampling_mask = None
         self.rollout_routed_experts = None
+        self.rollout_routed_value_experts = None
         self.status = Sample.Status.ABORTED
         self.non_generation_time = 0.0
         self.spec_info = Sample.SpecInfo()