diff --git a/docker/Dockerfile b/docker/Dockerfile index 5af4ce11..15f9e2cc 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 05099e93..12ac5ca5 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 d005792c..403fc8aa 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 ff88c4d3..b0ba1ddd 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 ca61e39e..6d64a052 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 00000000..7ee9d4e6 --- /dev/null +++ b/miles/backends/fsdp_utils/sharded_checkpoint_fill.py @@ -0,0 +1,171 @@ +"""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: + """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: + """(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 3eca491c..fd3aa831 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 00000000..1d6fe4a4 --- /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/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2aab5ae8..2878ec8d 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,14 @@ 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(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(**_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. import addict @@ -68,20 +69,25 @@ def _launch_server_target(server_args, apply_rollout_patches: bool = False): def launch_server_process( - server_args: ServerArgs, + 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. + + ``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) - server_args.host = server_args.host.strip("[]") p = multiprocessing.Process( target=_launch_server_target, - args=(server_args, apply_rollout_patches), + args=(server_kwargs, sglang_attrs, apply_rollout_patches), ) p.start() _wait_server_healthy( - base_url=server_args.url(), + base_url=f"http://{server_kwargs['host']}:{server_kwargs['port']}", is_process_alive=lambda: p.is_alive(), ) @@ -105,7 +111,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): @@ -134,22 +140,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 +157,8 @@ 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), + _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), ) @@ -339,6 +338,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): @@ -357,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 = { @@ -384,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/placement_group.py b/miles/ray/placement_group.py index b87de71d..0932b8d5 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,45 @@ 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) + + +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/rollout.py b/miles/ray/rollout.py index 8cbcb677..98c273c6 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,7 +534,11 @@ 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 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 0b6ed046..7d37c891 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -124,6 +124,35 @@ 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( + "--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( "--skip-train-actor-gc-collect", action="store_true", diff --git a/miles/utils/startup_timing.py b/miles/utils/startup_timing.py new file mode 100644 index 00000000..07e6d202 --- /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 00000000..ada0c32a --- /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/backends/fsdp_utils/test_broadcast_full_state.py b/tests/fast/backends/fsdp_utils/test_broadcast_full_state.py new file mode 100644 index 00000000..d6aac4db --- /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/test_startup.py b/tests/fast/utils/test_startup.py new file mode 100644 index 00000000..cc3344d2 --- /dev/null +++ b/tests/fast/utils/test_startup.py @@ -0,0 +1,213 @@ +"""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 +import sys + +import pytest + + +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 _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() + + +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 + + +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 3e69b372..f4f647bb 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -3,7 +3,8 @@ 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 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 @@ -19,21 +20,19 @@ 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 - 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 rollout manager and training model") + with st.step("driver.startup"): + 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") - 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: + 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. - 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 +68,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 +89,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)