From 749057eece3b807075f31e80c869dda029cad2b7 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 22:01:31 +0530 Subject: [PATCH 1/3] perf: cut startup time to the first train step - Overlap engine boot with train actor init: the rollout manager no longer waits for engine init, and the startup offload/onload round trip is skipped. - Colocated train ranks gate their first large GPU allocation on engine init plus a memory budget (one shard + the largest tensor + margin). If that does not fit next to the resident engines, the engines are offloaded once and onloaded again before the first weight sync. - --startup-engine-reserve-gib lets the train GPU load overlap engine boot too, when an upper bound on engine memory is known. - --serial-startup restores the old order. - Load engine weights with mmap instead of the RunAI streamer. - Keep sglang out of the rollout manager and engine actor processes: ServerArgs is built in the sglang server process, and module-level sglang imports are deferred. - Create all train ranks at once instead of waiting for rank 0 to boot. - Bake .pyc files into the image. --- docker/Dockerfile | 9 + miles/backends/fsdp_utils/actor.py | 5 +- miles/backends/fsdp_utils/model_loader.py | 9 + miles/backends/fsdp_utils/startup_gate.py | 94 +++++++++++ .../sglang_diffusion_utils/arguments.py | 5 +- .../sglang_diffusion_engine.py | 68 +++++--- miles/ray/actor_group.py | 27 +-- miles/ray/placement_group.py | 9 +- miles/ray/rollout.py | 71 ++++++-- miles/ray/train_actor.py | 5 + miles/utils/arguments.py | 19 +++ tests/fast/utils/test_startup_gate.py | 157 ++++++++++++++++++ train_diffusion.py | 5 +- 13 files changed, 433 insertions(+), 50 deletions(-) create mode 100644 miles/backends/fsdp_utils/startup_gate.py create mode 100644 tests/fast/utils/test_startup_gate.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 5af4ce118..3d5edf775 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -88,3 +88,12 @@ RUN git clone https://github.com/radixark/miles_diffusion.git /root/miles_diffus # Pre-download PaddleOCR (en) det+rec+cls models so the OCR-reward actors don't # race-download them from bcebos at runtime. RUN python -c "from paddleocr import PaddleOCR; PaddleOCR(use_angle_cls=False, lang='en', use_gpu=False, show_log=False)" + +# ====================================== Bytecode ====================================================== + +# Bake .pyc files. About half of the package dirs ship without __pycache__, so every fresh container +# (every CI job) recompiles bytecode on its first import. Measured on a fresh container: +# `import sglang.multimodal_gen.runtime.launch_server` took 65 s without this step and 50 s with it +# (38 s once warm). Some vendored files are not valid Python 3 source, so compile errors are ignored. +RUN python -m compileall -q -j0 "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \ + /sgl-workspace/sglang/python /root/miles_diffusion || true diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 1aebbeb05..94c533044 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -66,7 +66,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, rollout_manager=None) -> int: # type: ignore[override] super().init(args, role, with_ref) if args.deterministic_mode: @@ -113,6 +113,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, + rollout_manager=rollout_manager, ) self.reference_models = {} if args.ref_load is not None: @@ -124,6 +125,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, + rollout_manager=rollout_manager, ) self.teacher_models = {} if args.teacher_load is not None: @@ -135,6 +137,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, + rollout_manager=rollout_manager, ) # A natively CPU-offloaded model never sleeps: its parameters stay on the host and its buffers on the GPU. self.actor_models_to_sleep = [] if args.fsdp_cpu_offload else list(self.models.values()) diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index ca61e39ec..80b9820c6 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/miles/backends/fsdp_utils/model_loader.py @@ -12,6 +12,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 +27,7 @@ def load_fsdp_models( lora_adapter_paths: list[str] | None = None, trainable: bool = False, cpu_offload: bool = False, + rollout_manager=None, ) -> dict[str, torch.nn.Module]: materialize_weights = dist.get_rank() == 0 master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) @@ -54,6 +56,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 rollout_manager is not None: + wait_for_gpu_budget( + model, + shard_world_size=parallel_state.get_mesh("fsdp").size(), + rollout_manager=rollout_manager, + engine_reserve_gib=args.startup_engine_reserve_gib, + ) model = apply_fsdp2( model, model_backend.fsdp_parallel_plan(model), diff --git a/miles/backends/fsdp_utils/startup_gate.py b/miles/backends/fsdp_utils/startup_gate.py new file mode 100644 index 000000000..cbb965920 --- /dev/null +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -0,0 +1,94 @@ +"""Hold the train side's first large GPU allocation until the colocated engines leave enough room. + +Startup overlaps train init with engine boot, and the engines stay resident through it. Before +its first large GPU allocation a colocated train rank waits for engine init (so it cannot grab +memory the engine load is about to need), then checks a memory budget: one parameter shard plus +the largest full tensor (the broadcast from rank 0 moves one tensor at a time) plus a margin. +If that does not fit next to the engines, the engines are offloaded once, as the serial startup +did, and the driver onloads them again before the first weight sync. + +With --startup-engine-reserve-gib (an upper bound on one engine rank's memory), a train rank does not +wait for engine init at all when the budget fits next to that reserve, so its GPU load overlaps engine +boot too. +""" + +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 + + +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 offloads while another allocates. + if not dist.is_initialized(): + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=_gloo_group()) + return bool(flag.item()) + + +def wait_for_gpu_budget( + model: torch.nn.Module, *, shard_world_size: int, rollout_manager, engine_reserve_gib: float | None = None +) -> None: + margin = int(float(os.environ.get("MILES_STARTUP_GATE_MARGIN_GIB", "4")) * GIB) + timeout = float(os.environ.get("MILES_STARTUP_GATE_TIMEOUT_S", "900")) + shard, largest, need = required_bytes(model, shard_world_size, margin) + if engine_reserve_gib is not None: + # Free memory already excludes whatever the engine has allocated, so subtracting the full reserve + # can only over-count: if this fits, the engine's remaining load fits too. + free = torch.cuda.mem_get_info()[0] + reserve = int(engine_reserve_gib * GIB) + if _all_ranks_agree(free - reserve >= need): + return + ray.get(rollout_manager.wait_engines_ready.remote()) + free = torch.cuda.mem_get_info()[0] + fits = _all_ranks_agree(free >= need) + if not fits: + logger.info( + "startup gate: %.1f GiB free next to the engines, need %.1f GiB; offloading the engines", + free / GIB, + need / GIB, + ) + ray.get(rollout_manager.boot_offload_once.remote()) + free = _wait_free(need, timeout) + if free < need: + raise RuntimeError( + f"startup gate: only {free / GIB:.1f} GiB free after offloading the engines, need " + f"{need / GIB:.1f} GiB (shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} " + f"+ margin {margin / GIB:.1f})" + ) + + +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..90b7354ae 100644 --- a/miles/backends/sglang_diffusion_utils/arguments.py +++ b/miles/backends/sglang_diffusion_utils/arguments.py @@ -1,4 +1,3 @@ -from sglang.multimodal_gen.runtime.server_args import ServerArgs from miles.utils.http_utils import _wrap_ipv6 @@ -14,6 +13,10 @@ def add_sglang_diffusion_arguments(parser): """ Add arguments to the parser for the SGLang server. """ + # Imported here: sglang.multimodal_gen pulls in all of sglang (~14 s), and this module is imported by + # every process that touches miles.utils.arguments, including ones that never parse arguments. + from sglang.multimodal_gen.runtime.server_args import ServerArgs + parser = add_sglang_diffusion_router_arguments(parser) parser.add_argument("--sglang-server-concurrency", type=int, default=512) diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2aab5ae8d..25aebdc00 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -9,14 +9,13 @@ 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 + pass logger = logging.getLogger(__name__) @@ -49,7 +48,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(args, host, port, nccl_port, apply_rollout_patches: bool = False): + # ServerArgs is built here, not in the engine actor: this process imports sglang anyway, and keeping + # sglang out of the actor takes one full sglang import (~40 s) off the engine startup path. + from sglang.multimodal_gen.runtime.server_args import ServerArgs + + 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 @@ -67,21 +73,42 @@ def _launch_server_target(server_args, apply_rollout_patches: bool = False): launch_server(server_args) +def _picklable_args(args): + """Copy of ``args`` that survives the plain pickle of a spawn start (Ray used cloudpickle to get it here).""" + import argparse + import pickle + + kept = {} + for key, value in vars(args).items(): + try: + pickle.dumps(value) + except Exception: + continue + kept[key] = value + return argparse.Namespace(**kept) + + def launch_server_process( - server_args: ServerArgs, + args, + host: str, + port: int, + nccl_port: int, apply_rollout_patches: bool = False, ) -> multiprocessing.Process: + """Start the sglang-d server in a child process and wait until it is healthy. + + ``host`` is bracketed for IPv6 (as in the server URL); the child strips it for binding. + """ # use spawn to avoid potential risks of fork in terms of subthreads or CUDA. multiprocessing.set_start_method("spawn", force=True) - server_args.host = server_args.host.strip("[]") p = multiprocessing.Process( target=_launch_server_target, - args=(server_args, apply_rollout_patches), + args=(_picklable_args(args), host, port, nccl_port, apply_rollout_patches), ) 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 +161,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 +178,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 +361,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..3da651a05 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, rollout_manager=None): """ Allocate GPU resourced and initialize model, optimzier, local ckpt, etc. """ self.args = args - return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handlers] + kwargs = {} if rollout_manager is None else {"rollout_manager": rollout_manager} + return [actor.init.remote(args, role, with_ref=with_ref, **kwargs) for actor in self._actor_handlers] def async_train(self, rollout_id, rollout_data_ref): """Do one rollout training""" diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index b87de71d7..cf828b683 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -137,7 +137,11 @@ 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)) + # Colocated train ranks gate their GPU phase on the engines (see startup_gate.py). + gate_manager = rollout_manager if args.offload_rollout and not args.serial_startup else None + start_rollout_ids = ray.get( + actor_model.async_init(args, role="actor", with_ref=False, rollout_manager=gate_manager) + ) logger.info("Actor model initialized.") assert len(set(start_rollout_ids)) == 1 @@ -167,7 +171,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: + # With overlapped startup the engines stay resident; the train gate offloads them only if needed. + 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..14355c1e8 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,8 @@ def __init__(self, args, pg): self.pg = pg from miles.dashboard import hooks + self._engine_init_handles = [] + self._boot_offloaded = False hooks.register_rollout_manager(args) if not args.train_only: logger.info("RolloutManager: starting router...") @@ -100,11 +104,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 +161,45 @@ 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 boot_offload_once(self): + """Startup fallback when the train init peak does not fit next to the resident engines. + + Train ranks call this concurrently; the actor runs calls one at a time, so only the first offloads. + """ + if not self._boot_offloaded: + self._boot_offloaded = True + self.offload() + return True + + def onload_weights_if_boot_offloaded(self): + if self._boot_offloaded: + self.onload_weights() + return self._boot_offloaded + def get_rollout_engines_and_lock(self): + self._ensure_engines_ready() return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines def get_num_rollout_per_epoch(self): @@ -172,6 +210,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: @@ -494,8 +533,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 +577,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 +599,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..1737f64da 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -108,6 +108,25 @@ 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( + "--startup-engine-reserve-gib", + type=float, + default=None, + help=( + "Upper bound on the GPU memory one colocated rollout engine rank holds after boot. When set, " + "train ranks start their GPU load as soon as it fits next to that reserve instead of waiting " + "for engine init to finish. Unset: wait for engine init." + ), + ) parser.add_argument( "--skip-train-actor-gc-collect", action="store_true", diff --git a/tests/fast/utils/test_startup_gate.py b/tests/fast/utils/test_startup_gate.py new file mode 100644 index 000000000..0e975ddb0 --- /dev/null +++ b/tests/fast/utils/test_startup_gate.py @@ -0,0 +1,157 @@ +import subprocess +import sys + +import pytest +import torch + +from miles.backends.fsdp_utils import startup_gate + + +class _Remote: + def __init__(self, fn): + self.fn = fn + + def remote(self): + return self.fn() + + +class FakeRolloutManager: + def __init__(self, on_offload=None): + self.calls = [] + self.on_offload = on_offload + self.wait_engines_ready = _Remote(lambda: self.calls.append("wait") or True) + self.boot_offload_once = _Remote(self._offload) + + def _offload(self): + self.calls.append("offload") + if self.on_offload: + self.on_offload() + return True + + +@pytest.fixture +def no_ray(monkeypatch): + monkeypatch.setattr(startup_gate.ray, "get", lambda x: x) + + +def _free(monkeypatch, value): + state = {"free": value} + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (state["free"], 1 << 41), raising=False) + return state + + +def test_required_bytes_is_shard_plus_largest_plus_margin(): + 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_keeps_engines_resident_when_it_fits(monkeypatch, no_ray): + _free(monkeypatch, 1 << 40) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + assert rm.calls == ["wait"] + + +def test_gate_offloads_engines_when_it_does_not_fit(monkeypatch, no_ray): + state = _free(monkeypatch, 1) + rm = FakeRolloutManager(on_offload=lambda: state.update(free=1 << 40)) + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + assert rm.calls == ["wait", "offload"] + + +def test_gate_raises_when_offload_does_not_free_enough(monkeypatch, no_ray): + _free(monkeypatch, 1) + monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") + rm = FakeRolloutManager() + with pytest.raises(RuntimeError, match="startup gate"): + startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + + +def test_reserve_lets_train_go_before_engines_when_it_fits(monkeypatch, no_ray): + _free(monkeypatch, 100 << 30) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=60 + ) + assert rm.calls == [] + + +def test_reserve_falls_back_to_waiting_for_engines_when_it_does_not_fit(monkeypatch, no_ray): + _free(monkeypatch, 100 << 30) + rm = FakeRolloutManager() + startup_gate.wait_for_gpu_budget( + torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=99 + ) + assert rm.calls == ["wait"] + + +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), + ] diff --git a/train_diffusion.py b/train_diffusion.py index 3e69b372b..47098de8b 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -30,7 +30,10 @@ def train(args): logger.info("train: training model ready") if args.offload_rollout: - ray.get(rollout_manager.onload_weights.remote()) + if args.serial_startup: + ray.get(rollout_manager.onload_weights.remote()) + else: + ray.get(rollout_manager.onload_weights_if_boot_offloaded.remote()) # always update weight first so that sglang has the loaded weights from training. actor_model.update_weights() From 7ddb65fd1c01f4ffc0f7ed38d8e6ec3808f1c3c6 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sun, 11 Oct 2026 23:11:04 +0530 Subject: [PATCH 2/3] perf: optionally start sglang-d scheduler workers before the server imports sglang --rollout-prespawn-schedulers starts the scheduler workers in the sglang server process before its own sglang import, so the workers' import of the same stack runs in parallel instead of after it. Needs sglang's launch_server.register_prespawned_schedulers (OnePunchMonk/sglang#8); with an sglang that lacks it the waiting workers are told to exit and startup falls back to the old path. On H3 t2va (2x H200) engine boot went from about 138 s to 104 s. --- .../scheduler_prespawn.py | 73 +++++++++++++++++++ .../sglang_diffusion_engine.py | 18 ++++- miles/utils/arguments.py | 9 +++ tests/fast/utils/prespawn_fake_target.py | 8 ++ tests/fast/utils/test_startup_gate.py | 31 ++++++++ 5 files changed, 137 insertions(+), 2 deletions(-) create mode 100644 miles/backends/sglang_diffusion_utils/scheduler_prespawn.py create mode 100644 tests/fast/utils/prespawn_fake_target.py diff --git a/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py b/miles/backends/sglang_diffusion_utils/scheduler_prespawn.py 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 25aebdc00..733b992f8 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -49,6 +49,13 @@ def _bootstrap_scheduler_process_with_rollout_patches(spec): def _launch_server_target(args, host, port, nccl_port, apply_rollout_patches: bool = False): + prespawned = None + if getattr(args, "rollout_prespawn_schedulers", False) and not apply_rollout_patches: + # Before any sglang import, so the workers import the same stack in parallel with this process. + from miles.backends.sglang_diffusion_utils.scheduler_prespawn import prespawn + + prespawned = prespawn(args.rollout_num_gpus_per_engine) + # 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 @@ -68,9 +75,16 @@ def _launch_server_target(args, host, port, nccl_port, apply_rollout_patches: bo _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): diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 1737f64da..c9ef2b982 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -117,6 +117,15 @@ def add_cluster_arguments(parser): "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, 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 index 0e975ddb0..290a1d7a7 100644 --- a/tests/fast/utils/test_startup_gate.py +++ b/tests/fast/utils/test_startup_gate.py @@ -155,3 +155,34 @@ def options(self, **kwargs): ("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 From 2ffa78893e1391ae08ea1bb2423672a6d006e8d1 Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Mon, 12 Oct 2026 00:18:26 +0530 Subject: [PATCH 3/3] perf: keep the rollout manager out of train init; faster train-side startup The train backend no longer gets the rollout manager during init. The driver hands colocated train ranks an optional opaque gpu_ready handle to wait on before their first large GPU allocation, and decides what it stands for: with --startup-engine-reserve-gib the engines stay resident and ranks whose budget fits next to the reserve go right away (otherwise they wait for engine init); without it the engines are offloaded during train init and reloaded before the first weight sync. The H3 recipe sets the reserve. Train-side startup: - Each rank fills its FSDP shard straight from the safetensors checkpoint (its own dim-0 slice) for parameters that map 1:1 onto it; the rest (adapters, converted tensors) is broadcast from rank 0 as before, and rank 0 spot-checks the filled slices against its loaded weights. - Rank 0 no longer copies its own shard to the GPU just to have it overwritten by the broadcast. - The weight-sync helpers (which import sglang) are imported where the updater is built, and fsdp_utils resolves its exports lazily, so the driver no longer imports sglang for them. - sglang-d's CLI argument definitions are cached (primed at image build), so parsing miles args no longer imports sglang. --- docker/Dockerfile | 4 + miles/backends/fsdp_utils/__init__.py | 44 ++-- miles/backends/fsdp_utils/actor.py | 26 ++- miles/backends/fsdp_utils/checkpoint.py | 67 +++++- miles/backends/fsdp_utils/model_loader.py | 37 +++- .../fsdp_utils/sharded_checkpoint_fill.py | 170 +++++++++++++++ miles/backends/fsdp_utils/startup_gate.py | 69 +++--- .../sglang_diffusion_utils/arguments.py | 20 +- .../sglang_diffusion_utils/cli_cache.py | 201 ++++++++++++++++++ .../sglang_diffusion_engine.py | 5 - miles/ray/actor_group.py | 4 +- miles/ray/placement_group.py | 28 ++- miles/ray/rollout.py | 17 +- miles/utils/arguments.py | 7 +- scripts/run_diffusion_grpo_h3_t2va_2gpu.py | 3 + .../fsdp_utils/test_broadcast_full_state.py | 186 ++++++++++++++++ tests/fast/utils/test_startup_gate.py | 127 +++++++---- train_diffusion.py | 14 +- 18 files changed, 860 insertions(+), 169 deletions(-) create mode 100644 miles/backends/fsdp_utils/sharded_checkpoint_fill.py create mode 100644 miles/backends/sglang_diffusion_utils/cli_cache.py create mode 100644 tests/fast/backends/fsdp_utils/test_broadcast_full_state.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 3d5edf775..15f9e2cc3 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -97,3 +97,7 @@ RUN python -c "from paddleocr import PaddleOCR; PaddleOCR(use_angle_cls=False, l # (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 94c533044..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, rollout_manager=None) -> 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,7 +106,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag lora_adapter_paths=args.lora_adapter_paths, trainable=True, cpu_offload=args.fsdp_cpu_offload, - rollout_manager=rollout_manager, + gpu_ready=gpu_ready, ) self.reference_models = {} if args.ref_load is not None: @@ -125,7 +118,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag checkpoint_path=args.ref_load, lora_adapter_paths=args.ref_lora_adapter_paths, cpu_offload=args.ref_cpu_offload, - rollout_manager=rollout_manager, + gpu_ready=gpu_ready, ) self.teacher_models = {} if args.teacher_load is not None: @@ -137,7 +130,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag checkpoint_path=args.teacher_load, lora_adapter_paths=args.teacher_lora_adapter_paths, cpu_offload=args.teacher_cpu_offload, - rollout_manager=rollout_manager, + 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()) @@ -212,6 +205,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manag ) 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 80b9820c6..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 @@ -27,7 +28,7 @@ def load_fsdp_models( lora_adapter_paths: list[str] | None = None, trainable: bool = False, cpu_offload: bool = False, - rollout_manager=None, + gpu_ready=None, ) -> dict[str, torch.nn.Module]: materialize_weights = dist.get_rank() == 0 master_dtype = parse_dtype_from_str(args.fsdp_master_dtype) @@ -56,11 +57,11 @@ 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 rollout_manager is not None: + if gpu_ready is not None: wait_for_gpu_budget( model, shard_world_size=parallel_state.get_mesh("fsdp").size(), - rollout_manager=rollout_manager, + gpu_ready=gpu_ready, engine_reserve_gib=args.startup_engine_reserve_gib, ) model = apply_fsdp2( @@ -70,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: @@ -88,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 index cbb965920..b456c85a2 100644 --- a/miles/backends/fsdp_utils/startup_gate.py +++ b/miles/backends/fsdp_utils/startup_gate.py @@ -1,15 +1,10 @@ -"""Hold the train side's first large GPU allocation until the colocated engines leave enough room. - -Startup overlaps train init with engine boot, and the engines stay resident through it. Before -its first large GPU allocation a colocated train rank waits for engine init (so it cannot grab -memory the engine load is about to need), then checks a memory budget: one parameter shard plus -the largest full tensor (the broadcast from rank 0 moves one tensor at a time) plus a margin. -If that does not fit next to the engines, the engines are offloaded once, as the serial startup -did, and the driver onloads them again before the first weight sync. - -With --startup-engine-reserve-gib (an upper bound on one engine rank's memory), a train rank does not -wait for engine init at all when the budget fits next to that reserve, so its GPU load overlaps engine -boot too. +"""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 @@ -38,6 +33,20 @@ def required_bytes(model: torch.nn.Module, shard_world_size: int, margin_bytes: 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: @@ -48,7 +57,7 @@ def _wait_free(need: int, timeout: float) -> int: def _all_ranks_agree(ok: bool) -> bool: - # Every rank must take the same branch, or one rank offloads while another allocates. + # 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) @@ -57,35 +66,27 @@ def _all_ranks_agree(ok: bool) -> bool: def wait_for_gpu_budget( - model: torch.nn.Module, *, shard_world_size: int, rollout_manager, engine_reserve_gib: float | None = None + 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 engine has allocated, so subtracting the full reserve - # can only over-count: if this fits, the engine's remaining load fits too. + # 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] - reserve = int(engine_reserve_gib * GIB) - if _all_ranks_agree(free - reserve >= need): + if _all_ranks_agree(free - int(engine_reserve_gib * GIB) >= need): return - ray.get(rollout_manager.wait_engines_ready.remote()) - free = torch.cuda.mem_get_info()[0] - fits = _all_ranks_agree(free >= need) - if not fits: - logger.info( - "startup gate: %.1f GiB free next to the engines, need %.1f GiB; offloading the engines", - free / GIB, - need / GIB, + 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." ) - ray.get(rollout_manager.boot_offload_once.remote()) - free = _wait_free(need, timeout) - if free < need: - raise RuntimeError( - f"startup gate: only {free / GIB:.1f} GiB free after offloading the engines, need " - f"{need / GIB:.1f} GiB (shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} " - f"+ margin {margin / GIB:.1f})" - ) def _gloo_group(): diff --git a/miles/backends/sglang_diffusion_utils/arguments.py b/miles/backends/sglang_diffusion_utils/arguments.py index 90b7354ae..fd3aa8316 100644 --- a/miles/backends/sglang_diffusion_utils/arguments.py +++ b/miles/backends/sglang_diffusion_utils/arguments.py @@ -13,10 +13,6 @@ def add_sglang_diffusion_arguments(parser): """ Add arguments to the parser for the SGLang server. """ - # Imported here: sglang.multimodal_gen pulls in all of sglang (~14 s), and this module is imported by - # every process that touches miles.utils.arguments, including ones that never parse arguments. - from sglang.multimodal_gen.runtime.server_args import ServerArgs - parser = add_sglang_diffusion_router_arguments(parser) parser.add_argument("--sglang-server-concurrency", type=int, default=512) @@ -90,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/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 733b992f8..dc2e84f14 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -6,17 +6,12 @@ import multiprocessing import os import time -from typing import TYPE_CHECKING - import requests 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: - pass - logger = logging.getLogger(__name__) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 3da651a05..68ddcda8b 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -90,12 +90,12 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): [actor.set_master_addr_and_port.remote(master_addr, master_port) for actor in self._actor_handlers[1:]] ) - def async_init(self, args, role, with_ref=False, rollout_manager=None): + 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 - kwargs = {} if rollout_manager is None else {"rollout_manager": rollout_manager} + 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): diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index cf828b683..b74b29215 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -137,10 +137,8 @@ def create_training_models(args, pgs, rollout_manager): num_gpus_per_node=args.actor_num_gpus_per_node, pg=pgs["actor"], ) - # Colocated train ranks gate their GPU phase on the engines (see startup_gate.py). - gate_manager = rollout_manager if args.offload_rollout and not args.serial_startup else None start_rollout_ids = ray.get( - actor_model.async_init(args, role="actor", with_ref=False, rollout_manager=gate_manager) + actor_model.async_init(args, role="actor", with_ref=False, gpu_ready=_startup_gpu_ready(args, rollout_manager)) ) logger.info("Actor model initialized.") @@ -155,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( @@ -171,7 +191,7 @@ 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) - # With overlapped startup the engines stay resident; the train gate offloads them only if needed. + # 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()) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 14355c1e8..52f890830 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -60,7 +60,6 @@ def __init__(self, args, pg): from miles.dashboard import hooks self._engine_init_handles = [] - self._boot_offloaded = False hooks.register_rollout_manager(args) if not args.train_only: logger.info("RolloutManager: starting router...") @@ -183,21 +182,6 @@ def wait_engines_ready(self): self._ensure_engines_ready() return True - def boot_offload_once(self): - """Startup fallback when the train init peak does not fit next to the resident engines. - - Train ranks call this concurrently; the actor runs calls one at a time, so only the first offloads. - """ - if not self._boot_offloaded: - self._boot_offloaded = True - self.offload() - return True - - def onload_weights_if_boot_offloaded(self): - if self._boot_offloaded: - self.onload_weights() - return self._boot_offloaded - def get_rollout_engines_and_lock(self): self._ensure_engines_ready() return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines @@ -253,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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c9ef2b982..0d8048120 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -131,9 +131,10 @@ def add_cluster_arguments(parser): type=float, default=None, help=( - "Upper bound on the GPU memory one colocated rollout engine rank holds after boot. When set, " - "train ranks start their GPU load as soon as it fits next to that reserve instead of waiting " - "for engine init to finish. Unset: wait for engine init." + "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( 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/test_startup_gate.py b/tests/fast/utils/test_startup_gate.py index 290a1d7a7..d57129f52 100644 --- a/tests/fast/utils/test_startup_gate.py +++ b/tests/fast/utils/test_startup_gate.py @@ -7,31 +7,16 @@ from miles.backends.fsdp_utils import startup_gate -class _Remote: - def __init__(self, fn): - self.fn = fn +class _Handle: + """Stands in for the driver's GpuReady handle: records the wait, optionally frees memory.""" - def remote(self): - return self.fn() + def __init__(self, name, waited, on_wait=None): + self.name, self.waited, self.on_wait = name, waited, on_wait - -class FakeRolloutManager: - def __init__(self, on_offload=None): - self.calls = [] - self.on_offload = on_offload - self.wait_engines_ready = _Remote(lambda: self.calls.append("wait") or True) - self.boot_offload_once = _Remote(self._offload) - - def _offload(self): - self.calls.append("offload") - if self.on_offload: - self.on_offload() - return True - - -@pytest.fixture -def no_ray(monkeypatch): - monkeypatch.setattr(startup_gate.ray, "get", lambda x: x) + def wait(self): + self.waited.append(self.name) + if self.on_wait: + self.on_wait() def _free(monkeypatch, value): @@ -55,44 +40,102 @@ def test_required_bytes_works_on_meta_params(): assert startup_gate.required_bytes(model, 1, 0)[2] == 8 * 8 * 4 * 2 -def test_gate_keeps_engines_resident_when_it_fits(monkeypatch, no_ray): +def test_gate_waits_for_the_handle_then_passes_when_it_fits(monkeypatch): _free(monkeypatch, 1 << 40) - rm = FakeRolloutManager() - startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) - assert rm.calls == ["wait"] + 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_offloads_engines_when_it_does_not_fit(monkeypatch, no_ray): +def test_gate_passes_once_the_handle_frees_memory(monkeypatch): state = _free(monkeypatch, 1) - rm = FakeRolloutManager(on_offload=lambda: state.update(free=1 << 40)) - startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) - assert rm.calls == ["wait", "offload"] + 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_offload_does_not_free_enough(monkeypatch, no_ray): +def test_gate_raises_when_it_never_fits(monkeypatch): _free(monkeypatch, 1) monkeypatch.setenv("MILES_STARTUP_GATE_TIMEOUT_S", "0.3") - rm = FakeRolloutManager() with pytest.raises(RuntimeError, match="startup gate"): - startup_gate.wait_for_gpu_budget(torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm) + 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_engines_when_it_fits(monkeypatch, no_ray): +def test_reserve_lets_train_go_before_the_engines_when_it_fits(monkeypatch): _free(monkeypatch, 100 << 30) - rm = FakeRolloutManager() + waited = [] startup_gate.wait_for_gpu_budget( - torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=60 + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-ready", waited), engine_reserve_gib=60 ) - assert rm.calls == [] + assert waited == [] -def test_reserve_falls_back_to_waiting_for_engines_when_it_does_not_fit(monkeypatch, no_ray): +def test_reserve_that_does_not_fit_falls_back_to_waiting(monkeypatch): _free(monkeypatch, 100 << 30) - rm = FakeRolloutManager() + waited = [] startup_gate.wait_for_gpu_budget( - torch.nn.Linear(4, 4), shard_world_size=1, rollout_manager=rm, engine_reserve_gib=99 + torch.nn.Linear(4, 4), shard_world_size=1, gpu_ready=_Handle("engines-ready", waited), engine_reserve_gib=99 ) - assert rm.calls == ["wait"] + 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(): diff --git a/train_diffusion.py b/train_diffusion.py index 47098de8b..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,11 +34,8 @@ def train(args): actor_model = create_training_models(args, pgs, rollout_manager) logger.info("train: training model ready") - if args.offload_rollout: - if args.serial_startup: - ray.get(rollout_manager.onload_weights.remote()) - else: - ray.get(rollout_manager.onload_weights_if_boot_offloaded.remote()) + 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. actor_model.update_weights()