From 06bdcd7ff56714920678a091298c5c5e9ed19a6f Mon Sep 17 00:00:00 2001 From: wyeth Date: Mon, 6 Jul 2026 17:06:45 +0800 Subject: [PATCH] Support SGLang routed experts via Ray --- patch/README.md | 80 +++++++++++++++++ patch/sglang-0.5.10-routed-experts-ray.patch | 94 ++++++++++++++++++++ xtuner/v1/rl/rollout/sglang.py | 13 +++ 3 files changed, 187 insertions(+) create mode 100644 patch/README.md create mode 100644 patch/sglang-0.5.10-routed-experts-ray.patch diff --git a/patch/README.md b/patch/README.md new file mode 100644 index 000000000..61833c23c --- /dev/null +++ b/patch/README.md @@ -0,0 +1,80 @@ +# SGLang patches + +## sglang-0.5.10-routed-experts-ray.patch + +This patch is generated against the clean `sglang==0.5.10` wheel. + +It changes `sglang/srt/managers/tokenizer_manager.py` so `routed_experts` +is stored in a Ray named actor/object store and the HTTP response only returns +the shared-store key. XTuner can then fetch the real ndarray through Ray. + +Apply from the XTuner repository checkout: + +```bash +SITE_PACKAGES=$(python -c "import pathlib, sglang; print(pathlib.Path(sglang.__file__).resolve().parents[1])") +PATCH_FILE=$(python -c "import os, sys; print(os.path.relpath(sys.argv[1], sys.argv[2]))" patch/sglang-0.5.10-routed-experts-ray.patch "$SITE_PACKAGES") +cd "$SITE_PACKAGES" +git apply --check -p2 --include='sglang/**' --exclude='*' "$PATCH_FILE" +git apply -p2 --include='sglang/**' --exclude='*' "$PATCH_FILE" +``` + +The expected SGLang version is: + +```bash +python -c "import sglang; print(sglang.__version__)" +``` + +```text +0.5.10 +``` + +## Use Ray to transfer routed_experts + +This patch only changes the SGLang side. XTuner also needs the matching +`xtuner/v1/rl/rollout/sglang.py` logic that reads the returned key from the +SGLang `shared_store` actor with `namespace="sglang"` and then `ray.put`s the +real ndarray for trainer workers. + +Runtime requirements: + +- Use a MoE model that can return routed experts, for example Qwen3-A3B MoE. +- Use the SGLang rollout backend. +- Enable routed experts return: + +```bash +export ENABLE_RETURN_ROUTED_EXPERTS=1 +``` + +Then run XTuner normally, for example: + +```bash +bash examples/v1/scripts/run_rl.sh \ + examples/v1/config/rl_grpo_gsm8k_async.py \ + sglang \ + "$MODEL_PATH" "$DATA_PATH" "$EVAL_DATA_PATH" +``` + +Expected transfer path: + +```text +SGLang routed_experts tensor + -> SGLang _RoutedExpertsSharedStore.put() + -> ray.put(data) stores ndarray in Ray object store + -> HTTP meta_info["routed_experts"] returns shared-store key + -> XTuner SGLang rollout worker gets shared_store actor in namespace "sglang" + -> rollout worker actor.get(key) fetches real ndarray + -> rollout worker ray.put(np.asarray(...)) + -> trainer worker ray.get(...) consumes rollout_routed_expert +``` + +Validation signals in logs: + +- SGLang rollout config contains `"enable_return_routed_experts": true`. +- Ray actors contain `shared_store(_RoutedExpertsSharedStore)`. +- Training logs contain `rollout_routed_expert = torch.as_tensor(...)`. + +To disable this path, unset the flag or set: + +```bash +export ENABLE_RETURN_ROUTED_EXPERTS=0 +``` diff --git a/patch/sglang-0.5.10-routed-experts-ray.patch b/patch/sglang-0.5.10-routed-experts-ray.patch new file mode 100644 index 000000000..8d4173a1c --- /dev/null +++ b/patch/sglang-0.5.10-routed-experts-ray.patch @@ -0,0 +1,94 @@ +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 8142432..8dda24e 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -122,6 +122,78 @@ _REQUEST_STATE_WAIT_TIMEOUT = envs.SGLANG_REQUEST_STATE_WAIT_TIMEOUT.get() + + logger = logging.getLogger(__name__) + ++_ROUTED_EXPERTS_SHARED_STORE = None ++_ROUTED_EXPERTS_SHARED_STORE_NAME = "shared_store" ++_ROUTED_EXPERTS_SHARED_STORE_NAMESPACE = "sglang" ++ ++ ++class _RoutedExpertsSharedStore: ++ def __init__(self): ++ self._data = {} ++ ++ def put(self, data): ++ import ray ++ ++ ref = ray.put(data) ++ key = ref.hex() ++ self._data[key] = ref ++ return key ++ ++ def get(self, key): ++ import ray ++ ++ ref = self._data.pop(key) ++ return ray.get(ref) ++ ++ def clear(self): ++ import ray ++ ++ all_data = list(self._data.values()) ++ if len(all_data) > 0: ++ ray.internal.free(all_data, local_only=False) ++ self._data.clear() ++ ++ ++def _lazy_get_routed_experts_shared_store(): ++ global _ROUTED_EXPERTS_SHARED_STORE ++ if _ROUTED_EXPERTS_SHARED_STORE is not None: ++ return _ROUTED_EXPERTS_SHARED_STORE ++ ++ import ray ++ ++ if not ray.is_initialized(): ++ ray.init(address="auto", ignore_reinit_error=True) ++ ++ try: ++ _ROUTED_EXPERTS_SHARED_STORE = ray.get_actor( ++ _ROUTED_EXPERTS_SHARED_STORE_NAME, ++ namespace=_ROUTED_EXPERTS_SHARED_STORE_NAMESPACE, ++ ) ++ except ValueError: ++ try: ++ _ROUTED_EXPERTS_SHARED_STORE = ray.remote(num_cpus=0)(_RoutedExpertsSharedStore).options( ++ name=_ROUTED_EXPERTS_SHARED_STORE_NAME, ++ namespace=_ROUTED_EXPERTS_SHARED_STORE_NAMESPACE, ++ lifetime="detached", ++ ).remote() ++ except ray.exceptions.ActorAlreadyExistsError: ++ _ROUTED_EXPERTS_SHARED_STORE = ray.get_actor( ++ _ROUTED_EXPERTS_SHARED_STORE_NAME, ++ namespace=_ROUTED_EXPERTS_SHARED_STORE_NAMESPACE, ++ ) ++ return _ROUTED_EXPERTS_SHARED_STORE ++ ++ ++def _encode_routed_experts(routed_experts_tensor): ++ try: ++ import ray ++ ++ store = _lazy_get_routed_experts_shared_store() ++ return ray.get(store.put.remote(routed_experts_tensor.numpy())) ++ except Exception: ++ logger.exception("Failed to put routed_experts into Ray shared store; falling back to base64.") ++ return pybase64.b64encode(routed_experts_tensor.numpy().tobytes()).decode("utf-8") ++ + + @dataclasses.dataclass + class ReqState: +@@ -1593,9 +1665,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): + if getattr(recv_obj, "routed_experts", None): + routed_experts_tensor = recv_obj.routed_experts[i] + if routed_experts_tensor is not None: +- meta_info["routed_experts"] = pybase64.b64encode( +- routed_experts_tensor.numpy().tobytes() +- ).decode("utf-8") ++ meta_info["routed_experts"] = _encode_routed_experts(routed_experts_tensor) + if getattr(recv_obj, "customized_info", None): + for k, v in recv_obj.customized_info.items(): + meta_info[k] = v[i] diff --git a/xtuner/v1/rl/rollout/sglang.py b/xtuner/v1/rl/rollout/sglang.py index 6825fdcfb..050646f05 100644 --- a/xtuner/v1/rl/rollout/sglang.py +++ b/xtuner/v1/rl/rollout/sglang.py @@ -14,6 +14,9 @@ from .rollout_topology import RolloutEngine, RolloutServerProcess, RolloutTopology from .worker import RolloutConfig, RolloutWorker +SHARED_STORE = "shared_store" +SHARED_STORE_NAMESPACE = "sglang" + class SGLangWorker(RolloutWorker): def __init__( @@ -42,6 +45,7 @@ def __init__( 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 + self.sglang_actor = None @classmethod def build_rollout_topology( @@ -338,6 +342,15 @@ def reset_prefix_cache(self): async def _decode_routed_experts(self, routed_experts: Any): if isinstance(routed_experts, str): + try: + if self.sglang_actor is None: + self.sglang_actor = ray.get_actor(SHARED_STORE, namespace=SHARED_STORE_NAMESPACE) + routed_experts_data = await self.sglang_actor.get.remote(routed_experts) + if hasattr(routed_experts_data, "detach"): + routed_experts_data = routed_experts_data.detach().cpu().numpy() + return ray.put(np.asarray(routed_experts_data)) + except Exception: + self.logger.debug("Failed to resolve SGLang routed_experts from Ray shared store; trying base64.") routed_experts_flat = np.frombuffer(base64.b64decode(routed_experts), dtype=np.int32) routed_experts_array = routed_experts_flat.reshape( -1,