[DMD2 3/10] feat(ray): hand earlier trainer groups' outputs to later groups as external_data

docs/fsdp_model_roles.md

@@ -37,8 +37,9 @@ access `ctx.reference_models` by component name.
 
 A third source, the teacher (`--teacher-load`, with `--teacher-lora-adapter-path` for its
 adapters), is reserved for a DMD2-style loss that miles does not ship yet. Until then it
-has no built-in consumer: `--teacher-load` requires a custom prepare/loss hook, which reads
-it from `ctx.teacher_models`, and teacher inference should run inside `torch.no_grad()`.
+has no built-in consumer: `--teacher-load` requires a custom prepare/loss/train-output hook,
+which reads it from `ctx.teacher_models`, and teacher inference should run inside
+`torch.no_grad()`.
 
 Each source may load pretrained PEFT LoRA adapters, one plain PEFT directory (what
 `save_pretrained` writes) per `--update-weight-target-module` entry, in the same order:

docs/user-guide/cli-reference.md

@@ -273,11 +273,12 @@ See [Dtype Control](../advanced/dtype-control.md).
 | `--ref-load` | str | – | HF pipeline checkpoint of an independent frozen reference model. Requires `--ref-mode ref`. |
 | `--ref-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--ref-load`, one per `--update-weight-target-module` entry, in order. |
 | `--ref-cpu-offload` | flag | off | Keep the reference shards in pinned CPU memory between forwards (FSDP `CPUOffloadPolicy`). |
-| `--teacher-load` | str | – | HF pipeline checkpoint of a frozen teacher model, read only by custom prepare/loss hooks as `ctx.teacher_models`, so it requires `--custom-prepare-train-batch-path` or `--custom-loss-function-path`. |
+| `--teacher-load` | str | – | HF pipeline checkpoint of a frozen teacher model, read only by custom prepare/loss/train-output hooks as `ctx.teacher_models`, so it requires `--custom-prepare-train-batch-path`, `--custom-loss-function-path` or `--custom-train-output-function-path`. |
 | `--teacher-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--teacher-load`, one per `--update-weight-target-module` entry, in order. |
 | `--teacher-cpu-offload` | flag | off | Keep the teacher shards in pinned CPU memory between forwards. |
 | `--custom-prepare-train-batch-path` | str | – | Builds DiT inputs. |
 | `--custom-loss-function-path` | str | – | Loss **formula** only — the DiT forward stays in the actor. |
+| `--custom-train-output-function-path` | str | – | Runs without grad after a rollout's train steps; its return value reaches the paired rank of each later trainer group as `ctx.external_data[<role>]`. |
 | `--diffusion-nft-beta` | float | `1.0` | |
 | `--diffusion-nft-timestep-strategy-path` | str | `…nft.drop_random_steps` | `(args, num_steps, generator) -> list[int]` returning the denoising steps each sample trains on, in training order; `generator` is seeded per sample. Built in: `drop_random_steps` (random order), `drop_final_steps` (schedule order, noisiest first). The final clean output (σ = 0) is never a training step. |
 | `--diffusion-nft-num-dropped-timesteps` | int | `0` | Steps the built-in strategies drop: random ones per sample for `drop_random_steps`, the final (lowest-noise) ones for `drop_final_steps`. |

docs/user-guide/customization.md

@@ -25,6 +25,7 @@ signatures they expect, and the defaults they replace — in the same spirit as
 | | `--custom-convert-samples-to-train-data-path` | Full sample → train_data |
 | | `--custom-prepare-train-batch-path` | DiT batch preparation |
 | | `--custom-loss-function-path` | Loss formula only |
+| | `--custom-train-output-function-path` | Output handed to later trainer groups |
 | | `--sde-step-backend-path` | Train-side SDE dynamics |
 | **Logging** | `--custom-rollout-log-function-path` | Train-rollout logging |
 | | `--custom-eval-rollout-log-function-path` | Eval-rollout logging |
@@ -282,6 +283,17 @@ def loss_formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, metrics,
 `new_pred`, `old_pred` and `ref_pred` are keyed by the same stream names as `prepared.latents`.
 `old_pred` is `None` unless `--loss-type nft`; `ref_pred` is `None` with `--ref-mode none`.
 
+### `--custom-train-output-function-path`
+
+```python
+def train_output(ctx, train_pairs) -> Any:
+    ...
+```
+
+Runs without grad after this trainer group's train steps on a rollout. Every trainer group that
+trains later on the same rollout reads the return value of its paired rank (same rank, same
+train pairs) as `ctx.external_data[<role>]`.
+
 ### `--sde-step-backend-path`
 
 **Class** implementing `SdeStepBackend`. Auto-selected from

miles/backends/fsdp_utils/actor.py

@@ -1,6 +1,7 @@
 import logging
 from argparse import Namespace
 from contextlib import nullcontext
+from typing import Any
 
 import ray
 import torch
@@ -19,6 +20,7 @@
 from miles.utils.metric_buffer import MetricBuffer
 from miles.utils.metric_utils import compute_rollout_step
 from miles.utils.output_store import connect_output_store, materialize_stored_tensors, mooncake_client
+from miles.utils.ray_utils import Box
 from miles.utils.timer import Timer, inverse_timer, timer
 from miles.utils.tracking_utils import init_tracking
 from miles.utils.train_data_utils import (
@@ -175,6 +177,11 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         self.custom_loss_formula_func = (
             load_function(args.custom_loss_function_path) if args.custom_loss_function_path is not None else None
         )
+        self.custom_train_output_func = (
+            load_function(args.custom_train_output_function_path)
+            if args.custom_train_output_function_path is not None
+            else None
+        )
 
         if args.optimizer == "adam":
             self.optimizer = torch.optim.AdamW(
@@ -323,7 +330,7 @@ def _log_metrics(self, rollout_id: int, log_dict: dict[str, float], step: int) -
             )
         )
 
-    def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[override]
+    def train(self, rollout_id: int, rollout_data_ref, external_data) -> Any:  # type: ignore[override]
         if self.args.offload_train:
             self.wake_up()
 
@@ -333,7 +340,10 @@ def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[over
                 return
             if self.args.use_output_store:
                 materialize_stored_tensors(rollout_data["train_data"], store=self._output_store)
-            self._train_core(rollout_id=rollout_id, rollout_data=rollout_data)
+            external_data = {role: ray.get(box.inner) for role, box in external_data.items()}
+            train_output = self._train_core(
+                rollout_id=rollout_id, rollout_data=rollout_data, external_data=external_data
+            )
             if self.ema_optimizer is not None:
                 self.ema_optimizer.step(self.global_step)
 
@@ -342,8 +352,11 @@ def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[over
             args=self.args,
             is_primary_rank=dist.get_rank() == 0,
         )
+        # Boxed like the rollout shards: the driver's ray.get on this task only waits, the payload stays in the
+        # object store until the paired rank of a later group reads it.
+        return Box(ray.put(train_output))
 
-    def _train_core(self, rollout_id: int, rollout_data) -> None:
+    def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, Any]) -> Any:
         """Run the shared diffusion training loop."""
         device = torch.device("cuda", torch.cuda.current_device())
 
@@ -389,6 +402,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
             device=device,
             rollout_id=rollout_id,
             dp_rank=self.parallel_state.dp_rank,
+            external_data=external_data,
         )
 
         # ------------- Recompute old log-probs (impl-consistent PPO ratio) -------------
@@ -471,6 +485,10 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
                 reduced = {f"train/{key}": value for key, value in metrics.reduce().items()}
                 self._log_metrics(rollout_id, reduced, step=self.global_step)
 
+        if self.custom_train_output_func is not None:
+            with torch.no_grad():
+                return self.custom_train_output_func(loss_ctx, train_pairs)
+
     def _maybe_legacy_window_pad_len(self, train_pairs: list, microbatch_ranges: list) -> int | None:
         """LEGACY 2D parity: the whole-window max cond seq_len (like the legacy tile path), or
         None unless the legacy --micro-batch-size-sample>1 path is active. TODO: remove with it."""

miles/backends/fsdp_utils/loss_hub/types.py

@@ -25,6 +25,7 @@ class DiffusionLossContext:
     dp_rank: int = 0
     reference_models: dict[str, torch.nn.Module] = field(default_factory=dict)
     teacher_models: dict[str, torch.nn.Module] = field(default_factory=dict)
+    external_data: dict[str, Any] = field(default_factory=dict)
 
 
 @dataclass(frozen=True)

miles/ray/actor_group.py

@@ -100,9 +100,15 @@ def async_init(self, args, role, with_ref=False):
         self.args = args
         return [actor.init.remote(args, role, with_ref=with_ref) for actor in self._actor_handlers]
 
-    def async_train(self, rollout_id, rollout_data_ref):
-        """Do one rollout training"""
-        return [actor.train.remote(rollout_id, rollout_data_ref) for actor in self._actor_handlers]
+    def async_train(self, rollout_id, rollout_data_ref, external_data=None):
+        """Do one rollout training; rank i receives rank i's output of each group in ``external_data``."""
+        external_data = external_data or {}
+        return [
+            actor.train.remote(
+                rollout_id, rollout_data_ref, {role: refs[rank] for role, refs in external_data.items()}
+            )
+            for rank, actor in enumerate(self._actor_handlers)
+        ]
 
     def save_model(self, rollout_id, force_sync=False):
         """Save actor model"""

miles/ray/train_actor.py

@@ -122,7 +122,7 @@ def wake_up(self, tags):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def train(self, rollout_id, rollout_data_ref):
+    def train(self, rollout_id, rollout_data_ref, external_data):
         raise NotImplementedError
 
     @abc.abstractmethod

miles/utils/arguments.py

@@ -971,6 +971,16 @@ def add_algo_arguments(parser):
                     "DiT forward is owned by the actor; this hook only computes the objective."
                 ),
             )
+            parser.add_argument(
+                "--custom-train-output-function-path",
+                type=str,
+                default=None,
+                help=(
+                    "Dotted path to `fn(ctx, train_pairs) -> Any`, run without grad after the train steps of "
+                    "a rollout. Its return value reaches the paired rank of every trainer group that trains "
+                    "later on the same rollout as `ctx.external_data[<role>]`."
+                ),
+            )
             parser.add_argument(
                 "--diffusion-nft-beta",
                 type=float,
@@ -1798,8 +1808,11 @@ def validate_reference_model_args(args) -> None:
         args.teacher_load is not None
         and args.custom_prepare_train_batch_path is None
         and args.custom_loss_function_path is None
+        and args.custom_train_output_function_path is None
     ):
-        raise ValueError("--teacher-load has no consumer; only custom prepare/loss hooks read ctx.teacher_models")
+        raise ValueError(
+            "--teacher-load has no consumer; only custom prepare/loss/train-output hooks read ctx.teacher_models"
+        )
     if args.ref_mode == "lora_base" and not args.use_lora:
         raise ValueError("--ref-mode lora_base requires --use-lora")
 

tests/fast/backends/fsdp_utils/test_ema_optimizer.py

@@ -21,6 +21,7 @@
 from contextlib import nullcontext
 from pathlib import Path
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -29,6 +30,7 @@
 
 from miles.backends.fsdp_utils import checkpoint
 from miles.backends.fsdp_utils.ema import EMAOptimizer
+from miles.utils.ray_utils import Box
 
 
 class Model(nn.Module):
@@ -146,7 +148,9 @@ def test_actor_train_updates_ema_once_and_publication_only_reads_it():
     for method in methods:
         method.decorator_list = []
     namespace = dict(
-        ray=SimpleNamespace(get=lambda value: value),
+        Any=Any,
+        Box=Box,
+        ray=SimpleNamespace(get=lambda value: value, put=lambda value: value),
         dist=SimpleNamespace(get_rank=lambda: 0),
         timer=lambda name: nullcontext(),
         inverse_timer=lambda name: nullcontext(),
@@ -181,7 +185,7 @@ def train_core(**kwargs):
         rollout_manager=SimpleNamespace(get_rollout_engines_and_lock=SimpleNamespace(remote=lambda: (None, None, 0))),
         weight_updater=SimpleNamespace(update_weights=lambda: published_weights.append(model.base.detach().clone())),
     )
-    namespace["train"](actor, rollout_id=0, rollout_data_ref=[SimpleNamespace(inner={})])
+    namespace["train"](actor, rollout_id=0, rollout_data_ref=[SimpleNamespace(inner={})], external_data={})
     assert model.base.item() == 4.0
     assert actor.ema_optimizer.state[model.base]["ema"].item() == 3.0
 

tests/fast/backends/fsdp_utils/test_train_output_handoff.py

@@ -0,0 +1,67 @@
+"""FSDPTrainRayActor.train hands its output to later groups through the object store, not through the driver."""
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
+
+import ast
+from contextlib import nullcontext
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+from miles.utils.ray_utils import Box
+
+
+def _load_train():
+    path = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/actor.py"
+    actor_class = next(
+        node
+        for node in ast.parse(path.read_text()).body
+        if isinstance(node, ast.ClassDef) and node.name == "FSDPTrainRayActor"
+    )
+    train = next(node for node in actor_class.body if isinstance(node, ast.FunctionDef) and node.name == "train")
+    train.decorator_list = []
+    store = {}
+
+    def put(value):
+        ref = f"ref{len(store)}"
+        store[ref] = value
+        return ref
+
+    namespace = dict(
+        Any=Any,
+        Box=Box,
+        ray=SimpleNamespace(get=store.__getitem__, put=put),
+        dist=SimpleNamespace(get_rank=lambda: 0),
+        timer=lambda name: nullcontext(),
+        inverse_timer=lambda name: nullcontext(),
+        train_metric_utils=SimpleNamespace(log_perf_data_raw=lambda **kwargs: None),
+    )
+    exec(compile(ast.Module(body=[train], type_ignores=[]), str(path), "exec"), namespace)
+    return namespace["train"], store, put
+
+
+def test_train_unboxes_external_data_and_boxes_its_output():
+    train, store, put = _load_train()
+    seen = {}
+
+    def train_core(rollout_id, rollout_data, external_data):
+        seen.update(rollout_data=rollout_data, external_data=external_data)
+        return {"scores": [rollout_id]}
+
+    actor = SimpleNamespace(
+        args=SimpleNamespace(offload_train=False, debug_rollout_only=False, use_output_store=False),
+        ema_optimizer=None,
+        _train_core=train_core,
+        parallel_state=SimpleNamespace(get_mesh=lambda name: SimpleNamespace(get_local_rank=lambda: 0)),
+    )
+    shard = Box(put({"train_data": []}))
+    critic_output = Box(put({"scores": [0]}))
+
+    result = train(actor, rollout_id=3, rollout_data_ref=[shard], external_data={"critic": critic_output})
+
+    assert seen == {"rollout_data": {"train_data": []}, "external_data": {"critic": {"scores": [0]}}}
+    assert isinstance(result, Box)
+    assert store[result.inner] == {"scores": [3]}
+    assert result.inner not in (shard.inner, critic_output.inner)

tests/fast/rollout/test_async_training.py

@@ -94,7 +94,7 @@ def update_weights(self, previous_ema=False):
             self.actor.FSDPTrainRayActor.update_weights(self, previous_ema=previous_ema)
         return self.published, self.global_step
 
-    def _train_core(self, rollout_id, rollout_data):
+    def _train_core(self, rollout_id, rollout_data, external_data):
         start = time.monotonic()
         with self.ema_optimizer.use_weights(self.model):
             reference = self.param.item()
@@ -109,7 +109,7 @@ def train(self, rollout_id, batch):
         with patch.object(self.actor.dist, "get_rank", return_value=0), patch.object(
             self.actor.train_metric_utils, "log_perf_data_raw"
         ):
-            self.actor.FSDPTrainRayActor.train(self, rollout_id, batch)
+            self.actor.FSDPTrainRayActor.train(self, rollout_id, batch, {})
 
     def save_model(self, rollout_id, force_sync=False):
         from copy import deepcopy

tests/fast/rollout/test_placement_group.py

@@ -3,10 +3,12 @@
 register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
 
 from argparse import Namespace
+from types import SimpleNamespace
 
 import pytest
 
 import miles.ray.placement_group as placement_module
+from miles.ray.actor_group import RayTrainGroup
 from miles.ray.utils import TRAIN_ACTOR_GPU
 from miles.utils.trainer_config import TrainerConfig
 
@@ -107,3 +109,19 @@ def test_trainer_groups_split_the_train_gpu_share(monkeypatch, roles):
 def test_trainers_restored_to_different_rollouts_are_rejected(monkeypatch):
     with pytest.raises(AssertionError, match="trainers restored different rollouts"):
         _create_training_models(monkeypatch, {"critic": 2, "actor": 3})
+
+
+def test_each_rank_receives_the_same_rank_output_of_earlier_groups():
+    class _Handle:
+        def __init__(self, rank):
+            self.train = SimpleNamespace(remote=lambda *args: (rank, *args))
+
+    group = object.__new__(RayTrainGroup)
+    group._actor_handlers = [_Handle(rank) for rank in range(2)]
+    external_data = {"critic": ["critic_0", "critic_1"], "reward": ["reward_0", "reward_1"]}
+
+    assert group.async_train(7, "data", external_data) == [
+        (0, 7, "data", {"critic": "critic_0", "reward": "reward_0"}),
+        (1, 7, "data", {"critic": "critic_1", "reward": "reward_1"}),
+    ]
+    assert group.async_train(7, "data") == [(0, 7, "data", {}), (1, 7, "data", {})]

tests/fast/test_train_diffusion.py

@@ -1,4 +1,5 @@
-"""train_diffusion.py trains every trainer group on each rollout in declared order, then releases the rollout."""
+"""train_diffusion.py trains every trainer group on each rollout in declared order, handing each the earlier
+outputs, then releases the rollout."""
 
 from tests.ci.ci_register import register_cpu_ci
 
@@ -32,8 +33,9 @@ def __init__(self, role, events):
         self._role = role
         self._events = events
 
-    def async_train(self, rollout_id, rollout_data_ref):
-        self._events.append(("train_start", self._role, rollout_id, rollout_data_ref))
+    def async_train(self, rollout_id, rollout_data_ref, external_data):
+        received = {role: [(ref.role, ref.rollout_id) for ref in refs] for role, refs in external_data.items()}
+        self._events.append(("train_start", self._role, rollout_id, rollout_data_ref, received))
         return [_Ref(self._role, rollout_id)]
 
     def save_model(self, rollout_id, force_sync=False):
@@ -86,9 +88,10 @@ def test_every_trainer_group_trains_each_rollout_in_order(monkeypatch):
 
     expected = [("update_weights", "actor")]
     for rollout_id in range(2):
-        for role in roles:
+        for i, role in enumerate(roles):
+            received = {earlier: [(earlier, rollout_id)] for earlier in roles[:i]}
             expected += [
-                ("train_start", role, rollout_id, ("rollout_data", rollout_id)),
+                ("train_start", role, rollout_id, ("rollout_data", rollout_id), received),
                 ("train_done", role, rollout_id),
             ]
         expected += [("release", rollout_id)]