Reference: startup time experiments for radixark/miles_diffusion#310

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

miles/backends/fsdp_utils/actor.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from argparse import Namespace
 from contextlib import nullcontext
 
@@ -11,6 +12,7 @@
 import miles.backends.fsdp_utils.configs.qwen_image  # noqa: F401 — register pipeline config
 import miles.backends.fsdp_utils.configs.sd3  # noqa: F401 — register pipeline config
 import miles.backends.fsdp_utils.configs.wan2_2  # noqa: F401 — register pipeline config
+import miles.utils.startup_timing as st
 from miles.ray.train_actor import TrainRayActor
 from miles.utils import tracking_utils, train_metric_utils
 from miles.utils.context_utils import with_defer
@@ -66,8 +68,10 @@ class FSDPTrainRayActor(TrainRayActor):
     """
 
     @with_defer(lambda: Timer().start("train_wait"))
-    def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # type: ignore[override]
-        super().init(args, role, with_ref)
+    def init(self, args: Namespace, role: str, with_ref: bool = False, rollout_manager=None) -> int:  # type: ignore[override]
+        st.mark("train.init_enter", rank=os.environ.get("RANK", "?"))
+        with st.step("train.super_init"):
+            super().init(args, role, with_ref)
 
         if args.deterministic_mode:
             _enable_deterministic_training(args)
@@ -104,16 +108,18 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             # flash-attn is opaque to torch's determinism flag; backends patch their own dispatch.
             self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend)
         self.scheduler = self.model_backend.load_scheduler(args)
-        self.models = load_fsdp_models(
-            args,
-            self.model_backend,
-            self.train_pipeline_config,
-            self.parallel_state,
-            checkpoint_path=args.hf_checkpoint,
-            lora_adapter_paths=args.lora_adapter_paths,
-            trainable=True,
-            cpu_offload=args.fsdp_cpu_offload,
-        )
+        with st.step("train.load_fsdp_models"):
+            self.models = load_fsdp_models(
+                args,
+                self.model_backend,
+                self.train_pipeline_config,
+                self.parallel_state,
+                checkpoint_path=args.hf_checkpoint,
+                lora_adapter_paths=args.lora_adapter_paths,
+                trainable=True,
+                cpu_offload=args.fsdp_cpu_offload,
+                rollout_manager=rollout_manager,
+            )
         self.reference_models = {}
         if args.ref_load is not None:
             self.reference_models = load_fsdp_models(
@@ -124,6 +130,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 +142,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())
@@ -229,7 +237,9 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         checkpoint.finalize_load(self, checkpoint_payload)
 
         if self.args.offload_train:
-            self.sleep()
+            with st.step("train.offload_after_init"):
+                self.sleep()
+        st.mark("train.init_done")
 
         return self.args.start_rollout_id
 
@@ -283,12 +293,14 @@ def update_weights(self, previous_ema: bool = False) -> None:  # type: ignore[ov
             dist.barrier(group=get_gloo_group())
             return
 
-        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
-            self.rollout_manager.get_rollout_engines_and_lock.remote()
-        )
+        with st.step("train.uw.get_engines"):
+            rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+                self.rollout_manager.get_rollout_engines_and_lock.remote()
+            )
         if num_new_engines > 0:
-            self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
-            dist.barrier(group=get_gloo_group())
+            with st.step("train.uw.connect_engines"):
+                self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+                dist.barrier(group=get_gloo_group())
             if dist.get_rank() == 0:
                 ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
@@ -297,7 +309,7 @@ def update_weights(self, previous_ema: bool = False) -> None:  # type: ignore[ov
             if self.args.rollout_weights == "ema"
             else nullcontext()
         )
-        with rollout_weight_context:
+        with rollout_weight_context, st.step("train.uw.push_weights"):
             self.weight_updater.update_weights()
         clear_memory()
 

miles/backends/fsdp_utils/model_loader.py

@@ -10,6 +10,7 @@
 
 from . import checkpoint
 from .device_move import move_model
+from .startup_gate import wait_for_gpu_budget
 from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str
 from .sequence_parallel.plan import apply_sequence_parallel
 
@@ -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),

miles/backends/fsdp_utils/startup_gate.py

@@ -0,0 +1,102 @@
+"""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.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+import ray
+import torch
+import torch.distributed as dist
+
+import miles.utils.startup_timing as st
+
+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)
+    with st.step(
+        "train.budget_gate",
+        need_gib=f"{need / GIB:.1f}",
+        shard_gib=f"{shard / GIB:.1f}",
+        largest_gib=f"{largest / GIB:.2f}",
+    ):
+        if engine_reserve_gib is not None:
+            # Free memory already excludes whatever the engine has allocated, so subtracting the full reserve
+            # can only over-count: if this fits, the engine's remaining load fits too.
+            free = torch.cuda.mem_get_info()[0]
+            reserve = int(engine_reserve_gib * GIB)
+            if _all_ranks_agree(free - reserve >= need):
+                st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}", engines_resident=-1)
+                return
+        with st.step("train.budget_gate.wait_engines"):
+            ray.get(rollout_manager.wait_engines_ready.remote())
+        free = torch.cuda.mem_get_info()[0]
+        fits = _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,
+            )
+            with st.step("train.budget_gate.offload_engines"):
+                ray.get(rollout_manager.boot_offload_once.remote())
+            free = _wait_free(need, timeout)
+            if free < need:
+                raise RuntimeError(
+                    f"startup gate: only {free / GIB:.1f} GiB free after offloading the engines, need "
+                    f"{need / GIB:.1f} GiB (shard {shard / GIB:.1f} + largest tensor {largest / GIB:.2f} "
+                    f"+ margin {margin / GIB:.1f})"
+                )
+        st.mark("train.budget_gate.free", free_gib=f"{free / GIB:.1f}", engines_resident=int(fits))
+
+
+def _gloo_group():
+    from miles.utils.distributed_utils import get_gloo_group
+
+    return get_gloo_group()

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)