support moe model deterministic

Author: YanhuiDua

xtuner/v1/ray/environment/single_turn_env.py

@@ -16,6 +16,7 @@
     update_rollout_item,
 )
 from xtuner.v1.ray.environment.base_env import BaseEnvironment
+from xtuner.v1.ray.utils import build_deterministic_session_id, deterministic_item_sort_key
 from xtuner.v1.utils import XTUNER_DETERMINISTIC, get_logger, ray_method
 
 
@@ -93,14 +94,19 @@ async def generate(  # type: ignore[override]
         if extra_params is None:
             extra_params = {}
         if self.rollout_controller:
+            if XTUNER_DETERMINISTIC:
+                group_data_items = sorted(group_data_items, key=deterministic_item_sort_key)
             response_future = []
             for i, sample in enumerate(group_data_items):
                 rollout_extra_info = dict(sample.data.extra_info)
                 rollout_extra_info["root_id"] = sample.uid.root_id
                 rollout_extra_info["action_id"] = sample.uid.action_id
+                rollout_extra_info["observation_id"] = sample.uid.observation_id
                 update_sample_params = sample_params
+                session_id = None
                 if XTUNER_DETERMINISTIC:
                     update_sample_params.sampling_seed = self.rollout_cfg.random_seed + i
+                    session_id = build_deterministic_session_id(self.environment, sample)
 
                 if "partial_rollout_input_ids" in sample.env.rollout.extra_info:
                     input_ids_length = len(sample.data.input_ids) if sample.data.input_ids is not None else 0
@@ -125,6 +131,7 @@ async def generate(  # type: ignore[override]
                     input_ids=sample.data.input_ids,
                     sample_params=update_sample_params,
                     extra_params=extra_params,
+                    session_id=session_id,
                     extra_info=rollout_extra_info,
                 )
                 del rollout_extra_info

xtuner/v1/ray/rollout/controller.py

@@ -5,7 +5,6 @@
 import time
 from collections import OrderedDict
 from dataclasses import dataclass
-from itertools import cycle
 from typing import Any, Dict, List, Optional, Union
 from uuid import uuid4
 
@@ -52,7 +51,6 @@ def __init__(
 
         # OrderedDict: key=session_id -> value=(worker, last_used_ts)
         self._map: OrderedDict[int, tuple[Any, float]] = OrderedDict()
-        self._worker_cycler = cycle(self._workers)
         self._lock = asyncio.Lock()
         self.logger = get_logger()
 
@@ -80,7 +78,16 @@ def _evict_lru_to_capacity(self):
     def update_active_workers(self, worker_status: Dict[Any, bool]):
         self._workers = list(worker_status.items())
         self.logger.debug(f"SessionRouter update active workers: {self._workers}")
-        self._worker_cycler = cycle(self._workers)
+
+    def _get_healthy_workers(self) -> List[tuple[Any, bool]]:
+        return [worker for worker in self._workers if worker[1]]
+
+    def _select_worker_for_session(self, session_id: int) -> tuple[Any, bool]:
+        healthy_workers = self._get_healthy_workers()
+        if not healthy_workers:
+            raise RuntimeError("No healthy rollout workers available for SessionRouter.")
+        worker_idx = session_id % len(healthy_workers)
+        return healthy_workers[worker_idx]
 
     async def get_worker(self, session_id: int) -> Any:
         async with self._lock:
@@ -92,9 +99,7 @@ async def get_worker(self, session_id: int) -> Any:
                 if worker[1]:  # worker is healthy
                     return worker[0]
 
-            worker = next(self._worker_cycler)
-            while worker[1] is False:
-                worker = next(self._worker_cycler)
+            worker = self._select_worker_for_session(session_id)
             self._map[session_id] = (worker, self._now())
 
             self._evict_lru_to_capacity()

xtuner/v1/ray/rollout/lmdeploy.py

@@ -6,6 +6,7 @@
 
 import ray
 import requests
+import torch
 from ray.util.placement_group import placement_group_table
 
 from transformers import AutoTokenizer
@@ -208,6 +209,14 @@ def reset_prefix_cache(self):
         """It will implemented for LMDeploy worker in the future."""
         pass
 
+    def _decode_routed_experts(self, routed_experts: Any, meta_info: Dict[str, Any]):
+        if isinstance(routed_experts, str):
+            import base64
+
+            data = base64.b64decode(routed_experts)
+            return ray.cloudpickle.loads(data)
+        return torch.tensor(routed_experts)
+
     def _transform_rollout_config_to_server_configs(self) -> Namespace:
         """Transform the RolloutConfig into a Namespace suitable for the
         LMDeploy server.

xtuner/v1/ray/rollout/sglang.py

@@ -1,10 +1,13 @@
+import base64
 import os
 from typing import Any, Dict, List, Union
 
+import numpy as np
 import requests
+import torch
 from urllib3.exceptions import NewConnectionError
 
-from transformers import AutoTokenizer
+from transformers import AutoConfig, AutoTokenizer
 from xtuner.v1.ray.config import RolloutConfig
 from xtuner.v1.utils import XTUNER_DETERMINISTIC
 
@@ -29,6 +32,11 @@ def __init__(
         self.endpoints["generate"] = "generate"
         self.endpoints["v1/chat/completions"] = "v1/chat/completions"
         self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path, trust_remote_code=True)
+        self.model_config = AutoConfig.from_pretrained(self.config.model_path, trust_remote_code=True)
+        text_config = getattr(self.model_config, "text_config", self.model_config)
+        self.model_type = getattr(text_config, "model_type", getattr(self.model_config, "model_type", None))
+        self.routed_experts_num_hidden_layers = getattr(text_config, "num_hidden_layers", None)
+        self.routed_experts_num_experts_per_tok = getattr(text_config, "num_experts_per_tok", None)
         self.api_keys = self.config.api_key
         self.model_name = self.config.model_name
         self.enable_return_routed_experts = self.config.enable_return_routed_experts
@@ -141,6 +149,37 @@ def reset_prefix_cache(self):
         self.flush_cache()
         return self._make_request("release_memory_occupation")
 
+    def _decode_routed_experts(self, routed_experts: Any, meta_info: Dict[str, Any]):
+        if not isinstance(routed_experts, str):
+            return super()._decode_routed_experts(routed_experts, meta_info)
+
+        prompt_tokens = meta_info.get("prompt_tokens", 0)
+        completion_tokens = meta_info.get("completion_tokens", 0)
+        num_tokens = prompt_tokens + completion_tokens - 1
+        assert num_tokens > 0, (
+            f"Unexpected routed_experts token count: prompt_tokens={prompt_tokens}, completion_tokens={completion_tokens}"
+        )
+        assert self.routed_experts_num_hidden_layers is not None, (
+            "num_hidden_layers is required to decode routed_experts"
+        )
+        assert self.routed_experts_num_experts_per_tok is not None, (
+            "num_experts_per_tok is required to decode routed_experts"
+        )
+
+        routed_experts_flat = np.frombuffer(base64.b64decode(routed_experts), dtype=np.int32)
+        expected_size = num_tokens * self.routed_experts_num_hidden_layers * self.routed_experts_num_experts_per_tok
+        assert routed_experts_flat.size == expected_size, (
+            f"Unexpected routed_experts size {routed_experts_flat.size}, expected {expected_size}. "
+            f"num_tokens={num_tokens}, num_hidden_layers={self.routed_experts_num_hidden_layers}, "
+            f"num_experts_per_tok={self.routed_experts_num_experts_per_tok}"
+        )
+        routed_experts_array = routed_experts_flat.reshape(
+            num_tokens,
+            self.routed_experts_num_hidden_layers,
+            self.routed_experts_num_experts_per_tok,
+        )
+        return torch.from_numpy(routed_experts_array.copy())
+
     def _transform_rollout_config_to_server_configs(self):
         # remove the CUDA_VISIBLE_DEVICES set by ray and use base_gpu_id
         os.environ.pop("CUDA_VISIBLE_DEVICES", None)
@@ -150,55 +189,70 @@ def _transform_rollout_config_to_server_configs(self):
         sglang_config_kwargs = {
             k.replace("sglang_", ""): v for k, v in extra_config.items() if k.startswith("sglang_")
         }
-        grammar_backend = sglang_config_kwargs.get(
-            "grammar_backend", None
-        )  # for intern-s1 series models, have to set the grammar_backend to "none"
         log_level = sglang_config_kwargs.get("log_level", "error")
         log_level_http = sglang_config_kwargs.get("log_level_http", "error")
-        sglang_server_args = ServerArgs(model_path=self.config.model_path, trust_remote_code=True)
         num_gpus_per_engine = (
             self.config.expert_parallel_size
             if self.config.expert_parallel_size > 1
             else self.config.tensor_parallel_size
         )
-        sglang_server_args.host = self.host
-        sglang_server_args.port = self.server_port
-        sglang_server_args.nccl_port = self.nccl_port
-        sglang_server_args.dist_init_addr = self.dist_init_addr
-        sglang_server_args.base_gpu_id = self.rank % self.config.gpus_per_node
-        sglang_server_args.gpu_id_step = 1
-        sglang_server_args.nnodes = max(1, num_gpus_per_engine // self.config.gpus_per_node)
-        sglang_server_args.skip_server_warmup = True
-        sglang_server_args.mem_fraction_static = self.config.gpu_memory_utilization
-        # note: 非共卡模式下无需设置,共卡模式下需要offload必须设置，否则显存释放不了
-        sglang_server_args.enable_memory_saver = True
-
+        tp_size = num_gpus_per_engine if self.config.expert_parallel_size > 1 else self.config.tensor_parallel_size
+        ep_size = num_gpus_per_engine if self.config.expert_parallel_size > 1 else self.config.expert_parallel_size
+        nnodes = max(1, num_gpus_per_engine // self.config.gpus_per_node)
+        node_rank = self.rank // self.config.gpus_per_node if nnodes > 1 else 0
+        init_kwargs = dict(
+            model_path=self.config.model_path,
+            trust_remote_code=True,
+            host=self.host,
+            port=self.server_port,
+            nccl_port=self.nccl_port,
+            dist_init_addr=self.dist_init_addr,
+            base_gpu_id=self.rank % self.config.gpus_per_node,
+            gpu_id_step=1,
+            nnodes=nnodes,
+            node_rank=node_rank,
+            skip_server_warmup=True,
+            mem_fraction_static=self.config.gpu_memory_utilization,
+            enable_memory_saver=True,
+            max_running_requests=self.config.rollout_max_batch_size_per_instance,
+            log_level=log_level,
+            log_level_http=log_level_http,
+            tp_size=tp_size,
+            ep_size=ep_size,
+        )
         if self.enable_return_routed_experts:
-            sglang_server_args.enable_return_routed_experts = True
-
-        sglang_server_args.max_running_requests = self.config.rollout_max_batch_size_per_instance
-        sglang_server_args.log_level = log_level
-        sglang_server_args.log_level_http = log_level_http
+            init_kwargs["enable_return_routed_experts"] = True
         if XTUNER_DETERMINISTIC:
-            sglang_server_args.enable_deterministic_inference = True
-            sglang_server_args.rl_on_policy_target = True
-        if self.config.expert_parallel_size > 1:
-            sglang_server_args.tp_size = num_gpus_per_engine
-            sglang_server_args.ep_size = num_gpus_per_engine
-        else:
-            sglang_server_args.tp_size = self.config.tensor_parallel_size
-            sglang_server_args.ep_size = self.config.expert_parallel_size
+            init_kwargs["enable_deterministic_inference"] = True
+            init_kwargs["rl_on_policy_target"] = "fsdp"
+            init_kwargs["attention_backend"] = "fa3"
+            init_kwargs["random_seed"] = self.config.random_seed
+            # SGLang's deterministic mode does not currently force-disable every
+            # performance-oriented runtime path. For long MoE rollouts we still
+            # observed rare trajectory divergence, so explicitly turn off the
+            # scheduler/cache/graph features that can perturb execution order.
+            init_kwargs["disable_radix_cache"] = True
+            init_kwargs["disable_overlap_schedule"] = True
+            init_kwargs["disable_cuda_graph"] = True
 
-        if grammar_backend is not None:
-            sglang_server_args.grammar_backend = grammar_backend
+        # Forward supported sglang_* extra configs to ServerArgs directly.
+        server_arg_fields = getattr(ServerArgs, "__dataclass_fields__", {})
+        for key, value in sglang_config_kwargs.items():
+            if key in server_arg_fields:
+                init_kwargs[key] = value
+            else:
+                self.logger.warning(f"Ignore unknown SGLang server arg: {key}={value!r}")
+
+        # Qwen3-MoE in sglang 0.5.9 can hit native rotary + fused KV buffer incompatibility
+        # during server startup unless fused qk_norm_rope is enabled.
+        if self.model_type == "qwen3_moe" and "enable_fused_qk_norm_rope" not in sglang_config_kwargs:
+            init_kwargs["enable_fused_qk_norm_rope"] = True
+            self.logger.info("Auto enable SGLang enable_fused_qk_norm_rope for qwen3_moe.")
 
         if self.config.context_length is not None:
-            sglang_server_args.context_length = self.config.context_length
+            init_kwargs["context_length"] = self.config.context_length
 
-        if sglang_server_args.nnodes > 1:
-            sglang_server_args.node_rank = self.rank // self.config.gpus_per_node
-        else:
-            sglang_server_args.node_rank = 0
+        sglang_server_args = ServerArgs(**init_kwargs)
 
         return sglang_server_args
 

xtuner/v1/ray/rollout/vllm.py

@@ -101,6 +101,9 @@ def reset_prefix_cache(self):
         # todo
         pass
 
+    def _decode_routed_experts(self, routed_experts: Any, meta_info: Dict[str, Any]):
+        raise NotImplementedError
+
     def _transform_rollout_config_to_server_configs(self) -> Namespace:
         # use vllm FlexibleArgumentParser to parse the config
         # and return the args as the default server config

xtuner/v1/ray/rollout/worker.py

@@ -13,8 +13,8 @@
 import numpy as np
 import ray
 import requests  # type: ignore[import-untyped]
-import torch
 from packaging.version import Version
+from ray import ObjectRef
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
 from transformers import AutoTokenizer
@@ -153,6 +153,9 @@ def init(self, dist_init_addr: str = ""):
         self.launch_server()
         return (self.rank, self.server_url)
 
+    def _decode_routed_experts(self, routed_experts: Any, meta_info: Dict[str, Any]) -> Any:
+        return routed_experts
+
     def set_engine_rank_mesh_array(self, engine_rank_mesh_array: list[list[int]]):
         self.engine_rank_mesh_array = engine_rank_mesh_array
 
@@ -360,7 +363,7 @@ async def rollout_task(
     ) -> RLRolloutResponseItem:
         uid = extra_info.get("action_id", str(uuid.uuid4()))
         action_id = extra_info.get("action_id", str(uuid.uuid4()))
-        root_id = extra_info.get("action_id", str(uuid.uuid4()))
+        root_id = extra_info.get("root_id", str(uuid.uuid4()))
         response = None
         cur_retry_times = 0
 
@@ -568,28 +571,18 @@ async def _handle_non_stream_response(
                     routed_experts = response["meta_info"].pop("routed_experts")  # token[layer[expert]]
                     if routed_experts is not None and not exist_history_routed_experts:
                         # 不存在历史专家，先把当前专家存起来
-                        if isinstance(routed_experts, str):
-                            import base64
-
-                            data = base64.b64decode(routed_experts)
-                            routed_experts = ray.cloudpickle.loads(data)
-                            del data
-                        else:
-                            routed_experts = torch.tensor(routed_experts)  # n,layer,expert
+                        routed_experts = self._decode_routed_experts(routed_experts, response["meta_info"])
+                        if not isinstance(routed_experts, ObjectRef):
                             routed_experts = ray.put(routed_experts)
                         extra_info["routed_experts"] = routed_experts
                     elif routed_experts is not None and exist_history_routed_experts:
                         # 存在历史专家，则不进行put 操作，直接进行concat
-                        if isinstance(routed_experts, str):
-                            import base64
-
-                            data = base64.b64decode(routed_experts)
-                            routed_experts = ray.cloudpickle.loads(data)
+                        routed_experts = self._decode_routed_experts(routed_experts, response["meta_info"])
+                        if isinstance(routed_experts, ObjectRef):
                             cur_routed_experts = await routed_experts  # n,layer,expert
                             ray.internal.free(routed_experts, local_only=False)
                             del data
                         else:
-                            routed_experts = torch.tensor(routed_experts)  # n,layer,expert
                             cur_routed_experts = routed_experts
 
                         history_routed_experts = await input_extra_info["routed_experts"]  # n, layer, expert

xtuner/v1/ray/utils.py

@@ -1,4 +1,5 @@
 import asyncio
+import hashlib
 import importlib
 import socket
 from asyncio import AbstractEventLoop, Task
@@ -71,7 +72,9 @@ def _is_port_available(check_socket: socket.socket, port: int) -> bool:
 
 
 @ray.remote
-def find_master_addr_and_port(nums=1, start_port=None, end_port=None):
+def find_master_addr_and_port(
+    nums: int = 1, start_port: Optional[int] = None, end_port: Optional[int] = None
+) -> tuple[str, int] | tuple[str, list[int]]:
     """Finds an available master address and a specified number of ports.
 
     This remote function gets the node's IP address and binds to one or more
@@ -219,3 +222,16 @@ def free_object_refs(refs: List[ObjectRef]) -> None:
         ray._private.internal_api.free(valid_refs, local_only=False)
     except Exception:
         ray.internal.free(valid_refs, local_only=False)
+def deterministic_item_sort_key(sample) -> tuple[int, int, int, int]:
+    return (
+        sample.uid.root_id,
+        sample.uid.action_id,
+        sample.uid.observation_id,
+        sample.uid.version,
+    )
+
+
+def build_deterministic_session_id(environment: str, sample) -> int:
+    session_key = f"{environment}|{sample.uid.root_id}|{sample.uid.action_id}|{sample.uid.observation_id}"
+    session_id = int.from_bytes(hashlib.sha256(session_key.encode("utf-8")).digest()[:8], "big")
+    return session_id or 1