S24: sampler island, trainer start/stop, decoupled PPO, identity tiering

docker/miles-ports/Dockerfile

@@ -19,7 +19,7 @@
 FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c
 
 ARG MILES_REPOSITORY=https://github.com/michaellchung/miles
-ARG MILES_COMMIT=386ae3b07138ccd45ce4a07fc0383587f7968a3a
+ARG MILES_COMMIT=44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9
 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang
 ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121
 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b);
@@ -47,5 +47,17 @@ RUN cd /sgl-workspace/sglang \
  && SETUPTOOLS_SCM_PRETEND_VERSION="$SGLANG_VERSION" \
     SGLANG_BUILD_RUST_EXTS=none pip install -e "python[all]" --no-deps
 
+# peft at the pinned MILES_PEFT_VERSION, no deps (as the build script's
+# wheel unpack); island setup then reads it from the manifest, not by import.
+ARG PEFT_VERSION=0.20.0
+RUN pip install --no-deps "peft==$PEFT_VERSION"
+
+# Lazy modelopt (rl-trainer-cold-start-cost L2.2), as the crane layer: the
+# Megatron-Bridge package __init__ and the Megatron-LM patch.
+COPY megatron_bridge_init.py /opt/sglang/lib/python3.12/site-packages/megatron/bridge/__init__.py
+COPY megatron_lm_lazy_modelopt.patch /tmp/
+RUN cd /root/Megatron-LM && git apply /tmp/megatron_lm_lazy_modelopt.patch \
+ && rm /tmp/megatron_lm_lazy_modelopt.patch
+
 # Provenance (written by the build script: component SHAs + base digests).
 COPY --from=manifest image-manifest.json /opt/yeto/image-manifest.json

docker/miles-ports/make_layer.py

@@ -12,6 +12,15 @@
   fork repository (opaque whiteout hides the base history),
 * rewrites SGLang's editable-install metadata (dist-info, .pth, finder,
   ``sglang/_version.py``) for the fork's setuptools-scm version,
+* unpacks extra pure-Python wheels (``--wheel``, e.g. peft) into the base's
+  site-packages, as ``pip install --no-deps`` would,
+* replaces the base's ``megatron/bridge/__init__.py`` (Megatron-Bridge
+  8cd3466) with ``megatron_bridge_init.py``: the same file without its eager
+  ``import modelopt.torch`` (rl-trainer-cold-start-cost L2.2),
+* rewrites the Megatron-LM files (/root/Megatron-LM, base commit a84b105)
+  touched by ``megatron_lm_lazy_modelopt.patch`` -- the patch applied to the
+  commit's files from ``--megatron-lm-src`` -- so a trainer process does not
+  import modelopt unless a model uses it (L2.2),
 * adds /opt/yeto/image-manifest.json with every component SHA.
 
 Approximate docker equivalent (unverified): docker/miles-ports/Dockerfile.
@@ -30,6 +39,7 @@
 import subprocess
 import tarfile
 import tempfile
+import zipfile
 from pathlib import Path
 
 SITE = "opt/sglang/lib/python3.12/site-packages"
@@ -38,6 +48,14 @@
 # pyproject.toml, setup.cfg, requirements*, sgl-kernel/**, rust/**, csrc,
 # .cu/.cpp -- changes what a real install builds and needs a docker build.
 EXEMPT_PREFIXES = ("test/", "tests/", "docs/")
+# Megatron-Bridge as the base installs it (pip, not editable) and the file the
+# layer writes over its package __init__ (see megatron_bridge_init.py).
+MEGATRON_BRIDGE_COMMIT = "8cd3466d14d2337c8492827b3712482c2b3e4866"
+MEGATRON_BRIDGE_INIT = Path(__file__).with_name("megatron_bridge_init.py")
+# Megatron-LM as the base checks it out (/root/Megatron-LM, on sys.path) and
+# the patch the layer applies to it (see the comments it adds).
+MEGATRON_LM_COMMIT = "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e"
+MEGATRON_LM_PATCH = Path(__file__).with_name("megatron_lm_lazy_modelopt.patch")
 
 
 def git(repo: Path, *argv: str) -> str:
@@ -93,6 +111,45 @@ def version_file(version: str) -> bytes:
     ).encode()
 
 
+def wheel_files(wheel: Path) -> dict[str, bytes]:
+    """``{site-packages relative path: bytes}`` of a pure-Python wheel.
+
+    Only ``*-none-any`` wheels without a ``.data`` directory are accepted:
+    those install by unpacking alone (no scripts, headers or compiled code).
+    An ``INSTALLER`` file is added, as pip writes one.
+    """
+
+    if not wheel.name.endswith("-none-any.whl"):
+        raise SystemExit(f"{wheel.name}: only pure-Python (*-none-any) wheels are supported")
+    files: dict[str, bytes] = {}
+    with zipfile.ZipFile(wheel) as archive:
+        for name in archive.namelist():
+            if name.endswith("/"):
+                continue
+            if ".data/" in name.split(".dist-info/")[0]:
+                raise SystemExit(f"{wheel.name}: has a .data directory ({name}); not supported")
+            files[name] = archive.read(name)
+    dist_info = {n.split("/")[0] for n in files if n.split("/")[0].endswith(".dist-info")}
+    if len(dist_info) != 1:
+        raise SystemExit(f"{wheel.name}: expected one .dist-info directory, found {sorted(dist_info)}")
+    files[f"{dist_info.pop()}/INSTALLER"] = b"yeto-make-layer\n"
+    return files
+
+
+def patched_files(src: Path, commit: str, patch: Path) -> dict[str, bytes]:
+    """``{repo path: bytes}`` of the files ``patch`` touches, as ``commit`` plus the patch."""
+    paths = [line.split("\t")[-1] for line in subprocess.run(
+        ["git", "apply", "--numstat", str(patch)], check=True, capture_output=True, text=True
+    ).stdout.splitlines()]
+    with tempfile.TemporaryDirectory(prefix="img-patch-") as tmp:
+        archive = subprocess.run(["git", "-C", str(src), "archive", commit, *paths],
+                                 check=True, capture_output=True).stdout
+        with tarfile.open(fileobj=io.BytesIO(archive)) as tar:
+            tar.extractall(tmp, filter="data")
+        subprocess.run(["git", "apply", str(patch.resolve())], cwd=tmp, check=True, capture_output=True)
+        return {p: Path(tmp, p).read_bytes() for p in paths}
+
+
 def main() -> None:
     p = argparse.ArgumentParser()
     p.add_argument("--manifest-json", required=True, help="image-manifest.json content file")
@@ -108,6 +165,10 @@ def main() -> None:
     p.add_argument("--sglang-version", required=True)
     p.add_argument("--sglang-base-metadata", type=Path, required=True,
                    help="dir with the base image's sglang site-packages metadata")
+    p.add_argument("--wheel", type=Path, action="append", default=[],
+                   help="pure-Python wheel unpacked into site-packages (repeatable)")
+    p.add_argument("--megatron-lm-src", type=Path, required=True,
+                   help=f"clone containing Megatron-LM {MEGATRON_LM_COMMIT[:7]} (the base's checkout)")
     p.add_argument("--mtime", type=int, required=True)
     p.add_argument("--out", type=Path, required=True)
     a = p.parse_args()
@@ -184,6 +245,14 @@ def add_tree(root: Path, dest: str) -> None:
         add_file(f"{SITE}/{name}", data)
     add_file("sgl-workspace/sglang/python/sglang/_version.py", version_file(new))
 
+    for wheel in a.wheel:
+        for name, data in sorted(wheel_files(wheel).items()):
+            add_file(f"{SITE}/{name}", data)
+
+    add_file(f"{SITE}/megatron/bridge/__init__.py", MEGATRON_BRIDGE_INIT.read_bytes())
+    for name, data in patched_files(a.megatron_lm_src, MEGATRON_LM_COMMIT, MEGATRON_LM_PATCH).items():
+        add_file(f"root/Megatron-LM/{name}", data)
+
     add_dir("opt/yeto")
     add_file("opt/yeto/image-manifest.json", Path(a.manifest_json).read_bytes())
 

docker/miles-ports/megatron_bridge_init.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Megatron Bridge - A component of the Megatron ecosystem."""
+
+# yeto miles-ports image: the eager ``import modelopt.torch`` that Megatron-Bridge
+# 8cd3466 runs here is removed (rl-trainer-cold-start-cost L2.2).  It cost
+# 21-30 s per trainer process (coldstart-gpu-review §6.2).  Every bridge module
+# that uses modelopt (training.checkpointing, distillation, conversion.modelopt_utils)
+# imports it itself, so its plugins still load when a modelopt feature is used.
+# Written over the base image's file by docker/miles-ports/make_layer.py.
+
+import megatron.bridge.diffusion.models  # noqa: F401 — registers diffusion bridges
+import megatron.bridge.models  # noqa: F401 — triggers all bridge and HF Auto class registrations
+from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+from megatron.bridge.package_info import (
+    __contact_emails__,
+    __contact_names__,
+    __download_url__,
+    __homepage__,
+    __package_name__,
+    __repository_url__,
+    __version__,
+)
+
+
+__all__ = [
+    "__version__",
+    "__package_name__",
+    "__contact_names__",
+    "__contact_emails__",
+    "__homepage__",
+    "__repository_url__",
+    "__download_url__",
+    "AutoBridge",
+]

docker/miles-ports/megatron_lm_lazy_modelopt.patch

@@ -0,0 +1,190 @@
+diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py
+index 19df61b..70106de 100644
+--- a/megatron/post_training/checkpointing.py
++++ b/megatron/post_training/checkpointing.py
+@@ -5,15 +5,7 @@ import os
+ from pathlib import Path
+ from typing import Optional, Tuple, Union
+
+-import modelopt
+-import modelopt.torch.opt as mto
+ import torch.nn as nn
+-from modelopt.torch.opt.plugins import (
+-    restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy,
+-)
+-from modelopt.torch.opt.plugins.mcore_dist_checkpointing import (
+-    _load_extra_state_from_sharded_checkpoint,
+-)
+
+ from megatron.core import dist_checkpointing
+ from megatron.core.dist_checkpointing.serialization import _legacy_common_state_exists
+@@ -22,7 +14,11 @@ from megatron.training import get_args
+ from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint
+ from megatron.training.utils import print_rank_0
+
+-from .utils import print_distributed_quant_summary
++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported
++# modelopt (and .utils, which imports it) at module top.  megatron.training's
++# get_model imports this module for has_modelopt_state, a file check that needs
++# no modelopt, so every trainer paid for modelopt.  The functions that use
++# modelopt import it themselves.
+
+ logger = logging.getLogger(__name__)
+
+@@ -109,6 +105,11 @@ def load_modelopt_state(model: nn.Module, load_dir: Optional[str] = None) -> Non
+         model: the model to load the modelopt_state into
+         load_dir: optionally provide a different loading path
+     """
++    import modelopt.torch.opt as mto
++    from modelopt.torch.opt.plugins import (
++        restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy,
++    )
++
+     args = get_args()
+     load_dir = load_dir or args.load
+
+@@ -162,6 +163,8 @@ def load_modelopt_checkpoint(
+             When loading an .nemo sharded checkpoint, this is usually `model.`. Otherwise, this is
+             typically an empty string.
+     """
++    from .utils import print_distributed_quant_summary
++
+
+     def _remove_prefix_state_dict_pre_hook(
+         state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+@@ -213,6 +216,12 @@ def restore_sharded_modelopt_state(model: list[nn.Module], checkpoint_name: str
+     """Temporary function. Copy of modelopt.torch.opt.plugins.restore_sharded_modelopt_state.
+     Will be removed once modelopt.torch.opt.plugins.restore_sharded_modelopt_state is up to date.
+     """
++    import modelopt
++    import modelopt.torch.opt as mto
++    from modelopt.torch.opt.plugins.mcore_dist_checkpointing import (
++        _load_extra_state_from_sharded_checkpoint,
++    )
++
+     if len(model) > 1:
+         raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!")
+
+diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
+index 41865c5..8d4cc10 100644
+--- a/megatron/training/checkpointing.py
++++ b/megatron/training/checkpointing.py
+@@ -3,6 +3,7 @@
+ """Input/output checkpointing."""
+
+ import contextlib
++import importlib.util
+ import inspect
+ import multiprocessing
+ import os
+@@ -67,14 +68,31 @@ except ImportError:
+
+
+ # [ModelOpt]: Import
+-try:
+-    from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state
+-
+-    from megatron.post_training.utils import print_distributed_quant_summary
+-
+-    has_nvidia_modelopt = True
+-except Exception:
+-    has_nvidia_modelopt = False
++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported
++# modelopt at module top here, which every trainer process paid for (21-30 s).
++# Every use below is a no-op for a model ModelOpt has not converted
++# (save_modelopt_state / save_sharded_modelopt_state return early on
++# ``not is_converted``; the quant summary has nothing to report), and only
++# modelopt's own code converts a model.  So while modelopt is not imported in
++# this process, skipping them is exact; once it is, they run as before.
++has_nvidia_modelopt = importlib.util.find_spec("modelopt") is not None
++
++
++def _modelopt_api():
++    """ModelOpt checkpoint helpers, or None when no model here can carry ModelOpt state."""
++    if not has_nvidia_modelopt or "modelopt" not in sys.modules:
++        return None
++    try:
++        from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state
++
++        from megatron.post_training.utils import print_distributed_quant_summary
++    except Exception:
++        return None
++    return types.SimpleNamespace(
++        save_modelopt_state=save_modelopt_state,
++        save_sharded_modelopt_state=save_sharded_modelopt_state,
++        print_distributed_quant_summary=print_distributed_quant_summary,
++    )
+
+
+ _CHECKPOINT_VERSION = None
+@@ -882,8 +900,8 @@ def save_checkpoint(
+                 verify_integrity=args.verify_integrity,
+             )
+             # [ModelOpt]: save sharded modelopt_state
+-            if has_nvidia_modelopt:
+-                save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1))
++            if (modelopt_api := _modelopt_api()) is not None:
++                modelopt_api.save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1))
+         elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ["torch_dcp", "fsdp_dtensor"]:
+             if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+                 # TODO Handle non-empty directories (e.g., after a crash during saving).
+@@ -940,11 +958,11 @@ def save_checkpoint(
+                 )
+         else:
+             # [ModelOpt]: Inject modelopt_state into state_dict
+-            if has_nvidia_modelopt:
++            if (modelopt_api := _modelopt_api()) is not None:
+                 if ckpt_type == CheckpointType.LOCAL:
+                     print_rank_0('WARNING: Local checkpointing does not support nvidia_modelopt.')
+                 else:
+-                    save_modelopt_state(model, state_dict)
++                    modelopt_api.save_modelopt_state(model, state_dict)
+
+             end_ckpt = time()
+             logger.debug(
+@@ -2570,7 +2588,8 @@ def load_checkpoint(
+                 log_printed = True
+
+     if has_nvidia_modelopt:
+-        print_distributed_quant_summary(model, msg="After loading checkpoint")
++        if (modelopt_api := _modelopt_api()) is not None:
++            modelopt_api.print_distributed_quant_summary(model, msg="After loading checkpoint")
+
+         # Load teacher model in Distillation mode.
+         if getattr(args, "export_kd_teacher_load", None):
+diff --git a/megatron/training/training.py b/megatron/training/training.py
+index d64ae84..eb51d7d 100644
+--- a/megatron/training/training.py
++++ b/megatron/training/training.py
+@@ -13,6 +13,7 @@ import argparse
+ import copy
+ import dataclasses
+ import functools
++import importlib.util
+ import gc
+ import inspect
+ import logging
+@@ -204,12 +205,18 @@ try:
+ except ImportError:
+     has_rl_utils = False
+
+-try:
+-    from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation
++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported
++# modelopt at module top here, which every trainer process paid for (21-30 s).
++# The distillation helper is now imported where a modelopt-enabled run uses it.
++has_nvidia_modelopt = importlib.util.find_spec("modelopt") is not None
+
+-    has_nvidia_modelopt = True
+-except ImportError:
+-    has_nvidia_modelopt = False
++
++def get_tensor_shapes_adjust_fn_for_distillation(*args, **kwargs):
++    from modelopt.torch.distill.plugins.megatron import (
++        get_tensor_shapes_adjust_fn_for_distillation as adjust_fn,
++    )
++
++    return adjust_fn(*args, **kwargs)
+
+ try:
+     from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP

docs/MILES_RL.md

@@ -1021,6 +1021,8 @@ the run exit 3. A synced run still fetches its checkpoint first.
 | 4 | a fixed-roster RL island could not be recovered |
 | 5 | the Modal app was not confirmed stopped after teardown |
 | 6 | the run stalled: no island event for `--rl-stall-timeout` seconds (default 900, 0 disables) and not every island finalized |
+| 7 | Modal moved an island to a different container mid-run |
+| 8 | the trainer scheduler gave up: `--rl-trainer-max-failures` (default 3) trainer incarnations failed in a row |
 
 For exit 5, every row of `modal app list` with the run's app name must be
 `stopped` with 0 tasks; an earlier run's row with the same name counts too.