From 2cbd605e2db1630f6a3007d65cb4975a5e624a4b Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Wed, 1 Jul 2026 08:31:13 +0000 Subject: [PATCH 01/11] feat: add Mooncake RDMA transport for rollout data Enable rollout data transfer via Mooncake structured object store as an alternative to Ray object store, activated by --rollout-data-transport=mooncake. Co-Authored-By: Claude Opus 4.6 --- slime/ray/rollout.py | 6 ++- slime/utils/arguments.py | 2 +- slime/utils/data.py | 8 +++- slime/utils/data_transfer.py | 84 ++++++++++++++++++++++++++++++++++++ train.py | 3 ++ train_async.py | 3 ++ 6 files changed, 103 insertions(+), 3 deletions(-) create mode 100644 slime/utils/data_transfer.py diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index abd8e35d70..cac5b98e96 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -929,7 +929,11 @@ def _split_train_data_by_dp(self, data): rollout_data["micro_batch_indices"] = micro_batch_indices[r] _tensorize_rollout_data_for_training(rollout_data) transport = getattr(self.args, "rollout_data_transport", "object-store") - if transport == "nixl": + if transport == "mooncake": + from slime.utils.data_transfer import put_mooncake_rollout_data + + rollout_data_refs.append(put_mooncake_rollout_data(self.args, rollout_data, partition=f"dp{r}")) + elif transport == "nixl": rollout_data_refs.append(Box(ray.put(rollout_data, _tensor_transport="nixl"))) elif transport == "object-store": rollout_data_refs.append(Box(ray.put(rollout_data))) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 96be544576..6f39db84c8 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -557,7 +557,7 @@ def add_rollout_arguments(parser): parser.add_argument( "--rollout-data-transport", type=str, - choices=["object-store", "nixl"], + choices=["object-store", "nixl", "mooncake"], default="object-store", help=( "Transport for rollout data refs sent from rollout manager to trainer. Large rollout " diff --git a/slime/utils/data.py b/slime/utils/data.py index adac4b95a6..ae245170f5 100644 --- a/slime/utils/data.py +++ b/slime/utils/data.py @@ -304,7 +304,13 @@ def __len__(self): def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): assert len(rollout_data_ref) == dp_size - rollout_data = ray.get(rollout_data_ref[dp_rank].inner) + ref = rollout_data_ref[dp_rank] + if getattr(args, "rollout_data_transport", "object-store") == "mooncake": + from slime.utils.data_transfer import get_mooncake_rollout_data + + rollout_data = get_mooncake_rollout_data(args, ref) + else: + rollout_data = ray.get(ref.inner) # Keep `partition` in rollout_data: each local sample's position in the # flattened rollout batch (== its index in the rollout debug dump's diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py new file mode 100644 index 0000000000..8321fc5173 --- /dev/null +++ b/slime/utils/data_transfer.py @@ -0,0 +1,84 @@ +"""Thin Mooncake rollout data transport: put/get/cleanup.""" + +import os +from functools import cache +from typing import Any + +from slime.utils.misc import Box + + +def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) -> Box: + from mooncake.structured_object_store import export_dataproto_ref + + ref = _mooncake_transfer(args).put_legacy_dict( + data, + namespace="slime", + partition=partition, + stage="rollout", + field_schemas=_rollout_field_schemas(), + ) + return Box(export_dataproto_ref(ref)) + + +@cache +def _rollout_field_schemas() -> dict: + from mooncake.structured_object_store import FieldSchema + from slime.ray.rollout import _ROLLOUT_DATA_TENSOR_DTYPES + + ragged = FieldSchema(codec="typed_ragged", nullable=False) + return {k: ragged for k in _ROLLOUT_DATA_TENSOR_DTYPES} + + +def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: + from mooncake.structured_object_store import import_dataproto_ref + + return _mooncake_transfer(args).get_legacy_dict(import_dataproto_ref(ref.inner)) + + +def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: + from mooncake.structured_object_store import import_dataproto_ref + + _mooncake_transfer(args).remove_legacy_dict(import_dataproto_ref(ref.inner)) + + +def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box] | None) -> None: + if getattr(args, "rollout_data_transport", "object-store") != "mooncake" or refs is None: + return + for ref in refs: + cleanup_mooncake_rollout_data(args, ref) + + +@cache +def _mooncake_transfer(args: Any): + from mooncake.store import MooncakeDistributedStore + from mooncake.structured_object_store import MooncakeBundleTransfer + + store = MooncakeDistributedStore() + mc_kwargs = getattr(args, "mooncake_store_init_kwargs", None) or {} + ret = store.setup( + { + "local_hostname": mc_kwargs.get("local_hostname") or _local_hostname(), + "metadata_server": mc_kwargs.get("metadata_server") or os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"), + "global_segment_size": int( + mc_kwargs.get("global_segment_size") or os.getenv("MC_SEGMENT_SIZE", str(8 * 1024**3)) + ), + "local_buffer_size": int( + mc_kwargs.get("local_buffer_size") or os.getenv("MC_BUFFER_SIZE", str(32 * 1024**3)) + ), + "protocol": mc_kwargs.get("protocol") or os.getenv("MC_PROTOCOL", "rdma"), + "rdma_devices": mc_kwargs.get("rdma_devices") or os.getenv("MC_DEVICE", ""), + "master_server_addr": mc_kwargs.get("master_server_addr") or os.getenv("MC_MASTER_SERVER", ""), + } + ) + if ret: + raise RuntimeError(f"Mooncake store setup failed: {ret}") + return MooncakeBundleTransfer(store, key_prefix="slime-rollout") + + +def _local_hostname() -> str: + value = os.getenv("MC_LOCAL_HOSTNAME") or os.getenv("LOCAL_HOSTNAME") + if value: + return value + import ray + + return ray.util.get_node_ip_address() diff --git a/train.py b/train.py index 9cb2968866..33b3e4d702 100644 --- a/train.py +++ b/train.py @@ -2,6 +2,7 @@ from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from slime.utils.arguments import parse_args +from slime.utils.data_transfer import cleanup_mooncake_rollout_refs from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking from slime.utils.misc import should_run_periodic_action @@ -68,6 +69,8 @@ def offload_train(actor_trains_this_step): else: ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) + cleanup_mooncake_rollout_refs(args, rollout_data_ref) + if release_train or should_run_periodic_action( rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout ): diff --git a/train_async.py b/train_async.py index 9141b612fe..7feeecf4b4 100644 --- a/train_async.py +++ b/train_async.py @@ -2,6 +2,7 @@ from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from slime.utils.arguments import parse_args +from slime.utils.data_transfer import cleanup_mooncake_rollout_refs from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking from slime.utils.misc import should_run_periodic_action @@ -52,6 +53,8 @@ def train(args): else: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) + cleanup_mooncake_rollout_refs(args, rollout_data_curr_ref) + if release_train or should_run_periodic_action( rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout ): From 164d7fb815b6e119cd5f58a887dd0c451a290b09 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 3 Jul 2026 02:21:55 +0000 Subject: [PATCH 02/11] fix: mooncake import try/except, parse-time check, and pool buffer release Co-Authored-By: Claude Opus 4.6 --- slime/utils/arguments.py | 5 +++++ slime/utils/data_transfer.py | 40 ++++++++++++++++++++++++++---------- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 6f39db84c8..d60d5a4bec 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -2065,3 +2065,8 @@ def slime_validate_args(args): "--update-weight-mode=delta requires --update-weight-local-checkpoint-dir " "(a rollout-host-local NVMe directory)." ) + + if getattr(args, "rollout_data_transport", "object-store") == "mooncake": + from slime.utils.data_transfer import check_mooncake_available + + check_mooncake_available() diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index 8321fc5173..a20211c4ee 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -1,15 +1,38 @@ """Thin Mooncake rollout data transport: put/get/cleanup.""" +import logging import os from functools import cache from typing import Any from slime.utils.misc import Box +try: + from mooncake.structured_object_store import ( + FieldSchema, + MooncakeBundleTransfer, + export_dataproto_ref, + import_dataproto_ref, + ) + from mooncake.store import MooncakeDistributedStore -def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) -> Box: - from mooncake.structured_object_store import export_dataproto_ref + _MOONCAKE_AVAILABLE = True +except ImportError: + _MOONCAKE_AVAILABLE = False + +logger = logging.getLogger(__name__) + + +def check_mooncake_available() -> None: + """Call during argument parsing to fail fast if mooncake is not installed.""" + if not _MOONCAKE_AVAILABLE: + raise ImportError( + "rollout-data-transport='mooncake' requires the mooncake package. " + "Install it with: pip install mooncake" + ) + +def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) -> Box: ref = _mooncake_transfer(args).put_legacy_dict( data, namespace="slime", @@ -22,7 +45,6 @@ def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) - @cache def _rollout_field_schemas() -> dict: - from mooncake.structured_object_store import FieldSchema from slime.ray.rollout import _ROLLOUT_DATA_TENSOR_DTYPES ragged = FieldSchema(codec="typed_ragged", nullable=False) @@ -30,14 +52,13 @@ def _rollout_field_schemas() -> dict: def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: - from mooncake.structured_object_store import import_dataproto_ref - - return _mooncake_transfer(args).get_legacy_dict(import_dataproto_ref(ref.inner)) + transfer = _mooncake_transfer(args) + result = transfer.get_legacy_dict(import_dataproto_ref(ref.inner)) + transfer.release_result(result) + return result def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: - from mooncake.structured_object_store import import_dataproto_ref - _mooncake_transfer(args).remove_legacy_dict(import_dataproto_ref(ref.inner)) @@ -50,9 +71,6 @@ def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box] | None) -> None: @cache def _mooncake_transfer(args: Any): - from mooncake.store import MooncakeDistributedStore - from mooncake.structured_object_store import MooncakeBundleTransfer - store = MooncakeDistributedStore() mc_kwargs = getattr(args, "mooncake_store_init_kwargs", None) or {} ret = store.setup( From 6dd5d5b8664871ff4dd4c1dc57821837414653b6 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 3 Jul 2026 02:50:49 +0000 Subject: [PATCH 03/11] refactor: use export_ref/import_ref aliases instead of dataproto naming Switch imports from export_dataproto_ref/import_dataproto_ref to export_ref/import_ref to align with mooncake's parallel API refactoring. Co-Authored-By: Claude Opus 4.6 --- slime/utils/data_transfer.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index a20211c4ee..d5f66cfde6 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -11,8 +11,8 @@ from mooncake.structured_object_store import ( FieldSchema, MooncakeBundleTransfer, - export_dataproto_ref, - import_dataproto_ref, + export_ref, + import_ref, ) from mooncake.store import MooncakeDistributedStore @@ -40,7 +40,7 @@ def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) - stage="rollout", field_schemas=_rollout_field_schemas(), ) - return Box(export_dataproto_ref(ref)) + return Box(export_ref(ref)) @cache @@ -53,13 +53,13 @@ def _rollout_field_schemas() -> dict: def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: transfer = _mooncake_transfer(args) - result = transfer.get_legacy_dict(import_dataproto_ref(ref.inner)) + result = transfer.get_legacy_dict(import_ref(ref.inner)) transfer.release_result(result) return result def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: - _mooncake_transfer(args).remove_legacy_dict(import_dataproto_ref(ref.inner)) + _mooncake_transfer(args).remove_legacy_dict(import_ref(ref.inner)) def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box] | None) -> None: From d1c1b4ae0f39e8c2be657761e0f67c31ca7043ad Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 3 Jul 2026 07:36:12 +0000 Subject: [PATCH 04/11] fix: defer pool buffer release and move transport check to call site - Move release_result from get_mooncake_rollout_data to after training completes in actor.py, preventing premature buffer corruption. - Move transport type check from cleanup_mooncake_rollout_refs to call sites in train.py and train_async.py per reviewer feedback. Co-Authored-By: Claude Opus 4.6 --- slime/backends/megatron_utils/actor.py | 5 +++++ slime/utils/data_transfer.py | 16 +++++++++------- train.py | 3 ++- train_async.py | 3 ++- 4 files changed, 18 insertions(+), 9 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index ea2601a9d4..d181a83963 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -387,6 +387,11 @@ def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): self.train_actor(rollout_id, rollout_data, external_data=external_data) result = None + if getattr(self.args, "rollout_data_transport", "object-store") == "mooncake": + from slime.utils.data_transfer import release_mooncake_rollout_data + + release_mooncake_rollout_data(self.args, rollout_data) + if self.args.offload_train: del rollout_data self.sleep() diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index d5f66cfde6..2d02643a49 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -52,19 +52,21 @@ def _rollout_field_schemas() -> dict: def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: - transfer = _mooncake_transfer(args) - result = transfer.get_legacy_dict(import_ref(ref.inner)) - transfer.release_result(result) - return result + return _mooncake_transfer(args).get_legacy_dict(import_ref(ref.inner)) + + +def release_mooncake_rollout_data(args: Any, data: dict[str, Any]) -> None: + """Release pool-backed buffers after training has fully consumed the data.""" + from mooncake.structured_object_store import MooncakeBundleTransfer + + MooncakeBundleTransfer.release_result(data) def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: _mooncake_transfer(args).remove_legacy_dict(import_ref(ref.inner)) -def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box] | None) -> None: - if getattr(args, "rollout_data_transport", "object-store") != "mooncake" or refs is None: - return +def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box]) -> None: for ref in refs: cleanup_mooncake_rollout_data(args, ref) diff --git a/train.py b/train.py index 33b3e4d702..258f19bdfc 100644 --- a/train.py +++ b/train.py @@ -69,7 +69,8 @@ def offload_train(actor_trains_this_step): else: ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) - cleanup_mooncake_rollout_refs(args, rollout_data_ref) + if getattr(args, "rollout_data_transport", "object-store") == "mooncake": + cleanup_mooncake_rollout_refs(args, rollout_data_ref) if release_train or should_run_periodic_action( rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout diff --git a/train_async.py b/train_async.py index 7feeecf4b4..d511d1e67e 100644 --- a/train_async.py +++ b/train_async.py @@ -53,7 +53,8 @@ def train(args): else: ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref)) - cleanup_mooncake_rollout_refs(args, rollout_data_curr_ref) + if getattr(args, "rollout_data_transport", "object-store") == "mooncake": + cleanup_mooncake_rollout_refs(args, rollout_data_curr_ref) if release_train or should_run_periodic_action( rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout From 1ef4122dcda7fabe0fa93d86520470ffb1c6d0e4 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Tue, 7 Jul 2026 10:41:04 +0800 Subject: [PATCH 05/11] fix: refine Mooncake rollout transfer config Use static rollout field schemas and explicit producer/consumer segment contribution so Mooncake rollout transfer avoids runtime schema inference while keeping deployment-specific store settings externalized. Co-Authored-By: Claude Opus 4.6 --- slime/utils/data_transfer.py | 113 ++++++++++++++++++++++++----------- 1 file changed, 79 insertions(+), 34 deletions(-) diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index 2d02643a49..43064cbb1f 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -8,12 +8,7 @@ from slime.utils.misc import Box try: - from mooncake.structured_object_store import ( - FieldSchema, - MooncakeBundleTransfer, - export_ref, - import_ref, - ) + from mooncake.structured_object_store import FieldSchema, MooncakeBundleTransfer, export_ref, import_ref from mooncake.store import MooncakeDistributedStore _MOONCAKE_AVAILABLE = True @@ -22,6 +17,37 @@ logger = logging.getLogger(__name__) +_ROLLOUT_FIELD_SCHEMA_SPECS = { + # rollout.py tensorizes these row-aligned fields before transport. + "tokens": ("ragged_tensor", None), + "loss_masks": ("ragged_tensor", None), + "rollout_log_probs": ("ragged_tensor", None), + "rollout_top_p_token_ids": ("ragged_tensor", None), + "rollout_top_p_token_offsets": ("ragged_tensor", None), + "teacher_log_probs": ("ragged_tensor", None), + "rollout_routed_experts": ("ragged_tensor", None), + # Row-aligned scalar fields. + "partition": ("ndarray", "int64"), + "response_lengths": ("ndarray", "int64"), + "rewards": ("ndarray", "float32"), + "truncated": ("ndarray", "int64"), + "round_number": ("ndarray", "int64"), + "sample_indices": ("ndarray", "int64"), + "rollout_ids": ("ndarray", "int64"), + "rollout_mask_sums": ("ndarray", "int64"), + # Optional row-aligned text fields. + "prompt": ("utf8_ragged", None), +} + +_ROLLOUT_FIELD_SCHEMAS = ( + { + key: FieldSchema(codec=codec, nullable=False, metadata=({"dtype": dtype} if dtype else {})) + for key, (codec, dtype) in _ROLLOUT_FIELD_SCHEMA_SPECS.items() + } + if _MOONCAKE_AVAILABLE + else {} +) + def check_mooncake_available() -> None: """Call during argument parsing to fail fast if mooncake is not installed.""" @@ -33,26 +59,22 @@ def check_mooncake_available() -> None: def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) -> Box: - ref = _mooncake_transfer(args).put_legacy_dict( + ref = _mooncake_transfer(args, contribute_segment=True).put_legacy_dict( data, namespace="slime", partition=partition, stage="rollout", - field_schemas=_rollout_field_schemas(), + field_schemas=_rollout_field_schemas_for_data(data), ) return Box(export_ref(ref)) -@cache -def _rollout_field_schemas() -> dict: - from slime.ray.rollout import _ROLLOUT_DATA_TENSOR_DTYPES - - ragged = FieldSchema(codec="typed_ragged", nullable=False) - return {k: ragged for k in _ROLLOUT_DATA_TENSOR_DTYPES} +def _rollout_field_schemas_for_data(data: dict[str, Any]) -> dict: + return {key: schema for key, schema in _ROLLOUT_FIELD_SCHEMAS.items() if key in data} def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: - return _mooncake_transfer(args).get_legacy_dict(import_ref(ref.inner)) + return _mooncake_transfer(args, contribute_segment=False).get_legacy_dict(import_ref(ref.inner)) def release_mooncake_rollout_data(args: Any, data: dict[str, Any]) -> None: @@ -63,7 +85,7 @@ def release_mooncake_rollout_data(args: Any, data: dict[str, Any]) -> None: def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: - _mooncake_transfer(args).remove_legacy_dict(import_ref(ref.inner)) + _mooncake_transfer(args, contribute_segment=False).remove_legacy_dict(import_ref(ref.inner)) def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box]) -> None: @@ -71,32 +93,55 @@ def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box]) -> None: cleanup_mooncake_rollout_data(args, ref) +def _mooncake_transfer(args: Any, contribute_segment: bool): + config = _mooncake_store_config(args, contribute_segment=contribute_segment) + return _cached_mooncake_transfer(tuple(sorted(config.items()))) + + @cache -def _mooncake_transfer(args: Any): +def _cached_mooncake_transfer(config_items: tuple[tuple[str, Any], ...]): store = MooncakeDistributedStore() - mc_kwargs = getattr(args, "mooncake_store_init_kwargs", None) or {} - ret = store.setup( - { - "local_hostname": mc_kwargs.get("local_hostname") or _local_hostname(), - "metadata_server": mc_kwargs.get("metadata_server") or os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"), - "global_segment_size": int( - mc_kwargs.get("global_segment_size") or os.getenv("MC_SEGMENT_SIZE", str(8 * 1024**3)) - ), - "local_buffer_size": int( - mc_kwargs.get("local_buffer_size") or os.getenv("MC_BUFFER_SIZE", str(32 * 1024**3)) - ), - "protocol": mc_kwargs.get("protocol") or os.getenv("MC_PROTOCOL", "rdma"), - "rdma_devices": mc_kwargs.get("rdma_devices") or os.getenv("MC_DEVICE", ""), - "master_server_addr": mc_kwargs.get("master_server_addr") or os.getenv("MC_MASTER_SERVER", ""), - } - ) + ret = store.setup(dict(config_items)) if ret: raise RuntimeError(f"Mooncake store setup failed: {ret}") return MooncakeBundleTransfer(store, key_prefix="slime-rollout") +def _mooncake_store_config(args: Any, contribute_segment: bool) -> dict[str, Any]: + mc_kwargs = getattr(args, "mooncake_store_init_kwargs", None) or {} + global_segment_size = _parse_size( + mc_kwargs.get("global_segment_size") or os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", str(8 * 1024**3)) + ) + if not contribute_segment: + global_segment_size = 0 + return { + "local_hostname": str(mc_kwargs.get("local_hostname") or _local_hostname()), + "metadata_server": str( + mc_kwargs.get("metadata_server") or os.getenv("MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE") + ), + "global_segment_size": global_segment_size, + "local_buffer_size": _parse_size( + mc_kwargs.get("local_buffer_size") or os.getenv("MOONCAKE_LOCAL_BUFFER_SIZE", str(32 * 1024**3)) + ), + "protocol": str(mc_kwargs.get("protocol") or os.getenv("MOONCAKE_PROTOCOL", "rdma")), + "rdma_devices": str(mc_kwargs.get("device_name") or os.getenv("MOONCAKE_DEVICE", "")), + "master_server_addr": str(mc_kwargs.get("master_server_address") or os.getenv("MOONCAKE_MASTER", "")), + } + + +def _parse_size(value: Any) -> int: + if isinstance(value, int): + return value + text = str(value).strip().lower() + units = {"kb": 1024, "mb": 1024**2, "gb": 1024**3, "k": 1024, "m": 1024**2, "g": 1024**3} + for suffix, multiplier in units.items(): + if text.endswith(suffix): + return int(float(text[: -len(suffix)]) * multiplier) + return int(text) + + def _local_hostname() -> str: - value = os.getenv("MC_LOCAL_HOSTNAME") or os.getenv("LOCAL_HOSTNAME") + value = os.getenv("MOONCAKE_LOCAL_HOSTNAME") or os.getenv("LOCAL_HOSTNAME") if value: return value import ray From 3aa47505d85a1ff1cf2d39d103732c3a97d5ae67 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Wed, 8 Jul 2026 12:55:37 +0800 Subject: [PATCH 06/11] fix: declare Mooncake rollout schema sections Explicitly assign rollout fields to DataProto sections so metadata lists do not affect row-count detection and tensor fields follow the batch path. Co-Authored-By: Claude Opus 4.6 --- slime/utils/data_transfer.py | 46 ++++++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index 43064cbb1f..baa3e77f7c 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -19,30 +19,40 @@ _ROLLOUT_FIELD_SCHEMA_SPECS = { # rollout.py tensorizes these row-aligned fields before transport. - "tokens": ("ragged_tensor", None), - "loss_masks": ("ragged_tensor", None), - "rollout_log_probs": ("ragged_tensor", None), - "rollout_top_p_token_ids": ("ragged_tensor", None), - "rollout_top_p_token_offsets": ("ragged_tensor", None), - "teacher_log_probs": ("ragged_tensor", None), - "rollout_routed_experts": ("ragged_tensor", None), + "tokens": ("ragged_tensor", None, "non_tensor_batch"), + "loss_masks": ("ragged_tensor", None, "non_tensor_batch"), + "rollout_log_probs": ("ragged_tensor", None, "non_tensor_batch"), + "rollout_top_p_token_ids": ("ragged_tensor", None, "non_tensor_batch"), + "rollout_top_p_token_offsets": ("ragged_tensor", None, "non_tensor_batch"), + "teacher_log_probs": ("ragged_tensor", None, "non_tensor_batch"), + "rollout_routed_experts": ("ragged_tensor", None, "non_tensor_batch"), # Row-aligned scalar fields. - "partition": ("ndarray", "int64"), - "response_lengths": ("ndarray", "int64"), - "rewards": ("ndarray", "float32"), - "truncated": ("ndarray", "int64"), - "round_number": ("ndarray", "int64"), - "sample_indices": ("ndarray", "int64"), - "rollout_ids": ("ndarray", "int64"), - "rollout_mask_sums": ("ndarray", "int64"), + "partition": ("ndarray", "int64", "non_tensor_batch"), + "response_lengths": ("ndarray", "int64", "non_tensor_batch"), + "rewards": ("ndarray", "float32", "non_tensor_batch"), + "truncated": ("ndarray", "int64", "non_tensor_batch"), + "round_number": ("ndarray", "int64", "non_tensor_batch"), + "sample_indices": ("ndarray", "int64", "non_tensor_batch"), + "rollout_ids": ("ndarray", "int64", "non_tensor_batch"), + "rollout_mask_sums": ("tensor", "float32", "batch"), # Optional row-aligned text fields. - "prompt": ("utf8_ragged", None), + "prompt": ("utf8_ragged", None, "non_tensor_batch"), + # Metadata fields carried with each DP partition. + "raw_reward": ("auto", None, "meta_info"), + "total_lengths": ("auto", None, "meta_info"), + "global_batch_sizes": ("auto", None, "meta_info"), + "num_microbatches": ("auto", None, "meta_info"), + "micro_batch_indices": ("auto", None, "meta_info"), } _ROLLOUT_FIELD_SCHEMAS = ( { - key: FieldSchema(codec=codec, nullable=False, metadata=({"dtype": dtype} if dtype else {})) - for key, (codec, dtype) in _ROLLOUT_FIELD_SCHEMA_SPECS.items() + key: FieldSchema( + codec=codec, + nullable=False, + metadata={"section": section, **({"dtype": dtype} if dtype else {})}, + ) + for key, (codec, dtype, section) in _ROLLOUT_FIELD_SCHEMA_SPECS.items() } if _MOONCAKE_AVAILABLE else {} From 2065324824ac3f91b8968b649f23a86f16c907c0 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Mon, 20 Jul 2026 07:59:18 +0000 Subject: [PATCH 07/11] fix: align Mooncake rollout transfer API --- slime/utils/data_transfer.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index baa3e77f7c..16bb5024cb 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -69,8 +69,9 @@ def check_mooncake_available() -> None: def put_mooncake_rollout_data(args: Any, data: dict[str, Any], partition: str) -> Box: - ref = _mooncake_transfer(args, contribute_segment=True).put_legacy_dict( + ref = _mooncake_transfer(args, contribute_segment=True).put( data, + type="dict", namespace="slime", partition=partition, stage="rollout", @@ -84,7 +85,7 @@ def _rollout_field_schemas_for_data(data: dict[str, Any]) -> dict: def get_mooncake_rollout_data(args: Any, ref: Box) -> dict[str, Any]: - return _mooncake_transfer(args, contribute_segment=False).get_legacy_dict(import_ref(ref.inner)) + return _mooncake_transfer(args, contribute_segment=False).get(import_ref(ref.inner), type="dict") def release_mooncake_rollout_data(args: Any, data: dict[str, Any]) -> None: @@ -95,7 +96,7 @@ def release_mooncake_rollout_data(args: Any, data: dict[str, Any]) -> None: def cleanup_mooncake_rollout_data(args: Any, ref: Box) -> None: - _mooncake_transfer(args, contribute_segment=False).remove_legacy_dict(import_ref(ref.inner)) + _mooncake_transfer(args, contribute_segment=False).cleanup_dataproto(import_ref(ref.inner)) def cleanup_mooncake_rollout_refs(args: Any, refs: list[Box]) -> None: From 1d1433791ae50ef47d5ad9919448e6c7eea8ab37 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 14 Aug 2026 12:13:42 +0800 Subject: [PATCH 08/11] docs: add Mooncake rollout transfer guide --- docs/en/advanced/mooncake-rollout-transfer.md | 273 ++++++++++++++++++ docs/en/index.rst | 1 + docs/zh/advanced/mooncake-rollout-transfer.md | 265 +++++++++++++++++ docs/zh/index.rst | 1 + slime/utils/arguments.py | 3 +- 5 files changed, 542 insertions(+), 1 deletion(-) create mode 100644 docs/en/advanced/mooncake-rollout-transfer.md create mode 100644 docs/zh/advanced/mooncake-rollout-transfer.md diff --git a/docs/en/advanced/mooncake-rollout-transfer.md b/docs/en/advanced/mooncake-rollout-transfer.md new file mode 100644 index 0000000000..ab66dc4c2d --- /dev/null +++ b/docs/en/advanced/mooncake-rollout-transfer.md @@ -0,0 +1,273 @@ +# Mooncake Rollout Data Transfer + +slime normally passes rollout data from the rollout manager to the trainer through +Ray's object store. Mooncake Store is an alternative for this handoff. It keeps the +rollout dictionary contract unchanged while moving the payload through Mooncake's +data plane. + +This setting only controls rollout data transfer: + +- `--rollout-data-transport` selects how rollout data reaches the trainer. +- Weight synchronization is configured separately by the `--update-weight-*` + options. +- SGLang KV-cache disaggregation uses its own transfer settings. + +Both `train.py` and `train_async.py` support the Mooncake path. + +## Requirements + +Before starting a slime job: + +- Use the same slime revision and Mooncake wheel on every Ray node. The Mooncake + package must provide the structured-object `put`/`get` APIs used by slime. +- Start `mooncake_master`, or provide a Mooncake HA endpoint. slime connects as a + client and does not manage the endpoint lifecycle. +- Use data-network addresses that are reachable from every Ray node. +- Reserve enough host memory for `MOONCAKE_GLOBAL_SEGMENT_SIZE` and + `MOONCAKE_LOCAL_BUFFER_SIZE`. +- For RDMA, expose the RDMA device to the runtime, allow memory locking, and set the + node-local device name before starting Ray. +- Make the model, Megatron checkpoint, dataset, and Python environment available on + every node that may run the corresponding worker. + +Follow the [Mooncake installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html) +for packages and platform requirements. + +## Configure the backend + +Choose TCP or RDMA before starting Ray. TCP only needs a routable data network. +RDMA also needs a local device on each node; device names may differ across nodes. + +Add this option to an existing slime recipe: + +```bash +--rollout-data-transport mooncake +``` + +Mooncake reads its connection settings from the environment. Export them before +starting Ray so that Ray workers inherit the node-local values: + +```bash +export MOONCAKE_MASTER=":50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" + +# Optional when the Ray node IP is already the desired data-network address. +# export MOONCAKE_LOCAL_HOSTNAME="" + +# RDMA only. Set the device attached to this node's data-network address. +# export MOONCAKE_DEVICE="" +``` + +`MOONCAKE_LOCAL_HOSTNAME` and `MOONCAKE_DEVICE` are node-local. Do not put one +node's values into a cluster-wide Ray runtime environment. + +## Two-node walkthrough + +The following example runs two synchronous rollout and training iterations with +Qwen3-4B. It uses one eight-GPU node for training and one eight-GPU node for rollout. +Set the variables for your cluster, then complete the steps in order. + +Both nodes must use the same Python environment, slime revision, and Mooncake wheel. + +### 1. Choose the protocol and set node-local values + +On the head node: + +```bash +export HEAD_IP="" +export MOONCAKE_MASTER="${HEAD_IP}:50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" +export MOONCAKE_LOCAL_HOSTNAME="${HEAD_IP}" + +# RDMA only. +# export MOONCAKE_DEVICE="" +``` + +On the worker node: + +```bash +export HEAD_IP="" +export WORKER_IP="" +export MOONCAKE_MASTER="${HEAD_IP}:50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" +export MOONCAKE_LOCAL_HOSTNAME="${WORKER_IP}" + +# RDMA only. +# export MOONCAKE_DEVICE="" +``` + +The 2 GiB settings keep this small example easy to run. For a production job, +size both values for the rollout payloads and the number of partitions that may +remain in flight at the same time. + +### 2. Start Mooncake and Ray on the head + +Activate the slime environment, then run: + +```bash +mooncake_master --rpc_address=0.0.0.0 --rpc_port=50051 \ + >mooncake_master.log 2>&1 & + +ray stop --force +ray start --head \ + --node-ip-address="${HEAD_IP}" \ + --port=6379 \ + --num-gpus=8 \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 +``` + +An existing Mooncake master can be replaced with a configured HA endpoint; in that +case, do not launch another local master. + +### 3. Join the worker + +Activate the same slime environment on the worker, then run: + +```bash +ray stop --force +ray start \ + --address="${HEAD_IP}:6379" \ + --node-ip-address="${WORKER_IP}" \ + --num-gpus=8 \ + --disable-usage-stats +``` + +Check `ray status` on the head before submitting training. The cluster should report +16 GPUs. + +### 4. Submit training from the head + +Set paths that are valid on the nodes where the corresponding workers run: + +```bash +export SLIME_HOME="" +export MEGATRON_HOME="" +export HF_CHECKPOINT="" +export REF_LOAD="" +export PROMPT_DATA="" +export RAY_DASHBOARD_ADDR="http://127.0.0.1:8265" + +cd "${SLIME_HOME}" +source scripts/models/qwen3-4B.sh + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${MEGATRON_HOME}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"0\", + \"MOONCAKE_MASTER\": \"${MOONCAKE_MASTER}\", + \"MOONCAKE_PROTOCOL\": \"${MOONCAKE_PROTOCOL}\", + \"MOONCAKE_TE_META_DATA_SERVER\": \"${MOONCAKE_TE_META_DATA_SERVER}\", + \"MOONCAKE_GLOBAL_SEGMENT_SIZE\": \"${MOONCAKE_GLOBAL_SEGMENT_SIZE}\", + \"MOONCAKE_LOCAL_BUFFER_SIZE\": \"${MOONCAKE_LOCAL_BUFFER_SIZE}\" + } +}" + +ray job submit --address="${RAY_DASHBOARD_ADDR}" \ + --working-dir="${SLIME_HOME}" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --rollout-num-gpus 8 \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${HF_CHECKPOINT}" \ + --ref-load "${REF_LOAD}" \ + --load "${REF_LOAD}" \ + --prompt-data "${PROMPT_DATA}" \ + --input-key prompt \ + --label-key label \ + --apply-chat-template \ + --rollout-shuffle \ + --rm-type deepscaler \ + --start-rollout-id 0 \ + --num-rollout 2 \ + --rollout-batch-size 2 \ + --n-samples-per-prompt 2 \ + --rollout-max-response-len 256 \ + --rollout-temperature 1 \ + --global-batch-size 4 \ + --balance-data \ + --rollout-data-transport mooncake \ + --tensor-model-parallel-size 8 \ + --sequence-parallel \ + --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --expert-model-parallel-size 1 \ + --expert-tensor-parallel-size 1 \ + --use-dynamic-batch-size \ + --max-tokens-per-gpu 1024 \ + --advantage-estimator grpo \ + --use-kl-loss \ + --kl-loss-coef 0.0 \ + --kl-loss-type low_var_kl \ + --entropy-coef 0.0 \ + --eps-clip 0.2 \ + --eps-clip-high 0.28 \ + --optimizer adam \ + --lr 1e-6 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + --rollout-num-gpus-per-engine 8 \ + --sglang-mem-fraction-static 0.35 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash \ + --no-gradient-accumulation-fusion \ + --bf16 \ + --distributed-backend nccl +``` + +For fully asynchronous training, keep the same Mooncake environment and transport +option, then use the regular async entrypoint and settings: + +```diff +- -- python3 train.py ... ++ -- python3 train_async.py ... +``` + +See the [fully asynchronous example](../_examples_synced/fully_async/README.md) for +the remaining async arguments. + +## Configuration reference + +| Setting | Default | Purpose | +|---|---|---| +| `--rollout-data-transport` | `object-store` | Set to `mooncake` to use Mooncake for rollout data. | +| `MOONCAKE_MASTER` | none | Address of `mooncake_master` or the HA metadata endpoint. | +| `MOONCAKE_LOCAL_HOSTNAME` | Ray node IP | Data-network address advertised by the local client. | +| `MOONCAKE_TE_META_DATA_SERVER` | `P2PHANDSHAKE` | Transfer Engine metadata service. | +| `MOONCAKE_PROTOCOL` | `rdma` | Transfer protocol, normally `tcp` or `rdma`. | +| `MOONCAKE_DEVICE` | auto-discovery | Local RDMA device name. | +| `MOONCAKE_GLOBAL_SEGMENT_SIZE` | `8gb` | Store capacity contributed by rollout writers. | +| `MOONCAKE_LOCAL_BUFFER_SIZE` | `32gb` | Local transfer and staging capacity. | + +The trainer-side GET clients do not contribute a global Store segment. slime still +uses `MOONCAKE_LOCAL_BUFFER_SIZE` there for transfer and pool-backed results, and +releases those buffers after the training step consumes them. + +## Troubleshooting + +- **Mooncake import fails during argument parsing:** install a compatible Mooncake + wheel in the Python environment used by every Ray worker. +- **Store setup fails:** verify `MOONCAKE_MASTER`, endpoint reachability, and that all + nodes use the same Mooncake version. +- **RDMA setup fails:** verify `MOONCAKE_DEVICE`, locked-memory limits, device access, + and that `MOONCAKE_LOCAL_HOSTNAME` belongs to the selected RDMA network. +- **Allocation fails:** increase available host memory or lower the two Mooncake size + settings. Account for concurrent rollout partitions and in-flight steps. diff --git a/docs/en/index.rst b/docs/en/index.rst index 83230af1fa..0b60020746 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -84,6 +84,7 @@ Start by Use Case advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/mooncake-rollout-transfer.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/docs/zh/advanced/mooncake-rollout-transfer.md b/docs/zh/advanced/mooncake-rollout-transfer.md new file mode 100644 index 0000000000..0729a2397e --- /dev/null +++ b/docs/zh/advanced/mooncake-rollout-transfer.md @@ -0,0 +1,265 @@ +# 使用 Mooncake 传输 Rollout Data + +slime 默认通过 Ray object store 将 rollout data 从 rollout manager 传给 trainer。 +Mooncake Store 可以替换这段传输路径：rollout dict 的接口和训练语义保持不变， +数据 payload 则通过 Mooncake data plane 传输。 + +这个配置只影响 rollout data： + +- `--rollout-data-transport` 选择 rollout data 到达 trainer 的方式； +- 模型权重同步由 `--update-weight-*` 参数单独配置； +- SGLang KV cache disaggregation 也使用独立的传输配置。 + +同步入口 `train.py` 和异步入口 `train_async.py` 都支持 Mooncake。 + +## 环境要求 + +启动 slime 任务前，请确认： + +- 所有 Ray 节点使用相同的 slime revision 和 Mooncake wheel。Mooncake 包需要 + 提供 slime 使用的 structured-object `put`/`get` 接口； +- 已经启动 `mooncake_master`，或者准备好 Mooncake HA endpoint。slime 只作为 + client 连接，不负责管理 endpoint 生命周期； +- 所有 Ray 节点都能访问用于数据传输的网络地址； +- 为 `MOONCAKE_GLOBAL_SEGMENT_SIZE` 和 `MOONCAKE_LOCAL_BUFFER_SIZE` + 预留足够的主机内存； +- 使用 RDMA 时，运行环境能够访问 RDMA device、允许锁定内存，并且在启动 + Ray 前配置好每个节点自己的 device name； +- 模型、Megatron checkpoint、数据集和 Python 环境在使用它们的节点上均可用。 + +Mooncake 的包和平台要求请参考 +[安装文档](https://kvcache-ai.github.io/Mooncake/getting_started/build.html)。 + +## 配置传输后端 + +启动 Ray 前先选择 TCP 或 RDMA。TCP 只要求数据网络可达；RDMA 还需要每个 +节点配置本地 device，不同节点的 device name 可以不同。 + +在已有 slime recipe 中增加： + +```bash +--rollout-data-transport mooncake +``` + +Mooncake 从环境变量中读取连接配置。请在启动 Ray 前导出这些变量，使 Ray +worker 继承节点本地配置： + +```bash +export MOONCAKE_MASTER=":50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" + +# Ray node IP 已经是数据网络地址时可以省略。 +# export MOONCAKE_LOCAL_HOSTNAME="" + +# 仅 RDMA 需要。填写当前节点数据网卡对应的 device。 +# export MOONCAKE_DEVICE="" +``` + +`MOONCAKE_LOCAL_HOSTNAME` 和 `MOONCAKE_DEVICE` 是节点本地配置。不要把 +某一个节点的值写入整个集群共用的 Ray runtime environment。 + +## 双机运行示例 + +下面的示例使用 Qwen3-4B 完成两轮同步 rollout 和训练。一个八卡节点用于 +训练，另一个八卡节点用于 rollout。请先根据集群环境设置变量，再按顺序执行。 + +两个节点必须使用相同的 Python 环境、slime revision 和 Mooncake wheel。 + +### 1. 选择协议并配置节点变量 + +在 head 节点执行： + +```bash +export HEAD_IP="" +export MOONCAKE_MASTER="${HEAD_IP}:50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" +export MOONCAKE_LOCAL_HOSTNAME="${HEAD_IP}" + +# 仅 RDMA 需要。 +# export MOONCAKE_DEVICE="" +``` + +在 worker 节点执行： + +```bash +export HEAD_IP="" +export WORKER_IP="" +export MOONCAKE_MASTER="${HEAD_IP}:50051" +export MOONCAKE_PROTOCOL="" +export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" +export MOONCAKE_GLOBAL_SEGMENT_SIZE="2gb" +export MOONCAKE_LOCAL_BUFFER_SIZE="2gb" +export MOONCAKE_LOCAL_HOSTNAME="${WORKER_IP}" + +# 仅 RDMA 需要。 +# export MOONCAKE_DEVICE="" +``` + +这里使用 2 GiB 是为了方便运行这个小规模示例。生产任务需要根据 rollout +payload 大小和同时处于 in-flight 状态的 partition 数量重新评估这两个值。 + +### 2. 在 head 节点启动 Mooncake 和 Ray + +激活 slime Python 环境后执行： + +```bash +mooncake_master --rpc_address=0.0.0.0 --rpc_port=50051 \ + >mooncake_master.log 2>&1 & + +ray stop --force +ray start --head \ + --node-ip-address="${HEAD_IP}" \ + --port=6379 \ + --num-gpus=8 \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 +``` + +如果使用已有的 Mooncake master 或 HA endpoint，不要再启动本地 master。 + +### 3. 将 worker 加入 Ray 集群 + +在 worker 节点激活相同的 Python 环境，然后执行： + +```bash +ray stop --force +ray start \ + --address="${HEAD_IP}:6379" \ + --node-ip-address="${WORKER_IP}" \ + --num-gpus=8 \ + --disable-usage-stats +``` + +提交训练前，在 head 节点运行 `ray status`，确认集群中有 16 张 GPU。 + +### 4. 从 head 节点提交训练 + +以下路径必须在实际使用它们的节点上有效： + +```bash +export SLIME_HOME="" +export MEGATRON_HOME="" +export HF_CHECKPOINT="" +export REF_LOAD="" +export PROMPT_DATA="" +export RAY_DASHBOARD_ADDR="http://127.0.0.1:8265" + +cd "${SLIME_HOME}" +source scripts/models/qwen3-4B.sh + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${MEGATRON_HOME}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"0\", + \"MOONCAKE_MASTER\": \"${MOONCAKE_MASTER}\", + \"MOONCAKE_PROTOCOL\": \"${MOONCAKE_PROTOCOL}\", + \"MOONCAKE_TE_META_DATA_SERVER\": \"${MOONCAKE_TE_META_DATA_SERVER}\", + \"MOONCAKE_GLOBAL_SEGMENT_SIZE\": \"${MOONCAKE_GLOBAL_SEGMENT_SIZE}\", + \"MOONCAKE_LOCAL_BUFFER_SIZE\": \"${MOONCAKE_LOCAL_BUFFER_SIZE}\" + } +}" + +ray job submit --address="${RAY_DASHBOARD_ADDR}" \ + --working-dir="${SLIME_HOME}" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 8 \ + --rollout-num-gpus 8 \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint "${HF_CHECKPOINT}" \ + --ref-load "${REF_LOAD}" \ + --load "${REF_LOAD}" \ + --prompt-data "${PROMPT_DATA}" \ + --input-key prompt \ + --label-key label \ + --apply-chat-template \ + --rollout-shuffle \ + --rm-type deepscaler \ + --start-rollout-id 0 \ + --num-rollout 2 \ + --rollout-batch-size 2 \ + --n-samples-per-prompt 2 \ + --rollout-max-response-len 256 \ + --rollout-temperature 1 \ + --global-batch-size 4 \ + --balance-data \ + --rollout-data-transport mooncake \ + --tensor-model-parallel-size 8 \ + --sequence-parallel \ + --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --expert-model-parallel-size 1 \ + --expert-tensor-parallel-size 1 \ + --use-dynamic-batch-size \ + --max-tokens-per-gpu 1024 \ + --advantage-estimator grpo \ + --use-kl-loss \ + --kl-loss-coef 0.0 \ + --kl-loss-type low_var_kl \ + --entropy-coef 0.0 \ + --eps-clip 0.2 \ + --eps-clip-high 0.28 \ + --optimizer adam \ + --lr 1e-6 \ + --lr-decay-style constant \ + --weight-decay 0.1 \ + --adam-beta1 0.9 \ + --adam-beta2 0.98 \ + --rollout-num-gpus-per-engine 8 \ + --sglang-mem-fraction-static 0.35 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 \ + --attention-softmax-in-fp32 \ + --attention-backend flash \ + --no-gradient-accumulation-fusion \ + --bf16 \ + --distributed-backend nccl +``` + +使用 fully async 训练时，保留相同的 Mooncake 环境变量和传输参数，切换到 +slime 原有的异步入口和配置： + +```diff +- -- python3 train.py ... ++ -- python3 train_async.py ... +``` + +其余异步参数请参考 [fully async 示例](../_examples_synced/fully_async/README.md)。 + +## 配置项说明 + +| 配置 | 默认值 | 作用 | +|---|---|---| +| `--rollout-data-transport` | `object-store` | 设为 `mooncake` 后使用 Mooncake 传输 rollout data。 | +| `MOONCAKE_MASTER` | 无 | `mooncake_master` 或 HA metadata endpoint 的地址。 | +| `MOONCAKE_LOCAL_HOSTNAME` | Ray node IP | 当前 client 对外发布的数据网络地址。 | +| `MOONCAKE_TE_META_DATA_SERVER` | `P2PHANDSHAKE` | Transfer Engine metadata service。 | +| `MOONCAKE_PROTOCOL` | `rdma` | 传输协议，通常使用 `tcp` 或 `rdma`。 | +| `MOONCAKE_DEVICE` | 自动发现 | 当前节点的 RDMA device name。 | +| `MOONCAKE_GLOBAL_SEGMENT_SIZE` | `8gb` | rollout writer 贡献的 Store 容量。 | +| `MOONCAKE_LOCAL_BUFFER_SIZE` | `32gb` | 本地传输和 staging buffer 容量。 | + +trainer 侧的 GET client 不会贡献 global Store segment，但仍会使用 +`MOONCAKE_LOCAL_BUFFER_SIZE` 完成传输和管理 pool-backed result。训练消费完 +数据后，slime 会释放这些 buffer。 + +## 常见问题 + +- **参数解析时无法导入 Mooncake：**所有 Ray worker 的 Python 环境都需要安装 + 兼容的 Mooncake wheel； +- **Store setup 失败：**检查 `MOONCAKE_MASTER`、endpoint 可达性，以及所有 + 节点使用的 Mooncake 版本是否一致； +- **RDMA 初始化失败：**检查 `MOONCAKE_DEVICE`、locked-memory limit、device + 权限，以及 `MOONCAKE_LOCAL_HOSTNAME` 是否属于所选 RDMA 网络； +- **内存分配失败：**释放主机内存或下调两个 Mooncake size 参数，并为并发的 + rollout partition 和 in-flight step 预留容量。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 747deddf63..a5763e741e 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -84,6 +84,7 @@ slime 的设计目标，是让这两大能力彼此强化，同时避免把系 advanced/pd-disaggregation.md advanced/external-rollout-engines.md advanced/delta-weight-sync.md + advanced/mooncake-rollout-transfer.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index d60d5a4bec..40e00a0633 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -562,7 +562,8 @@ def add_rollout_arguments(parser): help=( "Transport for rollout data refs sent from rollout manager to trainer. Large rollout " "fields are tensorized on CPU before the refs are stored. Set to nixl to transfer " - "those torch tensors via Ray NIXL." + "those torch tensors via Ray NIXL, or mooncake to transfer the rollout dictionary " + "through Mooncake Store." ), ) parser.add_argument( From c11b925193794af13b1171fb3c37dfc87779b415 Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Fri, 14 Aug 2026 13:32:45 +0800 Subject: [PATCH 09/11] style: format Mooncake data transfer helper --- slime/utils/data_transfer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/slime/utils/data_transfer.py b/slime/utils/data_transfer.py index 16bb5024cb..c57dd6f3ce 100644 --- a/slime/utils/data_transfer.py +++ b/slime/utils/data_transfer.py @@ -8,8 +8,8 @@ from slime.utils.misc import Box try: - from mooncake.structured_object_store import FieldSchema, MooncakeBundleTransfer, export_ref, import_ref from mooncake.store import MooncakeDistributedStore + from mooncake.structured_object_store import FieldSchema, MooncakeBundleTransfer, export_ref, import_ref _MOONCAKE_AVAILABLE = True except ImportError: @@ -63,8 +63,7 @@ def check_mooncake_available() -> None: """Call during argument parsing to fail fast if mooncake is not installed.""" if not _MOONCAKE_AVAILABLE: raise ImportError( - "rollout-data-transport='mooncake' requires the mooncake package. " - "Install it with: pip install mooncake" + "rollout-data-transport='mooncake' requires the mooncake package. " "Install it with: pip install mooncake" ) From 80795381d2b5ccee8e953ed3993acbf8a812076a Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Sat, 15 Aug 2026 00:38:10 +0800 Subject: [PATCH 10/11] ci: retrigger flaky megatron test From c71ae4eb69ae27b379b6921c4e32d59cdf41950e Mon Sep 17 00:00:00 2001 From: Cruz Zhao Date: Sat, 15 Aug 2026 01:36:51 +0800 Subject: [PATCH 11/11] ci: retry transient model loading failure