From de4c35a65bbafa836bab2b2c1b24b6abce716bc1 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 3 Aug 2026 06:37:19 +0000 Subject: [PATCH 01/19] feat(diffusion): add SFT loss hub and pre-encoded data manager --- miles/backends/fsdp_utils/loss_hub/sft.py | 101 ++++++++++++++ miles/backends/fsdp_utils/metrics.py | 1 + miles/ray/placement_group.py | 6 +- miles/ray/sft_data_manager.py | 81 +++++++++++ miles/utils/arguments.py | 35 +++++ scripts/run-diffusion-sft-wan22.sh | 75 ++++++++++ scripts/sft_encode_wan.py | 126 +++++++++++++++++ .../backends/fsdp_utils/test_loss_hub_sft.py | 130 ++++++++++++++++++ 8 files changed, 553 insertions(+), 2 deletions(-) create mode 100644 miles/backends/fsdp_utils/loss_hub/sft.py create mode 100644 miles/ray/sft_data_manager.py create mode 100644 scripts/run-diffusion-sft-wan22.sh create mode 100644 scripts/sft_encode_wan.py create mode 100644 tests/fast/backends/fsdp_utils/test_loss_hub_sft.py diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py new file mode 100644 index 00000000..ee32b44c --- /dev/null +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -0,0 +1,101 @@ +"""SFT batch preparation and loss formula (rectified-flow velocity MSE).""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch +from miles.backends.fsdp_utils.loss_hub.utils import cast_cond_to_dtype +from miles.utils.metric_buffer import MetricBuffer + + +def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]: + """Pick one DiT component per micro-batch (phase-pure), then grid indices within its range.""" + num_grid = len(ctx.scheduler.timesteps) + if len(ctx.models) == 1: + component_name, model = next(iter(ctx.models.items())) + return component_name, model, torch.randint(num_grid, (bsz,)) + + num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) + config = ctx.train_pipeline_config + components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps] + names = sorted(set(components)) + counts = torch.tensor([components.count(name) for name in names], dtype=torch.float32) + component_name = names[int(torch.multinomial(counts, 1))] + pool = torch.tensor([i for i, name in enumerate(components) if name == component_name]) + return component_name, ctx.models[component_name], pool[torch.randint(len(pool), (bsz,))] + + +def prepare_sft_batch( + ctx: DiffusionLossContext, + batch: list[dict], + *, + pad_to_len: int | None = None, +) -> PreparedBatch: + """Corrupt cached clean latents at sampled grid sigmas; CFG-free cached cond.""" + device = ctx.device + config = ctx.train_pipeline_config + bsz = len(batch) + + x0 = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32) + component_name, model, idx = sample_grid_indices(ctx, bsz) + idx = idx.to(device) + timesteps = ctx.scheduler.timesteps[idx].to(dtype=torch.float32) + sigmas = ctx.scheduler.sigmas[idx].to(dtype=torch.float32) + + noise = torch.randn(x0.shape, device=device, dtype=torch.float32) + sigma_exp = sigmas.view(bsz, *([1] * (x0.ndim - 1))) + latents = (1.0 - sigma_exp) * x0 + sigma_exp * noise + + num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) + if config.needs_timestep_scaling: + timesteps_for_model = timesteps / float(num_train_timesteps) + else: + timesteps_for_model = timesteps + + cond_list = [{key: value.to(device) for key, value in pair["cond_kwargs"].items()} for pair in batch] + pos_cond = cast_cond_to_dtype( + config.collate_cond_for_sample_batch(cond_list, device, pad_to_len=pad_to_len), + ctx.forward_dtype, + ) + + return PreparedBatch( + latents=latents, + timesteps=timesteps, + timesteps_for_model=timesteps_for_model, + model=model, + component_name=component_name, + guidance_scale=0.0, + use_cfg=False, + cfg_batching=False, + true_cfg_scale=None, + pos_cond=pos_cond, + neg_cond=None, + joint_cond=None, + advantage=torch.ones(bsz, device=device, dtype=torch.float32), + extras={"target": noise - x0}, + ) + + +def sft_loss_formula( + ctx: DiffusionLossContext, + batch: list[dict], + prepared: PreparedBatch, + *, + new_pred: torch.Tensor, + ref_pred: torch.Tensor | None, + metrics: MetricBuffer, + write_old_log_prob: bool = False, + old_log_prob_from_new: bool = False, +) -> torch.Tensor: + """Velocity-target MSE: ``||pred - (eps - x0)||^2`` averaged per pair.""" + target = prepared.extras["target"] + per_pair = ((new_pred.float() - target) ** 2).mean(dim=tuple(range(1, target.ndim))) + loss_sum = per_pair.sum() + + bsz = len(batch) + with torch.no_grad(): + metrics.emit_mean("loss", total=loss_sum, count=bsz) + metrics.emit_mean("sft_t_mean", total=prepared.timesteps.sum(), count=bsz) + return loss_sum diff --git a/miles/backends/fsdp_utils/metrics.py b/miles/backends/fsdp_utils/metrics.py index 66af7ffe..b76cd785 100644 --- a/miles/backends/fsdp_utils/metrics.py +++ b/miles/backends/fsdp_utils/metrics.py @@ -31,6 +31,7 @@ "nft_adv_mean": MetricReduce.MEAN, "nft_t_mean": MetricReduce.MEAN, "nft_num_timesteps": MetricReduce.MEAN, + "sft_t_mean": MetricReduce.MEAN, } diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 94d5ab2e..1fda2ae5 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -7,6 +7,7 @@ from .actor_group import RayTrainGroup from .rollout import RolloutManager +from .sft_data_manager import SftDataManager logger = logging.getLogger(__name__) @@ -151,8 +152,9 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): - logger.info("Creating rollout manager (num_gpus=%s)", 0) - rollout_manager = RolloutManager.options( + manager_cls = SftDataManager if args.loss_type == "sft_loss" else RolloutManager + logger.info("Creating rollout manager (%s, num_gpus=%s)", manager_cls.__ray_metadata__.class_name, 0) + rollout_manager = manager_cls.options( num_cpus=1, num_gpus=0, ).remote(args, pg) diff --git a/miles/ray/sft_data_manager.py b/miles/ray/sft_data_manager.py new file mode 100644 index 00000000..9cd9a0ec --- /dev/null +++ b/miles/ray/sft_data_manager.py @@ -0,0 +1,81 @@ +"""Serves pre-encoded SFT samples through the RolloutManager driver interface.""" + +import logging +from pathlib import Path + +import ray +import torch + +from miles.utils.logging_utils import configure_logger +from miles.utils.misc import load_function +from miles.utils.ray_utils import Box +from miles.utils.train_data_utils import TrainDataDPSplitter + +logger = logging.getLogger(__name__) + + +@ray.remote +class SftDataManager: + """Deterministically shuffles --sft-data-path pairs per epoch; generate() is stateless in rollout_id.""" + + def __init__(self, args, pg): + configure_logger() + self.args = args + self.files = sorted(Path(args.sft_data_path).glob("*.pt")) + if len(self.files) < args.rollout_batch_size: + raise ValueError( + f"--sft-data-path holds {len(self.files)} samples, " + f"fewer than rollout_batch_size={args.rollout_batch_size}" + ) + self.train_data_dp_splitter = TrainDataDPSplitter() + + train_pipeline_config = load_function(args.train_pipeline_config_path)() + scheduler = load_function(args.model_backend_path)(train_pipeline_config).load_scheduler(args) + num_train_timesteps = int(scheduler.config.num_train_timesteps) + shift = args.diffusion_flow_shift + sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps, dtype=torch.float64) + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + self.scheduler_timesteps = (sigmas * num_train_timesteps).to(torch.float32) + self.scheduler_sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]).to(torch.float32) + logger.info( + "SftDataManager: %d samples, flow_shift=%s, num_train_timesteps=%d", + len(self.files), + shift, + num_train_timesteps, + ) + + def set_train_parallel_config(self, config: dict): + self.train_parallel_config = config + + def get_num_rollout_per_epoch(self): + return len(self.files) // self.args.rollout_batch_size + + def generate(self, rollout_id): + batch_size = self.args.rollout_batch_size + epoch, slot = divmod(rollout_id, len(self.files) // batch_size) + generator = torch.Generator().manual_seed(self.args.seed + epoch) + perm = torch.randperm(len(self.files), generator=generator) + indices = perm[slot * batch_size : (slot + 1) * batch_size].tolist() + pairs = [torch.load(self.files[i], map_location="cpu") for i in indices] + data = { + "train_data": pairs, + "scheduler_timesteps": self.scheduler_timesteps, + "scheduler_sigmas": self.scheduler_sigmas, + } + shards = self.train_data_dp_splitter.split_by_dp(data, self.train_parallel_config["dp_size"]) + return [Box(ray.put(shard)) for shard in shards] + + def save(self, rollout_id): + pass + + def load(self, rollout_id=None): + pass + + def offload(self): + pass + + def onload_weights(self): + pass + + def dispose(self): + pass diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index e0c9c4ac..4f90d1fa 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -690,6 +690,15 @@ def add_data_arguments(parser): ) parser.add_argument("--input-key", type=str, default="input", help="JSON dataset key") parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key") + parser.add_argument( + "--sft-data-path", + type=str, + default=None, + help=( + "Directory of pre-encoded SFT samples (one .pt per sample holding latent + " + "cond_kwargs, see scripts/sft_encode_wan.py) for --loss-type sft_loss." + ), + ) parser.add_argument( "--start-rollout-id", @@ -1473,6 +1482,14 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def set_default_diffusion_args(args) -> None: + if args.loss_type == "sft_loss": + if args.custom_prepare_train_batch_path is None: + args.custom_prepare_train_batch_path = "miles.backends.fsdp_utils.loss_hub.sft.prepare_sft_batch" + if args.custom_loss_function_path is None: + args.custom_loss_function_path = "miles.backends.fsdp_utils.loss_hub.sft.sft_loss_formula" + # SFT needs no sglang engines; reuse the train-only wiring end to end. + args.debug_train_only = True + is_nft = args.loss_type == "nft" if is_nft: if args.custom_expand_samples_to_train_pairs_path is None: @@ -1585,6 +1602,24 @@ def miles_validate_args(args): if args.ema_rollout_policy == "ema" and not args.ema_shadow: raise ValueError("--ema-rollout-policy ema requires --ema-shadow") + if args.loss_type == "sft_loss": + if args.sft_data_path is None: + raise ValueError("--loss-type sft_loss requires --sft-data-path") + if args.diffusion_flow_shift is None: + raise ValueError("--loss-type sft_loss requires --diffusion-flow-shift for the training sigma grid") + if args.n_samples_per_prompt != 1: + raise ValueError("--loss-type sft_loss requires --n-samples-per-prompt 1") + if args.eval_interval is not None: + raise ValueError("--loss-type sft_loss does not support --eval-interval (no rollout engines)") + if args.diffusion_kl_beta > 0: + raise ValueError("--loss-type sft_loss does not support --diffusion-kl-beta") + if args.diffusion_recompute_old_log_prob: + raise ValueError("--loss-type sft_loss does not support --diffusion-recompute-old-log-prob") + if args.ref_mode != "none": + raise ValueError("--loss-type sft_loss does not use a reference model; drop --ref-mode") + if args.ema_shadow: + raise ValueError("--loss-type sft_loss does not support --ema-shadow (EMA updates run in weight sync)") + is_nft = args.loss_type == "nft" if is_nft: if args.diffusion_noise_level == 0 and args.diffusion_sde_type != "ode": diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh new file mode 100644 index 00000000..ff27993f --- /dev/null +++ b/scripts/run-diffusion-sft-wan22.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# 4-GPU Wan2.2-T2V-A14B dual-expert LoRA SFT on a pre-encoded dataset. +# No sglang engines: the SftDataManager serves cached (latent, cond) pairs. +# +# Encode the raw (video, prompt) jsonl once before training, e.g.: +# python scripts/sft_encode_wan.py \ +# --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ +# --data-path /path/to/train.jsonl --output-dir "${SFT_DATA_DIR}" \ +# --height 480 --width 832 --num-frames 81 --num-gpus 4 +# +# Per rollout step: 64 samples, num_steps_per_rollout=4 +# -> 16 samples/optim step / 4 dp ranks = 4 samples/rank at mbs=1. + +set -euo pipefail +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +RUN_NAME="${RUN_NAME:-diffusion_sft_wan22_$(date +%Y%m%d_%H%M%S)}" +SAVE_DIR="${ROOT_DIR}/logs/${RUN_NAME}/ckpt" + +SFT_DATA_DIR="${SFT_DATA_DIR:?set SFT_DATA_DIR to the sft_encode_wan.py output dir}" + +WANDB_ARGS=() +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS+=( + --use-wandb + --wandb-project miles-diffusion-sft + --wandb-group "${RUN_NAME}" + --wandb-key "${WANDB_API_KEY}" + --disable-wandb-random-suffix + ) +fi + +PYTHON_BIN="${PYTHON_BIN:-python}" + +RESUME_ARGS=() +if [[ -n "${RESUME_CKPT:-}" ]]; then + RESUME_ARGS+=(--load "${RESUME_CKPT}") + [[ -n "${START_ROLLOUT:-}" ]] && RESUME_ARGS+=(--start-rollout-id "${START_ROLLOUT}") +fi + +WAN_LORA_TARGET_MODULES=( + attn1.to_q attn1.to_k attn1.to_v attn1.to_out.0 + attn2.to_q attn2.to_k attn2.to_v attn2.to_out.0 + ffn.net.0.proj ffn.net.2 +) + +"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ + --train-backend fsdp \ + --loss-type sft_loss \ + --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --sft-data-path "${SFT_DATA_DIR}" \ + --rollout-batch-size 64 \ + --num-epoch 3 \ + --num-steps-per-rollout 4 \ + --micro-batch-size 1 \ + --actor-num-gpus-per-node 4 \ + --num-gpus-per-node 4 \ + --use-lora \ + --lora-rank 64 \ + --lora-alpha 128 \ + --lora-target-modules "${WAN_LORA_TARGET_MODULES[@]}" \ + --diffusion-init-lora-weight gaussian \ + --lr 1e-4 \ + --adam-beta2 0.999 \ + --weight-decay 1e-4 \ + --update-weight-target-module transformer,transformer_2 \ + --fsdp-master-dtype fp32 \ + --fsdp-reduce-dtype fp32 \ + --diffusion-forward-dtype bf16 \ + --diffusion-flow-shift 3.0 \ + --save "${SAVE_DIR}" \ + --save-interval 20 \ + "${RESUME_ARGS[@]}" \ + "${WANDB_ARGS[@]}" diff --git a/scripts/sft_encode_wan.py b/scripts/sft_encode_wan.py new file mode 100644 index 00000000..faf4397c --- /dev/null +++ b/scripts/sft_encode_wan.py @@ -0,0 +1,126 @@ +"""Pre-encode a (video, prompt) jsonl dataset into Wan SFT train pairs. + +Each output ``{index:08d}.pt`` holds ``{"latent": [C,T,H,W] fp16 (VAE-normalized), +"cond_kwargs": {"encoder_hidden_states": [1, 512, D] bf16}, "prompt": str}``, +the pair schema consumed by miles.backends.fsdp_utils.loss_hub.sft. +""" + +import argparse +import json +from pathlib import Path + +import ray +import torch + + +def read_video_clip(path: str, *, height: int, width: int, num_frames: int) -> torch.Tensor: + import torchvision + + frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") + if frames.shape[0] < num_frames: + raise ValueError(f"{path} has {frames.shape[0]} frames, need {num_frames}") + start = (frames.shape[0] - num_frames) // 2 + frames = frames[start : start + num_frames].float() / 127.5 - 1.0 + + scale = max(height / frames.shape[2], width / frames.shape[3]) + new_h = max(height, round(frames.shape[2] * scale)) + new_w = max(width, round(frames.shape[3] * scale)) + frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) + top = (new_h - height) // 2 + left = (new_w - width) // 2 + return frames[:, :, top : top + height, left : left + width].permute(1, 0, 2, 3) + + +@ray.remote(num_gpus=1) +class WanEncodeActor: + def __init__(self, checkpoint: str): + from diffusers import AutoencoderKLWan + from transformers import AutoTokenizer, UMT5EncoderModel + + self.device = torch.device("cuda") + self.tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="tokenizer") + self.text_encoder = UMT5EncoderModel.from_pretrained( + checkpoint, subfolder="text_encoder", torch_dtype=torch.bfloat16 + ).to(self.device) + self.vae = AutoencoderKLWan.from_pretrained(checkpoint, subfolder="vae", torch_dtype=torch.float32).to( + self.device + ) + view = (1, self.vae.config.z_dim, 1, 1, 1) + self.latents_mean = torch.tensor(self.vae.config.latents_mean).view(view).to(self.device) + self.latents_std = torch.tensor(self.vae.config.latents_std).view(view).to(self.device) + + @torch.no_grad() + def encode(self, items: list[dict], output_dir: str, height: int, width: int, num_frames: int) -> int: + from diffusers.pipelines.wan.pipeline_wan import prompt_clean + + done = 0 + for item in items: + out_path = Path(output_dir) / f"{item['index']:08d}.pt" + if out_path.exists(): + continue + video = read_video_clip(item["video"], height=height, width=width, num_frames=num_frames) + latent = self.vae.encode(video.unsqueeze(0).to(self.device)).latent_dist.sample() + latent = (latent - self.latents_mean) / self.latents_std + + inputs = self.tokenizer( + [prompt_clean(item["prompt"])], + padding="max_length", + max_length=512, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + embeds = self.text_encoder( + inputs.input_ids.to(self.device), inputs.attention_mask.to(self.device) + ).last_hidden_state + embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 + + torch.save( + { + "latent": latent[0].to(torch.float16).cpu(), + "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, + "prompt": item["prompt"], + }, + out_path, + ) + done += 1 + return done + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--hf-checkpoint", required=True) + parser.add_argument("--data-path", required=True, help="jsonl with one {video, prompt} object per line") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--video-key", default="video") + parser.add_argument("--prompt-key", default="prompt") + parser.add_argument("--height", type=int, required=True) + parser.add_argument("--width", type=int, required=True) + parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--num-gpus", type=int, default=1) + args = parser.parse_args() + + if (args.num_frames - 1) % 4 != 0: + parser.error("--num-frames must be 4k+1 for the Wan VAE temporal stride") + + items = [] + with open(args.data_path) as f: + for index, line in enumerate(f): + row = json.loads(line) + items.append({"index": index, "video": row[args.video_key], "prompt": row[args.prompt_key]}) + + Path(args.output_dir).mkdir(parents=True, exist_ok=True) + ray.init() + actors = [WanEncodeActor.remote(args.hf_checkpoint) for _ in range(args.num_gpus)] + done = ray.get( + [ + actor.encode.remote(items[i :: args.num_gpus], args.output_dir, args.height, args.width, args.num_frames) + for i, actor in enumerate(actors) + ] + ) + print(f"encoded {sum(done)} new samples into {args.output_dir} ({len(items)} total)") + + +if __name__ == "__main__": + main() diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py new file mode 100644 index 00000000..d1420baa --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py @@ -0,0 +1,130 @@ +"""Smoke tests for diffusion SFT hooks (prepare + loss formula; actor owns DiT).""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +import torch +import torch.nn as nn + +from miles.backends.fsdp_utils.loss_hub.sft import prepare_sft_batch, sample_grid_indices, sft_loss_formula +from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext + +NUM_TRAIN_TIMESTEPS = 1000 +NUM_GRID = 8 + + +class _Config: + needs_timestep_scaling = False + + def collate_cond_for_sample_batch(self, per_sample_cond_kwargs, device, pad_to_len=None): + return {"encoder_hidden_states": torch.cat([kw["encoder_hidden_states"] for kw in per_sample_cond_kwargs])} + + def component_for_timestep(self, timestep, num_train_timesteps): + return "transformer" if timestep >= 0.875 * num_train_timesteps else "transformer_2" + + +def _scheduler(): + sigmas = torch.linspace(1.0, 1.0 / NUM_GRID, NUM_GRID) + return Namespace( + timesteps=sigmas * NUM_TRAIN_TIMESTEPS, + sigmas=torch.cat([sigmas, torch.zeros(1)]), + config=Namespace(num_train_timesteps=NUM_TRAIN_TIMESTEPS), + ) + + +def _ctx(models): + return DiffusionLossContext( + models=models, + train_pipeline_config=_Config(), + sde_backend=None, + scheduler=_scheduler(), + args=Namespace(), + forward_dtype=torch.float32, + device=torch.device("cpu"), + ) + + +def _batch(bsz=4): + return [ + {"latent": torch.randn(16, 2, 4, 4), "cond_kwargs": {"encoder_hidden_states": torch.randn(1, 6, 8)}} + for _ in range(bsz) + ] + + +class _Metrics: + def __init__(self): + self.seen = {} + + def emit_mean(self, key, *, total, count): + self.seen[key] = (float(total), count) + + +class TestPrepareSftBatch: + def test_corruption_and_target_identity(self): + torch.manual_seed(0) + ctx = _ctx({"transformer": nn.Identity()}) + batch = _batch() + prepared = prepare_sft_batch(ctx, batch) + + x0 = torch.stack([pair["latent"] for pair in batch]).float() + sigma = (prepared.timesteps / NUM_TRAIN_TIMESTEPS).view(-1, 1, 1, 1, 1) + assert torch.allclose(prepared.latents, x0 + sigma * prepared.extras["target"], atol=1e-5) + assert not prepared.use_cfg + assert prepared.pos_cond["encoder_hidden_states"].shape == (4, 6, 8) + assert torch.equal(prepared.timesteps_for_model, prepared.timesteps) + + def test_timestep_scaling(self): + torch.manual_seed(0) + ctx = _ctx({"transformer": nn.Identity()}) + ctx.train_pipeline_config.needs_timestep_scaling = True + prepared = prepare_sft_batch(ctx, _batch()) + assert torch.allclose(prepared.timesteps_for_model, prepared.timesteps / NUM_TRAIN_TIMESTEPS) + + def test_dual_expert_micro_batch_is_phase_pure(self): + torch.manual_seed(0) + models = {"transformer": nn.Identity(), "transformer_2": nn.Identity()} + ctx = _ctx(models) + config = ctx.train_pipeline_config + picked = set() + for _ in range(20): + name, model, idx = sample_grid_indices(ctx, bsz=4) + picked.add(name) + assert model is models[name] + for i in idx.tolist(): + t = float(ctx.scheduler.timesteps[i]) + assert config.component_for_timestep(t, NUM_TRAIN_TIMESTEPS) == name + assert picked == {"transformer", "transformer_2"} + + +class TestSftLossFormula: + def test_zero_loss_on_exact_velocity(self): + torch.manual_seed(0) + ctx = _ctx({"transformer": nn.Identity()}) + batch = _batch() + prepared = prepare_sft_batch(ctx, batch) + metrics = _Metrics() + loss = sft_loss_formula( + ctx, batch, prepared, new_pred=prepared.extras["target"], ref_pred=None, metrics=metrics + ) + assert torch.allclose(loss, torch.zeros(())) + + def test_unit_offset_loss(self): + torch.manual_seed(0) + ctx = _ctx({"transformer": nn.Identity()}) + batch = _batch() + prepared = prepare_sft_batch(ctx, batch) + metrics = _Metrics() + loss = sft_loss_formula( + ctx, + batch, + prepared, + new_pred=prepared.extras["target"] + 1.0, + ref_pred=None, + metrics=metrics, + ) + assert torch.allclose(loss, torch.tensor(float(len(batch)))) + assert metrics.seen["loss"] == (float(len(batch)), len(batch)) + assert "sft_t_mean" in metrics.seen From 4c9ff871f74fa74e7dc32b132ed58d4595d2104d Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 3 Aug 2026 07:10:00 +0000 Subject: [PATCH 02/19] feat(diffusion): add frame stride to SFT encode script --- scripts/sft_encode_wan.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/scripts/sft_encode_wan.py b/scripts/sft_encode_wan.py index faf4397c..b02c9adc 100644 --- a/scripts/sft_encode_wan.py +++ b/scripts/sft_encode_wan.py @@ -13,14 +13,15 @@ import torch -def read_video_clip(path: str, *, height: int, width: int, num_frames: int) -> torch.Tensor: +def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: import torchvision frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") - if frames.shape[0] < num_frames: - raise ValueError(f"{path} has {frames.shape[0]} frames, need {num_frames}") - start = (frames.shape[0] - num_frames) // 2 - frames = frames[start : start + num_frames].float() / 127.5 - 1.0 + span = (num_frames - 1) * frame_stride + 1 + if frames.shape[0] < span: + raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") + start = (frames.shape[0] - span) // 2 + frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 scale = max(height / frames.shape[2], width / frames.shape[3]) new_h = max(height, round(frames.shape[2] * scale)) @@ -50,7 +51,9 @@ def __init__(self, checkpoint: str): self.latents_std = torch.tensor(self.vae.config.latents_std).view(view).to(self.device) @torch.no_grad() - def encode(self, items: list[dict], output_dir: str, height: int, width: int, num_frames: int) -> int: + def encode( + self, items: list[dict], output_dir: str, height: int, width: int, num_frames: int, frame_stride: int + ) -> int: from diffusers.pipelines.wan.pipeline_wan import prompt_clean done = 0 @@ -58,7 +61,9 @@ def encode(self, items: list[dict], output_dir: str, height: int, width: int, nu out_path = Path(output_dir) / f"{item['index']:08d}.pt" if out_path.exists(): continue - video = read_video_clip(item["video"], height=height, width=width, num_frames=num_frames) + video = read_video_clip( + item["video"], height=height, width=width, num_frames=num_frames, frame_stride=frame_stride + ) latent = self.vae.encode(video.unsqueeze(0).to(self.device)).latent_dist.sample() latent = (latent - self.latents_mean) / self.latents_std @@ -98,6 +103,7 @@ def main(): parser.add_argument("--height", type=int, required=True) parser.add_argument("--width", type=int, required=True) parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--frame-stride", type=int, default=1) parser.add_argument("--num-gpus", type=int, default=1) args = parser.parse_args() @@ -115,7 +121,9 @@ def main(): actors = [WanEncodeActor.remote(args.hf_checkpoint) for _ in range(args.num_gpus)] done = ray.get( [ - actor.encode.remote(items[i :: args.num_gpus], args.output_dir, args.height, args.width, args.num_frames) + actor.encode.remote( + items[i :: args.num_gpus], args.output_dir, args.height, args.width, args.num_frames, args.frame_stride + ) for i, actor in enumerate(actors) ] ) From 1616581eaa06a5f70a0372b4b488e26050069772 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 4 Aug 2026 00:02:54 +0000 Subject: [PATCH 03/19] feat(diffusion): auto-build SFT cache via family encode hooks --- .../configs/train_pipeline_config.py | 8 ++ miles/backends/fsdp_utils/configs/wan2_2.py | 54 +++++++ miles/ray/sft_data_manager.py | 37 ++++- miles/ray/sft_encode.py | 77 ++++++++++ miles/utils/arguments.py | 22 ++- scripts/run-diffusion-sft-wan22.sh | 20 +-- scripts/sft_encode_wan.py | 134 ------------------ tests/fast/ray/__init__.py | 0 tests/fast/ray/test_sft_cache_dir.py | 31 ++++ 9 files changed, 230 insertions(+), 153 deletions(-) create mode 100644 miles/ray/sft_encode.py delete mode 100644 scripts/sft_encode_wan.py create mode 100644 tests/fast/ray/__init__.py create mode 100644 tests/fast/ray/test_sft_cache_dir.py diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c392ea57..e7df0d0f 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -189,3 +189,11 @@ def cfg_combine( def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: """Postprocess the model after FSDP wrap + weight materialization (default: no-op).""" return None + + def load_sft_encoder(self, args, device: torch.device): + """Load this family's frozen encode components (tokenizer/text encoder/VAE) for SFT caching.""" + raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") + + def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str) -> dict: + """Encode one (pixels [C,T,H,W] in [-1,1], prompt) into the train-pair dict for prepare_sft_batch.""" + raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 7e227cd2..9a307774 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -66,3 +66,57 @@ def cfg_combine( ) -> torch.Tensor: scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) + + @classmethod + def validate_args(cls, args) -> None: + if args.loss_type == "sft_loss" and (args.sft_num_frames - 1) % 4 != 0: + raise ValueError("--sft-num-frames must be 4k+1 for the Wan VAE temporal stride") + + def load_sft_encoder(self, args, device: torch.device): + from diffusers import AutoencoderKLWan + from transformers import AutoTokenizer, UMT5EncoderModel + + tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, subfolder="tokenizer") + text_encoder = UMT5EncoderModel.from_pretrained( + args.hf_checkpoint, subfolder="text_encoder", torch_dtype=torch.bfloat16 + ).to(device) + vae = AutoencoderKLWan.from_pretrained(args.hf_checkpoint, subfolder="vae", torch_dtype=torch.float32).to( + device + ) + view = (1, vae.config.z_dim, 1, 1, 1) + return { + "device": device, + "tokenizer": tokenizer, + "text_encoder": text_encoder, + "vae": vae, + "latents_mean": torch.tensor(vae.config.latents_mean).view(view).to(device), + "latents_std": torch.tensor(vae.config.latents_std).view(view).to(device), + } + + @torch.no_grad() + def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str) -> dict: + from diffusers.pipelines.wan.pipeline_wan import prompt_clean + + device = encoder["device"] + latent = encoder["vae"].encode(pixels.unsqueeze(0).to(device, torch.float32)).latent_dist.sample() + latent = (latent - encoder["latents_mean"]) / encoder["latents_std"] + + inputs = encoder["tokenizer"]( + [prompt_clean(prompt)], + padding="max_length", + max_length=512, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + embeds = encoder["text_encoder"]( + inputs.input_ids.to(device), inputs.attention_mask.to(device) + ).last_hidden_state + embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 + + return { + "latent": latent[0].to(torch.float16).cpu(), + "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, + "prompt": prompt, + } diff --git a/miles/ray/sft_data_manager.py b/miles/ray/sft_data_manager.py index 9cd9a0ec..8b66079d 100644 --- a/miles/ray/sft_data_manager.py +++ b/miles/ray/sft_data_manager.py @@ -1,11 +1,14 @@ -"""Serves pre-encoded SFT samples through the RolloutManager driver interface.""" +"""Serves auto-cached SFT samples through the RolloutManager driver interface.""" +import hashlib +import json import logging from pathlib import Path import ray import torch +from miles.ray.sft_encode import build_sft_cache from miles.utils.logging_utils import configure_logger from miles.utils.misc import load_function from miles.utils.ray_utils import Box @@ -14,19 +17,38 @@ logger = logging.getLogger(__name__) +def sft_cache_dir(args, data_path: Path) -> Path: + key = hashlib.sha256( + f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}" + f"|{args.sft_num_frames}s{args.sft_frame_stride}".encode() + data_path.read_bytes() + ).hexdigest()[:12] + return data_path.parent / ".sft_cache" / key + + @ray.remote class SftDataManager: - """Deterministically shuffles --sft-data-path pairs per epoch; generate() is stateless in rollout_id.""" + """Encodes --sft-data-path (jsonl) into a content-addressed cache on first run, + then serves per-epoch shuffled pairs; generate() is stateless in rollout_id.""" def __init__(self, args, pg): configure_logger() self.args = args - self.files = sorted(Path(args.sft_data_path).glob("*.pt")) - if len(self.files) < args.rollout_batch_size: + data_path = Path(args.sft_data_path) + items = [ + {"index": i, "video": row[args.sft_video_key], "prompt": row[args.sft_prompt_key]} + for i, row in enumerate(json.loads(line) for line in data_path.read_text().splitlines() if line.strip()) + ] + if len(items) < args.rollout_batch_size: raise ValueError( - f"--sft-data-path holds {len(self.files)} samples, " - f"fewer than rollout_batch_size={args.rollout_batch_size}" + f"--sft-data-path holds {len(items)} samples, fewer than rollout_batch_size={args.rollout_batch_size}" ) + + cache_dir = sft_cache_dir(args, data_path) + if len(list(cache_dir.glob("*.pt"))) < len(items): + logger.info("SftDataManager: building cache %s for %d samples", cache_dir, len(items)) + build_sft_cache(args, items, cache_dir, pg[0]) + self.files = sorted(cache_dir.glob("*.pt")) + assert len(self.files) == len(items) self.train_data_dp_splitter = TrainDataDPSplitter() train_pipeline_config = load_function(args.train_pipeline_config_path)() @@ -38,8 +60,9 @@ def __init__(self, args, pg): self.scheduler_timesteps = (sigmas * num_train_timesteps).to(torch.float32) self.scheduler_sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]).to(torch.float32) logger.info( - "SftDataManager: %d samples, flow_shift=%s, num_train_timesteps=%d", + "SftDataManager: %d samples, cache=%s, flow_shift=%s, num_train_timesteps=%d", len(self.files), + cache_dir, shift, num_train_timesteps, ) diff --git a/miles/ray/sft_encode.py b/miles/ray/sft_encode.py new file mode 100644 index 00000000..ef49019c --- /dev/null +++ b/miles/ray/sft_encode.py @@ -0,0 +1,77 @@ +"""Family-generic SFT cache building: encode actors driven by TrainPipelineConfig hooks.""" + +import logging +from pathlib import Path + +import ray +import torch +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from miles.utils.misc import load_function + +logger = logging.getLogger(__name__) + + +def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: + import torchvision + + frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") + span = (num_frames - 1) * frame_stride + 1 + if frames.shape[0] < span: + raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") + start = (frames.shape[0] - span) // 2 + frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 + + scale = max(height / frames.shape[2], width / frames.shape[3]) + new_h = max(height, round(frames.shape[2] * scale)) + new_w = max(width, round(frames.shape[3] * scale)) + frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) + top = (new_h - height) // 2 + left = (new_w - width) // 2 + return frames[:, :, top : top + height, left : left + width].permute(1, 0, 2, 3) + + +@ray.remote +class SftEncodeActor: + def __init__(self, args): + self.args = args + self.config = load_function(args.train_pipeline_config_path)() + self.encoder = self.config.load_sft_encoder(args, torch.device("cuda")) + + def encode(self, items: list[dict], cache_dir: str) -> int: + args = self.args + done = 0 + for item in items: + out_path = Path(cache_dir) / f"{item['index']:08d}.pt" + if out_path.exists(): + continue + pixels = read_video_clip( + item["video"], + height=args.sft_height, + width=args.sft_width, + num_frames=args.sft_num_frames, + frame_stride=args.sft_frame_stride, + ) + torch.save(self.config.encode_sft_sample(self.encoder, pixels, item["prompt"]), out_path) + done += 1 + return done + + +def build_sft_cache(args, items: list[dict], cache_dir: Path, pg) -> None: + cache_dir.mkdir(parents=True, exist_ok=True) + num_workers = min(len(pg.bundle_specs), len(items)) + actors = [ + SftEncodeActor.options( + num_cpus=1, + num_gpus=1, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ), + ).remote(args) + for i in range(num_workers) + ] + done = ray.get([actor.encode.remote(items[i::num_workers], str(cache_dir)) for i, actor in enumerate(actors)]) + for actor in actors: + ray.kill(actor) + logger.info("SFT cache: encoded %d new samples into %s (%d workers)", sum(done), cache_dir, num_workers) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4f90d1fa..57c5c253 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -695,10 +695,17 @@ def add_data_arguments(parser): type=str, default=None, help=( - "Directory of pre-encoded SFT samples (one .pt per sample holding latent + " - "cond_kwargs, see scripts/sft_encode_wan.py) for --loss-type sft_loss." + "SFT dataset jsonl for --loss-type sft_loss, one object per line with the video " + "path and prompt. Encoded pairs are cached next to it under .sft_cache/ and " + "rebuilt automatically whenever the data or encode settings change." ), ) + parser.add_argument("--sft-video-key", type=str, default="video", help="SFT jsonl video path key") + parser.add_argument("--sft-prompt-key", type=str, default="prompt", help="SFT jsonl prompt key") + parser.add_argument("--sft-height", type=int, default=None, help="SFT encode height (center crop)") + parser.add_argument("--sft-width", type=int, default=None, help="SFT encode width (center crop)") + parser.add_argument("--sft-num-frames", type=int, default=None, help="SFT encode frames per clip") + parser.add_argument("--sft-frame-stride", type=int, default=1, help="SFT encode temporal stride") parser.add_argument( "--start-rollout-id", @@ -1605,6 +1612,17 @@ def miles_validate_args(args): if args.loss_type == "sft_loss": if args.sft_data_path is None: raise ValueError("--loss-type sft_loss requires --sft-data-path") + if args.sft_height is None or args.sft_width is None or args.sft_num_frames is None: + raise ValueError("--loss-type sft_loss requires --sft-height, --sft-width and --sft-num-frames") + from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig + from miles.utils.misc import load_function + + sft_cfg_cls = load_function(args.train_pipeline_config_path) + if sft_cfg_cls.encode_sft_sample is TrainPipelineConfig.encode_sft_sample: + raise ValueError( + f"--loss-type sft_loss is not supported for {sft_cfg_cls.__name__}: " + "it does not implement load_sft_encoder/encode_sft_sample" + ) if args.diffusion_flow_shift is None: raise ValueError("--loss-type sft_loss requires --diffusion-flow-shift for the training sigma grid") if args.n_samples_per_prompt != 1: diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index ff27993f..578b199a 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -1,12 +1,8 @@ #!/usr/bin/env bash -# 4-GPU Wan2.2-T2V-A14B dual-expert LoRA SFT on a pre-encoded dataset. -# No sglang engines: the SftDataManager serves cached (latent, cond) pairs. -# -# Encode the raw (video, prompt) jsonl once before training, e.g.: -# python scripts/sft_encode_wan.py \ -# --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ -# --data-path /path/to/train.jsonl --output-dir "${SFT_DATA_DIR}" \ -# --height 480 --width 832 --num-frames 81 --num-gpus 4 +# 4-GPU Wan2.2-T2V-A14B dual-expert LoRA SFT on a (video, prompt) jsonl dataset. +# No sglang engines: the SftDataManager encodes the dataset into .sft_cache/ +# next to the jsonl on first run (re-encoding automatically when data or encode +# settings change), then serves cached (latent, cond) pairs. # # Per rollout step: 64 samples, num_steps_per_rollout=4 # -> 16 samples/optim step / 4 dp ranks = 4 samples/rank at mbs=1. @@ -17,7 +13,7 @@ export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" RUN_NAME="${RUN_NAME:-diffusion_sft_wan22_$(date +%Y%m%d_%H%M%S)}" SAVE_DIR="${ROOT_DIR}/logs/${RUN_NAME}/ckpt" -SFT_DATA_DIR="${SFT_DATA_DIR:?set SFT_DATA_DIR to the sft_encode_wan.py output dir}" +SFT_DATA_JSONL="${SFT_DATA_JSONL:?set SFT_DATA_JSONL to a jsonl with one {video, prompt} object per line}" WANDB_ARGS=() if [[ -n "${WANDB_API_KEY:-}" ]]; then @@ -49,7 +45,11 @@ WAN_LORA_TARGET_MODULES=( --loss-type sft_loss \ --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --sft-data-path "${SFT_DATA_DIR}" \ + --sft-data-path "${SFT_DATA_JSONL}" \ + --sft-height 480 \ + --sft-width 832 \ + --sft-num-frames 81 \ + --sft-frame-stride 2 \ --rollout-batch-size 64 \ --num-epoch 3 \ --num-steps-per-rollout 4 \ diff --git a/scripts/sft_encode_wan.py b/scripts/sft_encode_wan.py deleted file mode 100644 index b02c9adc..00000000 --- a/scripts/sft_encode_wan.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Pre-encode a (video, prompt) jsonl dataset into Wan SFT train pairs. - -Each output ``{index:08d}.pt`` holds ``{"latent": [C,T,H,W] fp16 (VAE-normalized), -"cond_kwargs": {"encoder_hidden_states": [1, 512, D] bf16}, "prompt": str}``, -the pair schema consumed by miles.backends.fsdp_utils.loss_hub.sft. -""" - -import argparse -import json -from pathlib import Path - -import ray -import torch - - -def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: - import torchvision - - frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") - span = (num_frames - 1) * frame_stride + 1 - if frames.shape[0] < span: - raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") - start = (frames.shape[0] - span) // 2 - frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 - - scale = max(height / frames.shape[2], width / frames.shape[3]) - new_h = max(height, round(frames.shape[2] * scale)) - new_w = max(width, round(frames.shape[3] * scale)) - frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) - top = (new_h - height) // 2 - left = (new_w - width) // 2 - return frames[:, :, top : top + height, left : left + width].permute(1, 0, 2, 3) - - -@ray.remote(num_gpus=1) -class WanEncodeActor: - def __init__(self, checkpoint: str): - from diffusers import AutoencoderKLWan - from transformers import AutoTokenizer, UMT5EncoderModel - - self.device = torch.device("cuda") - self.tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="tokenizer") - self.text_encoder = UMT5EncoderModel.from_pretrained( - checkpoint, subfolder="text_encoder", torch_dtype=torch.bfloat16 - ).to(self.device) - self.vae = AutoencoderKLWan.from_pretrained(checkpoint, subfolder="vae", torch_dtype=torch.float32).to( - self.device - ) - view = (1, self.vae.config.z_dim, 1, 1, 1) - self.latents_mean = torch.tensor(self.vae.config.latents_mean).view(view).to(self.device) - self.latents_std = torch.tensor(self.vae.config.latents_std).view(view).to(self.device) - - @torch.no_grad() - def encode( - self, items: list[dict], output_dir: str, height: int, width: int, num_frames: int, frame_stride: int - ) -> int: - from diffusers.pipelines.wan.pipeline_wan import prompt_clean - - done = 0 - for item in items: - out_path = Path(output_dir) / f"{item['index']:08d}.pt" - if out_path.exists(): - continue - video = read_video_clip( - item["video"], height=height, width=width, num_frames=num_frames, frame_stride=frame_stride - ) - latent = self.vae.encode(video.unsqueeze(0).to(self.device)).latent_dist.sample() - latent = (latent - self.latents_mean) / self.latents_std - - inputs = self.tokenizer( - [prompt_clean(item["prompt"])], - padding="max_length", - max_length=512, - truncation=True, - add_special_tokens=True, - return_attention_mask=True, - return_tensors="pt", - ) - embeds = self.text_encoder( - inputs.input_ids.to(self.device), inputs.attention_mask.to(self.device) - ).last_hidden_state - embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 - - torch.save( - { - "latent": latent[0].to(torch.float16).cpu(), - "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, - "prompt": item["prompt"], - }, - out_path, - ) - done += 1 - return done - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--hf-checkpoint", required=True) - parser.add_argument("--data-path", required=True, help="jsonl with one {video, prompt} object per line") - parser.add_argument("--output-dir", required=True) - parser.add_argument("--video-key", default="video") - parser.add_argument("--prompt-key", default="prompt") - parser.add_argument("--height", type=int, required=True) - parser.add_argument("--width", type=int, required=True) - parser.add_argument("--num-frames", type=int, required=True) - parser.add_argument("--frame-stride", type=int, default=1) - parser.add_argument("--num-gpus", type=int, default=1) - args = parser.parse_args() - - if (args.num_frames - 1) % 4 != 0: - parser.error("--num-frames must be 4k+1 for the Wan VAE temporal stride") - - items = [] - with open(args.data_path) as f: - for index, line in enumerate(f): - row = json.loads(line) - items.append({"index": index, "video": row[args.video_key], "prompt": row[args.prompt_key]}) - - Path(args.output_dir).mkdir(parents=True, exist_ok=True) - ray.init() - actors = [WanEncodeActor.remote(args.hf_checkpoint) for _ in range(args.num_gpus)] - done = ray.get( - [ - actor.encode.remote( - items[i :: args.num_gpus], args.output_dir, args.height, args.width, args.num_frames, args.frame_stride - ) - for i, actor in enumerate(actors) - ] - ) - print(f"encoded {sum(done)} new samples into {args.output_dir} ({len(items)} total)") - - -if __name__ == "__main__": - main() diff --git a/tests/fast/ray/__init__.py b/tests/fast/ray/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fast/ray/test_sft_cache_dir.py b/tests/fast/ray/test_sft_cache_dir.py new file mode 100644 index 00000000..920a4b8e --- /dev/null +++ b/tests/fast/ray/test_sft_cache_dir.py @@ -0,0 +1,31 @@ +"""Cache addressing for the SFT data manager.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +from miles.ray.sft_data_manager import sft_cache_dir + + +def _args(**overrides): + base = dict(hf_checkpoint="ckpt", sft_height=480, sft_width=832, sft_num_frames=81, sft_frame_stride=2) + base.update(overrides) + return Namespace(**base) + + +def test_cache_dir_is_stable_and_content_addressed(tmp_path): + data = tmp_path / "train.jsonl" + data.write_text('{"video": "a.mp4", "prompt": "p"}\n') + + base = sft_cache_dir(_args(), data) + assert base == sft_cache_dir(_args(), data) + assert base.parent == tmp_path / ".sft_cache" + + assert sft_cache_dir(_args(sft_height=512), data) != base + assert sft_cache_dir(_args(sft_frame_stride=1), data) != base + assert sft_cache_dir(_args(hf_checkpoint="other"), data) != base + + data.write_text('{"video": "b.mp4", "prompt": "p"}\n') + assert sft_cache_dir(_args(), data) != base From e1e871001a73d26f9f9ae64cfe9a0195c11604ae Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 4 Aug 2026 03:13:36 +0000 Subject: [PATCH 04/19] fix(sft): per-sample content-addressed cache, atomic writes, seeded VAE sampling --- .../configs/train_pipeline_config.py | 7 ++- miles/backends/fsdp_utils/configs/wan2_2.py | 4 +- miles/ray/sft_data_manager.py | 42 +++++++++----- miles/ray/sft_encode.py | 16 +++--- tests/fast/ray/test_sft_cache_dir.py | 31 ---------- tests/fast/ray/test_sft_sample_key.py | 57 +++++++++++++++++++ 6 files changed, 101 insertions(+), 56 deletions(-) delete mode 100644 tests/fast/ray/test_sft_cache_dir.py create mode 100644 tests/fast/ray/test_sft_sample_key.py diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index e7df0d0f..9b5a20e5 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -194,6 +194,9 @@ def load_sft_encoder(self, args, device: torch.device): """Load this family's frozen encode components (tokenizer/text encoder/VAE) for SFT caching.""" raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") - def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str) -> dict: - """Encode one (pixels [C,T,H,W] in [-1,1], prompt) into the train-pair dict for prepare_sft_batch.""" + def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict: + """Encode one (pixels [C,T,H,W] in [-1,1], prompt) into the train-pair dict for prepare_sft_batch. + + ``generator`` seeds any stochastic encode step (e.g. VAE posterior sampling) so a cache + entry's content is a deterministic function of its cache key.""" raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 9a307774..2fad8f52 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -94,11 +94,11 @@ def load_sft_encoder(self, args, device: torch.device): } @torch.no_grad() - def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str) -> dict: + def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict: from diffusers.pipelines.wan.pipeline_wan import prompt_clean device = encoder["device"] - latent = encoder["vae"].encode(pixels.unsqueeze(0).to(device, torch.float32)).latent_dist.sample() + latent = encoder["vae"].encode(pixels.unsqueeze(0).to(device, torch.float32)).latent_dist.sample(generator) latent = (latent - encoder["latents_mean"]) / encoder["latents_std"] inputs = encoder["tokenizer"]( diff --git a/miles/ray/sft_data_manager.py b/miles/ray/sft_data_manager.py index 8b66079d..ffb30fb6 100644 --- a/miles/ray/sft_data_manager.py +++ b/miles/ray/sft_data_manager.py @@ -17,12 +17,14 @@ logger = logging.getLogger(__name__) -def sft_cache_dir(args, data_path: Path) -> Path: - key = hashlib.sha256( - f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}" - f"|{args.sft_num_frames}s{args.sft_frame_stride}".encode() + data_path.read_bytes() - ).hexdigest()[:12] - return data_path.parent / ".sft_cache" / key +def sft_sample_key(args, item: dict) -> tuple[str, int]: + """Content-addressed cache filename and latent-sampling seed for one (video, prompt) item.""" + stat = Path(item["video"]).stat() + digest = hashlib.sha256( + f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}|{args.sft_num_frames}s{args.sft_frame_stride}" + f"|{item['video']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() + ).digest() + return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 @ray.remote @@ -35,20 +37,32 @@ def __init__(self, args, pg): self.args = args data_path = Path(args.sft_data_path) items = [ - {"index": i, "video": row[args.sft_video_key], "prompt": row[args.sft_prompt_key]} - for i, row in enumerate(json.loads(line) for line in data_path.read_text().splitlines() if line.strip()) + {"video": row[args.sft_video_key], "prompt": row[args.sft_prompt_key]} + for row in (json.loads(line) for line in data_path.read_text().splitlines() if line.strip()) ] if len(items) < args.rollout_batch_size: raise ValueError( f"--sft-data-path holds {len(items)} samples, fewer than rollout_batch_size={args.rollout_batch_size}" ) + dropped = len(items) % args.rollout_batch_size + if dropped: + logger.warning( + "SFT drop-last: %d of %d samples unused per epoch (rollout_batch_size=%d); " + "the per-epoch reshuffle rotates which samples are dropped", + dropped, + len(items), + args.rollout_batch_size, + ) - cache_dir = sft_cache_dir(args, data_path) - if len(list(cache_dir.glob("*.pt"))) < len(items): - logger.info("SftDataManager: building cache %s for %d samples", cache_dir, len(items)) - build_sft_cache(args, items, cache_dir, pg[0]) - self.files = sorted(cache_dir.glob("*.pt")) - assert len(self.files) == len(items) + cache_dir = data_path.parent / ".sft_cache" + for item in items: + item["cache_name"], item["latent_seed"] = sft_sample_key(args, item) + self.files = [cache_dir / item["cache_name"] for item in items] + missing = [item for item in items if not (cache_dir / item["cache_name"]).exists()] + if missing: + logger.info("SftDataManager: encoding %d of %d samples into %s", len(missing), len(items), cache_dir) + build_sft_cache(args, missing, cache_dir, pg[0]) + assert all(f.exists() for f in self.files) self.train_data_dp_splitter = TrainDataDPSplitter() train_pipeline_config = load_function(args.train_pipeline_config_path)() diff --git a/miles/ray/sft_encode.py b/miles/ray/sft_encode.py index ef49019c..007c6a92 100644 --- a/miles/ray/sft_encode.py +++ b/miles/ray/sft_encode.py @@ -1,6 +1,7 @@ """Family-generic SFT cache building: encode actors driven by TrainPipelineConfig hooks.""" import logging +import os from pathlib import Path import ray @@ -40,11 +41,7 @@ def __init__(self, args): def encode(self, items: list[dict], cache_dir: str) -> int: args = self.args - done = 0 for item in items: - out_path = Path(cache_dir) / f"{item['index']:08d}.pt" - if out_path.exists(): - continue pixels = read_video_clip( item["video"], height=args.sft_height, @@ -52,9 +49,14 @@ def encode(self, items: list[dict], cache_dir: str) -> int: num_frames=args.sft_num_frames, frame_stride=args.sft_frame_stride, ) - torch.save(self.config.encode_sft_sample(self.encoder, pixels, item["prompt"]), out_path) - done += 1 - return done + generator = torch.Generator().manual_seed(item["latent_seed"]) + pair = self.config.encode_sft_sample(self.encoder, pixels, item["prompt"], generator) + out_path = Path(cache_dir) / item["cache_name"] + # Temp-then-rename so an interrupted write never leaves a loadable-looking cache entry. + tmp_path = out_path.with_name(out_path.name + ".tmp") + torch.save(pair, tmp_path) + os.replace(tmp_path, out_path) + return len(items) def build_sft_cache(args, items: list[dict], cache_dir: Path, pg) -> None: diff --git a/tests/fast/ray/test_sft_cache_dir.py b/tests/fast/ray/test_sft_cache_dir.py deleted file mode 100644 index 920a4b8e..00000000 --- a/tests/fast/ray/test_sft_cache_dir.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Cache addressing for the SFT data manager.""" - -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) - -from argparse import Namespace - -from miles.ray.sft_data_manager import sft_cache_dir - - -def _args(**overrides): - base = dict(hf_checkpoint="ckpt", sft_height=480, sft_width=832, sft_num_frames=81, sft_frame_stride=2) - base.update(overrides) - return Namespace(**base) - - -def test_cache_dir_is_stable_and_content_addressed(tmp_path): - data = tmp_path / "train.jsonl" - data.write_text('{"video": "a.mp4", "prompt": "p"}\n') - - base = sft_cache_dir(_args(), data) - assert base == sft_cache_dir(_args(), data) - assert base.parent == tmp_path / ".sft_cache" - - assert sft_cache_dir(_args(sft_height=512), data) != base - assert sft_cache_dir(_args(sft_frame_stride=1), data) != base - assert sft_cache_dir(_args(hf_checkpoint="other"), data) != base - - data.write_text('{"video": "b.mp4", "prompt": "p"}\n') - assert sft_cache_dir(_args(), data) != base diff --git a/tests/fast/ray/test_sft_sample_key.py b/tests/fast/ray/test_sft_sample_key.py new file mode 100644 index 00000000..c85d09d0 --- /dev/null +++ b/tests/fast/ray/test_sft_sample_key.py @@ -0,0 +1,57 @@ +"""Per-sample content addressing for the SFT cache.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +import os +from argparse import Namespace + +from miles.ray.sft_data_manager import sft_sample_key + + +def _args(**overrides): + base = dict(hf_checkpoint="ckpt", sft_height=480, sft_width=832, sft_num_frames=81, sft_frame_stride=2) + base.update(overrides) + return Namespace(**base) + + +def test_key_is_deterministic_with_seed(tmp_path): + video = tmp_path / "a.mp4" + video.write_bytes(b"x" * 100) + item = {"video": str(video), "prompt": "p"} + + name, seed = sft_sample_key(_args(), item) + assert (name, seed) == sft_sample_key(_args(), item) + assert name.endswith(".pt") + assert 0 <= seed < 2**63 + + +def test_key_invalidates_per_axis(tmp_path): + video = tmp_path / "a.mp4" + video.write_bytes(b"x" * 100) + item = {"video": str(video), "prompt": "p"} + base_name, base_seed = sft_sample_key(_args(), item) + + assert sft_sample_key(_args(sft_height=512), item)[0] != base_name + assert sft_sample_key(_args(sft_frame_stride=1), item)[0] != base_name + assert sft_sample_key(_args(hf_checkpoint="other"), item)[0] != base_name + assert sft_sample_key(_args(), {"video": str(video), "prompt": "q"})[0] != base_name + + video.write_bytes(b"y" * 101) + assert sft_sample_key(_args(), item)[0] != base_name + + video.write_bytes(b"x" * 100) + os.utime(video, ns=(1, 1)) + replaced_name, replaced_seed = sft_sample_key(_args(), item) + assert replaced_name != base_name + assert replaced_seed != base_seed + + +def test_key_is_per_sample(tmp_path): + a, b = tmp_path / "a.mp4", tmp_path / "b.mp4" + a.write_bytes(b"x" * 100) + b.write_bytes(b"x" * 100) + name_a, _ = sft_sample_key(_args(), {"video": str(a), "prompt": "p"}) + name_b, _ = sft_sample_key(_args(), {"video": str(b), "prompt": "p"}) + assert name_a != name_b From ce29764416d67936620b8e45acb8a93edb763324 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 4 Aug 2026 23:30:55 +0000 Subject: [PATCH 05/19] refactor(sft): plug into RolloutManager via rollout-function/convert/log hooks with lazy encoder pool --- miles/backends/fsdp_utils/actor.py | 4 +- miles/backends/fsdp_utils/loss_hub/sft.py | 4 +- miles/backends/fsdp_utils/metrics.py | 1 - miles/ray/placement_group.py | 12 +- miles/ray/rollout.py | 17 +- miles/ray/sft_data_manager.py | 118 ------------ miles/ray/sft_encode.py | 79 -------- miles/rollout/rm_hub/core.py | 17 +- miles/rollout/sft_rollout.py | 181 ++++++++++++++++++ miles/utils/arguments.py | 31 ++- scripts/run-diffusion-sft-wan22.sh | 13 +- .../backends/fsdp_utils/test_loss_hub_sft.py | 9 +- tests/fast/ray/__init__.py | 0 .../{ray => rollout}/test_sft_sample_key.py | 2 +- 14 files changed, 241 insertions(+), 247 deletions(-) delete mode 100644 miles/ray/sft_data_manager.py delete mode 100644 miles/ray/sft_encode.py create mode 100644 miles/rollout/sft_rollout.py delete mode 100644 tests/fast/ray/__init__.py rename tests/fast/{ray => rollout}/test_sft_sample_key.py (97%) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d9430378..26fc494d 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -223,7 +223,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty ) # sglang-d now supports /update_weights_from_tensor (PR #20464). - if self.args.debug_train_only: + if self.args.train_only: self.weight_updater = None elif self.args.use_lora and self.args.lora_ipc_weight_sync: self.weight_updater = DiffusionUpdateWeightFromTensorLoRAIPC(self.args, self.models) @@ -292,7 +292,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: # type @timer def update_weights(self) -> None: # type: ignore[override] - if self.args.debug_train_only or self.args.debug_rollout_only: + if self.args.train_only or self.args.debug_rollout_only: return if self.weight_updater is None: diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py index ee32b44c..72034dfa 100644 --- a/miles/backends/fsdp_utils/loss_hub/sft.py +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -94,8 +94,6 @@ def sft_loss_formula( per_pair = ((new_pred.float() - target) ** 2).mean(dim=tuple(range(1, target.ndim))) loss_sum = per_pair.sum() - bsz = len(batch) with torch.no_grad(): - metrics.emit_mean("loss", total=loss_sum, count=bsz) - metrics.emit_mean("sft_t_mean", total=prepared.timesteps.sum(), count=bsz) + metrics.emit_mean("loss", total=loss_sum, count=len(batch)) return loss_sum diff --git a/miles/backends/fsdp_utils/metrics.py b/miles/backends/fsdp_utils/metrics.py index b76cd785..66af7ffe 100644 --- a/miles/backends/fsdp_utils/metrics.py +++ b/miles/backends/fsdp_utils/metrics.py @@ -31,7 +31,6 @@ "nft_adv_mean": MetricReduce.MEAN, "nft_t_mean": MetricReduce.MEAN, "nft_num_timesteps": MetricReduce.MEAN, - "sft_t_mean": MetricReduce.MEAN, } diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 1fda2ae5..a1fcf03f 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -7,7 +7,6 @@ from .actor_group import RayTrainGroup from .rollout import RolloutManager -from .sft_data_manager import SftDataManager logger = logging.getLogger(__name__) @@ -89,7 +88,7 @@ def create_placement_groups(args): train and rollout each own a disjoint GPU pool — avoids bundle overlap / scheduling deadlock when running side-by-side. """ - if not args.colocate and not args.debug_train_only and not args.debug_rollout_only: + if not args.colocate and not args.train_only and not args.debug_rollout_only: logger.info("Creating placement groups (separate actor/rollout)...") actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node rollout_gpus = args.rollout_num_gpus @@ -110,8 +109,8 @@ def create_placement_groups(args): actor_pg_reordered_bundle_indices = all_reordered_bundle_indices actor_pg_reordered_gpu_ids = all_reordered_gpu_ids - rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.debug_train_only else [] - rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.debug_train_only else [] + rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.train_only else [] + rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.train_only else [] return { "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids), @@ -152,9 +151,8 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): - manager_cls = SftDataManager if args.loss_type == "sft_loss" else RolloutManager - logger.info("Creating rollout manager (%s, num_gpus=%s)", manager_cls.__ray_metadata__.class_name, 0) - rollout_manager = manager_cls.options( + logger.info("Creating rollout manager (num_gpus=%s)", 0) + rollout_manager = RolloutManager.options( num_cpus=1, num_gpus=0, ).remote(args, pg) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 2b49513a..abcf70df 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -17,7 +17,7 @@ expand_samples_to_train_pairs as flow_grpo_expand_samples_to_train_pairs, ) from miles.rollout.base_types import call_rollout_fn -from miles.rollout.rm_hub.core import set_reward_placement_group +from miles.rollout.rm_hub.core import set_manager_placement_group from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor from miles.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client @@ -54,9 +54,10 @@ def __init__(self, args, pg): from miles.dashboard import hooks hooks.register_rollout_manager(args) - set_reward_placement_group(pg) - logger.info("RolloutManager: starting router...") - _start_router(args) + set_manager_placement_group(pg) + if not args.train_only: + logger.info("RolloutManager: starting router...") + _start_router(args) logger.info("RolloutManager: router started, init tracking...") # TODO make args immutable init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") @@ -91,10 +92,10 @@ def __init__(self, args, pg): logger.info("RolloutManager rollout_num_gpus=%s", self.args.rollout_num_gpus) - if self.args.debug_train_only: + if self.args.train_only: self.all_rollout_engines = [] self.num_new_engines = 0 - logger.info("RolloutManager using no sglang engines (debug_train_only).") + logger.info("RolloutManager using no sglang engines (train_only).") else: num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine @@ -200,7 +201,7 @@ def generate(self, rollout_id): return [Box(ray.put(shard)) for shard in shards] def eval(self, rollout_id): - if self.args.debug_train_only: + if self.args.train_only: # if debug train only, we don't generate evaluation data return self.health_monitoring_resume() @@ -474,7 +475,7 @@ def set_train_parallel_config(self, config: dict): def init_rollout_engines(args, pg, all_rollout_engines): - if args.debug_train_only: + if args.train_only: return 0 num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) diff --git a/miles/ray/sft_data_manager.py b/miles/ray/sft_data_manager.py deleted file mode 100644 index ffb30fb6..00000000 --- a/miles/ray/sft_data_manager.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Serves auto-cached SFT samples through the RolloutManager driver interface.""" - -import hashlib -import json -import logging -from pathlib import Path - -import ray -import torch - -from miles.ray.sft_encode import build_sft_cache -from miles.utils.logging_utils import configure_logger -from miles.utils.misc import load_function -from miles.utils.ray_utils import Box -from miles.utils.train_data_utils import TrainDataDPSplitter - -logger = logging.getLogger(__name__) - - -def sft_sample_key(args, item: dict) -> tuple[str, int]: - """Content-addressed cache filename and latent-sampling seed for one (video, prompt) item.""" - stat = Path(item["video"]).stat() - digest = hashlib.sha256( - f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}|{args.sft_num_frames}s{args.sft_frame_stride}" - f"|{item['video']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() - ).digest() - return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 - - -@ray.remote -class SftDataManager: - """Encodes --sft-data-path (jsonl) into a content-addressed cache on first run, - then serves per-epoch shuffled pairs; generate() is stateless in rollout_id.""" - - def __init__(self, args, pg): - configure_logger() - self.args = args - data_path = Path(args.sft_data_path) - items = [ - {"video": row[args.sft_video_key], "prompt": row[args.sft_prompt_key]} - for row in (json.loads(line) for line in data_path.read_text().splitlines() if line.strip()) - ] - if len(items) < args.rollout_batch_size: - raise ValueError( - f"--sft-data-path holds {len(items)} samples, fewer than rollout_batch_size={args.rollout_batch_size}" - ) - dropped = len(items) % args.rollout_batch_size - if dropped: - logger.warning( - "SFT drop-last: %d of %d samples unused per epoch (rollout_batch_size=%d); " - "the per-epoch reshuffle rotates which samples are dropped", - dropped, - len(items), - args.rollout_batch_size, - ) - - cache_dir = data_path.parent / ".sft_cache" - for item in items: - item["cache_name"], item["latent_seed"] = sft_sample_key(args, item) - self.files = [cache_dir / item["cache_name"] for item in items] - missing = [item for item in items if not (cache_dir / item["cache_name"]).exists()] - if missing: - logger.info("SftDataManager: encoding %d of %d samples into %s", len(missing), len(items), cache_dir) - build_sft_cache(args, missing, cache_dir, pg[0]) - assert all(f.exists() for f in self.files) - self.train_data_dp_splitter = TrainDataDPSplitter() - - train_pipeline_config = load_function(args.train_pipeline_config_path)() - scheduler = load_function(args.model_backend_path)(train_pipeline_config).load_scheduler(args) - num_train_timesteps = int(scheduler.config.num_train_timesteps) - shift = args.diffusion_flow_shift - sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps, dtype=torch.float64) - sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) - self.scheduler_timesteps = (sigmas * num_train_timesteps).to(torch.float32) - self.scheduler_sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]).to(torch.float32) - logger.info( - "SftDataManager: %d samples, cache=%s, flow_shift=%s, num_train_timesteps=%d", - len(self.files), - cache_dir, - shift, - num_train_timesteps, - ) - - def set_train_parallel_config(self, config: dict): - self.train_parallel_config = config - - def get_num_rollout_per_epoch(self): - return len(self.files) // self.args.rollout_batch_size - - def generate(self, rollout_id): - batch_size = self.args.rollout_batch_size - epoch, slot = divmod(rollout_id, len(self.files) // batch_size) - generator = torch.Generator().manual_seed(self.args.seed + epoch) - perm = torch.randperm(len(self.files), generator=generator) - indices = perm[slot * batch_size : (slot + 1) * batch_size].tolist() - pairs = [torch.load(self.files[i], map_location="cpu") for i in indices] - data = { - "train_data": pairs, - "scheduler_timesteps": self.scheduler_timesteps, - "scheduler_sigmas": self.scheduler_sigmas, - } - shards = self.train_data_dp_splitter.split_by_dp(data, self.train_parallel_config["dp_size"]) - return [Box(ray.put(shard)) for shard in shards] - - def save(self, rollout_id): - pass - - def load(self, rollout_id=None): - pass - - def offload(self): - pass - - def onload_weights(self): - pass - - def dispose(self): - pass diff --git a/miles/ray/sft_encode.py b/miles/ray/sft_encode.py deleted file mode 100644 index 007c6a92..00000000 --- a/miles/ray/sft_encode.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Family-generic SFT cache building: encode actors driven by TrainPipelineConfig hooks.""" - -import logging -import os -from pathlib import Path - -import ray -import torch -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy - -from miles.utils.misc import load_function - -logger = logging.getLogger(__name__) - - -def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: - import torchvision - - frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") - span = (num_frames - 1) * frame_stride + 1 - if frames.shape[0] < span: - raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") - start = (frames.shape[0] - span) // 2 - frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 - - scale = max(height / frames.shape[2], width / frames.shape[3]) - new_h = max(height, round(frames.shape[2] * scale)) - new_w = max(width, round(frames.shape[3] * scale)) - frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) - top = (new_h - height) // 2 - left = (new_w - width) // 2 - return frames[:, :, top : top + height, left : left + width].permute(1, 0, 2, 3) - - -@ray.remote -class SftEncodeActor: - def __init__(self, args): - self.args = args - self.config = load_function(args.train_pipeline_config_path)() - self.encoder = self.config.load_sft_encoder(args, torch.device("cuda")) - - def encode(self, items: list[dict], cache_dir: str) -> int: - args = self.args - for item in items: - pixels = read_video_clip( - item["video"], - height=args.sft_height, - width=args.sft_width, - num_frames=args.sft_num_frames, - frame_stride=args.sft_frame_stride, - ) - generator = torch.Generator().manual_seed(item["latent_seed"]) - pair = self.config.encode_sft_sample(self.encoder, pixels, item["prompt"], generator) - out_path = Path(cache_dir) / item["cache_name"] - # Temp-then-rename so an interrupted write never leaves a loadable-looking cache entry. - tmp_path = out_path.with_name(out_path.name + ".tmp") - torch.save(pair, tmp_path) - os.replace(tmp_path, out_path) - return len(items) - - -def build_sft_cache(args, items: list[dict], cache_dir: Path, pg) -> None: - cache_dir.mkdir(parents=True, exist_ok=True) - num_workers = min(len(pg.bundle_specs), len(items)) - actors = [ - SftEncodeActor.options( - num_cpus=1, - num_gpus=1, - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_bundle_index=i, - ), - ).remote(args) - for i in range(num_workers) - ] - done = ray.get([actor.encode.remote(items[i::num_workers], str(cache_dir)) for i, actor in enumerate(actors)]) - for actor in actors: - ray.kill(actor) - logger.info("SFT cache: encoded %d new samples into %s (%d workers)", sum(done), cache_dir, num_workers) diff --git a/miles/rollout/rm_hub/core.py b/miles/rollout/rm_hub/core.py index ba42b4dc..7dad519f 100644 --- a/miles/rollout/rm_hub/core.py +++ b/miles/rollout/rm_hub/core.py @@ -8,17 +8,22 @@ import ray from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -_reward_placement_group = None +_manager_placement_group = None logger = logging.getLogger(__name__) -def set_reward_placement_group(pg) -> None: - global _reward_placement_group - _reward_placement_group = pg +def set_manager_placement_group(pg) -> None: + """Publish the manager's (pg, bundle_indices, gpu_ids) for colocated actor pools.""" + global _manager_placement_group + _manager_placement_group = pg -def get_reward_placement_group(): - return _reward_placement_group +def get_manager_placement_group(): + return _manager_placement_group + + +set_reward_placement_group = set_manager_placement_group +get_reward_placement_group = get_manager_placement_group class AsyncRewardActorPool: diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py new file mode 100644 index 00000000..190bde35 --- /dev/null +++ b/miles/rollout/sft_rollout.py @@ -0,0 +1,181 @@ +"""SFT rollout plugin: lazy content-addressed encode cache behind the standard RolloutManager flow. + +Plugged via --loss-type sft_loss (see set_default_diffusion_args): generate_rollout replaces the +sglang rollout, convert_samples_to_train_data replaces the reward/advantage conversion, and +log_rollout_data replaces the reward logging. The encoder actor pool plays the architectural role +sglang engines play in RL: the GPU data producer behind the rollout function, colocated on the +manager placement group. +""" + +import hashlib +import logging +import os +import time +from pathlib import Path + +import ray +import torch +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from miles.rollout.base_types import RolloutFnTrainOutput +from miles.rollout.rm_hub.core import get_manager_placement_group +from miles.utils import tracking_utils +from miles.utils.metric_utils import compute_rollout_step +from miles.utils.misc import load_function +from miles.utils.types import Sample + +logger = logging.getLogger(__name__) + +ENCODE_GPU_FRACTION = 0.2 + + +def sft_sample_key(args, item: dict) -> tuple[str, int]: + """Content-addressed cache filename and latent-sampling seed for one (video, prompt) item.""" + stat = Path(item["video"]).stat() + digest = hashlib.sha256( + f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}|{args.sft_num_frames}s{args.sft_frame_stride}" + f"|{item['video']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() + ).digest() + return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 + + +def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: + import torchvision + + frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") + span = (num_frames - 1) * frame_stride + 1 + if frames.shape[0] < span: + raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") + start = (frames.shape[0] - span) // 2 + frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 + + scale = max(height / frames.shape[2], width / frames.shape[3]) + new_h = max(height, round(frames.shape[2] * scale)) + new_w = max(width, round(frames.shape[3] * scale)) + frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True) + top = (new_h - height) // 2 + left = (new_w - width) // 2 + return frames[:, :, top : top + height, left : left + width].permute(1, 0, 2, 3) + + +@ray.remote +class SftEncodeActor: + def __init__(self, args): + self.args = args + self.config = load_function(args.train_pipeline_config_path)() + self.encoder = self.config.load_sft_encoder(args, torch.device("cuda")) + + def encode(self, items: list[dict], cache_dir: str) -> int: + args = self.args + for item in items: + pixels = read_video_clip( + item["video"], + height=args.sft_height, + width=args.sft_width, + num_frames=args.sft_num_frames, + frame_stride=args.sft_frame_stride, + ) + generator = torch.Generator().manual_seed(item["latent_seed"]) + pair = self.config.encode_sft_sample(self.encoder, pixels, item["prompt"], generator) + out_path = Path(cache_dir) / item["cache_name"] + # Temp-then-rename so an interrupted write never leaves a loadable-looking cache entry. + tmp_path = out_path.with_name(out_path.name + ".tmp") + torch.save(pair, tmp_path) + os.replace(tmp_path, out_path) + return len(items) + + +_encode_actors: list | None = None +_scheduler_grid: tuple[torch.Tensor, torch.Tensor] | None = None + + +def _encode_pool(args) -> list: + global _encode_actors + if _encode_actors is None: + pg, _, _ = get_manager_placement_group() + _encode_actors = [ + SftEncodeActor.options( + num_cpus=ENCODE_GPU_FRACTION, + num_gpus=ENCODE_GPU_FRACTION, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ), + ).remote(args) + for i in range(len(pg.bundle_specs)) + ] + logger.info("SFT encode pool: %d workers at %.2f GPU each", len(_encode_actors), ENCODE_GPU_FRACTION) + return _encode_actors + + +def _get_scheduler_grid(args) -> tuple[torch.Tensor, torch.Tensor]: + global _scheduler_grid + if _scheduler_grid is None: + config = load_function(args.train_pipeline_config_path)() + scheduler = load_function(args.model_backend_path)(config).load_scheduler(args) + num_train_timesteps = int(scheduler.config.num_train_timesteps) + shift = args.diffusion_flow_shift + sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps, dtype=torch.float64) + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + _scheduler_grid = ( + (sigmas * num_train_timesteps).to(torch.float32), + torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]).to(torch.float32), + ) + return _scheduler_grid + + +def generate_rollout(args, rollout_id, data_source, evaluation: bool = False) -> RolloutFnTrainOutput: + assert not evaluation, "sft_loss does not support eval rollouts" + # Deterministic per epoch and idempotent; non-divisible datasets may repeat a few + # boundary samples across an epoch wrap. + data_source.dataset.shuffle(data_source.epoch_id) + groups = data_source.get_samples(args.rollout_batch_size) + samples = [sample for group in groups for sample in group] + + cache_dir = Path(args.prompt_data).parent / ".sft_cache" + items = [] + for sample in samples: + item = {"video": sample.metadata["video"], "prompt": sample.prompt} + item["cache_name"], item["latent_seed"] = sft_sample_key(args, item) + items.append(item) + + missing = {item["cache_name"]: item for item in items if not (cache_dir / item["cache_name"]).exists()} + encode_seconds = 0.0 + if missing: + cache_dir.mkdir(parents=True, exist_ok=True) + start = time.time() + actors = _encode_pool(args) + miss_items = list(missing.values()) + shards = [miss_items[i :: len(actors)] for i in range(len(actors))] + ray.get( + [actor.encode.remote(shard, str(cache_dir)) for actor, shard in zip(actors, shards, strict=True) if shard] + ) + encode_seconds = time.time() - start + + for sample, item in zip(samples, items, strict=True): + sample.train_metadata = {"sft_pair": torch.load(cache_dir / item["cache_name"], map_location="cpu")} + sample.status = Sample.Status.COMPLETED + + metrics = { + "sft_cache_miss": len(missing), + "sft_encode_seconds": round(encode_seconds, 3), + "sft_epoch": data_source.epoch_id, + } + return RolloutFnTrainOutput(samples=groups, metrics=metrics) + + +def convert_samples_to_train_data(args, samples: list[Sample]) -> dict: + scheduler_timesteps, scheduler_sigmas = _get_scheduler_grid(args) + return { + "train_data": [sample.train_metadata["sft_pair"] for sample in samples], + "scheduler_timesteps": scheduler_timesteps, + "scheduler_sigmas": scheduler_sigmas, + } + + +def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool: + log_dict = {f"rollout/{key}": value for key, value in (rollout_extra_metrics or {}).items()} + log_dict["rollout/step"] = compute_rollout_step(args, rollout_id) + tracking_utils.log(args, log_dict, step_key="rollout/step") + logger.info("sft rollout %d: %s (%.1fs)", rollout_id, log_dict, rollout_time) + return True diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 57c5c253..4bba7872 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -690,18 +690,9 @@ def add_data_arguments(parser): ) parser.add_argument("--input-key", type=str, default="input", help="JSON dataset key") parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key") - parser.add_argument( - "--sft-data-path", - type=str, - default=None, - help=( - "SFT dataset jsonl for --loss-type sft_loss, one object per line with the video " - "path and prompt. Encoded pairs are cached next to it under .sft_cache/ and " - "rebuilt automatically whenever the data or encode settings change." - ), - ) - parser.add_argument("--sft-video-key", type=str, default="video", help="SFT jsonl video path key") - parser.add_argument("--sft-prompt-key", type=str, default="prompt", help="SFT jsonl prompt key") + # SFT (--loss-type sft_loss) reads --prompt-data/--input-key like RL; the video path + # lives in each row's metadata dict under "video". Encoded pairs are cached next to + # the jsonl under .sft_cache/, one content-addressed file per sample. parser.add_argument("--sft-height", type=int, default=None, help="SFT encode height (center crop)") parser.add_argument("--sft-width", type=int, default=None, help="SFT encode width (center crop)") parser.add_argument("--sft-num-frames", type=int, default=None, help="SFT encode frames per clip") @@ -1490,12 +1481,16 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def set_default_diffusion_args(args) -> None: if args.loss_type == "sft_loss": + if args.rollout_function_path == "miles.rollout.sglang_rollout.generate_rollout": + args.rollout_function_path = "miles.rollout.sft_rollout.generate_rollout" + if args.custom_convert_samples_to_train_data_path is None: + args.custom_convert_samples_to_train_data_path = "miles.rollout.sft_rollout.convert_samples_to_train_data" + if args.custom_rollout_log_function_path is None: + args.custom_rollout_log_function_path = "miles.rollout.sft_rollout.log_rollout_data" if args.custom_prepare_train_batch_path is None: args.custom_prepare_train_batch_path = "miles.backends.fsdp_utils.loss_hub.sft.prepare_sft_batch" if args.custom_loss_function_path is None: args.custom_loss_function_path = "miles.backends.fsdp_utils.loss_hub.sft.sft_loss_formula" - # SFT needs no sglang engines; reuse the train-only wiring end to end. - args.debug_train_only = True is_nft = args.loss_type == "nft" if is_nft: @@ -1610,8 +1605,8 @@ def miles_validate_args(args): raise ValueError("--ema-rollout-policy ema requires --ema-shadow") if args.loss_type == "sft_loss": - if args.sft_data_path is None: - raise ValueError("--loss-type sft_loss requires --sft-data-path") + if args.prompt_data is None: + raise ValueError("--loss-type sft_loss requires --prompt-data (jsonl with prompt + metadata.video)") if args.sft_height is None or args.sft_width is None or args.sft_num_frames is None: raise ValueError("--loss-type sft_loss requires --sft-height, --sft-width and --sft-num-frames") from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig @@ -1680,6 +1675,10 @@ def miles_validate_args(args): args.offload_rollout = True del args.offload + # Formal "no rollout engines" mode: skips engine/router startup, weight sync, and + # the rollout placement view. Implied by debug_train_only and by SFT. + args.train_only = args.debug_train_only or args.loss_type == "sft_loss" + if args.debug_rollout_only: if args.colocate and (not args.rollout_num_gpus): args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index 578b199a..aa9567be 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -1,8 +1,10 @@ #!/usr/bin/env bash # 4-GPU Wan2.2-T2V-A14B dual-expert LoRA SFT on a (video, prompt) jsonl dataset. -# No sglang engines: the SftDataManager encodes the dataset into .sft_cache/ -# next to the jsonl on first run (re-encoding automatically when data or encode -# settings change), then serves cached (latent, cond) pairs. +# No sglang engines: the sft_rollout plugin lazily encodes each round's cache +# misses via a colocated encoder actor pool, writing one content-addressed file +# per sample into .sft_cache/ next to the jsonl. Epoch 2+ is all cache hits. +# +# Dataset rows: {"prompt": "...", "metadata": {"video": "/abs/path.mp4"}} # # Per rollout step: 64 samples, num_steps_per_rollout=4 # -> 16 samples/optim step / 4 dp ranks = 4 samples/rank at mbs=1. @@ -13,7 +15,7 @@ export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" RUN_NAME="${RUN_NAME:-diffusion_sft_wan22_$(date +%Y%m%d_%H%M%S)}" SAVE_DIR="${ROOT_DIR}/logs/${RUN_NAME}/ckpt" -SFT_DATA_JSONL="${SFT_DATA_JSONL:?set SFT_DATA_JSONL to a jsonl with one {video, prompt} object per line}" +SFT_DATA_JSONL="${SFT_DATA_JSONL:?set SFT_DATA_JSONL to a jsonl with prompt + metadata.video per line}" WANDB_ARGS=() if [[ -n "${WANDB_API_KEY:-}" ]]; then @@ -45,7 +47,8 @@ WAN_LORA_TARGET_MODULES=( --loss-type sft_loss \ --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --sft-data-path "${SFT_DATA_JSONL}" \ + --prompt-data "${SFT_DATA_JSONL}" \ + --input-key prompt \ --sft-height 480 \ --sft-width 832 \ --sft-num-frames 81 \ diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py index d1420baa..30e36b4b 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py @@ -83,6 +83,14 @@ def test_timestep_scaling(self): prepared = prepare_sft_batch(ctx, _batch()) assert torch.allclose(prepared.timesteps_for_model, prepared.timesteps / NUM_TRAIN_TIMESTEPS) + def test_single_model_indices_cover_grid_uniformly(self): + torch.manual_seed(0) + ctx = _ctx({"transformer": nn.Identity()}) + _, _, idx = sample_grid_indices(ctx, bsz=20000) + counts = torch.bincount(idx, minlength=NUM_GRID).float() + assert counts.min() > 0 + assert ((counts / 20000) - 1 / NUM_GRID).abs().max() < 0.02 + def test_dual_expert_micro_batch_is_phase_pure(self): torch.manual_seed(0) models = {"transformer": nn.Identity(), "transformer_2": nn.Identity()} @@ -127,4 +135,3 @@ def test_unit_offset_loss(self): ) assert torch.allclose(loss, torch.tensor(float(len(batch)))) assert metrics.seen["loss"] == (float(len(batch)), len(batch)) - assert "sft_t_mean" in metrics.seen diff --git a/tests/fast/ray/__init__.py b/tests/fast/ray/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/fast/ray/test_sft_sample_key.py b/tests/fast/rollout/test_sft_sample_key.py similarity index 97% rename from tests/fast/ray/test_sft_sample_key.py rename to tests/fast/rollout/test_sft_sample_key.py index c85d09d0..ab80cb5d 100644 --- a/tests/fast/ray/test_sft_sample_key.py +++ b/tests/fast/rollout/test_sft_sample_key.py @@ -7,7 +7,7 @@ import os from argparse import Namespace -from miles.ray.sft_data_manager import sft_sample_key +from miles.rollout.sft_rollout import sft_sample_key def _args(**overrides): From 2fb0a6c94ba76b21e4e2c0d8b700077bbf9702d4 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 00:15:49 +0000 Subject: [PATCH 06/19] feat(sft): image datasets for Wan (single-frame media branch) --- miles/rollout/sft_rollout.py | 43 +++++++++++++++-------- miles/utils/arguments.py | 7 ++-- tests/fast/rollout/test_sft_read_media.py | 40 +++++++++++++++++++++ tests/fast/rollout/test_sft_sample_key.py | 10 +++--- 4 files changed, 78 insertions(+), 22 deletions(-) create mode 100644 tests/fast/rollout/test_sft_read_media.py diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index 190bde35..d0b6c6c7 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -30,24 +30,36 @@ def sft_sample_key(args, item: dict) -> tuple[str, int]: - """Content-addressed cache filename and latent-sampling seed for one (video, prompt) item.""" - stat = Path(item["video"]).stat() + """Content-addressed cache filename and latent-sampling seed for one (media, prompt) item.""" + stat = Path(item["media"]).stat() digest = hashlib.sha256( f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}|{args.sft_num_frames}s{args.sft_frame_stride}" - f"|{item['video']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() + f"|{item['media']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() ).digest() return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 -def read_video_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: - import torchvision +IMAGE_EXTENSIONS = {".bmp", ".jpeg", ".jpg", ".png", ".webp"} - frames, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") - span = (num_frames - 1) * frame_stride + 1 - if frames.shape[0] < span: - raise ValueError(f"{path} has {frames.shape[0]} frames, need {span}") - start = (frames.shape[0] - span) // 2 - frames = frames[start : start + span : frame_stride].float() / 127.5 - 1.0 + +def read_media_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: + if Path(path).suffix.lower() in IMAGE_EXTENSIONS: + if num_frames != 1: + raise ValueError(f"{path} is an image, which requires --sft-num-frames 1") + import numpy as np + from PIL import Image + + frames = torch.from_numpy(np.asarray(Image.open(path).convert("RGB"))).permute(2, 0, 1)[None].float() + else: + import torchvision + + video, _, _ = torchvision.io.read_video(path, pts_unit="sec", output_format="TCHW") + span = (num_frames - 1) * frame_stride + 1 + if video.shape[0] < span: + raise ValueError(f"{path} has {video.shape[0]} frames, need {span}") + start = (video.shape[0] - span) // 2 + frames = video[start : start + span : frame_stride].float() + frames = frames / 127.5 - 1.0 scale = max(height / frames.shape[2], width / frames.shape[3]) new_h = max(height, round(frames.shape[2] * scale)) @@ -68,8 +80,8 @@ def __init__(self, args): def encode(self, items: list[dict], cache_dir: str) -> int: args = self.args for item in items: - pixels = read_video_clip( - item["video"], + pixels = read_media_clip( + item["media"], height=args.sft_height, width=args.sft_width, num_frames=args.sft_num_frames, @@ -135,7 +147,10 @@ def generate_rollout(args, rollout_id, data_source, evaluation: bool = False) -> cache_dir = Path(args.prompt_data).parent / ".sft_cache" items = [] for sample in samples: - item = {"video": sample.metadata["video"], "prompt": sample.prompt} + media = sample.metadata.get("video") or sample.metadata.get("image") + if media is None: + raise ValueError(f"sample {sample.index} metadata has neither 'video' nor 'image': {sample.metadata}") + item = {"media": media, "prompt": sample.prompt} item["cache_name"], item["latent_seed"] = sft_sample_key(args, item) items.append(item) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4bba7872..82dca659 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -690,9 +690,10 @@ def add_data_arguments(parser): ) parser.add_argument("--input-key", type=str, default="input", help="JSON dataset key") parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key") - # SFT (--loss-type sft_loss) reads --prompt-data/--input-key like RL; the video path - # lives in each row's metadata dict under "video". Encoded pairs are cached next to - # the jsonl under .sft_cache/, one content-addressed file per sample. + # SFT (--loss-type sft_loss) reads --prompt-data/--input-key like RL; the media path + # lives in each row's metadata dict under "video" or "image" (images train as single + # frames and require --sft-num-frames 1). Encoded pairs are cached next to the jsonl + # under .sft_cache/, one content-addressed file per sample. parser.add_argument("--sft-height", type=int, default=None, help="SFT encode height (center crop)") parser.add_argument("--sft-width", type=int, default=None, help="SFT encode width (center crop)") parser.add_argument("--sft-num-frames", type=int, default=None, help="SFT encode frames per clip") diff --git a/tests/fast/rollout/test_sft_read_media.py b/tests/fast/rollout/test_sft_read_media.py new file mode 100644 index 00000000..72608aad --- /dev/null +++ b/tests/fast/rollout/test_sft_read_media.py @@ -0,0 +1,40 @@ +"""Image branch of the SFT media reader.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +import numpy as np +import pytest +import torch +from PIL import Image + +from miles.rollout.sft_rollout import read_media_clip + + +def _write_png(path, height, width): + rng = np.random.default_rng(0) + Image.fromarray(rng.integers(0, 256, (height, width, 3), dtype=np.uint8)).save(path) + + +def test_image_reads_as_single_frame(tmp_path): + path = tmp_path / "a.png" + _write_png(path, 120, 200) + clip = read_media_clip(str(path), height=64, width=64, num_frames=1, frame_stride=1) + assert clip.shape == (3, 1, 64, 64) + assert clip.min() >= -1.0 and clip.max() <= 1.0 + + +def test_image_rejects_multi_frame(tmp_path): + path = tmp_path / "a.png" + _write_png(path, 64, 64) + with pytest.raises(ValueError, match="requires --sft-num-frames 1"): + read_media_clip(str(path), height=64, width=64, num_frames=5, frame_stride=1) + + +def test_image_no_resize_when_exact(tmp_path): + path = tmp_path / "a.png" + _write_png(path, 64, 64) + clip = read_media_clip(str(path), height=64, width=64, num_frames=1, frame_stride=1) + original = torch.from_numpy(np.asarray(Image.open(path))).permute(2, 0, 1).float() / 127.5 - 1.0 + assert torch.allclose(clip[:, 0], original) diff --git a/tests/fast/rollout/test_sft_sample_key.py b/tests/fast/rollout/test_sft_sample_key.py index ab80cb5d..3cc69f01 100644 --- a/tests/fast/rollout/test_sft_sample_key.py +++ b/tests/fast/rollout/test_sft_sample_key.py @@ -19,7 +19,7 @@ def _args(**overrides): def test_key_is_deterministic_with_seed(tmp_path): video = tmp_path / "a.mp4" video.write_bytes(b"x" * 100) - item = {"video": str(video), "prompt": "p"} + item = {"media": str(video), "prompt": "p"} name, seed = sft_sample_key(_args(), item) assert (name, seed) == sft_sample_key(_args(), item) @@ -30,13 +30,13 @@ def test_key_is_deterministic_with_seed(tmp_path): def test_key_invalidates_per_axis(tmp_path): video = tmp_path / "a.mp4" video.write_bytes(b"x" * 100) - item = {"video": str(video), "prompt": "p"} + item = {"media": str(video), "prompt": "p"} base_name, base_seed = sft_sample_key(_args(), item) assert sft_sample_key(_args(sft_height=512), item)[0] != base_name assert sft_sample_key(_args(sft_frame_stride=1), item)[0] != base_name assert sft_sample_key(_args(hf_checkpoint="other"), item)[0] != base_name - assert sft_sample_key(_args(), {"video": str(video), "prompt": "q"})[0] != base_name + assert sft_sample_key(_args(), {"media": str(video), "prompt": "q"})[0] != base_name video.write_bytes(b"y" * 101) assert sft_sample_key(_args(), item)[0] != base_name @@ -52,6 +52,6 @@ def test_key_is_per_sample(tmp_path): a, b = tmp_path / "a.mp4", tmp_path / "b.mp4" a.write_bytes(b"x" * 100) b.write_bytes(b"x" * 100) - name_a, _ = sft_sample_key(_args(), {"video": str(a), "prompt": "p"}) - name_b, _ = sft_sample_key(_args(), {"video": str(b), "prompt": "p"}) + name_a, _ = sft_sample_key(_args(), {"media": str(a), "prompt": "p"}) + name_b, _ = sft_sample_key(_args(), {"media": str(b), "prompt": "p"}) assert name_a != name_b From e5d7afc4b0f22a644dcdbed965ef197b7981dc07 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 00:18:57 +0000 Subject: [PATCH 07/19] refactor(sft): pick dual-expert via uniform anchor index --- miles/backends/fsdp_utils/loss_hub/sft.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py index 72034dfa..581605da 100644 --- a/miles/backends/fsdp_utils/loss_hub/sft.py +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -17,12 +17,12 @@ def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Mo component_name, model = next(iter(ctx.models.items())) return component_name, model, torch.randint(num_grid, (bsz,)) + # A uniform anchor index picks each expert with probability equal to its share of the + # grid, keeping the marginal over indices uniform while the micro-batch stays single-expert. num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) config = ctx.train_pipeline_config components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps] - names = sorted(set(components)) - counts = torch.tensor([components.count(name) for name in names], dtype=torch.float32) - component_name = names[int(torch.multinomial(counts, 1))] + component_name = components[int(torch.randint(num_grid, (1,)))] pool = torch.tensor([i for i, name in enumerate(components) if name == component_name]) return component_name, ctx.models[component_name], pool[torch.randint(len(pool), (bsz,))] From 0f026e89840e6e2f13deb6e6e0c624b113289ba8 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 02:51:01 +0000 Subject: [PATCH 08/19] refactor(sft): seat encoder pool via the rollout placement view --- miles/ray/placement_group.py | 5 +++-- miles/rollout/sft_rollout.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index a1fcf03f..f972356c 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -109,8 +109,9 @@ def create_placement_groups(args): actor_pg_reordered_bundle_indices = all_reordered_bundle_indices actor_pg_reordered_gpu_ids = all_reordered_gpu_ids - rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.train_only else [] - rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.train_only else [] + # SFT keeps the rollout seats: its encoder pool is the rollout-side producer. + rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.debug_train_only else [] + rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.debug_train_only else [] return { "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids), diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index d0b6c6c7..6b5e1d72 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -104,7 +104,8 @@ def encode(self, items: list[dict], cache_dir: str) -> int: def _encode_pool(args) -> list: global _encode_actors if _encode_actors is None: - pg, _, _ = get_manager_placement_group() + # Encode is SFT's rollout: the pool takes the rollout placement seats sglang engines use in RL. + pg, bundle_indices, _ = get_manager_placement_group() _encode_actors = [ SftEncodeActor.options( num_cpus=ENCODE_GPU_FRACTION, @@ -114,7 +115,7 @@ def _encode_pool(args) -> list: placement_group_bundle_index=i, ), ).remote(args) - for i in range(len(pg.bundle_specs)) + for i in bundle_indices ] logger.info("SFT encode pool: %d workers at %.2f GPU each", len(_encode_actors), ENCODE_GPU_FRACTION) return _encode_actors From abcc73a817cc2d430accaea92ff7e7c7b95faaee Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 03:51:38 +0000 Subject: [PATCH 09/19] chore(sft): encoder pool inherits the engine's 0.3 GPU share --- miles/rollout/sft_rollout.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index 6b5e1d72..098f90db 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -26,7 +26,7 @@ logger = logging.getLogger(__name__) -ENCODE_GPU_FRACTION = 0.2 +ENCODE_GPU_FRACTION = 0.3 def sft_sample_key(args, item: dict) -> tuple[str, int]: From 34f1cb2bd8417a664a9016d3ff6a1b30f0198189 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 23:32:19 +0000 Subject: [PATCH 10/19] refactor(sft): reuse --diffusion-height/width/output-num-frames; add --fsdp-flow-shift Review feedback (PR #90): the media geometry args describe the generated media regardless of whether it comes from preprocessing or the rollout engine, so SFT reuses --diffusion-height/--diffusion-width/ --diffusion-output-num-frames instead of its own --sft-* trio. The SFT training sigma grid is regenerated on the training side, so its shift is now --fsdp-flow-shift (fsdp_* namespace), leaving --diffusion-flow-shift as the rollout-engine launch parameter. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/configs/wan2_2.py | 4 +-- miles/rollout/sft_rollout.py | 13 ++++----- miles/utils/arguments.py | 30 ++++++++++++--------- scripts/run-diffusion-sft-wan22.sh | 8 +++--- tests/fast/rollout/test_sft_read_media.py | 2 +- tests/fast/rollout/test_sft_sample_key.py | 10 +++++-- 6 files changed, 40 insertions(+), 27 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 2fad8f52..d6462fb0 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -69,8 +69,8 @@ def cfg_combine( @classmethod def validate_args(cls, args) -> None: - if args.loss_type == "sft_loss" and (args.sft_num_frames - 1) % 4 != 0: - raise ValueError("--sft-num-frames must be 4k+1 for the Wan VAE temporal stride") + if args.loss_type == "sft_loss" and (args.diffusion_output_num_frames - 1) % 4 != 0: + raise ValueError("--diffusion-output-num-frames must be 4k+1 for the Wan VAE temporal stride") def load_sft_encoder(self, args, device: torch.device): from diffusers import AutoencoderKLWan diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index 098f90db..81e97c5d 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -33,7 +33,8 @@ def sft_sample_key(args, item: dict) -> tuple[str, int]: """Content-addressed cache filename and latent-sampling seed for one (media, prompt) item.""" stat = Path(item["media"]).stat() digest = hashlib.sha256( - f"{args.hf_checkpoint}|{args.sft_height}x{args.sft_width}|{args.sft_num_frames}s{args.sft_frame_stride}" + f"{args.hf_checkpoint}|{args.diffusion_height}x{args.diffusion_width}" + f"|{args.diffusion_output_num_frames}s{args.sft_frame_stride}" f"|{item['media']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() ).digest() return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 @@ -45,7 +46,7 @@ def sft_sample_key(args, item: dict) -> tuple[str, int]: def read_media_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor: if Path(path).suffix.lower() in IMAGE_EXTENSIONS: if num_frames != 1: - raise ValueError(f"{path} is an image, which requires --sft-num-frames 1") + raise ValueError(f"{path} is an image, which requires --diffusion-output-num-frames 1") import numpy as np from PIL import Image @@ -82,9 +83,9 @@ def encode(self, items: list[dict], cache_dir: str) -> int: for item in items: pixels = read_media_clip( item["media"], - height=args.sft_height, - width=args.sft_width, - num_frames=args.sft_num_frames, + height=args.diffusion_height, + width=args.diffusion_width, + num_frames=args.diffusion_output_num_frames, frame_stride=args.sft_frame_stride, ) generator = torch.Generator().manual_seed(item["latent_seed"]) @@ -127,7 +128,7 @@ def _get_scheduler_grid(args) -> tuple[torch.Tensor, torch.Tensor]: config = load_function(args.train_pipeline_config_path)() scheduler = load_function(args.model_backend_path)(config).load_scheduler(args) num_train_timesteps = int(scheduler.config.num_train_timesteps) - shift = args.diffusion_flow_shift + shift = args.fsdp_flow_shift sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps, dtype=torch.float64) sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) _scheduler_grid = ( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 82dca659..cd11cbcc 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -189,6 +189,16 @@ def add_train_arguments(parser): "of bf16 add-non-associativity noise across ranks." ), ) + parser.add_argument( + "--fsdp-flow-shift", + type=float, + default=None, + help=( + "Flow-matching shift for the training-side sigma grid, regenerated by " + "the trainer when no rollout engine supplies scheduler meta (SFT). " + "Distinct from --diffusion-flow-shift, which configures the rollout engine." + ), + ) parser.add_argument( "--diffusion-forward-dtype", type=str, @@ -326,7 +336,7 @@ def add_rollout_arguments(parser): "--diffusion-output-num-frames", type=int, default=1, - help="Requested decoded video frame count for diffusion video rollout. Wan2.2 MVP uses 1.", + help="Frame count of the trained media: decoded rollout frames, SFT encode frames per clip. Wan2.2 MVP uses 1.", ) parser.add_argument( "--diffusion-guidance-scale", @@ -350,13 +360,13 @@ def add_rollout_arguments(parser): "--diffusion-height", type=int, default=512, - help="Output image height for diffusion rollout.", + help="Height of the trained media: rollout output height, SFT encode center-crop height.", ) parser.add_argument( "--diffusion-width", type=int, default=512, - help="Output image width for diffusion rollout.", + help="Width of the trained media: rollout output width, SFT encode center-crop width.", ) parser.add_argument( "--diffusion-negative-prompt", @@ -692,11 +702,9 @@ def add_data_arguments(parser): parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key") # SFT (--loss-type sft_loss) reads --prompt-data/--input-key like RL; the media path # lives in each row's metadata dict under "video" or "image" (images train as single - # frames and require --sft-num-frames 1). Encoded pairs are cached next to the jsonl - # under .sft_cache/, one content-addressed file per sample. - parser.add_argument("--sft-height", type=int, default=None, help="SFT encode height (center crop)") - parser.add_argument("--sft-width", type=int, default=None, help="SFT encode width (center crop)") - parser.add_argument("--sft-num-frames", type=int, default=None, help="SFT encode frames per clip") + # frames and require --diffusion-output-num-frames 1). Encode geometry comes from + # --diffusion-height/--diffusion-width/--diffusion-output-num-frames. Encoded pairs are + # cached next to the jsonl under .sft_cache/, one content-addressed file per sample. parser.add_argument("--sft-frame-stride", type=int, default=1, help="SFT encode temporal stride") parser.add_argument( @@ -1608,8 +1616,6 @@ def miles_validate_args(args): if args.loss_type == "sft_loss": if args.prompt_data is None: raise ValueError("--loss-type sft_loss requires --prompt-data (jsonl with prompt + metadata.video)") - if args.sft_height is None or args.sft_width is None or args.sft_num_frames is None: - raise ValueError("--loss-type sft_loss requires --sft-height, --sft-width and --sft-num-frames") from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig from miles.utils.misc import load_function @@ -1619,8 +1625,8 @@ def miles_validate_args(args): f"--loss-type sft_loss is not supported for {sft_cfg_cls.__name__}: " "it does not implement load_sft_encoder/encode_sft_sample" ) - if args.diffusion_flow_shift is None: - raise ValueError("--loss-type sft_loss requires --diffusion-flow-shift for the training sigma grid") + if args.fsdp_flow_shift is None: + raise ValueError("--loss-type sft_loss requires --fsdp-flow-shift for the training sigma grid") if args.n_samples_per_prompt != 1: raise ValueError("--loss-type sft_loss requires --n-samples-per-prompt 1") if args.eval_interval is not None: diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index aa9567be..fa151d87 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -49,9 +49,9 @@ WAN_LORA_TARGET_MODULES=( --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt-data "${SFT_DATA_JSONL}" \ --input-key prompt \ - --sft-height 480 \ - --sft-width 832 \ - --sft-num-frames 81 \ + --diffusion-height 480 \ + --diffusion-width 832 \ + --diffusion-output-num-frames 81 \ --sft-frame-stride 2 \ --rollout-batch-size 64 \ --num-epoch 3 \ @@ -71,7 +71,7 @@ WAN_LORA_TARGET_MODULES=( --fsdp-master-dtype fp32 \ --fsdp-reduce-dtype fp32 \ --diffusion-forward-dtype bf16 \ - --diffusion-flow-shift 3.0 \ + --fsdp-flow-shift 3.0 \ --save "${SAVE_DIR}" \ --save-interval 20 \ "${RESUME_ARGS[@]}" \ diff --git a/tests/fast/rollout/test_sft_read_media.py b/tests/fast/rollout/test_sft_read_media.py index 72608aad..b1503c54 100644 --- a/tests/fast/rollout/test_sft_read_media.py +++ b/tests/fast/rollout/test_sft_read_media.py @@ -28,7 +28,7 @@ def test_image_reads_as_single_frame(tmp_path): def test_image_rejects_multi_frame(tmp_path): path = tmp_path / "a.png" _write_png(path, 64, 64) - with pytest.raises(ValueError, match="requires --sft-num-frames 1"): + with pytest.raises(ValueError, match="requires --diffusion-output-num-frames 1"): read_media_clip(str(path), height=64, width=64, num_frames=5, frame_stride=1) diff --git a/tests/fast/rollout/test_sft_sample_key.py b/tests/fast/rollout/test_sft_sample_key.py index 3cc69f01..91e030a3 100644 --- a/tests/fast/rollout/test_sft_sample_key.py +++ b/tests/fast/rollout/test_sft_sample_key.py @@ -11,7 +11,13 @@ def _args(**overrides): - base = dict(hf_checkpoint="ckpt", sft_height=480, sft_width=832, sft_num_frames=81, sft_frame_stride=2) + base = dict( + hf_checkpoint="ckpt", + diffusion_height=480, + diffusion_width=832, + diffusion_output_num_frames=81, + sft_frame_stride=2, + ) base.update(overrides) return Namespace(**base) @@ -33,7 +39,7 @@ def test_key_invalidates_per_axis(tmp_path): item = {"media": str(video), "prompt": "p"} base_name, base_seed = sft_sample_key(_args(), item) - assert sft_sample_key(_args(sft_height=512), item)[0] != base_name + assert sft_sample_key(_args(diffusion_height=512), item)[0] != base_name assert sft_sample_key(_args(sft_frame_stride=1), item)[0] != base_name assert sft_sample_key(_args(hf_checkpoint="other"), item)[0] != base_name assert sft_sample_key(_args(), {"media": str(video), "prompt": "q"})[0] != base_name From d974b2746d75b627be5d61d746cb60c5b2646935 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 5 Aug 2026 23:50:31 +0000 Subject: [PATCH 11/19] refactor(sft): validate explicit plugin paths instead of auto-wiring defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback (PR #90): follow Miles LLM convention — the sample script passes the five SFT plugin paths explicitly and the framework validates the combination instead of silently rewriting args. Misconfigured runs now fail at argument validation with the exact flag to set, rather than deep in the RL data path. Co-Authored-By: Claude Fable 5 --- miles/utils/arguments.py | 28 ++++++++++++++++------------ scripts/run-diffusion-sft-wan22.sh | 5 +++++ 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index cd11cbcc..18d5a5f4 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1489,18 +1489,6 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def set_default_diffusion_args(args) -> None: - if args.loss_type == "sft_loss": - if args.rollout_function_path == "miles.rollout.sglang_rollout.generate_rollout": - args.rollout_function_path = "miles.rollout.sft_rollout.generate_rollout" - if args.custom_convert_samples_to_train_data_path is None: - args.custom_convert_samples_to_train_data_path = "miles.rollout.sft_rollout.convert_samples_to_train_data" - if args.custom_rollout_log_function_path is None: - args.custom_rollout_log_function_path = "miles.rollout.sft_rollout.log_rollout_data" - if args.custom_prepare_train_batch_path is None: - args.custom_prepare_train_batch_path = "miles.backends.fsdp_utils.loss_hub.sft.prepare_sft_batch" - if args.custom_loss_function_path is None: - args.custom_loss_function_path = "miles.backends.fsdp_utils.loss_hub.sft.sft_loss_formula" - is_nft = args.loss_type == "nft" if is_nft: if args.custom_expand_samples_to_train_pairs_path is None: @@ -1614,6 +1602,22 @@ def miles_validate_args(args): raise ValueError("--ema-rollout-policy ema requires --ema-shadow") if args.loss_type == "sft_loss": + if args.rollout_function_path == "miles.rollout.sglang_rollout.generate_rollout": + raise ValueError( + "--loss-type sft_loss does not run rollout engines; pass " + "--rollout-function-path miles.rollout.sft_rollout.generate_rollout (or your own)" + ) + for name in ( + "custom_convert_samples_to_train_data_path", + "custom_rollout_log_function_path", + "custom_prepare_train_batch_path", + "custom_loss_function_path", + ): + if getattr(args, name) is None: + raise ValueError( + f"--loss-type sft_loss requires --{name.replace('_', '-')}; " + "see scripts/run-diffusion-sft-wan22.sh" + ) if args.prompt_data is None: raise ValueError("--loss-type sft_loss requires --prompt-data (jsonl with prompt + metadata.video)") from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index fa151d87..a4ec7642 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -45,6 +45,11 @@ WAN_LORA_TARGET_MODULES=( "${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ --train-backend fsdp \ --loss-type sft_loss \ + --rollout-function-path miles.rollout.sft_rollout.generate_rollout \ + --custom-convert-samples-to-train-data-path miles.rollout.sft_rollout.convert_samples_to_train_data \ + --custom-rollout-log-function-path miles.rollout.sft_rollout.log_rollout_data \ + --custom-prepare-train-batch-path miles.backends.fsdp_utils.loss_hub.sft.prepare_sft_batch \ + --custom-loss-function-path miles.backends.fsdp_utils.loss_hub.sft.sft_loss_formula \ --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt-data "${SFT_DATA_JSONL}" \ From f1ceb9d04396dbc12a66af7fd0c15868832517cd Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 00:10:49 +0000 Subject: [PATCH 12/19] refactor: make --train-only an explicit flag; --debug-train-only becomes its alias Review feedback (PR #90): no silent inference of train_only from loss_type. --train-only is now a real user-facing flag (argparse alias keeps --debug-train-only working) and sft_loss validates it is set instead of setting it. The rollout placement view keeps its seats unconditionally: engine startup is gated by args.train_only, and rollout-side actor pools (the SFT encoder pool) seat there. The debug_rollout_only/train_only exclusion assert now runs before the debug_rollout_only reconfiguration so the combo fails with the intended message. Co-Authored-By: Claude Fable 5 --- miles/ray/placement_group.py | 14 +++++--------- miles/utils/arguments.py | 21 +++++++++++---------- scripts/run-diffusion-sft-wan22.sh | 1 + 3 files changed, 17 insertions(+), 19 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index f972356c..c56e9dad 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -82,7 +82,7 @@ def create_placement_groups(args): """Create placement groups for actor and rollout engines. Two topologies: - - Colocate (or --debug-{train,rollout}-only): one combined placement + - Colocate (or --train-only / --debug-rollout-only): one combined placement group; both roles see the same bundle list. - Disaggregate (the else branch): two separate placement groups so train and rollout each own a disjoint GPU pool — avoids bundle @@ -107,15 +107,11 @@ def create_placement_groups(args): logger.info(f"Creating placement group with {num_gpus} GPUs...") pg, all_reordered_bundle_indices, all_reordered_gpu_ids = _create_placement_group(num_gpus) - actor_pg_reordered_bundle_indices = all_reordered_bundle_indices - actor_pg_reordered_gpu_ids = all_reordered_gpu_ids - # SFT keeps the rollout seats: its encoder pool is the rollout-side producer. - rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.debug_train_only else [] - rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.debug_train_only else [] - + # The rollout view keeps its seats under train_only: engine startup is gated by + # args.train_only, and rollout-side actor pools (e.g. the SFT encoder pool) seat here. return { - "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids), - "rollout": (pg, rollout_pg_reordered_bundle_indices, rollout_pg_reordered_gpu_ids), + "actor": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids), + "rollout": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids), } diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 18d5a5f4..4831d0a4 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1101,12 +1101,15 @@ def add_debug_arguments(parser): ), ) parser.add_argument( + "--train-only", "--debug-train-only", + dest="train_only", action="store_true", default=False, help=( - "Whether to only run the training without sglang servers. " - "This is useful for debugging the rollout generation function." + "Run without rollout engines: skips engine/router startup, weight sync, " + "and eval. Used by SFT and for debugging the training side. " + "--debug-train-only is a legacy alias." ), ) parser.add_argument( @@ -1602,6 +1605,8 @@ def miles_validate_args(args): raise ValueError("--ema-rollout-policy ema requires --ema-shadow") if args.loss_type == "sft_loss": + if not args.train_only: + raise ValueError("--loss-type sft_loss runs no rollout engines; pass --train-only") if args.rollout_function_path == "miles.rollout.sglang_rollout.generate_rollout": raise ValueError( "--loss-type sft_loss does not run rollout engines; pass " @@ -1679,16 +1684,16 @@ def miles_validate_args(args): f"load_debug_rollout_data {args.load_debug_rollout_data} is set, " "will not instantiate sglang servers and will only run the training process." ) - args.debug_train_only = True + args.train_only = True if args.offload: args.offload_train = True args.offload_rollout = True del args.offload - # Formal "no rollout engines" mode: skips engine/router startup, weight sync, and - # the rollout placement view. Implied by debug_train_only and by SFT. - args.train_only = args.debug_train_only or args.loss_type == "sft_loss" + assert not (args.debug_rollout_only and args.train_only), ( + "debug_rollout_only and train_only cannot be set at the same time, please set only one of them." + ) if args.debug_rollout_only: if args.colocate and (not args.rollout_num_gpus): @@ -1702,10 +1707,6 @@ def miles_validate_args(args): logger.warning("Force train_memory_margin_bytes=0 since debug_rollout_only does not support it") args.train_memory_margin_bytes = 0 - assert not (args.debug_rollout_only and args.debug_train_only), ( - "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." - ) - if getattr(args, "diffusion_model", None): from miles.backends.fsdp_utils.arguments import validate_hybrid_shard_args, validate_sp_args diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index a4ec7642..4543eabf 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -45,6 +45,7 @@ WAN_LORA_TARGET_MODULES=( "${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ --train-backend fsdp \ --loss-type sft_loss \ + --train-only \ --rollout-function-path miles.rollout.sft_rollout.generate_rollout \ --custom-convert-samples-to-train-data-path miles.rollout.sft_rollout.convert_samples_to_train_data \ --custom-rollout-log-function-path miles.rollout.sft_rollout.log_rollout_data \ From 795e5d7d616fd6cb65d63a8d539b41081da3b406 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 00:48:59 +0000 Subject: [PATCH 13/19] refactor(sft): move frozen-encoder logic to encoder_hub; TPC back to training-only Review feedback (PR #90): TrainPipelineConfig drives the training backend only. Encoder loading/encoding now lives in miles/rollout/encoder_hub (stacked base PR #96), dispatched by args.diffusion_model_family; the Wan-specific 4k+1 frame constraint validates there too. Encoders load from the explicit --sft-encoder-checkpoint, which also replaces hf_checkpoint in the per-sample cache key since it is what determines cache content. Co-Authored-By: Claude Fable 5 --- .../configs/train_pipeline_config.py | 11 ---- miles/backends/fsdp_utils/configs/wan2_2.py | 54 ------------------- miles/rollout/sft_rollout.py | 10 ++-- miles/utils/arguments.py | 19 ++++--- scripts/run-diffusion-sft-wan22.sh | 1 + tests/fast/rollout/test_sft_sample_key.py | 4 +- 6 files changed, 21 insertions(+), 78 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 9b5a20e5..c392ea57 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -189,14 +189,3 @@ def cfg_combine( def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: """Postprocess the model after FSDP wrap + weight materialization (default: no-op).""" return None - - def load_sft_encoder(self, args, device: torch.device): - """Load this family's frozen encode components (tokenizer/text encoder/VAE) for SFT caching.""" - raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") - - def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict: - """Encode one (pixels [C,T,H,W] in [-1,1], prompt) into the train-pair dict for prepare_sft_batch. - - ``generator`` seeds any stochastic encode step (e.g. VAE posterior sampling) so a cache - entry's content is a deterministic function of its cache key.""" - raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding") diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index d6462fb0..7e227cd2 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -66,57 +66,3 @@ def cfg_combine( ) -> torch.Tensor: scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) - - @classmethod - def validate_args(cls, args) -> None: - if args.loss_type == "sft_loss" and (args.diffusion_output_num_frames - 1) % 4 != 0: - raise ValueError("--diffusion-output-num-frames must be 4k+1 for the Wan VAE temporal stride") - - def load_sft_encoder(self, args, device: torch.device): - from diffusers import AutoencoderKLWan - from transformers import AutoTokenizer, UMT5EncoderModel - - tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, subfolder="tokenizer") - text_encoder = UMT5EncoderModel.from_pretrained( - args.hf_checkpoint, subfolder="text_encoder", torch_dtype=torch.bfloat16 - ).to(device) - vae = AutoencoderKLWan.from_pretrained(args.hf_checkpoint, subfolder="vae", torch_dtype=torch.float32).to( - device - ) - view = (1, vae.config.z_dim, 1, 1, 1) - return { - "device": device, - "tokenizer": tokenizer, - "text_encoder": text_encoder, - "vae": vae, - "latents_mean": torch.tensor(vae.config.latents_mean).view(view).to(device), - "latents_std": torch.tensor(vae.config.latents_std).view(view).to(device), - } - - @torch.no_grad() - def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict: - from diffusers.pipelines.wan.pipeline_wan import prompt_clean - - device = encoder["device"] - latent = encoder["vae"].encode(pixels.unsqueeze(0).to(device, torch.float32)).latent_dist.sample(generator) - latent = (latent - encoder["latents_mean"]) / encoder["latents_std"] - - inputs = encoder["tokenizer"]( - [prompt_clean(prompt)], - padding="max_length", - max_length=512, - truncation=True, - add_special_tokens=True, - return_attention_mask=True, - return_tensors="pt", - ) - embeds = encoder["text_encoder"]( - inputs.input_ids.to(device), inputs.attention_mask.to(device) - ).last_hidden_state - embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 - - return { - "latent": latent[0].to(torch.float16).cpu(), - "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, - "prompt": prompt, - } diff --git a/miles/rollout/sft_rollout.py b/miles/rollout/sft_rollout.py index 81e97c5d..96e6bb16 100644 --- a/miles/rollout/sft_rollout.py +++ b/miles/rollout/sft_rollout.py @@ -33,7 +33,7 @@ def sft_sample_key(args, item: dict) -> tuple[str, int]: """Content-addressed cache filename and latent-sampling seed for one (media, prompt) item.""" stat = Path(item["media"]).stat() digest = hashlib.sha256( - f"{args.hf_checkpoint}|{args.diffusion_height}x{args.diffusion_width}" + f"{args.sft_encoder_checkpoint}|{args.diffusion_height}x{args.diffusion_width}" f"|{args.diffusion_output_num_frames}s{args.sft_frame_stride}" f"|{item['media']}|{stat.st_size}|{stat.st_mtime_ns}|{item['prompt']}".encode() ).digest() @@ -74,9 +74,11 @@ def read_media_clip(path: str, *, height: int, width: int, num_frames: int, fram @ray.remote class SftEncodeActor: def __init__(self, args): + from miles.rollout.encoder_hub import get_encoder + self.args = args - self.config = load_function(args.train_pipeline_config_path)() - self.encoder = self.config.load_sft_encoder(args, torch.device("cuda")) + self.encoder_module = get_encoder(args.diffusion_model_family) + self.encoder = self.encoder_module.load_encoder(args, torch.device("cuda")) def encode(self, items: list[dict], cache_dir: str) -> int: args = self.args @@ -89,7 +91,7 @@ def encode(self, items: list[dict], cache_dir: str) -> int: frame_stride=args.sft_frame_stride, ) generator = torch.Generator().manual_seed(item["latent_seed"]) - pair = self.config.encode_sft_sample(self.encoder, pixels, item["prompt"], generator) + pair = self.encoder_module.encode_sample(self.encoder, pixels, item["prompt"], generator) out_path = Path(cache_dir) / item["cache_name"] # Temp-then-rename so an interrupted write never leaves a loadable-looking cache entry. tmp_path = out_path.with_name(out_path.name + ".tmp") diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4831d0a4..ad85742b 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -706,6 +706,12 @@ def add_data_arguments(parser): # --diffusion-height/--diffusion-width/--diffusion-output-num-frames. Encoded pairs are # cached next to the jsonl under .sft_cache/, one content-addressed file per sample. parser.add_argument("--sft-frame-stride", type=int, default=1, help="SFT encode temporal stride") + parser.add_argument( + "--sft-encoder-checkpoint", + type=str, + default=None, + help="HF name or path for the frozen SFT encoders (tokenizer/text_encoder/vae subfolders).", + ) parser.add_argument( "--start-rollout-id", @@ -1625,15 +1631,14 @@ def miles_validate_args(args): ) if args.prompt_data is None: raise ValueError("--loss-type sft_loss requires --prompt-data (jsonl with prompt + metadata.video)") - from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig - from miles.utils.misc import load_function - - sft_cfg_cls = load_function(args.train_pipeline_config_path) - if sft_cfg_cls.encode_sft_sample is TrainPipelineConfig.encode_sft_sample: + if args.sft_encoder_checkpoint is None: raise ValueError( - f"--loss-type sft_loss is not supported for {sft_cfg_cls.__name__}: " - "it does not implement load_sft_encoder/encode_sft_sample" + "--loss-type sft_loss requires --sft-encoder-checkpoint " + "(HF name or path holding the family's tokenizer/text_encoder/vae)" ) + from miles.rollout.encoder_hub import get_encoder + + get_encoder(args.diffusion_model_family).validate_args(args) if args.fsdp_flow_shift is None: raise ValueError("--loss-type sft_loss requires --fsdp-flow-shift for the training sigma grid") if args.n_samples_per_prompt != 1: diff --git a/scripts/run-diffusion-sft-wan22.sh b/scripts/run-diffusion-sft-wan22.sh index 4543eabf..d4f25888 100644 --- a/scripts/run-diffusion-sft-wan22.sh +++ b/scripts/run-diffusion-sft-wan22.sh @@ -53,6 +53,7 @@ WAN_LORA_TARGET_MODULES=( --custom-loss-function-path miles.backends.fsdp_utils.loss_hub.sft.sft_loss_formula \ --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --sft-encoder-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt-data "${SFT_DATA_JSONL}" \ --input-key prompt \ --diffusion-height 480 \ diff --git a/tests/fast/rollout/test_sft_sample_key.py b/tests/fast/rollout/test_sft_sample_key.py index 91e030a3..970eece4 100644 --- a/tests/fast/rollout/test_sft_sample_key.py +++ b/tests/fast/rollout/test_sft_sample_key.py @@ -12,7 +12,7 @@ def _args(**overrides): base = dict( - hf_checkpoint="ckpt", + sft_encoder_checkpoint="ckpt", diffusion_height=480, diffusion_width=832, diffusion_output_num_frames=81, @@ -41,7 +41,7 @@ def test_key_invalidates_per_axis(tmp_path): assert sft_sample_key(_args(diffusion_height=512), item)[0] != base_name assert sft_sample_key(_args(sft_frame_stride=1), item)[0] != base_name - assert sft_sample_key(_args(hf_checkpoint="other"), item)[0] != base_name + assert sft_sample_key(_args(sft_encoder_checkpoint="other"), item)[0] != base_name assert sft_sample_key(_args(), {"media": str(video), "prompt": "q"})[0] != base_name video.write_bytes(b"y" * 101) From cd44bab5229db4fb8f974b4852bf55ff2510260d Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 08:12:04 +0000 Subject: [PATCH 14/19] style: black-format the train_only exclusion assert Co-Authored-By: Claude Fable 5 --- miles/utils/arguments.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index ad85742b..fb1db8b7 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1696,9 +1696,9 @@ def miles_validate_args(args): args.offload_rollout = True del args.offload - assert not (args.debug_rollout_only and args.train_only), ( - "debug_rollout_only and train_only cannot be set at the same time, please set only one of them." - ) + assert not ( + args.debug_rollout_only and args.train_only + ), "debug_rollout_only and train_only cannot be set at the same time, please set only one of them." if args.debug_rollout_only: if args.colocate and (not args.rollout_num_gpus): From 68c88f86df53e3c2568fc7697e021ad14aae6be0 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 7 Aug 2026 11:41:26 +0000 Subject: [PATCH 15/19] feat(sft): support disaggregated encoder GPUs Treat explicit --rollout-num-gpus in train-only jobs as a dedicated rollout-side producer pool while preserving the colocated default. Co-authored-by: Cursor --- miles/ray/placement_group.py | 15 +++--- miles/utils/arguments.py | 6 +-- tests/fast/rollout/test_placement_group.py | 55 ++++++++++++++++++++++ 3 files changed, 67 insertions(+), 9 deletions(-) create mode 100644 tests/fast/rollout/test_placement_group.py diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index c56e9dad..82bae912 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -82,13 +82,16 @@ def create_placement_groups(args): """Create placement groups for actor and rollout engines. Two topologies: - - Colocate (or --train-only / --debug-rollout-only): one combined placement - group; both roles see the same bundle list. - - Disaggregate (the else branch): two separate placement groups so - train and rollout each own a disjoint GPU pool — avoids bundle - overlap / scheduling deadlock when running side-by-side. + - Colocate: one combined placement group; both roles see the same bundles. + Train-only jobs use this by default. + - Disaggregate: separate actor and rollout placement groups. A train-only + job opts into this by setting --rollout-num-gpus, reserving those GPUs for + its rollout-side producer (for example, the SFT encoder pool). """ - if not args.colocate and not args.train_only and not args.debug_rollout_only: + disaggregate = ( + not args.colocate and not args.debug_rollout_only and (not args.train_only or bool(args.rollout_num_gpus)) + ) + if disaggregate: logger.info("Creating placement groups (separate actor/rollout)...") actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node rollout_gpus = args.rollout_num_gpus diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index fb1db8b7..18ba848a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -43,9 +43,9 @@ def add_cluster_arguments(parser): type=int, default=None, help=( - "Number of GPUs for inference. Note that when using --colocate, " - "i.e. the training and the inference engines are on the same gpus, this param will be ignored and will be set as " - "actor_num_gpus_per_node * actor_num_nodes." + "Number of GPUs for rollout-side work. For train-only SFT, leave unset to colocate encoders " + "with training or set it to reserve dedicated encoder GPUs. Under --colocate this is overridden " + "to actor_num_gpus_per_node * actor_num_nodes." ), ) parser.add_argument( diff --git a/tests/fast/rollout/test_placement_group.py b/tests/fast/rollout/test_placement_group.py new file mode 100644 index 00000000..de559d6e --- /dev/null +++ b/tests/fast/rollout/test_placement_group.py @@ -0,0 +1,55 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +import miles.ray.placement_group as placement_module + + +def _args(**overrides): + values = dict( + actor_num_nodes=1, + actor_num_gpus_per_node=4, + rollout_num_gpus=None, + colocate=False, + train_only=True, + debug_rollout_only=False, + ) + values.update(overrides) + return Namespace(**values) + + +def _record_created_groups(monkeypatch): + calls = [] + + def fake_create(num_gpus): + calls.append(num_gpus) + return object(), list(range(num_gpus)), list(range(num_gpus)) + + monkeypatch.setattr(placement_module, "_create_placement_group", fake_create) + return calls + + +def test_train_only_reuses_actor_group_by_default(monkeypatch): + calls = _record_created_groups(monkeypatch) + groups = placement_module.create_placement_groups(_args()) + + assert calls == [4] + assert groups["actor"][0] is groups["rollout"][0] + + +def test_train_only_reserves_explicit_rollout_group(monkeypatch): + calls = _record_created_groups(monkeypatch) + groups = placement_module.create_placement_groups(_args(rollout_num_gpus=1)) + + assert calls == [4, 1] + assert groups["actor"][0] is not groups["rollout"][0] + assert len(groups["rollout"][1]) == 1 + + +def test_regular_disaggregated_training_is_unchanged(monkeypatch): + calls = _record_created_groups(monkeypatch) + placement_module.create_placement_groups(_args(train_only=False, rollout_num_gpus=2)) + + assert calls == [4, 2] From 4d2450c6ce0457e0e1a0157ba0233d12d8ef5a6f Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 7 Aug 2026 11:55:15 +0000 Subject: [PATCH 16/19] fix(ci): lazy-load rollout manager from placement helper Keep CPU placement tests from importing optional SGLang runtime modules during collection. Co-authored-by: Cursor --- miles/ray/placement_group.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 82bae912..3c10a213 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -6,7 +6,6 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from .actor_group import RayTrainGroup -from .rollout import RolloutManager logger = logging.getLogger(__name__) @@ -151,6 +150,8 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): + from .rollout import RolloutManager + logger.info("Creating rollout manager (num_gpus=%s)", 0) rollout_manager = RolloutManager.options( num_cpus=1, From 8f792a08b0864f997b6567f7aa10341b2f40d510 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 7 Aug 2026 12:39:37 +0000 Subject: [PATCH 17/19] Revert "fix(ci): lazy-load rollout manager from placement helper" This reverts commit 4d2450c6ce0457e0e1a0157ba0233d12d8ef5a6f. --- miles/ray/placement_group.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 3c10a213..82bae912 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -6,6 +6,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from .actor_group import RayTrainGroup +from .rollout import RolloutManager logger = logging.getLogger(__name__) @@ -150,8 +151,6 @@ def create_training_models(args, pgs, rollout_manager): def create_rollout_manager(args, pg): - from .rollout import RolloutManager - logger.info("Creating rollout manager (num_gpus=%s)", 0) rollout_manager = RolloutManager.options( num_cpus=1, From 6ebacad2ccc4c9a82f037037e907dcdd6c9d712b Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 7 Aug 2026 19:59:52 +0000 Subject: [PATCH 18/19] fix(ci): expose SGLang source package to CPU tests Put the checked-out Python package ahead of the workspace root so multimodal_gen imports resolve under SGLang's current layout. Co-authored-by: Cursor --- .github/workflows/_run-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml index e3e45255..c56c27c0 100644 --- a/.github/workflows/_run-ci.yml +++ b/.github/workflows/_run-ci.yml @@ -188,7 +188,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 60 env: - PYTHONPATH: ${{ github.workspace }} + PYTHONPATH: ${{ github.workspace }}/sglang/python:${{ github.workspace }} steps: - name: Free disk space shell: bash From f049028082cd34438abbf0a0ae6911d1b242d965 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 7 Aug 2026 20:10:45 +0000 Subject: [PATCH 19/19] fix(sglang): defer diffusion-only imports until engine startup Use SGLang's common process helper at module load so CPU-side orchestration can import without GPU-only diffusion dependencies. Co-authored-by: Cursor --- .github/workflows/_run-ci.yml | 2 +- .../sglang_diffusion_engine.py | 13 +++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml index c56c27c0..e3e45255 100644 --- a/.github/workflows/_run-ci.yml +++ b/.github/workflows/_run-ci.yml @@ -188,7 +188,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 60 env: - PYTHONPATH: ${{ github.workspace }}/sglang/python:${{ github.workspace }} + PYTHONPATH: ${{ github.workspace }} steps: - name: Free disk space shell: bash diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 153223d2..3de4dd1a 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -1,17 +1,22 @@ +from __future__ import annotations + import dataclasses import ipaddress import logging import multiprocessing import os import time +from typing import TYPE_CHECKING import requests -from sglang.multimodal_gen.runtime.launch_server import kill_process_tree -from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.srt.utils.common import kill_process_tree from miles.ray.ray_actor import RayActor from miles.utils.http_utils import get_host_info +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.server_args import ServerArgs + logger = logging.getLogger(__name__) @@ -139,6 +144,8 @@ def _format_v6_uri(addr): self._init_normal(server_args_dict) def _init_normal(self, server_args_dict): + from sglang.multimodal_gen.runtime.server_args import ServerArgs + logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}") self._pin_to_assigned_gpu() from miles.backends.sglang_diffusion_utils.monkey_patches import ROLLOUT_PATCH_GROUPS_ENV @@ -300,6 +307,8 @@ def simulate_crash(self): def _compute_server_args(args, host, port, nccl_port): + from sglang.multimodal_gen.runtime.server_args import ServerArgs + # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg. kwargs = {