feat: async DiffusionNFT training with NCCL weight sync

miles/backends/fsdp_utils/actor.py

@@ -30,9 +30,11 @@
 
 from . import checkpoint
 from .diffusion_update_weight_utils import (
+    DiffusionUpdateWeightFromDistributed,
     DiffusionUpdateWeightFromTensor,
     DiffusionUpdateWeightFromTensorLoRA,
     DiffusionUpdateWeightFromTensorLoRAIPC,
+    DiffusionUpdateWeightLoRADistributed,
 )
 from .ema import EmaShadow
 from .input_dtype_policy import apply_input_dtype_policy
@@ -220,13 +222,24 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 uprate=self.args.ema_decay_ramp,
                 uphold=self.args.ema_decay_max,
                 flat_steps=self.args.ema_decay_flat_steps,
+                # Async prefetch uses the EMA from before the concurrent training update.
+                keep_previous_ema=(
+                    getattr(self.args, "train_async", False)
+                    and self.args.ref_mode == "ema"
+                    and self.args.ema_rollout_policy == "ema"
+                ),
             )
 
         # sglang-d now supports /update_weights_from_tensor (PR #20464).
         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)
+        elif not self.args.colocate:
+            updater = (
+                DiffusionUpdateWeightLoRADistributed if self.args.use_lora else DiffusionUpdateWeightFromDistributed
+            )
+            self.weight_updater = updater(self.args, self.models)
         elif self.args.use_lora:
             self.weight_updater = DiffusionUpdateWeightFromTensorLoRA(self.args, self.models)
         else:
@@ -304,10 +317,7 @@ def update_weights(self) -> None:  # type: ignore[override]
                 ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
         ema_shadow = self.ema_shadow
-        if ema_shadow is not None:
-            delta = ema_shadow.update()
-            if dist.get_rank() == 0:
-                logger.info("EMA shadow updated (decay=%.4f step=%d)", delta, ema_shadow.step)
+        # Publish the current EMA; the previous EMA is only a training reference.
         rollout_weight_context = (
             ema_shadow.swap_in() if ema_shadow is not None and self.args.ema_rollout_policy == "ema" else nullcontext()
         )
@@ -343,6 +353,10 @@ def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[over
             if self.args.debug_rollout_only:
                 return
             self._train_core(rollout_id=rollout_id, rollout_data=rollout_data)
+            if self.ema_shadow is not None:
+                delta = self.ema_shadow.update()
+                if dist.get_rank() == 0:
+                    logger.info("EMA shadow updated (decay=%.4f step=%d)", delta, self.ema_shadow.step)
 
         train_metric_utils.log_perf_data_raw(
             rollout_id=rollout_id,
@@ -555,7 +569,8 @@ def _compute_noise_pred() -> torch.Tensor:
         ref_mode = self.args.ref_mode
         if ref_mode != "none":
             if ref_mode == "ema":
-                ref_ctx = self.ema_shadow.swap_in()
+                # Match the EMA that generated the prefetched batch in async mode.
+                ref_ctx = self.ema_shadow.swap_in(use_previous_ema=self.ema_shadow.previous_ema is not None)
             else:
                 ref_ctx = prepared.model.disable_adapter()
             with torch.no_grad(), ref_ctx:

miles/backends/fsdp_utils/checkpoint.py

@@ -256,6 +256,7 @@ def load(actor: Any) -> dict[str, Any] | None:
         "rng": rng_state,
         "metadata": metadata,
         "iteration": target_step,
+        "checkpoint_dir": checkpoint_dir,
     }
 
 
@@ -264,6 +265,14 @@ def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None
         dist.barrier()
         return
 
+    if actor.ema_shadow is not None:
+        ema_dir = checkpoint_payload["checkpoint_dir"] / "ema"
+        if ema_dir.exists():
+            dcp.load({"ema": actor.ema_shadow}, checkpoint_id=str(ema_dir))
+        else:
+            logger.warning("Checkpoint has no EMA state; initializing EMA from the loaded model.")
+            actor.ema_shadow.step = checkpoint_payload["iteration"]
+
     if checkpoint_payload.get("rng") is not None and not actor.args.no_load_rng:
         rng_state = checkpoint_payload["rng"]
         if "torch" in rng_state:
@@ -315,6 +324,9 @@ def save(actor: Any, iteration: int) -> None:
     state_dict = {"model_state": model_state}
     dcp.save(state_dict, checkpoint_id=str(model_dir))
 
+    if actor.ema_shadow is not None:
+        dcp.save({"ema": actor.ema_shadow}, checkpoint_id=str(checkpoint_dir / "ema"))
+
     # --no-save-optim drops both the optimizer and the LR scheduler.
     if not actor.args.no_save_optim:
         allowed_missing = actor.train_pipeline_config.optimizer_state_allowed_missing

miles/backends/fsdp_utils/diffusion_update_weight_utils.py

@@ -2,8 +2,10 @@
 import logging
 import os
 import re
+import socket
 from argparse import Namespace
 from collections.abc import Mapping, Sequence
+from datetime import timedelta
 
 import ray
 import torch
@@ -16,7 +18,8 @@
 except ImportError:
     from sglang.srt.patch_torch import monkey_patch_torch_reductions  # type: ignore[import]
 
-from sglang.srt.utils import MultiprocessingSerializer
+from sglang.srt.utils import MultiprocessingSerializer, init_custom_process_group
+from sglang.srt.utils.network import NetworkAddress
 
 try:
     from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]
@@ -33,10 +36,9 @@
 
 from miles.ray.utils import get_physical_gpu_id
 
-
 logger = logging.getLogger(__name__)
 
-LORA_IPC_WEIGHT_UPDATE_MODE = "lora_merge"
+LORA_WEIGHT_UPDATE_MODE = "lora_merge"
 
 
 class PeftLoRAKeyMapper:
@@ -444,7 +446,7 @@ def _verify_weight_sync(self, pairs: list[tuple[str, torch.Tensor]], target_modu
         logger.warning(f"[weight_sync verify v{self.weight_version} cross-engine] " f"all_equal={all_equal}  {pretty}")
 
 
-class DiffusionUpdateWeightFromTensorLoRAIPC(DiffusionUpdateWeightFromTensor):
+class DiffusionUpdateWeightLoRA(DiffusionUpdateWeight):
     """Push only lora_A/lora_B tensors; rollout merges locally via weight_update_mode=lora_merge."""
 
     def _prepare_lora_param(self, param: torch.Tensor) -> torch.Tensor:
@@ -459,7 +461,7 @@ def _prepare_lora_param(self, param: torch.Tensor) -> torch.Tensor:
     def _collect_layer_groups(
         self, model: torch.nn.Module
     ) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]:
-        """Group PEFT LoRA tensors so each layer's A/B pair stays in one IPC bucket.
+        """Group PEFT LoRA tensors so each layer's A/B pair stays in one transfer bucket.
 
         Names stay PEFT/diffusers-shaped (``transformer_blocks.0.attn.to_q.lora_A``).
         sglang-d's ``lora_merge`` path applies ``param_names_mapping`` and the
@@ -483,7 +485,7 @@ def update_weights(self) -> None:
                     self.wait_and_update_bucket_weights(
                         bucket,
                         target_module,
-                        weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
+                        weight_update_mode=LORA_WEIGHT_UPDATE_MODE,
                     )
                     num_buckets += 1
                     bucket = []
@@ -497,7 +499,7 @@ def update_weights(self) -> None:
                 self.wait_and_update_bucket_weights(
                     bucket,
                     target_module,
-                    weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
+                    weight_update_mode=LORA_WEIGHT_UPDATE_MODE,
                 )
                 num_buckets += 1
 
@@ -507,7 +509,7 @@ def update_weights(self) -> None:
                 num_layers = len(layer_groups)
                 sample_layers = [PeftLoRAKeyMapper.layer_prefix(group[0][0]) for group in layer_groups[:3]]
                 logger.info(
-                    "LoRA IPC weight sync v%s [%s]: pushed %d lora tensors, "
+                    "LoRA weight sync v%s [%s]: pushed %d lora tensors, "
                     "%d layer prefixes in %d buckets (unmapped=%d)",
                     self.weight_version,
                     target_module,
@@ -518,18 +520,119 @@ def update_weights(self) -> None:
                 )
                 if sample_layers:
                     logger.info(
-                        "LoRA IPC [%s] sample layer prefixes: %s",
+                        "LoRA weight sync [%s] sample layer prefixes: %s",
                         target_module,
                         sample_layers,
                     )
                 if unmapped_keys:
                     logger.warning(
-                        "LoRA IPC unmapped PEFT keys [%s] (first 5): %s",
+                        "LoRA weight sync unmapped PEFT keys [%s] (first 5): %s",
                         target_module,
                         unmapped_keys[:5],
                     )
                 if num_lora_keys == 0:
                     logger.error(
-                        "LoRA IPC [%s]: no lora tensors found in training state_dict",
+                        "LoRA weight sync [%s]: no lora tensors found in training state_dict",
                         target_module,
                     )
+
+
+class DiffusionUpdateWeightFromTensorLoRAIPC(DiffusionUpdateWeightLoRA, DiffusionUpdateWeightFromTensor):
+    pass
+
+
+def connect_rollout_engines_from_distributed(rollout_engines, engine_gpu_counts, group_name, timeout):
+    if len(rollout_engines) != len(engine_gpu_counts) or any(count <= 0 for count in engine_gpu_counts):
+        raise ValueError("Each engine requires a positive GPU count")
+    master_address = ray._private.services.get_node_ip_address()
+    with socket.socket() as sock:
+        sock.bind(("", 0))
+        master_port = sock.getsockname()[1]
+    world_size = 1 + sum(engine_gpu_counts)
+    refs = []
+    rank_offset = 1
+    for engine, count in zip(rollout_engines, engine_gpu_counts, strict=True):
+        refs.append(
+            engine.init_weights_update_group.remote(
+                master_address=master_address,
+                master_port=master_port,
+                rank_offset=rank_offset,
+                world_size=world_size,
+                group_name=group_name,
+                backend="nccl",
+            )
+        )
+        rank_offset += count
+    options = dist.ProcessGroupNCCL.Options()
+    group = init_custom_process_group(
+        backend="nccl",
+        init_method=NetworkAddress(master_address, master_port).to_tcp(),
+        world_size=world_size,
+        rank=0,
+        group_name=group_name,
+        timeout=timeout,
+        pg_options=options,
+    )
+    # Custom groups span independent worlds and cannot split the default communicator.
+    options.split_from = None
+    ray.get(refs)
+    return group
+
+
+def broadcast_bucket(rollout_engines, group, group_name, named_tensors, target_module, **kwargs):
+    refs = [
+        engine.update_weights_from_distributed.remote(
+            names=[name for name, _ in named_tensors],
+            dtypes=[str(tensor.dtype).removeprefix("torch.") for _, tensor in named_tensors],
+            shapes=[list(tensor.shape) for _, tensor in named_tensors],
+            group_name=group_name,
+            target_modules=[target_module],
+            **kwargs,
+        )
+        for engine in rollout_engines
+    ]
+    tensors = [tensor.contiguous() for _, tensor in named_tensors]
+    handles = [dist.broadcast(tensor, src=0, group=group, async_op=True) for tensor in tensors]
+    for handle in handles:
+        handle.wait()
+    ray.get(refs)
+
+
+class DiffusionUpdateWeightFromDistributed(DiffusionUpdateWeight):
+    def __init__(self, args, models):
+        super().__init__(args, models)
+        self._model_update_group = None
+        self._group_name = "diffusion-weight-update"
+
+    def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
+        if dist.get_rank() != 0:
+            return
+        if self._model_update_group is not None:
+            refs = [engine.destroy_weights_update_group.remote(self._group_name) for engine in rollout_engines]
+            dist.destroy_process_group(self._model_update_group)
+            ray.get(refs)
+        self.rollout_engines = rollout_engines
+        self._model_update_group = connect_rollout_engines_from_distributed(
+            rollout_engines=rollout_engines,
+            engine_gpu_counts=[self.args.rollout_num_gpus_per_engine] * len(rollout_engines),
+            group_name=self._group_name,
+            timeout=timedelta(minutes=self.args.distributed_timeout_minutes),
+        )
+
+    def update_bucket_weights(self, named_tensors, target_module, weight_version=None, weight_update_mode=None):
+        if dist.get_rank() != 0:
+            return
+        broadcast_bucket(
+            rollout_engines=self.rollout_engines,
+            group=self._model_update_group,
+            group_name=self._group_name,
+            named_tensors=named_tensors,
+            target_module=target_module,
+            weight_update_mode=weight_update_mode,
+            lora_alpha=self.args.lora_alpha,
+            lora_rank=self.args.lora_rank,
+        )
+
+
+class DiffusionUpdateWeightLoRADistributed(DiffusionUpdateWeightLoRA, DiffusionUpdateWeightFromDistributed):
+    pass

miles/backends/fsdp_utils/ema.py

@@ -15,7 +15,13 @@ def _local(t: torch.Tensor) -> torch.Tensor:
 
 
 class EmaShadow:
-    """EMA shadow of trainable parameters."""
+    """EMA shadow of trainable parameters.
+
+    ``shadow`` holds the current EMA. With ``keep_previous_ema=True``,
+    ``previous_ema`` preserves the EMA from before the most recent ``update()``
+    for the async trainer's prefetched-batch reference. At initialization and
+    checkpoint restore, both snapshots start from the same weights.
+    """
 
     def __init__(
         self,
@@ -25,6 +31,7 @@ def __init__(
         uprate: float = 0.001,
         uphold: float = 0.5,
         flat_steps: int = 0,
+        keep_previous_ema: bool = False,
     ) -> None:
         self.decay = float(decay)
         self.uprate = float(uprate)
@@ -37,6 +44,7 @@ def __init__(
         if not self.params:
             raise ValueError("EmaShadow: model has no trainable parameters")
         self.shadow = [_local(p.detach()).clone() for p in self.params]
+        self.previous_ema = [sh.clone() for sh in self.shadow] if keep_previous_ema else None
 
     def decay_at(self, t: int) -> float:
         if t <= self.flat_steps:
@@ -50,24 +58,56 @@ def update(self) -> float:
             raise RuntimeError("EmaShadow.update called while swapped in")
         self.step += 1
         delta = self.decay_at(self.step)
+        if self.previous_ema is not None:
+            for previous_ema, current_ema in zip(self.previous_ema, self.shadow, strict=True):
+                previous_ema.copy_(current_ema)
         for live, sh in zip(self.params, self.shadow, strict=True):
             sh.mul_(delta).add_(_local(live.detach()).to(sh.device), alpha=1.0 - delta)
         return delta
 
+    def state_dict(self) -> dict:
+        # Preserve FSDP shard metadata so DCP can restore on a different mesh.
+        shadow = [
+            (
+                DTensor.from_local(
+                    sh,
+                    device_mesh=param.device_mesh,
+                    placements=param.placements,
+                    shape=param.shape,
+                    stride=param.stride(),
+                )
+                if isinstance(param, DTensor)
+                else sh
+            )
+            for param, sh in zip(self.params, self.shadow, strict=True)
+        ]
+        return {"shadow": shadow, "step": self.step}
+
+    @torch.no_grad()
+    def load_state_dict(self, state_dict: dict) -> None:
+        for sh, restored in zip(self.shadow, state_dict["shadow"], strict=True):
+            sh.copy_(_local(restored))
+        self.step = int(state_dict["step"])
+        # Resume starts a fresh pipeline: its first two batches use the restored EMA.
+        if self.previous_ema is not None:
+            for previous_ema, current_ema in zip(self.previous_ema, self.shadow, strict=True):
+                previous_ema.copy_(current_ema)
+
     @contextmanager
-    def swap_in(self):
-        """Temporarily expose EMA weights as the live parameters."""
-        self._swap()
+    def swap_in(self, use_previous_ema: bool = False):
+        """Temporarily use current EMA weights, or the snapshot before the last update."""
+        buffers = self.previous_ema if use_previous_ema else self.shadow
+        self._swap(buffers)
         self._swapped = True
         try:
             yield
         finally:
-            self._swap()
+            self._swap(buffers)
             self._swapped = False
 
     @torch.no_grad()
-    def _swap(self) -> None:
-        for live, sh in zip(self.params, self.shadow, strict=True):
+    def _swap(self, buffers: list[torch.Tensor]) -> None:
+        for live, sh in zip(self.params, buffers, strict=True):
             live_local = _local(live.data)
             tmp = live_local.clone()
             live_local.copy_(sh)