diff --git a/.gitignore b/.gitignore index 03c124b1..a017d4c6 100644 --- a/.gitignore +++ b/.gitignore @@ -177,6 +177,8 @@ settings.json .pypirc wandb/ +logs/ +dist/ outputs/ data/ local/ @@ -192,4 +194,4 @@ local/ glm/ _examples_synced/ .env -.DS_Store \ No newline at end of file +.DS_Store diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 103a5b2e..49bda3bc 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,13 +1,12 @@ import logging from argparse import Namespace from collections import defaultdict -from contextlib import nullcontext import ray import torch import torch.distributed as dist -from diffusers import DiffusionPipeline +import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.qwen_image # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.sd3 # noqa: F401 — register pipeline config from miles.ray.train_actor import TrainRayActor @@ -17,7 +16,6 @@ from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.metric_utils import compute_rollout_step from miles.utils.profile_utils import TrainProfiler -from miles.utils.sde_log_prob import sde_step_with_logprob from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils import init_tracking @@ -65,38 +63,33 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self._master_dtype = _resolve_dtype(args.fsdp_master_dtype) self._forward_dtype = _resolve_dtype(args.diffusion_forward_dtype) - with self._get_init_weight_context_manager(): - pipeline = DiffusionPipeline.from_pretrained( - self.args.hf_checkpoint, - torch_dtype=self._master_dtype, - trust_remote_code=True, - text_encoder=None, - vae=None, - tokenizer=None, - ) - model = pipeline.transformer - self.scheduler = pipeline.scheduler - del pipeline + diffusion_model_id = args.diffusion_model or args.hf_checkpoint + self.train_pipeline_config = get_train_pipeline_config(diffusion_model_id) + self.train_step_backend = self.train_pipeline_config.get_train_step_backend() - self.train_pipeline_config = get_train_pipeline_config(args.diffusion_model) + model, self.scheduler = self.train_step_backend.load_model_and_scheduler( + self.args, + init_context_factory=self._get_init_weight_context_manager, + master_dtype=self._master_dtype, + ) if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) model.train() - - if args.gradient_checkpointing: - model.enable_gradient_checkpointing() + self.train_step_backend.apply_gradient_checkpointing(model, args) model.to(torch.cuda.current_device()) self.train_pipeline_config.preprocess_model_before_fsdp(model) + fsdp_wrap = self.train_step_backend.get_fsdp_wrap_classes() model = apply_fsdp2( model, mesh=self.parallel_state.dp_mesh, cpu_offload=self.args.fsdp_cpu_offload, args=self.args, + fsdp_wrap_classes=fsdp_wrap, ) # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -298,8 +291,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: # ------------- CFG Scale ------------- guidance_scale = self.args.diffusion_guidance_scale true_cfg_scale = self.args.diffusion_true_cfg_scale - cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale - use_cfg = cfg_scale > 0 + use_cfg = self.train_step_backend.should_use_cfg(self.args) # ------------- KL loss ------------- kl_beta = float(self.args.diffusion_kl_beta) @@ -325,15 +317,16 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: advantages = torch.clamp(advantages, -adv_clip_max, adv_clip_max) # ------------- scheduler ------------- - # Use rollout's exact sigmas snapshot; fall back to reconstruction if unavailable. - num_train_timesteps = self.scheduler.config.num_train_timesteps timesteps_ref = dit_trajectories[0].timesteps.to(device).float() sigmas_snapshot = getattr(dit_trajectories[0], "sigmas", None) - if sigmas_snapshot is not None: - sigmas_ref = sigmas_snapshot.to(device).float() - else: - sigmas_ref = timesteps_ref / float(num_train_timesteps) - sigmas_ref = torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + sched_config = getattr(self.scheduler, "config", None) + num_train_timesteps = int(sched_config.num_train_timesteps) if sched_config is not None else 1000 + sigmas_ref = self.train_step_backend.resolve_sigmas_ref( + timesteps_ref, + sigmas_snapshot, + self.scheduler, + num_train_timesteps=num_train_timesteps, + ) self.scheduler.timesteps = timesteps_ref self.scheduler.sigmas = sigmas_ref @@ -434,6 +427,7 @@ def _build_train_grids( log_prob_old_list, advantage_list = [], [] positive_cond_kwargs_list, negative_cond_kwargs_list = [], [] rollout_model_outputs_list: list[torch.Tensor] = [] + sde_indices_per_sample_list: list[torch.Tensor] = [] sde_window_size: int | None = None for traj_idx in range(traj_start, traj_end): @@ -445,7 +439,12 @@ def _build_train_grids( # prepare cond kwargs (denoising_env) denoising_env = denoising_envs[traj_idx] positive_cond_kwargs_list.append( - train_pipeline_config.prepare_cond_kwargs(denoising_env.pos_cond_kwargs, device) + train_pipeline_config.build_train_cond_kwargs( + denoising_env.pos_cond_kwargs, + latents=latents, + args=self.args, + device=device, + ) ) if use_cfg: negative_cond_kwargs_list.append( @@ -464,6 +463,7 @@ def _build_train_grids( ) sde_step_indices = sde_step_indices_list[traj_idx] + sde_indices_tensor = None if sde_step_indices is not None: sde_indices_tensor = torch.as_tensor(sde_step_indices, device=device, dtype=torch.long) latents = latents[sde_indices_tensor] @@ -472,7 +472,10 @@ def _build_train_grids( log_prob_old = log_prob_old[sde_indices_tensor] advantage = advantage[: sde_indices_tensor.numel()] if rollout_model_output is not None: - rollout_model_output = rollout_model_output[sde_indices_tensor] + n_mo = int(rollout_model_output.shape[0]) + n_win = int(sde_indices_tensor.numel()) + if n_mo != n_win: + rollout_model_output = rollout_model_output[sde_indices_tensor] current_window_size = int(sde_indices_tensor.numel()) else: current_window_size = default_window_size @@ -491,6 +494,8 @@ def _build_train_grids( advantage_list.append(advantage) if rollout_model_output is not None: rollout_model_outputs_list.append(rollout_model_output) + if sde_indices_tensor is not None: + sde_indices_per_sample_list.append(sde_indices_tensor) latents_window = torch.stack(latents_list, dim=0) next_latents_window = torch.stack(next_latents_list, dim=0) @@ -532,6 +537,11 @@ def _build_train_grids( "num_samples_in_window": num_samples_in_window, "sde_window_size": int(sde_window_size or 0), "rollout_model_outputs": rollout_model_outputs_window, + "sde_step_indices_window": ( + torch.stack(sde_indices_per_sample_list, dim=0) + if sde_indices_per_sample_list and len(sde_indices_per_sample_list) == (traj_end - traj_start) + else None + ), } def _run_optim_window( @@ -609,6 +619,7 @@ def _forward_tile( flattened to batch = tile_sample * tile_tstep.""" forward_dtype = self._forward_dtype train_pipeline_config = self.train_pipeline_config + train_step_backend = self.train_step_backend tile_sample_count = int(sample_indices.numel()) tile_tstep_count = int(tstep_indices.numel()) num_samples_in_window = grids["num_samples_in_window"] @@ -621,6 +632,7 @@ def _forward_tile( latents_flat = latents_tile.reshape(tile_sample_count * tile_tstep_count, *latents_tile.shape[2:]) timesteps_flat = timesteps_tile.reshape(tile_sample_count * tile_tstep_count) + timesteps_for_sde = train_step_backend.scale_timesteps_for_sde(timesteps_flat) # sgl-d's Qwen DiT divides timestep by num_train_timesteps inside # forward; diffusers' does not. SD3 already expects raw timesteps. @@ -664,52 +676,33 @@ def _forward_tile( # than rollout → systematic noise_pred drift. latents_input = latents_flat.to(forward_dtype) timesteps_input = timesteps_for_model.to(forward_dtype) + prev_sample_flat = next_latents_tile.reshape( + tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] + ) - def _forward(cond: dict) -> torch.Tensor: - return self.model( - hidden_states=latents_input, - timestep=timesteps_input, - return_dict=False, - **cond, - )[0] - - cfg_batching = bool(self.args.fsdp_cfg_batching) - - def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: - adapter_ctx = self.model.disable_adapter() if disable_adapter else nullcontext() - with adapter_ctx: - if not use_cfg: - return _forward(pos_cond_tile) - if cfg_batching: - joint_cond = _pack_cond_for_joint_cfg(pos_cond_tile, neg_cond_tile) - joint_out = self.model( - hidden_states=torch.cat([latents_input, latents_input], dim=0), - timestep=torch.cat([timesteps_input, timesteps_input], dim=0), - return_dict=False, - **joint_cond, - )[0] - noise_pred_pos, noise_pred_neg = joint_out.chunk(2, dim=0) - else: - noise_pred_pos = _forward(pos_cond_tile) - noise_pred_neg = _forward(neg_cond_tile) - return train_pipeline_config.cfg_combine( - noise_pred_pos, - noise_pred_neg, - guidance_scale, - true_cfg_scale=true_cfg_scale, - ) - - noise_pred_flat = _compute_noise_pred() - - _, log_prob_new_flat, prev_sample_mean, std_dev_t = sde_step_with_logprob( - self.scheduler, - noise_pred_flat.float(), - timesteps_flat, - latents_flat.float(), - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ).float(), + noise_pred_flat = train_step_backend.compute_noise_pred( + model=self.model, + latents_input=latents_input, + timesteps_input=timesteps_input, + pos_cond=pos_cond_tile, + neg_cond=neg_cond_tile, + use_cfg=use_cfg, + guidance_scale=guidance_scale, + true_cfg_scale=true_cfg_scale, + fsdp_cfg_batching=bool(self.args.fsdp_cfg_batching), + ) + log_prob_new_flat, prev_sample_mean, std_dev_t = train_step_backend.sde_step_logprob( + scheduler=self.scheduler, + noise_pred=noise_pred_flat, + timesteps_for_sde=timesteps_for_sde, + timesteps_flat=timesteps_flat, + latents_flat=latents_flat, + prev_sample=prev_sample_flat, noise_level=noise_level, + grids=grids, + sample_indices=sample_indices, + tstep_indices=tstep_indices, + args=self.args, ) # TODO: revamp and gather all loss logics @@ -722,17 +715,30 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: kl_loss = log_prob_new.new_zeros(()) if kl_beta > 0: with torch.no_grad(): - ref_noise_pred_flat = _compute_noise_pred(disable_adapter=True) - # TODO: unify sde_step_with_logprob with rollout and trainer forward paths. - _, _, prev_sample_mean_ref, _ = sde_step_with_logprob( - self.scheduler, - ref_noise_pred_flat.float(), - timesteps_flat, - latents_flat.float(), - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ).float(), + ref_noise_pred_flat = train_step_backend.compute_noise_pred( + model=self.model, + latents_input=latents_input, + timesteps_input=timesteps_input, + pos_cond=pos_cond_tile, + neg_cond=neg_cond_tile, + use_cfg=use_cfg, + guidance_scale=guidance_scale, + true_cfg_scale=true_cfg_scale, + fsdp_cfg_batching=bool(self.args.fsdp_cfg_batching), + disable_adapter=True, + ) + _, prev_sample_mean_ref, _ = train_step_backend.sde_step_logprob( + scheduler=self.scheduler, + noise_pred=ref_noise_pred_flat, + timesteps_for_sde=timesteps_for_sde, + timesteps_flat=timesteps_flat, + latents_flat=latents_flat, + prev_sample=prev_sample_flat, noise_level=noise_level, + grids=grids, + sample_indices=sample_indices, + tstep_indices=tstep_indices, + args=self.args, ) kl_loss = ((prev_sample_mean - prev_sample_mean_ref) ** 2).mean( dim=tuple(range(1, prev_sample_mean.ndim)), @@ -767,6 +773,7 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: log_stats["model_output_max_abs_diff"].append(diff.max().detach()) log_stats["model_output_mean_abs_diff"].append(diff.mean().detach()) log_stats["model_output_rel_max"].append((diff.max() / ref_max).detach()) + train_step_backend.append_model_output_compare_stats(log_stats, noise_pred_flat, rollout_mo_flat) return loss @@ -814,20 +821,6 @@ def _tile_value(value, rows: torch.Tensor): return pos, neg -def _pack_cond_for_joint_cfg(pos: dict, neg: dict) -> dict: - """Pack pos and neg per-tile cond dicts into a single [pos | neg] dict - along the batch dim, for joint CFG forward.""" - out: dict = {} - for key, value in pos.items(): - if isinstance(value, torch.Tensor): - out[key] = torch.cat([value, neg[key]], dim=0) - elif isinstance(value, list): - out[key] = value + neg[key] - else: - out[key] = value - return out - - def _cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict: """Cast floating-point tensors to the model's compute dtype; leave bool masks / int / list / scalar values untouched. The bool @@ -891,13 +884,17 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, fsdp_wrap_classes=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None - layer_cls_to_wrap = model._no_split_modules - assert len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None + layer_cls_to_wrap = getattr(model, "_no_split_modules", None) + if not layer_cls_to_wrap: + layer_cls_to_wrap = fsdp_wrap_classes + assert ( + layer_cls_to_wrap and layer_cls_to_wrap[0] is not None + ), "apply_fsdp2 needs model._no_split_modules or fsdp_wrap_classes for LTX" modules = [module for name, module in model.named_modules() if module.__class__.__name__ in layer_cls_to_wrap] diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py new file mode 100644 index 00000000..156e6228 --- /dev/null +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -0,0 +1,688 @@ +"""LTX-2: model-family helpers + FSDP train pipeline config.""" + +from __future__ import annotations + +import logging +import os +from argparse import ArgumentParser, Namespace +from pathlib import Path +from typing import Any + +import torch + +from miles.backends.fsdp_utils.train_step_backend import LTXTrainStepBackend +from miles.utils.types import CondKwargs + +from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config + +logger = logging.getLogger(__name__) + +LTX_DEFAULT_HF_MODEL = "Lightricks/LTX-2.3" +LTX_DEFAULT_MODEL_ID = "LTX-2.3" + + +def is_ltx_model(args) -> bool: + model_type = (getattr(args, "diffusion_model_type", "auto") or "auto").lower() + if model_type == "ltx": + return True + if model_type != "auto": + return False + return _looks_like_ltx_ref(getattr(args, "diffusion_model", None)) + + +def _looks_like_ltx_ref(diffusion_model: str | None) -> bool: + if not diffusion_model: + return False + ref = str(diffusion_model).lower() + return "ltx" in ref or ref.endswith(".safetensors") + + +def _is_hf_model_id(ref: str | None) -> bool: + if not ref: + return False + text = str(ref) + if text.endswith(".safetensors") or os.path.exists(text): + return False + return "/" in text or "ltx" in text.lower() + + +def resolve_hf_model_id(args) -> str: + """HF hub id used for sglang ``model_path`` / overlay materialization.""" + diffusion_model = getattr(args, "diffusion_model", None) + if _is_hf_model_id(diffusion_model): + return str(diffusion_model) + if getattr(args, "sglang_model_path", None): + return str(args.sglang_model_path) + env_path = os.environ.get("MILES_LTX_ROLLOUT_MODEL_PATH") + if env_path: + return env_path + return LTX_DEFAULT_HF_MODEL + + +def resolve_model_id(args) -> str: + """Short registry id for sglang ``ServerArgs.model_id``.""" + if getattr(args, "sglang_model_id", None): + return str(args.sglang_model_id) + env_id = os.environ.get("MILES_LTX_MODEL_ID") + if env_id: + return env_id + return LTX_DEFAULT_MODEL_ID + + +def _diffusion_cache_root() -> Path: + return Path(os.environ.get("SGLANG_DIFFUSION_CACHE_ROOT", "/data/wenhao/sgl_diffusion_cache")) + + +def _find_cached_materialized_dir(hf_model_id: str) -> Path | None: + materialized = _diffusion_cache_root() / "materialized_models" + if not materialized.is_dir(): + return None + + prefix = hf_model_id.replace("/", "__") + "-" + candidates = sorted( + (d for d in materialized.iterdir() if d.is_dir() and d.name.startswith(prefix)), + key=lambda p: p.stat().st_mtime, + reverse=True, + ) + for directory in candidates: + checkpoint = directory / "transformer" / "model.safetensors" + if checkpoint.is_file(): + return directory + return None + + +def _transformer_checkpoint_in_dir(materialized_dir: Path) -> Path: + checkpoint = materialized_dir / "transformer" / "model.safetensors" + if not checkpoint.is_file(): + raise FileNotFoundError( + f"Materialized LTX model at {materialized_dir} is missing " f"transformer/model.safetensors" + ) + return checkpoint + + +def _materialized_config_path(checkpoint: Path) -> Path | None: + """Return sibling ``config.json`` for sglang overlay materialized DiT weights.""" + config_json = checkpoint.parent / "config.json" + return config_json if config_json.is_file() else None + + +def _is_materialized_diffusers_checkpoint(checkpoint: Path) -> bool: + return _materialized_config_path(checkpoint) is not None + + +def _read_materialized_transformer_config(checkpoint: Path) -> dict: + import json + + config_json = _materialized_config_path(checkpoint) + if config_json is None: + raise FileNotFoundError(f"Materialized LTX checkpoint {checkpoint} is missing sibling config.json") + transformer_cfg = json.loads(config_json.read_text()) + return {"transformer": transformer_cfg} + + +def load_ltx_transformer_for_train( + checkpoint_path: str | Path, + *, + device: str = "cpu", + dtype: Any = None, +): + """Load LTX DiT for FSDP train from materialized diffusers or comfy safetensors. + + Materialized overlay weights (``transformer/model.safetensors`` + ``config.json``) + use the same key layout as ltx_core / sglang and do not embed config in safetensors + metadata. Comfy-style single-file checkpoints keep using safetensors metadata. + """ + import torch + from ltx_core.loader.helpers import create_meta_model, load_state_dict + from ltx_core.loader.registry import DummyRegistry + from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader + from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder + from ltx_core.model.transformer.model_configurator import LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator + + checkpoint = Path(checkpoint_path).expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"LTX checkpoint not found: {checkpoint}") + + torch_device = torch.device(device) if isinstance(device, str) else device + if dtype is None: + dtype = torch.bfloat16 + + if _is_materialized_diffusers_checkpoint(checkpoint): + config = _read_materialized_transformer_config(checkpoint) + meta_model = create_meta_model(LTXModelConfigurator, config, ()) + loader = SafetensorsModelStateDictLoader() + sd = load_state_dict( + str(checkpoint), + loader, + DummyRegistry(), + torch.device("cpu"), + None, + ) + state_dict = sd.sd + if dtype is not None: + state_dict = {key: value.to(dtype=dtype) for key, value in state_dict.items()} + meta_model.load_state_dict(state_dict, strict=False, assign=True) + logger.info( + "LTX train: loaded materialized diffusers transformer from %s", + checkpoint, + ) + return meta_model.to(torch_device) + + return SingleGPUModelBuilder( + model_path=str(checkpoint), + model_class_configurator=LTXModelConfigurator, + model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, + ).build(device=torch_device, dtype=dtype) + + +def ensure_materialized_model(hf_model_id: str) -> Path: + """Materialize the overlay model via sglang (same pipeline as rollout). + + Downloads HF source weights + overlay metadata on first use, then caches + under ``SGLANG_DIFFUSION_CACHE_ROOT/materialized_models/``. + """ + cached = _find_cached_materialized_dir(hf_model_id) + if cached is not None: + return cached + + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model + + logger.info( + "LTX: materializing overlay model for %s (first run may download HF weights)", + hf_model_id, + ) + materialized = maybe_download_model( + hf_model_id, + download=True, + force_diffusers_model=True, + ) + materialized_dir = Path(materialized) + _transformer_checkpoint_in_dir(materialized_dir) + return materialized_dir + + +def resolve_materialized_model_dir( + hf_model_id: str, + *, + materialize: bool = True, +) -> Path | None: + cached = _find_cached_materialized_dir(hf_model_id) + if cached is not None: + return cached + if not materialize: + return None + return ensure_materialized_model(hf_model_id) + + +def resolve_transformer_checkpoint( + diffusion_model: str | None, + *, + explicit_path: str | None = None, + materialize: bool = True, +) -> str: + """Resolve the single-file DiT checkpoint used by FSDP train. + + Resolution order: + 1. Explicit ``--sglang-transformer-weights-path`` / env override + 2. ``--diffusion-model`` pointing at a ``.safetensors`` file + 3. Overlay materialized ``transformer/model.safetensors`` for a HF model id + (materializes via sglang on cache miss when ``materialize=True``) + """ + if explicit_path: + path = Path(explicit_path).expanduser() + if path.is_file(): + return str(path) + raise FileNotFoundError(f"LTX transformer checkpoint not found: {path}") + + env_path = os.environ.get("MILES_LTX_TRANSFORMER_WEIGHTS_PATH") + if env_path: + path = Path(env_path).expanduser() + if path.is_file(): + return str(path) + raise FileNotFoundError(f"LTX transformer checkpoint not found: {path}") + + if diffusion_model: + path = Path(str(diffusion_model)).expanduser() + if path.is_file() and path.suffix == ".safetensors": + return str(path) + + if _is_hf_model_id(str(diffusion_model)): + materialized_dir = resolve_materialized_model_dir( + str(diffusion_model), + materialize=materialize, + ) + if materialized_dir is not None: + checkpoint = _transformer_checkpoint_in_dir(materialized_dir) + logger.info( + "LTX train: using materialized transformer %s (from %s)", + checkpoint, + materialized_dir, + ) + return str(checkpoint) + + materialized_dir = resolve_materialized_model_dir( + LTX_DEFAULT_HF_MODEL, + materialize=materialize, + ) + if materialized_dir is not None: + checkpoint = _transformer_checkpoint_in_dir(materialized_dir) + logger.info("LTX train: using default materialized transformer %s", checkpoint) + return str(checkpoint) + + raise FileNotFoundError( + "Could not resolve LTX transformer checkpoint. Pass --diffusion-model " + "Lightricks/LTX-2.3 (recommended) or a .safetensors override." + ) + + +def server_kwargs_extras(args) -> dict: + """Extra ``ServerArgs`` kwargs; call only when ``is_ltx_model(args)``.""" + hf_model_id = resolve_hf_model_id(args) + extras: dict = {"model_id": resolve_model_id(args)} + + # Only override rollout DiT when user explicitly passes a safetensors path. + # For ``--diffusion-model Lightricks/LTX-2.3`` both sides use overlay defaults. + explicit = getattr(args, "sglang_transformer_weights_path", None) + weights_path = None + if explicit: + weights_path = resolve_transformer_checkpoint( + getattr(args, "diffusion_model", None), + explicit_path=explicit, + ) + elif getattr(args, "diffusion_model", None) and str(args.diffusion_model).endswith(".safetensors"): + weights_path = resolve_transformer_checkpoint(args.diffusion_model) + + if weights_path: + extras["transformer_weights_path"] = weights_path + logger.info( + "LTX rollout: model_path=%s model_id=%s transformer_weights_path=%s", + hf_model_id, + extras["model_id"], + weights_path, + ) + else: + logger.info( + "LTX rollout: model_path=%s model_id=%s (overlay default DiT weights)", + hf_model_id, + extras["model_id"], + ) + + gemma_path = getattr(args, "ltx_gemma_path", None) + if gemma_path: + extras["component_paths"] = {"text_encoder": gemma_path} + + return extras + + +def _normalize_ltx_dynamics_type(name: str) -> str: + key = str(name).strip().lower().replace("-", "_") + allowed = ("flow_sde", "cps", "ode", "dance_sde") + if key not in allowed: + raise ValueError(f"Unknown ltx dynamics_type {name!r}; expected one of {allowed}") + return key + + +def patch_rollout_sampling_params( + sampling_params: dict[str, Any], + args: Namespace, + *, + evaluation: bool, +) -> None: + """Apply LTX-specific rollout sampling fields in-place.""" + if getattr(args, "ltx_frames", None) is not None: + sampling_params["num_frames"] = int(args.ltx_frames) + if getattr(args, "ltx_fps", None) is not None: + sampling_params["fps"] = int(args.ltx_fps) + sampling_params["guidance_scale"] = 1.0 + sampling_params["negative_prompt"] = None + + if evaluation: + return + + dynamics = _normalize_ltx_dynamics_type(getattr(args, "ltx_dynamics_type", "cps")) + if dynamics == "dance_sde": + raise NotImplementedError("dance_sde rollout is not implemented in sglang-d flow_sde_sampling yet.") + sampling_params["rollout_sde_type"] = dynamics + if dynamics in ("cps", "ode"): + sampling_params["rollout_log_prob_no_const"] = True + elif dynamics == "flow_sde": + ltx_sigma_min = getattr(args, "ltx_sigma_min", None) + if ltx_sigma_min is not None: + sampling_params["rollout_sigma_min"] = float(ltx_sigma_min) + + +def patch_rollout_engine_env_vars(env_vars: dict[str, str], args) -> None: + """Add LTX-specific env vars for Ray rollout engine workers.""" + if not is_ltx_model(args): + return + + from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV + + if os.environ.get(LTX_ROLLOUT_PATCHES_ENV): + env_vars[LTX_ROLLOUT_PATCHES_ENV] = os.environ[LTX_ROLLOUT_PATCHES_ENV] + + +def register_args(parser: ArgumentParser) -> None: + parser.add_argument( + "--ltx-frames", + type=int, + default=25, + help="LTX video frame count (e.g. 57 for verl-omni default).", + ) + parser.add_argument( + "--ltx-fps", + type=float, + default=24.0, + help="LTX video fps for rollout VAE rescale.", + ) + parser.add_argument( + "--ltx-num-sde-steps", + type=int, + default=3, + help="Number of denoising steps with SDE noise + log_prob during LTX rollout.", + ) + parser.add_argument( + "--ltx-sde-step-candidates", + type=str, + default=None, + help="Comma-separated SDE step indices for LTX rollout (e.g. 0,1,...,9).", + ) + parser.add_argument( + "--ltx-dynamics-type", + type=str, + default="CPS", + choices=["Flow-SDE", "CPS", "ODE", "Dance-SDE"], + help="Stochastic dynamics for LTX SDE step during training.", + ) + parser.add_argument( + "--ltx-sigma-min", + type=float, + default=None, + help="Override σ_min for LTX SDE step.", + ) + parser.add_argument( + "--ltx-gemma-path", + type=str, + default=None, + help=( + "Deprecated: optional text_encoder override. " + "When unset, sglang overlay materializes text_encoder from --diffusion-model." + ), + ) + parser.add_argument( + "--pickscore-num-frames", + type=int, + default=3, + help="Number of evenly spaced frames to score per video (LTX PickScore reward).", + ) + + +def validate_args(args: Namespace) -> None: + ltx_gs = float(getattr(args, "diffusion_guidance_scale", 1.0)) + if ltx_gs != 1.0: + logger.warning( + "LTX rollout/train alignment expects --diffusion-guidance-scale 1.0 " + "(no CFG); using %s may break log_prob parity.", + ltx_gs, + ) + if getattr(args, "fsdp_master_dtype", "fp32") == "fp32": + logger.warning( + "diffusion_model_type=ltx with fsdp_master_dtype=fp32 is unlikely to fit " + "on small GPU counts; consider --fsdp-master-dtype bf16." + ) + + +# --- FSDP train pipeline config --- + + +@register_train_pipeline_config("ltx") +class LTXTrainPipelineConfig(TrainPipelineConfig): + """Training-side adapter for LTX-2.3 video DiT.""" + + train_step_backend_cls = LTXTrainStepBackend + needs_timestep_scaling = False + # Rollout stores σ×1000 in dit_trajectory.timesteps; CPS uses scheduler σ∈[0,1]. + sde_timestep_divisor = 1000.0 + + lora_target_modules = [ + "to_q", + "to_k", + "to_v", + "to_out.0", + "net.0.proj", + "net.2", + ] + + def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: + if cond is None: + return {} + kwargs: dict = {} + if cond.encoder_hidden_states: + ctx = torch.cat(cond.encoder_hidden_states).to(device) + if ctx.ndim == 2: + ctx = ctx.unsqueeze(0) + kwargs["context"] = ctx + if cond.audio_encoder_hidden_states: + audio_ctx = torch.cat(cond.audio_encoder_hidden_states).to(device) + if audio_ctx.ndim == 2: + audio_ctx = audio_ctx.unsqueeze(0) + kwargs["audio_context"] = audio_ctx + if cond.encoder_attention_mask is not None: + mask = cond.encoder_attention_mask.to(device) + if mask.ndim == 1: + mask = mask.unsqueeze(0) + kwargs["context_mask"] = mask + if cond.audio_encoder_attention_mask is not None: + audio_mask = cond.audio_encoder_attention_mask.to(device) + if audio_mask.ndim == 1: + audio_mask = audio_mask.unsqueeze(0) + kwargs["audio_context_mask"] = audio_mask + return kwargs + + def build_train_cond_kwargs( + self, + cond: CondKwargs | None, + *, + latents: torch.Tensor, + args, + device: torch.device, + ) -> dict: + """Merge rollout text embeds with locally rebuilt T2V geometry.""" + from miles.backends.fsdp_utils.ltx_geometry import build_ltx_t2v_geometry + + kwargs = self.prepare_cond_kwargs(cond, device) + if "context" not in kwargs: + raise ValueError("LTX train requires denoising_env.pos_cond_kwargs.encoder_hidden_states") + + ref = latents[0] if latents.ndim >= 2 else latents + if ref.ndim == 2: + batch_size, num_tokens, latent_dim = 1, ref.shape[0], ref.shape[1] + else: + batch_size, num_tokens, latent_dim = ref.shape[0], ref.shape[1], ref.shape[2] + + geom = build_ltx_t2v_geometry( + batch_size=batch_size, + num_tokens=num_tokens, + latent_dim=latent_dim, + height=int(getattr(args, "diffusion_height", 512)), + width=int(getattr(args, "diffusion_width", 512)), + num_frames=int(getattr(args, "ltx_frames", 25)), + fps=float(getattr(args, "ltx_fps", 24.0)), + device=device, + dtype=ref.dtype, + ) + kwargs.update(geom) + return kwargs + + def build_sde_extra( + self, + scheduler, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + args, + ) -> dict | None: + window = grids.get("sde_step_indices_window") + if window is None: + return None + idx = window[sample_indices][:, tstep_indices].reshape(-1).long() + return { + "sde_step_indices": idx, + "sigmas": scheduler.sigmas, + "dynamics_type": getattr(args, "ltx_dynamics_type", "cps"), + "sigma_min_override": getattr(args, "ltx_sigma_min", None), + } + + def expand_cond_for_timestep_batch(self, cond_kwargs: dict, batch_size: int) -> dict: + out: dict = {} + for k, v in cond_kwargs.items(): + if isinstance(v, torch.Tensor): + out[k] = v.expand(batch_size, *v.shape[1:]) if v.shape[0] == 1 else v + else: + out[k] = v + return out + + def collate_cond_for_sample_batch( + self, + per_sample_cond_kwargs: list[dict], + device: torch.device, + ) -> dict: + out: dict = {} + for key in per_sample_cond_kwargs[0]: + values = [kw[key] for kw in per_sample_cond_kwargs if key in kw] + if not values: + continue + if isinstance(values[0], torch.Tensor): + out[key] = torch.cat(values, dim=0).to(device) + else: + out[key] = values + return out + + def cfg_combine( + self, + noise_pred_pos: torch.Tensor, + noise_pred_neg: torch.Tensor, + guidance_scale: float, + true_cfg_scale: float | None = None, + ) -> torch.Tensor: + scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale + if scale == 1.0: + return noise_pred_pos + return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) + + def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: + return None + + @staticmethod + def _modality_timesteps_for_adaln(per_token_t: torch.Tensor) -> torch.Tensor: + """Collapse per-token sigma to batch-global AdaLN input when uniform. + + sglang rollout builds temb with shape ``[B, 1, D]`` (scheduler timestep + is batch-scalar expanded only for masking). ltx_core defaults to + ``[B, T, D]`` when ``Modality.timesteps`` has length T, which diverges + in AdaLN even when every active token shares the same sigma. + """ + if per_token_t.ndim != 2 or per_token_t.shape[1] == 1: + return per_token_t + ref = per_token_t[:, :1] + if torch.allclose(per_token_t, ref.expand_as(per_token_t), rtol=0.0, atol=0.0): + return ref + return per_token_t + + def forward_velocity( + self, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + cond: dict, + ) -> torch.Tensor: + from ltx_core.model.transformer.modality import Modality + from ltx_core.utils import to_denoised + + device = latents_input.device + dtype = latents_input.dtype + B = latents_input.shape[0] + + # dit_trajectory.timesteps are σ×1000; ltx_core AdaLN expects σ∈[0,1] and + # multiplies by timestep_scale_multiplier (1000) internally. + sigma_scaled = timesteps_input.to(latents_input.dtype) + sigma_unit = sigma_scaled / float(self.sde_timestep_divisor) + denoise_mask = cond["denoise_mask"].to(device) + denoise_mask_2d = denoise_mask.squeeze(-1) if denoise_mask.ndim == 3 else denoise_mask + denoise_mask_float = denoise_mask_2d.float() + + per_token_t = (sigma_unit.view(B, 1) * denoise_mask_2d).to(dtype) + adaln_timesteps = self._modality_timesteps_for_adaln(per_token_t) + + video_modality = Modality( + enabled=True, + latent=latents_input, + sigma=sigma_unit.reshape(B), + timesteps=adaln_timesteps, + positions=cond["positions"].to(dtype), + context=cond["context"].to(dtype), + context_mask=None, + ) + with torch.autocast(device_type=str(device).split(":")[0], dtype=dtype): + velocity, _ = model(video=video_modality, audio=None, perturbations=None) + + per_token_t_3d = per_token_t.unsqueeze(-1) if per_token_t.ndim == 2 else per_token_t + x0_pred = to_denoised(latents_input, velocity, per_token_t_3d).float() + + clean_latent = cond["clean_latent"].to(device).float() + denoise_mask_3d = denoise_mask_float.unsqueeze(-1) if denoise_mask_float.ndim == 2 else denoise_mask_float + x0_pred = x0_pred * denoise_mask_3d + clean_latent * (1.0 - denoise_mask_3d) + + sigma_safe = torch.clamp(sigma_unit, min=1e-8).view(B, 1, 1) + velocity_for_sde = (latents_input.float() - x0_pred) / sigma_safe + return velocity_for_sde.to(dtype) + + def forward_velocity_cfg_joint( + self, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + joint_cond: dict, + ) -> torch.Tensor: + raise NotImplementedError("LTX trains with guidance_scale=1.0; --fsdp-cfg-batching is not supported.") + + def sde_step( + self, + scheduler, + noise_pred: torch.Tensor, + timesteps: torch.Tensor, + sample: torch.Tensor, + prev_sample: torch.Tensor, + *, + noise_level: float, + extra: dict | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + from miles.utils.sde_log_prob import sde_step_with_logprob + + if extra is None or "sigmas" not in extra or "sde_step_indices" not in extra: + raise ValueError("LTXTrainPipelineConfig.sde_step requires extra={'sigmas','sde_step_indices',...}") + sigmas = extra["sigmas"].to(sample.device).float() + step_indices = extra["sde_step_indices"].to(sample.device).long() + sigma_view = timesteps.float() + sigma_next = sigmas[torch.clamp(step_indices + 1, max=len(sigmas) - 1)] + + dynamics_type = _normalize_ltx_dynamics_type(extra.get("dynamics_type", "cps")) + if dynamics_type != "cps": + raise NotImplementedError( + f"LTXTrainPipelineConfig.sde_step supports dynamics_type='cps' only " f"(got {dynamics_type!r})." + ) + + prev, log_prob, prev_mean, std_dev_t = sde_step_with_logprob( + None, + noise_pred.float(), + sigma_view, + sample.float(), + prev_sample.float(), + noise_level=noise_level, + sde_type="cps", + sigma=sigma_view, + sigma_prev=sigma_next, + ) + if std_dev_t.ndim > 1: + std_dev_t = std_dev_t.mean(dim=tuple(range(1, std_dev_t.ndim))) + return prev, log_prob, prev_mean, std_dev_t diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 622559ed..189c3c04 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -9,15 +9,22 @@ Each model (QwenImage, SD3, Flux, ...) subclasses TrainPipelineConfig and overrides the relevant methods. + +Trainer lifecycle (load, forward, SDE) lives in TrainStepBackend; see +``get_train_step_backend()``. """ from __future__ import annotations import abc +from typing import TYPE_CHECKING import torch from miles.utils.types import CondKwargs, DiTTrajectory +if TYPE_CHECKING: + from miles.backends.fsdp_utils.train_step_backend import TrainStepBackend + _REGISTRY: dict[str, type[TrainPipelineConfig]] = {} @@ -50,6 +57,14 @@ class TrainPipelineConfig(abc.ABC): lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] needs_timestep_scaling: bool = True optimizer_state_allowed_missing: list[str] = [] + sde_timestep_divisor: float = 1.0 + train_step_backend_cls: type[TrainStepBackend] | None = None + + def get_train_step_backend(self) -> TrainStepBackend: + from miles.backends.fsdp_utils.train_step_backend import DiffusersTrainStepBackend + + cls = type(self).train_step_backend_cls or DiffusersTrainStepBackend + return cls(self) def prepare_trajectory( self, @@ -75,6 +90,17 @@ def prepare_cond_kwargs( ) -> dict: """Convert CondKwargs to model-specific forward() kwargs.""" + def build_train_cond_kwargs( + self, + cond: CondKwargs | None, + *, + latents: torch.Tensor, + args, + device: torch.device, + ) -> dict: + """Build per-sample cond for training; default reuses rollout embeds only.""" + return self.prepare_cond_kwargs(cond, device) + def expand_cond_for_timestep_batch( self, cond_kwargs: dict, diff --git a/miles/backends/fsdp_utils/ltx_geometry.py b/miles/backends/fsdp_utils/ltx_geometry.py new file mode 100644 index 00000000..9aef9c8a --- /dev/null +++ b/miles/backends/fsdp_utils/ltx_geometry.py @@ -0,0 +1,118 @@ +"""Deterministic LTX-2.3 video geometry for miles train forward. + +sglang rollout returns text embeds via ``denoising_env`` only; RoPE coords and +TI2V masks are rebuilt here from the same request-level constants sglang uses +(see ``MILES_ROLLOUT_HANDOFF.md``). +""" + +from __future__ import annotations + +import torch + +# LTX-2.3 defaults (match sglang ``LTXVideoRotaryPositionalEmbeddings`` / VAE). +_LTX_VAE_SPATIAL_COMPRESSION = 32 +_LTX_VAE_TEMPORAL_COMPRESSION = 8 +_LTX_PATCH_SIZE = 1 +_LTX_PATCH_SIZE_T = 1 +_LTX_SCALE_FACTORS = (8, 32, 32) +_LTX_CAUSAL_OFFSET = 1 + + +def latent_grid_shape( + *, + height: int, + width: int, + num_frames: int, +) -> tuple[int, int, int]: + """Return ``(latent_frames, latent_height, latent_width)``.""" + latent_height = height // _LTX_VAE_SPATIAL_COMPRESSION + latent_width = width // _LTX_VAE_SPATIAL_COMPRESSION + latent_frames = (num_frames - 1) // _LTX_VAE_TEMPORAL_COMPRESSION + 1 + return latent_frames, latent_height, latent_width + + +def prepare_ltx_video_positions( + *, + batch_size: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + fps: float, + device: torch.device, + start_frame: int = 0, +) -> torch.Tensor: + """Build video position grid ``[B, 3, T, 2]`` for ltx_core (middle-index RoPE).""" + grid_f = torch.arange( + start=int(start_frame), + end=int(num_latent_frames) + int(start_frame), + step=_LTX_PATCH_SIZE_T, + dtype=torch.float32, + device=device, + ) + grid_h = torch.arange( + start=0, + end=latent_height, + step=_LTX_PATCH_SIZE, + dtype=torch.float32, + device=device, + ) + grid_w = torch.arange( + start=0, + end=latent_width, + step=_LTX_PATCH_SIZE, + dtype=torch.float32, + device=device, + ) + grid = torch.stack(torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij"), dim=0) + + patch_size = (_LTX_PATCH_SIZE_T, _LTX_PATCH_SIZE, _LTX_PATCH_SIZE) + patch_ends = grid + torch.tensor(patch_size, dtype=grid.dtype, device=grid.device).view(3, 1, 1, 1) + latent_coords = torch.stack([grid, patch_ends], dim=-1) + latent_coords = latent_coords.flatten(1, 3).unsqueeze(0).expand(batch_size, -1, -1, -1) + + scale_tensor = torch.tensor(_LTX_SCALE_FACTORS, device=latent_coords.device) + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + _LTX_CAUSAL_OFFSET - _LTX_SCALE_FACTORS[0]).clamp(min=0) + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) + return pixel_coords + + +def build_ltx_t2v_geometry( + *, + batch_size: int, + num_tokens: int, + latent_dim: int, + height: int, + width: int, + num_frames: int, + fps: float, + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + """Pure text-to-video geometry: all tokens denoise, clean latent is zero.""" + latent_frames, latent_height, latent_width = latent_grid_shape(height=height, width=width, num_frames=num_frames) + expected_tokens = latent_frames * latent_height * latent_width + if expected_tokens != num_tokens: + raise ValueError( + f"LTX latent token count mismatch: trajectory has T={num_tokens} but " + f"geometry from {height}x{width}x{num_frames}@{fps}fps expects " + f"T={expected_tokens} ({latent_frames=}, {latent_height=}, {latent_width=})" + ) + + positions = prepare_ltx_video_positions( + batch_size=batch_size, + num_latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + fps=fps, + device=device, + ) + denoise_mask = torch.ones(batch_size, num_tokens, device=device, dtype=torch.float32) + clean_latent = torch.zeros(batch_size, num_tokens, latent_dim, device=device, dtype=dtype) + return { + "positions": positions.to(dtype), + "denoise_mask": denoise_mask, + "clean_latent": clean_latent, + } diff --git a/miles/backends/fsdp_utils/train_step_backend.py b/miles/backends/fsdp_utils/train_step_backend.py new file mode 100644 index 00000000..5a085fd5 --- /dev/null +++ b/miles/backends/fsdp_utils/train_step_backend.py @@ -0,0 +1,365 @@ +"""Training-step backends for diffusion GRPO. + +TrainPipelineConfig holds model-specific cond/trajectory/CFG logic. +TrainStepBackend holds trainer lifecycle + forward + SDE paths so actor +stays model-agnostic without ``if is_ltx`` branches. +""" + +from __future__ import annotations + +import abc +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F +from diffusers import DiffusionPipeline + +from miles.utils.sde_log_prob import sde_step_with_logprob + +if TYPE_CHECKING: + from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig + + +def _pack_cond_for_joint_cfg(pos: dict, neg: dict) -> dict: + out: dict = {} + for key, value in pos.items(): + if isinstance(value, torch.Tensor): + out[key] = torch.cat([value, neg[key]], dim=0) + elif isinstance(value, list): + out[key] = value + neg[key] + else: + out[key] = value + return out + + +class TrainStepBackend(abc.ABC): + """Orchestrates load / forward / SDE for one model family.""" + + supports_cfg_training: bool = True + fsdp_wrap_classes: list[str] | None = None + + def __init__(self, config: TrainPipelineConfig) -> None: + self.config = config + + @abc.abstractmethod + def load_model_and_scheduler( + self, + args, + init_context_factory, + *, + master_dtype: torch.dtype, + ) -> tuple[torch.nn.Module, object]: ... + + def apply_gradient_checkpointing(self, model: torch.nn.Module, args) -> None: + if args.gradient_checkpointing: + model.enable_gradient_checkpointing() + + def get_fsdp_wrap_classes(self) -> list[str] | None: + return self.fsdp_wrap_classes + + def should_use_cfg(self, args) -> bool: + if not self.supports_cfg_training: + return False + guidance_scale = args.diffusion_guidance_scale + true_cfg_scale = args.diffusion_true_cfg_scale + cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale + return cfg_scale > 0 + + def resolve_sigmas_ref( + self, + timesteps_ref: torch.Tensor, + sigmas_snapshot: torch.Tensor | None, + scheduler, + *, + num_train_timesteps: int, + ) -> torch.Tensor: + if sigmas_snapshot is not None: + return sigmas_snapshot.to(timesteps_ref.device).float() + sigmas_ref = timesteps_ref / float(num_train_timesteps) + return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + + def scale_timesteps_for_sde(self, timesteps_flat: torch.Tensor) -> torch.Tensor: + return timesteps_flat / float(self.config.sde_timestep_divisor) + + @abc.abstractmethod + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict, + neg_cond: dict | None, + use_cfg: bool, + guidance_scale: float, + true_cfg_scale: float | None, + fsdp_cfg_batching: bool, + disable_adapter: bool = False, + ) -> torch.Tensor: ... + + @abc.abstractmethod + def sde_step_logprob( + self, + *, + scheduler, + noise_pred: torch.Tensor, + timesteps_for_sde: torch.Tensor, + timesteps_flat: torch.Tensor, + latents_flat: torch.Tensor, + prev_sample: torch.Tensor, + noise_level: float, + grids: dict | None = None, + sample_indices: torch.Tensor | None = None, + tstep_indices: torch.Tensor | None = None, + args=None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (log_prob, prev_sample_mean, std_dev_t).""" + + def append_model_output_compare_stats( + self, + log_stats: dict[str, list[torch.Tensor]], + noise_pred: torch.Tensor, + rollout_mo_flat: torch.Tensor, + ) -> None: + """Optional hook for comparing rollout vs train noise predictions.""" + return + + +class DiffusersTrainStepBackend(TrainStepBackend): + """Default path: diffusers DiT + generic SDE logprob.""" + + def load_model_and_scheduler( + self, + args, + init_context_factory, + *, + master_dtype: torch.dtype, + ) -> tuple[torch.nn.Module, object]: + with init_context_factory(): + pipeline = DiffusionPipeline.from_pretrained( + args.hf_checkpoint, + torch_dtype=master_dtype, + trust_remote_code=True, + text_encoder=None, + vae=None, + tokenizer=None, + ) + model = pipeline.transformer + scheduler = pipeline.scheduler + del pipeline + return model, scheduler + + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict, + neg_cond: dict | None, + use_cfg: bool, + guidance_scale: float, + true_cfg_scale: float | None, + fsdp_cfg_batching: bool, + disable_adapter: bool = False, + ) -> torch.Tensor: + def _forward(cond: dict) -> torch.Tensor: + return model( + hidden_states=latents_input, + timestep=timesteps_input, + return_dict=False, + **cond, + )[0] + + adapter_ctx = model.disable_adapter() if disable_adapter else nullcontext() + with adapter_ctx: + if not use_cfg: + return _forward(pos_cond) + if fsdp_cfg_batching: + joint_cond = _pack_cond_for_joint_cfg(pos_cond, neg_cond) + joint_out = model( + hidden_states=torch.cat([latents_input, latents_input], dim=0), + timestep=torch.cat([timesteps_input, timesteps_input], dim=0), + return_dict=False, + **joint_cond, + )[0] + noise_pred_pos, noise_pred_neg = joint_out.chunk(2, dim=0) + else: + noise_pred_pos = _forward(pos_cond) + noise_pred_neg = _forward(neg_cond) + return self.config.cfg_combine( + noise_pred_pos, + noise_pred_neg, + guidance_scale, + true_cfg_scale=true_cfg_scale, + ) + + def sde_step_logprob( + self, + *, + scheduler, + noise_pred: torch.Tensor, + timesteps_for_sde: torch.Tensor, + timesteps_flat: torch.Tensor, + latents_flat: torch.Tensor, + prev_sample: torch.Tensor, + noise_level: float, + grids: dict | None = None, + sample_indices: torch.Tensor | None = None, + tstep_indices: torch.Tensor | None = None, + args=None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + _, log_prob, prev_mean, std_dev_t = sde_step_with_logprob( + scheduler, + noise_pred.float(), + timesteps_flat, + latents_flat.float(), + prev_sample=prev_sample.float(), + noise_level=noise_level, + ) + return log_prob, prev_mean, std_dev_t + + +class LTXTrainStepBackend(TrainStepBackend): + """LTX-2.3: custom loader, velocity forward, CPS SDE.""" + + supports_cfg_training = False + fsdp_wrap_classes = ["BasicAVTransformerBlock"] + + def load_model_and_scheduler( + self, + args, + init_context_factory, + *, + master_dtype: torch.dtype | None = None, + ) -> tuple[torch.nn.Module, object]: + from dataclasses import dataclass, field + + from ltx_core.components.schedulers import LTX2Scheduler + + from miles.backends.fsdp_utils.configs.ltx import ( + load_ltx_transformer_for_train, + resolve_transformer_checkpoint, + ) + + @dataclass + class _LTXSchedulerHolder: + sigmas: torch.Tensor = field(default_factory=lambda: torch.tensor([])) + timesteps: torch.Tensor = field(default_factory=lambda: torch.tensor([])) + num_inference_steps: int = 0 + _step_index: int | None = None + _begin_index: int | None = None + + def to(self, device): + self.sigmas = self.sigmas.to(device) + self.timesteps = self.timesteps.to(device) + return self + + master_dtype_name = getattr(args, "fsdp_master_dtype", "bf16") + resolved_dtype = ( + master_dtype + or { + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + }[master_dtype_name] + ) + + checkpoint = resolve_transformer_checkpoint( + args.diffusion_model, + explicit_path=getattr(args, "sglang_transformer_weights_path", None), + ) + model = load_ltx_transformer_for_train(checkpoint, device="cpu", dtype=resolved_dtype) + + num_steps = int(getattr(args, "diffusion_num_steps", 24)) + ltx_sched = LTX2Scheduler() + sigmas = ltx_sched.execute(steps=num_steps).float() + scheduler = _LTXSchedulerHolder( + sigmas=sigmas, + timesteps=sigmas[:num_steps], + num_inference_steps=num_steps, + ) + + if getattr(args, "gradient_checkpointing", False): + if hasattr(model, "set_gradient_checkpointing"): + model.set_gradient_checkpointing(True) + elif hasattr(model, "enable_gradient_checkpointing"): + model.enable_gradient_checkpointing() + + return model, scheduler + + def apply_gradient_checkpointing(self, model: torch.nn.Module, args) -> None: + # Applied inside load_model_and_scheduler for LTX. + pass + + def resolve_sigmas_ref( + self, + timesteps_ref: torch.Tensor, + sigmas_snapshot: torch.Tensor | None, + scheduler, + *, + num_train_timesteps: int = 1000, + ) -> torch.Tensor: + device = timesteps_ref.device + if sigmas_snapshot is not None: + return sigmas_snapshot.to(device).float() + sigmas_ref = timesteps_ref / float(self.config.sde_timestep_divisor) + return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict, + neg_cond: dict | None, + use_cfg: bool, + guidance_scale: float, + true_cfg_scale: float | None, + fsdp_cfg_batching: bool, + disable_adapter: bool = False, + ) -> torch.Tensor: + adapter_ctx = model.disable_adapter() if disable_adapter else nullcontext() + with adapter_ctx: + return self.config.forward_velocity(model, latents_input, timesteps_input, pos_cond) + + def sde_step_logprob( + self, + *, + scheduler, + noise_pred: torch.Tensor, + timesteps_for_sde: torch.Tensor, + timesteps_flat: torch.Tensor, + latents_flat: torch.Tensor, + prev_sample: torch.Tensor, + noise_level: float, + grids: dict | None = None, + sample_indices: torch.Tensor | None = None, + tstep_indices: torch.Tensor | None = None, + args=None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sde_extra = self.config.build_sde_extra(scheduler, grids, sample_indices, tstep_indices, args) + _, log_prob, prev_mean, std_dev_t = self.config.sde_step( + scheduler, + noise_pred, + timesteps_for_sde, + latents_flat, + prev_sample=prev_sample, + noise_level=noise_level, + extra=sde_extra, + ) + return log_prob, prev_mean, std_dev_t + + def append_model_output_compare_stats( + self, + log_stats: dict[str, list[torch.Tensor]], + noise_pred: torch.Tensor, + rollout_mo_flat: torch.Tensor, + ) -> None: + flat_train = noise_pred.float().reshape(noise_pred.shape[0], -1) + flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) + log_stats["model_output_cosine_sim"].append( + F.cosine_similarity(flat_train, flat_rollout, dim=1).mean().detach() + ) diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py new file mode 100644 index 00000000..da0fed31 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -0,0 +1,52 @@ +"""LTX-2 sglang-d rollout engine config (re-exports train-side model family helpers).""" + +from __future__ import annotations + +from miles.backends.fsdp_utils.configs.ltx import ( + LTX_DEFAULT_HF_MODEL, + LTX_DEFAULT_MODEL_ID, + ensure_materialized_model, + is_ltx_model, + resolve_hf_model_id, + resolve_materialized_model_dir, + resolve_model_id, + resolve_transformer_checkpoint, + server_kwargs_extras, +) + +__all__ = [ + "LTX_DEFAULT_HF_MODEL", + "LTX_DEFAULT_MODEL_ID", + "ensure_materialized_model", + "is_ltx_model", + "resolve_hf_model_id", + "resolve_materialized_model_dir", + "resolve_model_id", + "resolve_transformer_checkpoint", + "resolve_sglang_model_path", + "resolve_ltx_model_id", + "resolve_ltx_transformer_weights_path", + "server_kwargs_extras", +] + + +def resolve_sglang_model_path(args) -> str: + return resolve_hf_model_id(args) + + +def resolve_ltx_model_id(args) -> str: + return resolve_model_id(args) + + +def resolve_ltx_transformer_weights_path( + diffusion_model: str | None, + *, + explicit_path: str | None = None, +) -> str | None: + try: + return resolve_transformer_checkpoint( + diffusion_model, + explicit_path=explicit_path, + ) + except FileNotFoundError: + return None diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 42278690..29b4e2f9 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -1,29 +1,70 @@ -"""sgl-d → diffusers numerical-parity monkey patches. +"""sgl-d numerical-parity monkey patches for miles training alignment. -Patches sgl-d's generic op classes (RMSNorm, LayerNormScaleShift, MulAdd, -USPAttention, etc.) to match diffusers' bf16 cast/op order. Apply once at -sglang-d scheduler startup so DiT forwards on the rollout side agree with -diffusers-style training-side forwards down to bf16 ULPs. +Rollout engines select a patch group via ``resolve_rollout_patch_group(args)``; +the scheduler child reads ``MILES_ROLLOUT_PATCH_GROUP`` and calls +``apply_rollout_patch_group``. -Patches are at the op layer, not the model layer — they apply to every sgl-d -DiT that uses these generic classes. Adding alignment for a new op = drop a -new ``patch_.py`` file and add it to ``apply_sgld_monkey_patches``. +- ``sgld``: diffusers / SD3 op parity (RMSNorm, RoPE, attention, …). +- ``ltx``: LTX rollout cond kwargs + AV cross-off (video-only train parity). + +Patch modules are imported inside ``apply_*`` only so ``RolloutManager`` (a +CPU-only Ray actor) can import this package without pulling sglang triton kernels. """ -from miles.backends.sglang_diffusion_utils.monkey_patches import ( - patch_layernorm_scale_shift, - patch_mul_add, - patch_qk_norm_rope, - patch_rmsnorm, - patch_scale_residual_layernorm, - patch_usp_attention, -) +from __future__ import annotations + +ROLLOUT_PATCH_GROUP_ENV = "MILES_ROLLOUT_PATCH_GROUP" +PATCH_GROUP_SGLD = "sgld" +PATCH_GROUP_LTX = "ltx" + +# Propagated into Ray rollout workers (see miles/ray/rollout.py). +LTX_ROLLOUT_PATCHES_ENV = "MILES_APPLY_LTX_ROLLOUT_PATCHES" + + +def resolve_rollout_patch_group(args) -> str | None: + """Return the rollout patch group for this engine, or None.""" + if getattr(args, "apply_sgld_monkey_patches", False): + return PATCH_GROUP_SGLD + + from miles.backends.sglang_diffusion_utils.configs.ltx import is_ltx_model + + if is_ltx_model(args): + return PATCH_GROUP_LTX + + return None + + +def apply_rollout_patch_group(group: str | None) -> None: + if group == PATCH_GROUP_SGLD: + apply_sgld_monkey_patches() + elif group == PATCH_GROUP_LTX: + apply_ltx2_rollout_patches() def apply_sgld_monkey_patches() -> None: + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + patch_layernorm_scale_shift, + patch_mul_add, + patch_qk_norm_rope, + patch_rmsnorm, + patch_scale_residual_layernorm, + patch_usp_attention, + ) + patch_rmsnorm.apply() patch_layernorm_scale_shift.apply() patch_scale_residual_layernorm.apply() patch_mul_add.apply() patch_usp_attention.apply() patch_qk_norm_rope.apply() + + +def apply_ltx2_rollout_patches() -> None: + """LTX rollout: cond kwargs + disable AV cross-attn (video-only train parity).""" + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + patch_ltx2_disable_av_cross, + patch_ltx2_rollout_cond_kwargs, + ) + + patch_ltx2_rollout_cond_kwargs.apply() + patch_ltx2_disable_av_cross.apply() diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_disable_av_cross.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_disable_av_cross.py new file mode 100644 index 00000000..2cb073d5 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_disable_av_cross.py @@ -0,0 +1,29 @@ +"""Disable LTX A2V/V2A cross-attention in sglang rollout (video-only parity). + +Miles ltx_core train forward is video-only; the sglang LTX rollout otherwise +runs audio-video cross-attention. This injects the disable flags into the DiT +forward so the rollout video branch matches train. +""" + +from __future__ import annotations + +_APPLIED = False + + +def apply() -> None: + global _APPLIED + if _APPLIED: + return + + from sglang.multimodal_gen.runtime.models.dits import ltx_2 as ltx2_mod + + model_cls = ltx2_mod.LTX2VideoTransformer3DModel + orig_forward = model_cls.forward + + def forward(self, *args, **kwargs): + kwargs["disable_a2v_cross_attn"] = True + kwargs["disable_v2a_cross_attn"] = True + return orig_forward(self, *args, **kwargs) + + model_cls.forward = forward + _APPLIED = True diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_rollout_cond_kwargs.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_rollout_cond_kwargs.py new file mode 100644 index 00000000..99265003 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_rollout_cond_kwargs.py @@ -0,0 +1,63 @@ +"""Ensure LTX rollout denoising_env carries text/audio embeds for miles train replay. + +TODO(upstream): remove once sgl-d LTX rollout returns full cond kwargs in the +standard denoising_env schema without miles-side postprocessing. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +_APPLIED = False + + +def _first_batch_tensor(batch: Any, attr: str) -> Any | None: + value = getattr(batch, attr, None) + if value is None: + return None + return value[0] if isinstance(value, list) else value + + +def apply() -> None: + global _APPLIED + if _APPLIED: + return + + from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import LTX2DenoisingStage + + if not hasattr(LTX2DenoisingStage, "_prepare_denoising_loop"): + logger.warning( + "LTX2DenoisingStage._prepare_denoising_loop is missing; " + "rollout denoising_env may lack text/audio cond kwargs." + ) + _APPLIED = True + return + + orig_prepare = LTX2DenoisingStage._prepare_denoising_loop + + def _prepare_denoising_loop(self, batch, server_args): + ctx = orig_prepare(self, batch, server_args) + if not (batch.rollout and batch.rollout_return_denoising_env): + return ctx + ctx.pos_cond_kwargs = dict(ctx.pos_cond_kwargs) + if ctx.pos_cond_kwargs.get("encoder_hidden_states") is None: + embeds = _first_batch_tensor(batch, "prompt_embeds") + if embeds is not None: + ctx.pos_cond_kwargs["encoder_hidden_states"] = embeds + if ctx.pos_cond_kwargs.get("audio_encoder_hidden_states") is None: + audio_embeds = _first_batch_tensor(batch, "audio_prompt_embeds") + if audio_embeds is not None: + ctx.pos_cond_kwargs["audio_encoder_hidden_states"] = audio_embeds + attn_mask = getattr(batch, "prompt_attention_mask", None) + if attn_mask is not None: + if ctx.pos_cond_kwargs.get("encoder_attention_mask") is None: + ctx.pos_cond_kwargs["encoder_attention_mask"] = attn_mask + if ctx.pos_cond_kwargs.get("audio_encoder_attention_mask") is None: + ctx.pos_cond_kwargs["audio_encoder_attention_mask"] = attn_mask + return ctx + + LTX2DenoisingStage._prepare_denoising_loop = _prepare_denoising_loop + _APPLIED = True diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index c16e2f3f..217e4881 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -9,12 +9,35 @@ from sglang.multimodal_gen.runtime.launch_server import kill_process_tree from sglang.multimodal_gen.runtime.server_args import ServerArgs +from miles.backends.sglang_diffusion_utils.configs import ltx as ltx_config from miles.ray.ray_actor import RayActor from miles.utils.http_utils import get_host_info logger = logging.getLogger(__name__) +def build_rollout_engine_env_vars(args) -> dict[str, str]: + """Env vars forwarded to Ray-spawned sglang-diffusion rollout engine workers.""" + from miles.backends.fsdp_utils.configs.ltx import patch_rollout_engine_env_vars + from miles.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST + + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { + "SGL_JIT_DEEPGEMM_PRECOMPILE": "false", + "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false", + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true", + "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true", + "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", + "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false", + } + for passthrough in ("PYTHONPATH", "SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): + if os.environ.get(passthrough): + env_vars[passthrough] = os.environ[passthrough] + patch_rollout_engine_env_vars(env_vars, args) + return env_vars + + def _to_local_gpu_id(physical_gpu_id: int) -> int: cvd = os.environ.get("CUDA_VISIBLE_DEVICES") if not cvd: @@ -39,15 +62,23 @@ def _scheduler_process_with_sgld_monkey_patches(*args, **kwargs): # any monkey patches done in the middle child are gone. Apply them HERE, # before calling the real run_scheduler_process, so the DiT that's # constructed inside the grandchild sees the patched classes. - from miles.backends.sglang_diffusion_utils.monkey_patches import apply_sgld_monkey_patches + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + LTX_ROLLOUT_PATCHES_ENV, + apply_ltx2_rollout_patches, + apply_sgld_monkey_patches, + ) + + if os.environ.get("MILES_APPLY_SGLD_MONKEY_PATCHES") == "1": + apply_sgld_monkey_patches() + if os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "0") == "1": + apply_ltx2_rollout_patches() - apply_sgld_monkey_patches() from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process return run_scheduler_process(*args, **kwargs) -def _launch_server_target(server_args, apply_sgld_monkey_patches: bool = False): +def _launch_server_target(server_args, apply_rollout_patches: bool = False): # addict.Dict used by SGL-D loses its `__frozen` instance attribute across spawn pickle. # Reconstruct a fresh one from the unpickled (broken) instance import addict @@ -55,7 +86,7 @@ def _launch_server_target(server_args, apply_sgld_monkey_patches: bool = False): if server_args.attention_backend_config is not None: server_args.attention_backend_config = addict.Dict(server_args.attention_backend_config) - if apply_sgld_monkey_patches: + if apply_rollout_patches: # launch_server spawns its scheduler via mp.Process(target=run_scheduler_process). # Under spawn, target is pickled by qualname and re-imported in the grandchild, # so patching in THIS process doesn't help. Instead, rebind the name inside @@ -73,14 +104,14 @@ def _launch_server_target(server_args, apply_sgld_monkey_patches: bool = False): def launch_server_process( server_args: ServerArgs, - apply_sgld_monkey_patches: bool = False, + apply_rollout_patches: bool = False, ) -> multiprocessing.Process: # use spawn to avoid potential risks of fork in terms of subthreads or CUDA. multiprocessing.set_start_method("spawn", force=True) server_args.host = server_args.host.strip("[]") p = multiprocessing.Process( target=_launch_server_target, - args=(server_args, apply_sgld_monkey_patches), + args=(server_args, apply_rollout_patches), ) p.start() @@ -157,12 +188,20 @@ def _format_v6_uri(addr): def _init_normal(self, server_args_dict): logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}") self._pin_to_assigned_gpu() - apply_sgld_monkey_patches = self.args.apply_sgld_monkey_patches - if apply_sgld_monkey_patches: + from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV + + apply_sgld = bool(getattr(self.args, "apply_sgld_monkey_patches", False)) + apply_ltx = ltx_config.is_ltx_model(self.args) + use_rollout_patches = apply_sgld or apply_ltx + if apply_sgld: + os.environ["MILES_APPLY_SGLD_MONKEY_PATCHES"] = "1" logger.info("Launching sglang-d with sgl-d → diffusers monkey patches " "(--apply-sgld-monkey-patches)") + if apply_ltx: + os.environ[LTX_ROLLOUT_PATCHES_ENV] = "1" + logger.info("Launching sglang-d with LTX rollout monkey patches") self.process = launch_server_process( ServerArgs.from_kwargs(**server_args_dict), - apply_sgld_monkey_patches=apply_sgld_monkey_patches, + apply_rollout_patches=use_rollout_patches, ) if self.node_rank == 0 and self.router_ip and self.router_port: @@ -323,6 +362,11 @@ def _compute_server_args(args, host, port, nccl_port): "warmup": False, } + # LTX rollout: HF model id + overlay wrapper; optional DiT override via safetensors path. + if ltx_config.is_ltx_model(args): + kwargs["model_path"] = ltx_config.resolve_hf_model_id(args) + kwargs.update(ltx_config.server_kwargs_extras(args)) + # Forward every `args.sglang_` the user set via --sglang-* CLI for # ServerArgs fields not already hardcoded above. Picks up ulysses_degree / # ring_degree / dp_size / etc. without listing each one. diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 1516864c..9a0e54a1 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -6,13 +6,15 @@ from pathlib import Path from typing import Any -import numpy as np import ray import torch from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS -from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine +from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import ( + SGLangDiffusionEngine, + build_rollout_engine_env_vars, +) from miles.rollout.base_types import call_rollout_fn from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor @@ -26,7 +28,7 @@ from miles.utils.tracking_utils import init_tracking from miles.utils.types import Sample -from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock +from .utils import Lock logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) @@ -154,7 +156,6 @@ def generate(self, rollout_id): def eval(self, rollout_id): if self.args.debug_train_only: - # if debug train only, we don't generate evaluation data return self.health_monitoring_resume() @@ -380,7 +381,7 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "sample_indices": [sample.index for sample in samples], "prompt": [sample.prompt for sample in samples], # Per-sample training step indices (flow_grpo sde-window). None = train every step. - "sde_step_indices": [(sample.train_metadata or {}).get("sde_step_indices") for sample in samples], + "sde_step_indices": [sample.get_sde_step_indices() for sample in samples], } return train_data @@ -403,18 +404,18 @@ def _log_images( own namespace at least groups them in one UI section. """ import wandb + from miles.rollout.rm_hub.pickscore import first_frame_for_wandb log_dict: dict = {} for media_key, samples in media_key_to_samples.items(): images = [] for s in samples[:max_images]: t = s.generated_output - if t is None or t.ndim != 4: + if t is None: + continue + frame = first_frame_for_wandb(t) + if frame is None: continue - frame = t[:, 0, :, :].float().cpu().numpy().transpose(1, 2, 0) - if frame.max() <= 1.0 + 1e-3: - frame = frame * 255.0 - frame = np.clip(frame, 0, 255).astype(np.uint8) reward = s.reward if not reward_key else (s.reward or {}).get(reward_key) images.append(wandb.Image(frame, caption=f"{str(s.prompt)[:160]} | reward={reward}")) if images: @@ -479,16 +480,7 @@ def init_rollout_engines(args, pg, all_rollout_engines): placement_group_bundle_index=reordered_bundle_indices[i * num_gpu_per_engine], ) - env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { - "SGL_JIT_DEEPGEMM_PRECOMPILE": "false", - "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false", - "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", - "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", - "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true", - "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true", - "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", - "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false", - } + env_vars = build_rollout_engine_env_vars(args) rollout_engine = RolloutRayActor.options( num_cpus=num_cpus, diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index 76ddb122..872fde39 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -15,26 +15,100 @@ logger = logging.getLogger(__name__) +def sample_frame_indices(num_total_frames: int, num_frames: int) -> list[int]: + if num_total_frames <= 0: + raise ValueError(f"video has no frames: {num_total_frames}") + if num_total_frames <= num_frames: + return list(range(num_total_frames)) + if num_frames == 1: + return [num_total_frames // 2] + step = (num_total_frames - 1) / (num_frames - 1) + return [int(round(i * step)) for i in range(num_frames)] + + +def generated_output_to_fchw(t: torch.Tensor) -> torch.Tensor: + """Return ``[F, C, H, W]`` float tensor in ``[0, 1]``.""" + t = t.detach().cpu().float() + if t.ndim == 3: + if t.shape[0] not in (1, 3): + raise ValueError(f"expected [C, H, W] with C in {{1, 3}}, got {tuple(t.shape)}") + t = t.unsqueeze(0) + elif t.ndim == 4: + if t.shape[-1] in (1, 3): + t = t.permute(0, 3, 1, 2) + elif t.shape[0] in (1, 3): + t = t.permute(1, 0, 2, 3) + elif t.shape[1] not in (1, 3): + raise ValueError(f"unrecognized 4D video layout: {tuple(t.shape)}") + elif t.ndim == 5: + if t.shape[0] == 1 and t.shape[-1] in (1, 3): + t = t[0].permute(0, 3, 1, 2) + else: + raise ValueError(f"unrecognized 5D video layout: {tuple(t.shape)}") + else: + raise ValueError(f"generated_output must be 3D–5D, got {tuple(t.shape)}") + + if float(t.max()) > 1.0 + 1e-3: + t = t / 255.0 + return t.clamp(0.0, 1.0) + + +def fchw_frame_to_hwc_uint8(frame_chw: torch.Tensor) -> np.ndarray: + hwc = frame_chw.numpy().transpose(1, 2, 0) + if float(hwc.max()) <= 1.0 + 1e-3: + hwc = hwc * 255.0 + return np.ascontiguousarray(hwc.clip(0, 255).astype(np.uint8)) + + +def first_frame_for_wandb(t: torch.Tensor) -> np.ndarray | None: + """First frame as HWC uint8 for wandb logging; None if layout is unsupported.""" + try: + return fchw_frame_to_hwc_uint8(generated_output_to_fchw(t)[0]) + except (ValueError, TypeError): + if t.ndim != 4: + return None + frame = t[:, 0, :, :].float().cpu().numpy().transpose(1, 2, 0) + if float(frame.max()) <= 1.0 + 1e-3: + frame = frame * 255.0 + return np.clip(frame, 0, 255).astype(np.uint8) + + +def fchw_to_pil_frames(video_fchw: torch.Tensor, frame_indices: Sequence[int]) -> list[Image.Image]: + return [Image.fromarray(fchw_frame_to_hwc_uint8(video_fchw[idx])) for idx in frame_indices] + + +def is_video_generated_output(t: torch.Tensor) -> bool: + """True when output carries multiple temporal frames (LTX / sglang video).""" + fchw = generated_output_to_fchw(t) + return fchw.shape[0] > 1 + + def _feature_tensor(features): if isinstance(features, torch.Tensor): return features - return features.pooler_output + if hasattr(features, "pooler_output"): + pooled = features.pooler_output + if isinstance(pooled, torch.Tensor): + return pooled + for attr in ("image_embeds", "text_embeds"): + value = getattr(features, attr, None) + if isinstance(value, torch.Tensor): + return value + if isinstance(features, tuple): + for item in reversed(features): + if isinstance(item, torch.Tensor) and item.ndim == 2: + return item + raise TypeError(f"No 2-D tensor in model output tuple (len={len(features)})") + raise TypeError(f"Cannot extract embedding tensor from {type(features)!r}") def _sample_to_rgb_hwc_uint8(sample: Sample) -> np.ndarray: - frame_chw = sample.generated_output.detach().cpu()[:, 0, :, :] - hwc = frame_chw.float().numpy().transpose(1, 2, 0) - if float(hwc.max()) <= 1.0 + 1e-3: - hwc = hwc * 255.0 - return np.ascontiguousarray(hwc.clip(0, 255).astype(np.uint8)) + fchw = generated_output_to_fchw(sample.generated_output) + return fchw_frame_to_hwc_uint8(fchw[fchw.shape[0] // 2]) class PickScoreScorer(torch.nn.Module): - """Small local copy of Flow-GRPO's PickScore scorer. - - The scorer consumes final PIL images and prompt strings, then returns one - scalar reward per prompt/image pair. - """ + """PickScore for static images (SD3 / single-frame outputs).""" def __init__( self, @@ -80,6 +154,85 @@ def forward(self, prompts: Sequence[str], images: Sequence[Image.Image]) -> list return [float(score) for score in scores.detach().cpu()] +class VideoPickScoreScorer(torch.nn.Module): + """Multi-frame PickScore for LTX video (matches trainer-rollout / verl-omni).""" + + def __init__( + self, + *, + device: str = "cuda", + processor_path: str, + model_path: str, + dtype: torch.dtype = torch.float16, + ) -> None: + super().__init__() + from transformers import AutoModel, AutoProcessor + + self.device = torch.device(device) + self.dtype = dtype + self.processor = AutoProcessor.from_pretrained(processor_path) + self.model = AutoModel.from_pretrained(model_path).eval().to(device=self.device, dtype=dtype) + + @torch.no_grad() + def score_videos( + self, + videos_fchw: Sequence[torch.Tensor], + prompts: Sequence[str], + *, + num_frames: int, + batch_size: int, + ) -> list[float]: + if len(videos_fchw) != len(prompts): + raise ValueError(f"#videos ({len(videos_fchw)}) != #prompts ({len(prompts)})") + + flat_images: list[Image.Image] = [] + flat_prompts: list[str] = [] + per_sample_counts: list[int] = [] + + for video_fchw, prompt in zip(videos_fchw, prompts, strict=True): + frame_indices = sample_frame_indices(video_fchw.shape[0], num_frames) + per_sample_counts.append(len(frame_indices)) + flat_images.extend(fchw_to_pil_frames(video_fchw, frame_indices)) + flat_prompts.extend([prompt] * len(frame_indices)) + + logit_scale = self.model.logit_scale.exp() + flat_scores: list[torch.Tensor] = [] + for start in range(0, len(flat_images), batch_size): + image_chunk = flat_images[start : start + batch_size] + prompt_chunk = flat_prompts[start : start + batch_size] + + image_inputs = self.processor(images=image_chunk, return_tensors="pt", padding=True) + image_inputs = {k: v.to(device=self.device) for k, v in image_inputs.items()} + if "pixel_values" in image_inputs: + image_inputs["pixel_values"] = image_inputs["pixel_values"].to(self.dtype) + + text_inputs = self.processor( + text=prompt_chunk, + return_tensors="pt", + padding=True, + truncation=True, + max_length=77, + ) + text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()} + + image_embs = _feature_tensor(self.model.get_image_features(**image_inputs)) + image_embs = image_embs / image_embs.norm(p=2, dim=-1, keepdim=True).clamp_min(1e-12) + + text_embs = _feature_tensor(self.model.get_text_features(**text_inputs)) + text_embs = text_embs / text_embs.norm(p=2, dim=-1, keepdim=True).clamp_min(1e-12) + + chunk_scores = logit_scale * (text_embs * image_embs).sum(dim=-1) + flat_scores.append(chunk_scores.float()) + + all_scores = torch.cat(flat_scores, dim=0) + rewards: list[float] = [] + offset = 0 + for count in per_sample_counts: + rewards.append(float(all_scores[offset : offset + count].mean())) + offset += count + return rewards + + @ray.remote class PickScoreRewardActor: def __init__( @@ -93,15 +246,37 @@ def __init__( if use_cuda: torch.cuda.set_device(0) device = "cuda" if use_cuda else "cpu" - self.scorer = PickScoreScorer( + self.device = device + self.image_scorer = PickScoreScorer( + device=device, + processor_path=processor_path, + model_path=model_path, + ) + self.video_scorer = VideoPickScoreScorer( device=device, processor_path=processor_path, model_path=model_path, + dtype=torch.float16 if use_cuda else torch.float32, ) def score_batch(self, images: list[np.ndarray], prompts: list[str]) -> list[float]: pil_images = [Image.fromarray(image) for image in images] - return self.scorer(prompts, pil_images) + return self.image_scorer(prompts, pil_images) + + def score_videos_batch( + self, + videos_fchw: list[torch.Tensor], + prompts: list[str], + *, + num_frames: int, + batch_size: int, + ) -> list[float]: + return self.video_scorer.score_videos( + videos_fchw, + prompts, + num_frames=num_frames, + batch_size=batch_size, + ) class AsyncPickScorePool(metaclass=SingletonMeta): @@ -111,6 +286,7 @@ def __init__(self, args) -> None: num_workers = args.pickscore_num_workers num_gpus_per_worker = args.pickscore_num_gpus_per_worker self._batch_size = args.pickscore_batch_size + self._num_frames = int(getattr(args, "pickscore_num_frames", 3) or 3) self._actors = [ PickScoreRewardActor.options( num_cpus=1, @@ -145,9 +321,32 @@ async def score(self, images: list[np.ndarray], prompts: list[str]) -> list[floa chunked_scores = await loop.run_in_executor(None, ray.get, refs) return [float(score) for chunk in chunked_scores for score in chunk] + async def score_videos( + self, + videos_fchw: list[torch.Tensor], + prompts: list[str], + ) -> list[float]: + actor = self._next_actor() + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + ray.get, + actor.score_videos_batch.remote( + videos_fchw, + prompts, + num_frames=self._num_frames, + batch_size=self._batch_size, + ), + ) + async def pickscore_rm(args, samples: Sequence[Sample]) -> list[float]: pool = AsyncPickScorePool(args) + if any(is_video_generated_output(sample.generated_output) for sample in samples): + videos = [generated_output_to_fchw(sample.generated_output) for sample in samples] + prompts = [sample.prompt for sample in samples] + return await pool.score_videos(videos, prompts) + images = [_sample_to_rgb_hwc_uint8(sample) for sample in samples] prompts = [sample.prompt for sample in samples] return await pool.score(images, prompts) diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 6ad11abf..63398bb2 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -62,6 +62,13 @@ def build_rollout_sampling_params( } ) + # LTX dynamics must run after the generic rollout block so CPS / log_prob_no_const + # are not overwritten by SD3 defaults (rollout_sde_type="sde"). + from miles.backends.fsdp_utils.configs.ltx import is_ltx_model, patch_rollout_sampling_params + + if is_ltx_model(args): + patch_rollout_sampling_params(sampling_params, args, evaluation=evaluation) + if extra_sampling_params: sampling_params["extra_sampling_params"] = extra_sampling_params @@ -76,8 +83,8 @@ def build_rollout_generate_payload( ) -> dict[str, Any]: """Build full JSON payload for ``POST /rollout/generate`` (``RolloutImageRequest``).""" sampling_params["prompt"] = prompt - if sampling_params["negative_prompt"] is None: - sampling_params["negative_prompt"] = " " # FlowGRPO default + if sampling_params.get("negative_prompt") is None and float(sampling_params.get("guidance_scale", 1.0)) != 1.0: + sampling_params["negative_prompt"] = " " # FlowGRPO default when CFG is on sampling_params["num_outputs_per_prompt"] = num_outputs_per_prompt return sampling_params @@ -102,6 +109,7 @@ def __init__(self, args: Namespace) -> None: scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=self.node_id, soft=False) ).remote() + self.rollout_id = 0 self.reset() @contextmanager @@ -120,6 +128,7 @@ def reset(self) -> None: self.remaining_batch_size = 0 self.pendings = set() self.aborted = False + self.rollout_id = 0 def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: for group in samples: @@ -153,6 +162,7 @@ async def generate_microgroup( microgroup[0], int(sampling_params["num_inference_steps"]), int(sampling_params["seed"]), + rollout_id=int(getattr(state, "rollout_id", 0) or 0), ) sampling_params["rollout_sde_step_indices"] = sde_indices sampling_params["rollout_return_step_indices"] = return_indices @@ -275,6 +285,7 @@ async def generate_rollout_async( assert args.rollout_global_dataset state = GenerateState(args) + state.rollout_id = int(rollout_id) # instantiate data filters dynamic_filter = ( @@ -437,7 +448,7 @@ async def eval_rollout_single_dataset( reward_key = args.eval_reward_key return { dataset_config.name: { - "rewards": [sample.reward if not reward_key else sample.reward[reward_key] for sample in data], + "rewards": [sample.get_reward_value(args, reward_key=reward_key) for sample in data], "samples": data, } } diff --git a/miles/rollout/step_strategy_hub.py b/miles/rollout/step_strategy_hub.py index f638ae45..cf6846c3 100644 --- a/miles/rollout/step_strategy_hub.py +++ b/miles/rollout/step_strategy_hub.py @@ -3,6 +3,10 @@ Each function has signature ``(args, sample, num_steps, seed) -> (sde, ret)`` where ``sde`` and ``ret`` are ``list[int] | None`` (``None`` = all steps). +Strategies that must match trainer-rollout (``ltx_sde_candidates``) accept +``rollout_id`` as a keyword argument — see ``miles.rollout.sglang_diffusion_rollout``. +All strategies in this hub should accept ``*, rollout_id=0`` for a uniform call site. + Point ``--diffusion-step-strategy-path`` at any such function. """ @@ -11,20 +15,65 @@ from argparse import Namespace import numpy as np +import torch from miles.utils.types import Sample +def _normalize_sde_step_candidates(candidates, num_steps: int) -> list[int] | None: + if candidates is None or candidates == "": + return None + if isinstance(candidates, str): + candidates = [int(x.strip()) for x in candidates.split(",") if x.strip()] + else: + candidates = [int(x) for x in candidates] + invalid = [step for step in candidates if step < 0 or step >= num_steps] + if invalid: + raise ValueError(f"sde_step_candidates must be in [0, {num_steps}), got {invalid}") + return list(dict.fromkeys(candidates)) + + +def ltx_sde_candidates( + args: Namespace, + sample: Sample, + num_steps: int, + seed: int, + *, + rollout_id: int = 0, +) -> tuple[list[int] | None, list[int] | None]: + """verl-omni / trainer-rollout SDE step pick: ``--ltx-num-sde-steps`` random + draws from ``--ltx-sde-step-candidates``, keyed by ``rollout_seed + rollout_id``. + + Uses ``torch.randperm`` (not numpy) so the chosen indices match + ``miles.rollout.ltx_rollout._select_sde_step_set`` bit-for-bit. + + Non-candidate steps run as deterministic Euler in sglang; only listed indices + inject SDE noise and contribute log_probs — same as trainer-rollout. + """ + del sample, seed # trainer keys off rollout_id, not per-sample generation seed + candidates = _normalize_sde_step_candidates(getattr(args, "ltx_sde_step_candidates", None), num_steps) + if candidates is None: + raise ValueError("ltx_sde_candidates requires --ltx-sde-step-candidates " "(e.g. 0,1,2,3,4,5,6,7,8,9)") + num_sde = int(getattr(args, "ltx_num_sde_steps", 0) or len(candidates)) + num_sde = min(max(num_sde, 1), len(candidates)) + rng_seed = int(getattr(args, "rollout_seed", 42)) + int(rollout_id) + g = torch.Generator().manual_seed(rng_seed) + selected = torch.randperm(len(candidates), generator=g)[:num_sde] + indices = [candidates[i] for i in selected.tolist()] + return indices, None + + def sde_window( - args: Namespace, sample: Sample, num_steps: int, seed: int + args: Namespace, sample: Sample, num_steps: int, seed: int, *, rollout_id: int = 0 ) -> tuple[list[int] | None, list[int] | None]: """flow_grpo-style random contiguous SDE window. Returns (sde=window, return=None) so sglang-d returns the full trajectory and log_probs; the trainer then slices to the window for loss / backprop. Keeping the full trajectory avoids the sglang-d-side trailing ``x_final`` aliasing issue when the window ends before the last denoising step.""" - window_size = args.diffusion_sde_window_size - range_raw = args.diffusion_sde_window_range + del sample, rollout_id + window_size = int(args.diffusion_sde_window_size) + range_raw = getattr(args, "diffusion_sde_window_range", None) if range_raw: parts = [int(x) for x in str(range_raw).split(",")] lo, hi = parts[0], parts[1] diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 549dc71f..ab70bddb 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -395,10 +395,7 @@ def add_rollout_arguments(parser): default=False, help=( "Apply miles.backends.sglang_diffusion_utils.monkey_patches at " - "sglang-d startup so its DiT forward is bit-exact with diffusers' " - "implementation. Makes rollout (sglang-d path) and training-side " - "log-prob agree on noise_pred down to bf16 ULPs. Small perf hit on " - "the rollout engine." + "sglang-d startup for SD3 diffusers bf16 parity. Small perf hit." ), ) parser.add_argument( @@ -445,6 +442,19 @@ def add_rollout_arguments(parser): "Defaults to 'transformer', the DiT component for diffusers-based pipelines." ), ) + parser.add_argument( + "--diffusion-model-type", + type=str, + default="auto", + choices=["auto", "sd3", "ltx"], + help=( + "Override the diffusion model family. ``auto`` infers from --diffusion-model " + "(HF hub id with ``ltx`` → ltx, else sd3)." + ), + ) + from miles.backends.fsdp_utils.configs.ltx import register_args as register_ltx_args + + register_ltx_args(parser) parser.add_argument( "--rollout-seed", type=int, @@ -942,7 +952,7 @@ def add_debug_arguments(parser): default=False, help=( "Whether to only run the training without sglang servers. " - "This is useful for debugging the rollout generation function." + "Typically used with --load-debug-rollout-data to replay saved rollouts." ), ) parser.add_argument( @@ -1294,6 +1304,18 @@ def miles_validate_args(args): "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." ) + model_type_arg = (getattr(args, "diffusion_model_type", "auto") or "auto").lower() + if model_type_arg == "auto": + diff_model = (args.diffusion_model or "").lower() + if "ltx" in diff_model or diff_model.endswith(".safetensors"): + args.diffusion_model_type = "ltx" + else: + args.diffusion_model_type = "sd3" + if args.diffusion_model_type == "ltx": + from miles.backends.fsdp_utils.configs.ltx import validate_args as validate_ltx_args + + validate_ltx_args(args) + # always true on offload for colocate at the moment. if args.colocate: if args.offload_train is None: diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 8e825b36..f014dd4c 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -78,8 +78,15 @@ def _parse_cond_kwargs( freqs_cis=[deserialize_func(x) for x in data.get("freqs_cis", [])], img_shapes=data.get("img_shapes"), encoder_hidden_states=_parse_tensor_or_list( - data.get("encoder_hidden_states"), deserialize_func=deserialize_func + data.get("encoder_hidden_states") or data.get("context"), + deserialize_func=deserialize_func, ), + audio_encoder_hidden_states=_parse_tensor_or_list( + data.get("audio_encoder_hidden_states"), + deserialize_func=deserialize_func, + ), + encoder_attention_mask=deserialize_func(data.get("encoder_attention_mask")), + audio_encoder_attention_mask=deserialize_func(data.get("audio_encoder_attention_mask")), pooled_projections=_parse_tensor_or_list(data.get("pooled_projections"), deserialize_func=deserialize_func), ) @@ -110,6 +117,7 @@ def _parse_dit_trajectory( latents=deserialize_func(data.get("latents")), timesteps=deserialize_func(data.get("timesteps")), sigmas=deserialize_func(data.get("sigmas")), + sde_step_indices=deserialize_func(data.get("sde_step_indices")), ) diff --git a/miles/utils/sde_log_prob.py b/miles/utils/sde_log_prob.py index 60495dc1..1fad39cf 100644 --- a/miles/utils/sde_log_prob.py +++ b/miles/utils/sde_log_prob.py @@ -3,11 +3,20 @@ Adapted from flow_grpo/diffusers_patch/sd3_sde_with_logprob.py. """ +from __future__ import annotations + import math import torch +def _broadcast_sigma(sigma: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + sigma_view = sigma.float() + while sigma_view.ndim < sample.ndim: + sigma_view = sigma_view.unsqueeze(-1) + return sigma_view + + def sde_step_with_logprob( scheduler, model_output: torch.FloatTensor, @@ -15,16 +24,24 @@ def sde_step_with_logprob( sample: torch.FloatTensor, prev_sample: torch.FloatTensor, noise_level: float = 0.7, + *, + sde_type: str = "sde", + sigma: torch.FloatTensor | None = None, + sigma_prev: torch.FloatTensor | None = None, ): """Compute the log probability of `prev_sample` under one reverse-SDE step. Args: scheduler: A flow-matching scheduler with `sigmas` and `index_for_timestep`. - model_output: Predicted velocity from DiT, shape (B, C, H, W). - timestep: Current timestep(s), shape (B,). - sample: Current latent, shape (B, C, H, W). + Ignored when ``sigma`` and ``sigma_prev`` are both provided. + model_output: Predicted velocity from DiT. + timestep: Current timestep(s), shape (B,). Used only for scheduler lookup. + sample: Current latent. prev_sample: Recorded next-step latent to score under the SDE. noise_level: SDE noise scaling factor (eta). + sde_type: ``"sde"`` (default, SD3 flow-SDE) or ``"cps"``. + sigma: Optional current sigma(s), shape (B,). Bypasses scheduler lookup. + sigma_prev: Optional next sigma(s), shape (B,). Required with ``sigma``. Returns: (prev_sample, log_prob, prev_sample_mean, std_dev_t) @@ -32,26 +49,43 @@ def sde_step_with_logprob( model_output = model_output.float() sample = sample.float() prev_sample = prev_sample.float() + sde_type = str(sde_type).strip().lower() + + if sigma is not None and sigma_prev is not None: + sigma = _broadcast_sigma(sigma, sample) + sigma_prev = _broadcast_sigma(sigma_prev, sample) + else: + step_index = [scheduler.index_for_timestep(t) for t in timestep] + prev_step_index = [s + 1 for s in step_index] + sigma = scheduler.sigmas[step_index].view(-1, *([1] * (len(sample.shape) - 1))) + sigma_prev = scheduler.sigmas[prev_step_index].view(-1, *([1] * (len(sample.shape) - 1))) + + if sde_type == "cps": + std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2) + pred_original_sample = sample - sigma * model_output + noise_estimate = sample + model_output * (1.0 - sigma) + prev_sample_mean = pred_original_sample * (1.0 - sigma_prev) + noise_estimate * torch.sqrt( + torch.clamp(sigma_prev**2 - std_dev_t**2, min=1e-12) + ) + log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) + elif sde_type == "sde": + sigma_max = scheduler.sigmas[1].item() + dt = sigma_prev - sigma + + std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level + + prev_sample_mean = ( + sample * (1 + std_dev_t**2 / (2 * sigma) * dt) + + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt + ) - step_index = [scheduler.index_for_timestep(t) for t in timestep] - prev_step_index = [s + 1 for s in step_index] - sigma = scheduler.sigmas[step_index].view(-1, *([1] * (len(sample.shape) - 1))) - sigma_prev = scheduler.sigmas[prev_step_index].view(-1, *([1] * (len(sample.shape) - 1))) - sigma_max = scheduler.sigmas[1].item() - dt = sigma_prev - sigma - - std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level - - prev_sample_mean = ( - sample * (1 + std_dev_t**2 / (2 * sigma) * dt) - + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt - ) - - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) - - torch.log(std_dev_t * torch.sqrt(-1 * dt)) - - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) - ) + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) + - torch.log(std_dev_t * torch.sqrt(-1 * dt)) + - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) + ) + else: + raise ValueError(f"Unsupported sde_type {sde_type!r}; expected 'sde' or 'cps'.") # mean along all but batch dimension log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) diff --git a/miles/utils/types.py b/miles/utils/types.py index d7169d41..d9af6c1f 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -36,6 +36,9 @@ class CondKwargs: freqs_cis: list[torch.Tensor] | None = None img_shapes: list[list[tuple[int, int, int]]] | None = None encoder_hidden_states: list[torch.Tensor] | None = None + audio_encoder_hidden_states: list[torch.Tensor] | None = None + encoder_attention_mask: torch.Tensor | None = None + audio_encoder_attention_mask: torch.Tensor | None = None pooled_projections: list[torch.Tensor] | None = None @@ -57,6 +60,7 @@ class DiTTrajectory: # σ * 1000 / 1000 in fp32 and drifts 1-2 ULPs, amplifying to ~3e-5 # log_prob diff. sigmas: torch.Tensor | None = None + sde_step_indices: torch.Tensor | None = None @dataclass @@ -85,7 +89,9 @@ class Sample: inference_time_s: float | None = None peak_memory_mb: float | None = None - reward: dict[str, Any] | None = None + # Scalar from single RM (e.g. pickscore) or dict when combining multiple RMs + # (--reward-key selects the scalar used for GRPO / logging). + reward: float | dict[str, Any] | None = None weight_versions: list[str] = field(default_factory=list) class Status(Enum): @@ -124,5 +130,23 @@ def from_dict(data: dict): return sample - def get_reward_value(self, args) -> float: - return self.reward if not args.reward_key else self.reward[args.reward_key] + def get_reward_value(self, args, *, reward_key: str | None = None) -> float: + key = reward_key if reward_key is not None else getattr(args, "reward_key", None) + if isinstance(self.reward, dict): + if not key: + raise ValueError("sample.reward is a dict but no reward_key configured") + return float(self.reward[key]) + if self.reward is None: + raise ValueError("sample.reward is None") + return float(self.reward) + + def get_sde_step_indices(self) -> list[int] | None: + """Per-sample SDE training step indices from rollout metadata or trajectory.""" + md = self.train_metadata or {} + sde = md.get("sde_step_indices") + if sde is not None: + return sde + traj = self.dit_trajectory + if traj is not None and traj.sde_step_indices is not None: + return traj.sde_step_indices.tolist() + return None diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh new file mode 100644 index 00000000..925e8ea7 --- /dev/null +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -0,0 +1,146 @@ +#!/usr/bin/env bash +# LTX-2.3 video PickScore GRPO: sglang rollout + FSDP train (colocate). +# +# Default: 2-GPU colocate on CUDA 6,7 (train FSDP DP + one sglang engine / GPU). +# Override with CUDA_VISIBLE_DEVICES / NUM_GPUS. CPS dynamics, 3 SDE steps from +# candidates 0–9, clip 1e-4. +# +# Examples: +# # formal 2-GPU (default) +# bash scripts/run-diffusion-grpo-ltx23-sglang.sh +# +# # smoke +# CUDA_VISIBLE_DEVICES=6,7 NUM_GPUS=2 \ +# ROLLOUT_BATCH_SIZE=1 N_SAMPLES_PER_PROMPT=2 NUM_ROLLOUT=1 NUM_STEPS_PER_ROLLOUT=1 \ +# bash scripts/run-diffusion-grpo-ltx23-sglang.sh +# +# # single-GPU +# CUDA_VISIBLE_DEVICES=6 NUM_GPUS=1 bash scripts/run-diffusion-grpo-ltx23-sglang.sh +# +# Layout: train+rollout share the first NUM_GPUS in CUDA_VISIBLE_DEVICES; +# optional pickscore worker uses additional GPUs when configured. + +set -euo pipefail +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-6,7}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +if [[ -z "${NUM_GPUS:-}" ]]; then + IFS=',' read -ra _VISIBLE_GPUS <<< "${CUDA_VISIBLE_DEVICES}" + NUM_GPUS="${#_VISIBLE_GPUS[@]}" +fi +NUM_GPUS="${NUM_GPUS}" + +RUN_NAME="diffusion_grpo_ltx23_pickscore_$(date +%Y%m%d_%H%M%S)" +SAVE_DIR="${ROOT_DIR}/logs/${RUN_NAME}/ckpt" +mkdir -p "${SAVE_DIR}" + +PYTHON_BIN="${PYTHON_BIN:-python}" + +DATASETS_DIR="${DATASETS_DIR:-/root/datasets/miles-diffusion-datasets}" +if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then + hf download --repo-type dataset rockdu/miles-diffusion-datasets \ + --include "flowgrpo_pickscore/**" \ + --local-dir "${DATASETS_DIR}" +fi + +DIFFUSION_MODEL="${DIFFUSION_MODEL:-Lightricks/LTX-2.3}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-2}" +NUM_ROLLOUT="${NUM_ROLLOUT:-200}" +SAVE_INTERVAL="${SAVE_INTERVAL:-50}" +# One engine per GPU by default; concurrency scales with engines. +ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" +SGLANG_SERVER_CONCURRENCY="${SGLANG_SERVER_CONCURRENCY:-${NUM_GPUS}}" + +echo "[ltx23] CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} NUM_GPUS=${NUM_GPUS} engines=$((NUM_GPUS / ROLLOUT_NUM_GPUS_PER_ENGINE))" +echo "[ltx23] batch=${ROLLOUT_BATCH_SIZE}x${N_SAMPLES_PER_PROMPT} rollouts=${NUM_ROLLOUT} save_interval=${SAVE_INTERVAL}" +echo "[ltx23] run=${RUN_NAME}" + +WANDB_ARGS=() +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS+=( + --use-wandb + --wandb-project "${WANDB_PROJECT:-miles-diffusion-grpo}" + --wandb-group "${RUN_NAME}" + --wandb-key "${WANDB_API_KEY}" + --diffusion-log-images "${WANDB_LOG_IMAGES:-4}" + --diffusion-log-image-interval "${WANDB_LOG_IMAGE_INTERVAL:-10}" + --disable-wandb-random-suffix + ) +fi + +LORA_ARGS=() +if [[ "${USE_LORA:-1}" == "1" ]]; then + LORA_ARGS+=( + --use-lora + --lora-rank 64 + --lora-alpha 128 + --diffusion-init-lora-weight gaussian + ) +fi + +"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ + --train-backend fsdp \ + --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ + --diffusion-model "${DIFFUSION_MODEL}" \ + --diffusion-model-type ltx \ + --hf-checkpoint gpt2 \ + --prompt-data "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" \ + --input-key input \ + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" \ + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" \ + --num-steps-per-rollout "${NUM_STEPS_PER_ROLLOUT}" \ + --num-rollout "${NUM_ROLLOUT}" \ + --micro-batch-size-sample "${MICRO_BATCH_SIZE_SAMPLE:-1}" \ + --micro-batch-size-tstep "${MICRO_BATCH_SIZE_TSTEP:-1}" \ + --gradient-checkpointing \ + --colocate \ + --actor-num-gpus-per-node "${NUM_GPUS}" \ + --actor-num-nodes 1 \ + --num-gpus-per-node "${NUM_GPUS}" \ + --rollout-num-gpus "${NUM_GPUS}" \ + --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE}" \ + --use-miles-router \ + --rollout-health-check-interval "${ROLLOUT_HEALTH_CHECK_INTERVAL:-120}" \ + --miles-router-health-check-failure-threshold "${MILES_ROUTER_HEALTH_CHECK_FAILURE_THRESHOLD:-30}" \ + --sglang-server-concurrency "${SGLANG_SERVER_CONCURRENCY}" \ + --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ + "${LORA_ARGS[@]}" \ + --lr 2e-4 \ + --adam-beta2 0.999 \ + --weight-decay 1e-4 \ + --diffusion-clip-range "${CLIP_RANGE:-1e-4}" \ + --diffusion-kl-beta 0.0 \ + --diffusion-num-steps "${NUM_STEPS:-24}" \ + --diffusion-step-strategy-path miles.rollout.step_strategy_hub.ltx_sde_candidates \ + --ltx-num-sde-steps "${LTX_NUM_SDE_STEPS:-3}" \ + --ltx-sde-step-candidates "${LTX_SDE_STEP_CANDIDATES:-0,1,2,3,4,5,6,7,8,9}" \ + --ltx-dynamics-type CPS \ + --diffusion-noise-level 0.8 \ + --ltx-sigma-min 0.001 \ + --diffusion-guidance-scale 1.0 \ + --diffusion-height "${HEIGHT:-512}" \ + --diffusion-width "${WIDTH:-768}" \ + --ltx-frames "${FRAMES:-57}" \ + --ltx-fps "${LTX_FPS:-24}" \ + --diffusion-forward-dtype bf16 \ + --fsdp-master-dtype bf16 \ + --fsdp-reduce-dtype bf16 \ + --sglang-dit-precision bf16 \ + --advantage-estimator grpo \ + --globalize-reward-std \ + --rm-type pickscore \ + --diffusion-reward "pickscore:1.0" \ + --reward-key avg \ + --pickscore-processor-path "${PICKSCORE_PROCESSOR:-laion/CLIP-ViT-H-14-laion2B-s32B-b79K}" \ + --pickscore-model-path "${PICKSCORE_MODEL:-yuvalkirstain/PickScore_v1}" \ + --pickscore-num-frames "${PICKSCORE_NUM_FRAMES:-3}" \ + --pickscore-num-gpus-per-worker "${PICKSCORE_NUM_GPUS_PER_WORKER:-0}" \ + --pickscore-batch-size 8 \ + --update-weight-buffer-size 2147483648 \ + --save "${SAVE_DIR}" \ + --save-interval "${SAVE_INTERVAL}" \ + "${WANDB_ARGS[@]}" \ + "$@"