feat(diffusion): add SFT loss hub and pre-encoded data manager

miles/backends/fsdp_utils/actor.py

@@ -223,7 +223,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             )
 
         # sglang-d now supports /update_weights_from_tensor (PR #20464).
-        if self.args.debug_train_only:
+        if self.args.train_only:
             self.weight_updater = None
         elif self.args.use_lora and self.args.lora_ipc_weight_sync:
             self.weight_updater = DiffusionUpdateWeightFromTensorLoRAIPC(self.args, self.models)
@@ -292,7 +292,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:  # type
 
     @timer
     def update_weights(self) -> None:  # type: ignore[override]
-        if self.args.debug_train_only or self.args.debug_rollout_only:
+        if self.args.train_only or self.args.debug_rollout_only:
             return
 
         if self.weight_updater is None:

miles/backends/fsdp_utils/loss_hub/sft.py

@@ -0,0 +1,99 @@
+"""SFT batch preparation and loss formula (rectified-flow velocity MSE)."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch
+from miles.backends.fsdp_utils.loss_hub.utils import cast_cond_to_dtype
+from miles.utils.metric_buffer import MetricBuffer
+
+
+def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]:
+    """Pick one DiT component per micro-batch (phase-pure), then grid indices within its range."""
+    num_grid = len(ctx.scheduler.timesteps)
+    if len(ctx.models) == 1:
+        component_name, model = next(iter(ctx.models.items()))
+        return component_name, model, torch.randint(num_grid, (bsz,))
+
+    # A uniform anchor index picks each expert with probability equal to its share of the
+    # grid, keeping the marginal over indices uniform while the micro-batch stays single-expert.
+    num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps)
+    config = ctx.train_pipeline_config
+    components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps]
+    component_name = components[int(torch.randint(num_grid, (1,)))]
+    pool = torch.tensor([i for i, name in enumerate(components) if name == component_name])
+    return component_name, ctx.models[component_name], pool[torch.randint(len(pool), (bsz,))]
+
+
+def prepare_sft_batch(
+    ctx: DiffusionLossContext,
+    batch: list[dict],
+    *,
+    pad_to_len: int | None = None,
+) -> PreparedBatch:
+    """Corrupt cached clean latents at sampled grid sigmas; CFG-free cached cond."""
+    device = ctx.device
+    config = ctx.train_pipeline_config
+    bsz = len(batch)
+
+    x0 = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32)
+    component_name, model, idx = sample_grid_indices(ctx, bsz)
+    idx = idx.to(device)
+    timesteps = ctx.scheduler.timesteps[idx].to(dtype=torch.float32)
+    sigmas = ctx.scheduler.sigmas[idx].to(dtype=torch.float32)
+
+    noise = torch.randn(x0.shape, device=device, dtype=torch.float32)
+    sigma_exp = sigmas.view(bsz, *([1] * (x0.ndim - 1)))
+    latents = (1.0 - sigma_exp) * x0 + sigma_exp * noise
+
+    num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps)
+    if config.needs_timestep_scaling:
+        timesteps_for_model = timesteps / float(num_train_timesteps)
+    else:
+        timesteps_for_model = timesteps
+
+    cond_list = [{key: value.to(device) for key, value in pair["cond_kwargs"].items()} for pair in batch]
+    pos_cond = cast_cond_to_dtype(
+        config.collate_cond_for_sample_batch(cond_list, device, pad_to_len=pad_to_len),
+        ctx.forward_dtype,
+    )
+
+    return PreparedBatch(
+        latents=latents,
+        timesteps=timesteps,
+        timesteps_for_model=timesteps_for_model,
+        model=model,
+        component_name=component_name,
+        guidance_scale=0.0,
+        use_cfg=False,
+        cfg_batching=False,
+        true_cfg_scale=None,
+        pos_cond=pos_cond,
+        neg_cond=None,
+        joint_cond=None,
+        advantage=torch.ones(bsz, device=device, dtype=torch.float32),
+        extras={"target": noise - x0},
+    )
+
+
+def sft_loss_formula(
+    ctx: DiffusionLossContext,
+    batch: list[dict],
+    prepared: PreparedBatch,
+    *,
+    new_pred: torch.Tensor,
+    ref_pred: torch.Tensor | None,
+    metrics: MetricBuffer,
+    write_old_log_prob: bool = False,
+    old_log_prob_from_new: bool = False,
+) -> torch.Tensor:
+    """Velocity-target MSE: ``||pred - (eps - x0)||^2`` averaged per pair."""
+    target = prepared.extras["target"]
+    per_pair = ((new_pred.float() - target) ** 2).mean(dim=tuple(range(1, target.ndim)))
+    loss_sum = per_pair.sum()
+
+    with torch.no_grad():
+        metrics.emit_mean("loss", total=loss_sum, count=len(batch))
+    return loss_sum

miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py

@@ -1,17 +1,22 @@
+from __future__ import annotations
+
 import dataclasses
 import ipaddress
 import logging
 import multiprocessing
 import os
 import time
+from typing import TYPE_CHECKING
 
 import requests
-from sglang.multimodal_gen.runtime.launch_server import kill_process_tree
-from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.srt.utils.common import kill_process_tree
 
 from miles.ray.ray_actor import RayActor
 from miles.utils.http_utils import get_host_info
 
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
 logger = logging.getLogger(__name__)
 
 
@@ -139,6 +144,8 @@ def _format_v6_uri(addr):
         self._init_normal(server_args_dict)
 
     def _init_normal(self, server_args_dict):
+        from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
         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
@@ -300,6 +307,8 @@ def simulate_crash(self):
 
 
 def _compute_server_args(args, host, port, nccl_port):
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
     # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done
     # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg.
     kwargs = {

miles/ray/placement_group.py

@@ -82,13 +82,16 @@ def create_placement_groups(args):
     """Create placement groups for actor and rollout engines.
 
     Two topologies:
-    - Colocate (or --debug-{train,rollout}-only): one combined placement
-      group; both roles see the same bundle list.
-    - Disaggregate (the else branch): two separate placement groups so
-      train and rollout each own a disjoint GPU pool — avoids bundle
-      overlap / scheduling deadlock when running side-by-side.
+    - Colocate: one combined placement group; both roles see the same bundles.
+      Train-only jobs use this by default.
+    - Disaggregate: separate actor and rollout placement groups. A train-only
+      job opts into this by setting --rollout-num-gpus, reserving those GPUs for
+      its rollout-side producer (for example, the SFT encoder pool).
     """
-    if not args.colocate and not args.debug_train_only and not args.debug_rollout_only:
+    disaggregate = (
+        not args.colocate and not args.debug_rollout_only and (not args.train_only or bool(args.rollout_num_gpus))
+    )
+    if disaggregate:
         logger.info("Creating placement groups (separate actor/rollout)...")
         actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node
         rollout_gpus = args.rollout_num_gpus
@@ -107,14 +110,11 @@ def create_placement_groups(args):
     logger.info(f"Creating placement group with {num_gpus} GPUs...")
     pg, all_reordered_bundle_indices, all_reordered_gpu_ids = _create_placement_group(num_gpus)
 
-    actor_pg_reordered_bundle_indices = all_reordered_bundle_indices
-    actor_pg_reordered_gpu_ids = all_reordered_gpu_ids
-    rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.debug_train_only else []
-    rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.debug_train_only else []
-
+    # The rollout view keeps its seats under train_only: engine startup is gated by
+    # args.train_only, and rollout-side actor pools (e.g. the SFT encoder pool) seat here.
     return {
-        "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids),
-        "rollout": (pg, rollout_pg_reordered_bundle_indices, rollout_pg_reordered_gpu_ids),
+        "actor": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids),
+        "rollout": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids),
     }
 
 

miles/ray/rollout.py

@@ -17,7 +17,7 @@
     expand_samples_to_train_pairs as flow_grpo_expand_samples_to_train_pairs,
 )
 from miles.rollout.base_types import call_rollout_fn
-from miles.rollout.rm_hub.core import set_reward_placement_group
+from miles.rollout.rm_hub.core import set_manager_placement_group
 from miles.utils import tracking_utils
 from miles.utils.health_monitor import RolloutHealthMonitor
 from miles.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client
@@ -54,9 +54,10 @@ def __init__(self, args, pg):
         from miles.dashboard import hooks
 
         hooks.register_rollout_manager(args)
-        set_reward_placement_group(pg)
-        logger.info("RolloutManager: starting router...")
-        _start_router(args)
+        set_manager_placement_group(pg)
+        if not args.train_only:
+            logger.info("RolloutManager: starting router...")
+            _start_router(args)
         logger.info("RolloutManager: router started, init tracking...")
         # TODO make args immutable
         init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}")
@@ -91,10 +92,10 @@ def __init__(self, args, pg):
 
         logger.info("RolloutManager rollout_num_gpus=%s", self.args.rollout_num_gpus)
 
-        if self.args.debug_train_only:
+        if self.args.train_only:
             self.all_rollout_engines = []
             self.num_new_engines = 0
-            logger.info("RolloutManager using no sglang engines (debug_train_only).")
+            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
@@ -200,7 +201,7 @@ def generate(self, rollout_id):
         return [Box(ray.put(shard)) for shard in shards]
 
     def eval(self, rollout_id):
-        if self.args.debug_train_only:
+        if self.args.train_only:
             # if debug train only, we don't generate evaluation data
             return
         self.health_monitoring_resume()
@@ -474,7 +475,7 @@ def set_train_parallel_config(self, config: dict):
 
 
 def init_rollout_engines(args, pg, all_rollout_engines):
-    if args.debug_train_only:
+    if args.train_only:
         return 0
 
     num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node)

miles/rollout/rm_hub/core.py

@@ -8,17 +8,22 @@
 import ray
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
-_reward_placement_group = None
+_manager_placement_group = None
 logger = logging.getLogger(__name__)
 
 
-def set_reward_placement_group(pg) -> None:
-    global _reward_placement_group
-    _reward_placement_group = pg
+def set_manager_placement_group(pg) -> None:
+    """Publish the manager's (pg, bundle_indices, gpu_ids) for colocated actor pools."""
+    global _manager_placement_group
+    _manager_placement_group = pg
 
 
-def get_reward_placement_group():
-    return _reward_placement_group
+def get_manager_placement_group():
+    return _manager_placement_group
+
+
+set_reward_placement_group = set_manager_placement_group
+get_reward_placement_group = get_manager_placement_group
 
 
 class AsyncRewardActorPool: