[Rollout] feat: add Mooncake Store rollout transfer for disaggregated training.#300
[Rollout] feat: add Mooncake Store rollout transfer for disaggregated training.#300stmatengss wants to merge 2 commits into
Conversation
Use MooncakeDistributedStore (mooncake-transfer-engine) to transfer rollout tensors between rollout and training workers via --transfer-backend mooncake_store. Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Teng Ma <stmatengss@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Mooncake Store rollout transfer backend as an alternative to Ray object references for transferring rollout tensors when training and inference run on separate GPUs. The changes include new utility modules for managing remote batches, starting the Mooncake master, and splitting rollout data, along with integration into the training pipelines and new tests. The review feedback focuses on improving robustness, specifically by avoiding exception masking in finally blocks during buffer unregistration, adding defensive None checks during cleanup, guarding against empty tensors to prevent padding errors, and registering an atexit handler to clean up background processes.
| def _put_tensor(store: Any, key: str, tensor: torch.Tensor, config: Any) -> int: | ||
| arr = np.ascontiguousarray(tensor.numpy()) | ||
| region = _WritableRegion(memoryview(arr)) | ||
| try: | ||
| if store.register_buffer(region.ptr, region.size) != 0: | ||
| raise RuntimeError("register_buffer failed for put_from") | ||
| try: | ||
| return store.put_from(key=key, buffer_ptr=region.ptr, size=region.size, config=config) | ||
| finally: | ||
| if store.unregister_buffer(region.ptr) != 0: | ||
| raise RuntimeError("unregister_buffer failed for put_from") | ||
| finally: | ||
| region.close() |
There was a problem hiding this comment.
Avoid masking exceptions in the finally block when unregister_buffer fails. Additionally, add a guard for empty tensors to prevent potential ctypes/buffer errors with 0-byte buffers.
def _put_tensor(store: Any, key: str, tensor: torch.Tensor, config: Any) -> int:\n if tensor.numel() == 0:\n return 0\n arr = np.ascontiguousarray(tensor.numpy())\n region = _WritableRegion(memoryview(arr))\n try:\n if store.register_buffer(region.ptr, region.size) != 0:\n raise RuntimeError(\"register_buffer failed for put_from\")\n try:\n return store.put_from(key=key, buffer_ptr=region.ptr, size=region.size, config=config)\n finally:\n if store.unregister_buffer(region.ptr) != 0:\n import logging\n logging.getLogger(__name__).warning(\"unregister_buffer failed for put_from\")\n finally:\n region.close()| def _get_tensor(store: Any, key: str, shape: tuple[int, ...], dtype_name: str) -> torch.Tensor: | ||
| torch_dtype = getattr(torch, dtype_name.lower()) | ||
| nbytes = int(np.prod(shape, dtype=np.int64)) * torch_dtype.itemsize | ||
| if nbytes == 0: | ||
| return torch.empty(shape, dtype=torch_dtype) | ||
| region = _WritableRegion(bytearray(nbytes)) | ||
| try: | ||
| if store.register_buffer(region.ptr, region.size) != 0: | ||
| raise RuntimeError("register_buffer failed for get_into") | ||
| try: | ||
| got = store.get_into(key, region.ptr, nbytes) | ||
| if got != nbytes: | ||
| raise RuntimeError(f"get_into failed for {key}: expected {nbytes}, got {got}") | ||
| finally: | ||
| if store.unregister_buffer(region.ptr) != 0: | ||
| raise RuntimeError("unregister_buffer failed for get_into") | ||
| count = int(np.prod(shape, dtype=np.int64)) | ||
| return torch.frombuffer(region.buffer, dtype=torch_dtype, count=count).reshape(shape).clone() | ||
| finally: | ||
| region.close() |
There was a problem hiding this comment.
Avoid masking exceptions in the finally block when unregister_buffer fails. If get_into fails, raising a new exception in finally will discard the original error, making debugging difficult.
def _get_tensor(store: Any, key: str, shape: tuple[int, ...], dtype_name: str) -> torch.Tensor:\n torch_dtype = getattr(torch, dtype_name.lower())\n nbytes = int(np.prod(shape, dtype=np.int64)) * torch_dtype.itemsize\n if nbytes == 0:\n return torch.empty(shape, dtype=torch_dtype)\n region = _WritableRegion(bytearray(nbytes))\n try:\n if store.register_buffer(region.ptr, region.size) != 0:\n raise RuntimeError(\"register_buffer failed for get_into\")\n try:\n got = store.get_into(key, region.ptr, nbytes)\n if got != nbytes:\n raise RuntimeError(f\"get_into failed for {key}: expected {nbytes}, got {got}\")\n finally:\n if store.unregister_buffer(region.ptr) != 0:\n import logging\n logging.getLogger(__name__).warning(\"unregister_buffer failed for get_into\")\n count = int(np.prod(shape, dtype=np.int64))\n return torch.frombuffer(region.buffer, dtype=torch_dtype, count=count).reshape(shape).clone()\n finally:\n region.close()| def maybe_cleanup_mooncake_store_refs( | ||
| args: Any, refs: list[RolloutStoreBatch] | RolloutStoreBatch, suppress_errors: bool = False | ||
| ) -> None: | ||
| if getattr(args, "transfer_backend", "ray") != "mooncake_store": | ||
| return | ||
| batches = refs if isinstance(refs, list) else [refs] | ||
| try: | ||
| cleanup_mooncake_store_refs(batches) | ||
| except Exception: | ||
| if not suppress_errors: | ||
| raise | ||
|
|
There was a problem hiding this comment.
Add a defensive check for refs is None to prevent AttributeError during cleanup if training fails early or before references are initialized.
def maybe_cleanup_mooncake_store_refs(\n args: Any, refs: list[RolloutStoreBatch] | RolloutStoreBatch, suppress_errors: bool = False\n) -> None:\n if getattr(args, \"transfer_backend\", \"ray\") != \"mooncake_store\":\n return\n if refs is None:\n return\n batches = refs if isinstance(refs, list) else [refs]\n try:\n cleanup_mooncake_store_refs(batches)\n except Exception:\n if not suppress_errors:\n raise| def cleanup_mooncake_store_refs(refs: list[RolloutStoreBatch]) -> None: | ||
| keys: set[str] = set() | ||
| store_init_kwargs = None | ||
| for batch in refs: | ||
| keys.update(batch.meta_info.get("mooncake_cleanup_keys", [])) | ||
| if store_init_kwargs is None: | ||
| store_init_kwargs = batch.meta_info.get("mooncake_cleanup_store_kwargs") | ||
| if keys and store_init_kwargs is not None: | ||
| remove_mooncake_keys(get_cached_mooncake_store(store_init_kwargs), sorted(keys)) | ||
|
|
There was a problem hiding this comment.
Add a defensive check for batch is None inside the loop to prevent AttributeError during cleanup if some batches are uninitialized.
def cleanup_mooncake_store_refs(refs: list[RolloutStoreBatch]) -> None:\n keys: set[str] = set()\n store_init_kwargs = None\n for batch in refs:\n if batch is None:\n continue\n keys.update(batch.meta_info.get(\"mooncake_cleanup_keys\", []))\n if store_init_kwargs is None:\n store_init_kwargs = batch.meta_info.get(\"mooncake_cleanup_store_kwargs\")\n if keys and store_init_kwargs is not None:\n remove_mooncake_keys(get_cached_mooncake_store(store_init_kwargs), sorted(keys))| def _list_to_padded_tensor(values: list, dtype: torch.dtype) -> tuple[torch.Tensor, list[int]]: | ||
| if not values: | ||
| return torch.empty((0, 0), dtype=dtype), [] | ||
| tensors = [torch.as_tensor(value, dtype=dtype).reshape(-1) for value in values] | ||
| lengths = [int(tensor.numel()) for tensor in tensors] | ||
| return pad_sequence(tensors, batch_first=True, padding_value=0), lengths | ||
|
|
There was a problem hiding this comment.
Add a guard for when all sequences are empty (max(lengths) == 0) to prevent pad_sequence from raising a RuntimeError in older PyTorch versions.
def _list_to_padded_tensor(values: list, dtype: torch.dtype) -> tuple[torch.Tensor, list[int]]:\n if not values:\n return torch.empty((0, 0), dtype=dtype), []\n tensors = [torch.as_tensor(value, dtype=dtype).reshape(-1) for value in values]\n lengths = [int(tensor.numel()) for tensor in tensors]\n if max(lengths) == 0:\n return torch.empty((len(values), 0), dtype=dtype), lengths\n return pad_sequence(tensors, batch_first=True, padding_value=0), lengths| proc = subprocess.Popen( | ||
| [ | ||
| mooncake_master, | ||
| "--enable_http_metadata_server=true", | ||
| f"--http_metadata_server_host={host}", | ||
| f"--http_metadata_server_port={metadata_port}", | ||
| ], | ||
| stdout=subprocess.DEVNULL, | ||
| stderr=subprocess.DEVNULL, | ||
| start_new_session=True, | ||
| ) | ||
| deadline = time.time() + 15.0 | ||
| while time.time() < deadline: | ||
| if proc.poll() is not None: | ||
| raise RuntimeError("mooncake_master exited early") | ||
| if _port_open(host, rpc_port) and _port_open(host, metadata_port): | ||
| return | ||
| time.sleep(0.2) | ||
| proc.send_signal(signal.SIGTERM) | ||
| raise RuntimeError("mooncake_master did not become ready") |
There was a problem hiding this comment.
Register an atexit handler to terminate the background mooncake_master process when the Python interpreter exits. This prevents process leaks and port collisions in CI/CD environments.
proc = subprocess.Popen(\n [\n mooncake_master,\n \"--enable_http_metadata_server=true\",\n f\"--http_metadata_server_host={host}\",\n f\"--http_metadata_server_port={metadata_port}\",\n ],\n stdout=subprocess.DEVNULL,\n stderr=subprocess.DEVNULL,\n start_new_session=True,\n )\n import atexit\n atexit.register(proc.terminate)\n deadline = time.time() + 15.0\n while time.time() < deadline:\n if proc.poll() is not None:\n raise RuntimeError(\"mooncake_master exited early\")\n if _port_open(host, rpc_port) and _port_open(host, metadata_port):\n return\n time.sleep(0.2)\n proc.send_signal(signal.SIGTERM)\n raise RuntimeError(\"mooncake_master did not become ready\")
There was a problem hiding this comment.
Code Review
This pull request introduces support for transferring rollout tensors via Mooncake Store instead of Ray object references when rollout and training run on separate GPUs. The changes include adding CLI arguments, integrating transfer and cleanup logic in the training pipelines, implementing remote batching utilities, and adding documentation and tests. The review feedback highlights several improvement opportunities: handling None values in the cleanup utility to prevent AttributeError, adding type validation for store initialization arguments, registering an atexit handler to prevent background process leaks of the Mooncake master, and logging warnings instead of raising exceptions in finally blocks to avoid masking active exceptions.
| def maybe_cleanup_mooncake_store_refs( | ||
| args: Any, refs: list[RolloutStoreBatch] | RolloutStoreBatch, suppress_errors: bool = False | ||
| ) -> None: | ||
| if getattr(args, "transfer_backend", "ray") != "mooncake_store": | ||
| return |
There was a problem hiding this comment.
The maybe_cleanup_mooncake_store_refs function does not handle the case where refs is None. If refs is None (which can happen in train_async.py if rollout_data_curr_ref is None), cleanup_mooncake_store_refs will be called with [None], leading to an AttributeError when trying to access batch.meta_info.
Adding an early return when refs is None prevents this potential crash.
| def maybe_cleanup_mooncake_store_refs( | |
| args: Any, refs: list[RolloutStoreBatch] | RolloutStoreBatch, suppress_errors: bool = False | |
| ) -> None: | |
| if getattr(args, "transfer_backend", "ray") != "mooncake_store": | |
| return | |
| def maybe_cleanup_mooncake_store_refs( | |
| args: Any, refs: list[RolloutStoreBatch] | RolloutStoreBatch | None, suppress_errors: bool = False | |
| ) -> None: | |
| if getattr(args, "transfer_backend", "ray") != "mooncake_store" or refs is None: | |
| return |
| def normalize_store_init_kwargs(store_init_kwargs: dict[str, Any] | None) -> dict[str, Any]: | ||
| if not store_init_kwargs: | ||
| return default_store_init_kwargs() | ||
| return default_store_init_kwargs() | dict(store_init_kwargs) |
There was a problem hiding this comment.
If store_init_kwargs is passed as a non-dict JSON value (e.g., a list or string) via the CLI, calling dict(store_init_kwargs) will raise a TypeError or ValueError.
Adding an explicit type check ensures robust handling of invalid CLI inputs.
| def normalize_store_init_kwargs(store_init_kwargs: dict[str, Any] | None) -> dict[str, Any]: | |
| if not store_init_kwargs: | |
| return default_store_init_kwargs() | |
| return default_store_init_kwargs() | dict(store_init_kwargs) | |
| def normalize_store_init_kwargs(store_init_kwargs: dict[str, Any] | None) -> dict[str, Any]: | |
| if not store_init_kwargs: | |
| return default_store_init_kwargs() | |
| if not isinstance(store_init_kwargs, dict): | |
| raise TypeError(f"Expected dict for store_init_kwargs, got {type(store_init_kwargs).__name__}") | |
| return default_store_init_kwargs() | store_init_kwargs |
| proc = subprocess.Popen( | ||
| [ | ||
| mooncake_master, | ||
| "--enable_http_metadata_server=true", | ||
| f"--http_metadata_server_host={host}", | ||
| f"--http_metadata_server_port={metadata_port}", | ||
| ], | ||
| stdout=subprocess.DEVNULL, | ||
| stderr=subprocess.DEVNULL, | ||
| start_new_session=True, | ||
| ) |
There was a problem hiding this comment.
The mooncake_master process is started with start_new_session=True and left running in the background. If the test suite or training script exits, this background process will leak and remain running on the host machine, potentially causing port conflicts (50051 or 18080) in subsequent runs.
Registering an atexit handler to terminate the process when the Python interpreter exits will prevent these resource leaks.
import atexit
proc = subprocess.Popen(
[
mooncake_master,
"--enable_http_metadata_server=true",
f"--http_metadata_server_host={host}",
f"--http_metadata_server_port={metadata_port}",
],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
start_new_session=True,
)
atexit.register(lambda: proc.terminate())| try: | ||
| return store.put_from(key=key, buffer_ptr=region.ptr, size=region.size, config=config) | ||
| finally: | ||
| if store.unregister_buffer(region.ptr) != 0: | ||
| raise RuntimeError("unregister_buffer failed for put_from") |
There was a problem hiding this comment.
Raising a RuntimeError inside the finally block when unregister_buffer fails will discard and mask any active exception raised in the inner try block (such as a failure in put_from). This makes debugging extremely difficult because the root cause of the failure is lost.
Consider logging a warning instead of raising a fatal exception during unregistration cleanup.
| try: | |
| return store.put_from(key=key, buffer_ptr=region.ptr, size=region.size, config=config) | |
| finally: | |
| if store.unregister_buffer(region.ptr) != 0: | |
| raise RuntimeError("unregister_buffer failed for put_from") | |
| try: | |
| return store.put_from(key=key, buffer_ptr=region.ptr, size=region.size, config=config) | |
| finally: | |
| if store.unregister_buffer(region.ptr) != 0: | |
| import logging | |
| logging.getLogger(__name__).warning("unregister_buffer failed for put_from") |
1 participant
Summary
This PR adds an optional rollout data transfer path for disaggregated training (separate rollout and actor GPUs). When enabled, vime transfers rollout tensors (
tokens,loss_masks) through Mooncake Store instead of the Ray object store.This is useful when rollout and training run on different GPUs or nodes and you want a dedicated high-bandwidth transfer layer rather than serializing large tensors through Ray.
Ported from slime PR #1709.
What changed
--transfer-backendflag:ray(default) ormooncake_storemooncake-transfer-engine(MooncakeDistributedStore.setup+put_from/get_into)RolloutStoreBatchwrapper carries remote tensor refs plus non-tensor rollout metadata across DP ranks13 files, +708 / −20 lines. Default behavior is unchanged (
--transfer-backend ray).Usage
pip install mooncake-transfer-engine mooncake_master \ --enable_http_metadata_server=true \ --http_metadata_server_host=127.0.0.1 \ --http_metadata_server_port=18080 python train_async.py \ --transfer-backend mooncake_store \ --actor-num-gpus-per-node 1 \ --rollout-num-gpus 1 \ # ... other training argsStore connection defaults to
MOONCAKE_*env vars. Override with--mooncake-store-init-kwargsJSON if needed.Design notes
tokensandloss_masksgo through Mooncake Store; other rollout fields stay in the batch wrappermooncake_store_service.ensure_mooncake_master()is provided for local smoke tests onlyTest plan
pytest tests/utils/test_mooncake_store_transfer.py(3 tests, realmooncake_master)CUDA_VISIBLE_DEVICES=4,5 python tests/run_disaggregate_2gpu.pyraymooncake_storestep_timerollout_timedata_preprocess_timeRollout throughput is effectively the same; Mooncake materialization adds ~10ms/step, negligible at this scale.
Dependencies
mooncake-transfer-engine(only required when--transfer-backend mooncake_storeis set)