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 new file mode 100644 index 00000000..581605da --- /dev/null +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -0,0 +1,99 @@ +"""SFT batch preparation and loss formula (rectified-flow velocity MSE).""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch +from miles.backends.fsdp_utils.loss_hub.utils import cast_cond_to_dtype +from miles.utils.metric_buffer import MetricBuffer + + +def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]: + """Pick one DiT component per micro-batch (phase-pure), then grid indices within its range.""" + num_grid = len(ctx.scheduler.timesteps) + if len(ctx.models) == 1: + component_name, model = next(iter(ctx.models.items())) + return component_name, model, torch.randint(num_grid, (bsz,)) + + # A uniform anchor index picks each expert with probability equal to its share of the + # grid, keeping the marginal over indices uniform while the micro-batch stays single-expert. + num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) + config = ctx.train_pipeline_config + components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps] + component_name = components[int(torch.randint(num_grid, (1,)))] + pool = torch.tensor([i for i, name in enumerate(components) if name == component_name]) + return component_name, ctx.models[component_name], pool[torch.randint(len(pool), (bsz,))] + + +def prepare_sft_batch( + ctx: DiffusionLossContext, + batch: list[dict], + *, + pad_to_len: int | None = None, +) -> PreparedBatch: + """Corrupt cached clean latents at sampled grid sigmas; CFG-free cached cond.""" + device = ctx.device + config = ctx.train_pipeline_config + bsz = len(batch) + + x0 = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32) + component_name, model, idx = sample_grid_indices(ctx, bsz) + idx = idx.to(device) + timesteps = ctx.scheduler.timesteps[idx].to(dtype=torch.float32) + sigmas = ctx.scheduler.sigmas[idx].to(dtype=torch.float32) + + noise = torch.randn(x0.shape, device=device, dtype=torch.float32) + sigma_exp = sigmas.view(bsz, *([1] * (x0.ndim - 1))) + latents = (1.0 - sigma_exp) * x0 + sigma_exp * noise + + num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) + if config.needs_timestep_scaling: + timesteps_for_model = timesteps / float(num_train_timesteps) + else: + timesteps_for_model = timesteps + + cond_list = [{key: value.to(device) for key, value in pair["cond_kwargs"].items()} for pair in batch] + pos_cond = cast_cond_to_dtype( + config.collate_cond_for_sample_batch(cond_list, device, pad_to_len=pad_to_len), + ctx.forward_dtype, + ) + + return PreparedBatch( + latents=latents, + timesteps=timesteps, + timesteps_for_model=timesteps_for_model, + model=model, + component_name=component_name, + guidance_scale=0.0, + use_cfg=False, + cfg_batching=False, + true_cfg_scale=None, + pos_cond=pos_cond, + neg_cond=None, + joint_cond=None, + advantage=torch.ones(bsz, device=device, dtype=torch.float32), + extras={"target": noise - x0}, + ) + + +def sft_loss_formula( + ctx: DiffusionLossContext, + batch: list[dict], + prepared: PreparedBatch, + *, + new_pred: torch.Tensor, + ref_pred: torch.Tensor | None, + metrics: MetricBuffer, + write_old_log_prob: bool = False, + old_log_prob_from_new: bool = False, +) -> torch.Tensor: + """Velocity-target MSE: ``||pred - (eps - x0)||^2`` averaged per pair.""" + target = prepared.extras["target"] + per_pair = ((new_pred.float() - target) ** 2).mean(dim=tuple(range(1, target.ndim))) + loss_sum = per_pair.sum() + + with torch.no_grad(): + metrics.emit_mean("loss", total=loss_sum, count=len(batch)) + return loss_sum 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 = { diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 94d5ab2e..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 --debug-{train,rollout}-only): one combined placement - group; both roles see the same bundle list. - - Disaggregate (the else branch): two separate placement groups so - train and rollout each own a disjoint GPU pool — avoids bundle - overlap / scheduling deadlock when running side-by-side. + - Colocate: one combined placement group; both roles see the same bundles. + Train-only jobs use this by default. + - Disaggregate: separate actor and rollout placement groups. A train-only + job opts into this by setting --rollout-num-gpus, reserving those GPUs for + its rollout-side producer (for example, the SFT encoder pool). """ - if not args.colocate and not args.debug_train_only and not args.debug_rollout_only: + disaggregate = ( + not args.colocate and not args.debug_rollout_only and (not args.train_only or bool(args.rollout_num_gpus)) + ) + if disaggregate: logger.info("Creating placement groups (separate actor/rollout)...") actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node rollout_gpus = args.rollout_num_gpus @@ -107,14 +110,11 @@ def create_placement_groups(args): logger.info(f"Creating placement group with {num_gpus} GPUs...") pg, all_reordered_bundle_indices, all_reordered_gpu_ids = _create_placement_group(num_gpus) - actor_pg_reordered_bundle_indices = all_reordered_bundle_indices - actor_pg_reordered_gpu_ids = all_reordered_gpu_ids - rollout_pg_reordered_bundle_indices = all_reordered_bundle_indices if not args.debug_train_only else [] - rollout_pg_reordered_gpu_ids = all_reordered_gpu_ids if not args.debug_train_only else [] - + # The rollout view keeps its seats under train_only: engine startup is gated by + # args.train_only, and rollout-side actor pools (e.g. the SFT encoder pool) seat here. return { - "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids), - "rollout": (pg, rollout_pg_reordered_bundle_indices, rollout_pg_reordered_gpu_ids), + "actor": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids), + "rollout": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids), } 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/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..96e6bb16 --- /dev/null +++ b/miles/rollout/sft_rollout.py @@ -0,0 +1,200 @@ +"""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.3 + + +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.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() + return digest.hex()[:16] + ".pt", int.from_bytes(digest[8:16], "big") % 2**63 + + +IMAGE_EXTENSIONS = {".bmp", ".jpeg", ".jpg", ".png", ".webp"} + + +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 --diffusion-output-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)) + 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): + from miles.rollout.encoder_hub import get_encoder + + self.args = args + 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 + for item in items: + pixels = read_media_clip( + item["media"], + 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"]) + 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") + 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: + # 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, + num_gpus=ENCODE_GPU_FRACTION, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ), + ).remote(args) + 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 + + +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.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 = ( + (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: + 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) + + 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 e0c9c4ac..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( @@ -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", @@ -690,6 +700,18 @@ 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 media path + # lives in each row's metadata dict under "video" or "image" (images train as single + # 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( + "--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", @@ -1085,12 +1107,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( @@ -1585,6 +1610,50 @@ 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 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 " + "--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)") + if args.sft_encoder_checkpoint is None: + raise ValueError( + "--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: + 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": @@ -1620,13 +1689,17 @@ 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 + 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): args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes @@ -1639,10 +1712,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 new file mode 100644 index 00000000..d4f25888 --- /dev/null +++ b/scripts/run-diffusion-sft-wan22.sh @@ -0,0 +1,85 @@ +#!/usr/bin/env bash +# 4-GPU Wan2.2-T2V-A14B dual-expert LoRA SFT on a (video, prompt) jsonl dataset. +# 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. + +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_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 + 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 \ + --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 \ + --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 \ + --sft-encoder-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --prompt-data "${SFT_DATA_JSONL}" \ + --input-key prompt \ + --diffusion-height 480 \ + --diffusion-width 832 \ + --diffusion-output-num-frames 81 \ + --sft-frame-stride 2 \ + --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 \ + --fsdp-flow-shift 3.0 \ + --save "${SAVE_DIR}" \ + --save-interval 20 \ + "${RESUME_ARGS[@]}" \ + "${WANDB_ARGS[@]}" 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..30e36b4b --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py @@ -0,0 +1,137 @@ +"""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_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()} + 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)) 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] 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..b1503c54 --- /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 --diffusion-output-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 new file mode 100644 index 00000000..970eece4 --- /dev/null +++ b/tests/fast/rollout/test_sft_sample_key.py @@ -0,0 +1,63 @@ +"""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.rollout.sft_rollout import sft_sample_key + + +def _args(**overrides): + base = dict( + sft_encoder_checkpoint="ckpt", + diffusion_height=480, + diffusion_width=832, + diffusion_output_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 = {"media": 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 = {"media": str(video), "prompt": "p"} + base_name, base_seed = sft_sample_key(_args(), item) + + 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(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) + 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(), {"media": str(a), "prompt": "p"}) + name_b, _ = sft_sample_key(_args(), {"media": str(b), "prompt": "p"}) + assert name_a != name_b