From a7b5f4e5db75504c5364fe9b57efaf16ffbbb127 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sat, 17 Oct 2026 12:09:28 +0530 Subject: [PATCH 1/4] perf: cut startup time to the first train step without changing component interfaces The startup order stays as on main; each step gets cheaper. - Load engine weights with mmap instead of the RunAI streamer. - Keep sglang out of the processes that do not need it: cache sglang-d's CLI argument definitions for the driver (keyed on the installed sglang source, primed at image build), build ServerArgs in the sglang server process instead of the engine actor, defer module-level sglang imports in the rollout manager and engine actor, import the weight-sync helpers where the train actor builds its updater, and resolve fsdp_utils' exports lazily. - Create all train ranks at once; rank 0 still picks the rendezvous address. - Fill FSDP shards straight from the safetensors checkpoint, each rank reading only its dim-0 slice, for parameters that map 1:1 onto it; the rest is broadcast from rank 0 as before. Rank 0 spot-checks the filled slices, ranks agree on the files first, and any failure falls back to the full broadcast. Rank 0 no longer copies its own shard to the GPU before the broadcast. - --rollout-prespawn-schedulers (opt-in): start the sglang-d scheduler workers before the server process imports sglang, using sglang's register_prespawned_schedulers when it exists. - Bake .pyc files into the image. --- docker/Dockerfile | 13 ++ miles/backends/fsdp_utils/__init__.py | 44 ++-- miles/backends/fsdp_utils/actor.py | 18 +- miles/backends/fsdp_utils/checkpoint.py | 67 +++++- miles/backends/fsdp_utils/model_loader.py | 31 ++- .../fsdp_utils/sharded_checkpoint_fill.py | 170 +++++++++++++++ .../sglang_diffusion_utils/arguments.py | 17 +- .../sglang_diffusion_utils/cli_cache.py | 201 ++++++++++++++++++ .../scheduler_prespawn.py | 73 +++++++ .../sglang_diffusion_engine.py | 89 +++++--- miles/ray/actor_group.py | 20 +- miles/ray/rollout.py | 6 +- miles/ray/train_actor.py | 5 + miles/utils/arguments.py | 9 + .../fsdp_utils/test_broadcast_full_state.py | 186 ++++++++++++++++ tests/fast/utils/prespawn_fake_target.py | 8 + tests/fast/utils/test_startup.py | 155 ++++++++++++++ 17 files changed, 1027 insertions(+), 85 deletions(-) create mode 100644 miles/backends/fsdp_utils/sharded_checkpoint_fill.py create mode 100644 miles/backends/sglang_diffusion_utils/cli_cache.py create mode 100644 miles/backends/sglang_diffusion_utils/scheduler_prespawn.py create mode 100644 tests/fast/backends/fsdp_utils/test_broadcast_full_state.py create mode 100644 tests/fast/utils/prespawn_fake_target.py create mode 100644 tests/fast/utils/test_startup.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 5af4ce118..15f9e2cc3 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -88,3 +88,16 @@ 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 + +# Cache sglang-d's CLI argument definitions so parsing miles args does not import sglang (~25 s per +# launch). The cache is keyed on sglang's server_args.py, so a different sglang just rebuilds it. +RUN cd /root/miles_diffusion && python -m miles.backends.sglang_diffusion_utils.cli_cache || true diff --git a/miles/backends/fsdp_utils/__init__.py b/miles/backends/fsdp_utils/__init__.py index 05099e93c..12ac5ca52 100644 --- a/miles/backends/fsdp_utils/__init__.py +++ b/miles/backends/fsdp_utils/__init__.py @@ -1,33 +1,23 @@ import logging - -try: - _TORCH_MEMORY_SAVER_AVAILABLE = True -except ImportError: - logging.warning("torch_memory_saver is not installed, refer to : https://github.com/fzyzcjy/torch_memory_saver") - _TORCH_MEMORY_SAVER_AVAILABLE = False - -try: - _FSDP_AVAILABLE = True -except ImportError as e: - logging.warning(f"FSDP backend dependencies not available: {e}") - _FSDP_AVAILABLE = False - -if _FSDP_AVAILABLE: - from .actor import FSDPTrainRayActor - from .arguments import load_fsdp_args -else: - - def _raise_import_error(*args, **kwargs): - raise ImportError( - "FSDP backend is not available. " - "Please ensure PyTorch with FSDP2 support is installed. " - "For installation instructions, refer to: https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html" - ) - - FSDPTrainRayActor = _raise_import_error - load_fsdp_args = _raise_import_error +# Resolved lazily: importing the actor pulls in its whole training stack (and sglang for weight +# sync), which processes that only need fsdp_utils.arguments or a helper module do not want. +_TORCH_MEMORY_SAVER_AVAILABLE = True +_FSDP_AVAILABLE = True __all__ = ["load_fsdp_args", "FSDPTrainRayActor"] + +def __getattr__(name): + if name == "FSDPTrainRayActor": + from .actor import FSDPTrainRayActor + + return FSDPTrainRayActor + if name == "load_fsdp_args": + from .arguments import load_fsdp_args + + return load_fsdp_args + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + logging.getLogger().setLevel(logging.WARNING) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d005792c2..403fc8aa9 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -28,13 +28,6 @@ from . import checkpoint from .device_move import move_model, move_optimizer, sleep_frozen_model, wake_up_frozen_model -from .diffusion_update_weight_utils import ( - DiffusionUpdateWeightFromDistributed, - DiffusionUpdateWeightFromTensor, - DiffusionUpdateWeightFromTensorLoRA, - DiffusionUpdateWeightFromTensorLoRAIPC, - DiffusionUpdateWeightLoRADistributed, -) from .ema import EMAOptimizer, reshard_model from .input_dtype_policy import apply_input_dtype_policy from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch @@ -204,6 +197,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty ) checkpoint_payload = checkpoint.load(self) + # Imported here, not at module level: it pulls in sglang (~25 s), which the driver and other + # processes that only need this class or fsdp_utils' helpers should not pay for. Not in a + # background thread either: concurrent imports of transformers are not safe. + from .diffusion_update_weight_utils import ( + DiffusionUpdateWeightFromDistributed, + DiffusionUpdateWeightFromTensor, + DiffusionUpdateWeightFromTensorLoRA, + DiffusionUpdateWeightFromTensorLoRAIPC, + DiffusionUpdateWeightLoRADistributed, + ) + # sglang-d now supports /update_weights_from_tensor (PR #20464). if self.args.train_only: self.weight_updater = None diff --git a/miles/backends/fsdp_utils/checkpoint.py b/miles/backends/fsdp_utils/checkpoint.py index ff88c4d35..b0ba1ddd5 100644 --- a/miles/backends/fsdp_utils/checkpoint.py +++ b/miles/backends/fsdp_utils/checkpoint.py @@ -33,17 +33,31 @@ def broadcast_full_state_to_fsdp( full_state: dict, *, cpu_offload: bool, + device=None, + checkpoint_files: list[str] | None = None, ) -> None: - """Materialize FSDP2 shards from rank 0's full initial state dict.""" - from torch.distributed.checkpoint.state_dict import set_model_state_dict + """Materialize FSDP2 shards from rank 0's full initial state dict. - if dist.get_rank() == 0: - # Rank 0 was sharded on real CPU weights; move them (and real buffers) along. - model.to(device=torch.cuda.current_device(), non_blocking=True) - else: - # to_empty creates tensors on device without initializing memory. - model.to_empty(device=torch.cuda.current_device()) + With ``checkpoint_files``, parameters that map cleanly onto the checkpoint are read by every + rank straight from the files (its own slice only); only the rest are broadcast from rank 0. + """ + from torch.distributed.checkpoint.state_dict import set_model_state_dict + device = torch.cuda.current_device() if device is None else device + # Every parameter is overwritten from full_state below, so no rank needs its local values on the + # device first. Rank 0 used to move its real CPU shard over (a slow pageable copy of 1/N of the + # model); now it allocates like the other ranks and only carries its real buffers across. + real_buffers = dict(model.named_buffers()) if dist.get_rank() == 0 else None + # to_empty creates tensors on device without initializing memory. + model.to_empty(device=device) + if real_buffers is not None: + with torch.no_grad(): + for name, buffer in model.named_buffers(): + buffer.copy_(real_buffers[name], non_blocking=True) + + filled = _fill_from_checkpoint_files(model, full_state, checkpoint_files, device) if checkpoint_files else set() + if filled and dist.get_rank() == 0: + full_state = {k: v for k, v in full_state.items() if k not in filled} set_model_state_dict( model, full_state, @@ -51,6 +65,7 @@ def broadcast_full_state_to_fsdp( full_state_dict=True, cpu_offload=cpu_offload, broadcast_from_rank0=True, + strict=not filled, ), ) # State dict loading does not cover non-persistent buffers (e.g. Wan rope @@ -65,6 +80,42 @@ def broadcast_full_state_to_fsdp( buffer.data = buffer.data.to(torch.cuda.current_device()) +def _fill_from_checkpoint_files(model, full_state: dict, files: list[str], device) -> set[str]: + from . import sharded_checkpoint_fill + + try: + filled = sharded_checkpoint_fill.fill_from_checkpoint(model, files, device=device) + ok = True + except Exception as exc: + logger.warning("Checkpoint fill failed on rank %d (%s); broadcasting everything", dist.get_rank(), exc) + filled, ok = set(), False + # All ranks must agree on what still needs broadcasting; any failure falls back to the full broadcast, + # which also overwrites whatever was filled. + oks = [None] * dist.get_world_size() + dist.all_gather_object(oks, ok) + if not all(oks): + return set() + # Rank 0 holds the loaded weights: spot-check the first and last row of its slice of every filled + # parameter, in case a loader converted a tensor without changing its name or shape. + bad = [] + if dist.get_rank() == 0: + params = dict(model.named_parameters()) + for name in sorted(filled): + local = params[name].to_local() + if local.shape[0] == 0: + continue + first, _ = sharded_checkpoint_fill._local_rows(params[name]) + ref = full_state[name] + rows = (0, local.shape[0] - 1) + if not all(torch.equal(local[r].cpu(), ref[first + r].to(local.dtype)) for r in rows): + bad.append(name) + payload = [bad] + dist.broadcast_object_list(payload, src=0) + if payload[0]: + logger.warning("Checkpoint fill did not match rank 0 for %d parameters; broadcasting them", len(payload[0])) + return filled - set(payload[0]) + + def _is_lora_param_name(name: str) -> bool: return "lora_" in name diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index ca61e39ec..6d64a0529 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import os import warnings from argparse import Namespace from contextlib import contextmanager @@ -61,7 +62,12 @@ def load_fsdp_models( cpu_offload=cpu_offload, args=args, ) - checkpoint.broadcast_full_state_to_fsdp(model, full_state, cpu_offload=cpu_offload) + checkpoint.broadcast_full_state_to_fsdp( + model, + full_state, + cpu_offload=cpu_offload, + checkpoint_files=_checkpoint_files(checkpoint_path, component), + ) del full_state train_pipeline_config.postprocess_model_after_materialize(model) if parallel_state.get_optional_mesh("sp") is not None: @@ -79,6 +85,29 @@ def load_fsdp_models( return models +def _checkpoint_files(checkpoint_path: str, component: str) -> list[str] | None: + """Files for a sharded fill, or None; every rank gets the same answer.""" + from . import sharded_checkpoint_fill + + files = None + if sharded_checkpoint_fill.enabled(): + try: + files = sharded_checkpoint_fill.resolve_component_files(checkpoint_path, component) or None + except Exception as exc: # no local snapshot, odd layout: fall back to the broadcast + logger.info("No checkpoint files for a sharded fill of %s: %s", component, exc) + # Ranks decide what to fill on their own, so they must see the same files (e.g. a node without the + # snapshot would otherwise leave its ranks expecting a broadcast the others skip). + if dist.is_initialized() and dist.get_world_size() > 1: + from miles.utils.distributed_utils import get_gloo_group + + mine = [os.path.basename(f) for f in files] if files else None + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, mine, group=get_gloo_group()) + if any(g != gathered[0] for g in gathered) or gathered[0] is None: + return None + return files + + def finish_fsdp_lazy_init(model: torch.nn.Module) -> None: """Run FSDP's one-time lazy init now instead of at the first forward. diff --git a/miles/backends/fsdp_utils/sharded_checkpoint_fill.py b/miles/backends/fsdp_utils/sharded_checkpoint_fill.py new file mode 100644 index 000000000..3f9ed3040 --- /dev/null +++ b/miles/backends/fsdp_utils/sharded_checkpoint_fill.py @@ -0,0 +1,170 @@ +"""Fill FSDP2 shards straight from the safetensors checkpoint, each rank reading only its slice. + +The default initial load builds the full model on rank 0 and broadcasts every tensor +(set_model_state_dict(broadcast_from_rank0=True)): rank 0 copies the whole model from host to +device and the network carries all of it. For parameters that exist unchanged in the checkpoint, +each rank can instead read the byte range of its own dim-0 shard and copy it to the device. +Anything that does not map cleanly (adapters, renamed or converted tensors, other placements) is +left to the broadcast. +""" + +from __future__ import annotations + +import glob +import json +import logging +import os +import struct +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +import torch +from torch.distributed.tensor import DTensor + +logger = logging.getLogger(__name__) + +_MAX_INFLIGHT_BYTES = 2 << 30 + +_DTYPES = { + "BF16": torch.bfloat16, + "F16": torch.float16, + "F32": torch.float32, + "F64": torch.float64, + "I64": torch.int64, + "I32": torch.int32, + "I16": torch.int16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, +} + + +def enabled() -> bool: + return os.environ.get("MILES_SHARDED_CHECKPOINT_FILL", "1") == "1" + + +def resolve_component_files(checkpoint_path: str, component: str) -> list[str]: + """The safetensors files diffusers would load for ``component``, or [] if unsure.""" + if os.path.isdir(checkpoint_path): + root = checkpoint_path + else: + from huggingface_hub import snapshot_download + + root = snapshot_download(checkpoint_path, allow_patterns=[f"{component}/*"], local_files_only=True) + comp_dir = os.path.join(root, component) + indexes = glob.glob(os.path.join(comp_dir, "*.safetensors.index.json")) + if len(indexes) == 1: + with open(indexes[0]) as f: + names = sorted(set(json.load(f)["weight_map"].values())) + files = [os.path.join(comp_dir, n) for n in names] + return files if all(os.path.exists(f) for f in files) else [] + if indexes: + return [] # several variants: leave the choice to diffusers + files = sorted(glob.glob(os.path.join(comp_dir, "*.safetensors"))) + return files if len(files) == 1 else [] + + +def index_files(files: list[str]) -> dict[str, tuple[str, int, torch.dtype, list[int]]]: + """key -> (file, absolute byte offset of the tensor, dtype, shape); keys seen twice are dropped.""" + index, dup = {}, set() + for path in files: + with open(path, "rb") as f: + (n,) = struct.unpack(" str: + return name.replace(".base_layer.", ".") + + +def _local_rows(param: DTensor) -> tuple[int, int] | None: + """(first row, row count) of this rank's shard if it is a full-width dim-0 slice, else None.""" + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + + shape, offset = compute_local_shape_and_global_offset(param.shape, param.device_mesh, param.placements) + if any(o != 0 for o in offset[1:]) or tuple(shape[1:]) != tuple(param.shape[1:]): + return None + return offset[0], shape[0] + + +def fill_from_checkpoint(model: torch.nn.Module, files: list[str], *, device) -> set[str]: + """Copy every cleanly mapped parameter shard from ``files``; returns the state-dict keys filled.""" + index = index_files(files) + plan = [] + for name, param in model.named_parameters(): + entry = index.get(_checkpoint_key(name)) + if entry is None or not isinstance(param, DTensor) or list(param.shape) != list(entry[3]): + continue + rows = _local_rows(param) + if rows is None: + continue + path, base, dtype, shape = entry + row_bytes = dtype.itemsize + for dim in shape[1:]: + row_bytes *= dim + plan.append((name, param, path, base + rows[0] * row_bytes, rows[1] * row_bytes, dtype)) + + fds = {} + pool = ThreadPoolExecutor(8, thread_name_prefix="shard-fill") + use_pinned = torch.cuda.is_available() and torch.device(device).type == "cuda" + try: + for path in {p[2] for p in plan}: + fds[path] = os.open(path, os.O_RDONLY) + + def read(item): + name, param, path, start, nbytes, dtype = item + buf = torch.empty(nbytes, dtype=torch.uint8, pin_memory=use_pinned) + view, got = memoryview(buf.numpy()), 0 + while got < nbytes: + n = os.preadv(fds[path], [view[got:]], start + got) + if n <= 0: + raise OSError(f"short read in {path}") + got += n + return buf + + filled = set() + # Reads run ahead on the pool, bounded by bytes in flight; copies happen in order here. + pending, inflight_bytes, todo = deque(), 0, iter(plan) + + def submit_more(): + nonlocal inflight_bytes + while not pending or inflight_bytes < _MAX_INFLIGHT_BYTES: + item = next(todo, None) + if item is None: + return + pending.append((item, pool.submit(read, item))) + inflight_bytes += item[4] + + submit_more() + while pending: + item, future = pending.popleft() + buf = future.result() + inflight_bytes -= item[4] + submit_more() + name, param, _, _, _, dtype = item + local = param.to_local() + with torch.no_grad(): + src = buf.to(device, non_blocking=use_pinned).view(dtype).view(local.shape) + local.copy_(src) + filled.add(name) + if use_pinned: + torch.cuda.synchronize() + logger.info( + "Filled %d of %d parameters from the checkpoint files on every rank", + len(filled), + len(list(model.parameters())), + ) + return filled + finally: + pool.shutdown() + for fd in fds.values(): + os.close(fd) diff --git a/miles/backends/sglang_diffusion_utils/arguments.py b/miles/backends/sglang_diffusion_utils/arguments.py index 3eca491ca..fd3aa8316 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 @@ -87,9 +86,19 @@ def new_add_argument_wrapper(*name_or_flags, **kwargs): old_add_argument(*new_name_or_flags_list, **final_kwargs) - parser.add_argument = new_add_argument_wrapper - ServerArgs.add_cli_args(parser) - parser.add_argument = old_add_argument + def register(p): + # Imported only on a cache miss: sglang.multimodal_gen pulls in torch, transformers and diffusers. + from sglang.multimodal_gen.runtime.server_args import ServerArgs + + p.add_argument = new_add_argument_wrapper + try: + ServerArgs.add_cli_args(p) + finally: + p.add_argument = old_add_argument + + from miles.backends.sglang_diffusion_utils.cli_cache import add_cached + + add_cached(parser, register) return parser diff --git a/miles/backends/sglang_diffusion_utils/cli_cache.py b/miles/backends/sglang_diffusion_utils/cli_cache.py new file mode 100644 index 000000000..1d6fe4a44 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/cli_cache.py @@ -0,0 +1,201 @@ +"""Cache of the sglang-d server CLI arguments, so parsing miles args does not import sglang. + +Registering ServerArgs' CLI arguments imports sglang.multimodal_gen, which pulls in torch, +transformers and diffusers (~25 s). The resulting argparse actions are plain data, so they are +recorded once per sglang build and replayed from a JSON file afterwards. Anything that does not +round-trip through JSON disables the cache and falls back to the import. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import logging +import os + +logger = logging.getLogger(__name__) + +_CACHE_VERSION = 4 + + +def parse_int_list(value: str) -> list[int]: + """Same as sglang.multimodal_gen.configs.pipeline_configs.base.parse_int_list.""" + if not value: + return [] + return [int(x.strip()) for x in value.split(",")] + + +# argparse ``type`` callables the cache can restore without importing sglang, by name. +_TYPES = {None: None, "int": int, "str": str, "float": float, "parse_int_list": parse_int_list} + + +class StoreBoolean(argparse.Action): + """Same behavior as sglang.multimodal_gen.runtime.utils.argparse.StoreBoolean.""" + + def __init__(self, option_strings, dest, default=False, required=False, help=None): + super().__init__( + option_strings=option_strings, + dest=dest, + nargs="?", + const=True, + default=default, + required=required, + help=help, + ) + + def __call__(self, parser, namespace, values, option_string=None): + if values is None: + setattr(namespace, self.dest, True) + elif isinstance(values, str): + if values.lower() == "true": + setattr(namespace, self.dest, True) + elif values.lower() == "false": + setattr(namespace, self.dest, False) + else: + raise ValueError(f"Invalid boolean value: {values}. Expected 'true' or 'false'.") + else: + setattr(namespace, self.dest, bool(values)) + + +_ACTIONS = { + argparse._StoreAction: "store", + argparse._StoreTrueAction: "store_true", + argparse._StoreFalseAction: "store_false", + argparse._StoreConstAction: "store_const", + argparse._AppendAction: "append", +} + + +def _cache_path() -> str | None: + spec = importlib.util.find_spec("sglang") # top-level only: does not run sglang/__init__ + if spec is None or not spec.submodule_search_locations: + return None + root = os.path.join(list(spec.submodule_search_locations)[0], "multimodal_gen") + if not os.path.isdir(root): + return None + # ServerArgs.add_cli_args pulls definitions from several modules (pipeline configs, ...), so key + # on every source file of multimodal_gen rather than server_args.py alone. + digest = hashlib.sha1(f"{_CACHE_VERSION}:{root}".encode()) + for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = sorted(d for d in dirnames if d != "__pycache__") + for name in sorted(filenames): + if name.endswith(".py"): + st = os.stat(os.path.join(dirpath, name)) + digest.update(f"{dirpath}/{name}:{st.st_size}:{st.st_mtime_ns};".encode()) + cache_dir = os.environ.get("MILES_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "miles")) + return os.path.join(cache_dir, f"sglang_cli_args_{digest.hexdigest()[:16]}.json") + + +def _enc(value): + # JSON has no tuple; keep tuples (e.g. default=("fp32",)) distinguishable from lists. + if isinstance(value, tuple): + return {"__tuple__": [_enc(v) for v in value]} + if isinstance(value, list): + return [_enc(v) for v in value] + return value + + +def _dec(value): + if isinstance(value, dict) and set(value) == {"__tuple__"}: + return tuple(_dec(v) for v in value["__tuple__"]) + if isinstance(value, list): + return [_dec(v) for v in value] + return value + + +def _encode(action: argparse.Action) -> dict: + kind = "store_boolean" if type(action).__name__ == "StoreBoolean" else _ACTIONS.get(type(action)) + if kind is None: + raise TypeError(f"unsupported action {type(action)!r}") + type_name = None if action.type is None else getattr(action.type, "__name__", None) + if type_name not in _TYPES: + raise TypeError(f"unsupported type {action.type!r}") + entry = { + "kind": kind, + "option_strings": action.option_strings, + "dest": action.dest, + "nargs": action.nargs, + "const": _enc(action.const), + "default": _enc(action.default), + "type": type_name, + "choices": list(action.choices) if action.choices is not None else None, + "required": action.required, + "help": action.help, + "metavar": action.metavar, + } + json.dumps(entry) # raises on anything that would not round-trip + return entry + + +def _add(parser: argparse.ArgumentParser, e: dict) -> None: + e = {**e, "default": _dec(e["default"]), "const": _dec(e["const"])} + common = {"dest": e["dest"], "required": e["required"], "help": e["help"]} + kind = e["kind"] + if kind == "store_boolean": + parser.add_argument(*e["option_strings"], action=StoreBoolean, default=e["default"], **common) + elif kind in ("store_true", "store_false"): + parser.add_argument(*e["option_strings"], action=kind, default=e["default"], **common) + elif kind == "store_const": + parser.add_argument(*e["option_strings"], action=kind, const=e["const"], default=e["default"], **common) + else: + kwargs = {"default": e["default"], "type": _TYPES[e["type"]], "metavar": e["metavar"], **common} + if e["nargs"] is not None: + kwargs["nargs"] = e["nargs"] + if e["const"] is not None: + kwargs["const"] = e["const"] + if e["choices"] is not None: + kwargs["choices"] = e["choices"] + parser.add_argument(*e["option_strings"], action=kind, **kwargs) + + +def add_cached(parser: argparse.ArgumentParser, register) -> None: + """Add the sglang args to ``parser``: from the cache if present, else via ``register(parser)`` and cache them.""" + path = None + try: + path = _cache_path() + if path and os.path.exists(path): + with open(path) as f: + entries = json.load(f) + # Replay into a scratch parser first so a broken cache cannot leave ``parser`` half filled. + scratch = argparse.ArgumentParser(add_help=False) + for e in entries: + _add(scratch, e) + for e in entries: + _add(parser, e) + return + except Exception as exc: # a stale or broken cache must never break argument parsing + logger.warning("sglang CLI arg cache unusable (%s); importing sglang", exc) + if path: + try: + os.remove(path) + except OSError: + pass + + before = len(parser._actions) + register(parser) + if not path: + return + try: + entries = [_encode(a) for a in parser._actions[before:]] + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = f"{path}.{os.getpid()}.tmp" + with open(tmp, "w") as f: + json.dump(entries, f) + os.replace(tmp, path) + except Exception as exc: + logger.info("not caching sglang CLI args: %s", exc) + + +def prime() -> str | None: + """Write the cache for the installed sglang (e.g. at image build time); returns its path.""" + from miles.backends.sglang_diffusion_utils.arguments import add_sglang_diffusion_arguments + + add_sglang_diffusion_arguments(argparse.ArgumentParser(add_help=False)) + path = _cache_path() + return path if path and os.path.exists(path) else None + + +if __name__ == "__main__": + print(prime() or "sglang CLI args not cacheable") diff --git a/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py b/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py new file mode 100644 index 000000000..3166b4eee --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py @@ -0,0 +1,73 @@ +"""Start sglang-d scheduler workers before the server process imports sglang. + +A scheduler worker spends most of its start-up importing torch, transformers, diffusers and sglang, +and the server process imports the same stack before it starts the workers. Starting the workers +first lets both imports run in parallel. Each worker imports its target, then waits for the +arguments the server would have spawned it with. + +Keep this module stdlib only: it runs before anything heavy is imported. +""" + +import importlib +import multiprocessing as mp +import os + +DEFAULT_TARGET = "sglang.multimodal_gen.runtime.managers.gpu_worker:run_scheduler_process" + + +def _worker_main(conn, target: str) -> None: + module_name, func_name = target.split(":") + func = getattr(importlib.import_module(module_name), func_name) + msg = conn.recv() + conn.close() + if msg is None: + return + env, args = msg + # Match the environment the server process had when it would have spawned this worker. + os.environ.clear() + os.environ.update(env) + func(*args) + + +class PrespawnedProcess: + """Stands in for ``mp.Process(target=..., args=args)``: ``start()`` hands the args to a waiting worker.""" + + def __init__(self, process, conn, args, name): + self._process = process + self._conn = conn + self._args = args + self.name = name + + def start(self) -> None: + self._conn.send((dict(os.environ), self._args)) + self._conn.close() + + def __getattr__(self, attr): + return getattr(self._process, attr) + + +class WaitingWorker: + """A started worker waiting for its args; call it like ``factory(args=..., name=...)`` or cancel it.""" + + def __init__(self, process, conn): + self._process = process + self._conn = conn + + def __call__(self, args, name): + return PrespawnedProcess(self._process, self._conn, args, name) + + def cancel(self) -> None: + self._conn.send(None) + self._conn.close() + + +def prespawn(num_workers: int, target: str = DEFAULT_TARGET) -> list[WaitingWorker]: + ctx = mp.get_context("spawn") + workers = [] + for _ in range(num_workers): + reader, writer = ctx.Pipe(duplex=False) + process = ctx.Process(target=_worker_main, args=(reader, target), daemon=True) + process.start() + reader.close() + workers.append(WaitingWorker(process, writer)) + return workers diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2aab5ae8d..dc2e84f14 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -6,18 +6,12 @@ import multiprocessing import os import time -from typing import TYPE_CHECKING - import requests -from sglang.srt.utils.common import kill_process_tree from miles.ray.ray_actor import RayActor from miles.utils.arguments import rollout_merges_lora from miles.utils.http_utils import get_host_info -if TYPE_CHECKING: - from sglang.multimodal_gen.runtime.server_args import ServerArgs - logger = logging.getLogger(__name__) @@ -49,7 +43,21 @@ def _bootstrap_scheduler_process_with_rollout_patches(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): + prespawned = None + if getattr(args, "rollout_prespawn_schedulers", False) and not apply_rollout_patches: + # Before any sglang import, so the workers import the same stack in parallel with this process. + from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn + + prespawned = prespawn(args.rollout_num_gpus_per_engine) + + # 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 + + 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 @@ -62,26 +70,54 @@ def _launch_server_target(server_args, apply_rollout_patches: bool = False): _ls_mod.bootstrap_scheduler_process = _bootstrap_scheduler_process_with_rollout_patches - from sglang.multimodal_gen.runtime.launch_server import launch_server + import sglang.multimodal_gen.runtime.launch_server as launch_mod - launch_server(server_args) + if prespawned is not None: + if hasattr(launch_mod, "register_prespawned_schedulers"): + launch_mod.register_prespawned_schedulers(prespawned) + else: + logger.warning("--rollout-prespawn-schedulers needs sglang's register_prespawned_schedulers; ignoring it") + for worker in prespawned: + worker.cancel() + launch_mod.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), ) p.start() _wait_server_healthy( - base_url=server_args.url(), + base_url=f"http://{host}:{port}", is_process_alive=lambda: p.is_alive(), ) @@ -134,22 +170,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 = 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) + 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 - 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 @@ -159,7 +187,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), ) @@ -339,6 +370,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/actor_group.py b/miles/ray/actor_group.py index f36a1ef70..20cbc6d3c 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -77,21 +77,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): """ diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 8cbcb6774..ce6bb9760 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -11,7 +11,6 @@ import ray import torch from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine from miles.ray.data_conversion_hub.flow_grpo import ( @@ -44,6 +43,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: @@ -532,6 +534,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": "false", } if rollout_merges_lora(args): # Merge in the train forward dtype, not fp32, to cut train/rollout consistency error. 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/miles/utils/arguments.py b/miles/utils/arguments.py index 0b6ed0467..12dcd64aa 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -124,6 +124,15 @@ def add_cluster_arguments(parser): "This will always be true when --colocate is set." ), ) + parser.add_argument( + "--rollout-prespawn-schedulers", + action="store_true", + default=False, + help=( + "Start the sglang-d scheduler workers before the server process imports sglang, so the two " + "imports run in parallel. Needs sglang's launch_server.register_prespawned_schedulers." + ), + ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/tests/fast/backends/fsdp_utils/test_broadcast_full_state.py b/tests/fast/backends/fsdp_utils/test_broadcast_full_state.py new file mode 100644 index 000000000..d6aac4dbc --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_broadcast_full_state.py @@ -0,0 +1,186 @@ +"""broadcast_full_state_to_fsdp gives every rank rank 0's params and buffers (CPU, gloo, FSDP2).""" + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +class _Net(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(8, 16) + self.b = torch.nn.Linear(16, 4, bias=False) + self.register_buffer("persistent", torch.arange(6, dtype=torch.float32)) + self.register_buffer("table", torch.arange(5, dtype=torch.float32) * 0.5, persistent=False) + + +def _reference(): + torch.manual_seed(0) + net = _Net() + with torch.no_grad(): + net.persistent.mul_(3) + net.table.add_(7) + return net + + +def _worker(rank, world, port, out): + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + from miles.backends.fsdp_utils.checkpoint import broadcast_full_state_to_fsdp + + if rank == 0: + model = _reference() + full_state = model.state_dict() + else: + from accelerate import init_empty_weights + + with init_empty_weights(include_buffers=False): + model = _Net() # params on meta, buffers hold constructor values, not rank 0's + full_state = {} + mesh = init_device_mesh("cpu", (world,)) + fully_shard(model, mesh=mesh) + broadcast_full_state_to_fsdp(model, full_state, cpu_offload=False, device="cpu") + ref = _reference() + params_ok = all( + torch.equal(p.full_tensor(), dict(ref.named_parameters())[n]) for n, p in model.named_parameters() + ) + buffers_ok = all(torch.equal(b, dict(ref.named_buffers())[n]) for n, b in model.named_buffers()) + out[rank] = (params_ok, buffers_ok) + finally: + dist.destroy_process_group() + + +def test_every_rank_gets_rank0_params_and_buffers(): + pytest.importorskip("accelerate") + world = 2 + port = 29500 + os.getpid() % 1000 + with mp.Manager() as manager: + out = manager.dict() + mp.spawn(_worker, args=(world, port, out), nprocs=world, join=True) + assert dict(out) == {0: (True, True), 1: (True, True)} + + +class _CkptNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(6, 10) # weight + bias straight from the checkpoint (bias stored as F32) + self.extra = torch.nn.Linear(6, 3, bias=False) # not in the checkpoint, like a LoRA adapter + self.renamed = torch.nn.Linear(6, 4, bias=False) # stored under another key + self.permuted = torch.nn.Linear(6, 8, bias=False) # same key and shape, but the loader reorders rows + self.register_buffer("table", torch.arange(3, dtype=torch.float32), persistent=False) + + +def _write_checkpoint(root): + from safetensors.torch import save_file + + torch.manual_seed(1) + raw = { + "a.weight": torch.randn(10, 6).to(torch.bfloat16), + "a.bias": torch.randn(10, dtype=torch.float32), + "old_name.weight": torch.randn(4, 6).to(torch.bfloat16), + "permuted.weight": torch.randn(8, 6).to(torch.bfloat16), + } + comp = os.path.join(root, "transformer") + os.makedirs(comp, exist_ok=True) + save_file({k: raw[k] for k in ("a.weight", "a.bias")}, os.path.join(comp, "model-1.safetensors")) + save_file({k: raw[k] for k in ("old_name.weight", "permuted.weight")}, os.path.join(comp, "model-2.safetensors")) + weight_map = {"a.weight": "model-1.safetensors", "a.bias": "model-1.safetensors"} + weight_map.update({"old_name.weight": "model-2.safetensors", "permuted.weight": "model-2.safetensors"}) + with open(os.path.join(comp, "model.safetensors.index.json"), "w") as f: + import json + + json.dump({"weight_map": weight_map}, f) + return raw + + +def _loaded_reference(raw): + """What rank 0's loader produces: bf16 params, a converted tensor, a fresh adapter.""" + torch.manual_seed(2) + net = _CkptNet().to(torch.bfloat16) + with torch.no_grad(): + net.a.weight.copy_(raw["a.weight"]) + net.a.bias.copy_(raw["a.bias"].to(torch.bfloat16)) + net.renamed.weight.copy_(raw["old_name.weight"]) + net.permuted.weight.copy_(raw["permuted.weight"].flip(0)) + net.table.add_(5) + return net + + +def _ckpt_worker(rank, world, port, root, out): + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + from miles.backends.fsdp_utils import checkpoint as ck + from miles.backends.fsdp_utils import sharded_checkpoint_fill + from miles.backends.fsdp_utils.checkpoint import broadcast_full_state_to_fsdp + + seen = {} + raw_fill, checked_fill = sharded_checkpoint_fill.fill_from_checkpoint, ck._fill_from_checkpoint_files + + def record_raw(*a, **k): + seen["raw"] = sorted(raw_fill(*a, **k)) + return set(seen["raw"]) + + def record_checked(*a, **k): + seen["checked"] = sorted(checked_fill(*a, **k)) + return set(seen["checked"]) + + sharded_checkpoint_fill.fill_from_checkpoint = record_raw + ck._fill_from_checkpoint_files = record_checked + + raw = _write_checkpoint(root) if rank == 0 else None + dist.barrier() + if raw is None: + from safetensors.torch import load_file + + raw = {} + for name in ("model-1.safetensors", "model-2.safetensors"): + raw.update(load_file(os.path.join(root, "transformer", name))) + if rank == 0: + model = _loaded_reference(raw) + full_state = model.state_dict() + else: + from accelerate import init_empty_weights + + with init_empty_weights(include_buffers=False): + model = _CkptNet().to(torch.bfloat16) + full_state = {} + fully_shard(model, mesh=init_device_mesh("cpu", (world,))) + files = sharded_checkpoint_fill.resolve_component_files(root, "transformer") + broadcast_full_state_to_fsdp(model, full_state, cpu_offload=False, device="cpu", checkpoint_files=files) + ref = _loaded_reference(raw) + params_ok = { + n: torch.equal(p.full_tensor(), dict(ref.named_parameters())[n]) for n, p in model.named_parameters() + } + buffers_ok = all(torch.equal(b, dict(ref.named_buffers())[n]) for n, b in model.named_buffers()) + out[rank] = (params_ok, buffers_ok, len(files), seen["raw"], seen["checked"]) + finally: + dist.destroy_process_group() + + +def test_checkpoint_fill_matches_rank0_load_and_falls_back_for_the_rest(tmp_path): + pytest.importorskip("accelerate") + pytest.importorskip("safetensors") + world = 2 + port = 30500 + os.getpid() % 1000 + with mp.Manager() as manager: + out = manager.dict() + mp.spawn(_ckpt_worker, args=(world, port, str(tmp_path), out), nprocs=world, join=True) + result = dict(out) + for rank in range(world): + params_ok, buffers_ok, n_files, raw, checked = result[rank] + assert n_files == 2 + assert raw == ["a.bias", "a.weight", "permuted.weight"], raw # name and shape matched + assert checked == ["a.bias", "a.weight"], checked # rank 0's spot check rejected the converted one + assert all(params_ok.values()), (rank, params_ok) + assert buffers_ok diff --git a/tests/fast/utils/prespawn_fake_target.py b/tests/fast/utils/prespawn_fake_target.py new file mode 100644 index 000000000..b872d9b04 --- /dev/null +++ b/tests/fast/utils/prespawn_fake_target.py @@ -0,0 +1,8 @@ +"""Stand-in for run_scheduler_process in the prespawn tests.""" + +import os + + +def run(local_rank, writer, payload): + writer.send((local_rank, payload, os.environ.get("PRESPAWN_TEST_ENV"))) + writer.close() diff --git a/tests/fast/utils/test_startup.py b/tests/fast/utils/test_startup.py new file mode 100644 index 000000000..0ab13de97 --- /dev/null +++ b/tests/fast/utils/test_startup.py @@ -0,0 +1,155 @@ +"""Startup-time helpers: no sglang import in the light processes, train rank creation, scheduler prespawn, +and the sglang CLI argument cache.""" + +import argparse +import subprocess +import sys + +import pytest + + +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={}, colocate=False, 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), + ] + + +def test_prespawned_workers_get_args_pipes_and_env_at_start(monkeypatch): + import multiprocessing as mp + + from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn + + workers = prespawn(2, target="tests.fast.utils.prespawn_fake_target:run") + monkeypatch.setenv("PRESPAWN_TEST_ENV", "set-after-prespawn") + readers, procs = [], [] + for rank, worker in enumerate(workers): + reader, writer = mp.Pipe(duplex=False) + proc = worker(args=(rank, writer, {"rank": rank}), name=f"w{rank}") + proc.start() + writer.close() + readers.append(reader) + procs.append(proc) + got = sorted(reader.recv() for reader in readers) + for proc in procs: + proc.join(30) + assert got == [(0, {"rank": 0}, "set-after-prespawn"), (1, {"rank": 1}, "set-after-prespawn")] + assert [proc.exitcode for proc in procs] == [0, 0] + + +def test_cancelled_prespawned_worker_exits_cleanly(): + from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn + + (worker,) = prespawn(1, target="tests.fast.utils.prespawn_fake_target:run") + worker.cancel() + worker._process.join(30) + assert worker._process.exitcode == 0 + + +def _register(parser): + from miles.backends.sglang_diffusion_utils.cli_cache import StoreBoolean, parse_int_list + + parser.add_argument("--sglang-tp-size", type=int, default=1) + parser.add_argument("--sglang-name", type=str, default=None, help="h") + parser.add_argument("--sglang-precisions", type=str, nargs="+", default=("fp32",)) + parser.add_argument("--sglang-steps", type=parse_int_list, default=[]) + parser.add_argument("--sglang-mode", type=str, choices=["a", "b"], default="a") + parser.add_argument("--sglang-flag", action=StoreBoolean, default=False) + parser.add_argument("--sglang-on", action="store_true") + + +def _parse(argv, monkeypatch, tmp_path, register): + from miles.backends.sglang_diffusion_utils import cli_cache + + monkeypatch.setattr(cli_cache, "_cache_path", lambda: str(tmp_path / "args.json")) + parser = argparse.ArgumentParser() + cli_cache.add_cached(parser, register) + return vars(parser.parse_args(argv)) + + +@pytest.mark.parametrize( + "argv", + [ + [], + ["--sglang-tp-size", "2", "--sglang-steps", "1, 2,3", "--sglang-flag", "--sglang-mode", "b"], + ["--sglang-precisions", "bf16", "fp16", "--sglang-flag", "false", "--sglang-on"], + ], +) +def test_cli_cache_hit_parses_like_the_live_registration(argv, monkeypatch, tmp_path): + live = _parse(argv, monkeypatch, tmp_path, _register) # miss: registers and writes the cache + assert (tmp_path / "args.json").exists() + + def must_not_register(parser): + raise AssertionError("cache hit should not register") + + cached = _parse(argv, monkeypatch, tmp_path, must_not_register) + assert cached == live + assert type(cached["sglang_precisions"]) is type(live["sglang_precisions"]) + + +def test_cli_cache_falls_back_when_the_cache_is_broken(monkeypatch, tmp_path): + (tmp_path / "args.json").write_text('[{"kind": "bogus"}]') + assert _parse(["--sglang-tp-size", "4"], monkeypatch, tmp_path, _register)["sglang_tp_size"] == 4 + + +def test_cli_cache_is_skipped_for_unknown_types(monkeypatch, tmp_path): + def register(parser): + parser.add_argument("--sglang-x", type=lambda v: v.upper(), default="a") + + assert _parse(["--sglang-x", "q"], monkeypatch, tmp_path, register)["sglang_x"] == "Q" + assert not (tmp_path / "args.json").exists() From 13306b782aa5e5f1761f1293e3d5001de1244e92 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Tue, 20 Oct 2026 21:07:00 +0530 Subject: [PATCH 2/4] perf: overlap train start-up with engine boot from the driver; fix sharded fill under LoRA Add --startup-overlap {off,spawn,init} and --startup-keep-engines-resident. The driver starts the rollout manager without waiting, creates (and with init, initialises) the train ranks while the engines boot, then waits for both. Component interfaces are unchanged. Fix the sharded checkpoint fill missing every parameter under PEFT: keys now drop the base_model.model. prefix as well as .base_layer. Also: startup timing probes, 0.5 s engine health poll. --- .../fsdp_utils/sharded_checkpoint_fill.py | 3 +- .../sglang_diffusion_engine.py | 2 +- miles/ray/placement_group.py | 50 +++++++++-- miles/utils/arguments.py | 20 +++++ miles/utils/startup_timing.py | 55 ++++++++++++ scripts/parse_startup_timing.py | 88 +++++++++++++++++++ tests/fast/utils/test_startup.py | 82 +++++++++++++++++ train_diffusion.py | 66 +++++++++++--- 8 files changed, 347 insertions(+), 19 deletions(-) create mode 100644 miles/utils/startup_timing.py create mode 100644 scripts/parse_startup_timing.py diff --git a/miles/backends/fsdp_utils/sharded_checkpoint_fill.py b/miles/backends/fsdp_utils/sharded_checkpoint_fill.py index 3f9ed3040..7ee9d4e65 100644 --- a/miles/backends/fsdp_utils/sharded_checkpoint_fill.py +++ b/miles/backends/fsdp_utils/sharded_checkpoint_fill.py @@ -83,7 +83,8 @@ def index_files(files: list[str]) -> dict[str, tuple[str, int, torch.dtype, list def _checkpoint_key(name: str) -> str: - return name.replace(".base_layer.", ".") + """Checkpoint key of a train parameter: PEFT wraps the model (``base_model.model.``) and its linears (``.base_layer.``).""" + return name.removeprefix("base_model.model.").replace(".base_layer.", ".") def _local_rows(param: DTensor) -> tuple[int, int] | None: diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index dc2e84f14..d9473ddba 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -141,7 +141,7 @@ def _wait_server_healthy(base_url, is_process_alive): if not is_process_alive(): raise Exception("Server process terminated unexpectedly.") - time.sleep(2) + time.sleep(0.5) class SGLangDiffusionEngine(RayActor): diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b87de71d7..935a7fbe7 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -129,6 +129,26 @@ def allocate_train_group(args, num_nodes, num_gpus_per_node, pg): ) +def resolve_startup_plan(args): + """(overlap mode, engines stay resident) for this run: what --startup-overlap and its prerequisites allow.""" + mode = args.startup_overlap + if mode != "off" and (args.train_only or args.debug_rollout_only): + logger.info("--startup-overlap=%s ignored: no train/rollout pair to overlap", mode) + mode = "off" + # Without an offload there is nothing to skip, so the engines are resident either way. + resident = args.startup_keep_engines_resident and bool(args.offload_rollout) + if mode == "init": + if args.num_rollout is None: + logger.warning("--startup-overlap=init needs --num-rollout (train init reads it); using 'spawn'") + mode = "spawn" + elif args.colocate and args.offload_rollout and not resident: + logger.warning( + "--startup-overlap=init needs --startup-keep-engines-resident when colocated; using 'spawn'" + ) + mode = "spawn" + return mode, resident + + def create_training_models(args, pgs, rollout_manager): logger.info("Initializing actor model...") actor_model = allocate_train_group( @@ -137,7 +157,12 @@ def create_training_models(args, pgs, rollout_manager): 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)) + return init_training_models(args, actor_model, actor_model.async_init(args, role="actor", with_ref=False), rollout_manager) + + +def init_training_models(args, actor_model, init_refs, rollout_manager): + """Wait for the train ranks' ``async_init`` and finish wiring them to the rollout manager.""" + start_rollout_ids = ray.get(init_refs) logger.info("Actor model initialized.") assert len(set(start_rollout_ids)) == 1 @@ -151,13 +176,21 @@ def create_training_models(args, pgs, rollout_manager): return actor_model -def create_rollout_manager(args, pg): +def start_rollout_manager(args, pg): + """Launch the rollout manager (and with it the engines) without waiting for them.""" logger.info("Creating rollout manager (num_gpus=%s)", 0) - rollout_manager = RolloutManager.options( + return RolloutManager.options( num_cpus=1, num_gpus=0, ).remote(args, pg) + +def finish_rollout_manager(args, rollout_manager, *, offload=True, wait=False): + """Wait for a started rollout manager; returns ``num_rollout_per_epoch`` (None unless it was needed). + + ``offload=False`` leaves the engines resident (the caller skips the matching onload). ``wait`` blocks until the + manager is up even when nothing else here would. + """ # calculate num_rollout from num_epoch num_rollout_per_epoch = None if args.num_rollout is None: @@ -166,9 +199,16 @@ def create_rollout_manager(args, pg): args.num_rollout = num_rollout_per_epoch * args.num_epoch assert args.num_rollout > 0 logger.info("Computed num_rollout=%s (num_rollout_per_epoch=%s)", args.num_rollout, num_rollout_per_epoch) + elif wait: + ray.get(rollout_manager.get_rollout_engines_and_lock.remote()) - if args.offload_rollout: + if args.offload_rollout and offload: ray.get(rollout_manager.offload.remote()) logger.info("Rollout manager created.") - return rollout_manager, num_rollout_per_epoch + return num_rollout_per_epoch + + +def create_rollout_manager(args, pg): + rollout_manager = start_rollout_manager(args, pg) + return rollout_manager, finish_rollout_manager(args, rollout_manager) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 12dcd64aa..e58e4c977 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -124,6 +124,26 @@ def add_cluster_arguments(parser): "This will always be true when --colocate is set." ), ) + parser.add_argument( + "--startup-overlap", + choices=["off", "spawn", "init"], + default="off", + help=( + "How much of train start-up overlaps the rollout engines' boot. 'spawn': create the train ranks " + "(process start and imports) while the engines boot. 'init': also run their init; needs " + "--num-rollout and, when colocated, --startup-keep-engines-resident, else it falls back to " + "'spawn'. The order of the steps changes only in the driver." + ), + ) + parser.add_argument( + "--startup-keep-engines-resident", + action="store_true", + default=False, + help=( + "Skip the engines' offload after boot and the onload before the first weight sync. They stay on " + "the GPU through train init, so it needs room for both." + ), + ) parser.add_argument( "--rollout-prespawn-schedulers", action="store_true", 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.py b/tests/fast/utils/test_startup.py index 0ab13de97..59859803a 100644 --- a/tests/fast/utils/test_startup.py +++ b/tests/fast/utils/test_startup.py @@ -153,3 +153,85 @@ def register(parser): assert _parse(["--sglang-x", "q"], monkeypatch, tmp_path, register)["sglang_x"] == "Q" assert not (tmp_path / "args.json").exists() + + +def _plan_args(**overrides): + from types import SimpleNamespace + + base = dict( + startup_overlap="off", + startup_keep_engines_resident=False, + train_only=False, + debug_rollout_only=False, + num_rollout=5, + colocate=True, + offload_rollout=True, + ) + return SimpleNamespace(**{**base, **overrides}) + + +@pytest.mark.parametrize( + "overrides, expected", + [ + ({}, ("off", False)), + ({"startup_keep_engines_resident": True}, ("off", True)), + ({"startup_overlap": "spawn"}, ("spawn", False)), + ({"startup_overlap": "init", "startup_keep_engines_resident": True}, ("init", True)), + # init needs the engines resident when colocated, and --num-rollout + ({"startup_overlap": "init"}, ("spawn", False)), + ({"startup_overlap": "init", "startup_keep_engines_resident": True, "num_rollout": None}, ("spawn", True)), + # nothing offloaded: nothing to hold back + ({"startup_overlap": "init", "colocate": False, "offload_rollout": False}, ("init", False)), + ({"startup_overlap": "spawn", "train_only": True}, ("off", False)), + ({"startup_overlap": "init", "debug_rollout_only": True}, ("off", False)), + ], +) +def test_resolve_startup_plan(overrides, expected): + from miles.ray.placement_group import resolve_startup_plan + + assert resolve_startup_plan(_plan_args(**overrides)) == expected + + +def test_finish_rollout_manager_offloads_unless_resident(monkeypatch): + from miles.ray import placement_group + + calls = [] + + class Method: + def __init__(self, name): + self.name = name + + def remote(self): + calls.append(self.name) + return self.name + + class Manager: + offload = Method("offload") + get_num_rollout_per_epoch = Method("per_epoch") + get_rollout_engines_and_lock = Method("ready") + + monkeypatch.setattr(placement_group.ray, "get", lambda ref: 7 if ref == "per_epoch" else ref) + + args = _plan_args(num_rollout=None, num_epoch=2) + assert placement_group.finish_rollout_manager(args, Manager) == 7 + assert (args.num_rollout, calls) == (14, ["per_epoch", "offload"]) + + calls.clear() + args = _plan_args() + assert placement_group.finish_rollout_manager(args, Manager, offload=False, wait=True) is None + assert calls == ["ready"] + + +@pytest.mark.parametrize( + "train_name, key", + [ + ("blocks.0.attn.to_q.weight", "blocks.0.attn.to_q.weight"), + ("base_model.model.blocks.0.attn.to_q.base_layer.weight", "blocks.0.attn.to_q.weight"), + ("base_model.model.blocks.0.norm.weight", "blocks.0.norm.weight"), + ("blocks.0.attn.to_q.base_layer.weight", "blocks.0.attn.to_q.weight"), + ], +) +def test_checkpoint_key_strips_peft_wrapping(train_name, key): + from miles.backends.fsdp_utils.sharded_checkpoint_fill import _checkpoint_key + + assert _checkpoint_key(train_name) == key diff --git a/train_diffusion.py b/train_diffusion.py index 3e69b372b..ca9271cb9 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -3,7 +3,15 @@ import ray -from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +import miles.utils.startup_timing as st # first, so its import time anchors the driver's +from miles.ray.placement_group import ( + allocate_train_group, + create_placement_groups, + finish_rollout_manager, + init_training_models, + resolve_startup_plan, + start_rollout_manager, +) from miles.utils.arguments import parse_args from miles.utils.logging_utils import configure_logger from miles.utils.misc import should_run_periodic_action @@ -21,19 +29,48 @@ def train(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"]) - logger.info("train: rollout manager ready") - - logger.info("train: creating training model") - actor_model = create_training_models(args, pgs, rollout_manager) - logger.info("train: training model ready") - - if args.offload_rollout: + mode, resident = resolve_startup_plan(args) + logger.info("train: startup overlap=%s, engines resident=%s", mode, resident) + with st.step("driver.startup"): + # The engines boot inside the rollout manager actor; starting it does not wait for them. + rollout_manager = start_rollout_manager(args, pgs["rollout"]) + init_refs = None + if mode != "off": + # Spawn the train ranks (process start, imports) while the engines boot. + 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"], + ) + if mode == "init": + init_refs = actor_model.async_init(args, role="actor", with_ref=False) + with st.step("driver.rollout_manager"): + num_rollout_per_epoch = finish_rollout_manager( + args, rollout_manager, offload=not resident, wait=mode != "off" + ) + logger.info("train: rollout manager ready") + + logger.info("train: creating training model") + if mode == "off": + 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"], + ) + if init_refs is None: + init_refs = actor_model.async_init(args, role="actor", with_ref=False) + with st.step("driver.train_init"): + actor_model = init_training_models(args, actor_model, init_refs, rollout_manager) + logger.info("train: training model ready") + + if args.offload_rollout and not resident: 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() # special case for eval-only if args.num_rollout == 0 and args.eval_interval is not None: @@ -69,6 +106,8 @@ def save(rollout_id): logger.info(f"train: rollout {rollout_id} actor train start") ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) + if rollout_id == args.start_rollout_id: + st.mark("driver.first_train_step_done") logger.info(f"train: rollout {rollout_id} actor train done") if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): @@ -88,5 +127,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 585e213749a91ef43e835a43cb7e5b33d0211c7f Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Tue, 20 Oct 2026 21:41:27 +0530 Subject: [PATCH 3/4] perf: opt-in --startup-skip-p2p-check for the rollout engines Sets SGLANG_SKIP_P2P_CHECK=1 on the engines so sglang trusts the driver's P2P report instead of timing a copy between every GPU pair (about 21 to 28 s on a fresh container on 2x H200). --- miles/ray/rollout.py | 2 ++ miles/utils/arguments.py | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index ce6bb9760..98c273c6a 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -537,6 +537,8 @@ def init_rollout_engines(args, pg, all_rollout_engines): # 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": "false", } + if args.startup_skip_p2p_check: + env_vars["SGLANG_SKIP_P2P_CHECK"] = "1" if rollout_merges_lora(args): # Merge in the train forward dtype, not fp32, to cut train/rollout consistency error. env_vars["SGLANG_DIFFUSION_LORA_MERGE_FP32"] = "1" if args.diffusion_forward_dtype == "fp32" else "0" diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index e58e4c977..c241aa0ee 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -144,6 +144,15 @@ def add_cluster_arguments(parser): "the GPU through train init, so it needs room for both." ), ) + parser.add_argument( + "--startup-skip-p2p-check", + action="store_true", + default=False, + help=( + "Set SGLANG_SKIP_P2P_CHECK=1 on the rollout engines: sglang trusts the driver's P2P report " + "instead of timing a copy between every GPU pair (~28 s on a fresh container on 2x H200)." + ), + ) parser.add_argument( "--rollout-prespawn-schedulers", action="store_true", From dc9b35f8bc037ff9d4c88853128c56bf3d60c268 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Tue, 20 Oct 2026 21:55:44 +0530 Subject: [PATCH 4/4] refactor: keep start-up changes inside their components Drive the start-up order from one helper in placement_group.py (start_rollout_and_train) and leave the train actor and group as on main. Hand the engine server's args to its child as plain dicts instead of a filtered Namespace. Drop the scheduler prespawn path, which does nothing on the current sglang branch. --- .../scheduler_prespawn.py | 73 --------- .../sglang_diffusion_engine.py | 87 +++++----- miles/ray/actor_group.py | 20 ++- miles/ray/placement_group.py | 29 ++++ miles/ray/train_actor.py | 5 - miles/utils/arguments.py | 9 -- tests/fast/utils/prespawn_fake_target.py | 8 - tests/fast/utils/test_startup.py | 148 ++++++++---------- train_diffusion.py | 52 +----- 9 files changed, 142 insertions(+), 289 deletions(-) delete mode 100644 miles/backends/sglang_diffusion_utils/scheduler_prespawn.py delete mode 100644 tests/fast/utils/prespawn_fake_target.py diff --git a/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py b/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py deleted file mode 100644 index 3166b4eee..000000000 --- a/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Start sglang-d scheduler workers before the server process imports sglang. - -A scheduler worker spends most of its start-up importing torch, transformers, diffusers and sglang, -and the server process imports the same stack before it starts the workers. Starting the workers -first lets both imports run in parallel. Each worker imports its target, then waits for the -arguments the server would have spawned it with. - -Keep this module stdlib only: it runs before anything heavy is imported. -""" - -import importlib -import multiprocessing as mp -import os - -DEFAULT_TARGET = "sglang.multimodal_gen.runtime.managers.gpu_worker:run_scheduler_process" - - -def _worker_main(conn, target: str) -> None: - module_name, func_name = target.split(":") - func = getattr(importlib.import_module(module_name), func_name) - msg = conn.recv() - conn.close() - if msg is None: - return - env, args = msg - # Match the environment the server process had when it would have spawned this worker. - os.environ.clear() - os.environ.update(env) - func(*args) - - -class PrespawnedProcess: - """Stands in for ``mp.Process(target=..., args=args)``: ``start()`` hands the args to a waiting worker.""" - - def __init__(self, process, conn, args, name): - self._process = process - self._conn = conn - self._args = args - self.name = name - - def start(self) -> None: - self._conn.send((dict(os.environ), self._args)) - self._conn.close() - - def __getattr__(self, attr): - return getattr(self._process, attr) - - -class WaitingWorker: - """A started worker waiting for its args; call it like ``factory(args=..., name=...)`` or cancel it.""" - - def __init__(self, process, conn): - self._process = process - self._conn = conn - - def __call__(self, args, name): - return PrespawnedProcess(self._process, self._conn, args, name) - - def cancel(self) -> None: - self._conn.send(None) - self._conn.close() - - -def prespawn(num_workers: int, target: str = DEFAULT_TARGET) -> list[WaitingWorker]: - ctx = mp.get_context("spawn") - workers = [] - for _ in range(num_workers): - reader, writer = ctx.Pipe(duplex=False) - process = ctx.Process(target=_worker_main, args=(reader, target), daemon=True) - process.start() - reader.close() - workers.append(WaitingWorker(process, writer)) - return workers diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index d9473ddba..2878ec8d9 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -43,19 +43,12 @@ def _bootstrap_scheduler_process_with_rollout_patches(spec): bootstrap_scheduler_process(spec) -def _launch_server_target(args, host, port, nccl_port, apply_rollout_patches: bool = False): - prespawned = None - if getattr(args, "rollout_prespawn_schedulers", False) and not apply_rollout_patches: - # Before any sglang import, so the workers import the same stack in parallel with this process. - from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn - - prespawned = prespawn(args.rollout_num_gpus_per_engine) - +def _launch_server_target(server_kwargs, sglang_attrs, 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 - server_args = ServerArgs.from_kwargs(**_compute_server_args(args, host=host, port=port, nccl_port=nccl_port)) + server_args = ServerArgs.from_kwargs(**_add_sglang_fields(server_kwargs, sglang_attrs)) server_args.host = server_args.host.strip("[]") # addict.Dict loses its `__frozen` attribute across spawn pickle; rebuild it. @@ -70,54 +63,31 @@ def _launch_server_target(args, host, port, nccl_port, apply_rollout_patches: bo _ls_mod.bootstrap_scheduler_process = _bootstrap_scheduler_process_with_rollout_patches - import sglang.multimodal_gen.runtime.launch_server as launch_mod - - if prespawned is not None: - if hasattr(launch_mod, "register_prespawned_schedulers"): - launch_mod.register_prespawned_schedulers(prespawned) - else: - logger.warning("--rollout-prespawn-schedulers needs sglang's register_prespawned_schedulers; ignoring it") - for worker in prespawned: - worker.cancel() - launch_mod.launch_server(server_args) + from sglang.multimodal_gen.runtime.launch_server import launch_server - -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) + launch_server(server_args) def launch_server_process( - args, - host: str, - port: int, - nccl_port: int, + server_kwargs: dict, + sglang_attrs: dict, 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. + ``server_kwargs`` and ``sglang_attrs`` are plain data (see ``_base_server_kwargs``); the child turns them + into ``ServerArgs``. """ # use spawn to avoid potential risks of fork in terms of subthreads or CUDA. multiprocessing.set_start_method("spawn", force=True) p = multiprocessing.Process( target=_launch_server_target, - args=(_picklable_args(args), host, port, nccl_port, apply_rollout_patches), + args=(server_kwargs, sglang_attrs, apply_rollout_patches), ) p.start() _wait_server_healthy( - base_url=f"http://{host}:{port}", + base_url=f"http://{server_kwargs['host']}:{server_kwargs['port']}", is_process_alive=lambda: p.is_alive(), ) @@ -187,10 +157,8 @@ def _init_normal(self): 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( - self.args, - host=self.server_host, - port=self.server_port, - nccl_port=self._nccl_port, + _base_server_kwargs(self.args, host=self.server_host, port=self.server_port, nccl_port=self._nccl_port), + _sglang_attrs(self.args), apply_rollout_patches=bool(patch_groups), ) @@ -390,8 +358,16 @@ def simulate_crash(self): def _compute_server_args(args, host, port, nccl_port): - from sglang.multimodal_gen.runtime.server_args import ServerArgs + return _add_sglang_fields(_base_server_kwargs(args, host, port, nccl_port), _sglang_attrs(args)) + + +def _sglang_attrs(args) -> dict: + """The --sglang-* values, by ServerArgs field name.""" + return {k.removeprefix("sglang_"): v for k, v in vars(args).items() if k.startswith("sglang_")} + +def _base_server_kwargs(args, host, port, nccl_port) -> dict: + """ServerArgs kwargs that need no sglang import (see ``_add_sglang_fields`` for the rest).""" # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg. kwargs = { @@ -417,18 +393,25 @@ def _compute_server_args(args, host, port, nccl_port): if getattr(args, "diffusion_flow_shift", None) is not None: kwargs["flow_shift"] = float(args.diffusion_flow_shift) - # Forward remaining ServerArgs fields set via --sglang-* CLI (ulysses_degree, ring_degree, dp_size, etc.). - for attr in dataclasses.fields(ServerArgs): - if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs: - kwargs[attr.name] = getattr(args, f"sglang_{attr.name}") - if rollout_merges_lora(args): kwargs["lora_target_modules"] = args.lora_target_modules - # dit_precision / vae_precision are PipelineConfig fields, not ServerArgs, so forward them explicitly (only when changed from the class default, to avoid clobbering a subclass override). + return kwargs + + +def _add_sglang_fields(base_kwargs: dict, sglang_attrs: dict) -> dict: + """``base_kwargs`` plus the --sglang-* fields ServerArgs accepts; imports sglang.""" from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + from sglang.multimodal_gen.runtime.server_args import ServerArgs + kwargs = dict(base_kwargs) + # Forward remaining ServerArgs fields set via --sglang-* CLI (ulysses_degree, ring_degree, dp_size, etc.). + for attr in dataclasses.fields(ServerArgs): + if attr.name in sglang_attrs and attr.name not in kwargs: + kwargs[attr.name] = sglang_attrs[attr.name] + + # dit_precision / vae_precision are PipelineConfig fields, not ServerArgs, so forward them explicitly (only when changed from the class default, to avoid clobbering a subclass override). for field_name in ("dit_precision", "vae_precision"): - val = getattr(args, f"sglang_{field_name}", None) + val = sglang_attrs.get(field_name) if val is not None and val != getattr(PipelineConfig, field_name, None): kwargs[field_name] = val diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 20cbc6d3c..f36a1ef70 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -77,23 +77,21 @@ 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 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( + # Create worker actors + self._actor_handlers = [] + master_addr, master_port = None, None + for rank in range(world_size): + actor = 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, 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:]] - ) + ).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) def async_init(self, args, role, with_ref=False): """ diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 935a7fbe7..0932b8d59 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -212,3 +212,32 @@ def finish_rollout_manager(args, rollout_manager, *, offload=True, wait=False): def create_rollout_manager(args, pg): rollout_manager = start_rollout_manager(args, pg) return rollout_manager, finish_rollout_manager(args, rollout_manager) + + +def start_rollout_and_train(args, pgs): + """Create the rollout manager and the train group per ``--startup-overlap``. + + Returns ``(rollout_manager, actor_model, num_rollout_per_epoch, engines_resident)``. The caller skips the + engines' onload when they are resident. The order of the steps is the only thing that changes with the mode; + each component is created and initialised exactly as ``create_rollout_manager`` / ``create_training_models`` do. + """ + mode, resident = resolve_startup_plan(args) + logger.info("Startup overlap=%s, engines resident=%s", mode, resident) + # The engines boot inside the rollout manager actor; starting it does not wait for them. + rollout_manager = start_rollout_manager(args, pgs["rollout"]) + if mode == "off": + num_rollout_per_epoch = finish_rollout_manager(args, rollout_manager, offload=not resident) + return rollout_manager, create_training_models(args, pgs, rollout_manager), num_rollout_per_epoch, resident + + 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"], + ) + init_refs = actor_model.async_init(args, role="actor", with_ref=False) if mode == "init" else None + num_rollout_per_epoch = finish_rollout_manager(args, rollout_manager, offload=not resident, wait=True) + if init_refs is None: + init_refs = actor_model.async_init(args, role="actor", with_ref=False) + actor_model = init_training_models(args, actor_model, init_refs, rollout_manager) + return rollout_manager, actor_model, num_rollout_per_epoch, resident diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 983dc3fcf..bd5e7f889 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -49,11 +49,6 @@ 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/miles/utils/arguments.py b/miles/utils/arguments.py index c241aa0ee..7d37c891b 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -153,15 +153,6 @@ def add_cluster_arguments(parser): "instead of timing a copy between every GPU pair (~28 s on a fresh container on 2x H200)." ), ) - parser.add_argument( - "--rollout-prespawn-schedulers", - action="store_true", - default=False, - help=( - "Start the sglang-d scheduler workers before the server process imports sglang, so the two " - "imports run in parallel. Needs sglang's launch_server.register_prespawned_schedulers." - ), - ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/tests/fast/utils/prespawn_fake_target.py b/tests/fast/utils/prespawn_fake_target.py deleted file mode 100644 index b872d9b04..000000000 --- a/tests/fast/utils/prespawn_fake_target.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Stand-in for run_scheduler_process in the prespawn tests.""" - -import os - - -def run(local_rank, writer, payload): - writer.send((local_rank, payload, os.environ.get("PRESPAWN_TEST_ENV"))) - writer.close() diff --git a/tests/fast/utils/test_startup.py b/tests/fast/utils/test_startup.py index 59859803a..cc3344d2b 100644 --- a/tests/fast/utils/test_startup.py +++ b/tests/fast/utils/test_startup.py @@ -1,5 +1,5 @@ -"""Startup-time helpers: no sglang import in the light processes, train rank creation, scheduler prespawn, -and the sglang CLI argument cache.""" +"""Startup-time helpers: no sglang import in the light processes, the sglang CLI argument cache, the driver's +startup plan and the sharded checkpoint fill's key mapping.""" import argparse import subprocess @@ -8,18 +8,6 @@ import pytest -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; " @@ -29,78 +17,6 @@ def test_engine_module_does_not_import_sglang(): 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={}, colocate=False, 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), - ] - - -def test_prespawned_workers_get_args_pipes_and_env_at_start(monkeypatch): - import multiprocessing as mp - - from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn - - workers = prespawn(2, target="tests.fast.utils.prespawn_fake_target:run") - monkeypatch.setenv("PRESPAWN_TEST_ENV", "set-after-prespawn") - readers, procs = [], [] - for rank, worker in enumerate(workers): - reader, writer = mp.Pipe(duplex=False) - proc = worker(args=(rank, writer, {"rank": rank}), name=f"w{rank}") - proc.start() - writer.close() - readers.append(reader) - procs.append(proc) - got = sorted(reader.recv() for reader in readers) - for proc in procs: - proc.join(30) - assert got == [(0, {"rank": 0}, "set-after-prespawn"), (1, {"rank": 1}, "set-after-prespawn")] - assert [proc.exitcode for proc in procs] == [0, 0] - - -def test_cancelled_prespawned_worker_exits_cleanly(): - from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn - - (worker,) = prespawn(1, target="tests.fast.utils.prespawn_fake_target:run") - worker.cancel() - worker._process.join(30) - assert worker._process.exitcode == 0 - - def _register(parser): from miles.backends.sglang_diffusion_utils.cli_cache import StoreBoolean, parse_int_list @@ -235,3 +151,63 @@ def test_checkpoint_key_strips_peft_wrapping(train_name, key): from miles.backends.fsdp_utils.sharded_checkpoint_fill import _checkpoint_key assert _checkpoint_key(train_name) == key + + +def test_engine_server_kwargs_need_no_sglang(): + code = ( + "import sys; from types import SimpleNamespace as N; " + "from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import _base_server_kwargs, _sglang_attrs; " + "a = N(hf_checkpoint='m', rollout_num_gpus_per_engine=2, sglang_tp_size=2, sglang_sp_degree=1, " + "sglang_enable_cfg_parallel=False, sglang_dit_precision='bf16', diffusion_flow_shift=12.0, use_lora=False, " + "train_only=False, lora_ipc_weight_sync=False, colocate=True); " + "k = _base_server_kwargs(a, '[::1]', 15000, 15001); " + "assert k['host'] == '[::1]' and k['master_port'] == 25001 and k['flow_shift'] == 12.0 and k['num_gpus'] == 2, k; " + "assert _sglang_attrs(a) == {'tp_size': 2, 'sp_degree': 1, 'enable_cfg_parallel': False, 'dit_precision': 'bf16'}; " + "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_rollout_weights_tag_matches_sglang(): + constants = pytest.importorskip("sglang.srt.constants") + from miles.ray.rollout import GPU_MEMORY_TYPE_WEIGHTS + + assert GPU_MEMORY_TYPE_WEIGHTS == constants.GPU_MEMORY_TYPE_WEIGHTS + + +@pytest.mark.parametrize( + "overrides, order", + [ + ({}, ["start", "finish", "alloc", "init", "wire"]), + ({"startup_overlap": "spawn"}, ["start", "alloc", "finish", "init", "wire"]), + ( + {"startup_overlap": "init", "startup_keep_engines_resident": True}, + ["start", "alloc", "init", "finish", "wire"], + ), + ], +) +def test_start_rollout_and_train_orders_the_steps_by_mode(monkeypatch, overrides, order): + from miles.ray import placement_group as pg + + events = [] + + class Group: + def async_init(self, args, role, with_ref=False): + events.append("init") + return ["ref"] + + monkeypatch.setattr(pg, "start_rollout_manager", lambda args, p: events.append("start") or "rm") + monkeypatch.setattr( + pg, "finish_rollout_manager", lambda args, rm, **kw: events.append("finish") or (kw.get("offload"), 3)[1] + ) + monkeypatch.setattr(pg, "allocate_train_group", lambda **kw: events.append("alloc") or Group()) + monkeypatch.setattr(pg, "init_training_models", lambda args, group, refs, rm: events.append("wire") or group) + + args = _plan_args( + actor_num_nodes=1, actor_num_gpus_per_node=2, **{k: v for k, v in overrides.items()} + ) + rm, group, per_epoch, resident = pg.start_rollout_and_train(args, {"rollout": 0, "actor": 0}) + assert (rm, per_epoch) == ("rm", 3) + assert resident == bool(overrides.get("startup_keep_engines_resident")) + assert events == order diff --git a/train_diffusion.py b/train_diffusion.py index ca9271cb9..f4f647bb0 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -4,14 +4,7 @@ import ray import miles.utils.startup_timing as st # first, so its import time anchors the driver's -from miles.ray.placement_group import ( - allocate_train_group, - create_placement_groups, - finish_rollout_manager, - init_training_models, - resolve_startup_plan, - start_rollout_manager, -) +from miles.ray.placement_group import create_placement_groups, start_rollout_and_train from miles.utils.arguments import parse_args from miles.utils.logging_utils import configure_logger from miles.utils.misc import should_run_periodic_action @@ -27,45 +20,14 @@ def train(args): logger.info("train: placement groups ready") init_tracking(args) - # create the rollout manager, with sglang engines inside. + # create the rollout manager, with sglang engines inside, and the training model. # need to initialize rollout manager first to calculate num_rollout - mode, resident = resolve_startup_plan(args) - logger.info("train: startup overlap=%s, engines resident=%s", mode, resident) + logger.info("train: creating rollout manager and training model") with st.step("driver.startup"): - # The engines boot inside the rollout manager actor; starting it does not wait for them. - rollout_manager = start_rollout_manager(args, pgs["rollout"]) - init_refs = None - if mode != "off": - # Spawn the train ranks (process start, imports) while the engines boot. - 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"], - ) - if mode == "init": - init_refs = actor_model.async_init(args, role="actor", with_ref=False) - with st.step("driver.rollout_manager"): - num_rollout_per_epoch = finish_rollout_manager( - args, rollout_manager, offload=not resident, wait=mode != "off" - ) - logger.info("train: rollout manager ready") - - logger.info("train: creating training model") - if mode == "off": - 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"], - ) - if init_refs is None: - init_refs = actor_model.async_init(args, role="actor", with_ref=False) - with st.step("driver.train_init"): - actor_model = init_training_models(args, actor_model, init_refs, rollout_manager) - logger.info("train: training model ready") - - if args.offload_rollout and not resident: + rollout_manager, actor_model, num_rollout_per_epoch, engines_resident = start_rollout_and_train(args, pgs) + logger.info("train: rollout manager and training model ready") + + if args.offload_rollout and not engines_resident: ray.get(rollout_manager.onload_weights.remote()) # always update weight first so that sglang has the loaded weights from training.