From ffd1ba1bed7fb12097d3131e439cc96a3efc0558 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 17:03:23 +0530 Subject: [PATCH 1/5] perf: startup timing probes and env-gated startup optimizations Add STARTUP_TIMING probes across driver, rollout manager, engines, scheduler and train actors, plus scripts/parse_startup_timing.py to rebuild the timeline. Switch the sglang engine weight loader to mmap (SGLANG_USE_RUNAI_MODEL_STREAMER=false). Env-flag gated, off by default: MILES_STARTUP_OVERLAP, MILES_STARTUP_BUDGET_GATE, MILES_STARTUP_GATE_REQUIRE_ENGINES, MILES_STARTUP_SKIP_ROUNDTRIP, MILES_STARTUP_ENGINES_FIRST, MILES_STARTUP_ROUTER_POLL, MILES_STARTUP_BENCH_EXIT. --- miles/backends/fsdp_utils/actor.py | 48 ++++--- miles/backends/fsdp_utils/model_loader.py | 6 + miles/backends/fsdp_utils/startup_gate.py | 70 +++++++++++ .../sglang_diffusion_engine.py | 53 +++++--- miles/ray/actor_group.py | 5 +- miles/ray/placement_group.py | 49 +++++--- miles/ray/rollout.py | 118 ++++++++++++++---- miles/utils/startup_timing.py | 55 ++++++++ scripts/parse_startup_timing.py | 88 +++++++++++++ .../utils/test_startup_timing_and_gate.py | 104 +++++++++++++++ train_diffusion.py | 32 +++-- 11 files changed, 544 insertions(+), 84 deletions(-) create mode 100644 miles/backends/fsdp_utils/startup_gate.py create mode 100644 miles/utils/startup_timing.py create mode 100644 scripts/parse_startup_timing.py create mode 100644 tests/fast/utils/test_startup_timing_and_gate.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 1aebbeb05..264ecfefb 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,4 +1,5 @@ import logging +import os from argparse import Namespace from contextlib import nullcontext @@ -11,6 +12,7 @@ import miles.backends.fsdp_utils.configs.qwen_image # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.sd3 # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.wan2_2 # noqa: F401 — register pipeline config +import miles.utils.startup_timing as st from miles.ray.train_actor import TrainRayActor from miles.utils import tracking_utils, train_metric_utils from miles.utils.context_utils import with_defer @@ -66,8 +68,10 @@ class FSDPTrainRayActor(TrainRayActor): """ @with_defer(lambda: Timer().start("train_wait")) - def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # type: ignore[override] - super().init(args, role, with_ref) + def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manager=None) -> int: # type: ignore[override] + st.mark("train.init_enter", rank=os.environ.get("RANK", "?")) + with st.step("train.super_init"): + super().init(args, role, with_ref) if args.deterministic_mode: _enable_deterministic_training(args) @@ -104,16 +108,18 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty # flash-attn is opaque to torch's determinism flag; backends patch their own dispatch. self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend) self.scheduler = self.model_backend.load_scheduler(args) - self.models = load_fsdp_models( - args, - self.model_backend, - self.train_pipeline_config, - self.parallel_state, - checkpoint_path=args.hf_checkpoint, - lora_adapter_paths=args.lora_adapter_paths, - trainable=True, - cpu_offload=args.fsdp_cpu_offload, - ) + with st.step("train.load_fsdp_models"): + self.models = load_fsdp_models( + args, + self.model_backend, + self.train_pipeline_config, + self.parallel_state, + checkpoint_path=args.hf_checkpoint, + lora_adapter_paths=args.lora_adapter_paths, + trainable=True, + cpu_offload=args.fsdp_cpu_offload, + rollout_manager=rollout_manager, + ) self.reference_models = {} if args.ref_load is not None: self.reference_models = load_fsdp_models( @@ -229,7 +235,9 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty checkpoint.finalize_load(self, checkpoint_payload) if self.args.offload_train: - self.sleep() + with st.step("train.offload_after_init"): + self.sleep() + st.mark("train.init_done") return self.args.start_rollout_id @@ -283,12 +291,14 @@ def update_weights(self, previous_ema: bool = False) -> None: # type: ignore[ov dist.barrier(group=get_gloo_group()) return - rollout_engines, rollout_engine_lock, num_new_engines = ray.get( - self.rollout_manager.get_rollout_engines_and_lock.remote() - ) + with st.step("train.uw.get_engines"): + rollout_engines, rollout_engine_lock, num_new_engines = ray.get( + self.rollout_manager.get_rollout_engines_and_lock.remote() + ) if num_new_engines > 0: - self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) - dist.barrier(group=get_gloo_group()) + with st.step("train.uw.connect_engines"): + self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock) + dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_num_new_engines.remote()) @@ -297,7 +307,7 @@ def update_weights(self, previous_ema: bool = False) -> None: # type: ignore[ov if self.args.rollout_weights == "ema" else nullcontext() ) - with rollout_weight_context: + with rollout_weight_context, st.step("train.uw.push_weights"): self.weight_updater.update_weights() clear_memory() diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index ca61e39ec..569cb4bff 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -10,6 +10,7 @@ from . import checkpoint from .device_move import move_model +from .startup_gate import budget_gate_enabled, wait_for_gpu_budget from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str from .sequence_parallel.plan import apply_sequence_parallel @@ -26,6 +27,7 @@ def load_fsdp_models( lora_adapter_paths: list[str] | None = None, trainable: bool = False, cpu_offload: bool = False, + rollout_manager=None, ) -> dict[str, torch.nn.Module]: materialize_weights = dist.get_rank() == 0 master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) @@ -54,6 +56,10 @@ def load_fsdp_models( raise RuntimeError(f"{component} did not honor meta initialization") checkpoint.sync_model_dtypes(model) full_state = model.state_dict() if materialize_weights else {} + if trainable and budget_gate_enabled(): + wait_for_gpu_budget( + model, shard_world_size=parallel_state.get_mesh("fsdp").size(), rollout_manager=rollout_manager + ) model = apply_fsdp2( model, model_backend.fsdp_parallel_plan(model), diff --git a/miles/backends/fsdp_utils/startup_gate.py b/miles/backends/fsdp_utils/startup_gate.py new file mode 100644 index 000000000..f62360453 --- /dev/null +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -0,0 +1,70 @@ +"""Hold the train side's first large GPU allocation until the colocated engines leave enough room. + +Replaces a phase-based gate ("wait until the engines booted and offloaded") with a memory budget: +the train GPU phase only needs one parameter shard plus the largest full tensor (the broadcast +from rank 0 moves one tensor at a time) plus a margin. Sizes come from the module's parameter +shapes, so no weights are read. +""" + +from __future__ import annotations + +import logging +import os +import time + +import torch +import torch.distributed as dist + +import miles.utils.startup_timing as st + +logger = logging.getLogger(__name__) + +GIB = 1 << 30 + + +def budget_gate_enabled() -> bool: + return os.environ.get("MILES_STARTUP_BUDGET_GATE") == "1" + + +def required_bytes(model: torch.nn.Module, shard_world_size: int, margin_bytes: int) -> tuple[int, int, int]: + total = 0 + largest = 0 + for p in model.parameters(): + nbytes = p.numel() * p.element_size() + total += nbytes + largest = max(largest, nbytes) + shard = -(-total // shard_world_size) + return shard, largest, shard + largest + margin_bytes + + +def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollout_manager=None) -> None: + """Block until free GPU memory covers the init peak; with MILES_STARTUP_GATE_REQUIRE_ENGINES=1 also + wait for engine init to finish (guards the race where train allocates before the engines load).""" + margin = int(float(os.environ.get("MILES_STARTUP_GATE_MARGIN_GIB", "4")) * GIB) + timeout = float(os.environ.get("MILES_STARTUP_GATE_TIMEOUT_S", "900")) + shard, largest, need = required_bytes(model, shard_world_size, margin) + with st.step( + "train.budget_gate", + need_gib=f"{need / GIB:.1f}", + shard_gib=f"{shard / GIB:.1f}", + largest_gib=f"{largest / GIB:.2f}", + ): + if os.environ.get("MILES_STARTUP_GATE_REQUIRE_ENGINES") == "1" and rollout_manager is not None: + import ray + + ray.get(rollout_manager.wait_engines_ready.remote()) + start = time.time() + while True: + free = torch.cuda.mem_get_info()[0] + if free >= need: + break + if time.time() - start > timeout: + raise RuntimeError( + f"startup gate: only {free / GIB:.1f} GiB free after {timeout:.0f}s, need {need / GIB:.1f} GiB " + f"(shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} + margin {margin / GIB:.1f})" + ) + time.sleep(0.25) + st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}") + if dist.is_initialized(): + # Ranks gate on their own device; keep them in step so the broadcast does not start half-gated. + dist.barrier() diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2aab5ae8d..92f3ae53b 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -11,6 +11,7 @@ import requests from sglang.srt.utils.common import kill_process_tree +import miles.utils.startup_timing as st from miles.ray.ray_actor import RayActor from miles.utils.arguments import rollout_merges_lora from miles.utils.http_utils import get_host_info @@ -39,14 +40,28 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int: def _bootstrap_scheduler_process_with_rollout_patches(spec): # Re-apply patches here: the spawned scheduler grandchild re-imports modules fresh, losing the parent's patches. # Patch after platform init, as bootstrap orders it; its own init call is then a no-op. + st.set_role(f"scheduler{os.getpid()}") + st.mark("scheduler.enter") from sglang.multimodal_gen.runtime.managers.worker_bootstrap import bootstrap_scheduler_process from sglang.multimodal_gen.runtime.platforms import initialize_current_platform from miles.backends.sglang_diffusion_utils.monkey_patches import apply_env_selected_rollout_patches - initialize_current_platform() - apply_env_selected_rollout_patches() - bootstrap_scheduler_process(spec) + with st.step("scheduler.platform_and_patches"): + initialize_current_platform() + apply_env_selected_rollout_patches() + with st.step("scheduler.bootstrap"): + bootstrap_scheduler_process(spec) + + +def _bootstrap_scheduler_process_timed(spec): + # Timing-only wrapper for the no-patch path; keeps the grandchild visible to the startup probes. + st.set_role(f"scheduler{os.getpid()}") + st.mark("scheduler.enter") + from sglang.multimodal_gen.runtime.managers.worker_bootstrap import bootstrap_scheduler_process + + with st.step("scheduler.bootstrap"): + bootstrap_scheduler_process(spec) def _launch_server_target(server_args, apply_rollout_patches: bool = False): @@ -56,11 +71,12 @@ def _launch_server_target(server_args, apply_rollout_patches: bool = False): if server_args.attention_backend_config is not None: server_args.attention_backend_config = addict.Dict(server_args.attention_backend_config) - if apply_rollout_patches: - # Rebind the scheduler entrypoint in launch_server's module so spawn pickles our patched wrapper by qualname. - import sglang.multimodal_gen.runtime.launch_server as _ls_mod + # Rebind the scheduler entrypoint in launch_server's module so spawn pickles our wrapper by qualname. + import sglang.multimodal_gen.runtime.launch_server as _ls_mod - _ls_mod.bootstrap_scheduler_process = _bootstrap_scheduler_process_with_rollout_patches + _ls_mod.bootstrap_scheduler_process = ( + _bootstrap_scheduler_process_with_rollout_patches if apply_rollout_patches else _bootstrap_scheduler_process_timed + ) from sglang.multimodal_gen.runtime.launch_server import launch_server @@ -78,12 +94,14 @@ def launch_server_process( target=_launch_server_target, args=(server_args, apply_rollout_patches), ) - p.start() + with st.step("engine.spawn_server_process"): + p.start() - _wait_server_healthy( - base_url=server_args.url(), - is_process_alive=lambda: p.is_alive(), - ) + with st.step("engine.wait_server_healthy"): + _wait_server_healthy( + base_url=server_args.url(), + is_process_alive=lambda: p.is_alive(), + ) return p @@ -115,6 +133,8 @@ def __init__(self, args, rank: int, base_gpu_id: int | None = None): self.base_gpu_id = base_gpu_id def init(self, dist_init_addr, port, nccl_port, host=None): + st.set_role(f"engine{self.rank}") + st.mark("engine.init_enter") # SGL-D runs single-node per engine; accept dist_init_addr for caller compat, then drop. del dist_init_addr self.router_ip = self.args.sglang_router_ip @@ -165,10 +185,11 @@ def _init_normal(self, server_args_dict): if self.node_rank == 0 and self.router_ip and self.router_port: if self.args.use_miles_router: - response = requests.post( - f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}" - ) - response.raise_for_status() + with st.step("engine.router_register"): + response = requests.post( + f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}" + ) + response.raise_for_status() else: # SGL-D router TODO: add_worker path for the non-miles router logger.warning("Skipping router add_worker: only miles_router is supported for now") diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 7e61a1302..0eb24d3fa 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -88,12 +88,13 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) self._actor_handlers.append(actor) - def async_init(self, args, role, with_ref=False): + def async_init(self, args, role, with_ref=False, rollout_manager=None): """ Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. """ self.args = args - return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handlers] + kwargs = {} if rollout_manager is None else {"rollout_manager": rollout_manager} + return [actor.init.remote(args, role, with_ref=with_ref, **kwargs) for actor in self._actor_handlers] def async_train(self, rollout_id, rollout_data_ref): """Do one rollout training""" diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b87de71d7..b9e7c4465 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -1,10 +1,13 @@ import logging +import os import socket import ray from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +import miles.utils.startup_timing as st + from .actor_group import RayTrainGroup from .rollout import RolloutManager from .utils import TRAIN_ACTOR_GPU @@ -42,11 +45,12 @@ def sort_key(x): def _create_placement_group(num_gpus): """Create a placement group with the specified number of GPUs.""" bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] - pg = placement_group(bundles, strategy="PACK") - num_bundles = len(bundles) + with st.step("driver.pg.create_and_ready", num_gpus=num_gpus): + pg = placement_group(bundles, strategy="PACK") + num_bundles = len(bundles) - logger.info("Waiting for placement group to be ready...") - ray.get(pg.ready()) + logger.info("Waiting for placement group to be ready...") + ray.get(pg.ready()) logger.info("Placement group is ready.") # use info actor to get the GPU id info_actors = [] @@ -131,13 +135,18 @@ def allocate_train_group(args, num_nodes, num_gpus_per_node, pg): def create_training_models(args, pgs, rollout_manager): logger.info("Initializing actor model...") - actor_model = allocate_train_group( - args=args, - num_nodes=args.actor_num_nodes, - num_gpus_per_node=args.actor_num_gpus_per_node, - pg=pgs["actor"], - ) - start_rollout_ids = ray.get(actor_model.async_init(args, role="actor", with_ref=False)) + with st.step("driver.train_group.spawn_actors"): + actor_model = allocate_train_group( + args=args, + num_nodes=args.actor_num_nodes, + num_gpus_per_node=args.actor_num_gpus_per_node, + pg=pgs["actor"], + ) + with st.step("driver.train_group.init_wait"): + gate_manager = rollout_manager if os.environ.get("MILES_STARTUP_GATE_REQUIRE_ENGINES") == "1" else None + start_rollout_ids = ray.get( + actor_model.async_init(args, role="actor", with_ref=False, rollout_manager=gate_manager) + ) logger.info("Actor model initialized.") assert len(set(start_rollout_ids)) == 1 @@ -153,10 +162,11 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): logger.info("Creating rollout manager (num_gpus=%s)", 0) - rollout_manager = RolloutManager.options( - num_cpus=1, - num_gpus=0, - ).remote(args, pg) + with st.step("driver.rollout_manager.ctor"): + rollout_manager = RolloutManager.options( + num_cpus=1, + num_gpus=0, + ).remote(args, pg) # calculate num_rollout from num_epoch num_rollout_per_epoch = None @@ -167,8 +177,13 @@ def create_rollout_manager(args, pg): assert args.num_rollout > 0 logger.info("Computed num_rollout=%s (num_rollout_per_epoch=%s)", args.num_rollout, num_rollout_per_epoch) - if args.offload_rollout: - ray.get(rollout_manager.offload.remote()) + if args.offload_rollout and os.environ.get("MILES_STARTUP_SKIP_ROUNDTRIP") != "1": + if os.environ.get("MILES_STARTUP_OVERLAP") == "1": + # Don't block the driver: train init runs while engines boot and offload. + rollout_manager.offload.remote() + else: + with st.step("driver.rollout_manager.boot_offload"): + ray.get(rollout_manager.offload.remote()) logger.info("Rollout manager created.") return rollout_manager, num_rollout_per_epoch diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 8cbcb6774..d94053725 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -1,6 +1,7 @@ import itertools import logging import multiprocessing +import os import random import threading import time @@ -13,6 +14,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS +import miles.utils.startup_timing as st from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine from miles.ray.data_conversion_hub.flow_grpo import ( expand_samples_to_train_pairs as flow_grpo_expand_samples_to_train_pairs, @@ -45,6 +47,11 @@ logger = logging.getLogger(__name__) +def _flag(name: str) -> bool: + # Startup-benchmark switches; each startup optimization is toggled independently. + return os.environ.get(name) == "1" + + @ray.remote class RolloutManager: """The class to run rollout and convert rollout data to training data.""" @@ -52,24 +59,36 @@ class RolloutManager: def __init__(self, args, pg): configure_logger() + st.set_role("rollout_manager") + st.mark("rm.init_enter") logger.info("RolloutManager init start") self.args = args self.pg = pg from miles.dashboard import hooks - hooks.register_rollout_manager(args) + self._engine_init_handles = [] + self._engines_first = _flag("MILES_STARTUP_ENGINES_FIRST") and not args.train_only and not args.debug_rollout_only + if self._engines_first: + # Engines only need the router address, and add_worker once healthy (35+ s later): boot them first. + _pick_router_addr(args) + self._launch_engines() + with st.step("rm.dashboard_register"): + hooks.register_rollout_manager(args) if not args.train_only: logger.info("RolloutManager: starting router...") - _start_router(args) + with st.step("rm.start_router"): + _start_router(args) logger.info("RolloutManager: router started, init tracking...") # TODO make args immutable - init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") + with st.step("rm.init_tracking"): + init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") logger.info("RolloutManager: init http client...") init_http_client(args) logger.info("RolloutManager: loading data source...") - data_source_cls = load_function(self.args.data_source_path) - self.data_source = data_source_cls(args) + with st.step("rm.load_data_source"): + data_source_cls = load_function(self.args.data_source_path) + self.data_source = data_source_cls(args) logger.info("RolloutManager: data source loaded, loading rollout functions...") self.generate_rollout = load_function(self.args.rollout_function_path) @@ -99,18 +118,15 @@ def __init__(self, args, pg): self.all_rollout_engines = [] self.num_new_engines = 0 logger.info("RolloutManager using no sglang engines (train_only).") - else: - num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) - num_engines = args.rollout_num_gpus // num_gpu_per_engine - self.all_rollout_engines = [None] * num_engines - self.num_new_engines = init_rollout_engines(args, pg, self.all_rollout_engines) - logger.info("RolloutManager started %s rollout engines", len(self.all_rollout_engines)) + elif not self._engines_first: + self._launch_engines() # pools outlive every rollout call, so the manager seats them and the rm functions find them as singletons _, bundle_indices, gpu_ids = pg self.reward_slots = ColocatedRewardSlots( bundle_deal_order(bundle_indices, gpu_ids, args.num_gpus_per_node, args.rollout_num_gpus_per_engine) ) - self.reward_pools = create_colocated_reward_pools(args, pg, self.reward_slots) + with st.step("rm.create_reward_pools"): + self.reward_pools = create_colocated_reward_pools(args, pg, self.reward_slots) logger.info("RolloutManager colocated reward slots: %s", self.reward_slots) self.encode_pool = SftEncodePool(args, pg) if args.loss_type == "sft_loss" else None logger.info("RolloutManager: creating lock...") @@ -161,7 +177,32 @@ def rollout_engines(self): # when doing multi-node serving, we will only send request to node-0 for each engine. return self.all_rollout_engines[:: self.nodes_per_engine] + def _launch_engines(self): + args = self.args + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) + num_engines = args.rollout_num_gpus // num_gpu_per_engine + self.all_rollout_engines = [None] * num_engines + # --debug-rollout-only crashes unless engine init is synced (see init_rollout_engines). + defer = _flag("MILES_STARTUP_OVERLAP") and not args.debug_rollout_only + with st.step("rm.init_rollout_engines", num_engines=num_engines, deferred=int(defer)): + self.num_new_engines, self._engine_init_handles = _init_rollout_engines( + args, self.pg, self.all_rollout_engines, defer_init=defer + ) + logger.info("RolloutManager started %s rollout engines", len(self.all_rollout_engines)) + + def _ensure_engines_ready(self): + """Surface deferred engine init errors. Engine calls queue behind init on the engine actors anyway.""" + if self._engine_init_handles: + with st.step("rm.engines.init_wait"): + ray.get(self._engine_init_handles) + self._engine_init_handles = [] + + def wait_engines_ready(self): + self._ensure_engines_ready() + return True + def get_rollout_engines_and_lock(self): + self._ensure_engines_ready() return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines def get_num_rollout_per_epoch(self): @@ -172,6 +213,7 @@ def generate(self, rollout_id, train_in_flight=None): """`train_in_flight`: refs of the train step this generate overlaps; colocated SFT encoders wait for them.""" from miles.dashboard import hooks + self._ensure_engines_ready() start_time = time.time() self.rollout_id = rollout_id if self.encode_pool is not None: @@ -244,7 +286,7 @@ def load(self, rollout_id=None): def offload(self): self.health_monitoring_pause() - with timer("rollout_offload"): + with timer("rollout_offload"), st.step("rm.offload"): return ray.get( [engine.release_memory_occupation.remote() for engine in self.rollout_engines if engine is not None] ) @@ -494,8 +536,14 @@ def set_train_parallel_config(self, config: dict): def init_rollout_engines(args, pg, all_rollout_engines): + num_new_engines, _ = _init_rollout_engines(args, pg, all_rollout_engines, defer_init=False) + return num_new_engines + + +def _init_rollout_engines(args, pg, all_rollout_engines, defer_init): + """Create engines and start init. With defer_init the init handles are returned unawaited.""" if args.train_only: - return 0 + return 0, [] num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine @@ -532,6 +580,8 @@ def init_rollout_engines(args, pg, all_rollout_engines): "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true", "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false", + # The RunAI streamer is copy-bound (~2.4 GiB/s even from a warm page cache); mmap is much faster. + "SGLANG_USE_RUNAI_MODEL_STREAMER": "true" if os.environ.get("MILES_SGLANG_USE_RUNAI_STREAMER") == "1" else "false", } if rollout_merges_lora(args): # Merge in the train forward dtype, not fp32, to cut train/rollout consistency error. @@ -552,7 +602,7 @@ def init_rollout_engines(args, pg, all_rollout_engines): num_new_engines = len(rollout_engines) if num_new_engines == 0: - return num_new_engines + return num_new_engines, [] addr_and_ports = _allocate_rollout_engine_addr_and_ports_normal( args=args, num_engines=num_engines, rollout_engines=rollout_engines @@ -561,9 +611,11 @@ def init_rollout_engines(args, pg, all_rollout_engines): # TODO: don't ray.get here to overlap train actor init with rollout engine init. # somehow if we don't sync here, the --debug-rollout-only mode will crash. init_handles = [engine.init.remote(**(addr_and_ports[rank])) for rank, engine in rollout_engines] + if defer_init: + return num_new_engines, init_handles ray.get(init_handles) - return num_new_engines + return num_new_engines, [] def _allocate_rollout_engine_addr_and_ports_normal(*, args, num_engines, rollout_engines): @@ -636,15 +688,22 @@ def addr(): return addr_and_ports -def _start_router(args): - """start sgl router and miles router""" - if args.sglang_router_ip is not None: +def _pick_router_addr(args): + if not hasattr(args, "_router_external"): + args._router_external = args.sglang_router_ip is not None + if args._router_external or args.sglang_router_port is not None and args.sglang_router_ip is not None: return - args.sglang_router_ip = _wrap_ipv6(get_host_info()[1]) if args.sglang_router_port is None: args.sglang_router_port = find_available_port(random.randint(3000, 4000)) + +def _start_router(args): + """start sgl router and miles router""" + _pick_router_addr(args) + if args._router_external: + return + if args.use_miles_router: from miles.router.router import run_router @@ -658,12 +717,27 @@ def _start_router(args): ) process.daemon = True # Set the process as a daemon process.start() - # Wait 3 seconds - time.sleep(3) + if _flag("MILES_STARTUP_ROUTER_POLL"): + _wait_port_open(args.sglang_router_ip, args.sglang_router_port, process) + else: + time.sleep(3) assert process.is_alive() logger.info(f"Router launched at {args.sglang_router_ip}:{args.sglang_router_port}") +def _wait_port_open(host, port, process, timeout=30.0): + import socket + + deadline = time.time() + timeout + host = host.strip("[]") + while time.time() < deadline and process.is_alive(): + try: + with socket.create_connection((host, int(port)), timeout=0.2): + return + except OSError: + time.sleep(0.05) + + def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): if args.custom_eval_rollout_log_function_path is not None: custom_log_func = load_function(args.custom_eval_rollout_log_function_path) diff --git a/miles/utils/startup_timing.py b/miles/utils/startup_timing.py new file mode 100644 index 000000000..07e6d202d --- /dev/null +++ b/miles/utils/startup_timing.py @@ -0,0 +1,55 @@ +"""Wall-clock probes for the startup path (process launch -> first rollout). + +Every probe prints one machine-parseable line to stdout: + + STARTUP_TIMING role= pid= event= t= [dur=] [k=v ...] + +stdout of the driver and of every ray actor funnels into the ray job log, so +``tools/parse_startup_timing.py`` can rebuild the full cross-process timeline +from a single log file. ``print`` (not logging) keeps the probes independent of +per-process logger configuration; timestamps are embedded so log-line ordering +does not matter. This module must stay import-light (no torch/ray) because it +anchors each process's import-start time. +""" + +from __future__ import annotations + +import os +import time +from contextlib import contextmanager + +# Anchor: time this module was imported in the current process. When imported +# before the heavy imports (see train_diffusion.py), `mark("x.enter")` minus +# this anchor measures the process's import/bootstrap cost. +PROC_T0 = time.time() + +_role = "unknown" + + +def set_role(role: str) -> None: + global _role + _role = role + mark("role_set", since_import=f"{time.time() - PROC_T0:.3f}") + + +def _emit(event: str, t: float, dur: float | None = None, **extra: object) -> None: + parts = [f"STARTUP_TIMING role={_role}", f"pid={os.getpid()}", f"event={event}", f"t={t:.3f}"] + if dur is not None: + parts.append(f"dur={dur:.3f}") + parts.extend(f"{k}={v}" for k, v in extra.items()) + print(" ".join(parts), flush=True) + + +def mark(event: str, **extra: object) -> None: + _emit(event, time.time(), **extra) + + +@contextmanager +def step(event: str, **extra: object): + t0 = time.time() + _emit(f"{event}.begin", t0, **extra) + try: + yield + finally: + t1 = time.time() + _emit(f"{event}.end", t1, dur=t1 - t0, **extra) diff --git a/scripts/parse_startup_timing.py b/scripts/parse_startup_timing.py new file mode 100644 index 000000000..ada0c32a7 --- /dev/null +++ b/scripts/parse_startup_timing.py @@ -0,0 +1,88 @@ +"""Rebuild the startup timeline from a ray job log containing STARTUP_TIMING lines. + +Usage: + python3 scripts/parse_startup_timing.py /path/to/run.log + +Prints (1) a chronological cross-process timeline, (2) a per-step duration +rollup grouped by role, and (3) the driver-phase breakdown with percentages. +Only stdlib; safe to run anywhere. +""" + +import re +import sys +from collections import defaultdict + +LINE_RE = re.compile( + r"STARTUP_TIMING role=(?P\S+) pid=(?P\d+) event=(?P\S+) t=(?P[0-9.]+)(?: dur=(?P[0-9.]+))?(?P.*)" +) + + +def parse(path): + events = [] + with open(path, errors="replace") as f: + for line in f: + m = LINE_RE.search(line) + if m: + events.append( + { + "role": m.group("role"), + "pid": m.group("pid"), + "event": m.group("event"), + "t": float(m.group("t")), + "dur": float(m.group("dur")) if m.group("dur") else None, + "extra": m.group("extra").strip(), + } + ) + return sorted(events, key=lambda e: e["t"]) + + +def fmt_t(t, t0): + return f"+{t - t0:8.2f}s" + + +def main(path): + events = parse(path) + if not events: + print("no STARTUP_TIMING lines found") + return + t0 = events[0]["t"] + + print("=" * 100) + print("CHRONOLOGICAL TIMELINE (t=0 is the first probe)") + print("=" * 100) + for e in events: + dur = f" dur={e['dur']:8.2f}s" if e["dur"] is not None else "" + extra = f" {e['extra']}" if e["extra"] else "" + print(f"{fmt_t(e['t'], t0)} {e['role']:<16} {e['event']:<50}{dur}{extra}") + + print() + print("=" * 100) + print("PER-STEP DURATIONS (max across ranks per role-prefix, sorted desc)") + print("=" * 100) + durs = defaultdict(list) + for e in events: + if e["dur"] is not None and e["event"].endswith(".end"): + durs[e["event"][: -len(".end")]].append((e["role"], e["dur"])) + rows = [] + for name, vals in durs.items(): + worst_role, worst = max(vals, key=lambda rv: rv[1]) + rows.append((worst, name, worst_role, len(vals))) + for worst, name, worst_role, n in sorted(rows, reverse=True): + print(f"{worst:9.2f}s {name:<55} worst={worst_role} n={n}") + + print() + print("=" * 100) + print("DRIVER PHASES") + print("=" * 100) + driver = [e for e in events if e["role"] == "driver" and e["dur"] is not None] + total = sum(e["dur"] for e in driver if e["event"].count(".") == 2) # driver..end only + for e in driver: + if e["event"].count(".") == 2: + name = e["event"][: -len(".end")] + pct = 100 * e["dur"] / total if total else 0 + print(f"{e['dur']:9.2f}s {pct:5.1f}% {name}") + print(f"{total:9.2f}s 100.0% total (sum of driver top-level phases)") + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tests/fast/utils/test_startup_timing_and_gate.py b/tests/fast/utils/test_startup_timing_and_gate.py new file mode 100644 index 000000000..11f9b754c --- /dev/null +++ b/tests/fast/utils/test_startup_timing_and_gate.py @@ -0,0 +1,104 @@ +import socket +import subprocess +import sys +import textwrap +import threading +import time + +import pytest +import torch + +import miles.utils.startup_timing as st +from miles.backends.fsdp_utils import startup_gate +from miles.ray import rollout + + +def test_required_bytes_is_shard_plus_largest_plus_margin(): + model = torch.nn.Sequential(torch.nn.Linear(10, 10, bias=False), torch.nn.Linear(10, 40, bias=False)) + total = (100 + 400) * 4 + shard, largest, need = startup_gate.required_bytes(model, shard_world_size=2, margin_bytes=1000) + assert shard == total // 2 + assert largest == 400 * 4 + assert need == shard + largest + 1000 + + +def test_required_bytes_works_on_meta_params(): + with torch.device("meta"): + model = torch.nn.Linear(8, 8, bias=False) + assert startup_gate.required_bytes(model, 1, 0)[2] == 8 * 8 * 4 * 2 + + +def test_budget_gate_flag(monkeypatch): + monkeypatch.delenv("MILES_STARTUP_BUDGET_GATE", raising=False) + assert not startup_gate.budget_gate_enabled() + monkeypatch.setenv("MILES_STARTUP_BUDGET_GATE", "1") + assert startup_gate.budget_gate_enabled() + + +def test_gate_times_out_when_budget_never_fits(monkeypatch): + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (1, 2), raising=False) + monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") + with pytest.raises(RuntimeError, match="startup gate"): + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1) + + +def test_gate_passes_when_budget_fits(monkeypatch): + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (1 << 40, 1 << 41), raising=False) + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1) + + +def test_wait_port_open_returns_once_listening(): + srv = socket.socket() + srv.bind(("127.0.0.1", 0)) + port = srv.getsockname()[1] + threading.Timer(0.3, srv.listen).start() + + class Alive: + def is_alive(self): + return True + + t0 = time.time() + rollout._wait_port_open("127.0.0.1", port, Alive(), timeout=5) + assert 0.2 < time.time() - t0 < 3 + srv.close() + + +def test_wait_port_open_gives_up_when_process_dies(): + class Dead: + def is_alive(self): + return False + + t0 = time.time() + rollout._wait_port_open("127.0.0.1", 9, Dead(), timeout=5) + assert time.time() - t0 < 1 + + +def test_pick_router_addr_respects_external_router(): + args = type("A", (), {"sglang_router_ip": "10.0.0.1", "sglang_router_port": 1234})() + rollout._pick_router_addr(args) + assert (args.sglang_router_ip, args.sglang_router_port) == ("10.0.0.1", 1234) + assert args._router_external + + +def test_probe_lines_round_trip_through_parser(tmp_path): + log = tmp_path / "run.log" + script = textwrap.dedent( + """ + import miles.utils.startup_timing as st + st.set_role("driver") + with st.step("driver.x", k=1): + pass + st.mark("driver.startup_done") + """ + ) + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True).stdout + log.write_text(out) + import importlib.util + + spec = importlib.util.spec_from_file_location("parse_startup_timing", "scripts/parse_startup_timing.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + events = mod.parse(str(log)) + names = [e["event"] for e in events] + assert "driver.x.begin" in names and "driver.x.end" in names and "driver.startup_done" in names + assert next(e for e in events if e["event"] == "driver.x.end")["dur"] >= 0 diff --git a/train_diffusion.py b/train_diffusion.py index 3e69b372b..343f4e700 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -1,8 +1,10 @@ import logging +import os import sys import ray +import miles.utils.startup_timing as st # anchors PROC_T0 before the heavy import chain from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from miles.utils.arguments import parse_args from miles.utils.logging_utils import configure_logger @@ -15,25 +17,36 @@ def train(args): logger = logging.getLogger(__name__) # allocate the GPUs logger.info("train: creating placement groups") - pgs = create_placement_groups(args) + with st.step("driver.create_placement_groups"): + pgs = create_placement_groups(args) logger.info("train: placement groups ready") - init_tracking(args) + with st.step("driver.init_tracking"): + init_tracking(args) # create the rollout manager, with sglang engines inside. # need to initialize rollout manager first to calculate num_rollout logger.info("train: creating rollout manager") - rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + with st.step("driver.create_rollout_manager"): + rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) logger.info("train: rollout manager ready") logger.info("train: creating training model") - actor_model = create_training_models(args, pgs, rollout_manager) + with st.step("driver.create_training_models"): + actor_model = create_training_models(args, pgs, rollout_manager) logger.info("train: training model ready") - if args.offload_rollout: - ray.get(rollout_manager.onload_weights.remote()) + if args.offload_rollout and os.environ.get("MILES_STARTUP_SKIP_ROUNDTRIP") != "1": + with st.step("driver.onload_rollout_weights"): + ray.get(rollout_manager.onload_weights.remote()) # always update weight first so that sglang has the loaded weights from training. - actor_model.update_weights() + with st.step("driver.first_update_weights"): + actor_model.update_weights() + st.mark("driver.startup_done") + if os.environ.get("MILES_STARTUP_BENCH_EXIT") == "1": + # Startup benchmark: everything up to the first train step is done once weights are synced. + logger.info("train: MILES_STARTUP_BENCH_EXIT=1, exiting after startup") + return # special case for eval-only if args.num_rollout == 0 and args.eval_interval is not None: @@ -88,5 +101,8 @@ def save(rollout_id): if __name__ == "__main__": # Ensure stdout is line-buffered so nohup logs show progress immediately. sys.stdout.reconfigure(line_buffering=True) - args = parse_args() + st.set_role("driver") + st.mark("driver.main_enter") + with st.step("driver.parse_args"): + args = parse_args() train(args) From 900c755c36a9ca5c1558410401a1bc19c4f0f06f Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 18:49:58 +0530 Subject: [PATCH 2/5] perf: overlap engine boot and train init by default Engine init is no longer awaited by the rollout manager, the startup offload/onload round trip is skipped, and colocated train ranks gate their first large GPU allocation on engine init plus a memory budget (shard + largest tensor + margin). If the budget does not fit next to the resident engines, the engines are offloaded once and onloaded again before the first weight sync. --serial-startup restores the old order. Drops the engines-first and router polling experiments, which did not measure. Bakes .pyc files into the image. --- docker/Dockerfile | 9 ++ miles/backends/fsdp_utils/actor.py | 2 + miles/backends/fsdp_utils/model_loader.py | 4 +- miles/backends/fsdp_utils/startup_gate.py | 72 ++++++++----- miles/ray/placement_group.py | 15 ++- miles/ray/rollout.py | 66 +++++------- miles/utils/arguments.py | 9 ++ .../utils/test_startup_timing_and_gate.py | 100 +++++++++--------- train_diffusion.py | 7 +- 9 files changed, 152 insertions(+), 132 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 5af4ce118..3d5edf775 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -88,3 +88,12 @@ RUN git clone https://github.com/radixark/miles_diffusion.git /root/miles_diffus # Pre-download PaddleOCR (en) det+rec+cls models so the OCR-reward actors don't # race-download them from bcebos at runtime. RUN python -c "from paddleocr import PaddleOCR; PaddleOCR(use_angle_cls=False, lang='en', use_gpu=False, show_log=False)" + +# ====================================== Bytecode ====================================================== + +# Bake .pyc files. About half of the package dirs ship without __pycache__, so every fresh container +# (every CI job) recompiles bytecode on its first import. Measured on a fresh container: +# `import sglang.multimodal_gen.runtime.launch_server` took 65 s without this step and 50 s with it +# (38 s once warm). Some vendored files are not valid Python 3 source, so compile errors are ignored. +RUN python -m compileall -q -j0 "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \ + /sgl-workspace/sglang/python /root/miles_diffusion || true diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 264ecfefb..59a8bd22d 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -130,6 +130,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag checkpoint_path=args.ref_load, lora_adapter_paths=args.ref_lora_adapter_paths, cpu_offload=args.ref_cpu_offload, + rollout_manager=rollout_manager, ) self.teacher_models = {} if args.teacher_load is not None: @@ -141,6 +142,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag checkpoint_path=args.teacher_load, lora_adapter_paths=args.teacher_lora_adapter_paths, cpu_offload=args.teacher_cpu_offload, + rollout_manager=rollout_manager, ) # A natively CPU-offloaded model never sleeps: its parameters stay on the host and its buffers on the GPU. self.actor_models_to_sleep = [] if args.fsdp_cpu_offload else list(self.models.values()) diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index 569cb4bff..afac4fc41 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -10,7 +10,7 @@ from . import checkpoint from .device_move import move_model -from .startup_gate import budget_gate_enabled, wait_for_gpu_budget +from .startup_gate import wait_for_gpu_budget from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str from .sequence_parallel.plan import apply_sequence_parallel @@ -56,7 +56,7 @@ def load_fsdp_models( raise RuntimeError(f"{component} did not honor meta initialization") checkpoint.sync_model_dtypes(model) full_state = model.state_dict() if materialize_weights else {} - if trainable and budget_gate_enabled(): + if rollout_manager is not None: wait_for_gpu_budget( model, shard_world_size=parallel_state.get_mesh("fsdp").size(), rollout_manager=rollout_manager ) diff --git a/miles/backends/fsdp_utils/startup_gate.py b/miles/backends/fsdp_utils/startup_gate.py index f62360453..10b3ce5c0 100644 --- a/miles/backends/fsdp_utils/startup_gate.py +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -1,9 +1,11 @@ """Hold the train side's first large GPU allocation until the colocated engines leave enough room. -Replaces a phase-based gate ("wait until the engines booted and offloaded") with a memory budget: -the train GPU phase only needs one parameter shard plus the largest full tensor (the broadcast -from rank 0 moves one tensor at a time) plus a margin. Sizes come from the module's parameter -shapes, so no weights are read. +Startup overlaps train init with engine boot, and the engines stay resident through it. Before +its first large GPU allocation a colocated train rank waits for engine init (so it cannot grab +memory the engine load is about to need), then checks a memory budget: one parameter shard plus +the largest full tensor (the broadcast from rank 0 moves one tensor at a time) plus a margin. +If that does not fit next to the engines, the engines are offloaded once, as the serial startup +did, and the driver onloads them again before the first weight sync. """ from __future__ import annotations @@ -12,6 +14,7 @@ import os import time +import ray import torch import torch.distributed as dist @@ -22,10 +25,6 @@ GIB = 1 << 30 -def budget_gate_enabled() -> bool: - return os.environ.get("MILES_STARTUP_BUDGET_GATE") == "1" - - def required_bytes(model: torch.nn.Module, shard_world_size: int, margin_bytes: int) -> tuple[int, int, int]: total = 0 largest = 0 @@ -37,9 +36,16 @@ def required_bytes(model: torch.nn.Module, shard_world_size: int, margin_bytes: return shard, largest, shard + largest + margin_bytes -def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollout_manager=None) -> None: - """Block until free GPU memory covers the init peak; with MILES_STARTUP_GATE_REQUIRE_ENGINES=1 also - wait for engine init to finish (guards the race where train allocates before the engines load).""" +def _wait_free(need: int, timeout: float) -> int: + start = time.time() + while True: + free = torch.cuda.mem_get_info()[0] + if free >= need or time.time() - start > timeout: + return free + time.sleep(0.25) + + +def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollout_manager) -> None: margin = int(float(os.environ.get("MILES_STARTUP_GATE_MARGIN_GIB", "4")) * GIB) timeout = float(os.environ.get("MILES_STARTUP_GATE_TIMEOUT_S", "900")) shard, largest, need = required_bytes(model, shard_world_size, margin) @@ -49,22 +55,34 @@ def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollou shard_gib=f"{shard / GIB:.1f}", largest_gib=f"{largest / GIB:.2f}", ): - if os.environ.get("MILES_STARTUP_GATE_REQUIRE_ENGINES") == "1" and rollout_manager is not None: - import ray - + with st.step("train.budget_gate.wait_engines"): ray.get(rollout_manager.wait_engines_ready.remote()) - start = time.time() - while True: - free = torch.cuda.mem_get_info()[0] - if free >= need: - break - if time.time() - start > timeout: + free = torch.cuda.mem_get_info()[0] + fits = free >= need + # Every rank must agree, or one rank offloads while another allocates. + if dist.is_initialized(): + flag = torch.tensor([int(fits)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=_gloo_group()) + fits = bool(flag.item()) + if not fits: + logger.info( + "startup gate: %.1f GiB free next to the engines, need %.1f GiB; offloading the engines", + free / GIB, + need / GIB, + ) + with st.step("train.budget_gate.offload_engines"): + ray.get(rollout_manager.boot_offload_once.remote()) + free = _wait_free(need, timeout) + if free < need: raise RuntimeError( - f"startup gate: only {free / GIB:.1f} GiB free after {timeout:.0f}s, need {need / GIB:.1f} GiB " - f"(shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} + margin {margin / GIB:.1f})" + f"startup gate: only {free / GIB:.1f} GiB free after offloading the engines, need " + f"{need / GIB:.1f} GiB (shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} " + f"+ margin {margin / GIB:.1f})" ) - time.sleep(0.25) - st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}") - if dist.is_initialized(): - # Ranks gate on their own device; keep them in step so the broadcast does not start half-gated. - dist.barrier() + st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}", engines_resident=int(fits)) + + +def _gloo_group(): + from miles.utils.distributed_utils import get_gloo_group + + return get_gloo_group() diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b9e7c4465..feb4eeeda 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -1,5 +1,4 @@ import logging -import os import socket import ray @@ -143,7 +142,8 @@ def create_training_models(args, pgs, rollout_manager): pg=pgs["actor"], ) with st.step("driver.train_group.init_wait"): - gate_manager = rollout_manager if os.environ.get("MILES_STARTUP_GATE_REQUIRE_ENGINES") == "1" else None + # Colocated train ranks gate their GPU phase on the engines (see startup_gate.py). + gate_manager = rollout_manager if args.offload_rollout and not args.serial_startup else None start_rollout_ids = ray.get( actor_model.async_init(args, role="actor", with_ref=False, rollout_manager=gate_manager) ) @@ -177,13 +177,10 @@ def create_rollout_manager(args, pg): assert args.num_rollout > 0 logger.info("Computed num_rollout=%s (num_rollout_per_epoch=%s)", args.num_rollout, num_rollout_per_epoch) - if args.offload_rollout and os.environ.get("MILES_STARTUP_SKIP_ROUNDTRIP") != "1": - if os.environ.get("MILES_STARTUP_OVERLAP") == "1": - # Don't block the driver: train init runs while engines boot and offload. - rollout_manager.offload.remote() - else: - with st.step("driver.rollout_manager.boot_offload"): - ray.get(rollout_manager.offload.remote()) + # With overlapped startup the engines stay resident; the train gate offloads them only if needed. + if args.offload_rollout and args.serial_startup: + with st.step("driver.rollout_manager.boot_offload"): + ray.get(rollout_manager.offload.remote()) logger.info("Rollout manager created.") return rollout_manager, num_rollout_per_epoch diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index d94053725..677544c1c 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -47,11 +47,6 @@ logger = logging.getLogger(__name__) -def _flag(name: str) -> bool: - # Startup-benchmark switches; each startup optimization is toggled independently. - return os.environ.get(name) == "1" - - @ray.remote class RolloutManager: """The class to run rollout and convert rollout data to training data.""" @@ -67,11 +62,7 @@ def __init__(self, args, pg): from miles.dashboard import hooks self._engine_init_handles = [] - self._engines_first = _flag("MILES_STARTUP_ENGINES_FIRST") and not args.train_only and not args.debug_rollout_only - if self._engines_first: - # Engines only need the router address, and add_worker once healthy (35+ s later): boot them first. - _pick_router_addr(args) - self._launch_engines() + self._boot_offloaded = False with st.step("rm.dashboard_register"): hooks.register_rollout_manager(args) if not args.train_only: @@ -118,7 +109,7 @@ def __init__(self, args, pg): self.all_rollout_engines = [] self.num_new_engines = 0 logger.info("RolloutManager using no sglang engines (train_only).") - elif not self._engines_first: + else: self._launch_engines() # pools outlive every rollout call, so the manager seats them and the rm functions find them as singletons _, bundle_indices, gpu_ids = pg @@ -182,8 +173,8 @@ def _launch_engines(self): num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine self.all_rollout_engines = [None] * num_engines - # --debug-rollout-only crashes unless engine init is synced (see init_rollout_engines). - defer = _flag("MILES_STARTUP_OVERLAP") and not args.debug_rollout_only + # Defer so train init overlaps engine boot. --debug-rollout-only crashes unless init is synced. + defer = not args.serial_startup and not args.debug_rollout_only with st.step("rm.init_rollout_engines", num_engines=num_engines, deferred=int(defer)): self.num_new_engines, self._engine_init_handles = _init_rollout_engines( args, self.pg, self.all_rollout_engines, defer_init=defer @@ -201,6 +192,21 @@ def wait_engines_ready(self): self._ensure_engines_ready() return True + def boot_offload_once(self): + """Startup fallback when the train init peak does not fit next to the resident engines. + + Train ranks call this concurrently; the actor runs calls one at a time, so only the first offloads. + """ + if not self._boot_offloaded: + self._boot_offloaded = True + self.offload() + return True + + def onload_weights_if_boot_offloaded(self): + if self._boot_offloaded: + self.onload_weights() + return self._boot_offloaded + def get_rollout_engines_and_lock(self): self._ensure_engines_ready() return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines @@ -688,22 +694,15 @@ def addr(): return addr_and_ports -def _pick_router_addr(args): - if not hasattr(args, "_router_external"): - args._router_external = args.sglang_router_ip is not None - if args._router_external or args.sglang_router_port is not None and args.sglang_router_ip is not None: +def _start_router(args): + """start sgl router and miles router""" + if args.sglang_router_ip is not None: return + args.sglang_router_ip = _wrap_ipv6(get_host_info()[1]) if args.sglang_router_port is None: args.sglang_router_port = find_available_port(random.randint(3000, 4000)) - -def _start_router(args): - """start sgl router and miles router""" - _pick_router_addr(args) - if args._router_external: - return - if args.use_miles_router: from miles.router.router import run_router @@ -717,27 +716,12 @@ def _start_router(args): ) process.daemon = True # Set the process as a daemon process.start() - if _flag("MILES_STARTUP_ROUTER_POLL"): - _wait_port_open(args.sglang_router_ip, args.sglang_router_port, process) - else: - time.sleep(3) + # Wait 3 seconds + time.sleep(3) assert process.is_alive() logger.info(f"Router launched at {args.sglang_router_ip}:{args.sglang_router_port}") -def _wait_port_open(host, port, process, timeout=30.0): - import socket - - deadline = time.time() + timeout - host = host.strip("[]") - while time.time() < deadline and process.is_alive(): - try: - with socket.create_connection((host, int(port)), timeout=0.2): - return - except OSError: - time.sleep(0.05) - - def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): if args.custom_eval_rollout_log_function_path is not None: custom_log_func = load_function(args.custom_eval_rollout_log_function_path) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4390b123f..91dd9b704 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -108,6 +108,15 @@ def add_cluster_arguments(parser): "This will always be true when --colocate is set." ), ) + parser.add_argument( + "--serial-startup", + action="store_true", + default=False, + help=( + "Boot the rollout engines and the train actors one after the other, as before, " + "instead of overlapping them. Only for debugging startup." + ), + ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/tests/fast/utils/test_startup_timing_and_gate.py b/tests/fast/utils/test_startup_timing_and_gate.py index 11f9b754c..b5a38be93 100644 --- a/tests/fast/utils/test_startup_timing_and_gate.py +++ b/tests/fast/utils/test_startup_timing_and_gate.py @@ -1,16 +1,44 @@ -import socket import subprocess import sys import textwrap -import threading -import time import pytest import torch -import miles.utils.startup_timing as st from miles.backends.fsdp_utils import startup_gate -from miles.ray import rollout + + +class _Remote: + def __init__(self, fn): + self.fn = fn + + def remote(self): + return self.fn() + + +class FakeRolloutManager: + def __init__(self, on_offload=None): + self.calls = [] + self.on_offload = on_offload + self.wait_engines_ready = _Remote(lambda: self.calls.append("wait") or True) + self.boot_offload_once = _Remote(self._offload) + + def _offload(self): + self.calls.append("offload") + if self.on_offload: + self.on_offload() + return True + + +@pytest.fixture +def no_ray(monkeypatch): + monkeypatch.setattr(startup_gate.ray, "get", lambda x: x) + + +def _free(monkeypatch, value): + state = {"free": value} + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (state["free"], 1 << 41), raising=False) + return state def test_required_bytes_is_shard_plus_largest_plus_margin(): @@ -28,56 +56,26 @@ def test_required_bytes_works_on_meta_params(): assert startup_gate.required_bytes(model, 1, 0)[2] == 8 * 8 * 4 * 2 -def test_budget_gate_flag(monkeypatch): - monkeypatch.delenv("MILES_STARTUP_BUDGET_GATE", raising=False) - assert not startup_gate.budget_gate_enabled() - monkeypatch.setenv("MILES_STARTUP_BUDGET_GATE", "1") - assert startup_gate.budget_gate_enabled() +def test_gate_keeps_engines_resident_when_it_fits(monkeypatch, no_ray): + _free(monkeypatch, 1 << 40) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + assert rm.calls == ["wait"] -def test_gate_times_out_when_budget_never_fits(monkeypatch): - monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (1, 2), raising=False) - monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") - with pytest.raises(RuntimeError, match="startup gate"): - startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1) +def test_gate_offloads_engines_when_it_does_not_fit(monkeypatch, no_ray): + state = _free(monkeypatch, 1) + rm = FakeRolloutManager(on_offload=lambda: state.update(free=1 << 40)) + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + assert rm.calls == ["wait", "offload"] -def test_gate_passes_when_budget_fits(monkeypatch): - monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (1 << 40, 1 << 41), raising=False) - startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1) - - -def test_wait_port_open_returns_once_listening(): - srv = socket.socket() - srv.bind(("127.0.0.1", 0)) - port = srv.getsockname()[1] - threading.Timer(0.3, srv.listen).start() - - class Alive: - def is_alive(self): - return True - - t0 = time.time() - rollout._wait_port_open("127.0.0.1", port, Alive(), timeout=5) - assert 0.2 < time.time() - t0 < 3 - srv.close() - - -def test_wait_port_open_gives_up_when_process_dies(): - class Dead: - def is_alive(self): - return False - - t0 = time.time() - rollout._wait_port_open("127.0.0.1", 9, Dead(), timeout=5) - assert time.time() - t0 < 1 - - -def test_pick_router_addr_respects_external_router(): - args = type("A", (), {"sglang_router_ip": "10.0.0.1", "sglang_router_port": 1234})() - rollout._pick_router_addr(args) - assert (args.sglang_router_ip, args.sglang_router_port) == ("10.0.0.1", 1234) - assert args._router_external +def test_gate_raises_when_offload_does_not_free_enough(monkeypatch, no_ray): + _free(monkeypatch, 1) + monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") + rm = FakeRolloutManager() + with pytest.raises(RuntimeError, match="startup gate"): + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) def test_probe_lines_round_trip_through_parser(tmp_path): diff --git a/train_diffusion.py b/train_diffusion.py index 343f4e700..821b12c16 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -35,9 +35,12 @@ def train(args): actor_model = create_training_models(args, pgs, rollout_manager) logger.info("train: training model ready") - if args.offload_rollout and os.environ.get("MILES_STARTUP_SKIP_ROUNDTRIP") != "1": + if args.offload_rollout: with st.step("driver.onload_rollout_weights"): - ray.get(rollout_manager.onload_weights.remote()) + if args.serial_startup: + ray.get(rollout_manager.onload_weights.remote()) + else: + ray.get(rollout_manager.onload_weights_if_boot_offloaded.remote()) # always update weight first so that sglang has the loaded weights from training. with st.step("driver.first_update_weights"): From 53a24fddfc316ab864871c4f3b5ef89bdfc48d17 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 19:51:11 +0530 Subject: [PATCH 3/5] perf: lighter actor imports, optional engine memory reserve for the startup gate The rollout manager and engine actor processes no longer import sglang at module load: GPU_MEMORY_TYPE_WEIGHTS is a local constant, kill_process_tree and ServerArgs are imported where used. Importing miles.ray.rollout drops from 29 s to 5.5 s and the engine actor module from 23 s to 1.4 s, both on the startup critical path. --startup-engine-reserve-gib lets colocated train ranks start their GPU load before engine init finishes when free memory minus that reserve covers the init peak. Unset keeps waiting for the engines. Also stubs the startup probes in test_ema_optimizer, which execs actor methods in a hand-built namespace. --- miles/backends/fsdp_utils/model_loader.py | 5 +++- miles/backends/fsdp_utils/startup_gate.py | 28 ++++++++++++++----- .../sglang_diffusion_utils/arguments.py | 5 +++- .../sglang_diffusion_engine.py | 3 +- miles/ray/rollout.py | 4 ++- miles/utils/arguments.py | 10 +++++++ .../backends/fsdp_utils/test_ema_optimizer.py | 1 + .../utils/test_startup_timing_and_gate.py | 18 ++++++++++++ 8 files changed, 63 insertions(+), 11 deletions(-) diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index afac4fc41..2d7357f2d 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -58,7 +58,10 @@ def load_fsdp_models( full_state = model.state_dict() if materialize_weights else {} if rollout_manager is not None: wait_for_gpu_budget( - model, shard_world_size=parallel_state.get_mesh("fsdp").size(), rollout_manager=rollout_manager + model, + shard_world_size=parallel_state.get_mesh("fsdp").size(), + rollout_manager=rollout_manager, + engine_reserve_gib=args.startup_engine_reserve_gib, ) model = apply_fsdp2( model, diff --git a/miles/backends/fsdp_utils/startup_gate.py b/miles/backends/fsdp_utils/startup_gate.py index 10b3ce5c0..1e7cde099 100644 --- a/miles/backends/fsdp_utils/startup_gate.py +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -45,7 +45,18 @@ def _wait_free(need: int, timeout: float) -> int: time.sleep(0.25) -def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollout_manager) -> None: +def _all_ranks_agree(ok: bool) -> bool: + # Every rank must take the same branch, or one rank offloads while another allocates. + if not dist.is_initialized(): + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=_gloo_group()) + return bool(flag.item()) + + +def wait_for_gpu_budget( + model: torch.nn.Module, *, shard_world_size: int, rollout_manager, engine_reserve_gib: float | None = None +) -> None: margin = int(float(os.environ.get("MILES_STARTUP_GATE_MARGIN_GIB", "4")) * GIB) timeout = float(os.environ.get("MILES_STARTUP_GATE_TIMEOUT_S", "900")) shard, largest, need = required_bytes(model, shard_world_size, margin) @@ -55,15 +66,18 @@ def wait_for_gpu_budget(model: torch.nn.Module, *, shard_world_size: int, rollou shard_gib=f"{shard / GIB:.1f}", largest_gib=f"{largest / GIB:.2f}", ): + if engine_reserve_gib is not None: + # Free memory already excludes whatever the engine has allocated, so subtracting the full reserve + # can only over-count: if this fits, the engine's remaining load fits too. + free = torch.cuda.mem_get_info()[0] + reserve = int(engine_reserve_gib * GIB) + if _all_ranks_agree(free - reserve >= need): + st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}", engines_resident=-1) + return with st.step("train.budget_gate.wait_engines"): ray.get(rollout_manager.wait_engines_ready.remote()) free = torch.cuda.mem_get_info()[0] - fits = free >= need - # Every rank must agree, or one rank offloads while another allocates. - if dist.is_initialized(): - flag = torch.tensor([int(fits)], dtype=torch.int32) - dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=_gloo_group()) - fits = bool(flag.item()) + fits = _all_ranks_agree(free >= need) if not fits: logger.info( "startup gate: %.1f GiB free next to the engines, need %.1f GiB; offloading the engines", diff --git a/miles/backends/sglang_diffusion_utils/arguments.py b/miles/backends/sglang_diffusion_utils/arguments.py index 3eca491ca..90b7354ae 100644 --- a/miles/backends/sglang_diffusion_utils/arguments.py +++ b/miles/backends/sglang_diffusion_utils/arguments.py @@ -1,4 +1,3 @@ -from sglang.multimodal_gen.runtime.server_args import ServerArgs from miles.utils.http_utils import _wrap_ipv6 @@ -14,6 +13,10 @@ def add_sglang_diffusion_arguments(parser): """ Add arguments to the parser for the SGLang server. """ + # Imported here: sglang.multimodal_gen pulls in all of sglang (~14 s), and this module is imported by + # every process that touches miles.utils.arguments, including ones that never parse arguments. + from sglang.multimodal_gen.runtime.server_args import ServerArgs + parser = add_sglang_diffusion_router_arguments(parser) parser.add_argument("--sglang-server-concurrency", type=int, default=512) diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 92f3ae53b..f61701039 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING import requests -from sglang.srt.utils.common import kill_process_tree import miles.utils.startup_timing as st from miles.ray.ray_actor import RayActor @@ -360,6 +359,8 @@ def shutdown(self): if response is not None: response.raise_for_status() + from sglang.srt.utils.common import kill_process_tree + kill_process_tree(self.process.pid) def release_memory_occupation(self): diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 677544c1c..3e2f8076a 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -12,7 +12,6 @@ import ray import torch from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS import miles.utils.startup_timing as st from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine @@ -46,6 +45,9 @@ logger = logging.getLogger(__name__) +# sglang.srt.constants.GPU_MEMORY_TYPE_WEIGHTS; importing it pulls in all of sglang (~10 s) at actor start. +GPU_MEMORY_TYPE_WEIGHTS = "weights" + @ray.remote class RolloutManager: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 91dd9b704..1737f64da 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -117,6 +117,16 @@ def add_cluster_arguments(parser): "instead of overlapping them. Only for debugging startup." ), ) + parser.add_argument( + "--startup-engine-reserve-gib", + type=float, + default=None, + help=( + "Upper bound on the GPU memory one colocated rollout engine rank holds after boot. When set, " + "train ranks start their GPU load as soon as it fits next to that reserve instead of waiting " + "for engine init to finish. Unset: wait for engine init." + ), + ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/tests/fast/backends/fsdp_utils/test_ema_optimizer.py b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py index 5c68a7732..cf1b4f0d4 100644 --- a/tests/fast/backends/fsdp_utils/test_ema_optimizer.py +++ b/tests/fast/backends/fsdp_utils/test_ema_optimizer.py @@ -153,6 +153,7 @@ def test_actor_train_updates_ema_once_and_publication_only_reads_it(): nullcontext=nullcontext, clear_memory=lambda: None, train_metric_utils=SimpleNamespace(log_perf_data_raw=lambda **kwargs: None), + st=SimpleNamespace(step=lambda *args, **kwargs: nullcontext(), mark=lambda *args, **kwargs: None), ) exec(compile(ast.Module(body=methods, type_ignores=[]), str(path), "exec"), namespace) diff --git a/tests/fast/utils/test_startup_timing_and_gate.py b/tests/fast/utils/test_startup_timing_and_gate.py index b5a38be93..d73803672 100644 --- a/tests/fast/utils/test_startup_timing_and_gate.py +++ b/tests/fast/utils/test_startup_timing_and_gate.py @@ -100,3 +100,21 @@ def test_probe_lines_round_trip_through_parser(tmp_path): names = [e["event"] for e in events] assert "driver.x.begin" in names and "driver.x.end" in names and "driver.startup_done" in names assert next(e for e in events if e["event"] == "driver.x.end")["dur"] >= 0 + + +def test_reserve_lets_train_go_before_engines_when_it_fits(monkeypatch, no_ray): + _free(monkeypatch, 100 << 30) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=60 + ) + assert rm.calls == [] + + +def test_reserve_falls_back_to_waiting_for_engines_when_it_does_not_fit(monkeypatch, no_ray): + _free(monkeypatch, 100 << 30) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=99 + ) + assert rm.calls == ["wait"] From c95ce79f57eadf61838cac425b8a28c5aa3f3ad6 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 20:43:52 +0530 Subject: [PATCH 4/5] perf: build sglang ServerArgs in the server process, not the engine actor The launch_server child imports sglang anyway; building ServerArgs there keeps sglang out of the engine actor, which saves one full sglang import on the engine startup path. The actor passes a plain-picklable copy of args, the host and the ports. --- .../sglang_diffusion_engine.py | 64 +++++++++++++------ 1 file changed, 44 insertions(+), 20 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index f61701039..6623fffa0 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -63,7 +63,15 @@ def _bootstrap_scheduler_process_timed(spec): bootstrap_scheduler_process(spec) -def _launch_server_target(server_args, apply_rollout_patches: bool = False): +def _launch_server_target(args, host, port, nccl_port, apply_rollout_patches: bool = False): + # ServerArgs is built here, not in the engine actor: this process imports sglang anyway, and keeping + # sglang out of the actor takes one full sglang import (~40 s) off the engine startup path. + from sglang.multimodal_gen.runtime.server_args import ServerArgs + + with st.step("server.build_server_args"): + server_args = ServerArgs.from_kwargs(**_compute_server_args(args, host=host, port=port, nccl_port=nccl_port)) + server_args.host = server_args.host.strip("[]") + # addict.Dict loses its `__frozen` attribute across spawn pickle; rebuild it. import addict @@ -82,23 +90,44 @@ def _launch_server_target(server_args, apply_rollout_patches: bool = False): launch_server(server_args) +def _picklable_args(args): + """Copy of ``args`` that survives the plain pickle of a spawn start (Ray used cloudpickle to get it here).""" + import argparse + import pickle + + kept = {} + for key, value in vars(args).items(): + try: + pickle.dumps(value) + except Exception: + continue + kept[key] = value + return argparse.Namespace(**kept) + + def launch_server_process( - server_args: ServerArgs, + args, + host: str, + port: int, + nccl_port: int, apply_rollout_patches: bool = False, ) -> multiprocessing.Process: + """Start the sglang-d server in a child process and wait until it is healthy. + + ``host`` is bracketed for IPv6 (as in the server URL); the child strips it for binding. + """ # use spawn to avoid potential risks of fork in terms of subthreads or CUDA. multiprocessing.set_start_method("spawn", force=True) - server_args.host = server_args.host.strip("[]") p = multiprocessing.Process( target=_launch_server_target, - args=(server_args, apply_rollout_patches), + args=(_picklable_args(args), host, port, nccl_port, apply_rollout_patches), ) with st.step("engine.spawn_server_process"): p.start() with st.step("engine.wait_server_healthy"): _wait_server_healthy( - base_url=server_args.url(), + base_url=f"http://{host}:{port}", is_process_alive=lambda: p.is_alive(), ) @@ -153,22 +182,14 @@ def _format_v6_uri(addr): host = _format_v6_uri(host) - server_args_dict = _compute_server_args( - self.args, - host=host, - port=port, - nccl_port=nccl_port, - ) + self.node_rank = getattr(self.args, "sglang_node_rank", None) or 0 + self.server_host = host + self.server_port = port + self._nccl_port = nccl_port - self.node_rank = server_args_dict.get("node_rank", 0) - self.server_host = server_args_dict["host"] - self.server_port = server_args_dict["port"] - - self._init_normal(server_args_dict) - - def _init_normal(self, server_args_dict): - from sglang.multimodal_gen.runtime.server_args import ServerArgs + self._init_normal() + def _init_normal(self): logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}") self._pin_to_assigned_gpu() from miles.backends.sglang_diffusion_utils.monkey_patches import ROLLOUT_PATCH_GROUPS_ENV @@ -178,7 +199,10 @@ def _init_normal(self, server_args_dict): os.environ[ROLLOUT_PATCH_GROUPS_ENV] = ",".join(patch_groups) logger.info("Launching sglang-d with rollout patch groups: %s", patch_groups) self.process = launch_server_process( - ServerArgs.from_kwargs(**server_args_dict), + self.args, + host=self.server_host, + port=self.server_port, + nccl_port=self._nccl_port, apply_rollout_patches=bool(patch_groups), ) From 85605b11c34b0badc7ee2af2b2a4bf690705f14a Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 20:43:52 +0530 Subject: [PATCH 5/5] perf: create all train ranks at once Rank 0 used to finish booting (about 50 s, mostly imports) before ranks 1..N were created. Now all ranks start together and get rank 0's rendezvous address before init. --- miles/ray/actor_group.py | 20 +++--- miles/ray/train_actor.py | 5 ++ .../utils/test_startup_timing_and_gate.py | 62 +++++++++++++++++++ 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 0eb24d3fa..3da651a05 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -72,21 +72,23 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl) - # Create worker actors - self._actor_handlers = [] - master_addr, master_port = None, None - for rank in range(world_size): - actor = TrainRayActor.options( + # Create every rank at once so their process start-up overlaps; rank 0 picks the rendezvous address + # in its constructor and the other ranks get it before init (actor calls run in submission order). + self._actor_handlers = [ + TrainRayActor.options( num_cpus=num_gpus_per_actor, num_gpus=num_gpus_per_actor, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=reordered_bundle_indices[rank], ), - ).remote(world_size, rank, master_addr, master_port) - if rank == 0: - master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote()) - self._actor_handlers.append(actor) + ).remote(world_size, rank, None, None) + for rank in range(world_size) + ] + master_addr, master_port = ray.get(self._actor_handlers[0].get_master_addr_and_port.remote()) + ray.get( + [actor.set_master_addr_and_port.remote(master_addr, master_port) for actor in self._actor_handlers[1:]] + ) def async_init(self, args, role, with_ref=False, rollout_manager=None): """ diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index bd5e7f889..983dc3fcf 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -49,6 +49,11 @@ def __init__(self, world_size, rank, master_addr, master_port): os.environ["CUDA_VISIBLE_DEVICES"] = str(local_gpu_id) os.environ["LOCAL_RANK"] = "0" + def set_master_addr_and_port(self, master_addr, master_port): + self.master_addr, self.master_port = master_addr, master_port + os.environ["MASTER_ADDR"] = self.master_addr + os.environ["MASTER_PORT"] = str(self.master_port) + def init(self, args, role, with_ref=False): self.args = args self.role = role diff --git a/tests/fast/utils/test_startup_timing_and_gate.py b/tests/fast/utils/test_startup_timing_and_gate.py index d73803672..edaae1f90 100644 --- a/tests/fast/utils/test_startup_timing_and_gate.py +++ b/tests/fast/utils/test_startup_timing_and_gate.py @@ -118,3 +118,65 @@ def test_reserve_falls_back_to_waiting_for_engines_when_it_does_not_fit(monkeypa torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=99 ) assert rm.calls == ["wait"] + + +def test_picklable_args_drops_only_unpicklable_values(): + import argparse + import pickle + + from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import _picklable_args + + args = argparse.Namespace(hf_checkpoint="m", sglang_tp_size=2, hook=lambda: None) + copy = _picklable_args(args) + assert vars(copy) == {"hf_checkpoint": "m", "sglang_tp_size": 2} + pickle.dumps(copy) + + +def test_engine_module_does_not_import_sglang(): + code = ( + "import sys; import miles.backends.sglang_diffusion_utils.sglang_diffusion_engine, miles.ray.rollout; " + "print('sglang' in sys.modules)" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True).stdout.strip() + assert out == "False" + + +def test_train_ranks_are_created_together_then_get_rank0_address(monkeypatch): + import argparse + from types import SimpleNamespace + + from miles.ray import actor_group + + events = [] + + class FakeActor: + def __init__(self, rank): + self.rank = rank + self.get_master_addr_and_port = SimpleNamespace( + remote=lambda: events.append(("get_addr", self.rank)) or ("10.0.0.1", 29500) + ) + self.set_master_addr_and_port = SimpleNamespace( + remote=lambda addr, port: events.append(("set_addr", self.rank, addr, port)) + ) + + class FakeRemoteClass: + def options(self, **kwargs): + return SimpleNamespace( + remote=lambda world_size, rank, addr, port: events.append(("create", rank, addr)) or FakeActor(rank) + ) + + monkeypatch.setattr(actor_group.ray, "remote", lambda **kwargs: (lambda cls: FakeRemoteClass())) + monkeypatch.setattr(actor_group.ray, "get", lambda x: x) + monkeypatch.setattr(actor_group, "PlacementGroupSchedulingStrategy", lambda **kwargs: None) + args = argparse.Namespace(train_env_vars={}, deterministic_mode=False, train_backend="fsdp") + + actor_group.RayTrainGroup(args, num_nodes=1, num_gpus_per_node=3, pg=(object(), [0, 1, 2], [0, 1, 2])) + + assert events == [ + ("create", 0, None), + ("create", 1, None), + ("create", 2, None), + ("get_addr", 0), + ("set_addr", 1, "10.0.0.1", 29500), + ("set_addr", 2, "10.0.0.1", 29500), + ]