Startup: overlap engine boot and train init with a GPU-ready gate (invasive version)

docker/Dockerfile

@@ -88,3 +88,16 @@ RUN git clone https://github.com/radixark/miles_diffusion.git /root/miles_diffus
 # Pre-download PaddleOCR (en) det+rec+cls models so the OCR-reward actors don't
 # race-download them from bcebos at runtime. 
 RUN python -c "from paddleocr import PaddleOCR; PaddleOCR(use_angle_cls=False, lang='en', use_gpu=False, show_log=False)"
+
+# ====================================== Bytecode ======================================================
+
+# Bake .pyc files. About half of the package dirs ship without __pycache__, so every fresh container
+# (every CI job) recompiles bytecode on its first import. Measured on a fresh container:
+# `import sglang.multimodal_gen.runtime.launch_server` took 65 s without this step and 50 s with it
+# (38 s once warm). Some vendored files are not valid Python 3 source, so compile errors are ignored.
+RUN python -m compileall -q -j0 "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \
+        /sgl-workspace/sglang/python /root/miles_diffusion || true
+
+# Cache sglang-d's CLI argument definitions so parsing miles args does not import sglang (~25 s per
+# launch). The cache is keyed on sglang's server_args.py, so a different sglang just rebuilds it.
+RUN cd /root/miles_diffusion && python -m miles.backends.sglang_diffusion_utils.cli_cache || true

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)

miles/backends/fsdp_utils/actor.py

@@ -29,13 +29,6 @@
 
 from . import checkpoint
 from .device_move import move_model, move_optimizer, sleep_frozen_model, wake_up_frozen_model
-from .diffusion_update_weight_utils import (
-    DiffusionUpdateWeightFromDistributed,
-    DiffusionUpdateWeightFromTensor,
-    DiffusionUpdateWeightFromTensorLoRA,
-    DiffusionUpdateWeightFromTensorLoRAIPC,
-    DiffusionUpdateWeightLoRADistributed,
-)
 from .ema import EMAOptimizer, reshard_model
 from .input_dtype_policy import apply_input_dtype_policy
 from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch
@@ -66,7 +59,7 @@ class FSDPTrainRayActor(TrainRayActor):
     """
 
     @with_defer(lambda: Timer().start("train_wait"))
-    def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # type: ignore[override]
+    def init(self, args: Namespace, role: str, with_ref: bool = False, gpu_ready=None) -> int:  # type: ignore[override]
         super().init(args, role, with_ref)
 
         if args.deterministic_mode:
@@ -113,6 +106,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             lora_adapter_paths=args.lora_adapter_paths,
             trainable=True,
             cpu_offload=args.fsdp_cpu_offload,
+            gpu_ready=gpu_ready,
         )
         self.reference_models = {}
         if args.ref_load is not None:
@@ -124,6 +118,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 checkpoint_path=args.ref_load,
                 lora_adapter_paths=args.ref_lora_adapter_paths,
                 cpu_offload=args.ref_cpu_offload,
+                gpu_ready=gpu_ready,
             )
         self.teacher_models = {}
         if args.teacher_load is not None:
@@ -135,6 +130,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 checkpoint_path=args.teacher_load,
                 lora_adapter_paths=args.teacher_lora_adapter_paths,
                 cpu_offload=args.teacher_cpu_offload,
+                gpu_ready=gpu_ready,
             )
         # A natively CPU-offloaded model never sleeps: its parameters stay on the host and its buffers on the GPU.
         self.actor_models_to_sleep = [] if args.fsdp_cpu_offload else list(self.models.values())
@@ -209,6 +205,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             )
         checkpoint_payload = checkpoint.load(self)
 
+        # Imported here, not at module level: it pulls in sglang (~25 s), which the driver and other
+        # processes that only need this class or fsdp_utils' helpers should not pay for. Not in a
+        # background thread either: concurrent imports of transformers are not safe.
+        from .diffusion_update_weight_utils import (
+            DiffusionUpdateWeightFromDistributed,
+            DiffusionUpdateWeightFromTensor,
+            DiffusionUpdateWeightFromTensorLoRA,
+            DiffusionUpdateWeightFromTensorLoRAIPC,
+            DiffusionUpdateWeightLoRADistributed,
+        )
+
         # sglang-d now supports /update_weights_from_tensor (PR #20464).
         if self.args.train_only:
             self.weight_updater = None

miles/backends/fsdp_utils/checkpoint.py

@@ -33,24 +33,39 @@ 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,
         options=StateDictOptions(
             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
 

miles/backends/fsdp_utils/model_loader.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import warnings
 from argparse import Namespace
 from contextlib import contextmanager
@@ -12,6 +13,7 @@
 from .device_move import move_model
 from .mixed_precision import compile_param_dtype_maps, parse_dtype_from_str
 from .sequence_parallel.plan import apply_sequence_parallel
+from .startup_gate import wait_for_gpu_budget
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +28,7 @@ def load_fsdp_models(
     lora_adapter_paths: list[str] | None = None,
     trainable: bool = False,
     cpu_offload: bool = False,
+    gpu_ready=None,
 ) -> dict[str, torch.nn.Module]:
     materialize_weights = dist.get_rank() == 0
     master_dtype = parse_dtype_from_str(args.fsdp_master_dtype)
@@ -54,14 +57,26 @@ def load_fsdp_models(
             raise RuntimeError(f"{component} did not honor meta initialization")
         checkpoint.sync_model_dtypes(model)
         full_state = model.state_dict() if materialize_weights else {}
+        if gpu_ready is not None:
+            wait_for_gpu_budget(
+                model,
+                shard_world_size=parallel_state.get_mesh("fsdp").size(),
+                gpu_ready=gpu_ready,
+                engine_reserve_gib=args.startup_engine_reserve_gib,
+            )
         model = apply_fsdp2(
             model,
             model_backend.fsdp_parallel_plan(model),
             mesh=parallel_state.get_mesh("fsdp"),
             cpu_offload=cpu_offload,
             args=args,
         )
-        checkpoint.broadcast_full_state_to_fsdp(model, full_state, cpu_offload=cpu_offload)
+        checkpoint.broadcast_full_state_to_fsdp(
+            model,
+            full_state,
+            cpu_offload=cpu_offload,
+            checkpoint_files=_checkpoint_files(checkpoint_path, component),
+        )
         del full_state
         train_pipeline_config.postprocess_model_after_materialize(model)
         if parallel_state.get_optional_mesh("sp") is not None:
@@ -79,6 +94,29 @@ def load_fsdp_models(
     return models
 
 
+def _checkpoint_files(checkpoint_path: str, component: str) -> list[str] | None:
+    """Files for a sharded fill, or None; every rank gets the same answer."""
+    from . import sharded_checkpoint_fill
+
+    files = None
+    if sharded_checkpoint_fill.enabled():
+        try:
+            files = sharded_checkpoint_fill.resolve_component_files(checkpoint_path, component) or None
+        except Exception as exc:  # no local snapshot, odd layout: fall back to the broadcast
+            logger.info("No checkpoint files for a sharded fill of %s: %s", component, exc)
+    # Ranks decide what to fill on their own, so they must see the same files (e.g. a node without the
+    # snapshot would otherwise leave its ranks expecting a broadcast the others skip).
+    if dist.is_initialized() and dist.get_world_size() > 1:
+        from miles.utils.distributed_utils import get_gloo_group
+
+        mine = [os.path.basename(f) for f in files] if files else None
+        gathered = [None] * dist.get_world_size()
+        dist.all_gather_object(gathered, mine, group=get_gloo_group())
+        if any(g != gathered[0] for g in gathered) or gathered[0] is None:
+            return None
+    return files
+
+
 def finish_fsdp_lazy_init(model: torch.nn.Module) -> None:
     """Run FSDP's one-time lazy init now instead of at the first forward.