From 29fa6afa998e1d64b12a21df51a381d3e6ae6bd6 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 18:55:05 +0000 Subject: [PATCH 01/14] feat(fsdp): CPU-staged LoRA weight sync for disaggregated rollout engines --- miles/backends/fsdp_utils/actor.py | 3 + .../diffusion_update_weight_utils.py | 57 +++++++++++++++++-- .../sglang_diffusion_engine.py | 4 +- miles/ray/rollout.py | 2 +- miles/utils/arguments.py | 6 ++ tests/fast/utils/test_lora_args.py | 6 ++ 6 files changed, 71 insertions(+), 7 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f3eab3694..5273fd204 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -32,6 +32,7 @@ from .diffusion_update_weight_utils import ( DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA, + DiffusionUpdateWeightFromTensorLoRACPU, DiffusionUpdateWeightFromTensorLoRAIPC, ) from .ema import EmaShadow @@ -227,6 +228,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty 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 self.args.use_lora and not self.args.colocate: + self.weight_updater = DiffusionUpdateWeightFromTensorLoRACPU(self.args, self.models) elif self.args.use_lora: self.weight_updater = DiffusionUpdateWeightFromTensorLoRA(self.args, self.models) else: diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index ff45c2136..e4dd2eb9a 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -524,7 +524,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, @@ -535,18 +535,67 @@ 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 DiffusionUpdateWeightFromTensorLoRACPU(DiffusionUpdateWeightFromTensorLoRAIPC): + """LoRA sync for disaggregated rollout: rank 0 posts CPU-staged adapters to every engine.""" + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle | None, + ) -> None: + self.rollout_engines = rollout_engines + + def update_bucket_weights( + self, + named_tensors, + target_module: str, + weight_version=None, + weight_update_mode: str | None = None, + ) -> None: + if dist.get_rank() != 0: + return + + named_tensors_by_dtype: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in named_tensors: + named_tensors_by_dtype.setdefault(tensor.dtype, []).append((name, tensor.cpu())) + + for group in named_tensors_by_dtype.values(): + bucket = FlattenedTensorBucket(named_tensors=group) + payload = { + target_module: { + "flattened_tensor": bucket.get_flattened_tensor(), + "metadata": bucket.get_metadata(), + } + } + serialized = MultiprocessingSerializer.serialize(payload, output_str=True) + kwargs = { + # A single unlabeled payload: every engine worker deserializes the same + # full adapter set (no shared-GPU visibility required). + "serialized_named_tensors": [serialized], + "payload_gpu_uuids": None, + "load_format": "flattened_bucket", + "target_modules": [target_module], + "weight_version": str(weight_version), + } + if weight_update_mode is not None: + kwargs["weight_update_mode"] = weight_update_mode + kwargs["lora_alpha"] = self.args.lora_alpha + kwargs["lora_rank"] = self.args.lora_rank + ray.get([engine.update_weights_from_tensor.remote(**kwargs) for engine in self.rollout_engines]) diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2f9cdcb87..2fee21091 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -234,7 +234,7 @@ def health_generate(self, timeout: float = 5.0) -> bool: def update_weights_from_tensor( self, serialized_named_tensors: list[str], - payload_gpu_uuids: list[str], + payload_gpu_uuids: list[str] | None, load_format: str | None = None, target_modules: list[str] | None = None, weight_version: str | None = None, @@ -339,7 +339,7 @@ def _compute_server_args(args, host, port, nccl_port): if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs: kwargs[attr.name] = getattr(args, f"sglang_{attr.name}") - if getattr(args, "use_lora", False) and getattr(args, "lora_ipc_weight_sync", False): + if args.use_lora and (args.lora_ipc_weight_sync or not args.colocate): kwargs["lora_target_modules"] = args.lora_target_modules # dit_precision / vae_precision are PipelineConfig fields, not ServerArgs, so forward them explicitly (only when changed from the class default, to avoid clobbering a subclass override). from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index cef95d3f0..2d1eabf2e 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -516,7 +516,7 @@ def init_rollout_engines(args, pg, all_rollout_engines): "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false", } - if args.lora_ipc_weight_sync: + if args.use_lora and (args.lora_ipc_weight_sync or not args.colocate): # Merge in the train forward dtype, not fp32, to cut train/rollout consistency error. env_vars["SGLANG_DIFFUSION_LORA_MERGE_FP32"] = "1" if args.diffusion_forward_dtype == "fp32" else "0" diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index d96eb6599..147535711 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1735,6 +1735,12 @@ def miles_validate_args(args): if args.offload_rollout is None: args.offload_rollout = False + if not args.colocate and not args.train_only and not args.debug_rollout_only: + if args.lora_ipc_weight_sync: + raise ValueError("--lora-ipc-weight-sync requires --colocate: CUDA IPC needs shared train/rollout GPUs") + if not args.use_lora: + raise ValueError("disaggregated train/rollout weight sync supports LoRA only; pass --use-lora or --colocate") + if args.colocate_reward: assert args.colocate, "--colocate-reward requires --colocate." assert args.pickscore_num_workers <= args.rollout_num_gpus, ( diff --git a/tests/fast/utils/test_lora_args.py b/tests/fast/utils/test_lora_args.py index a3615f779..8dcf3124c 100644 --- a/tests/fast/utils/test_lora_args.py +++ b/tests/fast/utils/test_lora_args.py @@ -18,6 +18,7 @@ def _server_args(**overrides): use_lora=True, lora_ipc_weight_sync=True, lora_target_modules=["to_q", "to_k"], + colocate=True, ) base.update(overrides) return Namespace(**base) @@ -33,3 +34,8 @@ def test_lora_ipc_omitted_when_disabled(self): args = _server_args(lora_ipc_weight_sync=False) kwargs = _compute_server_args(args, "127.0.0.1", 15000, 15001) assert "lora_target_modules" not in kwargs + + def test_lora_disaggregated_uses_resolved_args(self): + args = _server_args(lora_ipc_weight_sync=False, colocate=False) + kwargs = _compute_server_args(args, "127.0.0.1", 15000, 15001) + assert kwargs["lora_target_modules"] == ["to_q", "to_k"] From 5911ece786dd76c43deca187ad60b46c681b0107 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 18:58:02 +0000 Subject: [PATCH 02/14] feat(train): one-step async training entry with prefetched rollouts --- miles/ray/rollout.py | 1 + miles/rollout/data_source.py | 25 +++- .../fast/rollout/test_data_source_snapshot.py | 78 ++++++++++++ train_diffusion_async.py | 118 ++++++++++++++++++ 4 files changed, 216 insertions(+), 6 deletions(-) create mode 100644 tests/fast/rollout/test_data_source_snapshot.py create mode 100644 train_diffusion_async.py diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 2d1eabf2e..0654c4c4c 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -172,6 +172,7 @@ def generate(self, rollout_id): with timer("rollout"): data, metrics = self._get_rollout_data(rollout_id=rollout_id) + self.data_source.snapshot(rollout_id) with timer("save_debug_dump"): self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False) _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) diff --git a/miles/rollout/data_source.py b/miles/rollout/data_source.py index 7fdf6f737..bc24ea297 100644 --- a/miles/rollout/data_source.py +++ b/miles/rollout/data_source.py @@ -36,6 +36,12 @@ def load(self, rollout_id=None): Load the state of the data source """ + def snapshot(self, rollout_id): + """ + Record the cursor state right after generating rollout_id, so a later + save(rollout_id) persists it even if a prefetched rollout advanced the cursor + """ + # TODO may further refactor data-loading part later class RolloutDataSource(DataSource): @@ -48,6 +54,7 @@ def __init__(self, args): self.sample_offset = 0 # TODO remove this self.metadata = {} + self._cursor_snapshots: dict[int, dict] = {} if args.rollout_global_dataset: from miles.utils.diffusion_data import Dataset as DiffusionDataset @@ -89,17 +96,23 @@ def get_samples(self, num_samples): samples.append(group) return samples - def save(self, rollout_id): - if not self.args.rollout_global_dataset: - return - - state_dict = { + def snapshot(self, rollout_id): + self._cursor_snapshots = {rid: state for rid, state in self._cursor_snapshots.items() if rid > rollout_id - 2} + self._cursor_snapshots[rollout_id] = { "sample_offset": self.sample_offset, "epoch_id": self.epoch_id, "sample_group_index": self.sample_group_index, "sample_index": self.sample_index, - "metadata": self.metadata, + "metadata": copy.deepcopy(self.metadata), } + + def save(self, rollout_id): + if not self.args.rollout_global_dataset: + return + + if rollout_id not in self._cursor_snapshots: + raise ValueError(f"no cursor snapshot for rollout {rollout_id}: save must follow generate") + state_dict = self._cursor_snapshots[rollout_id] path = os.path.join(self.args.save, f"rollout/global_dataset_state_dict_{rollout_id}.pt") os.makedirs(os.path.dirname(path), exist_ok=True) torch.save(state_dict, path) diff --git a/tests/fast/rollout/test_data_source_snapshot.py b/tests/fast/rollout/test_data_source_snapshot.py new file mode 100644 index 000000000..710d7f0f0 --- /dev/null +++ b/tests/fast/rollout/test_data_source_snapshot.py @@ -0,0 +1,78 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) + +import json +from argparse import Namespace + +import pytest +import torch + +from miles.rollout.data_source import RolloutDataSourceWithBuffer + + +def _args(tmp_path, **overrides): + prompt_path = tmp_path / "prompts.jsonl" + with open(prompt_path, "w") as f: + for i in range(16): + f.write(json.dumps({"input": f"prompt {i}"}) + "\n") + values = dict( + rollout_global_dataset=True, + prompt_data=str(prompt_path), + input_key="input", + metadata_key="metadata", + rollout_seed=42, + n_samples_per_prompt=2, + save=str(tmp_path / "ckpt"), + load=str(tmp_path / "ckpt"), + buffer_filter_path=None, + ) + values.update(overrides) + return Namespace(**values) + + +def _cursor(source): + return (source.sample_offset, source.epoch_id, source.sample_group_index, source.sample_index) + + +class TestCursorSnapshot: + def test_save_uses_the_snapshot_not_the_live_cursor(self, tmp_path): + args = _args(tmp_path) + source = RolloutDataSourceWithBuffer(args) + + source.get_samples(4) + source.snapshot(0) + cursor_after_rollout_0 = _cursor(source) + + # A prefetched rollout advances the live cursor past the saved rollout. + source.get_samples(4) + source.snapshot(1) + assert _cursor(source) != cursor_after_rollout_0 + + source.save(0) + + restored = RolloutDataSourceWithBuffer(_args(tmp_path)) + restored.load(0) + assert _cursor(restored) == cursor_after_rollout_0 + + def test_save_without_snapshot_rejects(self, tmp_path): + source = RolloutDataSourceWithBuffer(_args(tmp_path)) + with pytest.raises(ValueError, match="no cursor snapshot"): + source.save(0) + + def test_snapshot_prunes_older_entries(self, tmp_path): + source = RolloutDataSourceWithBuffer(_args(tmp_path)) + for rollout_id in range(5): + source.get_samples(2) + source.snapshot(rollout_id) + assert sorted(source._cursor_snapshots) == [3, 4] + + def test_saved_state_matches_legacy_format(self, tmp_path): + args = _args(tmp_path) + source = RolloutDataSourceWithBuffer(args) + source.get_samples(4) + source.snapshot(0) + source.save(0) + + state = torch.load(f"{args.save}/rollout/global_dataset_state_dict_0.pt") + assert set(state) == {"sample_offset", "epoch_id", "sample_group_index", "sample_index", "metadata"} diff --git a/train_diffusion_async.py b/train_diffusion_async.py new file mode 100644 index 000000000..ab09755a5 --- /dev/null +++ b/train_diffusion_async.py @@ -0,0 +1,118 @@ +import logging +import sys +import time + +import ray + +from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +from miles.utils import tracking_utils +from miles.utils.arguments import parse_args +from miles.utils.logging_utils import configure_logger +from miles.utils.metric_utils import compute_rollout_step +from miles.utils.misc import should_run_periodic_action +from miles.utils.tracking_utils import init_tracking + + +def train(args): + configure_logger() + logger = logging.getLogger(__name__) + assert not args.colocate, "async training overlaps train and rollout; drop --colocate" + assert not args.offload_train and not args.offload_rollout, "async training keeps both pools resident" + + logger.info("train_async: creating placement groups") + pgs = create_placement_groups(args) + init_tracking(args) + + logger.info("train_async: creating rollout manager") + rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + + logger.info("train_async: creating training model") + actor_model = create_training_models(args, pgs, rollout_manager) + + # always update weight first so that sglang has the loaded weights from training. + actor_model.update_weights() + + # special case for eval-only + if args.num_rollout == 0 and args.eval_interval is not None: + ray.get(rollout_manager.eval.remote(rollout_id=0)) + + def save(rollout_id): + actor_model.save_model( + rollout_id, + force_sync=rollout_id == args.num_rollout - 1, + ) + if args.rollout_global_dataset: + ray.get(rollout_manager.save.remote(rollout_id)) + + def log_step_metrics(rollout_id, durations): + log_dict = {f"perf/{name}": value for name, value in durations.items()} + log_dict["rollout/step"] = compute_rollout_step(args, rollout_id) + tracking_utils.log(args, log_dict, step_key="rollout/step") + + if args.eval_interval is not None and not args.skip_eval_before_train and args.num_rollout > 0: + ray.get(rollout_manager.eval.remote(args.start_rollout_id)) + + # one-step overlap: generate(rollout_id + 1) runs while train(rollout_id) runs, so the + # trained batch is exactly one weight version stale. Weights are only pushed at the + # barrier below, after the in-flight generation drains, so every rollout sees a single + # weight version. + generate_future = None + if args.start_rollout_id < args.num_rollout: + generate_future = rollout_manager.generate.remote(args.start_rollout_id) + + rollout_data_ref = None + for rollout_id in range(args.start_rollout_id, args.num_rollout): + step_start = time.time() + + if generate_future is not None: + rollout_data_ref = ray.get(generate_future) + if rollout_id + 1 < args.num_rollout: + generate_future = rollout_manager.generate.remote(rollout_id + 1) + else: + generate_future = None + + logger.info(f"train_async: rollout {rollout_id} actor train start") + train_start = time.time() + ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) + train_wall = time.time() - train_start + + if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): + save(rollout_id) + + # drain the in-flight generation before updating weights so no rollout runs + # under a mid-generation weight swap + drain_start = time.time() + if generate_future is not None: + rollout_data_ref = ray.get(generate_future) + generate_future = None + drain_wait = time.time() - drain_start + + update_start = time.time() + actor_model.update_weights() + update_wall = time.time() - update_start + + if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): + ray.get(rollout_manager.eval.remote(rollout_id)) + + log_step_metrics( + rollout_id, + { + "train_wall": train_wall, + "drain_wait": drain_wait, + "update_weights_wall": update_wall, + "step_time": time.time() - step_start, + }, + ) + logger.info( + f"train_async: rollout {rollout_id} done " + f"train_wall={train_wall:.1f}s drain_wait={drain_wait:.1f}s update={update_wall:.1f}s" + ) + + ray.get(rollout_manager.dispose.remote()) + + +if __name__ == "__main__": + # Ensure stdout is line-buffered so nohup logs show progress immediately. + sys.stdout.reconfigure(line_buffering=True) + args = parse_args() + train(args) From 55a0fa59805e832cdda3e0bbbbed693dc414d12b Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 19:00:41 +0000 Subject: [PATCH 03/14] feat(fsdp): lagged EMA reference option for async training --- miles/backends/fsdp_utils/actor.py | 3 ++- miles/backends/fsdp_utils/ema.py | 18 ++++++++++++------ miles/utils/arguments.py | 12 ++++++++++++ .../backends/fsdp_utils/test_loss_hub_nft.py | 17 +++++++++++++++++ 4 files changed, 43 insertions(+), 7 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 5273fd204..4998fd6be 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -221,6 +221,7 @@ 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, + keep_lagged=self.args.ema_ref_lagged, ) # sglang-d now supports /update_weights_from_tensor (PR #20464). @@ -558,7 +559,7 @@ 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() + ref_ctx = self.ema_shadow.swap_in(lagged=self.args.ema_ref_lagged) else: ref_ctx = prepared.model.disable_adapter() with torch.no_grad(), ref_ctx: diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index 9120ef3b8..1bf9b8c25 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/miles/backends/fsdp_utils/ema.py @@ -25,6 +25,7 @@ def __init__( uprate: float = 0.001, uphold: float = 0.5, flat_steps: int = 0, + keep_lagged: bool = False, ) -> None: self.decay = float(decay) self.uprate = float(uprate) @@ -37,6 +38,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.lagged = [sh.clone() for sh in self.shadow] if keep_lagged else None def decay_at(self, t: int) -> float: if t <= self.flat_steps: @@ -50,24 +52,28 @@ def update(self) -> float: raise RuntimeError("EmaShadow.update called while swapped in") self.step += 1 delta = self.decay_at(self.step) + if self.lagged is not None: + for lg, sh in zip(self.lagged, self.shadow, strict=True): + lg.copy_(sh) 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 @contextmanager - def swap_in(self): - """Temporarily expose EMA weights as the live parameters.""" - self._swap() + def swap_in(self, lagged: bool = False): + """Temporarily expose EMA weights (or the pre-update EMA snapshot) as the live parameters.""" + buffers = self.lagged if lagged 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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 147535711..022aa9872 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1113,6 +1113,16 @@ def add_ema_arguments(parser): default=0, help="Steps the decay holds at --ema-decay-init before the ramp begins.", ) + parser.add_argument( + "--ema-ref-lagged", + action="store_true", + default=False, + help=( + "Use the pre-update EMA snapshot as the reference policy. Under one-step " + "async training this matches the reference to the weights that generated " + "the batch being trained." + ), + ) return parser def add_dashboard_arguments(parser): @@ -1612,6 +1622,8 @@ def miles_validate_args(args): raise ValueError("--use-ema has no consumer; set --ref-mode ema or --ema-rollout-policy ema") if args.ema_rollout_policy == "ema" and not args.use_ema: raise ValueError("--ema-rollout-policy ema requires --use-ema") + if args.ema_ref_lagged and args.ref_mode != "ema": + raise ValueError("--ema-ref-lagged requires --ref-mode ema") if args.loss_type == "sft_loss": if not args.train_only: diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index 8fe5ea642..d4f9d2a3a 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -286,3 +286,20 @@ def test_swap_in_restores_exactly(self): with ema.swap_in(): assert torch.equal(m.weight.detach(), live) assert torch.equal(m.weight.detach(), live + 2.0) + + def test_lagged_snapshot_tracks_pre_update_ema(self): + m = self._model() + ema = EmaShadow(m.parameters(), decay=0.5, uprate=0.001, uphold=0.5, flat_steps=10, keep_lagged=True) + init = m.weight.detach().clone() + with torch.no_grad(): + m.weight.add_(1.0) + ema.update() + assert torch.equal(ema.lagged[0], init) + assert torch.allclose(ema.shadow[0], init + 0.5) + with ema.swap_in(lagged=True): + assert torch.equal(m.weight.detach(), init) + assert torch.equal(m.weight.detach(), init + 1.0) + + def test_lagged_disabled_by_default(self): + ema = EmaShadow(self._model().parameters(), decay=0.1) + assert ema.lagged is None From f481c2a46bc072a098b364659a11ba68dc68c05e Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 19:02:27 +0000 Subject: [PATCH 04/14] feat(krea2): async training mode for the NFT recipe --- scripts/run_diffusion_nft_krea2.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/scripts/run_diffusion_nft_krea2.py b/scripts/run_diffusion_nft_krea2.py index 4a8844cf3..12c0e6e74 100644 --- a/scripts/run_diffusion_nft_krea2.py +++ b/scripts/run_diffusion_nft_krea2.py @@ -12,9 +12,14 @@ Smoke mode shrinks the batch for checking the pipeline end to end without a real run. +--train-async switches to one-step async training (train_diffusion_async.py): train and +rollout run on separate GPU pools (2+2) with CPU-staged LoRA weight sync, so the trained +batch is one weight version stale. + Usage: python3 scripts/run_diffusion_nft_krea2.py python3 scripts/run_diffusion_nft_krea2.py --reward pickscore + python3 scripts/run_diffusion_nft_krea2.py --train-async MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_krea2.py """ @@ -36,6 +41,7 @@ class ScriptArgs(U.ExecuteTrainConfig): data_dir: str = "/root/datasets" smoke: bool = False reward: str = "ocr" # ocr | pickscore + train_async: bool = False extra_args: str = "" @@ -48,7 +54,7 @@ def _subset(args: ScriptArgs) -> str: def _num_gpus(args: ScriptArgs) -> int: - return 2 if _use_ocr(args) else 3 + return (4 if args.train_async else 2) + (0 if _use_ocr(args) else 1) def prepare(args: ScriptArgs) -> str: @@ -57,7 +63,8 @@ def prepare(args: ScriptArgs) -> str: def execute(args: ScriptArgs, data_dir: str) -> None: - run_name = f"diffusion_nft_krea2_{args.reward}_{U.create_run_id()}" + mode = "async" if args.train_async else "colocate" + run_name = f"diffusion_nft_krea2_{args.reward}_{mode}_{U.create_run_id()}" num_rollout = args.num_rollout or (1 if args.smoke else 100) ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 " @@ -106,7 +113,9 @@ def execute(args: ScriptArgs, data_dir: str) -> None: optimizer_args = "--lr 3e-4 --adam-beta2 0.999 --weight-decay 1e-4 --clip-grad 1.0 " - lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " + lora_args = "--use-lora --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " + ( + "" if args.train_async else "--lora-ipc-weight-sync " + ) reward_args = ( "--rm-type ocr " @@ -142,8 +151,8 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--rollout-num-gpus 2 " "--rollout-num-gpus-per-engine 1 " f"--num-gpus-per-node {_num_gpus(args)} " - "--colocate " - "--deterministic-mode " + + ("" if args.train_async else "--colocate ") + + "--deterministic-mode " ) U.execute_train( @@ -153,6 +162,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: f"{train_backend_args} {perf_args} {misc_args} {args.extra_args}" ), num_gpus_per_node=_num_gpus(args), + train_script="train_diffusion_async.py" if args.train_async else "train_diffusion.py", config=args, extra_env_vars={ "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", From d12423940e94a5b39a97119ce4eccedb62cdc98c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 19:09:40 +0000 Subject: [PATCH 05/14] fix(fsdp): self-contained pickle for CPU-staged weight sync payloads --- miles/backends/fsdp_utils/diffusion_update_weight_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index e4dd2eb9a..2385eb0af 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -1,6 +1,8 @@ import abc +import base64 import logging import os +import pickle import re from argparse import Namespace from collections.abc import Mapping, Sequence @@ -584,7 +586,9 @@ def update_bucket_weights( "metadata": bucket.get_metadata(), } } - serialized = MultiprocessingSerializer.serialize(payload, output_str=True) + # Plain pickle embeds the CPU tensor bytes; ForkingPickler would ship + # shared-memory handles that unrelated engine processes cannot open. + serialized = base64.b64encode(pickle.dumps(payload)).decode() kwargs = { # A single unlabeled payload: every engine worker deserializes the same # full adapter set (no shared-GPU visibility required). From 056c479c83cd855a9618b159f6b115304a89dfdc Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 4 Sep 2026 19:10:54 +0000 Subject: [PATCH 06/14] refactor(train): log only drain_wait from the async driver --- train_diffusion_async.py | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/train_diffusion_async.py b/train_diffusion_async.py index ab09755a5..d333a6a87 100644 --- a/train_diffusion_async.py +++ b/train_diffusion_async.py @@ -44,9 +44,14 @@ def save(rollout_id): if args.rollout_global_dataset: ray.get(rollout_manager.save.remote(rollout_id)) - def log_step_metrics(rollout_id, durations): - log_dict = {f"perf/{name}": value for name, value in durations.items()} - log_dict["rollout/step"] = compute_rollout_step(args, rollout_id) + def log_drain_wait(rollout_id, drain_wait): + # The actor already logs perf/train_time, perf/step_time and perf/wait_time_ratio; + # drain_wait (time spent waiting for the prefetched rollout at the barrier) is the + # only phase invisible to it. + log_dict = { + "perf/drain_wait_time": drain_wait, + "rollout/step": compute_rollout_step(args, rollout_id), + } tracking_utils.log(args, log_dict, step_key="rollout/step") if args.eval_interval is not None and not args.skip_eval_before_train and args.num_rollout > 0: @@ -62,8 +67,6 @@ def log_step_metrics(rollout_id, durations): rollout_data_ref = None for rollout_id in range(args.start_rollout_id, args.num_rollout): - step_start = time.time() - if generate_future is not None: rollout_data_ref = ray.get(generate_future) if rollout_id + 1 < args.num_rollout: @@ -94,15 +97,7 @@ def log_step_metrics(rollout_id, durations): if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): ray.get(rollout_manager.eval.remote(rollout_id)) - log_step_metrics( - rollout_id, - { - "train_wall": train_wall, - "drain_wait": drain_wait, - "update_weights_wall": update_wall, - "step_time": time.time() - step_start, - }, - ) + log_drain_wait(rollout_id, drain_wait) logger.info( f"train_async: rollout {rollout_id} done " f"train_wall={train_wall:.1f}s drain_wait={drain_wait:.1f}s update={update_wall:.1f}s" From 3d6bc14d1a3fd08a40b1401069b1237b81d4bc6a Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Sat, 5 Sep 2026 00:03:17 +0000 Subject: [PATCH 07/14] fix(rollout): flush wandb in dispose so a final-step eval survives teardown --- miles/ray/rollout.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 0654c4c4c..0a48428a2 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -10,6 +10,7 @@ import numpy as np import ray import torch +import wandb from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS @@ -144,6 +145,9 @@ def dispose(self): self._metric_checker.dispose() if self._health_monitor is not None: self._health_monitor.stop() + if self.args.use_wandb: + # Flush buffered logs: a final-step eval logs moments before teardown. + wandb.finish() # TODO maybe rename "rollout_engines" and "all_rollout_engines" later @property From 26534cf59acb0084e6abbdef1f36e5dbf20233a8 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:09:41 +0000 Subject: [PATCH 08/14] fix: align async EMA references and simplify checkpoint flow Advance EMA after training, infer the lagged reference from async mode, and persist EMA weights and step with distributed checkpoints. Save the prompt cursor before prefetch and measure drain wait before checkpoint I/O. Remove cursor snapshot bookkeeping and the added wandb finish call. Run three async smoke rollouts and cover overlap, update barriers, cursor recovery, EMA resume, and checkpoint resharding. Leave the weight transport unchanged. --- docs/user-guide/launch-script.md | 20 ++ miles/backends/fsdp_utils/actor.py | 16 +- miles/backends/fsdp_utils/checkpoint.py | 12 + miles/backends/fsdp_utils/ema.py | 28 +++ miles/ray/rollout.py | 5 - miles/rollout/data_source.py | 25 +-- miles/utils/arguments.py | 16 +- scripts/run_diffusion_nft_krea2.py | 3 +- .../fsdp_utils/_ema_checkpoint_worker.py | 25 +++ .../fsdp_utils/test_ema_checkpoint.py | 97 ++++++++ tests/fast/rollout/test_async_training.py | 207 ++++++++++++++++++ .../fast/rollout/test_data_source_snapshot.py | 78 ------- train_diffusion_async.py | 146 +++++------- 13 files changed, 468 insertions(+), 210 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/_ema_checkpoint_worker.py create mode 100644 tests/fast/backends/fsdp_utils/test_ema_checkpoint.py create mode 100644 tests/fast/rollout/test_async_training.py delete mode 100644 tests/fast/rollout/test_data_source_snapshot.py diff --git a/docs/user-guide/launch-script.md b/docs/user-guide/launch-script.md index 7724a6d14..b2ac35e5a 100644 --- a/docs/user-guide/launch-script.md +++ b/docs/user-guide/launch-script.md @@ -88,6 +88,26 @@ live in `rollout_args`, a `--micro-batch-size-*` flag in `perf_args`). | `perf_args` | Gradient checkpointing, micro-batch tiling, parser workers | | `misc_args` | GPU layout, `--colocate`, `--deterministic-mode` | +## One-step async DiffusionNFT + +`python3 scripts/run_diffusion_nft_krea2.py --train-async` runs training and rollout on +separate resident GPU pools. It generates the next batch during training and waits for +that generation to finish before publishing new weights. Async smoke mode runs three +rollouts so it exercises both prefetch and weight updates. + +When both rollout and reference use EMA, the async entrypoint automatically retains the +sampling EMA for the reference forward. EMA advances once per completed training rollout; +initial synchronization and repeated weight publication do not advance its decay schedule. + +Checkpoints include the current EMA weights and step. Resume restores the prompt cursor and +EMA, then starts a fresh prefetch pipeline. It regenerates any prefetched batch, so resumed +training is not guaranteed to reproduce an uninterrupted trajectory. The first resumed batch +uses the restored EMA for both sampling and reference. Older checkpoints without EMA state +initialize it from the loaded model and emit a warning. + +`perf/drain_wait_time` measures the generation wait before checkpoint I/O. Compare throughput +with GPU counts included: this recipe uses four GPUs for async OCR versus two for colocated OCR. + ## Ways to override a recipe From lightest to heaviest: diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 4998fd6be..aa83a719b 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -221,7 +221,11 @@ 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, - keep_lagged=self.args.ema_ref_lagged, + keep_lagged=( + 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). @@ -308,10 +312,6 @@ 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) rollout_weight_context = ( ema_shadow.swap_in() if ema_shadow is not None and self.args.ema_rollout_policy == "ema" else nullcontext() ) @@ -347,6 +347,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, @@ -559,7 +563,7 @@ 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(lagged=self.args.ema_ref_lagged) + ref_ctx = self.ema_shadow.swap_in(lagged=self.ema_shadow.lagged is not None) else: ref_ctx = prepared.model.disable_adapter() with torch.no_grad(), ref_ctx: diff --git a/miles/backends/fsdp_utils/checkpoint.py b/miles/backends/fsdp_utils/checkpoint.py index 4137757b3..a4398e7e2 100644 --- a/miles/backends/fsdp_utils/checkpoint.py +++ b/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 diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index 1bf9b8c25..9c2927249 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/miles/backends/fsdp_utils/ema.py @@ -59,6 +59,34 @@ def update(self) -> float: 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.lagged is not None: + for lg, sh in zip(self.lagged, self.shadow, strict=True): + lg.copy_(sh) + @contextmanager def swap_in(self, lagged: bool = False): """Temporarily expose EMA weights (or the pre-update EMA snapshot) as the live parameters.""" diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 0a48428a2..2d1eabf2e 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -10,7 +10,6 @@ import numpy as np import ray import torch -import wandb from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS @@ -145,9 +144,6 @@ def dispose(self): self._metric_checker.dispose() if self._health_monitor is not None: self._health_monitor.stop() - if self.args.use_wandb: - # Flush buffered logs: a final-step eval logs moments before teardown. - wandb.finish() # TODO maybe rename "rollout_engines" and "all_rollout_engines" later @property @@ -176,7 +172,6 @@ def generate(self, rollout_id): with timer("rollout"): data, metrics = self._get_rollout_data(rollout_id=rollout_id) - self.data_source.snapshot(rollout_id) with timer("save_debug_dump"): self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=False) _log_rollout_data(rollout_id, self.args, data, metrics, time.time() - start_time) diff --git a/miles/rollout/data_source.py b/miles/rollout/data_source.py index bc24ea297..7fdf6f737 100644 --- a/miles/rollout/data_source.py +++ b/miles/rollout/data_source.py @@ -36,12 +36,6 @@ def load(self, rollout_id=None): Load the state of the data source """ - def snapshot(self, rollout_id): - """ - Record the cursor state right after generating rollout_id, so a later - save(rollout_id) persists it even if a prefetched rollout advanced the cursor - """ - # TODO may further refactor data-loading part later class RolloutDataSource(DataSource): @@ -54,7 +48,6 @@ def __init__(self, args): self.sample_offset = 0 # TODO remove this self.metadata = {} - self._cursor_snapshots: dict[int, dict] = {} if args.rollout_global_dataset: from miles.utils.diffusion_data import Dataset as DiffusionDataset @@ -96,23 +89,17 @@ def get_samples(self, num_samples): samples.append(group) return samples - def snapshot(self, rollout_id): - self._cursor_snapshots = {rid: state for rid, state in self._cursor_snapshots.items() if rid > rollout_id - 2} - self._cursor_snapshots[rollout_id] = { + def save(self, rollout_id): + if not self.args.rollout_global_dataset: + return + + state_dict = { "sample_offset": self.sample_offset, "epoch_id": self.epoch_id, "sample_group_index": self.sample_group_index, "sample_index": self.sample_index, - "metadata": copy.deepcopy(self.metadata), + "metadata": self.metadata, } - - def save(self, rollout_id): - if not self.args.rollout_global_dataset: - return - - if rollout_id not in self._cursor_snapshots: - raise ValueError(f"no cursor snapshot for rollout {rollout_id}: save must follow generate") - state_dict = self._cursor_snapshots[rollout_id] path = os.path.join(self.args.save, f"rollout/global_dataset_state_dict_{rollout_id}.pt") os.makedirs(os.path.dirname(path), exist_ok=True) torch.save(state_dict, path) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 022aa9872..bc70de88d 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1113,16 +1113,6 @@ def add_ema_arguments(parser): default=0, help="Steps the decay holds at --ema-decay-init before the ramp begins.", ) - parser.add_argument( - "--ema-ref-lagged", - action="store_true", - default=False, - help=( - "Use the pre-update EMA snapshot as the reference policy. Under one-step " - "async training this matches the reference to the weights that generated " - "the batch being trained." - ), - ) return parser def add_dashboard_arguments(parser): @@ -1622,8 +1612,6 @@ def miles_validate_args(args): raise ValueError("--use-ema has no consumer; set --ref-mode ema or --ema-rollout-policy ema") if args.ema_rollout_policy == "ema" and not args.use_ema: raise ValueError("--ema-rollout-policy ema requires --use-ema") - if args.ema_ref_lagged and args.ref_mode != "ema": - raise ValueError("--ema-ref-lagged requires --ref-mode ema") if args.loss_type == "sft_loss": if not args.train_only: @@ -1751,7 +1739,9 @@ def miles_validate_args(args): if args.lora_ipc_weight_sync: raise ValueError("--lora-ipc-weight-sync requires --colocate: CUDA IPC needs shared train/rollout GPUs") if not args.use_lora: - raise ValueError("disaggregated train/rollout weight sync supports LoRA only; pass --use-lora or --colocate") + raise ValueError( + "disaggregated train/rollout weight sync supports LoRA only; pass --use-lora or --colocate" + ) if args.colocate_reward: assert args.colocate, "--colocate-reward requires --colocate." diff --git a/scripts/run_diffusion_nft_krea2.py b/scripts/run_diffusion_nft_krea2.py index 12c0e6e74..37df2786b 100644 --- a/scripts/run_diffusion_nft_krea2.py +++ b/scripts/run_diffusion_nft_krea2.py @@ -65,7 +65,8 @@ def prepare(args: ScriptArgs) -> str: def execute(args: ScriptArgs, data_dir: str) -> None: mode = "async" if args.train_async else "colocate" run_name = f"diffusion_nft_krea2_{args.reward}_{mode}_{U.create_run_id()}" - num_rollout = args.num_rollout or (1 if args.smoke else 100) + smoke_rollouts = 3 if args.train_async else 1 + num_rollout = args.num_rollout or (smoke_rollouts if args.smoke else 100) ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 " diff --git a/tests/fast/backends/fsdp_utils/_ema_checkpoint_worker.py b/tests/fast/backends/fsdp_utils/_ema_checkpoint_worker.py new file mode 100644 index 000000000..20b0050ad --- /dev/null +++ b/tests/fast/backends/fsdp_utils/_ema_checkpoint_worker.py @@ -0,0 +1,25 @@ +"""Write distinct, uneven EMA shards for the CPU resharding test.""" + +import sys +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Shard, distribute_tensor + +from miles.backends.fsdp_utils.ema import EmaShadow + +if __name__ == "__main__": + dist.init_process_group("gloo") + mesh = init_device_mesh("cpu", (dist.get_world_size(),)) + full = torch.arange(15).reshape(5, 3).float() + param = torch.nn.Parameter(distribute_tensor(full, mesh, [Shard(0)])) + ema = EmaShadow([param], decay=0.5, flat_steps=10) + for _ in range(2): + with torch.no_grad(): + param.add_(1) + ema.update() + dcp.save({"ema": ema}, checkpoint_id=str(Path(sys.argv[1]) / "ema")) + dist.destroy_process_group() diff --git a/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py b/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py new file mode 100644 index 000000000..75d36bac3 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py @@ -0,0 +1,97 @@ +"""EMA checkpoint integration, including restoration of distributed shards.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="stage-a-cpu", labels=[]) + +import shutil +import subprocess +import sys +from argparse import Namespace + +import pytest +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp + +from miles.backends.fsdp_utils import checkpoint +from miles.backends.fsdp_utils.ema import EmaShadow + + +def make_actor(tmp_path): + model = torch.nn.Linear(3, 5, bias=False) + optimizer = torch.optim.AdamW(model.parameters()) + return Namespace( + model=model, + optimizer=optimizer, + lr_scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0), + ema_shadow=EmaShadow(model.parameters(), decay=0.5, flat_steps=10, keep_lagged=True), + global_step=2, + micro_step=0, + train_pipeline_config=Namespace(optimizer_state_allowed_missing=[]), + args=Namespace( + save=str(tmp_path), + load=str(tmp_path), + ckpt_step=None, + use_lora=False, + no_save_optim=False, + no_load_optim=False, + no_load_rng=True, + start_rollout_id=0, + ), + ) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_checkpoint_restores_ema_and_restarts_reference(tmp_path, monkeypatch, legacy): + dist.init_process_group("gloo", init_method=f"file://{tmp_path}/rendezvous", rank=0, world_size=1) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "get_rng_state_all", lambda: []) + try: + original = make_actor(tmp_path) + for _ in range(2): + with torch.no_grad(): + original.model.weight.add_(1) + original.ema_shadow.update() + checkpoint.save(original, iteration=1) + if legacy: + shutil.rmtree(tmp_path / "iter_0000002/ema") + restored = make_actor(tmp_path) + payload = checkpoint.load(restored) + # Actor initializes its EMA from the already-restored live model. + restored.ema_shadow = EmaShadow(restored.model.parameters(), decay=0.5, flat_steps=10, keep_lagged=True) + checkpoint.finalize_load(restored, payload) + assert restored.args.start_rollout_id == 2 + assert restored.ema_shadow.step == 2 + expected = original.model.weight if legacy else original.ema_shadow.shadow[0] + torch.testing.assert_close(restored.ema_shadow.shadow[0], expected) + torch.testing.assert_close(restored.ema_shadow.lagged[0], expected) + torch.testing.assert_close(restored.model.weight, original.model.weight) + if not legacy: + assert original.ema_shadow.update() == restored.ema_shadow.update() + torch.testing.assert_close(restored.ema_shadow.shadow[0], original.ema_shadow.shadow[0]) + finally: + dist.destroy_process_group() + + +def test_ema_checkpoint_reshards_to_single_process(tmp_path): + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc_per_node=2", + "--module", + "tests.fast.backends.fsdp_utils._ema_checkpoint_worker", + str(tmp_path), + ], + check=True, + timeout=180, + ) + param = torch.nn.Parameter(torch.zeros(5, 3)) + restored = EmaShadow([param], keep_lagged=True) + dcp.load({"ema": restored}, checkpoint_id=str(tmp_path / "ema")) + assert restored.step == 2 + torch.testing.assert_close(restored.shadow[0], torch.arange(15).reshape(5, 3).float() + 1.25) + torch.testing.assert_close(restored.lagged[0], restored.shadow[0]) diff --git a/tests/fast/rollout/test_async_training.py b/tests/fast/rollout/test_async_training.py new file mode 100644 index 000000000..e46f7c784 --- /dev/null +++ b/tests/fast/rollout/test_async_training.py @@ -0,0 +1,207 @@ +"""Exercise the async driver with real serial Ray actors and tiny CPU weights.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=90, suite="stage-a-cpu", labels=[]) + +import json +import tempfile +import time +from argparse import Namespace +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import ray +import torch +from train_diffusion_async import train_loop + +from miles.backends.fsdp_utils.ema import EmaShadow +from miles.rollout.data_source import RolloutDataSourceWithBuffer +from miles.utils.ray_utils import Box + + +@ray.remote +class RolloutProbe: + def __init__(self, args, delay): + self.source = RolloutDataSourceWithBuffer(args) + self.source.load(args.start_rollout_id - 1) + self.delay = delay + self.weight = 0.0 + self.events = [] + + def generate(self, rollout_id): + start = time.monotonic() + sample = self.source.get_samples(1)[0][0] + time.sleep(self.delay) + self.events.append((rollout_id, start, time.monotonic())) + return [Box(ray.put(dict(rollout_id=rollout_id, weight=self.weight, sample_index=sample.index)))] + + def save(self, rollout_id): + self.source.save(rollout_id) + + def install(self, weight): + self.weight = weight + + def get_rollout_engines_and_lock(self): + return [], None, 0 + + def get_events(self): + return self.events + + +@ray.remote +class TrainerProbe: + def __init__(self, manager, delay, restored=None): + from miles.backends.fsdp_utils import actor + from miles.utils.timer import Timer + + self.actor = actor + self.delay = delay + self.rollout_manager = manager + self.args = Namespace( + offload_train=False, debug_rollout_only=False, train_only=False, ema_rollout_policy="ema" + ) + self.parallel_state = SimpleNamespace(get_mesh=lambda name: SimpleNamespace(get_local_rank=lambda: 0)) + self.param = torch.nn.Parameter(torch.zeros(1)) + self.ema_shadow = EmaShadow([self.param], decay=0.5, flat_steps=100, keep_lagged=True) + if restored is not None: + with torch.no_grad(): + self.param.copy_(restored["param"]) + self.ema_shadow.load_state_dict(restored["ema"]) + self.weight_updater = SimpleNamespace(update_weights=self._capture_weight) + self.records = [] + self.saves = {} + Timer().start("train_wait") + + def _capture_weight(self): + self.published = self.param.item() + + def update_weights(self): + with patch.object(self.actor, "clear_memory"): + self.actor.FSDPTrainRayActor.update_weights(self) + return self.published, self.ema_shadow.step + + def _train_core(self, rollout_id, rollout_data): + start = time.monotonic() + with self.ema_shadow.swap_in(lagged=True): + reference = self.param.item() + assert rollout_data["rollout_id"] == rollout_id + assert reference == rollout_data["weight"] + time.sleep(self.delay) + with torch.no_grad(): + self.param.add_(1) + self.records.append((rollout_data, start, time.monotonic())) + + 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) + + def save_model(self, rollout_id, force_sync=False): + from copy import deepcopy + + self.saves[rollout_id] = deepcopy({"param": self.param.detach(), "ema": self.ema_shadow.state_dict()}) + + def result(self): + return self.records, self.saves + + +class TrainGroupProbe: + def __init__(self, trainer, manager): + self.trainer = trainer + self.manager = manager + self.updates = [] + + def async_train(self, rollout_id, batch): + return [self.trainer.train.remote(rollout_id, batch)] + + def update_weights(self): + requested_at = time.monotonic() + weight, ema_step = ray.get(self.trainer.update_weights.remote()) + self.updates.append((requested_at, ema_step)) + ray.get(self.manager.install.remote(weight)) + + def save_model(self, rollout_id, force_sync=False): + ray.get(self.trainer.save_model.remote(rollout_id, force_sync=force_sync)) + + +@pytest.fixture(scope="module", autouse=True) +def ray_runtime(): + with tempfile.TemporaryDirectory(prefix="pr232-ray-") as directory: + ray.init(address="local", num_cpus=2, num_gpus=0, include_dashboard=False, _temp_dir=directory) + yield + ray.shutdown() + + +def run_loop(tmp_path, monkeypatch, *, train_delay, rollout_delay, start=0, count=3, restored=None): + prompts = tmp_path / "prompts.jsonl" + prompts.write_text("".join(json.dumps({"input": str(i)}) + "\n" for i in range(16))) + args = Namespace( + start_rollout_id=start, + num_rollout=count, + save_interval=2, + eval_interval=None, + rollout_global_dataset=True, + prompt_data=str(prompts), + input_key="input", + metadata_key="metadata", + rollout_seed=42, + n_samples_per_prompt=1, + save=str(tmp_path / "ckpt"), + load=str(tmp_path / "ckpt") if restored else None, + buffer_filter_path=None, + use_wandb=False, + ) + metrics = [] + monkeypatch.setattr("train_diffusion_async.tracking_utils.log", lambda args, data, **kw: metrics.append(data)) + manager = RolloutProbe.remote(args, rollout_delay) + trainer = TrainerProbe.remote(manager, train_delay, restored) + group = TrainGroupProbe(trainer, manager) + try: + group.update_weights() + train_loop(args, group, manager, None) + records, saves = ray.get(trainer.result.remote()) + events = ray.get(manager.get_events.remote()) + for rid in saves: + saved = torch.load(f"{args.save}/rollout/global_dataset_state_dict_{rid}.pt") + assert saved["sample_index"] == rid + 1 + return records, saves, events, group.updates, metrics + finally: + ray.kill(trainer) + ray.kill(manager) + + +@pytest.mark.parametrize("train_delay,rollout_delay", [(0.05, 0.25), (0.25, 0.05)]) +def test_overlap_reference_cursor_and_update_barrier(tmp_path, monkeypatch, train_delay, rollout_delay): + records, saves, events, updates, metrics = run_loop( + tmp_path, monkeypatch, train_delay=train_delay, rollout_delay=rollout_delay + ) + assert [record[0]["sample_index"] for record in records] == [0, 1, 2] + assert [record[0]["weight"] for record in records] == [0.0, 0.0, 0.5] + assert [step for _, step in updates] == [0, 1, 2, 3] + assert saves[1]["ema"]["step"] == 2 + for i in range(2): + assert max(records[i][1], events[i + 1][1]) < min(records[i][2], events[i + 1][2]) + assert updates[i + 1][0] >= events[i + 1][2] + if rollout_delay > train_delay: + # Rollout 1 saves a checkpoint; its drain wait must not disappear into save(). + assert metrics[1]["perf/drain_wait_time"] > 0.05 + + +def test_resume_rewarms_with_restored_ema(tmp_path, monkeypatch): + _, saves, _, _, _ = run_loop(tmp_path, monkeypatch, train_delay=0, rollout_delay=0) + records, _, _, updates, _ = run_loop( + tmp_path, monkeypatch, train_delay=0, rollout_delay=0, start=2, count=5, restored=saves[1] + ) + assert [r[0]["sample_index"] for r in records] == [2, 3, 4] + assert [r[0]["weight"] for r in records] == [1.25, 1.25, 2.125] + assert [step for _, step in updates] == [2, 3, 4, 5] + + +@pytest.mark.parametrize("count", [0, 1]) +def test_empty_and_single_rollout(tmp_path, monkeypatch, count): + records, _, events, updates, _ = run_loop(tmp_path, monkeypatch, train_delay=0, rollout_delay=0, count=count) + assert len(records) == len(events) == count + assert len(updates) == count + 1 diff --git a/tests/fast/rollout/test_data_source_snapshot.py b/tests/fast/rollout/test_data_source_snapshot.py deleted file mode 100644 index 710d7f0f0..000000000 --- a/tests/fast/rollout/test_data_source_snapshot.py +++ /dev/null @@ -1,78 +0,0 @@ -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) - -import json -from argparse import Namespace - -import pytest -import torch - -from miles.rollout.data_source import RolloutDataSourceWithBuffer - - -def _args(tmp_path, **overrides): - prompt_path = tmp_path / "prompts.jsonl" - with open(prompt_path, "w") as f: - for i in range(16): - f.write(json.dumps({"input": f"prompt {i}"}) + "\n") - values = dict( - rollout_global_dataset=True, - prompt_data=str(prompt_path), - input_key="input", - metadata_key="metadata", - rollout_seed=42, - n_samples_per_prompt=2, - save=str(tmp_path / "ckpt"), - load=str(tmp_path / "ckpt"), - buffer_filter_path=None, - ) - values.update(overrides) - return Namespace(**values) - - -def _cursor(source): - return (source.sample_offset, source.epoch_id, source.sample_group_index, source.sample_index) - - -class TestCursorSnapshot: - def test_save_uses_the_snapshot_not_the_live_cursor(self, tmp_path): - args = _args(tmp_path) - source = RolloutDataSourceWithBuffer(args) - - source.get_samples(4) - source.snapshot(0) - cursor_after_rollout_0 = _cursor(source) - - # A prefetched rollout advances the live cursor past the saved rollout. - source.get_samples(4) - source.snapshot(1) - assert _cursor(source) != cursor_after_rollout_0 - - source.save(0) - - restored = RolloutDataSourceWithBuffer(_args(tmp_path)) - restored.load(0) - assert _cursor(restored) == cursor_after_rollout_0 - - def test_save_without_snapshot_rejects(self, tmp_path): - source = RolloutDataSourceWithBuffer(_args(tmp_path)) - with pytest.raises(ValueError, match="no cursor snapshot"): - source.save(0) - - def test_snapshot_prunes_older_entries(self, tmp_path): - source = RolloutDataSourceWithBuffer(_args(tmp_path)) - for rollout_id in range(5): - source.get_samples(2) - source.snapshot(rollout_id) - assert sorted(source._cursor_snapshots) == [3, 4] - - def test_saved_state_matches_legacy_format(self, tmp_path): - args = _args(tmp_path) - source = RolloutDataSourceWithBuffer(args) - source.get_samples(4) - source.snapshot(0) - source.save(0) - - state = torch.load(f"{args.save}/rollout/global_dataset_state_dict_0.pt") - assert set(state) == {"sample_offset", "epoch_id", "sample_group_index", "sample_index", "metadata"} diff --git a/train_diffusion_async.py b/train_diffusion_async.py index d333a6a87..2318bf096 100644 --- a/train_diffusion_async.py +++ b/train_diffusion_async.py @@ -1,113 +1,83 @@ -import logging +"""One-rollout overlap. Resume discards prefetch and starts from the saved EMA.""" + import sys import time import ray -from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from miles.utils import tracking_utils -from miles.utils.arguments import parse_args -from miles.utils.logging_utils import configure_logger from miles.utils.metric_utils import compute_rollout_step from miles.utils.misc import should_run_periodic_action -from miles.utils.tracking_utils import init_tracking - -def train(args): - configure_logger() - logger = logging.getLogger(__name__) - assert not args.colocate, "async training overlaps train and rollout; drop --colocate" - assert not args.offload_train and not args.offload_rollout, "async training keeps both pools resident" - logger.info("train_async: creating placement groups") - pgs = create_placement_groups(args) - init_tracking(args) +def train_loop(args, actor_model, rollout_manager, num_rollout_per_epoch): + if args.start_rollout_id >= args.num_rollout: + return - logger.info("train_async: creating rollout manager") - rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + current_batch = ray.get(rollout_manager.generate.remote(args.start_rollout_id)) + for rollout_id in range(args.start_rollout_id, args.num_rollout): + save_checkpoint = should_run_periodic_action( + rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout + ) + # This serial actor saves the current cursor before the next generate advances it. + cursor_save = ( + rollout_manager.save.remote(rollout_id) if save_checkpoint and args.rollout_global_dataset else None + ) + next_future = rollout_manager.generate.remote(rollout_id + 1) if rollout_id + 1 < args.num_rollout else None - logger.info("train_async: creating training model") - actor_model = create_training_models(args, pgs, rollout_manager) + ray.get(actor_model.async_train(rollout_id, current_batch)) - # always update weight first so that sglang has the loaded weights from training. - actor_model.update_weights() + # Measure the exposed generation wait before checkpoint I/O can hide it. + drain_start = time.monotonic() + if next_future is not None: + current_batch = ray.get(next_future) + drain_wait = time.monotonic() - drain_start - # special case for eval-only - if args.num_rollout == 0 and args.eval_interval is not None: - ray.get(rollout_manager.eval.remote(rollout_id=0)) + if save_checkpoint: + if cursor_save is not None: + ray.get(cursor_save) + actor_model.save_model(rollout_id, force_sync=rollout_id == args.num_rollout - 1) - def save(rollout_id): - actor_model.save_model( - rollout_id, - force_sync=rollout_id == args.num_rollout - 1, - ) - if args.rollout_global_dataset: - ray.get(rollout_manager.save.remote(rollout_id)) - - def log_drain_wait(rollout_id, drain_wait): - # The actor already logs perf/train_time, perf/step_time and perf/wait_time_ratio; - # drain_wait (time spent waiting for the prefetched rollout at the barrier) is the - # only phase invisible to it. - log_dict = { - "perf/drain_wait_time": drain_wait, - "rollout/step": compute_rollout_step(args, rollout_id), - } - tracking_utils.log(args, log_dict, step_key="rollout/step") - - if args.eval_interval is not None and not args.skip_eval_before_train and args.num_rollout > 0: - ray.get(rollout_manager.eval.remote(args.start_rollout_id)) - - # one-step overlap: generate(rollout_id + 1) runs while train(rollout_id) runs, so the - # trained batch is exactly one weight version stale. Weights are only pushed at the - # barrier below, after the in-flight generation drains, so every rollout sees a single - # weight version. - generate_future = None - if args.start_rollout_id < args.num_rollout: - generate_future = rollout_manager.generate.remote(args.start_rollout_id) - - rollout_data_ref = None - for rollout_id in range(args.start_rollout_id, args.num_rollout): - if generate_future is not None: - rollout_data_ref = ray.get(generate_future) - if rollout_id + 1 < args.num_rollout: - generate_future = rollout_manager.generate.remote(rollout_id + 1) - else: - generate_future = None - - logger.info(f"train_async: rollout {rollout_id} actor train start") - train_start = time.time() - ray.get(actor_model.async_train(rollout_id, rollout_data_ref)) - train_wall = time.time() - train_start - - if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout): - save(rollout_id) - - # drain the in-flight generation before updating weights so no rollout runs - # under a mid-generation weight swap - drain_start = time.time() - if generate_future is not None: - rollout_data_ref = ray.get(generate_future) - generate_future = None - drain_wait = time.time() - drain_start - - update_start = time.time() + # No generation is in flight while the engines install the new weights. actor_model.update_weights() - update_wall = time.time() - update_start - if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): ray.get(rollout_manager.eval.remote(rollout_id)) - - log_drain_wait(rollout_id, drain_wait) - logger.info( - f"train_async: rollout {rollout_id} done " - f"train_wall={train_wall:.1f}s drain_wait={drain_wait:.1f}s update={update_wall:.1f}s" + tracking_utils.log( + args, + {"perf/drain_wait_time": drain_wait, "rollout/step": compute_rollout_step(args, rollout_id)}, + step_key="rollout/step", ) + +def train(args): + from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models + from miles.utils.logging_utils import configure_logger + from miles.utils.tracking_utils import init_tracking + + configure_logger() + if args.colocate or args.offload_train or args.offload_rollout: + raise ValueError("async training requires separate resident train/rollout GPU pools") + args.train_async = True + + pgs = create_placement_groups(args) + init_tracking(args) + rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"]) + actor_model = create_training_models(args, pgs, rollout_manager) + + # Publish initial/restored weights without advancing EMA or its decay schedule. + actor_model.update_weights() + if args.eval_interval is not None: + if args.num_rollout == 0: + ray.get(rollout_manager.eval.remote(rollout_id=0)) + elif not args.skip_eval_before_train: + ray.get(rollout_manager.eval.remote(args.start_rollout_id)) + + train_loop(args, actor_model, rollout_manager, num_rollout_per_epoch) ray.get(rollout_manager.dispose.remote()) if __name__ == "__main__": - # Ensure stdout is line-buffered so nohup logs show progress immediately. + from miles.utils.arguments import parse_args + sys.stdout.reconfigure(line_buffering=True) - args = parse_args() - train(args) + train(parse_args()) From 6c48945654f3105743bc1142bda110de29e6375d Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:15:24 +0000 Subject: [PATCH 09/14] docs: defer async training documentation --- docs/user-guide/launch-script.md | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/docs/user-guide/launch-script.md b/docs/user-guide/launch-script.md index b2ac35e5a..7724a6d14 100644 --- a/docs/user-guide/launch-script.md +++ b/docs/user-guide/launch-script.md @@ -88,26 +88,6 @@ live in `rollout_args`, a `--micro-batch-size-*` flag in `perf_args`). | `perf_args` | Gradient checkpointing, micro-batch tiling, parser workers | | `misc_args` | GPU layout, `--colocate`, `--deterministic-mode` | -## One-step async DiffusionNFT - -`python3 scripts/run_diffusion_nft_krea2.py --train-async` runs training and rollout on -separate resident GPU pools. It generates the next batch during training and waits for -that generation to finish before publishing new weights. Async smoke mode runs three -rollouts so it exercises both prefetch and weight updates. - -When both rollout and reference use EMA, the async entrypoint automatically retains the -sampling EMA for the reference forward. EMA advances once per completed training rollout; -initial synchronization and repeated weight publication do not advance its decay schedule. - -Checkpoints include the current EMA weights and step. Resume restores the prompt cursor and -EMA, then starts a fresh prefetch pipeline. It regenerates any prefetched batch, so resumed -training is not guaranteed to reproduce an uninterrupted trajectory. The first resumed batch -uses the restored EMA for both sampling and reference. Older checkpoints without EMA state -initialize it from the loaded model and emit a warning. - -`perf/drain_wait_time` measures the generation wait before checkpoint I/O. Compare throughput -with GPU counts included: this recipe uses four GPUs for async OCR versus two for colocated OCR. - ## Ways to override a recipe From lightest to heaviest: From bc437bb7595910bbaaa5ad9074416cccf554be46 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:08:48 +0000 Subject: [PATCH 10/14] Use NCCL for distributed diffusion weight updates --- miles/backends/fsdp_utils/actor.py | 10 +- .../diffusion_update_weight_utils.py | 150 ++++++++++++------ .../sglang_diffusion_engine.py | 43 +++++ miles/ray/actor_group.py | 5 - miles/utils/arguments.py | 4 - .../test_distributed_weight_update.py | 88 ++++++++++ 6 files changed, 238 insertions(+), 62 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_distributed_weight_update.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index aa83a719b..d3d875fe5 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -30,10 +30,11 @@ from . import checkpoint from .diffusion_update_weight_utils import ( + DiffusionUpdateWeightFromDistributed, DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA, - DiffusionUpdateWeightFromTensorLoRACPU, DiffusionUpdateWeightFromTensorLoRAIPC, + DiffusionUpdateWeightLoRADistributed, ) from .ema import EmaShadow from .input_dtype_policy import apply_input_dtype_policy @@ -233,8 +234,11 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty 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 self.args.use_lora and not self.args.colocate: - self.weight_updater = DiffusionUpdateWeightFromTensorLoRACPU(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: diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 2385eb0af..a6e7a5bc4 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -1,11 +1,11 @@ import abc -import base64 import logging import os -import pickle import re +import socket from argparse import Namespace from collections.abc import Mapping, Sequence +from datetime import timedelta import ray import torch @@ -18,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] @@ -35,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: @@ -446,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: @@ -502,7 +502,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 = [] @@ -516,7 +516,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 @@ -554,52 +554,102 @@ def update_weights(self) -> None: ) -class DiffusionUpdateWeightFromTensorLoRACPU(DiffusionUpdateWeightFromTensorLoRAIPC): - """LoRA sync for disaggregated rollout: rank 0 posts CPU-staged adapters to every engine.""" +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) - def connect_rollout_engines( - self, - rollout_engines: Sequence[ActorHandle], - rollout_engine_lock: ActorHandle | None, - ) -> None: + +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: str, - weight_version=None, - weight_update_mode: str | None = None, - ) -> None: + 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, + ) - named_tensors_by_dtype: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} - for name, tensor in named_tensors: - named_tensors_by_dtype.setdefault(tensor.dtype, []).append((name, tensor.cpu())) - for group in named_tensors_by_dtype.values(): - bucket = FlattenedTensorBucket(named_tensors=group) - payload = { - target_module: { - "flattened_tensor": bucket.get_flattened_tensor(), - "metadata": bucket.get_metadata(), - } - } - # Plain pickle embeds the CPU tensor bytes; ForkingPickler would ship - # shared-memory handles that unrelated engine processes cannot open. - serialized = base64.b64encode(pickle.dumps(payload)).decode() - kwargs = { - # A single unlabeled payload: every engine worker deserializes the same - # full adapter set (no shared-GPU visibility required). - "serialized_named_tensors": [serialized], - "payload_gpu_uuids": None, - "load_format": "flattened_bucket", - "target_modules": [target_module], - "weight_version": str(weight_version), - } - if weight_update_mode is not None: - kwargs["weight_update_mode"] = weight_update_mode - kwargs["lora_alpha"] = self.args.lora_alpha - kwargs["lora_rank"] = self.args.lora_rank - ray.get([engine.update_weights_from_tensor.remote(**kwargs) for engine in self.rollout_engines]) +class DiffusionUpdateWeightLoRADistributed(DiffusionUpdateWeightLoRA, DiffusionUpdateWeightFromDistributed): + pass diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2fee21091..366df047f 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -268,6 +268,49 @@ def update_weights_from_tensor( payload, ) + def init_weights_update_group( + self, master_address, master_port, rank_offset, world_size, group_name, backend="nccl" + ): + return self._make_request( + "init_weights_update_group", + { + "master_address": master_address, + "master_port": master_port, + "rank_offset": rank_offset, + "world_size": world_size, + "group_name": group_name, + "backend": backend, + }, + ) + + def destroy_weights_update_group(self, group_name): + return self._make_request("destroy_weights_update_group", {"group_name": group_name}) + + def update_weights_from_distributed( + self, + names, + dtypes, + shapes, + group_name, + target_modules, + weight_update_mode=None, + lora_alpha=None, + lora_rank=None, + ): + return self._make_request( + "update_weights_from_distributed", + { + "names": names, + "dtypes": dtypes, + "shapes": shapes, + "group_name": group_name, + "target_modules": target_modules, + "weight_update_mode": weight_update_mode, + "lora_alpha": lora_alpha, + "lora_rank": lora_rank, + }, + ) + def get_weights_checksum(self, module_names: list[str] | None = None) -> dict: """Query the live engine for SHA-256 checksums of the named pipeline modules. diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index 8ce4c3100..ebc2fead4 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -1,5 +1,3 @@ -import os - import ray from ray.util.placement_group import PlacementGroup from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy @@ -49,9 +47,6 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): pg, reordered_bundle_indices, _reordered_gpu_ids = pg env_vars = { - # because sglang will always set NCCL_CUMEM_ENABLE to 0 - # we need also set it to 0 to prevent nccl error. - "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", **self.args.train_env_vars, } diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index bc70de88d..93a82f697 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1738,10 +1738,6 @@ def miles_validate_args(args): if not args.colocate and not args.train_only and not args.debug_rollout_only: if args.lora_ipc_weight_sync: raise ValueError("--lora-ipc-weight-sync requires --colocate: CUDA IPC needs shared train/rollout GPUs") - if not args.use_lora: - raise ValueError( - "disaggregated train/rollout weight sync supports LoRA only; pass --use-lora or --colocate" - ) if args.colocate_reward: assert args.colocate, "--colocate-reward requires --colocate." diff --git a/tests/fast/backends/fsdp_utils/test_distributed_weight_update.py b/tests/fast/backends/fsdp_utils/test_distributed_weight_update.py new file mode 100644 index 000000000..95c638b55 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_distributed_weight_update.py @@ -0,0 +1,88 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) + +from datetime import timedelta +from types import SimpleNamespace + +import torch + +from miles.backends.fsdp_utils import diffusion_update_weight_utils as update + + +def test_connect_assigns_every_engine_rank_before_waiting(monkeypatch): + calls = [] + engines = [] + for index in range(2): + + def init(index=index, **kwargs): + calls.append((index, kwargs)) + return index + + engines.append(SimpleNamespace(init_weights_update_group=SimpleNamespace(remote=init))) + + monkeypatch.setattr(update.ray._private.services, "get_node_ip_address", lambda: "127.0.0.1") + + def join(**kwargs): + assert len(calls) == 2 + assert kwargs["rank"] == 0 + assert kwargs["world_size"] == 7 + return "group" + + monkeypatch.setattr(update, "init_custom_process_group", join) + monkeypatch.setattr(update.ray, "get", lambda refs: calls.append(("wait", refs))) + group = update.connect_rollout_engines_from_distributed(engines, [2, 4], "update", timedelta(seconds=30)) + assert group == "group" + assert [calls[i][1]["rank_offset"] for i in range(2)] == [1, 3] + assert calls[-1] == ("wait", [0, 1]) + + +def test_broadcast_matches_metadata_and_retains_contiguous_buffers(monkeypatch): + events = [] + payloads = [] + tensors = [("b", torch.arange(6).reshape(2, 3).t()), ("a", torch.ones(2, dtype=torch.bfloat16))] + + def receive(**kwargs): + payloads.append(kwargs) + events.append("rpc") + return len(payloads) + + engines = [SimpleNamespace(update_weights_from_distributed=SimpleNamespace(remote=receive)) for _ in range(2)] + + def broadcast(tensor, src, group, async_op): + assert len(payloads) == 2 + assert tensor.is_contiguous() + index = events.count("broadcast") + torch.testing.assert_close(tensor, tensors[index][1]) + events.append("broadcast") + return SimpleNamespace(wait=lambda: events.append("wait")) + + monkeypatch.setattr(update.dist, "broadcast", broadcast) + monkeypatch.setattr(update.ray, "get", lambda refs: events.append("ack")) + update.broadcast_bucket(engines, "group", "update", tensors, "transformer") + assert payloads[0]["names"] == ["b", "a"] + assert payloads[0]["dtypes"] == ["int64", "bfloat16"] + assert payloads[0]["shapes"] == [[3, 2], [2]] + assert events == ["rpc", "rpc", "broadcast", "broadcast", "wait", "wait", "ack"] + + +def test_reconnect_destroys_old_group_before_joining(monkeypatch): + args = SimpleNamespace(rollout_num_gpus_per_engine=2, distributed_timeout_minutes=1) + updater = update.DiffusionUpdateWeightFromDistributed(args, {}) + updater._model_update_group = "old" + events = [] + engine = SimpleNamespace( + destroy_weights_update_group=SimpleNamespace(remote=lambda name: events.append("remote destroy")) + ) + monkeypatch.setattr(update.dist, "get_rank", lambda: 0) + monkeypatch.setattr(update.dist, "destroy_process_group", lambda group: events.append(("destroy", group))) + monkeypatch.setattr(update.ray, "get", lambda refs: events.append("ack")) + + def connect(**kwargs): + events.append("connect") + return "new" + + monkeypatch.setattr(update, "connect_rollout_engines_from_distributed", connect) + updater.connect_rollout_engines([engine], None) + assert events == ["remote destroy", ("destroy", "old"), "ack", "connect"] + assert updater._model_update_group == "new" From d0cdb537637f660fc19bd87359f31a8220844d68 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:43:19 +0000 Subject: [PATCH 11/14] Defer LoRA documentation updates --- docs/advanced/lora.md | 55 ++++++++++++++++--------------------------- 1 file changed, 20 insertions(+), 35 deletions(-) diff --git a/docs/advanced/lora.md b/docs/advanced/lora.md index adfb0c58d..a90a8d939 100644 --- a/docs/advanced/lora.md +++ b/docs/advanced/lora.md @@ -1,6 +1,6 @@ --- title: LoRA Training and Weight Sync -description: PEFT LoRA on FSDP diffusion actors and IPC or NCCL weight sync to sglang-diffusion rollout engines. +description: PEFT LoRA on FSDP diffusion actors and CUDA-IPC weight sync to sglang-diffusion rollout engines. --- Miles-diffusion trains LoRA adapters on the FSDP actor and syncs them to sglang-diffusion rollout engines each iteration. The rollout engine has no PEFT @@ -21,67 +21,52 @@ For colocated LoRA training, prefer **IPC merge** — push only `lora_A` / ``` `--lora-rank` / `--lora-alpha` vary by recipe (e.g. SD3 uses 32/64; some others -use 64/128). With `--colocate` and without `--lora-ipc-weight-sync`, LoRA merges -on the train side and pushes full merged weights (§3). IPC merge requires -`--colocate`. With separate train and rollout GPU pools, omit both flags: -LoRA A/B tensors are sent through NCCL and merged on the rollout side. +use 64/128). Without `--lora-ipc-weight-sync`, LoRA still trains but merges on +the train side and pushes full merged weights (§3). IPC merge requires +`--colocate`. ## 2. Key flags | Flag | Purpose | |---|---| | `--use-lora` | Enable PEFT LoRA on the FSDP actor | -| `--lora-ipc-weight-sync` | Use IPC to push `lora_A`/`lora_B` with `--colocate`; rollout merges locally | +| `--lora-ipc-weight-sync` | Push only `lora_A`/`lora_B`; rollout merges locally | | `--lora-rank` | LoRA rank (recipe-specific; often 32 or 64) | | `--lora-alpha` | LoRA alpha (typically 2× rank) | | `--lora-target-modules` | Override family defaults (optional) | | `--lora-init-weights` | Init scheme, e.g. `gaussian` | -| `--update-weight-buffer-size` | Weight transfer bucket size in bytes (recipes use 2 GB) | +| `--update-weight-buffer-size` | IPC bucket size in bytes (recipes use 2 GB) | | `--update-weight-target-module` | Component to sync (SD3 default: `transformer`) | -| `--colocate` | Train and rollout share GPUs; omit for NCCL updates between separate GPU pools | +| `--colocate` | **Required** — train and rollout share GPU visibility | `--lora-ipc-weight-sync` requires both `--use-lora` and `--colocate`. Without -colocation, use the NCCL updater for separate train/rollout GPU pools. +colocation, CUDA IPC handles cannot cross the train/rollout process boundary. LoRA target modules default from the model family's `TrainPipelineConfig.lora_target_modules`. For SD3, see [SD3 model guide](../models/sd3/sd3.md). -## 3. Weight-sync strategies +## 3. Three weight-sync strategies Selection logic in `miles/backends/fsdp_utils/actor.py`: | Condition | Updater class | Behavior | |---|---|---| -| Separate GPU pools, no LoRA | `DiffusionUpdateWeightFromDistributed` | Full weights through NCCL | -| Separate GPU pools, `--use-lora` | `DiffusionUpdateWeightLoRADistributed` | LoRA A/B through NCCL; rollout merges locally | -| `--colocate`, no LoRA | `DiffusionUpdateWeightFromTensor` | Full base-weight IPC | -| `--colocate --use-lora`, no IPC flag | `DiffusionUpdateWeightFromTensorLoRA` | Merge `W + αBA/r` on train side, push merged weights | -| `--colocate --use-lora --lora-ipc-weight-sync` | `DiffusionUpdateWeightFromTensorLoRAIPC` | Push only `lora_A`/`lora_B`; rollout merges via `weight_update_mode=lora_merge` | +| No LoRA | `DiffusionUpdateWeightFromTensor` | Full base-weight IPC | +| `--use-lora`, no IPC | `DiffusionUpdateWeightFromTensorLoRA` | Merge `W + αBA/r` on train side, push merged weights | +| `--use-lora --lora-ipc-weight-sync` | `DiffusionUpdateWeightFromTensorLoRAIPC` | Push only `lora_A`/`lora_B`; rollout merges via `weight_update_mode=lora_merge` | Implementation: `miles/backends/fsdp_utils/diffusion_update_weight_utils.py`. -### NCCL sync (separate train/rollout GPUs) - -The training sender and rollout workers join a dedicated NCCL weight-update -group. Each bucket's names, shapes, and dtypes are sent through the engine API; -tensor contents are broadcast on the NCCL group. The engine uses its existing -weight loader or `lora_merge` path after receiving the tensors. This requires -sglang-d's distributed weight-update endpoints. - -LoRA parameter collection is shared by the IPC and NCCL updaters. The async -driver waits for training and the prefetched rollout to finish before publishing -weights, so generation never overlaps weight installation. - -### Full-weight IPC sync (colocated, no LoRA) +### Full-weight sync (no LoRA) FSDP shards are all-gathered into `FlattenedTensorBucket` objects, serialized via CUDA IPC, and sent to the rollout engine's `update_weights_from_tensor(load_format="flattened_bucket")`. -### Train-side merge (colocated LoRA, no IPC flag) +### Train-side merge (LoRA, no IPC) For each base layer with adapters, the updater computes: @@ -112,7 +97,7 @@ tensors. When adding the next group would exceed `--update-weight-buffer-size`, the current bucket is flushed first; the whole group (both `lora_A` and `lora_B`) then starts the next bucket. Pairs are never split across buckets. -Constant: `LORA_WEIGHT_UPDATE_MODE = "lora_merge"`. +Constant: `LORA_IPC_WEIGHT_UPDATE_MODE = "lora_merge"`. Rollout-side merge precision is controlled by environment variable `SGLANG_DIFFUSION_LORA_MERGE_FP32`: @@ -125,21 +110,21 @@ Set automatically in `RolloutManager` when spawning engines. On the first few syncs, rank 0 logs lines like: ```text -LoRA weight sync v1 [transformer]: pushed N lora tensors, M layer prefixes in K buckets (unmapped=0) +LoRA IPC weight sync v1 [transformer]: pushed N lora tensors, M layer prefixes in K buckets (unmapped=0) ``` -For IPC, after FSDP all-gather, serialized buckets are collected on the **gather-src +After FSDP all-gather, serialized buckets are collected on the **gather-src rank** only; that rank calls the rollout engine. If sync stalls or VRAM grows -across rollouts, check trainer logs for `LoRA weight sync` lines and Ray +across rollouts, check trainer logs for `LoRA IPC weight sync` lines and Ray worker stderr under `~/.ray/session_latest/logs/`. ## 4. Internals | File | Role | |---|---| -| `miles/backends/fsdp_utils/diffusion_update_weight_utils.py` | IPC/NCCL updaters, shared LoRA collection, and `PeftLoRAKeyMapper` | +| `miles/backends/fsdp_utils/diffusion_update_weight_utils.py` | Three updater classes + `PeftLoRAKeyMapper` | | `miles/backends/fsdp_utils/actor.py` | Updater selection, LoRA apply via PEFT | -| `miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py` | Engine API calls for IPC updates and NCCL group/update management | +| `miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py` | HTTP `update_weights_from_tensor` to rollout | | `miles/ray/rollout.py` | Engine env vars (`SGLANG_DIFFUSION_LORA_MERGE_FP32`) | ## 5. Limitations From 60660932336457fa9a42ba4736c5b50160fa7928 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:56:27 +0000 Subject: [PATCH 12/14] Split Krea async training into a separate recipe --- scripts/run_diffusion_nft_krea2.py | 23 +-- scripts/run_diffusion_nft_krea2_async.py | 175 +++++++++++++++++++++++ 2 files changed, 181 insertions(+), 17 deletions(-) create mode 100644 scripts/run_diffusion_nft_krea2_async.py diff --git a/scripts/run_diffusion_nft_krea2.py b/scripts/run_diffusion_nft_krea2.py index 76bace589..4a8844cf3 100644 --- a/scripts/run_diffusion_nft_krea2.py +++ b/scripts/run_diffusion_nft_krea2.py @@ -12,14 +12,9 @@ Smoke mode shrinks the batch for checking the pipeline end to end without a real run. ---train-async switches to one-step async training (train_diffusion_async.py): train and -rollout run on separate GPU pools (2+2) with NCCL LoRA weight sync, so the trained -batch is one weight version stale. - Usage: python3 scripts/run_diffusion_nft_krea2.py python3 scripts/run_diffusion_nft_krea2.py --reward pickscore - python3 scripts/run_diffusion_nft_krea2.py --train-async MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_krea2.py """ @@ -41,7 +36,6 @@ class ScriptArgs(U.ExecuteTrainConfig): data_dir: str = "/root/datasets" smoke: bool = False reward: str = "ocr" # ocr | pickscore - train_async: bool = False extra_args: str = "" @@ -54,7 +48,7 @@ def _subset(args: ScriptArgs) -> str: def _num_gpus(args: ScriptArgs) -> int: - return (4 if args.train_async else 2) + (0 if _use_ocr(args) else 1) + return 2 if _use_ocr(args) else 3 def prepare(args: ScriptArgs) -> str: @@ -63,10 +57,8 @@ def prepare(args: ScriptArgs) -> str: def execute(args: ScriptArgs, data_dir: str) -> None: - mode = "async" if args.train_async else "colocate" - run_name = f"diffusion_nft_krea2_{args.reward}_{mode}_{U.create_run_id()}" - smoke_rollouts = 3 if args.train_async else 1 - num_rollout = args.num_rollout or (smoke_rollouts if args.smoke else 100) + run_name = f"diffusion_nft_krea2_{args.reward}_{U.create_run_id()}" + num_rollout = args.num_rollout or (1 if args.smoke else 100) ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 " @@ -114,9 +106,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: optimizer_args = "--lr 3e-4 --adam-beta2 0.999 --weight-decay 1e-4 --clip-grad 1.0 " - lora_args = "--use-lora --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " + ( - "" if args.train_async else "--lora-ipc-weight-sync " - ) + lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " reward_args = ( "--rm-type ocr " @@ -152,8 +142,8 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--rollout-num-gpus 2 " "--rollout-num-gpus-per-engine 1 " f"--num-gpus-per-node {_num_gpus(args)} " - + ("" if args.train_async else "--colocate ") - + "--deterministic-mode " + "--colocate " + "--deterministic-mode " ) U.execute_train( @@ -163,7 +153,6 @@ def execute(args: ScriptArgs, data_dir: str) -> None: f"{train_backend_args} {perf_args} {misc_args} {args.extra_args}" ), num_gpus_per_node=_num_gpus(args), - train_script="train_diffusion_async.py" if args.train_async else "train_diffusion.py", config=args, extra_env_vars={ "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", diff --git a/scripts/run_diffusion_nft_krea2_async.py b/scripts/run_diffusion_nft_krea2_async.py new file mode 100644 index 000000000..a9050bdbd --- /dev/null +++ b/scripts/run_diffusion_nft_krea2_async.py @@ -0,0 +1,175 @@ +"""Async Krea-2-Raw DiffusionNFT training (OCR by default, PickScore via --reward). + +Same NFT shape as run_diffusion_nft_sd3_pickscore.py: EMA reference (--ref-mode ema), +rollout under pi_old (--ema-rollout-policy ema), deterministic ODE rollout +(noise_level=0, sde_type=ode) with no CFG. Krea-2 specifics: bf16, 1024px, and one +sample per rollout request (the engine's krea2 pipeline has no per-request output +expansion). Rollout debug tensors are collected (--diffusion-debug-mode). + +OCR is the default reward: text rendering improves visibly and its accuracy curve is +steep, so both the metric and the wandb images validate the run. It needs no reward +GPU. --reward pickscore switches to the aesthetic direction on one extra GPU. + +Smoke mode shrinks the batch for checking the pipeline end to end without a real run. + +Training and rollout run on separate GPU pools (2+2) with NCCL LoRA weight sync. +The one-step async pipeline trains each prefetched batch against its lagged EMA +reference. Smoke mode runs three rollouts to cover the first updated rollout batch. + +Usage: + python3 scripts/run_diffusion_nft_krea2_async.py + python3 scripts/run_diffusion_nft_krea2_async.py --reward pickscore + MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_krea2_async.py +""" + +import os +from dataclasses import dataclass + +import typer + +import miles.utils.external_utils.command_utils as U + +MODEL = "krea/Krea-2-Raw" +DATASET = "rockdu/miles-diffusion-datasets" +WANDB_PROJECT = "diffusionNFT" + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + num_rollout: int = 0 # 0 picks the smoke/full default + data_dir: str = "/root/datasets" + smoke: bool = False + reward: str = "ocr" # ocr | pickscore + extra_args: str = "" + + +def _use_ocr(args: ScriptArgs) -> bool: + return args.smoke or args.reward == "ocr" + + +def _subset(args: ScriptArgs) -> str: + return "flowgrpo_ocr" if _use_ocr(args) else "flowgrpo_pickscore" + + +def _num_gpus(args: ScriptArgs) -> int: + return 4 if _use_ocr(args) else 5 + + +def prepare(args: ScriptArgs) -> str: + local_dir = U.hf_download_dataset(DATASET, include=f"{_subset(args)}/**", data_dir=args.data_dir) + return f"{local_dir}/{_subset(args)}" + + +def execute(args: ScriptArgs, data_dir: str) -> None: + run_name = f"diffusion_nft_krea2_{args.reward}_async_{U.create_run_id()}" + num_rollout = args.num_rollout or (3 if args.smoke else 100) + + ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 " + + rollout_args = ( + "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout " + f"--prompt-data {data_dir}/train.jsonl " + "--input-key input " + f"--num-rollout {num_rollout} " + "--num-steps-per-rollout 1 " + "--diffusion-num-steps 10 " + "--diffusion-guidance-scale 1.0 " + "--diffusion-noise-level 0.0 " + "--diffusion-sde-type ode " + "--diffusion-height 1024 " + "--diffusion-width 1024 " + "--diffusion-debug-mode " + "--rollout-microgroup-size 1 " + ) + ( + "--rollout-batch-size 2 --n-samples-per-prompt 2 " + if args.smoke + else "--rollout-batch-size 8 --n-samples-per-prompt 8 " + ) + + eval_args = "--diffusion-eval-num-steps 52 --skip-eval-before-train " + ( + "" if args.smoke else f"--eval-prompt-data {args.reward}_test {data_dir}/test.jsonl --eval-interval 30 " + ) + + grpo_args = ( + "--loss-type nft " + "--diffusion-nft-beta 1.0 " + "--diffusion-nft-timestep-fraction 0.99 " + "--advantage-estimator grpo " + "--globalize-reward-std " + ) + + ema_args = ( + "--ref-mode ema " + "--use-ema " + "--ema-rollout-policy ema " + "--ema-decay-init 0.001 " + "--ema-decay-ramp 0.001 " + "--ema-decay-max 0.5 " + "--ema-decay-flat-steps 0 " + ) + + optimizer_args = "--lr 3e-4 --adam-beta2 0.999 --weight-decay 1e-4 --clip-grad 1.0 " + + lora_args = "--use-lora --lora-rank 32 --lora-alpha 64 --lora-init-weights gaussian " + + reward_args = ( + "--rm-type ocr " + if _use_ocr(args) + else ( + "--rm-type pickscore " + "--pickscore-num-workers 1 " + "--pickscore-num-gpus-per-worker 1.0 " + "--pickscore-batch-size 8 " + "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K " + "--pickscore-model-path yuvalkirstain/PickScore_v1 " + ) + ) + + wandb_args = U.get_default_wandb_args( + __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10 + ) + + sglang_args = ( + "--use-miles-router " + "--sglang-server-concurrency 8 " + "--sglang-dit-precision bf16 " + "--sglang-vae-slicing " + "--update-weight-buffer-size 2147483648 " + ) + + train_backend_args = "--train-backend fsdp --diffusion-forward-dtype bf16 " + + perf_args = "--gradient-checkpointing " + ("--micro-batch-size 1 " if args.smoke else "--micro-batch-size 2 ") + + misc_args = ( + "--actor-num-gpus-per-node 2 " + "--rollout-num-gpus 2 " + "--rollout-num-gpus-per-engine 1 " + f"--num-gpus-per-node {_num_gpus(args)} " + "--deterministic-mode " + ) + + U.execute_train( + train_args=( + f"{ckpt_args} {rollout_args} {eval_args} {grpo_args} {ema_args} " + f"{optimizer_args} {lora_args} {reward_args} {wandb_args} {sglang_args} " + f"{train_backend_args} {perf_args} {misc_args} {args.extra_args}" + ), + num_gpus_per_node=_num_gpus(args), + train_script="train_diffusion_async.py", + config=args, + extra_env_vars={ + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "HF_TOKEN": os.environ.get("HF_TOKEN", ""), + }, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs) -> None: + data_dir = prepare(args) + execute(args, data_dir) + + +if __name__ == "__main__": + typer.run(main) From 63c352f0a478e0578c82e0cb9247abdee4fa0f60 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:02:59 +0000 Subject: [PATCH 13/14] Clarify previous EMA snapshot naming --- miles/backends/fsdp_utils/actor.py | 7 +++-- miles/backends/fsdp_utils/ema.py | 30 +++++++++++-------- scripts/run_diffusion_nft_krea2_async.py | 4 +-- .../fsdp_utils/test_ema_checkpoint.py | 10 +++---- .../backends/fsdp_utils/test_loss_hub_nft.py | 12 ++++---- tests/fast/rollout/test_async_training.py | 4 +-- 6 files changed, 38 insertions(+), 29 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d3d875fe5..9e29faea6 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -222,7 +222,8 @@ 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, - keep_lagged=( + # 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" @@ -316,6 +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 + # 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() ) @@ -567,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(lagged=self.ema_shadow.lagged is not None) + # 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: diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index 9c2927249..47f8fa513 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/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,7 +31,7 @@ def __init__( uprate: float = 0.001, uphold: float = 0.5, flat_steps: int = 0, - keep_lagged: bool = False, + keep_previous_ema: bool = False, ) -> None: self.decay = float(decay) self.uprate = float(uprate) @@ -38,7 +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.lagged = [sh.clone() for sh in self.shadow] if keep_lagged else None + 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: @@ -52,9 +58,9 @@ def update(self) -> float: raise RuntimeError("EmaShadow.update called while swapped in") self.step += 1 delta = self.decay_at(self.step) - if self.lagged is not None: - for lg, sh in zip(self.lagged, self.shadow, strict=True): - lg.copy_(sh) + 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 @@ -83,14 +89,14 @@ def load_state_dict(self, state_dict: dict) -> None: 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.lagged is not None: - for lg, sh in zip(self.lagged, self.shadow, strict=True): - lg.copy_(sh) + 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, lagged: bool = False): - """Temporarily expose EMA weights (or the pre-update EMA snapshot) as the live parameters.""" - buffers = self.lagged if lagged else self.shadow + 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: diff --git a/scripts/run_diffusion_nft_krea2_async.py b/scripts/run_diffusion_nft_krea2_async.py index a9050bdbd..cb2b52455 100644 --- a/scripts/run_diffusion_nft_krea2_async.py +++ b/scripts/run_diffusion_nft_krea2_async.py @@ -13,8 +13,8 @@ Smoke mode shrinks the batch for checking the pipeline end to end without a real run. Training and rollout run on separate GPU pools (2+2) with NCCL LoRA weight sync. -The one-step async pipeline trains each prefetched batch against its lagged EMA -reference. Smoke mode runs three rollouts to cover the first updated rollout batch. +The one-step async pipeline uses the previous EMA as the reference for each +prefetched batch. Smoke mode runs three rollouts to cover the first updated rollout batch. Usage: python3 scripts/run_diffusion_nft_krea2_async.py diff --git a/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py b/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py index 75d36bac3..b2b68e03f 100644 --- a/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py +++ b/tests/fast/backends/fsdp_utils/test_ema_checkpoint.py @@ -25,7 +25,7 @@ def make_actor(tmp_path): model=model, optimizer=optimizer, lr_scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0), - ema_shadow=EmaShadow(model.parameters(), decay=0.5, flat_steps=10, keep_lagged=True), + ema_shadow=EmaShadow(model.parameters(), decay=0.5, flat_steps=10, keep_previous_ema=True), global_step=2, micro_step=0, train_pipeline_config=Namespace(optimizer_state_allowed_missing=[]), @@ -59,13 +59,13 @@ def test_checkpoint_restores_ema_and_restarts_reference(tmp_path, monkeypatch, l restored = make_actor(tmp_path) payload = checkpoint.load(restored) # Actor initializes its EMA from the already-restored live model. - restored.ema_shadow = EmaShadow(restored.model.parameters(), decay=0.5, flat_steps=10, keep_lagged=True) + restored.ema_shadow = EmaShadow(restored.model.parameters(), decay=0.5, flat_steps=10, keep_previous_ema=True) checkpoint.finalize_load(restored, payload) assert restored.args.start_rollout_id == 2 assert restored.ema_shadow.step == 2 expected = original.model.weight if legacy else original.ema_shadow.shadow[0] torch.testing.assert_close(restored.ema_shadow.shadow[0], expected) - torch.testing.assert_close(restored.ema_shadow.lagged[0], expected) + torch.testing.assert_close(restored.ema_shadow.previous_ema[0], expected) torch.testing.assert_close(restored.model.weight, original.model.weight) if not legacy: assert original.ema_shadow.update() == restored.ema_shadow.update() @@ -90,8 +90,8 @@ def test_ema_checkpoint_reshards_to_single_process(tmp_path): timeout=180, ) param = torch.nn.Parameter(torch.zeros(5, 3)) - restored = EmaShadow([param], keep_lagged=True) + restored = EmaShadow([param], keep_previous_ema=True) dcp.load({"ema": restored}, checkpoint_id=str(tmp_path / "ema")) assert restored.step == 2 torch.testing.assert_close(restored.shadow[0], torch.arange(15).reshape(5, 3).float() + 1.25) - torch.testing.assert_close(restored.lagged[0], restored.shadow[0]) + torch.testing.assert_close(restored.previous_ema[0], restored.shadow[0]) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index d4f9d2a3a..6ac239050 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -287,19 +287,19 @@ def test_swap_in_restores_exactly(self): assert torch.equal(m.weight.detach(), live) assert torch.equal(m.weight.detach(), live + 2.0) - def test_lagged_snapshot_tracks_pre_update_ema(self): + def test_previous_ema_tracks_pre_update_snapshot(self): m = self._model() - ema = EmaShadow(m.parameters(), decay=0.5, uprate=0.001, uphold=0.5, flat_steps=10, keep_lagged=True) + ema = EmaShadow(m.parameters(), decay=0.5, uprate=0.001, uphold=0.5, flat_steps=10, keep_previous_ema=True) init = m.weight.detach().clone() with torch.no_grad(): m.weight.add_(1.0) ema.update() - assert torch.equal(ema.lagged[0], init) + assert torch.equal(ema.previous_ema[0], init) assert torch.allclose(ema.shadow[0], init + 0.5) - with ema.swap_in(lagged=True): + with ema.swap_in(use_previous_ema=True): assert torch.equal(m.weight.detach(), init) assert torch.equal(m.weight.detach(), init + 1.0) - def test_lagged_disabled_by_default(self): + def test_previous_ema_disabled_by_default(self): ema = EmaShadow(self._model().parameters(), decay=0.1) - assert ema.lagged is None + assert ema.previous_ema is None diff --git a/tests/fast/rollout/test_async_training.py b/tests/fast/rollout/test_async_training.py index d67797b81..6967e361d 100644 --- a/tests/fast/rollout/test_async_training.py +++ b/tests/fast/rollout/test_async_training.py @@ -64,7 +64,7 @@ def __init__(self, manager, delay, restored=None): ) self.parallel_state = SimpleNamespace(get_mesh=lambda name: SimpleNamespace(get_local_rank=lambda: 0)) self.param = torch.nn.Parameter(torch.zeros(1)) - self.ema_shadow = EmaShadow([self.param], decay=0.5, flat_steps=100, keep_lagged=True) + self.ema_shadow = EmaShadow([self.param], decay=0.5, flat_steps=100, keep_previous_ema=True) if restored is not None: with torch.no_grad(): self.param.copy_(restored["param"]) @@ -84,7 +84,7 @@ def update_weights(self): def _train_core(self, rollout_id, rollout_data): start = time.monotonic() - with self.ema_shadow.swap_in(lagged=True): + with self.ema_shadow.swap_in(use_previous_ema=True): reference = self.param.item() assert rollout_data["rollout_id"] == rollout_id assert reference == rollout_data["weight"] From 0dc33d02aba4eb1fe47cbe7f493c674deabdf48e Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Sat, 12 Sep 2026 00:52:42 +0000 Subject: [PATCH 14/14] Tune async Krea OCR for eight H200 GPUs --- scripts/run_diffusion_nft_krea2_async.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/scripts/run_diffusion_nft_krea2_async.py b/scripts/run_diffusion_nft_krea2_async.py index cb2b52455..a963316c8 100644 --- a/scripts/run_diffusion_nft_krea2_async.py +++ b/scripts/run_diffusion_nft_krea2_async.py @@ -12,7 +12,7 @@ Smoke mode shrinks the batch for checking the pipeline end to end without a real run. -Training and rollout run on separate GPU pools (2+2) with NCCL LoRA weight sync. +Full OCR uses 6 training + 2 rollout H200 GPUs with NCCL LoRA sync; smoke uses 2+2. The one-step async pipeline uses the previous EMA as the reference for each prefetched batch. Smoke mode runs three rollouts to cover the first updated rollout batch. @@ -52,7 +52,7 @@ def _subset(args: ScriptArgs) -> str: def _num_gpus(args: ScriptArgs) -> int: - return 4 if _use_ocr(args) else 5 + return 4 if args.smoke else (8 if _use_ocr(args) else 5) def prepare(args: ScriptArgs) -> str: @@ -63,6 +63,7 @@ def prepare(args: ScriptArgs) -> str: def execute(args: ScriptArgs, data_dir: str) -> None: run_name = f"diffusion_nft_krea2_{args.reward}_async_{U.create_run_id()}" num_rollout = args.num_rollout or (3 if args.smoke else 100) + full_ocr = _use_ocr(args) and not args.smoke ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 20 " @@ -83,12 +84,13 @@ def execute(args: ScriptArgs, data_dir: str) -> None: ) + ( "--rollout-batch-size 2 --n-samples-per-prompt 2 " if args.smoke - else "--rollout-batch-size 8 --n-samples-per-prompt 8 " + else f"--rollout-batch-size {6 if full_ocr else 8} --n-samples-per-prompt 8 " ) - eval_args = "--diffusion-eval-num-steps 52 --skip-eval-before-train " + ( - "" if args.smoke else f"--eval-prompt-data {args.reward}_test {data_dir}/test.jsonl --eval-interval 30 " - ) + eval_args = "--diffusion-eval-num-steps 52 --skip-eval-before-train " + if not args.smoke: + eval_args += f"--eval-prompt-data {args.reward}_test {data_dir}/test.jsonl " + eval_args += f"--eval-interval {num_rollout if full_ocr else 30} " grpo_args = ( "--loss-type nft " @@ -139,10 +141,13 @@ def execute(args: ScriptArgs, data_dir: str) -> None: train_backend_args = "--train-backend fsdp --diffusion-forward-dtype bf16 " - perf_args = "--gradient-checkpointing " + ("--micro-batch-size 1 " if args.smoke else "--micro-batch-size 2 ") + micro_batch_size = 1 if args.smoke else (4 if full_ocr else 2) + perf_args = f"--micro-batch-size {micro_batch_size} " + if not full_ocr: + perf_args += "--gradient-checkpointing " misc_args = ( - "--actor-num-gpus-per-node 2 " + f"--actor-num-gpus-per-node {6 if full_ocr else 2} " "--rollout-num-gpus 2 " "--rollout-num-gpus-per-engine 1 " f"--num-gpus-per-node {_num_gpus(args)} " @@ -159,7 +164,6 @@ def execute(args: ScriptArgs, data_dir: str) -> None: train_script="train_diffusion_async.py", config=args, extra_env_vars={ - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "HF_TOKEN": os.environ.get("HF_TOKEN", ""), }, )