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 1aebbeb05..a575c6836 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -29,13 +29,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 @@ -66,7 +59,7 @@ class FSDPTrainRayActor(TrainRayActor): """ @with_defer(lambda: Timer().start("train_wait")) - def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # type: ignore[override] + def init(self, args: Namespace, role: str, with_ref: bool = False, gpu_ready=None) -> int: # type: ignore[override] super().init(args, role, with_ref) if args.deterministic_mode: @@ -113,6 +106,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty lora_adapter_paths=args.lora_adapter_paths, trainable=True, cpu_offload=args.fsdp_cpu_offload, + gpu_ready=gpu_ready, ) self.reference_models = {} if args.ref_load is not None: @@ -124,6 +118,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty checkpoint_path=args.ref_load, lora_adapter_paths=args.ref_lora_adapter_paths, cpu_offload=args.ref_cpu_offload, + gpu_ready=gpu_ready, ) self.teacher_models = {} if args.teacher_load is not None: @@ -135,6 +130,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty checkpoint_path=args.teacher_load, lora_adapter_paths=args.teacher_lora_adapter_paths, cpu_offload=args.teacher_cpu_offload, + gpu_ready=gpu_ready, ) # A natively CPU-offloaded model never sleeps: its parameters stay on the host and its buffers on the GPU. self.actor_models_to_sleep = [] if args.fsdp_cpu_offload else list(self.models.values()) @@ -209,6 +205,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..fb918b63f 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 @@ -12,6 +13,7 @@ from .device_move import move_model from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str from .sequence_parallel.plan import apply_sequence_parallel +from .startup_gate import wait_for_gpu_budget logger = logging.getLogger(__name__) @@ -26,6 +28,7 @@ def load_fsdp_models( lora_adapter_paths: list[str] | None = None, trainable: bool = False, cpu_offload: bool = False, + gpu_ready=None, ) -> dict[str, torch.nn.Module]: materialize_weights = dist.get_rank() == 0 master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) @@ -54,6 +57,13 @@ def load_fsdp_models( raise RuntimeError(f"{component} did not honor meta initialization") checkpoint.sync_model_dtypes(model) full_state = model.state_dict() if materialize_weights else {} + if gpu_ready is not None: + wait_for_gpu_budget( + model, + shard_world_size=parallel_state.get_mesh("fsdp").size(), + gpu_ready=gpu_ready, + engine_reserve_gib=args.startup_engine_reserve_gib, + ) model = apply_fsdp2( model, model_backend.fsdp_parallel_plan(model), @@ -61,7 +71,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 +94,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/fsdp_utils/startup_gate.py b/miles/backends/fsdp_utils/startup_gate.py new file mode 100644 index 000000000..b456c85a2 --- /dev/null +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -0,0 +1,95 @@ +"""Hold a colocated train rank's first large GPU allocation until it fits next to the rollout engines. + +Startup overlaps train init with engine boot. The train side only needs room for one parameter shard +plus the largest full tensor (the broadcast from rank 0 moves one tensor at a time) plus a margin. +What it waits for is decided by the driver and handed over as an opaque ``gpu_ready`` handle (a Ray +ObjectRef): engines initialized, or engines offloaded. With --startup-engine-reserve-gib, an upper +bound on what one engine rank holds, a rank whose budget fits next to that reserve does not wait. +""" + +from __future__ import annotations + +import logging +import os +import time + +import ray +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) + +GIB = 1 << 30 + + +def required_bytes(model: torch.nn.Module, shard_world_size: int, margin_bytes: int) -> tuple[int, int, int]: + total = 0 + largest = 0 + for p in model.parameters(): + nbytes = p.numel() * p.element_size() + total += nbytes + largest = max(largest, nbytes) + shard = -(-total // shard_world_size) + return shard, largest, shard + largest + margin_bytes + + +class GpuReady: + """What a train rank waits on before its first large GPU allocation, as handed over by the driver. + + Holds the Ray ObjectRef inside an object: Ray resolves refs passed as top-level task arguments + before the task starts, which would hold the whole train init until the engines are ready. + """ + + def __init__(self, ref): + self._ref = ref + + def wait(self) -> None: + ray.get(self._ref) + + +def _wait_free(need: int, timeout: float) -> int: + start = time.time() + while True: + free = torch.cuda.mem_get_info()[0] + if free >= need or time.time() - start > timeout: + return free + time.sleep(0.25) + + +def _all_ranks_agree(ok: bool) -> bool: + # Every rank must take the same branch, or one rank waits while another allocates. + if not dist.is_initialized(): + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=_gloo_group()) + return bool(flag.item()) + + +def wait_for_gpu_budget( + model: torch.nn.Module, *, shard_world_size: int, gpu_ready=None, engine_reserve_gib: float | None = None +) -> None: + margin = int(float(os.environ.get("MILES_STARTUP_GATE_MARGIN_GIB", "4")) * GIB) + timeout = float(os.environ.get("MILES_STARTUP_GATE_TIMEOUT_S", "900")) + shard, largest, need = required_bytes(model, shard_world_size, margin) + if engine_reserve_gib is not None: + # Free memory already excludes whatever the engines have allocated, so subtracting the full + # reserve can only over-count: if this fits, the engines' remaining load fits too. + free = torch.cuda.mem_get_info()[0] + if _all_ranks_agree(free - int(engine_reserve_gib * GIB) >= need): + return + if gpu_ready is not None: + gpu_ready.wait() + free = _wait_free(need, timeout) + if free < need: + raise RuntimeError( + f"startup gate: only {free / GIB:.1f} GiB free, need {need / GIB:.1f} GiB " + f"(shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} + margin {margin / GIB:.1f}). " + "If --startup-engine-reserve-gib is set, the engines may not fit next to the train shard; " + "unset it to offload the engines during train init." + ) + + +def _gloo_group(): + from miles.utils.distributed_utils import get_gloo_group + + return get_gloo_group() 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 7e61a1302..68ddcda8b 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -72,28 +72,31 @@ 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) - - def async_init(self, args, role, with_ref=False): + ).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, gpu_ready=None): """ Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. """ self.args = args - return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handlers] + kwargs = {} if gpu_ready is None else {"gpu_ready": gpu_ready} + return [actor.init.remote(args, role, with_ref=with_ref, **kwargs) for actor in self._actor_handlers] def async_train(self, rollout_id, rollout_data_ref): """Do one rollout training""" diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b87de71d7..b74b29215 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -137,7 +137,9 @@ 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)) + start_rollout_ids = ray.get( + actor_model.async_init(args, role="actor", with_ref=False, gpu_ready=_startup_gpu_ready(args, rollout_manager)) + ) logger.info("Actor model initialized.") assert len(set(start_rollout_ids)) == 1 @@ -151,6 +153,28 @@ def create_training_models(args, pgs, rollout_manager): return actor_model +def engines_stay_resident_at_startup(args) -> bool: + """Whether colocated engines keep their weights through train init (no startup offload/onload).""" + return args.offload_rollout and not args.serial_startup and args.startup_engine_reserve_gib is not None + + +def _startup_gpu_ready(args, rollout_manager): + """What colocated train ranks wait on before their first large GPU allocation (see startup_gate.py). + + The train backend only gets this opaque handle; deciding what it stands for stays here. + """ + if not args.offload_rollout or args.serial_startup: + return None + from miles.backends.fsdp_utils.startup_gate import GpuReady + + if engines_stay_resident_at_startup(args): + # Engine memory is bounded by the reserve: ranks whose budget fits next to it go right away, + # the others wait for the engines to finish loading. + return GpuReady(rollout_manager.wait_engines_ready.remote()) + # No bound on engine memory: release the engines first; they reload weights before the first sync. + return GpuReady(rollout_manager.offload.remote()) + + def create_rollout_manager(args, pg): logger.info("Creating rollout manager (num_gpus=%s)", 0) rollout_manager = RolloutManager.options( @@ -167,7 +191,8 @@ def create_rollout_manager(args, pg): assert args.num_rollout > 0 logger.info("Computed num_rollout=%s (num_rollout_per_epoch=%s)", args.num_rollout, num_rollout_per_epoch) - if args.offload_rollout: + # Overlapped startup decides the engines' startup offload in _startup_gpu_ready instead. + if args.offload_rollout and args.serial_startup: ray.get(rollout_manager.offload.remote()) logger.info("Rollout manager created.") diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 8cbcb6774..52f890830 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: @@ -57,6 +59,7 @@ def __init__(self, args, pg): self.pg = pg from miles.dashboard import hooks + self._engine_init_handles = [] hooks.register_rollout_manager(args) if not args.train_only: logger.info("RolloutManager: starting router...") @@ -100,11 +103,7 @@ def __init__(self, args, pg): self.num_new_engines = 0 logger.info("RolloutManager using no sglang engines (train_only).") else: - num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) - num_engines = args.rollout_num_gpus // num_gpu_per_engine - self.all_rollout_engines = [None] * num_engines - self.num_new_engines = init_rollout_engines(args, pg, self.all_rollout_engines) - logger.info("RolloutManager started %s rollout engines", len(self.all_rollout_engines)) + self._launch_engines() # pools outlive every rollout call, so the manager seats them and the rm functions find them as singletons _, bundle_indices, gpu_ids = pg self.reward_slots = ColocatedRewardSlots( @@ -161,7 +160,30 @@ def rollout_engines(self): # when doing multi-node serving, we will only send request to node-0 for each engine. return self.all_rollout_engines[:: self.nodes_per_engine] + def _launch_engines(self): + args = self.args + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) + num_engines = args.rollout_num_gpus // num_gpu_per_engine + self.all_rollout_engines = [None] * num_engines + # Defer so train init overlaps engine boot. --debug-rollout-only crashes unless init is synced. + defer = not args.serial_startup and not args.debug_rollout_only + self.num_new_engines, self._engine_init_handles = _init_rollout_engines( + args, self.pg, self.all_rollout_engines, defer_init=defer + ) + logger.info("RolloutManager started %s rollout engines", len(self.all_rollout_engines)) + + def _ensure_engines_ready(self): + """Surface deferred engine init errors. Engine calls queue behind init on the engine actors anyway.""" + if self._engine_init_handles: + ray.get(self._engine_init_handles) + self._engine_init_handles = [] + + def wait_engines_ready(self): + self._ensure_engines_ready() + return True + def get_rollout_engines_and_lock(self): + self._ensure_engines_ready() return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines def get_num_rollout_per_epoch(self): @@ -172,6 +194,7 @@ def generate(self, rollout_id, train_in_flight=None): """`train_in_flight`: refs of the train step this generate overlaps; colocated SFT encoders wait for them.""" from miles.dashboard import hooks + self._ensure_engines_ready() start_time = time.time() self.rollout_id = rollout_id if self.encode_pool is not None: @@ -214,6 +237,7 @@ def eval(self, rollout_id): if self.args.train_only: # if debug train only, we don't generate evaluation data return + self._ensure_engines_ready() self.health_monitoring_resume() result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True) @@ -494,8 +518,14 @@ def set_train_parallel_config(self, config: dict): def init_rollout_engines(args, pg, all_rollout_engines): + num_new_engines, _ = _init_rollout_engines(args, pg, all_rollout_engines, defer_init=False) + return num_new_engines + + +def _init_rollout_engines(args, pg, all_rollout_engines, defer_init): + """Create engines and start init. With defer_init the init handles are returned unawaited.""" if args.train_only: - return 0 + return 0, [] num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine @@ -532,6 +562,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. @@ -552,18 +584,20 @@ def init_rollout_engines(args, pg, all_rollout_engines): num_new_engines = len(rollout_engines) if num_new_engines == 0: - return num_new_engines + return num_new_engines, [] addr_and_ports = _allocate_rollout_engine_addr_and_ports_normal( args=args, num_engines=num_engines, rollout_engines=rollout_engines ) - # TODO: don't ray.get here to overlap train actor init with rollout engine init. - # somehow if we don't sync here, the --debug-rollout-only mode will crash. + # Deferred init lets train actor init overlap engine boot; --debug-rollout-only still syncs here, + # since it crashes otherwise. init_handles = [engine.init.remote(**(addr_and_ports[rank])) for rank, engine in rollout_engines] + if defer_init: + return num_new_engines, init_handles ray.get(init_handles) - return num_new_engines + return num_new_engines, [] def _allocate_rollout_engine_addr_and_ports_normal(*, args, num_engines, rollout_engines): 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 4390b123f..0d8048120 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -108,6 +108,35 @@ def add_cluster_arguments(parser): "This will always be true when --colocate is set." ), ) + parser.add_argument( + "--serial-startup", + action="store_true", + default=False, + help=( + "Boot the rollout engines and the train actors one after the other, as before, " + "instead of overlapping them. Only for debugging startup." + ), + ) + parser.add_argument( + "--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( + "--startup-engine-reserve-gib", + type=float, + default=None, + help=( + "Upper bound (GiB) on the GPU memory one colocated rollout engine rank holds after boot. When " + "set, the engines keep their weights through train init (no startup offload/onload) and a " + "train rank starts its GPU load as soon as it fits next to that reserve. Unset: the engines " + "are offloaded during train init and reloaded before the first weight sync." + ), + ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/scripts/run_diffusion_grpo_h3_t2va_2gpu.py b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py index 8a367df05..53ba97da3 100644 --- a/scripts/run_diffusion_grpo_h3_t2va_2gpu.py +++ b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py @@ -186,6 +186,9 @@ def execute(args: ScriptArgs, prompt_dir: str) -> None: "--deterministic-mode " "--rollout-health-check-interval 60 " "--miles-router-health-check-failure-threshold 30 " + # One tp2 engine rank holds about 62 GB of weights (about 65 GiB in use after boot), so the + # engines can stay resident while train init loads its shard next to them. + "--startup-engine-reserve-gib 72 " ) U.execute_train( 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_gate.py b/tests/fast/utils/test_startup_gate.py new file mode 100644 index 000000000..d57129f52 --- /dev/null +++ b/tests/fast/utils/test_startup_gate.py @@ -0,0 +1,231 @@ +import subprocess +import sys + +import pytest +import torch + +from miles.backends.fsdp_utils import startup_gate + + +class _Handle: + """Stands in for the driver's GpuReady handle: records the wait, optionally frees memory.""" + + def __init__(self, name, waited, on_wait=None): + self.name, self.waited, self.on_wait = name, waited, on_wait + + def wait(self): + self.waited.append(self.name) + if self.on_wait: + self.on_wait() + + +def _free(monkeypatch, value): + state = {"free": value} + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (state["free"], 1 << 41), raising=False) + return state + + +def test_required_bytes_is_shard_plus_largest_plus_margin(): + model = torch.nn.Sequential(torch.nn.Linear(10, 10, bias=False), torch.nn.Linear(10, 40, bias=False)) + total = (100 + 400) * 4 + shard, largest, need = startup_gate.required_bytes(model, shard_world_size=2, margin_bytes=1000) + assert shard == total // 2 + assert largest == 400 * 4 + assert need == shard + largest + 1000 + + +def test_required_bytes_works_on_meta_params(): + with torch.device("meta"): + model = torch.nn.Linear(8, 8, bias=False) + assert startup_gate.required_bytes(model, 1, 0)[2] == 8 * 8 * 4 * 2 + + +def test_gate_waits_for_the_handle_then_passes_when_it_fits(monkeypatch): + _free(monkeypatch, 1 << 40) + waited = [] + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-offloaded", waited) + ) + assert waited == ["engines-offloaded"] + + +def test_gate_passes_once_the_handle_frees_memory(monkeypatch): + state = _free(monkeypatch, 1) + waited = [] + handle = _Handle("engines-offloaded", waited, on_wait=lambda: state.update(free=1 << 40)) + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=handle) + assert waited == ["engines-offloaded"] + + +def test_gate_raises_when_it_never_fits(monkeypatch): + _free(monkeypatch, 1) + monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") + with pytest.raises(RuntimeError, match="startup gate"): + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-ready", []) + ) + + +def test_reserve_lets_train_go_before_the_engines_when_it_fits(monkeypatch): + _free(monkeypatch, 100 << 30) + waited = [] + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-ready", waited), engine_reserve_gib=60 + ) + assert waited == [] + + +def test_reserve_that_does_not_fit_falls_back_to_waiting(monkeypatch): + _free(monkeypatch, 100 << 30) + waited = [] + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-ready", waited), engine_reserve_gib=99 + ) + assert waited == ["engines-ready"] + + +def test_driver_picks_what_train_waits_for(): + import argparse + + from miles.ray import placement_group as pg + + class _RM: + wait_engines_ready = type("R", (), {"remote": staticmethod(lambda: "engines-ready")}) + offload = type("R", (), {"remote": staticmethod(lambda: "engines-offloaded")}) + + def args(**kw): + base = dict(offload_rollout=True, serial_startup=False, startup_engine_reserve_gib=None) + return argparse.Namespace(**{**base, **kw}) + + assert pg._startup_gpu_ready(args(offload_rollout=False), _RM) is None + assert pg._startup_gpu_ready(args(serial_startup=True), _RM) is None + assert pg._startup_gpu_ready(args(), _RM)._ref == "engines-offloaded" + assert pg._startup_gpu_ready(args(startup_engine_reserve_gib=72), _RM)._ref == "engines-ready" + assert pg.engines_stay_resident_at_startup(args(startup_engine_reserve_gib=72)) + assert not pg.engines_stay_resident_at_startup(args()) + + +def test_gpu_ready_reaches_the_task_unresolved_and_the_task_starts_first(): + """Ray resolves top-level ref args before a task runs; GpuReady must not let that happen.""" + import time + + ray = pytest.importorskip("ray") + ray.init(num_cpus=2, include_dashboard=False, ignore_reinit_error=True, log_to_driver=False) + try: + + @ray.remote + class Engines: + def ready(self, delay): + time.sleep(delay) + return time.time() + + @ray.remote + class Train: + def init(self, gpu_ready): + started = time.time() + unresolved = isinstance(gpu_ready._ref, ray.ObjectRef) + gpu_ready.wait() + return started, unresolved + + engines, train = Engines.remote(), Train.remote() + ray.get([engines.__ray_ready__.remote(), train.__ray_ready__.remote()]) # workers up first + ready = startup_gate.GpuReady(engines.ready.remote(8.0)) + started, unresolved = ray.get(train.init.remote(ready)) + resolved = ray.get(ready._ref) + assert unresolved + assert started < resolved - 4.0 # train init began well before the engines were ready + finally: + ray.shutdown() + + +def test_picklable_args_drops_only_unpicklable_values(): + import argparse + import pickle + + from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import _picklable_args + + args = argparse.Namespace(hf_checkpoint="m", sglang_tp_size=2, hook=lambda: None) + copy = _picklable_args(args) + assert vars(copy) == {"hf_checkpoint": "m", "sglang_tp_size": 2} + pickle.dumps(copy) + + +def test_engine_module_does_not_import_sglang(): + code = ( + "import sys; import miles.backends.sglang_diffusion_utils.sglang_diffusion_engine, miles.ray.rollout; " + "print('sglang' in sys.modules)" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True).stdout.strip() + assert out == "False" + + +def test_train_ranks_are_created_together_then_get_rank0_address(monkeypatch): + import argparse + from types import SimpleNamespace + + from miles.ray import actor_group + + events = [] + + class FakeActor: + def __init__(self, rank): + self.rank = rank + self.get_master_addr_and_port = SimpleNamespace( + remote=lambda: events.append(("get_addr", self.rank)) or ("10.0.0.1", 29500) + ) + self.set_master_addr_and_port = SimpleNamespace( + remote=lambda addr, port: events.append(("set_addr", self.rank, addr, port)) + ) + + class FakeRemoteClass: + def options(self, **kwargs): + return SimpleNamespace( + remote=lambda world_size, rank, addr, port: events.append(("create", rank, addr)) or FakeActor(rank) + ) + + monkeypatch.setattr(actor_group.ray, "remote", lambda **kwargs: (lambda cls: FakeRemoteClass())) + monkeypatch.setattr(actor_group.ray, "get", lambda x: x) + monkeypatch.setattr(actor_group, "PlacementGroupSchedulingStrategy", lambda **kwargs: None) + args = argparse.Namespace(train_env_vars={}, deterministic_mode=False, train_backend="fsdp") + + actor_group.RayTrainGroup(args, num_nodes=1, num_gpus_per_node=3, pg=(object(), [0, 1, 2], [0, 1, 2])) + + assert events == [ + ("create", 0, None), + ("create", 1, None), + ("create", 2, None), + ("get_addr", 0), + ("set_addr", 1, "10.0.0.1", 29500), + ("set_addr", 2, "10.0.0.1", 29500), + ] + + +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 diff --git a/train_diffusion.py b/train_diffusion.py index 3e69b372b..6e66ca750 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -3,7 +3,12 @@ import ray -from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +from miles.ray.placement_group import ( + create_placement_groups, + create_rollout_manager, + create_training_models, + engines_stay_resident_at_startup, +) from miles.utils.arguments import parse_args from miles.utils.logging_utils import configure_logger from miles.utils.misc import should_run_periodic_action @@ -29,7 +34,7 @@ def train(args): actor_model = create_training_models(args, pgs, rollout_manager) logger.info("train: training model ready") - if args.offload_rollout: + if args.offload_rollout and not engines_stay_resident_at_startup(args): ray.get(rollout_manager.onload_weights.remote()) # always update weight first so that sglang has the loaded weights from training.