From 87cf4b3b98af41fc805f72b5180988cd14a04ce5 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:17:52 +0000 Subject: [PATCH 01/31] chore: gitignore hygiene for debug artifacts Ignore core dumps, wandb logs, and local dist/ output. Stop ignoring docs/, scripts/, and tools/ so they can land in follow-up PRs. --- .gitignore | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 03c124b1..89451138 100644 --- a/.gitignore +++ b/.gitignore @@ -192,4 +192,14 @@ local/ glm/ _examples_synced/ .env -.DS_Store \ No newline at end of file +.DS_Store + +# crash dumps / runtime artifacts +core +core.* + +logs/ +wandb/ + +# local debug scripts / alignment tools (not shipped with the library) +dist/ \ No newline at end of file From b73844a63e64b1303d5cc38eb2cc5d3cb1777b93 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:17:57 +0000 Subject: [PATCH 02/31] fix: reward handling and debug_train_only rollout skip Unify scalar and dict rewards via Sample.get_reward_value() for eval and wandb logging. Skip sglang router and HTTP client init when debug_train_only is set so train-side rollout replay does not require live engines. --- miles/ray/placement_group.py | 5 +++-- miles/ray/rollout.py | 19 ++++++++++++------- miles/rollout/sglang_diffusion_rollout.py | 4 +++- miles/utils/types.py | 15 ++++++++++++--- 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index c47f8b68..58f20def 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -82,8 +82,9 @@ def create_placement_groups(args): """Create placement groups for actor and rollout engines. Two topologies: - - Colocate (or --debug-{train,rollout}-only): one combined placement - group; both roles see the same bundle list. + - Colocate (or ``--debug-train-only`` / ``--debug-rollout-only``): one combined + placement group; both roles see the same bundle list (train-only allocates + no rollout GPU bundles). - Disaggregate (the else branch): two separate placement groups so train and rollout each own a disjoint GPU pool — avoids bundle overlap / scheduling deadlock when running side-by-side. diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 036b0e84..11987db4 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -45,13 +45,19 @@ def __init__(self, args, pg): logger.info("RolloutManager init start") self.args = args self.pg = pg - logger.info("RolloutManager: starting router...") - _start_router(args) - logger.info("RolloutManager: router started, init tracking...") + if self.args.debug_train_only: + logger.info("RolloutManager: debug_train_only, skipping sglang router.") + router_addr = None + else: + logger.info("RolloutManager: starting router...") + _start_router(args) + logger.info("RolloutManager: router started, init tracking...") + router_addr = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" # TODO make args immutable - init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") - logger.info("RolloutManager: init http client...") - init_http_client(args) + init_tracking(args, primary=False, router_addr=router_addr) + if not self.args.debug_train_only: + logger.info("RolloutManager: init http client...") + init_http_client(args) logger.info("RolloutManager: loading data source...") data_source_cls = load_function(self.args.data_source_path) @@ -165,7 +171,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() diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 69b1caa9..cf89138e 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -444,7 +444,9 @@ async def eval_rollout_single_dataset( reward_key = args.eval_reward_key or args.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/utils/types.py b/miles/utils/types.py index f15c8588..39419f2d 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -82,7 +82,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): @@ -121,5 +123,12 @@ 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) From d0e42b49c0058094a3d8acd61ae658ad47882acd Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:18:04 +0000 Subject: [PATCH 03/31] refactor: extract rollout engine env var builders Factor sglang Ray worker env defaults into helpers and fall back to dit_trajectory.sde_step_indices when train_metadata omits SDE step lists. --- miles/ray/rollout.py | 60 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 58 insertions(+), 2 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 11987db4..77e35710 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -2,6 +2,7 @@ import itertools import logging import multiprocessing +import os import random import time from pathlib import Path @@ -397,7 +398,14 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "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 + (sample.train_metadata or {}).get("sde_step_indices") + if (sample.train_metadata or {}).get("sde_step_indices") is not None + else ( + sample.dit_trajectory.sde_step_indices.tolist() + if sample.dit_trajectory is not None and sample.dit_trajectory.sde_step_indices is not None + else None + ) + for sample in samples ], } @@ -465,10 +473,58 @@ def _split_train_data_by_dp(self, data, dp_size): rollout_data_refs.append(Box(ray.put(rollout_data))) return rollout_data_refs + +def _base_rollout_engine_env_vars() -> dict[str, str]: + return {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", + } + + +def _ltx_alignment_env_vars(args) -> dict[str, str]: + """Propagate LTX train/rollout alignment flags into sglang Ray workers. + + Ray ``runtime_env`` does not inherit the parent shell env; monkey patches in + ``sglang_diffusion_engine`` read these variables at scheduler startup. + + TODO(PR4): replace env bridging with explicit sglang ``ServerArgs`` fields + (or upstream train-mode guider / video-only flags) so miles does not depend + on opaque env + runtime monkey patches. + """ + env: dict[str, str] = {} + if getattr(args, "ltx_disable_av_cross_attn", False): + env["MILES_LTX_DISABLE_AV_CROSS"] = "1" + from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV + + for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): + if os.environ.get(name): + env[name] = os.environ[name] + return env + + +def _build_rollout_engine_env_vars(args) -> dict[str, str]: + env_vars = _base_rollout_engine_env_vars() + for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): + if os.environ.get(cache_var): + env_vars[cache_var] = os.environ[cache_var] + env_vars.update(_ltx_alignment_env_vars(args)) + return env_vars + + def init_rollout_engines(args, pg, all_rollout_engines): if args.debug_train_only: return 0 - + + from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import ( + SGLangDiffusionEngine, + ) + num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine assert len(all_rollout_engines) == num_engines From 7f19bc97e519f73665d693e1c3efc2820a4aad72 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:18:15 +0000 Subject: [PATCH 04/31] feat(ltx): add CLI args, types, and rollout sampling scaffold Add LTX-2 model-type detection, video rollout CLI flags, and CondKwargs fields for positions/denoise_mask/clean_latent. Wire basic LTX sampling params (frames, fps, no-CFG) into sglang diffusion rollout. --- miles/rollout/sglang_diffusion_rollout.py | 25 ++++++ miles/utils/arguments.py | 99 +++++++++++++++++++++-- miles/utils/types.py | 4 + 3 files changed, 122 insertions(+), 6 deletions(-) diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index cf89138e..887cfb18 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -29,6 +29,16 @@ logger = logging.getLogger(__name__) +def _resolve_diffusion_model_type(args: Namespace) -> str: + model_type = (getattr(args, "diffusion_model_type", "auto") or "auto").lower() + if model_type != "auto": + return model_type + diff_model = (getattr(args, "diffusion_model", None) or "").lower() + if "ltx" in diff_model or diff_model.endswith(".safetensors"): + return "ltx" + return "sd3" + + def build_rollout_sampling_params( args: Namespace, *, @@ -51,6 +61,21 @@ def build_rollout_sampling_params( "true_cfg_scale": getattr(args, "diffusion_true_cfg_scale", None), } + model_type = _resolve_diffusion_model_type(args) + if model_type == "ltx": + 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"] = " " + # LTX23 one-stage rollout uses a stage1 guider (CFG/STG/modality/rescale) + # whose params cannot be overridden via HTTP — sglang routes unknown + # SamplingParams kwargs through the base class. Train ``forward_velocity`` + # is video-only with no guidance, so the rollout engine forces an identity + # guider via the ``patch_ltx2_identity_guider`` monkey patch + # (MILES_LTX_IDENTITY_GUIDER, propagated in miles/ray/rollout.py). + if evaluation: sampling_params["rollout"] = False else: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 58c5eb19..9fdf99e1 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -396,10 +396,9 @@ 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. For LTX-2: also " + "applies ltx_core AdaLN/RoPE/attention parity patches " + "(disable via MILES_APPLY_LTX2_LTXCORE_PARITY=0). Small perf hit." ), ) parser.add_argument( @@ -437,6 +436,71 @@ def add_rollout_arguments(parser): default=1, help="Log diffusion images every N rollouts. Only used when diffusion-log-images > 0.", ) + 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 " + "(diffusers repo → sd3, single-file safetensors → ltx)." + ), + ) + parser.add_argument( + "--ltx-gemma-path", + type=str, + default=None, + help="Path to the Gemma-3 12B directory used as LTX's text encoder (rollout side).", + ) + 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-disable-av-cross-attn", + action="store_true", + default=False, + help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", + ) + parser.add_argument( + "--pickscore-num-frames", + type=int, + default=3, + help="Number of evenly spaced frames to score per video (LTX PickScore reward).", + ) parser.add_argument( "--rollout-seed", type=int, @@ -917,7 +981,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( @@ -1252,9 +1316,32 @@ def miles_validate_args(args): args.train_memory_margin_bytes = 0 assert not (args.debug_rollout_only and args.debug_train_only), ( - "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." + "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": + if float(getattr(args, "diffusion_guidance_scale", 1.0)) != 1.0: + raise ValueError( + "LTX requires --diffusion-guidance-scale 1.0 (no CFG)." + ) + if not args.debug_train_only and not getattr(args, "ltx_gemma_path", None): + logger.warning( + "--ltx-gemma-path is not set; sglang LTX rollout will need it once Phase 1 lands." + ) + 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." + ) + # always true on offload for colocate at the moment. if args.colocate: if args.offload_train is None: diff --git a/miles/utils/types.py b/miles/utils/types.py index 39419f2d..2457b578 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -34,6 +34,9 @@ class CondKwargs: img_shapes: list[list[tuple[int, int, int]]] | None = None encoder_hidden_states: list[torch.Tensor] | None = None pooled_projections: list[torch.Tensor] | None = None + ltx_positions: torch.Tensor | None = None + ltx_denoise_mask: torch.Tensor | None = None + ltx_clean_latent: torch.Tensor | None = None @dataclass @@ -54,6 +57,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 From 4c1f2ddca77e0cace55d2c02e890ed6620a9bb24 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:18:19 +0000 Subject: [PATCH 05/31] feat(ltx): train pipeline config, SDE dynamics, and FSDP actor Add LTXTrainPipelineConfig for ltx_core Modality forward and schedule-decoupled SDE log-prob (cps/flow_sde/ode/dance_sde). Wire FSDP actor to pluggable pipeline configs with colocate memory fixes and train/rollout alignment metrics. --- miles/backends/fsdp_utils/actor.py | 167 +++++++++--- miles/backends/fsdp_utils/configs/ltx.py | 246 ++++++++++++++++++ .../configs/train_pipeline_config.py | 57 +++- miles/utils/sde_log_prob.py | 179 +++++++++++-- 4 files changed, 590 insertions(+), 59 deletions(-) create mode 100644 miles/backends/fsdp_utils/configs/ltx.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f04b3525..6cce0bbf 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -15,7 +15,6 @@ from miles.utils.distributed_utils import get_gloo_group from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.metric_utils import compute_rollout_step -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 from miles.utils import tracking_utils @@ -24,6 +23,7 @@ from .configs.train_pipeline_config import get_train_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 +import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config from . import checkpoint from .lr_scheduler import get_lr_scheduler @@ -66,21 +66,27 @@ 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, + diffusion_model_id = args.diffusion_model or args.hf_checkpoint + + self.train_pipeline_config = get_train_pipeline_config(diffusion_model_id) + + if self.train_pipeline_config.is_diffusers_pipeline: + with self._get_init_weight_context_manager(): + pipeline = DiffusionPipeline.from_pretrained( + diffusion_model_id, + 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 + else: + model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( + self.args, init_context_factory=self._get_init_weight_context_manager, ) - model = pipeline.transformer - self.scheduler = pipeline.scheduler - del pipeline - - self.train_pipeline_config = get_train_pipeline_config(args.diffusion_model) if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) @@ -88,7 +94,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty model.train() if args.gradient_checkpointing: - model.enable_gradient_checkpointing() + if hasattr(model, "enable_gradient_checkpointing"): + model.enable_gradient_checkpointing() + elif hasattr(model, "set_gradient_checkpointing"): + model.set_gradient_checkpointing(True) + else: + logger.warning( + "gradient_checkpointing requested but model %s exposes neither " + "enable_gradient_checkpointing() nor set_gradient_checkpointing(); " + "skipping.", + type(model).__name__, + ) model.to(torch.cuda.current_device()) @@ -99,6 +115,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty mesh=self.parallel_state.dp_mesh, cpu_offload=self.args.fsdp_cpu_offload, args=self.args, + train_pipeline_config=self.train_pipeline_config, ) # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -158,6 +175,8 @@ def sleep(self) -> None: print_memory("before offload DiT") + self.optimizer.zero_grad(set_to_none=True) + _reshard_fsdp2_model(self.model) self.model.cpu() move_torch_optimizer(self.optimizer, "cpu") clear_memory() @@ -198,6 +217,7 @@ def update_weights(self) -> None: # type: ignore[override] ray.get(self.rollout_manager.clear_num_new_engines.remote()) self.weight_updater.update_weights() + dist.barrier(group=get_gloo_group()) clear_memory() def _get_init_weight_context_manager(self): @@ -298,6 +318,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: 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 + if not getattr(self.train_pipeline_config, "supports_cfg", True): + use_cfg = False # ------------- KL loss ------------- kl_beta = float(self.args.diffusion_kl_beta) @@ -322,9 +344,12 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: # ------------- 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) + sched_config = getattr(self.scheduler, "config", None) + num_train_timesteps = ( + int(sched_config.num_train_timesteps) if sched_config is not None else 1000 + ) if sigmas_snapshot is not None: sigmas_ref = sigmas_snapshot.to(device).float() else: @@ -430,6 +455,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): @@ -461,6 +487,7 @@ def _build_train_grids( ) sde_step_indices = sde_step_indices_list[traj_idx] + sde_indices_per_sample: torch.Tensor | None = 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] @@ -469,8 +496,14 @@ 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: + # Full-length debug tensors (legacy): index by global step. + rollout_model_output = rollout_model_output[sde_indices_tensor] + # else: sglang packs debug outputs in SDE-window order (0..W-1). current_window_size = int(sde_indices_tensor.numel()) + sde_indices_per_sample = sde_indices_tensor else: current_window_size = default_window_size @@ -488,6 +521,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_per_sample is not None: + sde_indices_per_sample_list.append(sde_indices_per_sample) latents_window = torch.stack(latents_list, dim=0) next_latents_window = torch.stack(next_latents_list, dim=0) @@ -532,6 +567,31 @@ 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 _build_sde_extra( + self, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + ) -> dict | None: + if grids.get("sde_step_indices_window") is None: + return None + + idx = grids["sde_step_indices_window"][sample_indices][:, tstep_indices] + idx = idx.reshape(-1).long() + + return { + "sigmas": self.scheduler.sigmas, + "sde_step_indices": idx, + "dynamics_type": getattr(self.args, "ltx_dynamics_type", "cps"), + "sigma_min_override": getattr(self.args, "ltx_sigma_min", None), } def _run_optim_window( @@ -670,12 +730,9 @@ def _forward_tile( timesteps_input = timesteps_for_model.to(forward_dtype) def _forward(cond: dict) -> torch.Tensor: - return self.model( - hidden_states=latents_input, - timestep=timesteps_input, - return_dict=False, - **cond, - )[0] + return train_pipeline_config.forward_velocity( + self.model, latents_input, timesteps_input, cond, + ) cfg_batching = bool(self.args.fsdp_cfg_batching) @@ -686,12 +743,9 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: 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] + joint_out = train_pipeline_config.forward_velocity_cfg_joint( + self.model, latents_input, timesteps_input, joint_cond, + ) noise_pred_pos, noise_pred_neg = joint_out.chunk(2, dim=0) else: noise_pred_pos = _forward(pos_cond_tile) @@ -702,16 +756,20 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: noise_pred_flat = _compute_noise_pred() - _, log_prob_new_flat, prev_sample_mean, std_dev_t = sde_step_with_logprob( + sde_extra = self._build_sde_extra(grids, sample_indices, tstep_indices) + + prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( self.scheduler, - noise_pred_flat.float(), + noise_pred_flat, timesteps_flat, - latents_flat.float(), + latents_flat, prev_sample=next_latents_tile.reshape( tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ).float(), + ), noise_level=noise_level, + extra=sde_extra, ) + del prev_sample_dummy # TODO: revamp and gather all loss logics log_prob_new = log_prob_new_flat.reshape(tile_sample_count, tile_tstep_count) @@ -724,16 +782,16 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: 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( + _, _, prev_sample_mean_ref, _ = train_pipeline_config.sde_step( self.scheduler, - ref_noise_pred_flat.float(), + ref_noise_pred_flat, timesteps_flat, - latents_flat.float(), + latents_flat, prev_sample=next_latents_tile.reshape( tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ).float(), + ), noise_level=noise_level, + extra=sde_extra, ) kl_loss = ((prev_sample_mean - prev_sample_mean_ref) ** 2).mean( dim=tuple(range(1, prev_sample_mean.ndim)), @@ -772,6 +830,13 @@ 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()) + flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) + flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) + log_stats["model_output_cosine_sim"].append( + torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1) + .mean() + .detach() + ) return loss @@ -846,6 +911,17 @@ def _cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict: return out +def _reshard_fsdp2_model(model: torch.nn.Module) -> None: + """Drop FSDP2 unsharded views so model.cpu() can release GPU memory.""" + if hasattr(model, "reshard"): + model.reshard() + return + for module in model.modules(): + reshard = getattr(module, "reshard", None) + if callable(reshard): + reshard() + + @torch.no_grad() def move_torch_optimizer(optimizer, device): """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py""" @@ -890,13 +966,20 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - model.print_trainable_parameters() return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, train_pipeline_config=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 = ( + getattr(train_pipeline_config, "fsdp_wrap_classes", None) if train_pipeline_config else None + ) + assert layer_cls_to_wrap and layer_cls_to_wrap[0] is not None, ( + "apply_fsdp2 needs either model._no_split_modules or " + "train_pipeline_config.fsdp_wrap_classes to know which submodules to shard." + ) modules = [ module diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py new file mode 100644 index 00000000..79c26bea --- /dev/null +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -0,0 +1,246 @@ +"""LTX-2.3 video diffusion training pipeline config. + +Adapts ltx_core's ``LTXModel`` (non-diffusers; Modality-keyed forward, patchified +``[B, T, D]`` token latents, per-token timesteps, custom ``LTX2Scheduler``) into +miles' FSDP GRPO training loop. +""" + +from __future__ import annotations + +import torch + +from miles.utils.types import CondKwargs + +from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config + + +@register_train_pipeline_config("ltx") +class LTXTrainPipelineConfig(TrainPipelineConfig): + """Training-side adapter for LTX-2.3 video DiT.""" + + is_diffusers_pipeline = False + needs_timestep_scaling = False + supports_cfg = False + + fsdp_wrap_classes = ["BasicAVTransformerBlock"] + + lora_target_modules = [ + "to_q", "to_k", "to_v", "to_out.0", + "net.0.proj", "net.2", + ] + + def load_model_and_scheduler(self, args, init_context_factory): + from dataclasses import dataclass, field + + from ltx_core.components.schedulers import LTX2Scheduler + from ltx_trainer.model_loader import load_transformer + + @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") + master_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[master_dtype_name] + + model = load_transformer(args.diffusion_model, device="cpu", dtype=master_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, + ) + return model, scheduler + + 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.ltx_positions is not None: + pos = cond.ltx_positions.to(device) + if pos.ndim == 2: + pos = pos.unsqueeze(0) + elif pos.ndim == 3: + pos = pos.unsqueeze(0) + kwargs["positions"] = pos + if cond.ltx_denoise_mask is not None: + mask = cond.ltx_denoise_mask.to(device) + if mask.ndim == 2 and mask.shape[-1] == 1: + mask = mask.squeeze(-1) + if mask.ndim == 1: + mask = mask.unsqueeze(0) + kwargs["denoise_mask"] = mask + if cond.ltx_clean_latent is not None: + cl = cond.ltx_clean_latent.to(device) + if cl.ndim == 2: + cl = cl.unsqueeze(0) + kwargs["clean_latent"] = cl + return kwargs + + 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] + + sigma = timesteps_input.to(latents_input.dtype) + 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.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, + 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, 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_dynamics + + 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 = extra.get("dynamics_type", "cps") + sigma_min_override = extra.get("sigma_min_override", None) + if sigma_min_override == 0.0: + sigma_min_override = None + + prev, log_prob, prev_mean, std_dev_t, _dt_sqrt = sde_step_with_logprob_dynamics( + model_output=noise_pred.float(), + sigma=sigma_view, + sigma_next=sigma_next, + sample=sample.float(), + sigmas=sigmas, + prev_sample=prev_sample.float(), + sigma_min_override=sigma_min_override, + noise_level=noise_level, + dynamics_type=dynamics_type, + ) + 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 9efbfa62..2ea48b67 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -14,6 +14,7 @@ from __future__ import annotations import abc +from typing import Any import torch from miles.utils.types import CondKwargs, DiTTrajectory @@ -46,6 +47,9 @@ def get_train_pipeline_config(model_name: str) -> "TrainPipelineConfig": class TrainPipelineConfig(abc.ABC): """Base class. Subclass per model family.""" + is_diffusers_pipeline: bool = True + supports_cfg: bool = True + fsdp_wrap_classes: list[str] | None = None 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] = [] @@ -118,4 +122,55 @@ def cfg_combine( @abc.abstractmethod def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: """Preprocess the model before FSDP.""" - pass \ No newline at end of file + pass + + def forward_velocity( + self, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + cond: dict, + ) -> torch.Tensor: + return model( + hidden_states=latents_input, + timestep=timesteps_input, + return_dict=False, + **cond, + )[0] + + def forward_velocity_cfg_joint( + self, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + joint_cond: dict, + ) -> torch.Tensor: + return 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] + + def sde_step( + self, + scheduler: Any, + 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 + + prev, log_prob, prev_mean, std_dev_t = sde_step_with_logprob( + scheduler, + noise_pred.float(), + timesteps, + sample.float(), + prev_sample=prev_sample.float(), + noise_level=noise_level, + ) + return prev, log_prob, prev_mean, std_dev_t \ No newline at end of file diff --git a/miles/utils/sde_log_prob.py b/miles/utils/sde_log_prob.py index 87856e0c..1324660d 100644 --- a/miles/utils/sde_log_prob.py +++ b/miles/utils/sde_log_prob.py @@ -1,13 +1,43 @@ """SDE step with log probability for flow matching schedulers. -Adapted from flow_grpo/diffusers_patch/sd3_sde_with_logprob.py. +This module exposes: + +- :func:`sde_step_with_logprob` — the original SD3 / flow-matching scheduler + contract used by miles' SD3 path. + +- :func:`sde_step_with_logprob_dynamics` — generic, schedule-decoupled version + used by the LTX-2.3 path which runs on patchified token latents and a + custom :class:`LTX2Scheduler`. """ import math -from typing import Union +from typing import Optional, Union import torch +# Canonical dynamics names. These match sglang-d ``rollout_sde_type`` so miles +# can pass them straight through to the rollout engine with no translation +# table — keeping train (this module) and rollout (sglang-d flow_sde_sampling) +# on a single shared vocabulary. +CANONICAL_DYNAMICS_TYPES = ("sde", "flow_sde", "cps", "ode", "dance_sde") + + +def normalize_dynamics_type(name: str) -> str: + """Map a dynamics-type alias (CLI / legacy casing) to its canonical name. + + Accepts any case and ``-``/``_`` spelling, e.g. ``"Flow-SDE"``, + ``"flow_sde"`` -> ``"flow_sde"``; ``"CPS"`` -> ``"cps"``; + ``"Dance-SDE"`` -> ``"dance_sde"``. Raises on unknown values rather than + silently falling back, so a typo can never mismatch train vs rollout. + """ + key = str(name).strip().lower().replace("-", "_") + if key not in CANONICAL_DYNAMICS_TYPES: + raise ValueError( + f"Unknown dynamics_type {name!r}; expected one of " + f"{CANONICAL_DYNAMICS_TYPES}" + ) + return key + def sde_step_with_logprob( scheduler, @@ -17,19 +47,7 @@ def sde_step_with_logprob( prev_sample: torch.FloatTensor, noise_level: float = 0.7, ): - """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). - prev_sample: Recorded next-step latent to score under the SDE. - noise_level: SDE noise scaling factor (eta). - - Returns: - (prev_sample, log_prob, prev_sample_mean, std_dev_t) - """ + """Compute the log probability of `prev_sample` under one reverse-SDE step.""" model_output = model_output.float() sample = sample.float() prev_sample = prev_sample.float() @@ -54,7 +72,136 @@ def sde_step_with_logprob( - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) ) - # mean along all but batch dimension log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) return prev_sample, log_prob, prev_sample_mean, std_dev_t + + +def sde_step_with_logprob_dynamics( + model_output: torch.FloatTensor, + sigma: torch.FloatTensor, + sigma_next: torch.FloatTensor, + sample: torch.FloatTensor, + sigmas: torch.FloatTensor, + prev_sample: Optional[torch.FloatTensor] = None, + generator: Optional[torch.Generator] = None, + deterministic: bool = False, + sigma_min_override: Optional[float] = None, + noise_level: float = 0.8, + dynamics_type: str = "flow_sde", +): + """Schedule-decoupled SDE step with log-prob for LTX-2.3 and similar models. + + ``dynamics_type`` accepts any alias (see :func:`normalize_dynamics_type`). + """ + dynamics_type = normalize_dynamics_type(dynamics_type) + model_output = model_output.float() + sample = sample.float() + if prev_sample is not None: + prev_sample = prev_sample.float() + + ndim = sample.ndim + sigma_view = sigma.float() + sigma_next_view = sigma_next.float() + while sigma_view.ndim < ndim: + sigma_view = sigma_view.unsqueeze(-1) + while sigma_next_view.ndim < ndim: + sigma_next_view = sigma_next_view.unsqueeze(-1) + + dt = sigma_next_view - sigma_view + + sigma_max = sigmas[0].float().item() + if sigma_min_override is not None: + sigma_min = sigma_min_override + else: + sigma_min = max(sigmas[-2].float().item(), 1e-4) if len(sigmas) > 1 else 1e-4 + + if dynamics_type == "ode": + prev_sample_mean = sample + dt * model_output + std_dev_t = torch.zeros_like(sigma_view) + if prev_sample is None: + prev_sample = prev_sample_mean + log_prob = torch.zeros(sample.shape[0], dtype=sample.dtype, device=sample.device) + + elif dynamics_type == "flow_sde": + std_dev_t = (sigma_min + (sigma_max - sigma_min) * sigma_view) * noise_level + sigma_safe = torch.clamp(sigma_view, min=1e-8) + + drift_sample = 1.0 + std_dev_t**2 / (2.0 * sigma_safe) * dt + drift_model = (1.0 + std_dev_t**2 * (1.0 - sigma_view) / (2.0 * sigma_safe)) * dt + prev_sample_mean = sample * drift_sample + model_output * drift_model + + noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) + + if prev_sample is None: + if deterministic: + prev_sample = sample + dt * model_output + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + noise_scale * variance_noise + + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) + - torch.log(noise_scale + 1e-12) + - 0.5 * math.log(2.0 * math.pi) + ) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + elif dynamics_type == "cps": + angle = torch.tensor(noise_level, dtype=sigma_next_view.dtype, device=sigma_next_view.device) * torch.pi / 2 + std_dev_t = sigma_next_view * torch.sin(angle) + + x0 = sample - sigma_view * model_output + x1 = sample + model_output * (1.0 - sigma_view) + sqrt_term = torch.sqrt(torch.clamp(sigma_next_view**2 - std_dev_t**2, min=1e-12)) + prev_sample_mean = x0 * (1.0 - sigma_next_view) + x1 * sqrt_term + + if prev_sample is None: + if deterministic: + prev_sample = prev_sample_mean + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + std_dev_t * variance_noise + + log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + elif dynamics_type == "dance_sde": + sigma_safe = torch.clamp(sigma_view, min=1e-8) + x0_pred = sample - sigma_safe * model_output + std_dev_t = torch.as_tensor(noise_level, dtype=sample.dtype, device=sample.device) + log_term = 0.5 * noise_level**2 * (sample - x0_pred * (1.0 - sigma_view)) / (sigma_safe**2) + prev_sample_mean = sample + (model_output + log_term) * dt + noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) + + if prev_sample is None: + if deterministic: + prev_sample = sample + dt * model_output + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + noise_scale * variance_noise + + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) + - torch.log(noise_scale + 1e-12) + - 0.5 * math.log(2.0 * math.pi) + ) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + else: + # ``sde`` (SD3-style, scheduler-based) is handled by + # ``sde_step_with_logprob`` above, not this schedule-decoupled path. + raise ValueError( + f"dynamics_type {dynamics_type!r} is not supported by the " + "schedule-decoupled path; use flow_sde / cps / ode / dance_sde " + "(or sde via sde_step_with_logprob)." + ) + + dt_sqrt = torch.sqrt(torch.clamp(-dt, min=1e-12)) + return prev_sample, log_prob, prev_sample_mean, std_dev_t, dt_sqrt From 69e3d737b4a5cf76af6ef51e01399c63b2b29a30 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:18:23 +0000 Subject: [PATCH 06/31] feat(ltx): sglang rollout patches, rollout config, and colocate GPU pin Add LTX-2 monkey-patch groups for ltx_core parity and RL alignment. Consolidate rollout ServerArgs helpers into sglang_diffusion_utils/configs/ltx.py. Keep full CUDA_VISIBLE_DEVICES for colocated weight sync and pin DiT via MILES_SGLANG_LOCAL_CUDA_RANK; lazy-import rollout engine on Ray workers. --- .../sglang_diffusion_utils/configs/ltx.py | 84 +++++ .../monkey_patches/__init__.py | 90 ++++- .../patch_ltx2_disable_av_cross.py | 37 +++ .../patch_ltx2_identity_guider.py | 76 +++++ .../patch_ltx2_ltxcore_parity.py | 311 ++++++++++++++++++ .../patch_ltx2_rollout_cond_kwargs.py | 50 +++ .../sglang_diffusion_engine.py | 76 +++-- miles/ray/rollout.py | 28 +- 8 files changed, 696 insertions(+), 56 deletions(-) create mode 100644 miles/backends/sglang_diffusion_utils/configs/ltx.py create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_disable_av_cross.py create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_rollout_cond_kwargs.py 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..990c2ccc --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -0,0 +1,84 @@ +"""LTX-2 sglang-d rollout engine config. + +Mirrors ``fsdp_utils/configs/ltx.py`` on the train side: model detection, +weight-path resolution, and extra ``ServerArgs`` fields for LTX2Pipeline. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path + +logger = logging.getLogger(__name__) + + +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 + diff_model = (getattr(args, "diffusion_model", None) or "").lower() + return "ltx" in diff_model or diff_model.endswith(".safetensors") + + +def resolve_ltx_transformer_weights_path( + diffusion_model: str | None, + *, + explicit_path: str | None = None, +) -> str | None: + """Return official safetensors path for sglang ``transformer_weights_path``. + + When miles train loads a single-file LTX checkpoint, sglang should load the + same safetensors via ``transformer_weights_path`` instead of the HF + materialized ``model.safetensors`` overlay (which can differ at ~1e-4 bf16). + """ + if explicit_path: + path = Path(explicit_path).expanduser() + if path.is_file(): + return str(path) + return None + + 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) + + if diffusion_model and str(diffusion_model).endswith(".safetensors"): + path = Path(diffusion_model).expanduser() + if path.is_file(): + return str(path) + return None + + +def resolve_sglang_model_path(args) -> str: + model_path = args.diffusion_model + if is_ltx_model(args) and model_path.endswith(".safetensors"): + return os.path.dirname(model_path) + return model_path + + +def server_kwargs_extras(args) -> dict: + """Extra ``ServerArgs`` kwargs; call only when ``is_ltx_model(args)``.""" + extras: dict = {"pipeline_class_name": "LTX2Pipeline"} + if getattr(args, "sglang_pipeline_class_name", None): + extras["pipeline_class_name"] = args.sglang_pipeline_class_name + + explicit = getattr(args, "sglang_transformer_weights_path", None) + weights_path = resolve_ltx_transformer_weights_path( + getattr(args, "diffusion_model", None), + explicit_path=explicit, + ) + if weights_path and not explicit: + extras["transformer_weights_path"] = weights_path + logger.info("LTX rollout: transformer_weights_path=%s", weights_path) + elif explicit: + extras["transformer_weights_path"] = explicit + + gemma_path = getattr(args, "ltx_gemma_path", None) + if gemma_path: + extras["component_paths"] = {"text_encoder": gemma_path} + + return extras diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 42278690..93759b79 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -1,29 +1,85 @@ -"""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-2 ltx_core parity + RL alignment (identity guider, AV-off, cond kwargs). + +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 + +import os + +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_LTX2_LTXCORE_PARITY" + + +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) and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1": + return PATCH_GROUP_LTX + + return None -def apply_sgld_monkey_patches() -> None: +def apply_rollout_patch_group(group: str | None) -> None: + if group == PATCH_GROUP_SGLD: + apply_sgld_monkey_patches(include_ltx2_ltxcore=False) + elif group == PATCH_GROUP_LTX: + apply_ltx2_rollout_patches() + + +def apply_sgld_monkey_patches(*, include_ltx2_ltxcore: bool | None = None) -> 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() + + if include_ltx2_ltxcore is None: + include_ltx2_ltxcore = os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" + if include_ltx2_ltxcore: + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + patch_ltx2_ltxcore_parity, + ) + + patch_ltx2_ltxcore_parity.apply() + + +def apply_ltx2_rollout_patches() -> None: + """LTX-2 ltx_core parity + RL train/rollout alignment patches.""" + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + patch_ltx2_disable_av_cross, + patch_ltx2_identity_guider, + patch_ltx2_ltxcore_parity, + patch_ltx2_rollout_cond_kwargs, + ) + + patch_ltx2_ltxcore_parity.apply() + patch_ltx2_disable_av_cross.apply() + patch_ltx2_rollout_cond_kwargs.apply() + patch_ltx2_identity_guider.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..f1172a57 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_disable_av_cross.py @@ -0,0 +1,37 @@ +"""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. Enabled via ``MILES_LTX_DISABLE_AV_CROSS=1`` +(set by miles rollout when ``--ltx-disable-av-cross-attn``), this injects the +disable flags into the DiT forward so the rollout video branch matches train. +""" + +from __future__ import annotations + +import os + +_APPLIED = False + + +def _enabled() -> bool: + return os.environ.get("MILES_LTX_DISABLE_AV_CROSS", "0") == "1" + + +def apply() -> None: + global _APPLIED + if _APPLIED or not _enabled(): + _APPLIED = True + 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_identity_guider.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py new file mode 100644 index 00000000..555aade3 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py @@ -0,0 +1,76 @@ +"""Force an identity LTX-2.3 stage1 guider for train/rollout alignment. + +Miles GRPO train side computes ``forward_velocity`` as a **video-only** forward +with no CFG / STG / modality / rescale. The sglang LTX2.3 one-stage rollout, +however, applies a stage1 guider whose parameters default to ``video_cfg_scale=3`` +etc. (see ``configs/sample/ltx_2.py``). Those parameters **cannot** be overridden +through ``POST /rollout/generate``: ``SamplingParams.from_user_sampling_params_args`` +routes unknown kwargs through the base ``SamplingParams`` class, which rejects +LTX23-only fields. So the rollout-side ``rollout_model_outputs`` are post-guider +velocities that diverge from the train forward (~0.94 cosine, scale≈0.86). + +This patch overrides ``LTX2DenoisingStage._get_ltx2_stage1_guider_params`` so the +guider becomes the identity transform: + + pred = cond + + (cfg-1)*(cond-uncond_text) # cfg=1 -> 0 + + stg*(cond-uncond_perturbed) # stg=0 -> 0 + + (modality-1)*(cond-uncond_mod) # mod=1 -> 0 + pred = rescale(cond, pred, 0.0) # rescale=0 -> pred unchanged + => pred == cond (video-only x0) => velocity == raw video velocity + +Controlled by ``MILES_LTX_IDENTITY_GUIDER`` (default ``"1"``). Set to ``"0"`` to +keep the official guider (e.g. for generation-quality eval, not RL alignment). +""" + +from __future__ import annotations + +import os +from typing import Any + +_APPLIED = False +_ORIG = None + +_IDENTITY_GUIDER: dict[str, Any] = { + "video_cfg_scale": 1.0, + "video_stg_scale": 0.0, + "video_rescale_scale": 0.0, + "video_modality_scale": 1.0, + "video_skip_step": 0, + "video_stg_blocks": [], + "audio_cfg_scale": 1.0, + "audio_stg_scale": 0.0, + "audio_rescale_scale": 0.0, + "audio_modality_scale": 1.0, + "audio_skip_step": 0, + "audio_stg_blocks": [], +} + + +def _identity_enabled() -> bool: + return os.environ.get("MILES_LTX_IDENTITY_GUIDER", "1") == "1" + + +def apply() -> None: + global _APPLIED, _ORIG + if _APPLIED: + return + + from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( + LTX2DenoisingStage, + ) + + _ORIG = LTX2DenoisingStage._get_ltx2_stage1_guider_params + + def _patched_get_guider(self, batch, server_args, stage): + result = _ORIG(self, batch, server_args, stage) + # Only override when guider is active (stage1 returns a dict) and the + # alignment flag is on. None (non-stage1 / official cfg path) is kept. + if result is None or not _identity_enabled(): + return result + merged = dict(result) + merged.update(_IDENTITY_GUIDER) + return merged + + LTX2DenoisingStage._get_ltx2_stage1_guider_params = _patched_get_guider + _APPLIED = True diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py new file mode 100644 index 00000000..00da367d --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py @@ -0,0 +1,311 @@ +"""LTX-2 DiT parity patches: align sglang ltx_2.py with miles/ltx_core.""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch +import torch.nn.functional as F + +_ORIGINALS: dict[str, Any] = {} +_APPLIED = False + + +def expand_temb_for_hidden(temb: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + """Broadcast batch-level temb ``[B, 1, D]`` to ``[B, T, D]`` when uniform.""" + if ( + temb.ndim == 3 + and temb.shape[1] == 1 + and hidden_states.ndim == 3 + and hidden_states.shape[1] > 1 + ): + return temb.expand(-1, hidden_states.shape[1], -1) + return temb + + +def _ltx_pytorch_sdpa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + heads: int, + attn_mask: torch.Tensor | None = None, +) -> torch.Tensor: + b, _, dim_head = q.shape + dim_head //= heads + q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v)) + mask = attn_mask + if mask is not None: + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False + ) + return out.transpose(1, 2).reshape(b, -1, heads * dim_head) + + +def _linear_out(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, module.weight, module.bias) + + +def _ltxcore_rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: + try: + from ltx_core.utils import rms_norm as ltx_rms_norm + + return ltx_rms_norm(x, eps=eps) + except ImportError: + return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps) + + +def _ltxcore_apply_split_rotary_emb( + x: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + cos, sin = freqs + try: + from ltx_core.model.transformer.rope import apply_split_rotary_emb as ltx_apply + + return ltx_apply(x, cos, sin) + except ImportError: + return _pytorch_apply_split_rotary_emb(x, cos, sin) + + +def _pytorch_apply_split_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + x_dtype = x.dtype + needs_reshape = False + if x.ndim != 4 and cos.ndim == 4: + b = x.shape[0] + _, h, t, _ = cos.shape + x = x.reshape(b, t, h, -1).swapaxes(1, 2) + needs_reshape = True + + last = x.shape[-1] + split_x = x.reshape(*x.shape[:-1], 2, last // 2) + first_x = split_x[..., :1, :] + second_x = split_x[..., 1:, :] + + cos_u = cos.unsqueeze(-2) + sin_u = sin.unsqueeze(-2) + + out = split_x * cos_u + first_out = out[..., :1, :] + second_out = out[..., 1:, :] + first_out.addcmul_(-sin_u, second_x) + second_out.addcmul_(sin_u, first_x) + + out = out.reshape(*out.shape[:-2], last) + if needs_reshape: + out = out.swapaxes(1, 2).reshape(b, t, -1) + return out.to(dtype=x_dtype) + + +def _patched_get_ada_values( + self, + scale_shift_table: torch.Tensor, + batch_size: int, + timestep: torch.Tensor, + indices: slice, +) -> tuple[torch.Tensor, ...]: + num_ada_params = int(scale_shift_table.shape[0]) + ada_values = ( + scale_shift_table[indices] + .unsqueeze(0) + .unsqueeze(0) + .to(device=timestep.device, dtype=timestep.dtype) + + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[ + :, :, indices, : + ] + ).unbind(dim=2) + return ada_values + + +def _patched_ltx2_adaln_single_forward( + self, + timestep: torch.Tensor, + hidden_dtype: torch.dtype | None = None, +): + """Match ltx_core AdaLayerNormSingle embedding path.""" + from ltx_core.model.transformer.timestep_embedding import get_timestep_embedding + + t = timestep.reshape(-1).to(dtype=torch.float32) + t_freq = get_timestep_embedding( + t, + 256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + if hidden_dtype is not None: + t_freq = t_freq.to(dtype=hidden_dtype) + + te = self.emb.timestep_embedder + x = F.silu(_linear_out(te.linear_1, t_freq)) + embedded_timestep = _linear_out(te.linear_2, x).to(dtype=self.linear.weight.dtype) + out = _linear_out(self.linear, F.silu(embedded_timestep)) + + if timestep.ndim == 0: + batch = 1 + elif timestep.ndim == 1: + batch = 1 + else: + batch = timestep.shape[0] + out = out.view(batch, -1, out.shape[-1]) + embedded_timestep = embedded_timestep.view(batch, -1, embedded_timestep.shape[-1]) + return out, embedded_timestep + + +def _make_patched_ltx2_attention_forward(orig_forward: Callable[..., torch.Tensor]): + def _patched_forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, + k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, + perturbation_mask: torch.Tensor | None = None, + all_perturbed: bool = False, + skip_sequence_parallel_override: bool = False, + gather_context_kv_for_sp: bool = False, + ) -> torch.Tensor: + from sglang.multimodal_gen.runtime.distributed import get_tp_world_size + from sglang.multimodal_gen.runtime.models.dits.ltx_2 import ( + apply_interleaved_rotary_emb, + ) + + if get_tp_world_size() > 1 or gather_context_kv_for_sp or self.use_local_attention: + return orig_forward( + self, + x, + context=context, + mask=mask, + pe=pe, + k_pe=k_pe, + perturbation_mask=perturbation_mask, + all_perturbed=all_perturbed, + skip_sequence_parallel_override=skip_sequence_parallel_override, + gather_context_kv_for_sp=gather_context_kv_for_sp, + ) + + gate_input = x + context_ = x if context is None else context + v = _linear_out(self.to_v, context_) + use_attention = not all_perturbed + + if use_attention: + q = _linear_out(self.to_q, x) + k = _linear_out(self.to_k, context_) + + if self.qk_norm: + assert self.q_norm is not None and self.k_norm is not None + q = self.q_norm(q) + k = self.k_norm(k) + + if pe is not None: + cos, sin = pe + k_cos, k_sin = pe if k_pe is None else k_pe + if cos.dim() == 3: + q = apply_interleaved_rotary_emb(q, (cos, sin)) + k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) + else: + q = _ltxcore_apply_split_rotary_emb(q, (cos, sin)) + k = _ltxcore_apply_split_rotary_emb(k, (k_cos, k_sin)) + + out = _ltx_pytorch_sdpa(q, k, v, self.local_heads, mask) + + if perturbation_mask is not None: + if perturbation_mask.ndim == out.ndim - 1: + perturbation_mask = perturbation_mask.unsqueeze(-1) + out = out * perturbation_mask + v * (1 - perturbation_mask) + else: + out = v + + if self.to_gate_logits is not None: + gate_logits = _linear_out(self.to_gate_logits, gate_input) + b, t = out.shape[:2] + out = out.view(b, t, self.local_heads, self.dim_head) + out = out * (2.0 * torch.sigmoid(gate_logits).unsqueeze(-1)) + out = out.view(b, t, self.local_heads * self.dim_head) + + return _linear_out(self.to_out[0], out) + + return _patched_forward + + +def _make_patched_ltx2_block_forward(orig_forward: Callable[..., tuple[torch.Tensor, torch.Tensor]]): + def _patched_forward(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: + args = list(args) + if len(args) >= 6 and isinstance(args[0], torch.Tensor) and isinstance(args[4], torch.Tensor): + args[4] = expand_temb_for_hidden(args[4], args[0]) + if len(args) >= 7 and isinstance(args[1], torch.Tensor) and isinstance(args[5], torch.Tensor): + args[5] = expand_temb_for_hidden(args[5], args[1]) + if "temb" in kwargs: + hidden_states = kwargs.get("hidden_states") + if hidden_states is None and args: + hidden_states = args[0] + if isinstance(hidden_states, torch.Tensor): + kwargs = dict(kwargs) + kwargs["temb"] = expand_temb_for_hidden(kwargs["temb"], hidden_states) + if "temb_audio" in kwargs: + audio_hidden_states = kwargs.get("audio_hidden_states") + if audio_hidden_states is None and len(args) >= 2: + audio_hidden_states = args[1] + if isinstance(audio_hidden_states, torch.Tensor): + kwargs = dict(kwargs) + kwargs["temb_audio"] = expand_temb_for_hidden( + kwargs["temb_audio"], audio_hidden_states + ) + return orig_forward(self, *args, **kwargs) + + return _patched_forward + + +def apply() -> None: + global _APPLIED + if _APPLIED: + return + + from sglang.multimodal_gen.runtime.models.dits import ltx_2 as ltx2_mod + + if "rms_norm" not in _ORIGINALS: + _ORIGINALS["rms_norm"] = ltx2_mod.rms_norm + + def _patched_rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: + return _ltxcore_rms_norm(x, eps=eps) + + ltx2_mod.rms_norm = _patched_rms_norm + + if "apply_split_rotary_emb" not in _ORIGINALS: + _ORIGINALS["apply_split_rotary_emb"] = ltx2_mod.apply_split_rotary_emb + + def _patched_apply_split_rotary_emb( + x: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + return _ltxcore_apply_split_rotary_emb(x, freqs) + + ltx2_mod.apply_split_rotary_emb = _patched_apply_split_rotary_emb + + adaln_cls = ltx2_mod.LTX2AdaLayerNormSingle + if "LTX2AdaLayerNormSingle.forward" not in _ORIGINALS: + _ORIGINALS["LTX2AdaLayerNormSingle.forward"] = adaln_cls.forward + adaln_cls.forward = _patched_ltx2_adaln_single_forward + + block_cls = ltx2_mod.LTX2TransformerBlock + if "LTX2TransformerBlock.get_ada_values" not in _ORIGINALS: + _ORIGINALS["LTX2TransformerBlock.get_ada_values"] = block_cls.get_ada_values + block_cls.get_ada_values = _patched_get_ada_values + + if "LTX2TransformerBlock.forward" not in _ORIGINALS: + _ORIGINALS["LTX2TransformerBlock.forward"] = block_cls.forward + block_cls.forward = _make_patched_ltx2_block_forward(block_cls.forward) + + attn_cls = ltx2_mod.LTX2Attention + if "LTX2Attention.forward" not in _ORIGINALS: + _ORIGINALS["LTX2Attention.forward"] = attn_cls.forward + attn_cls.forward = _make_patched_ltx2_attention_forward(attn_cls.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..9c139335 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_rollout_cond_kwargs.py @@ -0,0 +1,50 @@ +"""Ensure LTX rollout denoising_env carries text context for miles train replay.""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +_APPLIED = False + + +def _prompt_embeds_tensor(batch: Any) -> Any | None: + pe = getattr(batch, "prompt_embeds", None) + if pe is None: + return None + return pe[0] if isinstance(pe, list) else pe + + +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, "_attach_ltx_rollout_cond_kwargs"): + logger.warning( + "LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs is missing; " + "rollout denoising_env may lack encoder_hidden_states. " + "Upgrade sglang-diffusion or check the installed version." + ) + _APPLIED = True + return + + orig_attach = LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs + + def _attach_ltx_rollout_cond_kwargs(self, ctx, batch): + orig_attach(self, ctx, batch) + if not (batch.rollout and batch.rollout_return_denoising_env): + return + if ctx.pos_cond_kwargs.get("encoder_hidden_states") is None: + embeds = _prompt_embeds_tensor(batch) + if embeds is not None: + ctx.pos_cond_kwargs["encoder_hidden_states"] = embeds + + LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs = _attach_ltx_rollout_cond_kwargs + _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 87f04779..bee9ea10 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -9,6 +9,12 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.launch_server import kill_process_tree +from miles.backends.sglang_diffusion_utils.configs import ltx as ltx_config +from miles.backends.sglang_diffusion_utils.monkey_patches import ( + ROLLOUT_PATCH_GROUP_ENV, + apply_rollout_patch_group, + resolve_rollout_patch_group, +) from miles.ray.ray_actor import RayActor from miles.utils.http_utils import get_host_info @@ -39,13 +45,27 @@ 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 - apply_sgld_monkey_patches() + apply_rollout_patch_group(os.environ.get(ROLLOUT_PATCH_GROUP_ENV)) + + # Colocate weight sync keeps the full CUDA_VISIBLE_DEVICES (so CUDA IPC works + # across GPUs); pin the DiT to its assigned local cuda index instead. + local_cuda_rank = os.environ.get("MILES_SGLANG_LOCAL_CUDA_RANK") + if local_cuda_rank is not None: + from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker + + pinned_rank = int(local_cuda_rank) + _orig_init = GPUWorker.__init__ + + def _patched_init(self, local_rank, rank, master_port, server_args): + return _orig_init(self, pinned_rank, rank, master_port, server_args) + + GPUWorker.__init__ = _patched_init + 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, patch_group: str | None = None): # 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 @@ -53,15 +73,17 @@ 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: - # 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 - # launch_server's own module to point at our wrapper — pickle then carries - # the miles qualname across to the grandchild, which applies the patch before - # calling the real scheduler entrypoint. + # 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 + # launch_server's own module to point at our wrapper — pickle then carries + # the miles qualname across to the grandchild, which applies the patches (and + # colocate GPU pin) before calling the real scheduler entrypoint. + if patch_group is not None or os.environ.get("MILES_SGLANG_LOCAL_CUDA_RANK") is not None: import sglang.multimodal_gen.runtime.launch_server as _ls_mod _ls_mod.run_scheduler_process = _scheduler_process_with_sgld_monkey_patches + if patch_group is not None: + os.environ[ROLLOUT_PATCH_GROUP_ENV] = patch_group from sglang.multimodal_gen.runtime.launch_server import launch_server launch_server(server_args) @@ -69,14 +91,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, + patch_group: str | None = None, ) -> 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, patch_group), ) p.start() @@ -152,15 +174,12 @@ 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 = bool(getattr(self.args, "apply_sgld_monkey_patches", False)) - if apply_sgld_monkey_patches: - logger.info( - "Launching sglang-d with sgl-d → diffusers monkey patches " - "(--apply-sgld-monkey-patches)" - ) + patch_group = resolve_rollout_patch_group(self.args) + if patch_group is not None: + logger.info(f"Launching sglang-d with rollout patch group: {patch_group}") self.process = launch_server_process( ServerArgs.from_kwargs(**server_args_dict), - apply_sgld_monkey_patches=apply_sgld_monkey_patches, + patch_group=patch_group, ) if self.node_rank == 0 and self.router_ip and self.router_port: @@ -179,13 +198,14 @@ def _pin_to_assigned_gpu(self): cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not cvd: return - visible = [x.strip() for x in cvd.split(",") if x.strip()] - local_idx = _to_local_gpu_id(self.base_gpu_id) - pinned = visible[local_idx] - os.environ["CUDA_VISIBLE_DEVICES"] = pinned + local_id = _to_local_gpu_id(self.base_gpu_id) + # Keep the full CUDA_VISIBLE_DEVICES so colocated weight sync can CUDA-IPC + # buckets from FSDP ranks on other GPUs; pin the DiT to its local cuda + # index via MILES_SGLANG_LOCAL_CUDA_RANK (applied in the scheduler child). + os.environ["MILES_SGLANG_LOCAL_CUDA_RANK"] = str(local_id) logger.info( - f"Engine rank={self.rank}: pinned CUDA_VISIBLE_DEVICES={pinned} " - f"(base_gpu_id={self.base_gpu_id}, local_idx={local_idx})" + f"Engine rank={self.rank}: rollout cuda:{local_id} " + f"(base_gpu_id={self.base_gpu_id}, CUDA_VISIBLE_DEVICES={cvd})" ) def _make_request(self, endpoint: str, payload: dict | None = None): @@ -323,6 +343,12 @@ def _compute_server_args(args, host, port, nccl_port): "warmup": False, } + # LTX rollout loads the Lightricks repo dir + LTX2 pipeline / transformer + # weights / gemma text encoder. No-op for non-LTX models. + if ltx_config.is_ltx_model(args): + kwargs["model_path"] = ltx_config.resolve_sglang_model_path(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 77e35710..924fecc7 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -14,7 +14,6 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS -from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine from miles.rollout.base_types import call_rollout_fn from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor @@ -434,13 +433,23 @@ def _log_images( 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 = t[:, 0, :, :].float().cpu().numpy().transpose(1, 2, 0) + try: + from miles.rollout.rm_hub.video_pickscore import ( + fchw_frame_to_hwc_uint8, + generated_output_to_fchw, + ) + + frame = fchw_frame_to_hwc_uint8(generated_output_to_fchw(t)[0]) + except (ValueError, TypeError): + if t.ndim != 4: + 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) + reward = s.get_reward_value(self.args, reward_key=reward_key) images.append(wandb.Image(frame, caption=f"{str(s.prompt)[:160]} | reward={reward}")) if images: log_dict[media_key] = images @@ -553,16 +562,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, From d0d399c2b27082c8c7bef3dfacca105e2b9e817a Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:18:27 +0000 Subject: [PATCH 07/31] feat(ltx): rollout data path, SDE steps, dynamics, and video reward Parse LTX cond fields and sde_step_indices from rollout responses. Add ltx_sde_candidates step strategy and canonical dynamics passthrough to sglang-d. Support multi-frame video PickScore for LTX GRPO rewards. --- miles/rollout/rm_hub/pickscore.py | 149 ++++++++++++++++++++-- miles/rollout/rm_hub/video_pickscore.py | 67 ++++++++++ miles/rollout/sglang_diffusion_rollout.py | 46 ++++++- miles/rollout/step_strategy_hub.py | 52 ++++++++ miles/utils/diffusion_rollout_response.py | 7 +- 5 files changed, 307 insertions(+), 14 deletions(-) create mode 100644 miles/rollout/rm_hub/video_pickscore.py diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index 76ddb122..f6d61dd2 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -9,6 +9,13 @@ import torch from PIL import Image +from miles.rollout.rm_hub.video_pickscore import ( + fchw_frame_to_hwc_uint8, + fchw_to_pil_frames, + generated_output_to_fchw, + is_video_generated_output, + sample_frame_indices, +) from miles.utils.misc import SingletonMeta from miles.utils.types import Sample @@ -22,19 +29,12 @@ def _feature_tensor(features): 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 +80,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 = 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 = 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 +172,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 +212,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 +247,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/rm_hub/video_pickscore.py b/miles/rollout/rm_hub/video_pickscore.py new file mode 100644 index 00000000..73037671 --- /dev/null +++ b/miles/rollout/rm_hub/video_pickscore.py @@ -0,0 +1,67 @@ +"""Video PickScore helpers (LTX / sglang FHWC and trainer [C,F,H,W] layouts).""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import torch +from PIL import Image + + +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 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 diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 887cfb18..c1c40f18 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -90,6 +90,28 @@ def build_rollout_sampling_params( "rollout_return_dit_trajectory": True, } ) + if model_type == "ltx": + from miles.utils.sde_log_prob import normalize_dynamics_type + + # Canonical names match sglang-d rollout_sde_type, so pass through + # with no translation table (keeps train/rollout on one vocabulary). + dynamics = normalize_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 (train recompute supports it). " + "Add the sglang-d sampling branch before using it for rollout." + ) + 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) + # Disable flag is propagated via MILES_LTX_DISABLE_AV_CROSS on rollout engines + # (see miles/ray/rollout.py). Do not pass via extra_sampling_params — master + # sglang SamplingParams does not accept ltx2_disable_av_cross_attn. if extra_sampling_params: sampling_params["extra_sampling_params"] = extra_sampling_params @@ -132,6 +154,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 @@ -150,6 +173,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: @@ -166,6 +190,23 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: self.remaining_batch_size += len(samples) +def _call_step_strategy( + step_strategy_fn: Callable, + args: Namespace, + sample: Sample, + num_steps: int, + seed: int, + rollout_id: int, +) -> tuple[list[int] | None, list[int] | None]: + """Invoke a step-strategy hub function; pass ``rollout_id`` when supported.""" + params = inspect.signature(step_strategy_fn).parameters + if "rollout_id" in params: + return step_strategy_fn( + args, sample, num_steps, seed, rollout_id=rollout_id + ) + return step_strategy_fn(args, sample, num_steps, seed) + + async def generate_microgroup( args: Namespace, microgroup: list[Sample], sampling_params: dict[str, Any], *, evaluation: bool = False ) -> list[Sample]: @@ -178,11 +219,13 @@ async def generate_microgroup( # SGL-D TODO: support seed list for multiple samples in one request # currently only support assigning the first seed, SGL-D generates samples with seed, seed+1, seed+2, ... if not evaluation and state.step_strategy_fn is not None: - sde_indices, return_indices = state.step_strategy_fn( + sde_indices, return_indices = _call_step_strategy( + state.step_strategy_fn, args, microgroup[0], int(sampling_params["num_inference_steps"]), int(sampling_params["seed"]), + int(getattr(state, "rollout_id", 0) or 0), ) sampling_params["rollout_sde_step_indices"] = sde_indices sampling_params["rollout_return_step_indices"] = return_indices @@ -306,6 +349,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 = ( diff --git a/miles/rollout/step_strategy_hub.py b/miles/rollout/step_strategy_hub.py index 257c6ad4..922aa675 100644 --- a/miles/rollout/step_strategy_hub.py +++ b/miles/rollout/step_strategy_hub.py @@ -3,6 +3,9 @@ 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``) also +accept ``rollout_id`` via keyword — see ``miles.rollout.sglang_diffusion_rollout``. + Point ``--diffusion-step-strategy-path`` at any such function. """ from __future__ import annotations @@ -10,10 +13,59 @@ 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 ) -> tuple[list[int] | None, list[int] | None]: diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 0ddf9bc5..d456b09f 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -78,11 +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, ), pooled_projections=_parse_tensor_or_list( data.get("pooled_projections"), deserialize_func=deserialize_func ), + ltx_positions=deserialize_func(data.get("ltx_positions") or data.get("positions")), + ltx_denoise_mask=deserialize_func(data.get("ltx_denoise_mask") or data.get("denoise_mask")), + ltx_clean_latent=deserialize_func(data.get("ltx_clean_latent") or data.get("clean_latent")), ) @@ -112,6 +116,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")), ) From 7ddda31f7a2ea334a599b1497cd37e8db6f0f932 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 5 Jun 2026 14:20:20 +0000 Subject: [PATCH 08/31] chore: add LTX align skill and dev run script Add align-ltx-train-rollout skill for diagnosing train/rollout misalignment and a dev script to run LTX-2.3 CPS GRPO with sglang rollout + weight sync. --- .../skills/align-ltx-train-rollout/SKILL.md | 86 ++++++ scripts/run-diffusion-grpo-ltx23-sglang.sh | 267 ++++++++++++++++++ 2 files changed, 353 insertions(+) create mode 100644 .claude/skills/align-ltx-train-rollout/SKILL.md create mode 100644 scripts/run-diffusion-grpo-ltx23-sglang.sh diff --git a/.claude/skills/align-ltx-train-rollout/SKILL.md b/.claude/skills/align-ltx-train-rollout/SKILL.md new file mode 100644 index 00000000..8e5ac790 --- /dev/null +++ b/.claude/skills/align-ltx-train-rollout/SKILL.md @@ -0,0 +1,86 @@ +--- +name: align-ltx-train-rollout +description: Diagnose and fix train-vs-rollout forward numerical misalignment for LTX-2.3 diffusion GRPO on sglang rollout, driving model_output_cosine_sim to ~0.9998 and log_prob_mean_abs_diff small. Use when GRPO clipfrac stays near 1.0, train/log_prob_mean_abs_diff is large, model_output_cosine_sim is low, or when aligning the miles FSDP train forward with the sglang rollout forward for LTX video diffusion. +--- + +# align-ltx-train-rollout + +把 miles 训练侧重算的 `log_prob` / `model_output` 与 sglang live rollout 对齐,使 GRPO 的 importance ratio 可信(`clipfrac` 不再长期 = 1.0)。 + +## 核心心智模型(先读这一条) + +**对齐 gap 几乎从不是 checkpoint / 权重错。** 按以下三层顺序定位,不要一上来换权重: + +1. **log_prob 公式层** — SDE dynamics 类型、`sigma_min` 用错 → `log_prob_mean_abs_diff` 巨大(2~15) +2. **DiT forward 层** — temb/AdaLN 的 shape 与语义、算子 parity、AV cross-attn、attention backend → `cosine` 低(~0.96) +3. **live pipeline 后处理层** — guider(cfg/stg/modality/rescale)改了 velocity 语义 → offline 高但 live 低(~0.94) + +## 诊断决策树 + +开 `--diffusion-debug-mode`,看训练 log 的 `train/model_output_cosine_sim`、`train/log_prob_mean_abs_diff`、`train/clipfrac`(计算位置 `miles/backends/fsdp_utils/actor.py:836-847`)。然后: + +``` +cosine 高(>0.999) 但 log_prob_diff 大? + └─→ 第①层:SDE 公式 / sigma_min。查 sde_log_prob.py + sglang scheduler_rl_mixin。 + +cosine 低(<0.99)? + └─ 先做 offline injected 对比(同 latent/kwargs 喂两侧 DiT,排除 checkpoint): + bash scripts/capture-and-compare-ltx23-forward.sh + ├─ injected 也低 → 第②层:temb/parity/AV/attention(见下方必备开关) + └─ injected 高(>0.999) 但 live 低 → 第③层:guider 后处理(identity guider) +``` + +## 对齐必备开关(缺一项就掉精度) + +| 开关 | 作用 | 对应层 | +|------|------|--------| +| `MILES_APPLY_LTX2_LTXCORE_PARITY=1` | temb expand `[B,1,D]→[B,T,D]` + AdaLN/RMSNorm/RoPE/SDPA 数值对齐 | ② | +| `MILES_LTX_IDENTITY_GUIDER=1` | 强制 stage1 guider 为 identity(cfg=1/stg=0/modality=1/rescale=0) | ③ | +| `--ltx-disable-av-cross-attn` (+ `MILES_LTX_DISABLE_AV_CROSS=1`) | train video-only 与 rollout 算子图一致 | ② | +| `SGLANG_ATTENTION_BACKEND=torch_sdpa` | 避免 FlashAttention vs SDPA 数值差 | ② | +| train 与 sglang 用**同一份 `.safetensors`** | 避免 HF materialized overlay 与单文件差 ~1e-4(由 `configs/ltx.py` resolve) | ② | +| train 与 rollout **同 `--diffusion-num-steps` / `--ltx-dynamics-type` / `--ltx-sigma-min`** | 步数/动力学/σ_min 一致 | ①② | + +**经验**:`dev` ckpt + 24 步比 `distilled` + 8 步对齐好得多(完整路径 velocity 更平滑)。优先用 dev 验证。 + +## 验收标准 + +| 指标 | 达标 | dev/512×768×57f/24步实测 | +|------|------|--------------------------| +| `model_output_cosine_sim` | ≥ 0.999 | 0.9995 ~ 0.9999 | +| `log_prob_mean_abs_diff` | < 5e-3 | 6e-6 ~ 1.7e-3 | +| `clipfrac` | 明显 < 1.0 | 0 ~ 0.125 | + +## 快速命令 + +```bash +cd /sgl-workspace/master_miles/miles_diffusion +export PYTHONPATH=/sgl-workspace/master_sglang/sglang/python${PYTHONPATH:+:$PYTHONPATH} + +# 纯前向对齐验证(不更新权重,最快看 cosine / log_prob_diff) +CUDA_VISIBLE_DEVICES= MILES_DIFFUSION_DEBUG=1 \ +LTX_DISABLE_AV_CROSS_ATTN=1 MILES_LTX_IDENTITY_GUIDER=1 USE_LORA=1 SKIP_OPTIMIZER=1 \ +NUM_ROLLOUT=3 ROLLOUT_BATCH_SIZE=1 N_SAMPLES_PER_PROMPT=2 GLOBAL_BATCH_SIZE=2 \ +nohup bash scripts/run-diffusion-grpo-ltx23-sglang-dev-flowsde.sh > logs/align_verify.log 2>&1 & + +# 离线 capture + compare(定位 gap 在 DiT raw 还是后处理) +bash dist/scripts/capture-and-compare-ltx23-forward.sh +``` + +监控:`grep -E 'model_output_cosine_sim|log_prob_mean_abs_diff|clipfrac' logs/*.log` + +## 常见陷阱(实战踩过) + +- **temb 错 1000 倍**:AdaLN 输入应是 `σ×1000` 而非 `σ`;错了 cosine≈0(与"权重错"是不同量级)。 +- **temb shape**:ltx_core 走 `[B,T,D]`,sglang 构 `[B,1,D]`,即使 σ 均匀,block 内 broadcast 行为不同 → 必须 `expand_temb_for_hidden`。 +- **guider 默认值**:LTX23 one-stage 默认 `video_cfg_scale=3 / modality=3 / rescale=0.7`,在 x0 上改 velocity 语义 → live cosine 0.94。train `forward_velocity` 等价 cfg=1/无STG/无modality/无rescale。 +- **identity guider 不能经 `extra_sampling_params` 传**:sglang `SamplingParams` 基类拒绝 LTX23 guider 字段(`400 unexpected keyword 'video_cfg_scale'`)→ 必须用 monkey patch `patch_ltx2_identity_guider.py` override `_get_ltx2_stage1_guider_params`。 +- **Flow-SDE 公式 / sigma_min**:`--ltx-dynamics-type Flow-SDE` 时 rollout 误用 SD3 `sde` 公式、或误用 `sigmas[-2]=0.1` 作 σ_min → `log_prob_mean_abs_diff` 2~15。sglang 侧需 `flow_sde` + `rollout_sigma_min`。 +- **guider 修复后必须重新 capture dump 再 compare**:旧 dump 是默认 guider 下生成的。 + +## 参考文档(深度细节，本地 `dist/docs/`，不进 git) + +- [dist/docs/ltx23_changes_overview.md](../../../dist/docs/ltx23_changes_overview.md) — 全部代码改动按 DEBUG/TRAIN/ROLLOUT 三分类 +- [dist/docs/ltx23_train_rollout_alignment_journey.md](../../../dist/docs/ltx23_train_rollout_alignment_journey.md) — Phase A/B/C 排查历程(SDE→temb→guider) +- [dist/docs/ltx23_sglang_rollout_train_troubleshooting.md](../../../dist/docs/ltx23_sglang_rollout_train_troubleshooting.md) — 两侧工程问题 P1–P13 + 跨边界 C1–C8 +- [dist/docs/ltx23_forward_alignment_test_report.md](../../../dist/docs/ltx23_forward_alignment_test_report.md) — 数值实验矩阵与 block-wise 二分 diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh new file mode 100644 index 00000000..26fde4c6 --- /dev/null +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -0,0 +1,267 @@ +#!/usr/bin/env bash +# LTX-2.3 sglang-rollout GRPO — dev checkpoint, 512x768x57f, 24 steps, CPS. +# +# Mirrors the legacy trainer-rollout reward run +# (/sgl-workspace/miles/scripts/run-diffusion-grpo-ltx23-trainer-rollout.sh): +# CPS dynamics, 3 SDE steps from candidates 0–9, clip-range 1e-4. +# Rollout goes through sglang with weight sync; train/rollout forward alignment +# fixes stay on (ltxcore parity + AV-off + identity guider). +# +# GPU layout: single physical GPU colocate (train FSDP world_size=1 and sglang +# rollout time-share one GPU via offload). Set NUM_GPUS>1 for multi-GPU +# colocate if 512x768x57f OOMs on one card. +# +# Usage: +# CUDA_VISIBLE_DEVICES=1 USE_LORA=1 NUM_ROLLOUT=8 \ +# MILES_DIFFUSION_DEBUG=1 LTX_DISABLE_AV_CROSS_ATTN=1 \ +# MILES_LTX_IDENTITY_GUIDER=1 \ +# nohup bash scripts/run-diffusion-grpo-ltx23-sglang.sh \ +# > logs/ltx23_dev_cps_$(date +%Y%m%d_%H%M%S).log 2>&1 & +# +# Key overridable env: +# LTX_MODEL_PATH — dev 22B safetensors +# HEIGHT WIDTH FRAMES — 512 768 57 +# NUM_STEPS — 24 +# LTX_NUM_SDE_STEPS — 3 +# LTX_SDE_STEP_CANDIDATES — 0,1,2,3,4,5,6,7,8,9 +# CLIP_RANGE — 1e-4 +# ROLLOUT_BATCH_SIZE — unique prompts per rollout (default: 16) +# N_SAMPLES_PER_PROMPT — GRPO group size, aligned with verl (default: 8) +# NUM_STEPS_PER_ROLLOUT — optimizer steps per rollout (default: 2 → gbs=64) +# NUM_GPUS — 1 + +MILES_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +echo "[kill] hunting for stale miles processes under cwd=${MILES_ROOT}" +for pid in $(ls /proc 2>/dev/null | grep -E '^[0-9]+$'); do + # timeout guards against readlink hanging on a process whose cwd points at a + # stale/unresponsive mount — otherwise this loop can wedge the whole shell. + link=$(timeout 2 readlink "/proc/${pid}/cwd" 2>/dev/null) || continue + exe=$(timeout 2 readlink "/proc/${pid}/exe" 2>/dev/null) || continue + case "${link}" in + "${MILES_ROOT}"|"${MILES_ROOT}"/*) + case "${exe}" in + */python*|*/ray*) + echo "[kill] ${pid} (${exe}) cwd=${link}" + kill -9 "${pid}" 2>/dev/null || true + ;; + esac + ;; + esac +done +sleep 3 + +ps -eo ppid,state,comm --no-headers \ + | awk '$2=="Z" && $1!=1 && $3~/ray|python|sglang/ {print $1}' \ + | sort -u | xargs -r kill -9 2>/dev/null || true +sleep 2 + +set -euo pipefail + +ROOT_DIR="${MILES_ROOT}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +SGLANG_PYTHON="${SGLANG_PYTHON:-/sgl-workspace/master_sglang/sglang/python}" +export PYTHONPATH="${SGLANG_PYTHON}${PYTHONPATH:+:${PYTHONPATH}}" + +# All heavy I/O lives on /data — workspace overlay is small and often full. +MILES_DATA_DISK_ROOT="${MILES_DATA_DISK_ROOT:-/data/wenhao/miles_diffusion}" +RAY_BIG_TMP="${RAY_BIG_TMP:-/data/wenhao/miles_ray_tmp}" +TMP_BIG="${TMP_BIG:-/data/wenhao/tmp}" +SGL_DIFF_CACHE="${SGLANG_DIFFUSION_CACHE_ROOT:-/data/wenhao/sgl_diffusion_cache}" +HF_HOME="${HF_HOME:-/data/wenhao/hf_home}" +LOG_DIR="${LOG_DIR:-${MILES_DATA_DISK_ROOT}/logs}" +WANDB_DIR="${WANDB_DIR:-${MILES_DATA_DISK_ROOT}/wandb}" +CKPT_ROOT="${CKPT_ROOT:-${MILES_DATA_DISK_ROOT}/ckpt}" +mkdir -p "${MILES_DATA_DISK_ROOT}" "${RAY_BIG_TMP}" "${TMP_BIG}" "${SGL_DIFF_CACHE}" \ + "${HF_HOME}" "${LOG_DIR}" "${WANDB_DIR}" "${CKPT_ROOT}" +export RAY_TMPDIR="${RAY_BIG_TMP}" +export TMPDIR="${TMP_BIG}" +export SGLANG_DIFFUSION_CACHE_ROOT="${SGL_DIFF_CACHE}" +export HF_HOME +export WANDB_DIR +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" +mkdir -p "${HUGGINGFACE_HUB_CACHE}" +export MILES_APPLY_LTX2_LTXCORE_PARITY="${MILES_APPLY_LTX2_LTXCORE_PARITY:-1}" +export RAY_object_spilling_config="$(python -c "import json,os; print(json.dumps({'type':'filesystem','params':{'directory_path':[os.environ['RAY_TMPDIR']]}}))")" +ray stop --force 2>/dev/null || true +sleep 2 + +# ── dev checkpoint (borrowed from legacy reward run) ───────────────────── +LTX_MODEL_PATH="${LTX_MODEL_PATH:-/sgl-workspace/rollout_compare/models/LTX-2.3/ltx-2.3-22b-dev.safetensors}" +# sglang text_encoder: use materialized Lightricks overlay (local gemma_for_ltx23 +# symlinks often point at stale HF cache and break rollout startup). +LTX_MATERIALIZED_ROOT="${LTX_MATERIALIZED_ROOT:-/data/wenhao/sgl_diffusion_cache/materialized_models/Lightricks__LTX-2.3-10cce1713d7efa14}" +GEMMA_ROOT="${GEMMA_ROOT:-${LTX_MATERIALIZED_ROOT}/text_encoder}" +MILES_DATA_ROOT="${MILES_DATA_ROOT:-/sgl-workspace/miles}" +PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" + +NUM_ROLLOUT="${NUM_ROLLOUT:-8}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-16}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-2}" +SAMPLES_PER_ROLLOUT=$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT)) +DERIVED_GLOBAL_BATCH_SIZE=$((SAMPLES_PER_ROLLOUT / NUM_STEPS_PER_ROLLOUT)) +MICRO_BATCH_SIZE_SAMPLE="${MICRO_BATCH_SIZE_SAMPLE:-1}" +MICRO_BATCH_SIZE_TSTEP="${MICRO_BATCH_SIZE_TSTEP:-1}" + +# ── borrowed-from-legacy generation config ─────────────────────────────── +HEIGHT="${HEIGHT:-512}" +WIDTH="${WIDTH:-768}" +FRAMES="${FRAMES:-57}" +NUM_STEPS="${NUM_STEPS:-24}" +# ── trainer-rollout SDE config (CPS + candidate sampling) ──────────────── +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}" +CLIP_RANGE="${CLIP_RANGE:-1e-4}" + +NUM_GPUS="${NUM_GPUS:-1}" +# Multi-GPU colocate: one sglang engine PER GPU (each card runs both the sglang +# rollout engine and an FSDP trainer shard). per-engine=1 => num_engines=NUM_GPUS. +# Set ROLLOUT_NUM_GPUS_PER_ENGINE=NUM_GPUS instead for a single TP-sharded engine. +ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" +# Periodic checkpoint (LoRA adapter) so the run is resumable via LOAD_CKPT. +# (The earlier run had no --save-interval, so nothing was ever saved.) +SAVE_INTERVAL="${SAVE_INTERVAL:-5}" + +if [[ ! -f "${PROMPT_DATA}" ]]; then + python "${MILES_DATA_ROOT}/tools/prepare_vidgen_jsonl.py" +fi + +RUN_NAME="ltx23_dev_cps_${NUM_ROLLOUT}step_$(date +%Y%m%d_%H%M%S)" +SAVE_DIR="${CKPT_ROOT}/${RUN_NAME}" +mkdir -p "${SAVE_DIR}" + +echo "[run] dev+cps CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} NUM_GPUS=${NUM_GPUS}" +echo "[run] dit=${LTX_MODEL_PATH}" +echo "[run] gemma=${GEMMA_ROOT}" +echo "[run] log=${LOG_DIR}" +echo "[run] wandb=${WANDB_DIR}" +echo "[run] save=${SAVE_DIR}" +echo "[run] ${HEIGHT}x${WIDTH}x${FRAMES}f steps=${NUM_STEPS} sde_steps=${LTX_NUM_SDE_STEPS} candidates=${LTX_SDE_STEP_CANDIDATES} clip=${CLIP_RANGE}" +echo "[run] batch: rollout=${ROLLOUT_BATCH_SIZE} n_samples=${N_SAMPLES_PER_PROMPT} samples/rollout=${SAMPLES_PER_ROLLOUT} optim_steps/rollout=${NUM_STEPS_PER_ROLLOUT} gbs=${DERIVED_GLOBAL_BATCH_SIZE}" + +DEBUG_ARGS=() +if [[ "${MILES_DIFFUSION_DEBUG:-0}" == "1" ]]; then + DEBUG_ARGS+=(--diffusion-debug-mode) +fi + +DUMP_ARGS=() +if [[ -n "${LTX_FORWARD_DUMP_ROOT:-}" ]]; then + mkdir -p "${LTX_FORWARD_DUMP_ROOT}" + DUMP_ARGS+=(--dump-details "${LTX_FORWARD_DUMP_ROOT}") +fi + +LTX_AV_ARGS=() +if [[ "${LTX_DISABLE_AV_CROSS_ATTN:-0}" == "1" ]]; then + LTX_AV_ARGS+=(--ltx-disable-av-cross-attn) + export MILES_LTX_DISABLE_AV_CROSS=1 +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 + +SKIP_OPT_ARGS=() +if [[ "${SKIP_OPTIMIZER:-0}" == "1" ]]; then + SKIP_OPT_ARGS+=(--debug-skip-optimizer-step) +fi + +# Resume: point LOAD_CKPT at a previously saved --save dir (LoRA adapter). +LOAD_ARGS=() +if [[ -n "${LOAD_CKPT:-}" ]]; then + LOAD_ARGS+=(--load "${LOAD_CKPT}") +fi + +# WandB: enabled when WANDB_API_KEY is set. Mirrors the legacy reward run so the +# reward curve is directly comparable. +WANDB_ARGS=() +if [[ -n "${WANDB_API_KEY:-}" ]]; then + WANDB_ARGS+=( + --use-wandb + --wandb-dir "${WANDB_DIR}" + --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:-5}" + --disable-wandb-random-suffix + ) +fi + +python -u "${ROOT_DIR}/train_diffusion.py" \ + --train-backend fsdp \ + --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ + --diffusion-model "${LTX_MODEL_PATH}" \ + --diffusion-model-type ltx \ + --ltx-gemma-path "${GEMMA_ROOT}" \ + --hf-checkpoint gpt2 \ + --prompt-data "${PROMPT_DATA}" \ + --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}" \ + --micro-batch-size-tstep "${MICRO_BATCH_SIZE_TSTEP}" \ + --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 120 \ + --miles-router-health-check-failure-threshold 30 \ + --sglang-server-concurrency 1 \ + --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ + "${LORA_ARGS[@]}" \ + "${LTX_AV_ARGS[@]}" \ + --lr 2e-4 \ + --adam-beta2 0.999 \ + --weight-decay 1e-4 \ + --diffusion-clip-range "${CLIP_RANGE}" \ + --diffusion-kl-beta 0.0 \ + --diffusion-num-steps "${NUM_STEPS}" \ + --diffusion-step-strategy-path miles.rollout.step_strategy_hub.ltx_sde_candidates \ + --ltx-num-sde-steps "${LTX_NUM_SDE_STEPS}" \ + --ltx-sde-step-candidates "${LTX_SDE_STEP_CANDIDATES}" \ + --ltx-dynamics-type CPS \ + --diffusion-noise-level 0.8 \ + --ltx-sigma-min 0.001 \ + --diffusion-guidance-scale 1.0 \ + --diffusion-height "${HEIGHT}" \ + --diffusion-width "${WIDTH}" \ + --ltx-frames "${FRAMES}" \ + --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:-yuvalkirstain/PickScore_v1}" \ + --pickscore-model-path "${PICKSCORE_MODEL:-yuvalkirstain/PickScore_v1}" \ + --pickscore-num-frames 3 \ + --pickscore-batch-size 8 \ + --pickscore-num-gpus-per-worker 0 \ + --update-weight-buffer-size 2147483648 \ + --save "${SAVE_DIR}" \ + --save-interval "${SAVE_INTERVAL}" \ + "${LOAD_ARGS[@]}" \ + "${DEBUG_ARGS[@]}" \ + "${DUMP_ARGS[@]}" \ + "${SKIP_OPT_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "$@" From 097dbd46cf617cdf3298d835fdfad402abe0ad6c Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 9 Jun 2026 10:42:26 +0000 Subject: [PATCH 09/31] feat(ltx): parse text and audio cond kwargs from rollout denoising_env Extend CondKwargs and rollout JSON parsing so train can consume encoder hidden states and attention masks returned by sglang. --- miles/utils/diffusion_rollout_response.py | 9 +++++++++ miles/utils/types.py | 3 +++ 2 files changed, 12 insertions(+) diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index d456b09f..1dde0f01 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -81,9 +81,18 @@ def _parse_cond_kwargs( 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 ), + # Legacy rollout fields; miles train rebuilds geometry locally for LTX. ltx_positions=deserialize_func(data.get("ltx_positions") or data.get("positions")), ltx_denoise_mask=deserialize_func(data.get("ltx_denoise_mask") or data.get("denoise_mask")), ltx_clean_latent=deserialize_func(data.get("ltx_clean_latent") or data.get("clean_latent")), diff --git a/miles/utils/types.py b/miles/utils/types.py index 2457b578..86cc75c9 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -33,6 +33,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 ltx_positions: torch.Tensor | None = None ltx_denoise_mask: torch.Tensor | None = None From 44c1330a7f289de5558a9809fb6289d8854d3068 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 9 Jun 2026 10:42:28 +0000 Subject: [PATCH 10/31] feat(ltx): rebuild T2V geometry locally and fix sigma timestep scaling Rebuild positions/mask/clean_latent on train instead of trusting rollout env, and divide AdaLN-scale timesteps by 1000 for CPS SDE log_prob parity. --- miles/backends/fsdp_utils/actor.py | 29 +++- miles/backends/fsdp_utils/configs/ltx.py | 80 ++++++++--- .../configs/train_pipeline_config.py | 8 ++ miles/backends/fsdp_utils/ltx_geometry.py | 124 ++++++++++++++++++ 4 files changed, 214 insertions(+), 27 deletions(-) create mode 100644 miles/backends/fsdp_utils/ltx_geometry.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 6cce0bbf..88827131 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -352,6 +352,10 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: ) if sigmas_snapshot is not None: sigmas_ref = sigmas_snapshot.to(device).float() + elif not self.train_pipeline_config.needs_timestep_scaling: + # Trajectory timesteps may be AdaLN-scale (LTX: σ×1000). CPS expects σ∈[0,1]. + sigmas_ref = self.train_pipeline_config.scale_timesteps_for_sde(timesteps_ref) + sigmas_ref = torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) else: sigmas_ref = timesteps_ref / float(num_train_timesteps) sigmas_ref = torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) @@ -464,11 +468,23 @@ def _build_train_grids( dit_trajectories[traj_idx], device ) - # prepare cond kwargs (denoising_env) + # prepare cond kwargs (denoising_env text embeds + local geometry for LTX) denoising_env = denoising_envs[traj_idx] - positive_cond_kwargs_list.append( - train_pipeline_config.prepare_cond_kwargs(denoising_env.pos_cond_kwargs, device) - ) + if hasattr(train_pipeline_config, "build_train_cond_kwargs"): + positive_cond_kwargs_list.append( + train_pipeline_config.build_train_cond_kwargs( + denoising_env.pos_cond_kwargs, + video_latents=latents, + args=self.args, + device=device, + ) + ) + else: + positive_cond_kwargs_list.append( + train_pipeline_config.prepare_cond_kwargs( + denoising_env.pos_cond_kwargs, device + ) + ) if use_cfg: negative_cond_kwargs_list.append( train_pipeline_config.prepare_cond_kwargs(denoising_env.neg_cond_kwargs, device) @@ -683,6 +699,7 @@ def _forward_tile( 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_pipeline_config.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. @@ -761,7 +778,7 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( self.scheduler, noise_pred_flat, - timesteps_flat, + timesteps_for_sde, latents_flat, prev_sample=next_latents_tile.reshape( tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] @@ -785,7 +802,7 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: _, _, prev_sample_mean_ref, _ = train_pipeline_config.sde_step( self.scheduler, ref_noise_pred_flat, - timesteps_flat, + timesteps_for_sde, latents_flat, prev_sample=next_latents_tile.reshape( tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 79c26bea..f34c4617 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -20,6 +20,8 @@ class LTXTrainPipelineConfig(TrainPipelineConfig): is_diffusers_pipeline = False needs_timestep_scaling = False + # Rollout stores σ×1000 in dit_trajectory.timesteps; CPS uses scheduler σ∈[0,1]. + sde_timestep_divisor = 1000.0 supports_cfg = False fsdp_wrap_classes = ["BasicAVTransformerBlock"] @@ -70,25 +72,58 @@ def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> if ctx.ndim == 2: ctx = ctx.unsqueeze(0) kwargs["context"] = ctx - if cond.ltx_positions is not None: - pos = cond.ltx_positions.to(device) - if pos.ndim == 2: - pos = pos.unsqueeze(0) - elif pos.ndim == 3: - pos = pos.unsqueeze(0) - kwargs["positions"] = pos - if cond.ltx_denoise_mask is not None: - mask = cond.ltx_denoise_mask.to(device) - if mask.ndim == 2 and mask.shape[-1] == 1: - mask = mask.squeeze(-1) + 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["denoise_mask"] = mask - if cond.ltx_clean_latent is not None: - cl = cond.ltx_clean_latent.to(device) - if cl.ndim == 2: - cl = cl.unsqueeze(0) - kwargs["clean_latent"] = cl + 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, + *, + video_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 = video_latents[0] if video_latents.ndim >= 2 else video_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 expand_cond_for_timestep_batch(self, cond_kwargs: dict, batch_size: int) -> dict: @@ -161,18 +196,21 @@ def forward_velocity( dtype = latents_input.dtype B = latents_input.shape[0] - sigma = timesteps_input.to(latents_input.dtype) + # 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.view(B, 1) * denoise_mask_2d).to(dtype) + 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, + sigma=sigma_unit.reshape(B), timesteps=adaln_timesteps, positions=cond["positions"].to(dtype), context=cond["context"].to(dtype), @@ -188,7 +226,7 @@ def forward_velocity( 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, min=1e-8).view(B, 1, 1) + 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) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 2ea48b67..b0177aac 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -52,8 +52,16 @@ class TrainPipelineConfig(abc.ABC): fsdp_wrap_classes: list[str] | None = None lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] needs_timestep_scaling: bool = True + # When set, ``dit_trajectory.timesteps`` are on an AdaLN scale (e.g. σ×1000) + # but CPS/SDE log_prob expects σ in 0..1. Divide by this for sde_step only. + sde_timestep_divisor: float | None = None optimizer_state_allowed_missing: list[str] = [] + def scale_timesteps_for_sde(self, timesteps: torch.Tensor) -> torch.Tensor: + if self.sde_timestep_divisor is not None: + return timesteps / float(self.sde_timestep_divisor) + return timesteps + def prepare_trajectory( self, traj: DiTTrajectory, diff --git a/miles/backends/fsdp_utils/ltx_geometry.py b/miles/backends/fsdp_utils/ltx_geometry.py new file mode 100644 index 00000000..788fd352 --- /dev/null +++ b/miles/backends/fsdp_utils/ltx_geometry.py @@ -0,0 +1,124 @@ +"""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, + } From ee30443943299c3186a9f714d9158cbbbc1c2ca1 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 9 Jun 2026 10:42:30 +0000 Subject: [PATCH 11/31] feat(ltx): use model_id rollout engine config instead of pipeline_class_name Point sglang at Lightricks/LTX-2.3 via model_id and overlay materialization, while pinning DiT init to the same dev safetensors as FSDP train. --- .../sglang_diffusion_utils/configs/ltx.py | 70 ++++++++++++++----- .../sglang_diffusion_engine.py | 3 +- 2 files changed, 54 insertions(+), 19 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py index 990c2ccc..248dee85 100644 --- a/miles/backends/sglang_diffusion_utils/configs/ltx.py +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -1,7 +1,10 @@ """LTX-2 sglang-d rollout engine config. -Mirrors ``fsdp_utils/configs/ltx.py`` on the train side: model detection, -weight-path resolution, and extra ``ServerArgs`` fields for LTX2Pipeline. +Rollout engine uses ``model_path=Lightricks/LTX-2.3`` + ``model_id=LTX-2.3`` so +sglang's overlay wrapper materializes a full diffusers tree (``model_index.json``, +VAE, text encoder, connectors). Train FSDP still loads ``--diffusion-model`` as a +single official safetensors file; ``transformer_weights_path`` pins rollout DiT +to that same file for weight parity. """ from __future__ import annotations @@ -12,6 +15,10 @@ logger = logging.getLogger(__name__) +# sglang registry + overlay wrapper (see model_overlay.py). +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() @@ -23,16 +30,36 @@ def is_ltx_model(args) -> bool: return "ltx" in diff_model or diff_model.endswith(".safetensors") +def resolve_ltx_model_id(args) -> str: + """Short registry id for ``ServerArgs.model_id`` (matches ``Lightricks/LTX-2.3``).""" + 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 resolve_sglang_model_path(args) -> str: + """HF hub id for sglang pipeline skeleton (overlay materializes components).""" + 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_ltx_transformer_weights_path( diffusion_model: str | None, *, explicit_path: str | None = None, ) -> str | None: - """Return official safetensors path for sglang ``transformer_weights_path``. + """Return official single-file safetensors for sglang ``transformer_weights_path``. - When miles train loads a single-file LTX checkpoint, sglang should load the - same safetensors via ``transformer_weights_path`` instead of the HF - materialized ``model.safetensors`` overlay (which can differ at ~1e-4 bf16). + Overlay materialized ``transformer/model.safetensors`` is a different checkpoint + variant than dev 22B; miles train loads ``--diffusion-model`` via ltx_core. Point + rollout DiT init + weight sync at the same single-file ckpt as train. """ if explicit_path: path = Path(explicit_path).expanduser() @@ -53,18 +80,11 @@ def resolve_ltx_transformer_weights_path( return None -def resolve_sglang_model_path(args) -> str: - model_path = args.diffusion_model - if is_ltx_model(args) and model_path.endswith(".safetensors"): - return os.path.dirname(model_path) - return model_path - - def server_kwargs_extras(args) -> dict: """Extra ``ServerArgs`` kwargs; call only when ``is_ltx_model(args)``.""" - extras: dict = {"pipeline_class_name": "LTX2Pipeline"} - if getattr(args, "sglang_pipeline_class_name", None): - extras["pipeline_class_name"] = args.sglang_pipeline_class_name + extras: dict = { + "model_id": resolve_ltx_model_id(args), + } explicit = getattr(args, "sglang_transformer_weights_path", None) weights_path = resolve_ltx_transformer_weights_path( @@ -73,9 +93,25 @@ def server_kwargs_extras(args) -> dict: ) if weights_path and not explicit: extras["transformer_weights_path"] = weights_path - logger.info("LTX rollout: transformer_weights_path=%s", weights_path) + logger.info( + "LTX rollout: model_path=%s model_id=%s transformer_weights_path=%s", + resolve_sglang_model_path(args), + extras["model_id"], + weights_path, + ) elif explicit: extras["transformer_weights_path"] = explicit + logger.info( + "LTX rollout: model_path=%s model_id=%s transformer_weights_path=%s", + resolve_sglang_model_path(args), + extras["model_id"], + explicit, + ) + else: + logger.warning( + "LTX rollout: no transformer_weights_path resolved from --diffusion-model; " + "rollout DiT will use overlay default (may diverge from train ckpt)." + ) gemma_path = getattr(args, "ltx_gemma_path", None) if gemma_path: diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index bee9ea10..8d907987 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -343,8 +343,7 @@ def _compute_server_args(args, host, port, nccl_port): "warmup": False, } - # LTX rollout loads the Lightricks repo dir + LTX2 pipeline / transformer - # weights / gemma text encoder. No-op for non-LTX models. + # LTX rollout: HF model id + overlay wrapper; DiT weights via transformer_weights_path. if ltx_config.is_ltx_model(args): kwargs["model_path"] = ltx_config.resolve_sglang_model_path(args) kwargs.update(ltx_config.server_kwargs_extras(args)) From 6890479b754f49e9fb5c802ba9e64d31ac03b47d Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 9 Jun 2026 10:42:33 +0000 Subject: [PATCH 12/31] feat(ltx): rollout train alignment with gs=1 and cond kwargs patch Inject text/audio embeds into denoising_env via monkey patch, drop identity guider patch, force single-forward rollout at gs=1.0, and refresh the LTX23 sglang GRPO run script defaults. --- .../monkey_patches/__init__.py | 6 +-- .../patch_ltx2_rollout_cond_kwargs.py | 43 ++++++++++++------- miles/rollout/sglang_diffusion_rollout.py | 16 +++---- miles/utils/arguments.py | 9 ++-- scripts/run-diffusion-grpo-ltx23-sglang.sh | 30 +++++++------ 5 files changed, 58 insertions(+), 46 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 93759b79..e9127350 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -5,7 +5,7 @@ ``apply_rollout_patch_group``. - ``sgld``: diffusers / SD3 op parity (RMSNorm, RoPE, attention, …). -- ``ltx``: LTX-2 ltx_core parity + RL alignment (identity guider, AV-off, cond kwargs). +- ``ltx``: LTX-2 ltx_core parity + AV-off (rollout uses official gs=1 path). Patch modules are imported inside ``apply_*`` only so ``RolloutManager`` (a CPU-only Ray actor) can import this package without pulling sglang triton kernels. @@ -71,10 +71,9 @@ def apply_sgld_monkey_patches(*, include_ltx2_ltxcore: bool | None = None) -> No def apply_ltx2_rollout_patches() -> None: - """LTX-2 ltx_core parity + RL train/rollout alignment patches.""" + """LTX-2 ltx_core parity + video-only train alignment.""" from miles.backends.sglang_diffusion_utils.monkey_patches import ( patch_ltx2_disable_av_cross, - patch_ltx2_identity_guider, patch_ltx2_ltxcore_parity, patch_ltx2_rollout_cond_kwargs, ) @@ -82,4 +81,3 @@ def apply_ltx2_rollout_patches() -> None: patch_ltx2_ltxcore_parity.apply() patch_ltx2_disable_av_cross.apply() patch_ltx2_rollout_cond_kwargs.apply() - patch_ltx2_identity_guider.apply() 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 index 9c139335..711ffa68 100644 --- 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 @@ -1,4 +1,4 @@ -"""Ensure LTX rollout denoising_env carries text context for miles train replay.""" +"""Ensure LTX rollout denoising_env carries text/audio embeds for miles train replay.""" from __future__ import annotations @@ -10,11 +10,11 @@ _APPLIED = False -def _prompt_embeds_tensor(batch: Any) -> Any | None: - pe = getattr(batch, "prompt_embeds", None) - if pe is None: +def _first_batch_tensor(batch: Any, attr: str) -> Any | None: + value = getattr(batch, attr, None) + if value is None: return None - return pe[0] if isinstance(pe, list) else pe + return value[0] if isinstance(value, list) else value def apply() -> None: @@ -26,25 +26,36 @@ def apply() -> None: LTX2DenoisingStage, ) - if not hasattr(LTX2DenoisingStage, "_attach_ltx_rollout_cond_kwargs"): + if not hasattr(LTX2DenoisingStage, "_prepare_denoising_loop"): logger.warning( - "LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs is missing; " - "rollout denoising_env may lack encoder_hidden_states. " - "Upgrade sglang-diffusion or check the installed version." + "LTX2DenoisingStage._prepare_denoising_loop is missing; " + "rollout denoising_env may lack text/audio cond kwargs." ) _APPLIED = True return - orig_attach = LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs + orig_prepare = LTX2DenoisingStage._prepare_denoising_loop - def _attach_ltx_rollout_cond_kwargs(self, ctx, batch): - orig_attach(self, ctx, batch) + 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 + return ctx + ctx.pos_cond_kwargs = dict(ctx.pos_cond_kwargs) if ctx.pos_cond_kwargs.get("encoder_hidden_states") is None: - embeds = _prompt_embeds_tensor(batch) + embeds = _first_batch_tensor(batch, "prompt_embeds") if embeds is not None: ctx.pos_cond_kwargs["encoder_hidden_states"] = embeds - - LTX2DenoisingStage._attach_ltx_rollout_cond_kwargs = _attach_ltx_rollout_cond_kwargs + 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/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index c1c40f18..39078b78 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -67,14 +67,9 @@ def build_rollout_sampling_params( sampling_params["num_frames"] = int(args.ltx_frames) if getattr(args, "ltx_fps", None) is not None: sampling_params["fps"] = int(args.ltx_fps) + # Handoff: gs=1.0 + no negative prompt → single forward, aligned with train. sampling_params["guidance_scale"] = 1.0 - sampling_params["negative_prompt"] = " " - # LTX23 one-stage rollout uses a stage1 guider (CFG/STG/modality/rescale) - # whose params cannot be overridden via HTTP — sglang routes unknown - # SamplingParams kwargs through the base class. Train ``forward_velocity`` - # is video-only with no guidance, so the rollout engine forces an identity - # guider via the ``patch_ltx2_identity_guider`` monkey patch - # (MILES_LTX_IDENTITY_GUIDER, propagated in miles/ray/rollout.py). + sampling_params["negative_prompt"] = None if evaluation: sampling_params["rollout"] = False @@ -127,8 +122,11 @@ def build_rollout_generate_payload( """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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 9fdf99e1..5fbe4781 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1328,9 +1328,12 @@ def miles_validate_args(args): else: args.diffusion_model_type = "sd3" if args.diffusion_model_type == "ltx": - if float(getattr(args, "diffusion_guidance_scale", 1.0)) != 1.0: - raise ValueError( - "LTX requires --diffusion-guidance-scale 1.0 (no CFG)." + 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 not args.debug_train_only and not getattr(args, "ltx_gemma_path", None): logger.warning( diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index 26fde4c6..1baa7a80 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -12,23 +12,25 @@ # colocate if 512x768x57f OOMs on one card. # # Usage: -# CUDA_VISIBLE_DEVICES=1 USE_LORA=1 NUM_ROLLOUT=8 \ -# MILES_DIFFUSION_DEBUG=1 LTX_DISABLE_AV_CROSS_ATTN=1 \ -# MILES_LTX_IDENTITY_GUIDER=1 \ +# CUDA_VISIBLE_DEVICES=1 USE_LORA=1 NUM_ROLLOUT=200 \ +# LTX_DISABLE_AV_CROSS_ATTN=1 \ # nohup bash scripts/run-diffusion-grpo-ltx23-sglang.sh \ # > logs/ltx23_dev_cps_$(date +%Y%m%d_%H%M%S).log 2>&1 & # # Key overridable env: -# LTX_MODEL_PATH — dev 22B safetensors +# LTX_MODEL_PATH — dev 22B safetensors (train + rollout DiT via transformer_weights_path) +# MILES_LTX_ROLLOUT_MODEL_PATH — optional; default Lightricks/LTX-2.3 (sglang overlay) +# MILES_LTX_MODEL_ID — optional; default LTX-2.3 # HEIGHT WIDTH FRAMES — 512 768 57 # NUM_STEPS — 24 # LTX_NUM_SDE_STEPS — 3 # LTX_SDE_STEP_CANDIDATES — 0,1,2,3,4,5,6,7,8,9 # CLIP_RANGE — 1e-4 -# ROLLOUT_BATCH_SIZE — unique prompts per rollout (default: 16) -# N_SAMPLES_PER_PROMPT — GRPO group size, aligned with verl (default: 8) -# NUM_STEPS_PER_ROLLOUT — optimizer steps per rollout (default: 2 → gbs=64) -# NUM_GPUS — 1 +# ROLLOUT_BATCH_SIZE — unique prompts per rollout (default: 8) +# N_SAMPLES_PER_PROMPT — GRPO group size (default: 8) +# NUM_STEPS_PER_ROLLOUT — optimizer steps per rollout (default: 2 → gbs=32) +# NUM_ROLLOUT — 200 +# SAVE_INTERVAL — 50 MILES_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" echo "[kill] hunting for stale miles processes under cwd=${MILES_ROOT}" @@ -97,8 +99,8 @@ GEMMA_ROOT="${GEMMA_ROOT:-${LTX_MATERIALIZED_ROOT}/text_encoder}" MILES_DATA_ROOT="${MILES_DATA_ROOT:-/sgl-workspace/miles}" PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" -NUM_ROLLOUT="${NUM_ROLLOUT:-8}" -ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-16}" +NUM_ROLLOUT="${NUM_ROLLOUT:-200}" +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}" SAMPLES_PER_ROLLOUT=$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT)) @@ -123,7 +125,7 @@ NUM_GPUS="${NUM_GPUS:-1}" ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" # Periodic checkpoint (LoRA adapter) so the run is resumable via LOAD_CKPT. # (The earlier run had no --save-interval, so nothing was ever saved.) -SAVE_INTERVAL="${SAVE_INTERVAL:-5}" +SAVE_INTERVAL="${SAVE_INTERVAL:-50}" if [[ ! -f "${PROMPT_DATA}" ]]; then python "${MILES_DATA_ROOT}/tools/prepare_vidgen_jsonl.py" @@ -140,7 +142,7 @@ echo "[run] log=${LOG_DIR}" echo "[run] wandb=${WANDB_DIR}" echo "[run] save=${SAVE_DIR}" echo "[run] ${HEIGHT}x${WIDTH}x${FRAMES}f steps=${NUM_STEPS} sde_steps=${LTX_NUM_SDE_STEPS} candidates=${LTX_SDE_STEP_CANDIDATES} clip=${CLIP_RANGE}" -echo "[run] batch: rollout=${ROLLOUT_BATCH_SIZE} n_samples=${N_SAMPLES_PER_PROMPT} samples/rollout=${SAMPLES_PER_ROLLOUT} optim_steps/rollout=${NUM_STEPS_PER_ROLLOUT} gbs=${DERIVED_GLOBAL_BATCH_SIZE}" +echo "[run] batch: rollout=${ROLLOUT_BATCH_SIZE} n_samples=${N_SAMPLES_PER_PROMPT} samples/rollout=${SAMPLES_PER_ROLLOUT} optim_steps/rollout=${NUM_STEPS_PER_ROLLOUT} gbs=${DERIVED_GLOBAL_BATCH_SIZE} save_interval=${SAVE_INTERVAL}" DEBUG_ARGS=() if [[ "${MILES_DIFFUSION_DEBUG:-0}" == "1" ]]; then @@ -251,8 +253,8 @@ python -u "${ROOT_DIR}/train_diffusion.py" \ --rm-type pickscore \ --diffusion-reward "pickscore:1.0" \ --reward-key avg \ - --pickscore-processor-path "${PICKSCORE_PROCESSOR:-yuvalkirstain/PickScore_v1}" \ - --pickscore-model-path "${PICKSCORE_MODEL:-yuvalkirstain/PickScore_v1}" \ + --pickscore-processor-path "${PICKSCORE_PROCESSOR:-/data/wenhao/hf_home/pickscore}" \ + --pickscore-model-path "${PICKSCORE_MODEL:-/data/wenhao/hf_home/pickscore}" \ --pickscore-num-frames 3 \ --pickscore-batch-size 8 \ --pickscore-num-gpus-per-worker 0 \ From 9f40f70ca606cef561ab4da92b3bc3437bb729e1 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 10 Jun 2026 11:22:45 +0000 Subject: [PATCH 13/31] fix: new transformer version pickscore --- miles/rollout/rm_hub/pickscore.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index f6d61dd2..121bb84e 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -25,7 +25,20 @@ 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: @@ -141,10 +154,10 @@ def score_videos( ) text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()} - image_embs = self.model.get_image_features(**image_inputs) + 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 = self.model.get_text_features(**text_inputs) + 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) From f608937aff48ccc627904cde104488301024e3e6 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Thu, 11 Jun 2026 06:34:35 +0000 Subject: [PATCH 14/31] refactor(ltx): unify HF overlay loading and decouple actor from LTX hooks --- miles/backends/fsdp_utils/actor.py | 60 +-- miles/backends/fsdp_utils/configs/ltx.py | 38 +- .../configs/train_pipeline_config.py | 44 ++ miles/backends/model_families/__init__.py | 1 + miles/backends/model_families/ltx.py | 436 ++++++++++++++++++ .../sglang_diffusion_utils/configs/ltx.py | 141 ++---- .../patch_ltx2_identity_guider.py | 76 --- .../sglang_diffusion_engine.py | 4 +- miles/ray/rollout.py | 21 +- miles/rollout/sglang_diffusion_rollout.py | 44 +- miles/utils/arguments.py | 79 +--- miles/utils/diffusion_rollout_response.py | 4 - miles/utils/types.py | 3 - scripts/run-diffusion-grpo-ltx23-sglang.sh | 25 +- scripts/run-ltx23-grpo-local.sh | 82 ++++ train_diffusion.py | 6 + 16 files changed, 706 insertions(+), 358 deletions(-) create mode 100644 miles/backends/model_families/__init__.py create mode 100644 miles/backends/model_families/ltx.py delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py create mode 100755 scripts/run-ltx23-grpo-local.sh diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 88827131..db5cae7d 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -343,22 +343,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. + # Use rollout's exact sigmas snapshot; fall back to model-specific reconstruction. timesteps_ref = dit_trajectories[0].timesteps.to(device).float() sigmas_snapshot = getattr(dit_trajectories[0], "sigmas", None) sched_config = getattr(self.scheduler, "config", None) num_train_timesteps = ( int(sched_config.num_train_timesteps) if sched_config is not None else 1000 ) - if sigmas_snapshot is not None: - sigmas_ref = sigmas_snapshot.to(device).float() - elif not self.train_pipeline_config.needs_timestep_scaling: - # Trajectory timesteps may be AdaLN-scale (LTX: σ×1000). CPS expects σ∈[0,1]. - sigmas_ref = self.train_pipeline_config.scale_timesteps_for_sde(timesteps_ref) - sigmas_ref = torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) - else: - sigmas_ref = timesteps_ref / float(num_train_timesteps) - sigmas_ref = torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + sigmas_ref = self.train_pipeline_config.resolve_sigmas_ref( + timesteps_ref, sigmas_snapshot, self.scheduler, + ) self.scheduler.timesteps = timesteps_ref self.scheduler.sigmas = sigmas_ref @@ -468,23 +462,16 @@ def _build_train_grids( dit_trajectories[traj_idx], device ) - # prepare cond kwargs (denoising_env text embeds + local geometry for LTX) + # prepare cond kwargs (denoising_env + model-specific geometry when needed) denoising_env = denoising_envs[traj_idx] - if hasattr(train_pipeline_config, "build_train_cond_kwargs"): - positive_cond_kwargs_list.append( - train_pipeline_config.build_train_cond_kwargs( - denoising_env.pos_cond_kwargs, - video_latents=latents, - args=self.args, - device=device, - ) - ) - else: - positive_cond_kwargs_list.append( - train_pipeline_config.prepare_cond_kwargs( - denoising_env.pos_cond_kwargs, device - ) + positive_cond_kwargs_list.append( + 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( train_pipeline_config.prepare_cond_kwargs(denoising_env.neg_cond_kwargs, device) @@ -591,25 +578,6 @@ def _build_train_grids( ), } - def _build_sde_extra( - self, - grids: dict, - sample_indices: torch.Tensor, - tstep_indices: torch.Tensor, - ) -> dict | None: - if grids.get("sde_step_indices_window") is None: - return None - - idx = grids["sde_step_indices_window"][sample_indices][:, tstep_indices] - idx = idx.reshape(-1).long() - - return { - "sigmas": self.scheduler.sigmas, - "sde_step_indices": idx, - "dynamics_type": getattr(self.args, "ltx_dynamics_type", "cps"), - "sigma_min_override": getattr(self.args, "ltx_sigma_min", None), - } - def _run_optim_window( self, *, @@ -773,7 +741,9 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: noise_pred_flat = _compute_noise_pred() - sde_extra = self._build_sde_extra(grids, sample_indices, tstep_indices) + sde_extra = train_pipeline_config.build_sde_extra( + self.scheduler, grids, sample_indices, tstep_indices, self.args, + ) prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( self.scheduler, diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index f34c4617..ac2175f7 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -35,7 +35,10 @@ def load_model_and_scheduler(self, args, init_context_factory): from dataclasses import dataclass, field from ltx_core.components.schedulers import LTX2Scheduler - from ltx_trainer.model_loader import load_transformer + from miles.backends.model_families.ltx import ( + load_ltx_transformer_for_train, + resolve_transformer_checkpoint, + ) @dataclass class _LTXSchedulerHolder: @@ -53,7 +56,13 @@ def to(self, device): master_dtype_name = getattr(args, "fsdp_master_dtype", "bf16") master_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[master_dtype_name] - model = load_transformer(args.diffusion_model, device="cpu", dtype=master_dtype) + from miles.backends.model_families.ltx import resolve_transformer_checkpoint + + 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=master_dtype) num_steps = int(getattr(args, "diffusion_num_steps", 24)) ltx_sched = LTX2Scheduler() @@ -93,7 +102,7 @@ def build_train_cond_kwargs( self, cond: CondKwargs | None, *, - video_latents: torch.Tensor, + latents: torch.Tensor, args, device: torch.device, ) -> dict: @@ -106,7 +115,7 @@ def build_train_cond_kwargs( "LTX train requires denoising_env.pos_cond_kwargs.encoder_hidden_states" ) - ref = video_latents[0] if video_latents.ndim >= 2 else video_latents + 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: @@ -126,6 +135,27 @@ def build_train_cond_kwargs( kwargs.update(geom) return kwargs + def build_sde_extra( + self, + scheduler, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + args, + ) -> dict | None: + if grids.get("sde_step_indices_window") is None: + return None + + idx = grids["sde_step_indices_window"][sample_indices][:, tstep_indices] + idx = idx.reshape(-1).long() + + return { + "sigmas": scheduler.sigmas, + "sde_step_indices": idx, + "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(): diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index b0177aac..0a8d5506 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -14,6 +14,7 @@ from __future__ import annotations import abc +from argparse import Namespace from typing import Any import torch @@ -86,6 +87,49 @@ 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: Namespace, + device: torch.device, + ) -> dict: + """Build per-sample conditioning for the training forward pass.""" + return self.prepare_cond_kwargs(cond, device) + + def resolve_sigmas_ref( + self, + timesteps_ref: torch.Tensor, + sigmas_snapshot: torch.Tensor | None, + scheduler: Any, + ) -> torch.Tensor: + """Build ``[T+1]`` sigma reference for the training scheduler.""" + device = timesteps_ref.device + if sigmas_snapshot is not None: + return sigmas_snapshot.to(device).float() + + sched_config = getattr(scheduler, "config", None) + num_train_timesteps = ( + int(sched_config.num_train_timesteps) if sched_config is not None else 1000 + ) + if not self.needs_timestep_scaling: + sigmas_ref = self.scale_timesteps_for_sde(timesteps_ref) + else: + sigmas_ref = timesteps_ref / float(num_train_timesteps) + return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + + def build_sde_extra( + self, + scheduler: Any, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + args: Namespace, + ) -> dict | None: + """Optional per-tile metadata for model-specific SDE log_prob.""" + return None + def expand_cond_for_timestep_batch( self, cond_kwargs: dict, diff --git a/miles/backends/model_families/__init__.py b/miles/backends/model_families/__init__.py new file mode 100644 index 00000000..f03bee62 --- /dev/null +++ b/miles/backends/model_families/__init__.py @@ -0,0 +1 @@ +"""Per-model-family CLI, validation, and rollout hooks.""" diff --git a/miles/backends/model_families/ltx.py b/miles/backends/model_families/ltx.py new file mode 100644 index 00000000..985b392c --- /dev/null +++ b/miles/backends/model_families/ltx.py @@ -0,0 +1,436 @@ +"""LTX-2 model family: checkpoint resolution, rollout engine, and sampling hooks.""" + +from __future__ import annotations + +import logging +import os +from argparse import ArgumentParser, Namespace +from pathlib import Path +from typing import Any + +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 ( + LTXModelConfigurator, + LTXV_MODEL_COMFY_RENAMING_MAP, + ) + + 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 " + f"Lightricks/LTX-2.3 (recommended) or a .safetensors override." + ) + + +def preflight_ltx_models(args: Namespace) -> None: + """Ensure train + rollout can resolve the same model before Ray starts.""" + hf_model_id = resolve_hf_model_id(args) + checkpoint = resolve_transformer_checkpoint( + getattr(args, "diffusion_model", None), + explicit_path=getattr(args, "sglang_transformer_weights_path", None), + ) + ckpt_path = Path(checkpoint) + if _is_materialized_diffusers_checkpoint(ckpt_path): + _read_materialized_transformer_config(ckpt_path) + logger.info( + "LTX preflight ok: rollout model_path=%s train_checkpoint=%s", + hf_model_id, + checkpoint, + ) + + +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 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 + + from miles.utils.sde_log_prob import normalize_dynamics_type + + dynamics = normalize_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 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-disable-av-cross-attn", + action="store_true", + default=False, + help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", + ) + 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." + ) diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py index 248dee85..ca24b6d1 100644 --- a/miles/backends/sglang_diffusion_utils/configs/ltx.py +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -1,53 +1,43 @@ -"""LTX-2 sglang-d rollout engine config. - -Rollout engine uses ``model_path=Lightricks/LTX-2.3`` + ``model_id=LTX-2.3`` so -sglang's overlay wrapper materializes a full diffusers tree (``model_index.json``, -VAE, text encoder, connectors). Train FSDP still loads ``--diffusion-model`` as a -single official safetensors file; ``transformer_weights_path`` pins rollout DiT -to that same file for weight parity. -""" +"""LTX-2 sglang-d rollout engine config (re-exports model family helpers).""" from __future__ import annotations -import logging -import os -from pathlib import Path - -logger = logging.getLogger(__name__) - -# sglang registry + overlay wrapper (see model_overlay.py). -LTX_DEFAULT_HF_MODEL = "Lightricks/LTX-2.3" -LTX_DEFAULT_MODEL_ID = "LTX-2.3" +from miles.backends.model_families.ltx import ( + LTX_DEFAULT_HF_MODEL, + LTX_DEFAULT_MODEL_ID, + ensure_materialized_model, + is_ltx_model, + preflight_ltx_models, + 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", + "preflight_ltx_models", + "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 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 - diff_model = (getattr(args, "diffusion_model", None) or "").lower() - return "ltx" in diff_model or diff_model.endswith(".safetensors") +def resolve_sglang_model_path(args) -> str: + return resolve_hf_model_id(args) def resolve_ltx_model_id(args) -> str: - """Short registry id for ``ServerArgs.model_id`` (matches ``Lightricks/LTX-2.3``).""" - 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 resolve_sglang_model_path(args) -> str: - """HF hub id for sglang pipeline skeleton (overlay materializes components).""" - 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 + return resolve_model_id(args) def resolve_ltx_transformer_weights_path( @@ -55,66 +45,9 @@ def resolve_ltx_transformer_weights_path( *, explicit_path: str | None = None, ) -> str | None: - """Return official single-file safetensors for sglang ``transformer_weights_path``. - - Overlay materialized ``transformer/model.safetensors`` is a different checkpoint - variant than dev 22B; miles train loads ``--diffusion-model`` via ltx_core. Point - rollout DiT init + weight sync at the same single-file ckpt as train. - """ - if explicit_path: - path = Path(explicit_path).expanduser() - if path.is_file(): - return str(path) - return None - - 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) - - if diffusion_model and str(diffusion_model).endswith(".safetensors"): - path = Path(diffusion_model).expanduser() - if path.is_file(): - return str(path) - return None - - -def server_kwargs_extras(args) -> dict: - """Extra ``ServerArgs`` kwargs; call only when ``is_ltx_model(args)``.""" - extras: dict = { - "model_id": resolve_ltx_model_id(args), - } - - explicit = getattr(args, "sglang_transformer_weights_path", None) - weights_path = resolve_ltx_transformer_weights_path( - getattr(args, "diffusion_model", None), - explicit_path=explicit, - ) - if weights_path and not explicit: - extras["transformer_weights_path"] = weights_path - logger.info( - "LTX rollout: model_path=%s model_id=%s transformer_weights_path=%s", - resolve_sglang_model_path(args), - extras["model_id"], - weights_path, - ) - elif explicit: - extras["transformer_weights_path"] = explicit - logger.info( - "LTX rollout: model_path=%s model_id=%s transformer_weights_path=%s", - resolve_sglang_model_path(args), - extras["model_id"], - explicit, + try: + return resolve_transformer_checkpoint( + diffusion_model, explicit_path=explicit_path, ) - else: - logger.warning( - "LTX rollout: no transformer_weights_path resolved from --diffusion-model; " - "rollout DiT will use overlay default (may diverge from train ckpt)." - ) - - gemma_path = getattr(args, "ltx_gemma_path", None) - if gemma_path: - extras["component_paths"] = {"text_encoder": gemma_path} - - return extras + except FileNotFoundError: + return None diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py deleted file mode 100644 index 555aade3..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_identity_guider.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Force an identity LTX-2.3 stage1 guider for train/rollout alignment. - -Miles GRPO train side computes ``forward_velocity`` as a **video-only** forward -with no CFG / STG / modality / rescale. The sglang LTX2.3 one-stage rollout, -however, applies a stage1 guider whose parameters default to ``video_cfg_scale=3`` -etc. (see ``configs/sample/ltx_2.py``). Those parameters **cannot** be overridden -through ``POST /rollout/generate``: ``SamplingParams.from_user_sampling_params_args`` -routes unknown kwargs through the base ``SamplingParams`` class, which rejects -LTX23-only fields. So the rollout-side ``rollout_model_outputs`` are post-guider -velocities that diverge from the train forward (~0.94 cosine, scale≈0.86). - -This patch overrides ``LTX2DenoisingStage._get_ltx2_stage1_guider_params`` so the -guider becomes the identity transform: - - pred = cond - + (cfg-1)*(cond-uncond_text) # cfg=1 -> 0 - + stg*(cond-uncond_perturbed) # stg=0 -> 0 - + (modality-1)*(cond-uncond_mod) # mod=1 -> 0 - pred = rescale(cond, pred, 0.0) # rescale=0 -> pred unchanged - => pred == cond (video-only x0) => velocity == raw video velocity - -Controlled by ``MILES_LTX_IDENTITY_GUIDER`` (default ``"1"``). Set to ``"0"`` to -keep the official guider (e.g. for generation-quality eval, not RL alignment). -""" - -from __future__ import annotations - -import os -from typing import Any - -_APPLIED = False -_ORIG = None - -_IDENTITY_GUIDER: dict[str, Any] = { - "video_cfg_scale": 1.0, - "video_stg_scale": 0.0, - "video_rescale_scale": 0.0, - "video_modality_scale": 1.0, - "video_skip_step": 0, - "video_stg_blocks": [], - "audio_cfg_scale": 1.0, - "audio_stg_scale": 0.0, - "audio_rescale_scale": 0.0, - "audio_modality_scale": 1.0, - "audio_skip_step": 0, - "audio_stg_blocks": [], -} - - -def _identity_enabled() -> bool: - return os.environ.get("MILES_LTX_IDENTITY_GUIDER", "1") == "1" - - -def apply() -> None: - global _APPLIED, _ORIG - if _APPLIED: - return - - from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( - LTX2DenoisingStage, - ) - - _ORIG = LTX2DenoisingStage._get_ltx2_stage1_guider_params - - def _patched_get_guider(self, batch, server_args, stage): - result = _ORIG(self, batch, server_args, stage) - # Only override when guider is active (stage1 returns a dict) and the - # alignment flag is on. None (non-stage1 / official cfg path) is kept. - if result is None or not _identity_enabled(): - return result - merged = dict(result) - merged.update(_IDENTITY_GUIDER) - return merged - - LTX2DenoisingStage._get_ltx2_stage1_guider_params = _patched_get_guider - _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 8d907987..11a8c7be 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -343,9 +343,9 @@ def _compute_server_args(args, host, port, nccl_port): "warmup": False, } - # LTX rollout: HF model id + overlay wrapper; DiT weights via transformer_weights_path. + # 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_sglang_model_path(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 diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 924fecc7..ae286830 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -483,6 +483,20 @@ def _split_train_data_by_dp(self, data, dp_size): return rollout_data_refs +def _pythonpath_env_var() -> dict[str, str]: + """Ray workers do not inherit shell PYTHONPATH; propagate sglang source root.""" + sglang_python = os.environ.get("SGLANG_PYTHON") + parent_path = os.environ.get("PYTHONPATH", "") + if sglang_python: + prefix = sglang_python + if parent_path and not parent_path.startswith(prefix): + prefix = f"{sglang_python}:{parent_path}" + return {"PYTHONPATH": prefix} + if parent_path: + return {"PYTHONPATH": parent_path} + return {} + + def _base_rollout_engine_env_vars() -> dict[str, str]: return {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { "SGL_JIT_DEEPGEMM_PRECOMPILE": "false", @@ -511,7 +525,11 @@ def _ltx_alignment_env_vars(args) -> dict[str, str]: env["MILES_LTX_DISABLE_AV_CROSS"] = "1" from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV - for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): + for name in ( + LTX_ROLLOUT_PATCHES_ENV, + "MILES_LTX_DISABLE_AV_CROSS", + "MILES_LTX_IDENTITY_GUIDER", + ): if os.environ.get(name): env[name] = os.environ[name] return env @@ -519,6 +537,7 @@ def _ltx_alignment_env_vars(args) -> dict[str, str]: def _build_rollout_engine_env_vars(args) -> dict[str, str]: env_vars = _base_rollout_engine_env_vars() + env_vars.update(_pythonpath_env_var()) for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): if os.environ.get(cache_var): env_vars[cache_var] = os.environ[cache_var] diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 39078b78..7026cd2e 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -30,13 +30,12 @@ def _resolve_diffusion_model_type(args: Namespace) -> str: + from miles.backends.model_families.ltx import is_ltx_model + model_type = (getattr(args, "diffusion_model_type", "auto") or "auto").lower() if model_type != "auto": return model_type - diff_model = (getattr(args, "diffusion_model", None) or "").lower() - if "ltx" in diff_model or diff_model.endswith(".safetensors"): - return "ltx" - return "sd3" + return "ltx" if is_ltx_model(args) else "sd3" def build_rollout_sampling_params( @@ -62,14 +61,6 @@ def build_rollout_sampling_params( } model_type = _resolve_diffusion_model_type(args) - if model_type == "ltx": - 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) - # Handoff: gs=1.0 + no negative prompt → single forward, aligned with train. - sampling_params["guidance_scale"] = 1.0 - sampling_params["negative_prompt"] = None if evaluation: sampling_params["rollout"] = False @@ -85,28 +76,13 @@ def build_rollout_sampling_params( "rollout_return_dit_trajectory": True, } ) - if model_type == "ltx": - from miles.utils.sde_log_prob import normalize_dynamics_type - - # Canonical names match sglang-d rollout_sde_type, so pass through - # with no translation table (keeps train/rollout on one vocabulary). - dynamics = normalize_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 (train recompute supports it). " - "Add the sglang-d sampling branch before using it for rollout." - ) - 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) - # Disable flag is propagated via MILES_LTX_DISABLE_AV_CROSS on rollout engines - # (see miles/ray/rollout.py). Do not pass via extra_sampling_params — master - # sglang SamplingParams does not accept ltx2_disable_av_cross_attn. + + # 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"). + if model_type == "ltx": + from miles.backends.model_families.ltx import patch_rollout_sampling_params + + patch_rollout_sampling_params(sampling_params, args, evaluation=evaluation) if extra_sampling_params: sampling_params["extra_sampling_params"] = extra_sampling_params diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 5fbe4781..dcdc751c 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -443,64 +443,12 @@ def add_rollout_arguments(parser): choices=["auto", "sd3", "ltx"], help=( "Override the diffusion model family. ``auto`` infers from --diffusion-model " - "(diffusers repo → sd3, single-file safetensors → ltx)." + "(HF hub id with ``ltx`` → ltx, else sd3)." ), ) - parser.add_argument( - "--ltx-gemma-path", - type=str, - default=None, - help="Path to the Gemma-3 12B directory used as LTX's text encoder (rollout side).", - ) - 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-disable-av-cross-attn", - action="store_true", - default=False, - help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", - ) - parser.add_argument( - "--pickscore-num-frames", - type=int, - default=3, - help="Number of evenly spaced frames to score per video (LTX PickScore reward).", - ) + from miles.backends.model_families.ltx import register_args as register_ltx_args + + register_ltx_args(parser) parser.add_argument( "--rollout-seed", type=int, @@ -1328,22 +1276,9 @@ def miles_validate_args(args): else: args.diffusion_model_type = "sd3" if args.diffusion_model_type == "ltx": - 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 not args.debug_train_only and not getattr(args, "ltx_gemma_path", None): - logger.warning( - "--ltx-gemma-path is not set; sglang LTX rollout will need it once Phase 1 lands." - ) - 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." - ) + from miles.backends.model_families.ltx import validate_args as validate_ltx_args + + validate_ltx_args(args) # always true on offload for colocate at the moment. if args.colocate: diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 1dde0f01..b5199b7b 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -92,10 +92,6 @@ def _parse_cond_kwargs( pooled_projections=_parse_tensor_or_list( data.get("pooled_projections"), deserialize_func=deserialize_func ), - # Legacy rollout fields; miles train rebuilds geometry locally for LTX. - ltx_positions=deserialize_func(data.get("ltx_positions") or data.get("positions")), - ltx_denoise_mask=deserialize_func(data.get("ltx_denoise_mask") or data.get("denoise_mask")), - ltx_clean_latent=deserialize_func(data.get("ltx_clean_latent") or data.get("clean_latent")), ) diff --git a/miles/utils/types.py b/miles/utils/types.py index 86cc75c9..4bc53525 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -37,9 +37,6 @@ class CondKwargs: encoder_attention_mask: torch.Tensor | None = None audio_encoder_attention_mask: torch.Tensor | None = None pooled_projections: list[torch.Tensor] | None = None - ltx_positions: torch.Tensor | None = None - ltx_denoise_mask: torch.Tensor | None = None - ltx_clean_latent: torch.Tensor | None = None @dataclass diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index 1baa7a80..ef272b95 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -18,9 +18,8 @@ # > logs/ltx23_dev_cps_$(date +%Y%m%d_%H%M%S).log 2>&1 & # # Key overridable env: -# LTX_MODEL_PATH — dev 22B safetensors (train + rollout DiT via transformer_weights_path) -# MILES_LTX_ROLLOUT_MODEL_PATH — optional; default Lightricks/LTX-2.3 (sglang overlay) -# MILES_LTX_MODEL_ID — optional; default LTX-2.3 +# LTX_HF_MODEL — default Lightricks/LTX-2.3 (train + rollout via overlay) +# LTX_DEV_SAFETENSORS — optional dev .safetensors override for train + rollout DiT # HEIGHT WIDTH FRAMES — 512 768 57 # NUM_STEPS — 24 # LTX_NUM_SDE_STEPS — 3 @@ -90,12 +89,14 @@ export RAY_object_spilling_config="$(python -c "import json,os; print(json.dumps ray stop --force 2>/dev/null || true sleep 2 -# ── dev checkpoint (borrowed from legacy reward run) ───────────────────── -LTX_MODEL_PATH="${LTX_MODEL_PATH:-/sgl-workspace/rollout_compare/models/LTX-2.3/ltx-2.3-22b-dev.safetensors}" -# sglang text_encoder: use materialized Lightricks overlay (local gemma_for_ltx23 -# symlinks often point at stale HF cache and break rollout startup). -LTX_MATERIALIZED_ROOT="${LTX_MATERIALIZED_ROOT:-/data/wenhao/sgl_diffusion_cache/materialized_models/Lightricks__LTX-2.3-10cce1713d7efa14}" -GEMMA_ROOT="${GEMMA_ROOT:-${LTX_MATERIALIZED_ROOT}/text_encoder}" +# ── model: HF hub id by default; optional dev safetensors override ───────── +LTX_HF_MODEL="${LTX_HF_MODEL:-Lightricks/LTX-2.3}" +LTX_DEV_SAFETENSORS="${LTX_DEV_SAFETENSORS:-}" +if [[ -n "${LTX_DEV_SAFETENSORS}" ]]; then + DIFFUSION_MODEL="${LTX_DEV_SAFETENSORS}" +else + DIFFUSION_MODEL="${LTX_HF_MODEL}" +fi MILES_DATA_ROOT="${MILES_DATA_ROOT:-/sgl-workspace/miles}" PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" @@ -136,8 +137,7 @@ SAVE_DIR="${CKPT_ROOT}/${RUN_NAME}" mkdir -p "${SAVE_DIR}" echo "[run] dev+cps CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} NUM_GPUS=${NUM_GPUS}" -echo "[run] dit=${LTX_MODEL_PATH}" -echo "[run] gemma=${GEMMA_ROOT}" +echo "[run] model=${DIFFUSION_MODEL}" echo "[run] log=${LOG_DIR}" echo "[run] wandb=${WANDB_DIR}" echo "[run] save=${SAVE_DIR}" @@ -201,9 +201,8 @@ fi python -u "${ROOT_DIR}/train_diffusion.py" \ --train-backend fsdp \ --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ - --diffusion-model "${LTX_MODEL_PATH}" \ + --diffusion-model "${DIFFUSION_MODEL}" \ --diffusion-model-type ltx \ - --ltx-gemma-path "${GEMMA_ROOT}" \ --hf-checkpoint gpt2 \ --prompt-data "${PROMPT_DATA}" \ --input-key input \ diff --git a/scripts/run-ltx23-grpo-local.sh b/scripts/run-ltx23-grpo-local.sh new file mode 100755 index 00000000..e97a72e9 --- /dev/null +++ b/scripts/run-ltx23-grpo-local.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# LTX-2.3 GRPO 本地启动脚本（/data/wenhao 环境） +# +# 仅需 --diffusion-model Lightricks/LTX-2.3（train + rollout 共用 sglang overlay）。 +# 首次运行会自动 materialize 到 SGLANG_DIFFUSION_CACHE_ROOT。 +# +# 前置条件： +# - venv: /data/wenhao/.venvs/miles-diffusion +# - sglang: /data/wenhao/master_sglang/sglang/python +# - miles: /data/wenhao/master_miles/miles_diffusion +# +# 用法： +# # 前台调试 +# bash scripts/run-ltx23-grpo-local.sh +# +# # 冒烟测试（1 rollout，跳过 optimizer） +# NUM_ROLLOUT=1 ROLLOUT_BATCH_SIZE=1 N_SAMPLES_PER_PROMPT=2 NUM_STEPS_PER_ROLLOUT=1 \ +# SKIP_OPTIMIZER=1 bash scripts/run-ltx23-grpo-local.sh +# +# 可覆盖的环境变量： +# CUDA_VISIBLE_DEVICES GPU 编号（默认 2） +# NUM_ROLLOUT rollout 总数（默认 100） +# MILES_DIFFUSION_DEBUG 1=打印对齐指标（默认 1） +# SKIP_OPTIMIZER 1=冒烟模式，不更新权重（默认 0） +# LTX_DEV_SAFETENSORS 可选：dev .safetensors 覆盖 overlay 默认 DiT 权重 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MILES_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +# ── 路径 ────────────────────────────────────────────────────────────────── +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-2}" +export SGLANG_PYTHON="${SGLANG_PYTHON:-/data/wenhao/master_sglang/sglang/python}" +export PYTHONPATH="${SGLANG_PYTHON}${PYTHONPATH:+:${PYTHONPATH}}" + +export MILES_DATA_ROOT="${MILES_DATA_ROOT:-/data/wenhao/master_miles/miles_diffusion}" +export MILES_DATA_DISK_ROOT="${MILES_DATA_DISK_ROOT:-/data/wenhao/miles_diffusion}" +export LTX_HF_MODEL="${LTX_HF_MODEL:-Lightricks/LTX-2.3}" +export PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" +export HF_HOME="${HF_HOME:-/data/wenhao/hf_home}" +export SGLANG_DIFFUSION_CACHE_ROOT="${SGLANG_DIFFUSION_CACHE_ROOT:-/data/wenhao/sgl_diffusion_cache}" + +# ── 训练规模 ────────────────────────────────────────────────────────────── +export USE_LORA="${USE_LORA:-1}" +export NUM_ROLLOUT="${NUM_ROLLOUT:-100}" +export ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" +export N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +export NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-2}" +export SAVE_INTERVAL="${SAVE_INTERVAL:-50}" + +# ── 对齐开关 ────────────────────────────────────────────────────────────── +export MILES_DIFFUSION_DEBUG="${MILES_DIFFUSION_DEBUG:-1}" +export LTX_DISABLE_AV_CROSS_ATTN="${LTX_DISABLE_AV_CROSS_ATTN:-1}" +export MILES_APPLY_LTX2_LTXCORE_PARITY="${MILES_APPLY_LTX2_LTXCORE_PARITY:-1}" + +# ── 冒烟 / 恢复 ─────────────────────────────────────────────────────────── +export SKIP_OPTIMIZER="${SKIP_OPTIMIZER:-0}" + +# ── 激活环境 ────────────────────────────────────────────────────────────── +VENV="${VENV:-/data/wenhao/.venvs/miles-diffusion}" +if [[ -f "${VENV}/bin/activate" ]]; then + # shellcheck source=/dev/null + source "${VENV}/bin/activate" +else + echo "[warn] venv not found: ${VENV}" >&2 +fi + +echo "============================================" +echo " LTX-2.3 GRPO Local Launch" +echo "============================================" +echo " GPU: ${CUDA_VISIBLE_DEVICES}" +echo " Rollouts: ${NUM_ROLLOUT}" +echo " Model: ${LTX_DEV_SAFETENSORS:-${LTX_HF_MODEL}}" +echo " Debug: ${MILES_DIFFUSION_DEBUG}" +echo " Skip optim: ${SKIP_OPTIMIZER}" +echo " Logs: ${MILES_DATA_DISK_ROOT}/logs/" +echo " Ckpt: ${MILES_DATA_DISK_ROOT}/ckpt/" +echo "============================================" + +cd "${MILES_ROOT}" +exec bash scripts/run-diffusion-grpo-ltx23-sglang.sh diff --git a/train_diffusion.py b/train_diffusion.py index 4e7caadb..6cba5646 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -13,6 +13,12 @@ def train(args): configure_logger() logger = logging.getLogger(__name__) + + if getattr(args, "diffusion_model_type", None) == "ltx": + from miles.backends.model_families.ltx import preflight_ltx_models + + preflight_ltx_models(args) + # allocate the GPUs logger.info("train: creating placement groups") pgs = create_placement_groups(args) From 3be6d718490f2cb753278acd344835a8b7e0e05b Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Thu, 11 Jun 2026 11:00:35 +0000 Subject: [PATCH 15/31] chore(ltx): drop dev artifacts and simplify rollout env bridging Remove LTX preflight, local dev script/skill, and SGLANG_PYTHON-specific Ray env logic; rely on lazy model loading and generic PYTHONPATH passthrough. Co-authored-by: Cursor --- .../skills/align-ltx-train-rollout/SKILL.md | 86 ------------------- .gitignore | 4 +- miles/backends/model_families/ltx.py | 17 ---- .../sglang_diffusion_utils/configs/ltx.py | 2 - miles/ray/rollout.py | 18 +--- scripts/run-diffusion-grpo-ltx23-sglang.sh | 2 +- scripts/run-ltx23-grpo-local.sh | 82 ------------------ train_diffusion.py | 5 -- 8 files changed, 6 insertions(+), 210 deletions(-) delete mode 100644 .claude/skills/align-ltx-train-rollout/SKILL.md delete mode 100755 scripts/run-ltx23-grpo-local.sh diff --git a/.claude/skills/align-ltx-train-rollout/SKILL.md b/.claude/skills/align-ltx-train-rollout/SKILL.md deleted file mode 100644 index 8e5ac790..00000000 --- a/.claude/skills/align-ltx-train-rollout/SKILL.md +++ /dev/null @@ -1,86 +0,0 @@ ---- -name: align-ltx-train-rollout -description: Diagnose and fix train-vs-rollout forward numerical misalignment for LTX-2.3 diffusion GRPO on sglang rollout, driving model_output_cosine_sim to ~0.9998 and log_prob_mean_abs_diff small. Use when GRPO clipfrac stays near 1.0, train/log_prob_mean_abs_diff is large, model_output_cosine_sim is low, or when aligning the miles FSDP train forward with the sglang rollout forward for LTX video diffusion. ---- - -# align-ltx-train-rollout - -把 miles 训练侧重算的 `log_prob` / `model_output` 与 sglang live rollout 对齐,使 GRPO 的 importance ratio 可信(`clipfrac` 不再长期 = 1.0)。 - -## 核心心智模型(先读这一条) - -**对齐 gap 几乎从不是 checkpoint / 权重错。** 按以下三层顺序定位,不要一上来换权重: - -1. **log_prob 公式层** — SDE dynamics 类型、`sigma_min` 用错 → `log_prob_mean_abs_diff` 巨大(2~15) -2. **DiT forward 层** — temb/AdaLN 的 shape 与语义、算子 parity、AV cross-attn、attention backend → `cosine` 低(~0.96) -3. **live pipeline 后处理层** — guider(cfg/stg/modality/rescale)改了 velocity 语义 → offline 高但 live 低(~0.94) - -## 诊断决策树 - -开 `--diffusion-debug-mode`,看训练 log 的 `train/model_output_cosine_sim`、`train/log_prob_mean_abs_diff`、`train/clipfrac`(计算位置 `miles/backends/fsdp_utils/actor.py:836-847`)。然后: - -``` -cosine 高(>0.999) 但 log_prob_diff 大? - └─→ 第①层:SDE 公式 / sigma_min。查 sde_log_prob.py + sglang scheduler_rl_mixin。 - -cosine 低(<0.99)? - └─ 先做 offline injected 对比(同 latent/kwargs 喂两侧 DiT,排除 checkpoint): - bash scripts/capture-and-compare-ltx23-forward.sh - ├─ injected 也低 → 第②层:temb/parity/AV/attention(见下方必备开关) - └─ injected 高(>0.999) 但 live 低 → 第③层:guider 后处理(identity guider) -``` - -## 对齐必备开关(缺一项就掉精度) - -| 开关 | 作用 | 对应层 | -|------|------|--------| -| `MILES_APPLY_LTX2_LTXCORE_PARITY=1` | temb expand `[B,1,D]→[B,T,D]` + AdaLN/RMSNorm/RoPE/SDPA 数值对齐 | ② | -| `MILES_LTX_IDENTITY_GUIDER=1` | 强制 stage1 guider 为 identity(cfg=1/stg=0/modality=1/rescale=0) | ③ | -| `--ltx-disable-av-cross-attn` (+ `MILES_LTX_DISABLE_AV_CROSS=1`) | train video-only 与 rollout 算子图一致 | ② | -| `SGLANG_ATTENTION_BACKEND=torch_sdpa` | 避免 FlashAttention vs SDPA 数值差 | ② | -| train 与 sglang 用**同一份 `.safetensors`** | 避免 HF materialized overlay 与单文件差 ~1e-4(由 `configs/ltx.py` resolve) | ② | -| train 与 rollout **同 `--diffusion-num-steps` / `--ltx-dynamics-type` / `--ltx-sigma-min`** | 步数/动力学/σ_min 一致 | ①② | - -**经验**:`dev` ckpt + 24 步比 `distilled` + 8 步对齐好得多(完整路径 velocity 更平滑)。优先用 dev 验证。 - -## 验收标准 - -| 指标 | 达标 | dev/512×768×57f/24步实测 | -|------|------|--------------------------| -| `model_output_cosine_sim` | ≥ 0.999 | 0.9995 ~ 0.9999 | -| `log_prob_mean_abs_diff` | < 5e-3 | 6e-6 ~ 1.7e-3 | -| `clipfrac` | 明显 < 1.0 | 0 ~ 0.125 | - -## 快速命令 - -```bash -cd /sgl-workspace/master_miles/miles_diffusion -export PYTHONPATH=/sgl-workspace/master_sglang/sglang/python${PYTHONPATH:+:$PYTHONPATH} - -# 纯前向对齐验证(不更新权重,最快看 cosine / log_prob_diff) -CUDA_VISIBLE_DEVICES= MILES_DIFFUSION_DEBUG=1 \ -LTX_DISABLE_AV_CROSS_ATTN=1 MILES_LTX_IDENTITY_GUIDER=1 USE_LORA=1 SKIP_OPTIMIZER=1 \ -NUM_ROLLOUT=3 ROLLOUT_BATCH_SIZE=1 N_SAMPLES_PER_PROMPT=2 GLOBAL_BATCH_SIZE=2 \ -nohup bash scripts/run-diffusion-grpo-ltx23-sglang-dev-flowsde.sh > logs/align_verify.log 2>&1 & - -# 离线 capture + compare(定位 gap 在 DiT raw 还是后处理) -bash dist/scripts/capture-and-compare-ltx23-forward.sh -``` - -监控:`grep -E 'model_output_cosine_sim|log_prob_mean_abs_diff|clipfrac' logs/*.log` - -## 常见陷阱(实战踩过) - -- **temb 错 1000 倍**:AdaLN 输入应是 `σ×1000` 而非 `σ`;错了 cosine≈0(与"权重错"是不同量级)。 -- **temb shape**:ltx_core 走 `[B,T,D]`,sglang 构 `[B,1,D]`,即使 σ 均匀,block 内 broadcast 行为不同 → 必须 `expand_temb_for_hidden`。 -- **guider 默认值**:LTX23 one-stage 默认 `video_cfg_scale=3 / modality=3 / rescale=0.7`,在 x0 上改 velocity 语义 → live cosine 0.94。train `forward_velocity` 等价 cfg=1/无STG/无modality/无rescale。 -- **identity guider 不能经 `extra_sampling_params` 传**:sglang `SamplingParams` 基类拒绝 LTX23 guider 字段(`400 unexpected keyword 'video_cfg_scale'`)→ 必须用 monkey patch `patch_ltx2_identity_guider.py` override `_get_ltx2_stage1_guider_params`。 -- **Flow-SDE 公式 / sigma_min**:`--ltx-dynamics-type Flow-SDE` 时 rollout 误用 SD3 `sde` 公式、或误用 `sigmas[-2]=0.1` 作 σ_min → `log_prob_mean_abs_diff` 2~15。sglang 侧需 `flow_sde` + `rollout_sigma_min`。 -- **guider 修复后必须重新 capture dump 再 compare**:旧 dump 是默认 guider 下生成的。 - -## 参考文档(深度细节，本地 `dist/docs/`，不进 git) - -- [dist/docs/ltx23_changes_overview.md](../../../dist/docs/ltx23_changes_overview.md) — 全部代码改动按 DEBUG/TRAIN/ROLLOUT 三分类 -- [dist/docs/ltx23_train_rollout_alignment_journey.md](../../../dist/docs/ltx23_train_rollout_alignment_journey.md) — Phase A/B/C 排查历程(SDE→temb→guider) -- [dist/docs/ltx23_sglang_rollout_train_troubleshooting.md](../../../dist/docs/ltx23_sglang_rollout_train_troubleshooting.md) — 两侧工程问题 P1–P13 + 跨边界 C1–C8 -- [dist/docs/ltx23_forward_alignment_test_report.md](../../../dist/docs/ltx23_forward_alignment_test_report.md) — 数值实验矩阵与 block-wise 二分 diff --git a/.gitignore b/.gitignore index 89451138..b0fb9d67 100644 --- a/.gitignore +++ b/.gitignore @@ -202,4 +202,6 @@ logs/ wandb/ # local debug scripts / alignment tools (not shipped with the library) -dist/ \ No newline at end of file +dist/ +scripts/run-ltx23-grpo-local.sh +.claude/skills/align-ltx-train-rollout/ \ No newline at end of file diff --git a/miles/backends/model_families/ltx.py b/miles/backends/model_families/ltx.py index 985b392c..09bae519 100644 --- a/miles/backends/model_families/ltx.py +++ b/miles/backends/model_families/ltx.py @@ -270,23 +270,6 @@ def resolve_transformer_checkpoint( ) -def preflight_ltx_models(args: Namespace) -> None: - """Ensure train + rollout can resolve the same model before Ray starts.""" - hf_model_id = resolve_hf_model_id(args) - checkpoint = resolve_transformer_checkpoint( - getattr(args, "diffusion_model", None), - explicit_path=getattr(args, "sglang_transformer_weights_path", None), - ) - ckpt_path = Path(checkpoint) - if _is_materialized_diffusers_checkpoint(ckpt_path): - _read_materialized_transformer_config(ckpt_path) - logger.info( - "LTX preflight ok: rollout model_path=%s train_checkpoint=%s", - hf_model_id, - checkpoint, - ) - - def server_kwargs_extras(args) -> dict: """Extra ``ServerArgs`` kwargs; call only when ``is_ltx_model(args)``.""" hf_model_id = resolve_hf_model_id(args) diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py index ca24b6d1..1737fad2 100644 --- a/miles/backends/sglang_diffusion_utils/configs/ltx.py +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -7,7 +7,6 @@ LTX_DEFAULT_MODEL_ID, ensure_materialized_model, is_ltx_model, - preflight_ltx_models, resolve_hf_model_id, resolve_materialized_model_dir, resolve_model_id, @@ -20,7 +19,6 @@ "LTX_DEFAULT_MODEL_ID", "ensure_materialized_model", "is_ltx_model", - "preflight_ltx_models", "resolve_hf_model_id", "resolve_materialized_model_dir", "resolve_model_id", diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index ae286830..f12a38f4 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -483,20 +483,6 @@ def _split_train_data_by_dp(self, data, dp_size): return rollout_data_refs -def _pythonpath_env_var() -> dict[str, str]: - """Ray workers do not inherit shell PYTHONPATH; propagate sglang source root.""" - sglang_python = os.environ.get("SGLANG_PYTHON") - parent_path = os.environ.get("PYTHONPATH", "") - if sglang_python: - prefix = sglang_python - if parent_path and not parent_path.startswith(prefix): - prefix = f"{sglang_python}:{parent_path}" - return {"PYTHONPATH": prefix} - if parent_path: - return {"PYTHONPATH": parent_path} - return {} - - def _base_rollout_engine_env_vars() -> dict[str, str]: return {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { "SGL_JIT_DEEPGEMM_PRECOMPILE": "false", @@ -528,7 +514,6 @@ def _ltx_alignment_env_vars(args) -> dict[str, str]: for name in ( LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS", - "MILES_LTX_IDENTITY_GUIDER", ): if os.environ.get(name): env[name] = os.environ[name] @@ -537,7 +522,8 @@ def _ltx_alignment_env_vars(args) -> dict[str, str]: def _build_rollout_engine_env_vars(args) -> dict[str, str]: env_vars = _base_rollout_engine_env_vars() - env_vars.update(_pythonpath_env_var()) + if os.environ.get("PYTHONPATH"): + env_vars["PYTHONPATH"] = os.environ["PYTHONPATH"] for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): if os.environ.get(cache_var): env_vars[cache_var] = os.environ[cache_var] diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index ef272b95..a331a39f 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -5,7 +5,7 @@ # (/sgl-workspace/miles/scripts/run-diffusion-grpo-ltx23-trainer-rollout.sh): # CPS dynamics, 3 SDE steps from candidates 0–9, clip-range 1e-4. # Rollout goes through sglang with weight sync; train/rollout forward alignment -# fixes stay on (ltxcore parity + AV-off + identity guider). +# fixes stay on (ltxcore parity + AV-off + gs=1 rollout alignment). # # GPU layout: single physical GPU colocate (train FSDP world_size=1 and sglang # rollout time-share one GPU via offload). Set NUM_GPUS>1 for multi-GPU diff --git a/scripts/run-ltx23-grpo-local.sh b/scripts/run-ltx23-grpo-local.sh deleted file mode 100755 index e97a72e9..00000000 --- a/scripts/run-ltx23-grpo-local.sh +++ /dev/null @@ -1,82 +0,0 @@ -#!/usr/bin/env bash -# LTX-2.3 GRPO 本地启动脚本（/data/wenhao 环境） -# -# 仅需 --diffusion-model Lightricks/LTX-2.3（train + rollout 共用 sglang overlay）。 -# 首次运行会自动 materialize 到 SGLANG_DIFFUSION_CACHE_ROOT。 -# -# 前置条件： -# - venv: /data/wenhao/.venvs/miles-diffusion -# - sglang: /data/wenhao/master_sglang/sglang/python -# - miles: /data/wenhao/master_miles/miles_diffusion -# -# 用法： -# # 前台调试 -# bash scripts/run-ltx23-grpo-local.sh -# -# # 冒烟测试（1 rollout，跳过 optimizer） -# NUM_ROLLOUT=1 ROLLOUT_BATCH_SIZE=1 N_SAMPLES_PER_PROMPT=2 NUM_STEPS_PER_ROLLOUT=1 \ -# SKIP_OPTIMIZER=1 bash scripts/run-ltx23-grpo-local.sh -# -# 可覆盖的环境变量： -# CUDA_VISIBLE_DEVICES GPU 编号（默认 2） -# NUM_ROLLOUT rollout 总数（默认 100） -# MILES_DIFFUSION_DEBUG 1=打印对齐指标（默认 1） -# SKIP_OPTIMIZER 1=冒烟模式，不更新权重（默认 0） -# LTX_DEV_SAFETENSORS 可选：dev .safetensors 覆盖 overlay 默认 DiT 权重 - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -MILES_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" - -# ── 路径 ────────────────────────────────────────────────────────────────── -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-2}" -export SGLANG_PYTHON="${SGLANG_PYTHON:-/data/wenhao/master_sglang/sglang/python}" -export PYTHONPATH="${SGLANG_PYTHON}${PYTHONPATH:+:${PYTHONPATH}}" - -export MILES_DATA_ROOT="${MILES_DATA_ROOT:-/data/wenhao/master_miles/miles_diffusion}" -export MILES_DATA_DISK_ROOT="${MILES_DATA_DISK_ROOT:-/data/wenhao/miles_diffusion}" -export LTX_HF_MODEL="${LTX_HF_MODEL:-Lightricks/LTX-2.3}" -export PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" -export HF_HOME="${HF_HOME:-/data/wenhao/hf_home}" -export SGLANG_DIFFUSION_CACHE_ROOT="${SGLANG_DIFFUSION_CACHE_ROOT:-/data/wenhao/sgl_diffusion_cache}" - -# ── 训练规模 ────────────────────────────────────────────────────────────── -export USE_LORA="${USE_LORA:-1}" -export NUM_ROLLOUT="${NUM_ROLLOUT:-100}" -export ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}" -export N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" -export NUM_STEPS_PER_ROLLOUT="${NUM_STEPS_PER_ROLLOUT:-2}" -export SAVE_INTERVAL="${SAVE_INTERVAL:-50}" - -# ── 对齐开关 ────────────────────────────────────────────────────────────── -export MILES_DIFFUSION_DEBUG="${MILES_DIFFUSION_DEBUG:-1}" -export LTX_DISABLE_AV_CROSS_ATTN="${LTX_DISABLE_AV_CROSS_ATTN:-1}" -export MILES_APPLY_LTX2_LTXCORE_PARITY="${MILES_APPLY_LTX2_LTXCORE_PARITY:-1}" - -# ── 冒烟 / 恢复 ─────────────────────────────────────────────────────────── -export SKIP_OPTIMIZER="${SKIP_OPTIMIZER:-0}" - -# ── 激活环境 ────────────────────────────────────────────────────────────── -VENV="${VENV:-/data/wenhao/.venvs/miles-diffusion}" -if [[ -f "${VENV}/bin/activate" ]]; then - # shellcheck source=/dev/null - source "${VENV}/bin/activate" -else - echo "[warn] venv not found: ${VENV}" >&2 -fi - -echo "============================================" -echo " LTX-2.3 GRPO Local Launch" -echo "============================================" -echo " GPU: ${CUDA_VISIBLE_DEVICES}" -echo " Rollouts: ${NUM_ROLLOUT}" -echo " Model: ${LTX_DEV_SAFETENSORS:-${LTX_HF_MODEL}}" -echo " Debug: ${MILES_DIFFUSION_DEBUG}" -echo " Skip optim: ${SKIP_OPTIMIZER}" -echo " Logs: ${MILES_DATA_DISK_ROOT}/logs/" -echo " Ckpt: ${MILES_DATA_DISK_ROOT}/ckpt/" -echo "============================================" - -cd "${MILES_ROOT}" -exec bash scripts/run-diffusion-grpo-ltx23-sglang.sh diff --git a/train_diffusion.py b/train_diffusion.py index 6cba5646..289d56c4 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -14,11 +14,6 @@ def train(args): configure_logger() logger = logging.getLogger(__name__) - if getattr(args, "diffusion_model_type", None) == "ltx": - from miles.backends.model_families.ltx import preflight_ltx_models - - preflight_ltx_models(args) - # allocate the GPUs logger.info("train: creating placement groups") pgs = create_placement_groups(args) From 0eb442630c7c0d5f83555195d9737efc54ee6a7c Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 12 Jun 2026 02:16:41 +0000 Subject: [PATCH 16/31] clean code --- miles/backends/fsdp_utils/actor.py | 14 --- .../sglang_diffusion_engine.py | 83 ++++++++-------- miles/ray/rollout.py | 98 ++++++------------- 3 files changed, 73 insertions(+), 122 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index db5cae7d..0cc62687 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -175,8 +175,6 @@ def sleep(self) -> None: print_memory("before offload DiT") - self.optimizer.zero_grad(set_to_none=True) - _reshard_fsdp2_model(self.model) self.model.cpu() move_torch_optimizer(self.optimizer, "cpu") clear_memory() @@ -217,7 +215,6 @@ def update_weights(self) -> None: # type: ignore[override] ray.get(self.rollout_manager.clear_num_new_engines.remote()) self.weight_updater.update_weights() - dist.barrier(group=get_gloo_group()) clear_memory() def _get_init_weight_context_manager(self): @@ -898,17 +895,6 @@ def _cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict: return out -def _reshard_fsdp2_model(model: torch.nn.Module) -> None: - """Drop FSDP2 unsharded views so model.cpu() can release GPU memory.""" - if hasattr(model, "reshard"): - model.reshard() - return - for module in model.modules(): - reshard = getattr(module, "reshard", None) - if callable(reshard): - reshard() - - @torch.no_grad() def move_torch_optimizer(optimizer, device): """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py""" diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 11a8c7be..929e7022 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -10,11 +10,7 @@ from sglang.multimodal_gen.runtime.launch_server import kill_process_tree from miles.backends.sglang_diffusion_utils.configs import ltx as ltx_config -from miles.backends.sglang_diffusion_utils.monkey_patches import ( - ROLLOUT_PATCH_GROUP_ENV, - apply_rollout_patch_group, - resolve_rollout_patch_group, -) +from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV from miles.ray.ray_actor import RayActor from miles.utils.http_utils import get_host_info @@ -45,27 +41,21 @@ 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. - apply_rollout_patch_group(os.environ.get(ROLLOUT_PATCH_GROUP_ENV)) - - # Colocate weight sync keeps the full CUDA_VISIBLE_DEVICES (so CUDA IPC works - # across GPUs); pin the DiT to its assigned local cuda index instead. - local_cuda_rank = os.environ.get("MILES_SGLANG_LOCAL_CUDA_RANK") - if local_cuda_rank is not None: - from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker - - pinned_rank = int(local_cuda_rank) - _orig_init = GPUWorker.__init__ - - def _patched_init(self, local_rank, rank, master_port, server_args): - return _orig_init(self, pinned_rank, rank, master_port, server_args) + from miles.backends.sglang_diffusion_utils.monkey_patches import ( + apply_ltx2_rollout_patches, + apply_sgld_monkey_patches, + ) - GPUWorker.__init__ = _patched_init + 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() from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process return run_scheduler_process(*args, **kwargs) -def _launch_server_target(server_args, patch_group: str | None = None): +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 @@ -73,17 +63,15 @@ def _launch_server_target(server_args, patch_group: str | None = None): if server_args.attention_backend_config is not None: server_args.attention_backend_config = addict.Dict(server_args.attention_backend_config) - # 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 - # launch_server's own module to point at our wrapper — pickle then carries - # the miles qualname across to the grandchild, which applies the patches (and - # colocate GPU pin) before calling the real scheduler entrypoint. - if patch_group is not None or os.environ.get("MILES_SGLANG_LOCAL_CUDA_RANK") is not None: + 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 + # launch_server's own module to point at our wrapper — pickle then carries + # the miles qualname across to the grandchild, which applies the patch before + # calling the real scheduler entrypoint. import sglang.multimodal_gen.runtime.launch_server as _ls_mod _ls_mod.run_scheduler_process = _scheduler_process_with_sgld_monkey_patches - if patch_group is not None: - os.environ[ROLLOUT_PATCH_GROUP_ENV] = patch_group from sglang.multimodal_gen.runtime.launch_server import launch_server launch_server(server_args) @@ -91,14 +79,14 @@ def _launch_server_target(server_args, patch_group: str | None = None): def launch_server_process( server_args: ServerArgs, - patch_group: str | None = None, + 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, patch_group), + args=(server_args, apply_rollout_patches), ) p.start() @@ -174,12 +162,24 @@ 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() - patch_group = resolve_rollout_patch_group(self.args) - if patch_group is not None: - logger.info(f"Launching sglang-d with rollout patch group: {patch_group}") + apply_sgld = bool(getattr(self.args, "apply_sgld_monkey_patches", False)) + apply_ltx = ( + ltx_config.is_ltx_model(self.args) + and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" + ) + 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), - patch_group=patch_group, + apply_rollout_patches=use_rollout_patches, ) if self.node_rank == 0 and self.router_ip and self.router_port: @@ -198,14 +198,13 @@ def _pin_to_assigned_gpu(self): cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not cvd: return - local_id = _to_local_gpu_id(self.base_gpu_id) - # Keep the full CUDA_VISIBLE_DEVICES so colocated weight sync can CUDA-IPC - # buckets from FSDP ranks on other GPUs; pin the DiT to its local cuda - # index via MILES_SGLANG_LOCAL_CUDA_RANK (applied in the scheduler child). - os.environ["MILES_SGLANG_LOCAL_CUDA_RANK"] = str(local_id) + visible = [x.strip() for x in cvd.split(",") if x.strip()] + local_idx = _to_local_gpu_id(self.base_gpu_id) + pinned = visible[local_idx] + os.environ["CUDA_VISIBLE_DEVICES"] = pinned logger.info( - f"Engine rank={self.rank}: rollout cuda:{local_id} " - f"(base_gpu_id={self.base_gpu_id}, CUDA_VISIBLE_DEVICES={cvd})" + f"Engine rank={self.rank}: pinned CUDA_VISIBLE_DEVICES={pinned} " + f"(base_gpu_id={self.base_gpu_id}, local_idx={local_idx})" ) def _make_request(self, endpoint: str, payload: dict | None = None): diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index f12a38f4..f079e300 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -14,6 +14,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS +from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import SGLangDiffusionEngine from miles.rollout.base_types import call_rollout_fn from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor @@ -45,19 +46,13 @@ def __init__(self, args, pg): logger.info("RolloutManager init start") self.args = args self.pg = pg - if self.args.debug_train_only: - logger.info("RolloutManager: debug_train_only, skipping sglang router.") - router_addr = None - else: - logger.info("RolloutManager: starting router...") - _start_router(args) - logger.info("RolloutManager: router started, init tracking...") - router_addr = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + logger.info("RolloutManager: starting router...") + _start_router(args) + logger.info("RolloutManager: router started, init tracking...") # TODO make args immutable - init_tracking(args, primary=False, router_addr=router_addr) - if not self.args.debug_train_only: - logger.info("RolloutManager: init http client...") - init_http_client(args) + init_tracking(args, primary=False, router_addr=f"http://{args.sglang_router_ip}:{args.sglang_router_port}") + logger.info("RolloutManager: init http client...") + init_http_client(args) logger.info("RolloutManager: loading data source...") data_source_cls = load_function(self.args.data_source_path) @@ -449,7 +444,7 @@ def _log_images( if frame.max() <= 1.0 + 1e-3: frame = frame * 255.0 frame = np.clip(frame, 0, 255).astype(np.uint8) - reward = s.get_reward_value(self.args, reward_key=reward_key) + 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: log_dict[media_key] = images @@ -483,62 +478,10 @@ def _split_train_data_by_dp(self, data, dp_size): return rollout_data_refs -def _base_rollout_engine_env_vars() -> dict[str, str]: - return {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", - } - - -def _ltx_alignment_env_vars(args) -> dict[str, str]: - """Propagate LTX train/rollout alignment flags into sglang Ray workers. - - Ray ``runtime_env`` does not inherit the parent shell env; monkey patches in - ``sglang_diffusion_engine`` read these variables at scheduler startup. - - TODO(PR4): replace env bridging with explicit sglang ``ServerArgs`` fields - (or upstream train-mode guider / video-only flags) so miles does not depend - on opaque env + runtime monkey patches. - """ - env: dict[str, str] = {} - if getattr(args, "ltx_disable_av_cross_attn", False): - env["MILES_LTX_DISABLE_AV_CROSS"] = "1" - from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV - - for name in ( - LTX_ROLLOUT_PATCHES_ENV, - "MILES_LTX_DISABLE_AV_CROSS", - ): - if os.environ.get(name): - env[name] = os.environ[name] - return env - - -def _build_rollout_engine_env_vars(args) -> dict[str, str]: - env_vars = _base_rollout_engine_env_vars() - if os.environ.get("PYTHONPATH"): - env_vars["PYTHONPATH"] = os.environ["PYTHONPATH"] - for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): - if os.environ.get(cache_var): - env_vars[cache_var] = os.environ[cache_var] - env_vars.update(_ltx_alignment_env_vars(args)) - return env_vars - - def init_rollout_engines(args, pg, all_rollout_engines): if args.debug_train_only: return 0 - from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import ( - SGLangDiffusionEngine, - ) - num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node) num_engines = args.rollout_num_gpus // num_gpu_per_engine assert len(all_rollout_engines) == num_engines @@ -567,7 +510,30 @@ def init_rollout_engines(args, pg, all_rollout_engines): placement_group_bundle_index=reordered_bundle_indices[i * num_gpu_per_engine], ) - env_vars = _build_rollout_engine_env_vars(args) + 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", + } + if os.environ.get("PYTHONPATH"): + env_vars["PYTHONPATH"] = os.environ["PYTHONPATH"] + for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): + if os.environ.get(cache_var): + env_vars[cache_var] = os.environ[cache_var] + from miles.backends.sglang_diffusion_utils.configs.ltx import is_ltx_model + from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV + + if is_ltx_model(args): + if getattr(args, "ltx_disable_av_cross_attn", False): + env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" + for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): + if os.environ.get(name): + env_vars[name] = os.environ[name] rollout_engine = RolloutRayActor.options( num_cpus=num_cpus, From 1716ae64980f3f15fff411fc31a777cf7075cfc9 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 12 Jun 2026 06:51:39 +0000 Subject: [PATCH 17/31] refactor(ltx): slim actor and rollout diffs against main Move model loading, SDE window slicing, and debug metrics into shared helpers so LTX-specific logic stays in pipeline configs and model_families. --- miles/backends/fsdp_utils/actor.py | 111 ++++++------------ miles/backends/fsdp_utils/configs/ltx.py | 24 ++-- .../configs/train_pipeline_config.py | 108 ++++++++++++++++- miles/backends/model_families/ltx.py | 14 +++ .../sglang_diffusion_engine.py | 22 ++++ miles/ray/rollout.py | 63 ++-------- miles/rollout/rm_hub/video_pickscore.py | 13 ++ miles/rollout/sglang_diffusion_rollout.py | 37 +----- miles/rollout/step_strategy_hub.py | 8 +- miles/utils/train_metric_utils.py | 33 ++++++ miles/utils/types.py | 11 ++ 11 files changed, 268 insertions(+), 176 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 0cc62687..5ae1e54e 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -7,7 +7,6 @@ import ray import torch import torch.distributed as dist -from diffusers import DiffusionPipeline from miles.ray.train_actor import TrainRayActor from miles.utils.context_utils import with_defer @@ -70,41 +69,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config = get_train_pipeline_config(diffusion_model_id) - if self.train_pipeline_config.is_diffusers_pipeline: - with self._get_init_weight_context_manager(): - pipeline = DiffusionPipeline.from_pretrained( - diffusion_model_id, - 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 - else: - model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( - self.args, init_context_factory=self._get_init_weight_context_manager, - ) + model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( + self.args, init_context_factory=self._get_init_weight_context_manager, + ) if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) model.train() - if args.gradient_checkpointing: - if hasattr(model, "enable_gradient_checkpointing"): - model.enable_gradient_checkpointing() - elif hasattr(model, "set_gradient_checkpointing"): - model.set_gradient_checkpointing(True) - else: - logger.warning( - "gradient_checkpointing requested but model %s exposes neither " - "enable_gradient_checkpointing() nor set_gradient_checkpointing(); " - "skipping.", - type(model).__name__, - ) + if args.gradient_checkpointing and self.train_pipeline_config.is_diffusers_pipeline: + model.enable_gradient_checkpointing() model.to(torch.cuda.current_device()) @@ -487,25 +462,18 @@ def _build_train_grids( ) sde_step_indices = sde_step_indices_list[traj_idx] - sde_indices_per_sample: torch.Tensor | None = 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] - next_latents = next_latents[sde_indices_tensor] - timesteps = timesteps[sde_indices_tensor] - log_prob_old = log_prob_old[sde_indices_tensor] - advantage = advantage[: sde_indices_tensor.numel()] - if rollout_model_output is not None: - n_mo = int(rollout_model_output.shape[0]) - n_win = int(sde_indices_tensor.numel()) - if n_mo != n_win: - # Full-length debug tensors (legacy): index by global step. - rollout_model_output = rollout_model_output[sde_indices_tensor] - # else: sglang packs debug outputs in SDE-window order (0..W-1). - current_window_size = int(sde_indices_tensor.numel()) - sde_indices_per_sample = sde_indices_tensor - else: - current_window_size = default_window_size + window_batch = train_pipeline_config.apply_sde_step_window( + latents=latents, + next_latents=next_latents, + timesteps=timesteps, + log_prob_old=log_prob_old, + advantage=advantage, + rollout_model_output=rollout_model_output, + sde_step_indices=sde_step_indices, + default_window_size=default_window_size, + device=device, + ) + current_window_size = window_batch.window_size if sde_window_size is None: sde_window_size = current_window_size @@ -514,15 +482,15 @@ def _build_train_grids( f"for now per-sample SDE window length must match across microbatch " f"(got {sde_window_size} and {current_window_size})" ) - latents_list.append(latents) - next_latents_list.append(next_latents) - timesteps_list.append(timesteps) - log_prob_old_list.append(log_prob_old) - advantage_list.append(advantage) - if rollout_model_output is not None: - rollout_model_outputs_list.append(rollout_model_output) - if sde_indices_per_sample is not None: - sde_indices_per_sample_list.append(sde_indices_per_sample) + latents_list.append(window_batch.latents) + next_latents_list.append(window_batch.next_latents) + timesteps_list.append(window_batch.timesteps) + log_prob_old_list.append(window_batch.log_prob_old) + advantage_list.append(window_batch.advantage) + if window_batch.rollout_model_output is not None: + rollout_model_outputs_list.append(window_batch.rollout_model_output) + if window_batch.step_indices is not None: + sde_indices_per_sample_list.append(window_batch.step_indices) latents_window = torch.stack(latents_list, dim=0) next_latents_window = torch.stack(next_latents_list, dim=0) @@ -803,24 +771,15 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: torch.mean(torch.abs(log_prob_new - log_prob_old_tile)).detach() ) # To log model output diff, please enable --diffusion-debug-mode - rollout_mo_window = grids.get("rollout_model_outputs") - if rollout_mo_window is not None: - rollout_mo_tile = rollout_mo_window[sample_indices][:, tstep_indices] - rollout_mo_flat = rollout_mo_tile.reshape( - tile_sample_count * tile_tstep_count, *rollout_mo_tile.shape[2:] - ) - diff = (noise_pred_flat.float() - rollout_mo_flat.float()).abs() - ref_max = rollout_mo_flat.float().abs().max() + 1e-30 - 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()) - flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) - flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) - log_stats["model_output_cosine_sim"].append( - torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1) - .mean() - .detach() - ) + train_metric_utils.log_model_output_debug_metrics( + log_stats, + noise_pred_flat=noise_pred_flat, + grids=grids, + sample_indices=sample_indices, + tstep_indices=tstep_indices, + tile_sample_count=tile_sample_count, + tile_tstep_count=tile_tstep_count, + ) return loss diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index ac2175f7..b9b7e70f 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -70,6 +70,13 @@ def to(self, device): 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 prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: @@ -143,18 +150,13 @@ def build_sde_extra( tstep_indices: torch.Tensor, args, ) -> dict | None: - if grids.get("sde_step_indices_window") is None: + extra = super().build_sde_extra(scheduler, grids, sample_indices, tstep_indices, args) + if extra is None: return None - - idx = grids["sde_step_indices_window"][sample_indices][:, tstep_indices] - idx = idx.reshape(-1).long() - - return { - "sigmas": scheduler.sigmas, - "sde_step_indices": idx, - "dynamics_type": getattr(args, "ltx_dynamics_type", "cps"), - "sigma_min_override": getattr(args, "ltx_sigma_min", None), - } + extra["sigmas"] = scheduler.sigmas + extra["dynamics_type"] = getattr(args, "ltx_dynamics_type", "cps") + extra["sigma_min_override"] = getattr(args, "ltx_sigma_min", None) + return extra def expand_cond_for_timestep_batch(self, cond_kwargs: dict, batch_size: int) -> dict: out: dict = {} diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 0a8d5506..10e930da 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -15,12 +15,27 @@ import abc from argparse import Namespace -from typing import Any +from dataclasses import dataclass +from typing import Any, Callable import torch from miles.utils.types import CondKwargs, DiTTrajectory +@dataclass +class SdeWindowBatch: + """One sample's tensors after optional SDE-window slicing.""" + + latents: torch.Tensor + next_latents: torch.Tensor + timesteps: torch.Tensor + log_prob_old: torch.Tensor + advantage: torch.Tensor + rollout_model_output: torch.Tensor | None + window_size: int + step_indices: torch.Tensor | None + + _REGISTRY: dict[str, type["TrainPipelineConfig"]] = {} @@ -63,6 +78,92 @@ def scale_timesteps_for_sde(self, timesteps: torch.Tensor) -> torch.Tensor: return timesteps / float(self.sde_timestep_divisor) return timesteps + def load_model_and_scheduler( + self, + args: Namespace, + init_context_factory: Callable[[], Any], + ) -> tuple[torch.nn.Module, Any]: + """Load DiT + scheduler. Default: diffusers ``DiffusionPipeline`` (transformer only).""" + from diffusers import DiffusionPipeline + + diffusion_model_id = args.diffusion_model or args.hf_checkpoint + master_dtype_name = getattr(args, "fsdp_master_dtype", "bf16") + master_dtype = { + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + }[master_dtype_name] + + with init_context_factory(): + pipeline = DiffusionPipeline.from_pretrained( + diffusion_model_id, + 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 apply_sde_step_window( + self, + *, + latents: torch.Tensor, + next_latents: torch.Tensor, + timesteps: torch.Tensor, + log_prob_old: torch.Tensor, + advantage: torch.Tensor, + rollout_model_output: torch.Tensor | None, + sde_step_indices: list[int] | None, + default_window_size: int, + device: torch.device, + ) -> SdeWindowBatch: + """Slice trajectory/objective tensors to the rollout SDE window.""" + step_indices: torch.Tensor | None = None + if sde_step_indices is not None: + step_indices = torch.as_tensor(sde_step_indices, device=device, dtype=torch.long) + latents = latents[step_indices] + next_latents = next_latents[step_indices] + timesteps = timesteps[step_indices] + log_prob_old = log_prob_old[step_indices] + advantage = advantage[: step_indices.numel()] + if rollout_model_output is not None: + n_mo = int(rollout_model_output.shape[0]) + n_win = int(step_indices.numel()) + if n_mo != n_win: + # Full-length debug tensors (legacy): index by global step. + rollout_model_output = rollout_model_output[step_indices] + # else: sglang packs debug outputs in SDE-window order (0..W-1). + window_size = int(step_indices.numel()) + else: + window_size = default_window_size + + return SdeWindowBatch( + latents=latents, + next_latents=next_latents, + timesteps=timesteps, + log_prob_old=log_prob_old, + advantage=advantage, + rollout_model_output=rollout_model_output, + window_size=window_size, + step_indices=step_indices, + ) + + def resolve_tile_sde_step_indices( + self, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + ) -> torch.Tensor | None: + """Map a training tile to global denoising step indices.""" + window = grids.get("sde_step_indices_window") + if window is None: + return None + return window[sample_indices][:, tstep_indices].reshape(-1).long() + def prepare_trajectory( self, traj: DiTTrajectory, @@ -128,7 +229,10 @@ def build_sde_extra( args: Namespace, ) -> dict | None: """Optional per-tile metadata for model-specific SDE log_prob.""" - return None + idx = self.resolve_tile_sde_step_indices(grids, sample_indices, tstep_indices) + if idx is None: + return None + return {"sde_step_indices": idx} def expand_cond_for_timestep_batch( self, diff --git a/miles/backends/model_families/ltx.py b/miles/backends/model_families/ltx.py index 09bae519..cca6b42e 100644 --- a/miles/backends/model_families/ltx.py +++ b/miles/backends/model_families/ltx.py @@ -343,6 +343,20 @@ def patch_rollout_sampling_params( 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 getattr(args, "ltx_disable_av_cross_attn", False): + env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" + for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): + if os.environ.get(name): + env_vars[name] = os.environ[name] + + def register_args(parser: ArgumentParser) -> None: parser.add_argument( "--ltx-frames", diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 929e7022..ca107f62 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -17,6 +17,28 @@ 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.model_families.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: diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index f079e300..b48ed3ca 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -2,7 +2,6 @@ import itertools import logging import multiprocessing -import os import random import time from pathlib import Path @@ -14,7 +13,10 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, 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 @@ -28,7 +30,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) @@ -391,16 +393,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") - if (sample.train_metadata or {}).get("sde_step_indices") is not None - else ( - sample.dit_trajectory.sde_step_indices.tolist() - if sample.dit_trajectory is not None and sample.dit_trajectory.sde_step_indices is not None - else None - ) - for sample in samples - ], + "sde_step_indices": [sample.get_sde_step_indices() for sample in samples], } return train_data @@ -423,6 +416,8 @@ def _log_images( own namespace at least groups them in one UI section. """ import wandb + from miles.rollout.rm_hub.video_pickscore import first_frame_for_wandb + log_dict: dict = {} for media_key, samples in media_key_to_samples.items(): images = [] @@ -430,20 +425,9 @@ def _log_images( t = s.generated_output if t is None: continue - try: - from miles.rollout.rm_hub.video_pickscore import ( - fchw_frame_to_hwc_uint8, - generated_output_to_fchw, - ) - - frame = fchw_frame_to_hwc_uint8(generated_output_to_fchw(t)[0]) - except (ValueError, TypeError): - if t.ndim != 4: - 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) + frame = first_frame_for_wandb(t) + if frame is None: + continue 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: @@ -510,30 +494,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", - } - if os.environ.get("PYTHONPATH"): - env_vars["PYTHONPATH"] = os.environ["PYTHONPATH"] - for cache_var in ("SGLANG_DIFFUSION_CACHE_ROOT", "HF_HOME", "TMPDIR"): - if os.environ.get(cache_var): - env_vars[cache_var] = os.environ[cache_var] - from miles.backends.sglang_diffusion_utils.configs.ltx import is_ltx_model - from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV - - if is_ltx_model(args): - if getattr(args, "ltx_disable_av_cross_attn", False): - env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" - for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): - if os.environ.get(name): - env_vars[name] = os.environ[name] + env_vars = build_rollout_engine_env_vars(args) rollout_engine = RolloutRayActor.options( num_cpus=num_cpus, diff --git a/miles/rollout/rm_hub/video_pickscore.py b/miles/rollout/rm_hub/video_pickscore.py index 73037671..c4d2c3c4 100644 --- a/miles/rollout/rm_hub/video_pickscore.py +++ b/miles/rollout/rm_hub/video_pickscore.py @@ -54,6 +54,19 @@ def fchw_frame_to_hwc_uint8(frame_chw: torch.Tensor) -> np.ndarray: 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])) diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 7026cd2e..ab56a993 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -29,15 +29,6 @@ logger = logging.getLogger(__name__) -def _resolve_diffusion_model_type(args: Namespace) -> str: - from miles.backends.model_families.ltx import is_ltx_model - - model_type = (getattr(args, "diffusion_model_type", "auto") or "auto").lower() - if model_type != "auto": - return model_type - return "ltx" if is_ltx_model(args) else "sd3" - - def build_rollout_sampling_params( args: Namespace, *, @@ -60,8 +51,6 @@ def build_rollout_sampling_params( "true_cfg_scale": getattr(args, "diffusion_true_cfg_scale", None), } - model_type = _resolve_diffusion_model_type(args) - if evaluation: sampling_params["rollout"] = False else: @@ -79,9 +68,9 @@ 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"). - if model_type == "ltx": - from miles.backends.model_families.ltx import patch_rollout_sampling_params + from miles.backends.model_families.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: @@ -164,23 +153,6 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: self.remaining_batch_size += len(samples) -def _call_step_strategy( - step_strategy_fn: Callable, - args: Namespace, - sample: Sample, - num_steps: int, - seed: int, - rollout_id: int, -) -> tuple[list[int] | None, list[int] | None]: - """Invoke a step-strategy hub function; pass ``rollout_id`` when supported.""" - params = inspect.signature(step_strategy_fn).parameters - if "rollout_id" in params: - return step_strategy_fn( - args, sample, num_steps, seed, rollout_id=rollout_id - ) - return step_strategy_fn(args, sample, num_steps, seed) - - async def generate_microgroup( args: Namespace, microgroup: list[Sample], sampling_params: dict[str, Any], *, evaluation: bool = False ) -> list[Sample]: @@ -193,13 +165,12 @@ async def generate_microgroup( # SGL-D TODO: support seed list for multiple samples in one request # currently only support assigning the first seed, SGL-D generates samples with seed, seed+1, seed+2, ... if not evaluation and state.step_strategy_fn is not None: - sde_indices, return_indices = _call_step_strategy( - state.step_strategy_fn, + sde_indices, return_indices = state.step_strategy_fn( args, microgroup[0], int(sampling_params["num_inference_steps"]), int(sampling_params["seed"]), - int(getattr(state, "rollout_id", 0) or 0), + 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 diff --git a/miles/rollout/step_strategy_hub.py b/miles/rollout/step_strategy_hub.py index 922aa675..ce6d3e4a 100644 --- a/miles/rollout/step_strategy_hub.py +++ b/miles/rollout/step_strategy_hub.py @@ -3,8 +3,9 @@ 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``) also -accept ``rollout_id`` via keyword — see ``miles.rollout.sglang_diffusion_rollout``. +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. """ @@ -67,13 +68,14 @@ def ltx_sde_candidates( 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.""" + del sample, rollout_id window_size = int(args.diffusion_sde_window_size) range_raw = getattr(args, "diffusion_sde_window_range", None) if range_raw: diff --git a/miles/utils/train_metric_utils.py b/miles/utils/train_metric_utils.py index 184005a8..b39ab3e6 100644 --- a/miles/utils/train_metric_utils.py +++ b/miles/utils/train_metric_utils.py @@ -2,6 +2,8 @@ from argparse import Namespace from copy import deepcopy +import torch + from miles.utils import tracking_utils from miles.utils.metric_utils import compute_rollout_step from miles.utils.timer import Timer @@ -9,6 +11,37 @@ logger = logging.getLogger(__name__) +def log_model_output_debug_metrics( + log_stats: dict[str, list[torch.Tensor]], + *, + noise_pred_flat: torch.Tensor, + grids: dict, + sample_indices: torch.Tensor, + tstep_indices: torch.Tensor, + tile_sample_count: int, + tile_tstep_count: int, +) -> None: + """Compare train-side noise_pred against rollout debug model outputs.""" + rollout_mo_window = grids.get("rollout_model_outputs") + if rollout_mo_window is None: + return + + rollout_mo_tile = rollout_mo_window[sample_indices][:, tstep_indices] + rollout_mo_flat = rollout_mo_tile.reshape( + tile_sample_count * tile_tstep_count, *rollout_mo_tile.shape[2:] + ) + diff = (noise_pred_flat.float() - rollout_mo_flat.float()).abs() + ref_max = rollout_mo_flat.float().abs().max() + 1e-30 + 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()) + flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) + flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) + log_stats["model_output_cosine_sim"].append( + torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1).mean().detach() + ) + + def log_perf_data_raw(rollout_id: int, args: Namespace, is_primary_rank: bool) -> None: timer_instance = Timer() log_dict_raw = deepcopy(timer_instance.log_dict()) diff --git a/miles/utils/types.py b/miles/utils/types.py index 4bc53525..66075be6 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -136,3 +136,14 @@ def get_reward_value(self, args, *, reward_key: str | None = None) -> float: 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 From c45447d60c4be30b5c8c42943f50c6710402764b Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Fri, 12 Jun 2026 07:09:20 +0000 Subject: [PATCH 18/31] clean code --- train_diffusion.py | 1 - 1 file changed, 1 deletion(-) diff --git a/train_diffusion.py b/train_diffusion.py index 289d56c4..4e7caadb 100644 --- a/train_diffusion.py +++ b/train_diffusion.py @@ -13,7 +13,6 @@ def train(args): configure_logger() logger = logging.getLogger(__name__) - # allocate the GPUs logger.info("train: creating placement groups") pgs = create_placement_groups(args) From 590d47f8b520c0b5d833e6b2dac52824b6715f8c Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 22 Jun 2026 11:55:34 +0000 Subject: [PATCH 19/31] refactor(ltx): Wan-style train adapter without base config churn Move LTX load/forward/SDE into configs/ltx.py; restore main SD3 paths in actor via is_ltx_model branches. Revert train_pipeline_config.py to main. Extract LTX dynamics to ltx_sde.py; sde_log_prob.py unchanged vs main. --- miles/backends/fsdp_utils/actor.py | 306 +++++++---- miles/backends/fsdp_utils/configs/ltx.py | 478 +++++++++++++++++- .../configs/train_pipeline_config.py | 214 +------- miles/backends/fsdp_utils/ltx_sde.py | 144 ++++++ 4 files changed, 797 insertions(+), 345 deletions(-) create mode 100644 miles/backends/fsdp_utils/ltx_sde.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 636f0317..1a34cc67 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -6,9 +6,12 @@ import ray import torch import torch.distributed as dist +from diffusers import DiffusionPipeline 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 +import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config +from miles.backends.fsdp_utils.configs.ltx import is_ltx_model from miles.ray.train_actor import TrainRayActor from miles.utils import tracking_utils, train_metric_utils from miles.utils.context_utils import with_defer @@ -20,8 +23,6 @@ from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils import init_tracking -import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config - from . import checkpoint from .configs.train_pipeline_config import get_train_pipeline_config from .diffusion_update_weight_utils import DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA @@ -65,32 +66,48 @@ 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) - diffusion_model_id = args.diffusion_model or args.hf_checkpoint + diffusion_model_id = args.diffusion_model or args.hf_checkpoint self.train_pipeline_config = get_train_pipeline_config(diffusion_model_id) + self._is_ltx = is_ltx_model(args) - model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( - self.args, init_context_factory=self._get_init_weight_context_manager, - ) + if self._is_ltx: + model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( + self.args, init_context_factory=self._get_init_weight_context_manager, + ) + else: + 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 if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) model.train() - if args.gradient_checkpointing and self.train_pipeline_config.is_diffusers_pipeline: + if args.gradient_checkpointing and not self._is_ltx: model.enable_gradient_checkpointing() model.to(torch.cuda.current_device()) self.train_pipeline_config.preprocess_model_before_fsdp(model) + fsdp_wrap = getattr(self.train_pipeline_config, "fsdp_wrap_classes", None) if self._is_ltx else None model = apply_fsdp2( model, mesh=self.parallel_state.dp_mesh, cpu_offload=self.args.fsdp_cpu_offload, args=self.args, - train_pipeline_config=self.train_pipeline_config, + fsdp_wrap_classes=fsdp_wrap, ) # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() @@ -294,7 +311,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: 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 - if not getattr(self.train_pipeline_config, "supports_cfg", True): + if self._is_ltx: use_cfg = False # ------------- KL loss ------------- @@ -321,16 +338,22 @@ 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 model-specific reconstruction. timesteps_ref = dit_trajectories[0].timesteps.to(device).float() sigmas_snapshot = getattr(dit_trajectories[0], "sigmas", None) 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_pipeline_config.resolve_sigmas_ref( - timesteps_ref, sigmas_snapshot, self.scheduler, - ) + if self._is_ltx: + sigmas_ref = self.train_pipeline_config.resolve_sigmas_ref( + timesteps_ref, sigmas_snapshot, self.scheduler, + ) + else: + 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)]) self.scheduler.timesteps = timesteps_ref self.scheduler.sigmas = sigmas_ref @@ -440,16 +463,21 @@ def _build_train_grids( dit_trajectories[traj_idx], device ) - # prepare cond kwargs (denoising_env + model-specific geometry when needed) + # prepare cond kwargs (denoising_env) denoising_env = denoising_envs[traj_idx] - positive_cond_kwargs_list.append( - train_pipeline_config.build_train_cond_kwargs( - denoising_env.pos_cond_kwargs, - latents=latents, - args=self.args, - device=device, + if self._is_ltx: + positive_cond_kwargs_list.append( + train_pipeline_config.build_train_cond_kwargs( + denoising_env.pos_cond_kwargs, + latents=latents, + args=self.args, + device=device, + ) + ) + else: + positive_cond_kwargs_list.append( + train_pipeline_config.prepare_cond_kwargs(denoising_env.pos_cond_kwargs, device) ) - ) if use_cfg: negative_cond_kwargs_list.append( train_pipeline_config.prepare_cond_kwargs(denoising_env.neg_cond_kwargs, device) @@ -467,18 +495,22 @@ def _build_train_grids( ) sde_step_indices = sde_step_indices_list[traj_idx] - window_batch = train_pipeline_config.apply_sde_step_window( - latents=latents, - next_latents=next_latents, - timesteps=timesteps, - log_prob_old=log_prob_old, - advantage=advantage, - rollout_model_output=rollout_model_output, - sde_step_indices=sde_step_indices, - default_window_size=default_window_size, - device=device, - ) - current_window_size = window_batch.window_size + 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] + next_latents = next_latents[sde_indices_tensor] + timesteps = timesteps[sde_indices_tensor] + log_prob_old = log_prob_old[sde_indices_tensor] + advantage = advantage[: sde_indices_tensor.numel()] + if rollout_model_output is not None: + 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 if sde_window_size is None: sde_window_size = current_window_size @@ -487,15 +519,15 @@ def _build_train_grids( f"for now per-sample SDE window length must match across microbatch " f"(got {sde_window_size} and {current_window_size})" ) - latents_list.append(window_batch.latents) - next_latents_list.append(window_batch.next_latents) - timesteps_list.append(window_batch.timesteps) - log_prob_old_list.append(window_batch.log_prob_old) - advantage_list.append(window_batch.advantage) - if window_batch.rollout_model_output is not None: - rollout_model_outputs_list.append(window_batch.rollout_model_output) - if window_batch.step_indices is not None: - sde_indices_per_sample_list.append(window_batch.step_indices) + latents_list.append(latents) + next_latents_list.append(next_latents) + timesteps_list.append(timesteps) + log_prob_old_list.append(log_prob_old) + 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) @@ -632,7 +664,12 @@ 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_pipeline_config.scale_timesteps_for_sde(timesteps_flat) + if self._is_ltx: + timesteps_for_sde = timesteps_flat / float( + getattr(train_pipeline_config, "sde_timestep_divisor", 1000.0) + ) + else: + 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. @@ -677,52 +714,80 @@ def _forward_tile( latents_input = latents_flat.to(forward_dtype) timesteps_input = timesteps_for_model.to(forward_dtype) - def _forward(cond: dict) -> torch.Tensor: - return train_pipeline_config.forward_velocity( - self.model, latents_input, timesteps_input, cond, - ) - - 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 = train_pipeline_config.forward_velocity_cfg_joint( - self.model, latents_input, timesteps_input, joint_cond, - ) - 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, + if self._is_ltx: + def _forward(cond: dict) -> torch.Tensor: + return train_pipeline_config.forward_velocity( + self.model, latents_input, timesteps_input, cond, ) - noise_pred_flat = _compute_noise_pred() + def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: + adapter_ctx = self.model.disable_adapter() if disable_adapter else nullcontext() + with adapter_ctx: + return _forward(pos_cond_tile) - sde_extra = train_pipeline_config.build_sde_extra( - self.scheduler, grids, sample_indices, tstep_indices, self.args, - ) + noise_pred_flat = _compute_noise_pred() + sde_extra = train_pipeline_config.build_sde_extra( + self.scheduler, grids, sample_indices, tstep_indices, self.args, + ) + prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( + self.scheduler, + noise_pred_flat, + timesteps_for_sde, + latents_flat, + prev_sample=next_latents_tile.reshape( + tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] + ), + noise_level=noise_level, + extra=sde_extra, + ) + del prev_sample_dummy + else: + 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, + ) - prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( - self.scheduler, - noise_pred_flat, - timesteps_for_sde, - latents_flat, - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ), - noise_level=noise_level, - extra=sde_extra, - ) - del prev_sample_dummy + 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_level=noise_level, + ) # TODO: revamp and gather all loss logics log_prob_new = log_prob_new_flat.reshape(tile_sample_count, tile_tstep_count) @@ -735,17 +800,30 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: if kl_beta > 0: with torch.no_grad(): ref_noise_pred_flat = _compute_noise_pred(disable_adapter=True) - _, _, prev_sample_mean_ref, _ = train_pipeline_config.sde_step( - self.scheduler, - ref_noise_pred_flat, - timesteps_for_sde, - latents_flat, - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ), - noise_level=noise_level, - extra=sde_extra, - ) + # TODO: unify sde_step_with_logprob with rollout and trainer forward paths. + if self._is_ltx: + _, _, prev_sample_mean_ref, _ = train_pipeline_config.sde_step( + self.scheduler, + ref_noise_pred_flat, + timesteps_for_sde, + latents_flat, + prev_sample=next_latents_tile.reshape( + tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] + ), + noise_level=noise_level, + extra=sde_extra, + ) + else: + _, _, 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(), + noise_level=noise_level, + ) kl_loss = ((prev_sample_mean - prev_sample_mean_ref) ** 2).mean( dim=tuple(range(1, prev_sample_mean.ndim)), keepdim=True, @@ -768,15 +846,23 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: torch.mean(torch.abs(log_prob_new - log_prob_old_tile)).detach() ) # To log model output diff, please enable --diffusion-debug-mode - train_metric_utils.log_model_output_debug_metrics( - log_stats, - noise_pred_flat=noise_pred_flat, - grids=grids, - sample_indices=sample_indices, - tstep_indices=tstep_indices, - tile_sample_count=tile_sample_count, - tile_tstep_count=tile_tstep_count, - ) + rollout_mo_window = grids.get("rollout_model_outputs") + if rollout_mo_window is not None: + rollout_mo_tile = rollout_mo_window[sample_indices][:, tstep_indices] + rollout_mo_flat = rollout_mo_tile.reshape( + tile_sample_count * tile_tstep_count, *rollout_mo_tile.shape[2:] + ) + diff = (noise_pred_flat.float() - rollout_mo_flat.float()).abs() + ref_max = rollout_mo_flat.float().abs().max() + 1e-30 + 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()) + if self._is_ltx: + flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) + flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) + log_stats["model_output_cosine_sim"].append( + torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1).mean().detach() + ) return loss @@ -900,19 +986,17 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - model.print_trainable_parameters() return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, train_pipeline_config=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 = getattr(model, "_no_split_modules", None) if not layer_cls_to_wrap: - layer_cls_to_wrap = ( - getattr(train_pipeline_config, "fsdp_wrap_classes", None) if train_pipeline_config else None - ) + layer_cls_to_wrap = fsdp_wrap_classes assert layer_cls_to_wrap and layer_cls_to_wrap[0] is not None, ( - "apply_fsdp2 needs either model._no_split_modules or " - "train_pipeline_config.fsdp_wrap_classes to know which submodules to shard." + "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 index b9b7e70f..06b41f6a 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -1,28 +1,454 @@ -"""LTX-2.3 video diffusion training pipeline config. - -Adapts ltx_core's ``LTXModel`` (non-diffusers; Modality-keyed forward, patchified -``[B, T, D]`` token latents, per-token timesteps, custom ``LTX2Scheduler``) into -miles' FSDP GRPO training loop. -""" +"""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.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 ( + LTXModelConfigurator, + LTXV_MODEL_COMFY_RENAMING_MAP, + ) + + 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 " + f"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 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 + + from miles.backends.fsdp_utils.ltx_sde import normalize_dynamics_type + + dynamics = normalize_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 getattr(args, "ltx_disable_av_cross_attn", False): + env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" + for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): + if os.environ.get(name): + env_vars[name] = os.environ[name] + + +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-disable-av-cross-attn", + action="store_true", + default=False, + help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", + ) + 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.""" - is_diffusers_pipeline = False needs_timestep_scaling = False # Rollout stores σ×1000 in dit_trajectory.timesteps; CPS uses scheduler σ∈[0,1]. sde_timestep_divisor = 1000.0 - supports_cfg = False fsdp_wrap_classes = ["BasicAVTransformerBlock"] @@ -35,10 +461,7 @@ def load_model_and_scheduler(self, args, init_context_factory): from dataclasses import dataclass, field from ltx_core.components.schedulers import LTX2Scheduler - from miles.backends.model_families.ltx import ( - load_ltx_transformer_for_train, - resolve_transformer_checkpoint, - ) + # load_ltx_transformer_for_train / resolve_transformer_checkpoint defined above @dataclass class _LTXSchedulerHolder: @@ -56,8 +479,6 @@ def to(self, device): master_dtype_name = getattr(args, "fsdp_master_dtype", "bf16") master_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[master_dtype_name] - from miles.backends.model_families.ltx import resolve_transformer_checkpoint - checkpoint = resolve_transformer_checkpoint( args.diffusion_model, explicit_path=getattr(args, "sglang_transformer_weights_path", None), @@ -105,6 +526,18 @@ def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> kwargs["audio_context_mask"] = audio_mask return kwargs + def resolve_sigmas_ref( + self, + timesteps_ref: torch.Tensor, + sigmas_snapshot: torch.Tensor | None, + scheduler, + ) -> torch.Tensor: + device = timesteps_ref.device + if sigmas_snapshot is not None: + return sigmas_snapshot.to(device).float() + sigmas_ref = timesteps_ref / float(self.sde_timestep_divisor) + return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) + def build_train_cond_kwargs( self, cond: CondKwargs | None, @@ -150,13 +583,16 @@ def build_sde_extra( tstep_indices: torch.Tensor, args, ) -> dict | None: - extra = super().build_sde_extra(scheduler, grids, sample_indices, tstep_indices, args) - if extra is None: + window = grids.get("sde_step_indices_window") + if window is None: return None - extra["sigmas"] = scheduler.sigmas - extra["dynamics_type"] = getattr(args, "ltx_dynamics_type", "cps") - extra["sigma_min_override"] = getattr(args, "ltx_sigma_min", None) - return extra + 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 = {} @@ -284,7 +720,7 @@ def sde_step( 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_dynamics + from miles.backends.fsdp_utils.ltx_sde import sde_step_with_logprob_dynamics if extra is None or "sigmas" not in extra or "sde_step_indices" not in extra: raise ValueError( diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 47542a81..622559ed 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -14,29 +14,12 @@ from __future__ import annotations import abc -from argparse import Namespace -from dataclasses import dataclass -from typing import Any, Callable import torch from miles.utils.types import CondKwargs, DiTTrajectory -@dataclass -class SdeWindowBatch: - """One sample's tensors after optional SDE-window slicing.""" - - latents: torch.Tensor - next_latents: torch.Tensor - timesteps: torch.Tensor - log_prob_old: torch.Tensor - advantage: torch.Tensor - rollout_model_output: torch.Tensor | None - window_size: int - step_indices: torch.Tensor | None - - -_REGISTRY: dict[str, type["TrainPipelineConfig"]] = {} +_REGISTRY: dict[str, type[TrainPipelineConfig]] = {} def register_train_pipeline_config(*model_name_patterns: str): @@ -64,107 +47,10 @@ def get_train_pipeline_config(model_name: str) -> TrainPipelineConfig: class TrainPipelineConfig(abc.ABC): """Base class. Subclass per model family.""" - is_diffusers_pipeline: bool = True - supports_cfg: bool = True - fsdp_wrap_classes: list[str] | None = None lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] needs_timestep_scaling: bool = True - # When set, ``dit_trajectory.timesteps`` are on an AdaLN scale (e.g. σ×1000) - # but CPS/SDE log_prob expects σ in 0..1. Divide by this for sde_step only. - sde_timestep_divisor: float | None = None optimizer_state_allowed_missing: list[str] = [] - def scale_timesteps_for_sde(self, timesteps: torch.Tensor) -> torch.Tensor: - if self.sde_timestep_divisor is not None: - return timesteps / float(self.sde_timestep_divisor) - return timesteps - - def load_model_and_scheduler( - self, - args: Namespace, - init_context_factory: Callable[[], Any], - ) -> tuple[torch.nn.Module, Any]: - """Load DiT + scheduler. Default: diffusers ``DiffusionPipeline`` (transformer only).""" - from diffusers import DiffusionPipeline - - diffusion_model_id = args.diffusion_model or args.hf_checkpoint - master_dtype_name = getattr(args, "fsdp_master_dtype", "bf16") - master_dtype = { - "fp16": torch.float16, - "bf16": torch.bfloat16, - "fp32": torch.float32, - }[master_dtype_name] - - with init_context_factory(): - pipeline = DiffusionPipeline.from_pretrained( - diffusion_model_id, - 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 apply_sde_step_window( - self, - *, - latents: torch.Tensor, - next_latents: torch.Tensor, - timesteps: torch.Tensor, - log_prob_old: torch.Tensor, - advantage: torch.Tensor, - rollout_model_output: torch.Tensor | None, - sde_step_indices: list[int] | None, - default_window_size: int, - device: torch.device, - ) -> SdeWindowBatch: - """Slice trajectory/objective tensors to the rollout SDE window.""" - step_indices: torch.Tensor | None = None - if sde_step_indices is not None: - step_indices = torch.as_tensor(sde_step_indices, device=device, dtype=torch.long) - latents = latents[step_indices] - next_latents = next_latents[step_indices] - timesteps = timesteps[step_indices] - log_prob_old = log_prob_old[step_indices] - advantage = advantage[: step_indices.numel()] - if rollout_model_output is not None: - n_mo = int(rollout_model_output.shape[0]) - n_win = int(step_indices.numel()) - if n_mo != n_win: - # Full-length debug tensors (legacy): index by global step. - rollout_model_output = rollout_model_output[step_indices] - # else: sglang packs debug outputs in SDE-window order (0..W-1). - window_size = int(step_indices.numel()) - else: - window_size = default_window_size - - return SdeWindowBatch( - latents=latents, - next_latents=next_latents, - timesteps=timesteps, - log_prob_old=log_prob_old, - advantage=advantage, - rollout_model_output=rollout_model_output, - window_size=window_size, - step_indices=step_indices, - ) - - def resolve_tile_sde_step_indices( - self, - grids: dict, - sample_indices: torch.Tensor, - tstep_indices: torch.Tensor, - ) -> torch.Tensor | None: - """Map a training tile to global denoising step indices.""" - window = grids.get("sde_step_indices_window") - if window is None: - return None - return window[sample_indices][:, tstep_indices].reshape(-1).long() - def prepare_trajectory( self, traj: DiTTrajectory, @@ -189,52 +75,6 @@ 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: Namespace, - device: torch.device, - ) -> dict: - """Build per-sample conditioning for the training forward pass.""" - return self.prepare_cond_kwargs(cond, device) - - def resolve_sigmas_ref( - self, - timesteps_ref: torch.Tensor, - sigmas_snapshot: torch.Tensor | None, - scheduler: Any, - ) -> torch.Tensor: - """Build ``[T+1]`` sigma reference for the training scheduler.""" - device = timesteps_ref.device - if sigmas_snapshot is not None: - return sigmas_snapshot.to(device).float() - - sched_config = getattr(scheduler, "config", None) - num_train_timesteps = ( - int(sched_config.num_train_timesteps) if sched_config is not None else 1000 - ) - if not self.needs_timestep_scaling: - sigmas_ref = self.scale_timesteps_for_sde(timesteps_ref) - else: - sigmas_ref = timesteps_ref / float(num_train_timesteps) - return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) - - def build_sde_extra( - self, - scheduler: Any, - grids: dict, - sample_indices: torch.Tensor, - tstep_indices: torch.Tensor, - args: Namespace, - ) -> dict | None: - """Optional per-tile metadata for model-specific SDE log_prob.""" - idx = self.resolve_tile_sde_step_indices(grids, sample_indices, tstep_indices) - if idx is None: - return None - return {"sde_step_indices": idx} - def expand_cond_for_timestep_batch( self, cond_kwargs: dict, @@ -279,55 +119,3 @@ def cfg_combine( @abc.abstractmethod def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: """Preprocess the model before FSDP.""" - pass - - def forward_velocity( - self, - model: torch.nn.Module, - latents_input: torch.Tensor, - timesteps_input: torch.Tensor, - cond: dict, - ) -> torch.Tensor: - return model( - hidden_states=latents_input, - timestep=timesteps_input, - return_dict=False, - **cond, - )[0] - - def forward_velocity_cfg_joint( - self, - model: torch.nn.Module, - latents_input: torch.Tensor, - timesteps_input: torch.Tensor, - joint_cond: dict, - ) -> torch.Tensor: - return 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] - - def sde_step( - self, - scheduler: Any, - 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 - - prev, log_prob, prev_mean, std_dev_t = sde_step_with_logprob( - scheduler, - noise_pred.float(), - timesteps, - sample.float(), - prev_sample=prev_sample.float(), - noise_level=noise_level, - ) - return prev, log_prob, prev_mean, std_dev_t diff --git a/miles/backends/fsdp_utils/ltx_sde.py b/miles/backends/fsdp_utils/ltx_sde.py new file mode 100644 index 00000000..cadc9b1c --- /dev/null +++ b/miles/backends/fsdp_utils/ltx_sde.py @@ -0,0 +1,144 @@ +"""LTX-specific SDE dynamics (schedule-decoupled CPS / Flow-SDE / ODE paths).""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +CANONICAL_DYNAMICS_TYPES = ("sde", "flow_sde", "cps", "ode", "dance_sde") + + +def normalize_dynamics_type(name: str) -> str: + """Map a dynamics-type alias (CLI / legacy casing) to its canonical name.""" + key = str(name).strip().lower().replace("-", "_") + if key not in CANONICAL_DYNAMICS_TYPES: + raise ValueError( + f"Unknown dynamics_type {name!r}; expected one of {CANONICAL_DYNAMICS_TYPES}" + ) + return key + + +def sde_step_with_logprob_dynamics( + model_output: torch.FloatTensor, + sigma: torch.FloatTensor, + sigma_next: torch.FloatTensor, + sample: torch.FloatTensor, + sigmas: torch.FloatTensor, + prev_sample: Optional[torch.FloatTensor] = None, + generator: Optional[torch.Generator] = None, + deterministic: bool = False, + sigma_min_override: Optional[float] = None, + noise_level: float = 0.8, + dynamics_type: str = "flow_sde", +): + """Schedule-decoupled SDE step with log-prob for LTX-2.3 and similar models.""" + dynamics_type = normalize_dynamics_type(dynamics_type) + model_output = model_output.float() + sample = sample.float() + if prev_sample is not None: + prev_sample = prev_sample.float() + + ndim = sample.ndim + sigma_view = sigma.float() + sigma_next_view = sigma_next.float() + while sigma_view.ndim < ndim: + sigma_view = sigma_view.unsqueeze(-1) + while sigma_next_view.ndim < ndim: + sigma_next_view = sigma_next_view.unsqueeze(-1) + + dt = sigma_next_view - sigma_view + + sigma_max = sigmas[0].float().item() + if sigma_min_override is not None: + sigma_min = sigma_min_override + else: + sigma_min = max(sigmas[-2].float().item(), 1e-4) if len(sigmas) > 1 else 1e-4 + + if dynamics_type == "ode": + prev_sample_mean = sample + dt * model_output + std_dev_t = torch.zeros_like(sigma_view) + if prev_sample is None: + prev_sample = prev_sample_mean + log_prob = torch.zeros(sample.shape[0], dtype=sample.dtype, device=sample.device) + + elif dynamics_type == "flow_sde": + std_dev_t = (sigma_min + (sigma_max - sigma_min) * sigma_view) * noise_level + sigma_safe = torch.clamp(sigma_view, min=1e-8) + + drift_sample = 1.0 + std_dev_t**2 / (2.0 * sigma_safe) * dt + drift_model = (1.0 + std_dev_t**2 * (1.0 - sigma_view) / (2.0 * sigma_safe)) * dt + prev_sample_mean = sample * drift_sample + model_output * drift_model + + noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) + + if prev_sample is None: + if deterministic: + prev_sample = sample + dt * model_output + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + noise_scale * variance_noise + + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) + - torch.log(noise_scale + 1e-12) + - 0.5 * math.log(2.0 * math.pi) + ) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + elif dynamics_type == "cps": + angle = torch.tensor(noise_level, dtype=sigma_next_view.dtype, device=sigma_next_view.device) * torch.pi / 2 + std_dev_t = sigma_next_view * torch.sin(angle) + + x0 = sample - sigma_view * model_output + x1 = sample + model_output * (1.0 - sigma_view) + sqrt_term = torch.sqrt(torch.clamp(sigma_next_view**2 - std_dev_t**2, min=1e-12)) + prev_sample_mean = x0 * (1.0 - sigma_next_view) + x1 * sqrt_term + + if prev_sample is None: + if deterministic: + prev_sample = prev_sample_mean + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + std_dev_t * variance_noise + + log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + elif dynamics_type == "dance_sde": + sigma_safe = torch.clamp(sigma_view, min=1e-8) + x0_pred = sample - sigma_safe * model_output + std_dev_t = torch.as_tensor(noise_level, dtype=sample.dtype, device=sample.device) + log_term = 0.5 * noise_level**2 * (sample - x0_pred * (1.0 - sigma_view)) / (sigma_safe**2) + prev_sample_mean = sample + (model_output + log_term) * dt + noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) + + if prev_sample is None: + if deterministic: + prev_sample = sample + dt * model_output + else: + variance_noise = torch.randn( + sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, + ) + prev_sample = prev_sample_mean + noise_scale * variance_noise + + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) + - torch.log(noise_scale + 1e-12) + - 0.5 * math.log(2.0 * math.pi) + ) + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + else: + raise ValueError( + f"dynamics_type {dynamics_type!r} is not supported by the " + "schedule-decoupled path; use flow_sde / cps / ode / dance_sde." + ) + + dt_sqrt = torch.sqrt(torch.clamp(-dt, min=1e-12)) + return prev_sample, log_prob, prev_sample_mean, std_dev_t, dt_sqrt From 373591f7d97d13ee9a3230b82e907b1c71e40cf5 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 22 Jun 2026 11:55:34 +0000 Subject: [PATCH 20/31] feat(ltx): sglang rollout patches and engine env bridging Keep monkey patches with upstream TODO notes; wire LTX imports through fsdp configs/ltx.py instead of model_families/. --- miles/backends/sglang_diffusion_utils/configs/ltx.py | 4 ++-- .../monkey_patches/patch_ltx2_ltxcore_parity.py | 6 +++++- .../monkey_patches/patch_ltx2_rollout_cond_kwargs.py | 6 +++++- .../sglang_diffusion_utils/sglang_diffusion_engine.py | 2 +- miles/rollout/sglang_diffusion_rollout.py | 2 +- 5 files changed, 14 insertions(+), 6 deletions(-) diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py index 1737fad2..7f6ee79d 100644 --- a/miles/backends/sglang_diffusion_utils/configs/ltx.py +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -1,8 +1,8 @@ -"""LTX-2 sglang-d rollout engine config (re-exports model family helpers).""" +"""LTX-2 sglang-d rollout engine config (re-exports train-side model family helpers).""" from __future__ import annotations -from miles.backends.model_families.ltx import ( +from miles.backends.fsdp_utils.configs.ltx import ( LTX_DEFAULT_HF_MODEL, LTX_DEFAULT_MODEL_ID, ensure_materialized_model, diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py index 00da367d..fa02747a 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py @@ -1,4 +1,8 @@ -"""LTX-2 DiT parity patches: align sglang ltx_2.py with miles/ltx_core.""" +"""LTX-2 DiT parity patches: align sglang ltx_2.py with miles/ltx_core. + +TODO(upstream): remove once sgl-d LTX rollout matches ltx_core AdaLN / temb / +velocity-to-x0 paths natively (train/rollout alignment checks pass without patch). +""" from __future__ import annotations 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 index 711ffa68..1e32a580 100644 --- 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 @@ -1,4 +1,8 @@ -"""Ensure LTX rollout denoising_env carries text/audio embeds for miles train replay.""" +"""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 diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 1e0acccf..502abd36 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -19,7 +19,7 @@ def build_rollout_engine_env_vars(args) -> dict[str, str]: """Env vars forwarded to Ray-spawned sglang-diffusion rollout engine workers.""" - from miles.backends.model_families.ltx import patch_rollout_engine_env_vars + 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} | { diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 18c481df..51af6012 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -64,7 +64,7 @@ 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.model_families.ltx import is_ltx_model, patch_rollout_sampling_params + 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) From 4eda10fad63e1e53e4387637ce0acfd2fb6b59e7 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 22 Jun 2026 11:55:34 +0000 Subject: [PATCH 21/31] chore(ltx): rollout train data path and remove model_families module Consolidate model-family helpers into fsdp configs/ltx.py. --- miles/backends/model_families/__init__.py | 1 - miles/backends/model_families/ltx.py | 433 ---------------------- miles/ray/rollout.py | 2 +- 3 files changed, 1 insertion(+), 435 deletions(-) delete mode 100644 miles/backends/model_families/__init__.py delete mode 100644 miles/backends/model_families/ltx.py diff --git a/miles/backends/model_families/__init__.py b/miles/backends/model_families/__init__.py deleted file mode 100644 index f03bee62..00000000 --- a/miles/backends/model_families/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Per-model-family CLI, validation, and rollout hooks.""" diff --git a/miles/backends/model_families/ltx.py b/miles/backends/model_families/ltx.py deleted file mode 100644 index cca6b42e..00000000 --- a/miles/backends/model_families/ltx.py +++ /dev/null @@ -1,433 +0,0 @@ -"""LTX-2 model family: checkpoint resolution, rollout engine, and sampling hooks.""" - -from __future__ import annotations - -import logging -import os -from argparse import ArgumentParser, Namespace -from pathlib import Path -from typing import Any - -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 ( - LTXModelConfigurator, - LTXV_MODEL_COMFY_RENAMING_MAP, - ) - - 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 " - f"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 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 - - from miles.utils.sde_log_prob import normalize_dynamics_type - - dynamics = normalize_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 getattr(args, "ltx_disable_av_cross_attn", False): - env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" - for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): - if os.environ.get(name): - env_vars[name] = os.environ[name] - - -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-disable-av-cross-attn", - action="store_true", - default=False, - help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", - ) - 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." - ) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index ac88e83d..b69e8c77 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -405,7 +405,7 @@ def _log_images( own namespace at least groups them in one UI section. """ import wandb - from miles.rollout.rm_hub.video_pickscore import first_frame_for_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(): From 7a1df52d1b9b142f5af781866f4e5fda4c4d94b3 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 22 Jun 2026 11:55:34 +0000 Subject: [PATCH 22/31] feat(ltx): inline video PickScore helpers and trim shared utils Merge video_pickscore into pickscore.py (Wan-style). Drop unused debug metrics helper from train_metric_utils. --- miles/rollout/rm_hub/pickscore.py | 78 +++++++++++-- miles/utils/arguments.py | 4 +- miles/utils/sde_log_prob.py | 178 +++--------------------------- miles/utils/train_metric_utils.py | 33 ------ 4 files changed, 88 insertions(+), 205 deletions(-) diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index 121bb84e..c8b4245e 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -9,19 +9,83 @@ import torch from PIL import Image -from miles.rollout.rm_hub.video_pickscore import ( - fchw_frame_to_hwc_uint8, - fchw_to_pil_frames, - generated_output_to_fchw, - is_video_generated_output, - sample_frame_indices, -) from miles.utils.misc import SingletonMeta from miles.utils.types import Sample 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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 406f8925..8b1e7adf 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -454,7 +454,7 @@ def add_rollout_arguments(parser): "(HF hub id with ``ltx`` → ltx, else sd3)." ), ) - from miles.backends.model_families.ltx import register_args as register_ltx_args + from miles.backends.fsdp_utils.configs.ltx import register_args as register_ltx_args register_ltx_args(parser) parser.add_argument( @@ -1315,7 +1315,7 @@ def miles_validate_args(args): else: args.diffusion_model_type = "sd3" if args.diffusion_model_type == "ltx": - from miles.backends.model_families.ltx import validate_args as validate_ltx_args + from miles.backends.fsdp_utils.configs.ltx import validate_args as validate_ltx_args validate_ltx_args(args) diff --git a/miles/utils/sde_log_prob.py b/miles/utils/sde_log_prob.py index b90323b7..60495dc1 100644 --- a/miles/utils/sde_log_prob.py +++ b/miles/utils/sde_log_prob.py @@ -1,43 +1,12 @@ """SDE step with log probability for flow matching schedulers. -This module exposes: - -- :func:`sde_step_with_logprob` — the original SD3 / flow-matching scheduler - contract used by miles' SD3 path. - -- :func:`sde_step_with_logprob_dynamics` — generic, schedule-decoupled version - used by the LTX-2.3 path which runs on patchified token latents and a - custom :class:`LTX2Scheduler`. +Adapted from flow_grpo/diffusers_patch/sd3_sde_with_logprob.py. """ import math -from typing import Optional, Union import torch -# Canonical dynamics names. These match sglang-d ``rollout_sde_type`` so miles -# can pass them straight through to the rollout engine with no translation -# table — keeping train (this module) and rollout (sglang-d flow_sde_sampling) -# on a single shared vocabulary. -CANONICAL_DYNAMICS_TYPES = ("sde", "flow_sde", "cps", "ode", "dance_sde") - - -def normalize_dynamics_type(name: str) -> str: - """Map a dynamics-type alias (CLI / legacy casing) to its canonical name. - - Accepts any case and ``-``/``_`` spelling, e.g. ``"Flow-SDE"``, - ``"flow_sde"`` -> ``"flow_sde"``; ``"CPS"`` -> ``"cps"``; - ``"Dance-SDE"`` -> ``"dance_sde"``. Raises on unknown values rather than - silently falling back, so a typo can never mismatch train vs rollout. - """ - key = str(name).strip().lower().replace("-", "_") - if key not in CANONICAL_DYNAMICS_TYPES: - raise ValueError( - f"Unknown dynamics_type {name!r}; expected one of " - f"{CANONICAL_DYNAMICS_TYPES}" - ) - return key - def sde_step_with_logprob( scheduler, @@ -47,7 +16,19 @@ def sde_step_with_logprob( prev_sample: torch.FloatTensor, noise_level: float = 0.7, ): - """Compute the log probability of `prev_sample` under one reverse-SDE step.""" + """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). + prev_sample: Recorded next-step latent to score under the SDE. + noise_level: SDE noise scaling factor (eta). + + Returns: + (prev_sample, log_prob, prev_sample_mean, std_dev_t) + """ model_output = model_output.float() sample = sample.float() prev_sample = prev_sample.float() @@ -72,136 +53,7 @@ def sde_step_with_logprob( - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))) ) + # mean along all but batch dimension log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) return prev_sample, log_prob, prev_sample_mean, std_dev_t - - -def sde_step_with_logprob_dynamics( - model_output: torch.FloatTensor, - sigma: torch.FloatTensor, - sigma_next: torch.FloatTensor, - sample: torch.FloatTensor, - sigmas: torch.FloatTensor, - prev_sample: Optional[torch.FloatTensor] = None, - generator: Optional[torch.Generator] = None, - deterministic: bool = False, - sigma_min_override: Optional[float] = None, - noise_level: float = 0.8, - dynamics_type: str = "flow_sde", -): - """Schedule-decoupled SDE step with log-prob for LTX-2.3 and similar models. - - ``dynamics_type`` accepts any alias (see :func:`normalize_dynamics_type`). - """ - dynamics_type = normalize_dynamics_type(dynamics_type) - model_output = model_output.float() - sample = sample.float() - if prev_sample is not None: - prev_sample = prev_sample.float() - - ndim = sample.ndim - sigma_view = sigma.float() - sigma_next_view = sigma_next.float() - while sigma_view.ndim < ndim: - sigma_view = sigma_view.unsqueeze(-1) - while sigma_next_view.ndim < ndim: - sigma_next_view = sigma_next_view.unsqueeze(-1) - - dt = sigma_next_view - sigma_view - - sigma_max = sigmas[0].float().item() - if sigma_min_override is not None: - sigma_min = sigma_min_override - else: - sigma_min = max(sigmas[-2].float().item(), 1e-4) if len(sigmas) > 1 else 1e-4 - - if dynamics_type == "ode": - prev_sample_mean = sample + dt * model_output - std_dev_t = torch.zeros_like(sigma_view) - if prev_sample is None: - prev_sample = prev_sample_mean - log_prob = torch.zeros(sample.shape[0], dtype=sample.dtype, device=sample.device) - - elif dynamics_type == "flow_sde": - std_dev_t = (sigma_min + (sigma_max - sigma_min) * sigma_view) * noise_level - sigma_safe = torch.clamp(sigma_view, min=1e-8) - - drift_sample = 1.0 + std_dev_t**2 / (2.0 * sigma_safe) * dt - drift_model = (1.0 + std_dev_t**2 * (1.0 - sigma_view) / (2.0 * sigma_safe)) * dt - prev_sample_mean = sample * drift_sample + model_output * drift_model - - noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) - - if prev_sample is None: - if deterministic: - prev_sample = sample + dt * model_output - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + noise_scale * variance_noise - - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) - - torch.log(noise_scale + 1e-12) - - 0.5 * math.log(2.0 * math.pi) - ) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - elif dynamics_type == "cps": - angle = torch.tensor(noise_level, dtype=sigma_next_view.dtype, device=sigma_next_view.device) * torch.pi / 2 - std_dev_t = sigma_next_view * torch.sin(angle) - - x0 = sample - sigma_view * model_output - x1 = sample + model_output * (1.0 - sigma_view) - sqrt_term = torch.sqrt(torch.clamp(sigma_next_view**2 - std_dev_t**2, min=1e-12)) - prev_sample_mean = x0 * (1.0 - sigma_next_view) + x1 * sqrt_term - - if prev_sample is None: - if deterministic: - prev_sample = prev_sample_mean - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + std_dev_t * variance_noise - - log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - elif dynamics_type == "dance_sde": - sigma_safe = torch.clamp(sigma_view, min=1e-8) - x0_pred = sample - sigma_safe * model_output - std_dev_t = torch.as_tensor(noise_level, dtype=sample.dtype, device=sample.device) - log_term = 0.5 * noise_level**2 * (sample - x0_pred * (1.0 - sigma_view)) / (sigma_safe**2) - prev_sample_mean = sample + (model_output + log_term) * dt - noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) - - if prev_sample is None: - if deterministic: - prev_sample = sample + dt * model_output - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + noise_scale * variance_noise - - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) - - torch.log(noise_scale + 1e-12) - - 0.5 * math.log(2.0 * math.pi) - ) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - else: - # ``sde`` (SD3-style, scheduler-based) is handled by - # ``sde_step_with_logprob`` above, not this schedule-decoupled path. - raise ValueError( - f"dynamics_type {dynamics_type!r} is not supported by the " - "schedule-decoupled path; use flow_sde / cps / ode / dance_sde " - "(or sde via sde_step_with_logprob)." - ) - - dt_sqrt = torch.sqrt(torch.clamp(-dt, min=1e-12)) - return prev_sample, log_prob, prev_sample_mean, std_dev_t, dt_sqrt diff --git a/miles/utils/train_metric_utils.py b/miles/utils/train_metric_utils.py index b39ab3e6..184005a8 100644 --- a/miles/utils/train_metric_utils.py +++ b/miles/utils/train_metric_utils.py @@ -2,8 +2,6 @@ from argparse import Namespace from copy import deepcopy -import torch - from miles.utils import tracking_utils from miles.utils.metric_utils import compute_rollout_step from miles.utils.timer import Timer @@ -11,37 +9,6 @@ logger = logging.getLogger(__name__) -def log_model_output_debug_metrics( - log_stats: dict[str, list[torch.Tensor]], - *, - noise_pred_flat: torch.Tensor, - grids: dict, - sample_indices: torch.Tensor, - tstep_indices: torch.Tensor, - tile_sample_count: int, - tile_tstep_count: int, -) -> None: - """Compare train-side noise_pred against rollout debug model outputs.""" - rollout_mo_window = grids.get("rollout_model_outputs") - if rollout_mo_window is None: - return - - rollout_mo_tile = rollout_mo_window[sample_indices][:, tstep_indices] - rollout_mo_flat = rollout_mo_tile.reshape( - tile_sample_count * tile_tstep_count, *rollout_mo_tile.shape[2:] - ) - diff = (noise_pred_flat.float() - rollout_mo_flat.float()).abs() - ref_max = rollout_mo_flat.float().abs().max() + 1e-30 - 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()) - flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) - flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) - log_stats["model_output_cosine_sim"].append( - torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1).mean().detach() - ) - - def log_perf_data_raw(rollout_id: int, args: Namespace, is_primary_rank: bool) -> None: timer_instance = Timer() log_dict_raw = deepcopy(timer_instance.log_dict()) From 1673f8584915f8e9d92045a4cd997c735f1aa476 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 22 Jun 2026 11:55:57 +0000 Subject: [PATCH 23/31] chore(ltx): remove video_pickscore after inlining into pickscore --- miles/rollout/rm_hub/video_pickscore.py | 80 ------------------------- 1 file changed, 80 deletions(-) delete mode 100644 miles/rollout/rm_hub/video_pickscore.py diff --git a/miles/rollout/rm_hub/video_pickscore.py b/miles/rollout/rm_hub/video_pickscore.py deleted file mode 100644 index c4d2c3c4..00000000 --- a/miles/rollout/rm_hub/video_pickscore.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Video PickScore helpers (LTX / sglang FHWC and trainer [C,F,H,W] layouts).""" - -from __future__ import annotations - -from collections.abc import Sequence - -import numpy as np -import torch -from PIL import Image - - -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 From 34e2b8b161e62917227e80bf975a97ddcf4a5443 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 23 Jun 2026 02:13:32 +0000 Subject: [PATCH 24/31] chore(ltx): slim run recipe and revert placement_group docstring churn Align the LTX sglang script with other diffusion GRPO recipes by removing dev-cluster helpers and restoring colocate pickscore GPU defaults; drop an unrelated placement_group docstring-only edit from the PR. --- miles/ray/placement_group.py | 5 +- scripts/run-diffusion-grpo-ltx23-sglang.sh | 246 +++++---------------- 2 files changed, 54 insertions(+), 197 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 3fff7cfe..0200bb02 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -82,9 +82,8 @@ def create_placement_groups(args): """Create placement groups for actor and rollout engines. Two topologies: - - Colocate (or ``--debug-train-only`` / ``--debug-rollout-only``): one combined - placement group; both roles see the same bundle list (train-only allocates - no rollout GPU bundles). + - Colocate (or --debug-{train,rollout}-only): one combined placement + group; both roles see the same bundle list. - Disaggregate (the else branch): two separate placement groups so train and rollout each own a disjoint GPU pool — avoids bundle overlap / scheduling deadlock when running side-by-side. diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index a331a39f..3b8a2913 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -1,164 +1,50 @@ #!/usr/bin/env bash -# LTX-2.3 sglang-rollout GRPO — dev checkpoint, 512x768x57f, 24 steps, CPS. +# LTX-2.3 video PickScore GRPO: sglang rollout + FSDP train (colocate). # -# Mirrors the legacy trainer-rollout reward run -# (/sgl-workspace/miles/scripts/run-diffusion-grpo-ltx23-trainer-rollout.sh): -# CPS dynamics, 3 SDE steps from candidates 0–9, clip-range 1e-4. -# Rollout goes through sglang with weight sync; train/rollout forward alignment -# fixes stay on (ltxcore parity + AV-off + gs=1 rollout alignment). +# Default: 1-GPU colocate (train FSDP + sglang rollout). Override NUM_GPUS for +# multi-GPU colocate. CPS dynamics, 3 SDE steps from candidates 0–9, clip 1e-4. # -# GPU layout: single physical GPU colocate (train FSDP world_size=1 and sglang -# rollout time-share one GPU via offload). Set NUM_GPUS>1 for multi-GPU -# colocate if 512x768x57f OOMs on one card. -# -# Usage: -# CUDA_VISIBLE_DEVICES=1 USE_LORA=1 NUM_ROLLOUT=200 \ -# LTX_DISABLE_AV_CROSS_ATTN=1 \ -# nohup bash scripts/run-diffusion-grpo-ltx23-sglang.sh \ -# > logs/ltx23_dev_cps_$(date +%Y%m%d_%H%M%S).log 2>&1 & -# -# Key overridable env: -# LTX_HF_MODEL — default Lightricks/LTX-2.3 (train + rollout via overlay) -# LTX_DEV_SAFETENSORS — optional dev .safetensors override for train + rollout DiT -# HEIGHT WIDTH FRAMES — 512 768 57 -# NUM_STEPS — 24 -# LTX_NUM_SDE_STEPS — 3 -# LTX_SDE_STEP_CANDIDATES — 0,1,2,3,4,5,6,7,8,9 -# CLIP_RANGE — 1e-4 -# ROLLOUT_BATCH_SIZE — unique prompts per rollout (default: 8) -# N_SAMPLES_PER_PROMPT — GRPO group size (default: 8) -# NUM_STEPS_PER_ROLLOUT — optimizer steps per rollout (default: 2 → gbs=32) -# NUM_ROLLOUT — 200 -# SAVE_INTERVAL — 50 - -MILES_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -echo "[kill] hunting for stale miles processes under cwd=${MILES_ROOT}" -for pid in $(ls /proc 2>/dev/null | grep -E '^[0-9]+$'); do - # timeout guards against readlink hanging on a process whose cwd points at a - # stale/unresponsive mount — otherwise this loop can wedge the whole shell. - link=$(timeout 2 readlink "/proc/${pid}/cwd" 2>/dev/null) || continue - exe=$(timeout 2 readlink "/proc/${pid}/exe" 2>/dev/null) || continue - case "${link}" in - "${MILES_ROOT}"|"${MILES_ROOT}"/*) - case "${exe}" in - */python*|*/ray*) - echo "[kill] ${pid} (${exe}) cwd=${link}" - kill -9 "${pid}" 2>/dev/null || true - ;; - esac - ;; - esac -done -sleep 3 - -ps -eo ppid,state,comm --no-headers \ - | awk '$2=="Z" && $1!=1 && $3~/ray|python|sglang/ {print $1}' \ - | sort -u | xargs -r kill -9 2>/dev/null || true -sleep 2 +# Layout mirrors other scripts/run-diffusion-grpo-*.sh recipes: +# train+rollout share the first NUM_GPUS in CUDA_VISIBLE_DEVICES; +# optional pickscore worker uses additional GPUs when configured. set -euo pipefail - -ROOT_DIR="${MILES_ROOT}" -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}" +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" -SGLANG_PYTHON="${SGLANG_PYTHON:-/sgl-workspace/master_sglang/sglang/python}" -export PYTHONPATH="${SGLANG_PYTHON}${PYTHONPATH:+:${PYTHONPATH}}" +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}" -# All heavy I/O lives on /data — workspace overlay is small and often full. -MILES_DATA_DISK_ROOT="${MILES_DATA_DISK_ROOT:-/data/wenhao/miles_diffusion}" -RAY_BIG_TMP="${RAY_BIG_TMP:-/data/wenhao/miles_ray_tmp}" -TMP_BIG="${TMP_BIG:-/data/wenhao/tmp}" -SGL_DIFF_CACHE="${SGLANG_DIFFUSION_CACHE_ROOT:-/data/wenhao/sgl_diffusion_cache}" -HF_HOME="${HF_HOME:-/data/wenhao/hf_home}" -LOG_DIR="${LOG_DIR:-${MILES_DATA_DISK_ROOT}/logs}" -WANDB_DIR="${WANDB_DIR:-${MILES_DATA_DISK_ROOT}/wandb}" -CKPT_ROOT="${CKPT_ROOT:-${MILES_DATA_DISK_ROOT}/ckpt}" -mkdir -p "${MILES_DATA_DISK_ROOT}" "${RAY_BIG_TMP}" "${TMP_BIG}" "${SGL_DIFF_CACHE}" \ - "${HF_HOME}" "${LOG_DIR}" "${WANDB_DIR}" "${CKPT_ROOT}" -export RAY_TMPDIR="${RAY_BIG_TMP}" -export TMPDIR="${TMP_BIG}" -export SGLANG_DIFFUSION_CACHE_ROOT="${SGL_DIFF_CACHE}" -export HF_HOME -export WANDB_DIR -export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" -mkdir -p "${HUGGINGFACE_HUB_CACHE}" -export MILES_APPLY_LTX2_LTXCORE_PARITY="${MILES_APPLY_LTX2_LTXCORE_PARITY:-1}" -export RAY_object_spilling_config="$(python -c "import json,os; print(json.dumps({'type':'filesystem','params':{'directory_path':[os.environ['RAY_TMPDIR']]}}))")" -ray stop --force 2>/dev/null || true -sleep 2 +PYTHON_BIN="${PYTHON_BIN:-python}" -# ── model: HF hub id by default; optional dev safetensors override ───────── -LTX_HF_MODEL="${LTX_HF_MODEL:-Lightricks/LTX-2.3}" -LTX_DEV_SAFETENSORS="${LTX_DEV_SAFETENSORS:-}" -if [[ -n "${LTX_DEV_SAFETENSORS}" ]]; then - DIFFUSION_MODEL="${LTX_DEV_SAFETENSORS}" -else - DIFFUSION_MODEL="${LTX_HF_MODEL}" +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 -MILES_DATA_ROOT="${MILES_DATA_ROOT:-/sgl-workspace/miles}" -PROMPT_DATA="${PROMPT_DATA:-${MILES_DATA_ROOT}/data/vidgen/train.jsonl}" -NUM_ROLLOUT="${NUM_ROLLOUT:-200}" +DIFFUSION_MODEL="${DIFFUSION_MODEL:-Lightricks/LTX-2.3}" +NUM_GPUS="${NUM_GPUS:-1}" 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}" -SAMPLES_PER_ROLLOUT=$((ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT)) -DERIVED_GLOBAL_BATCH_SIZE=$((SAMPLES_PER_ROLLOUT / NUM_STEPS_PER_ROLLOUT)) -MICRO_BATCH_SIZE_SAMPLE="${MICRO_BATCH_SIZE_SAMPLE:-1}" -MICRO_BATCH_SIZE_TSTEP="${MICRO_BATCH_SIZE_TSTEP:-1}" - -# ── borrowed-from-legacy generation config ─────────────────────────────── -HEIGHT="${HEIGHT:-512}" -WIDTH="${WIDTH:-768}" -FRAMES="${FRAMES:-57}" -NUM_STEPS="${NUM_STEPS:-24}" -# ── trainer-rollout SDE config (CPS + candidate sampling) ──────────────── -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}" -CLIP_RANGE="${CLIP_RANGE:-1e-4}" - -NUM_GPUS="${NUM_GPUS:-1}" -# Multi-GPU colocate: one sglang engine PER GPU (each card runs both the sglang -# rollout engine and an FSDP trainer shard). per-engine=1 => num_engines=NUM_GPUS. -# Set ROLLOUT_NUM_GPUS_PER_ENGINE=NUM_GPUS instead for a single TP-sharded engine. -ROLLOUT_NUM_GPUS_PER_ENGINE="${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" -# Periodic checkpoint (LoRA adapter) so the run is resumable via LOAD_CKPT. -# (The earlier run had no --save-interval, so nothing was ever saved.) +NUM_ROLLOUT="${NUM_ROLLOUT:-200}" SAVE_INTERVAL="${SAVE_INTERVAL:-50}" -if [[ ! -f "${PROMPT_DATA}" ]]; then - python "${MILES_DATA_ROOT}/tools/prepare_vidgen_jsonl.py" -fi - -RUN_NAME="ltx23_dev_cps_${NUM_ROLLOUT}step_$(date +%Y%m%d_%H%M%S)" -SAVE_DIR="${CKPT_ROOT}/${RUN_NAME}" -mkdir -p "${SAVE_DIR}" - -echo "[run] dev+cps CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} NUM_GPUS=${NUM_GPUS}" -echo "[run] model=${DIFFUSION_MODEL}" -echo "[run] log=${LOG_DIR}" -echo "[run] wandb=${WANDB_DIR}" -echo "[run] save=${SAVE_DIR}" -echo "[run] ${HEIGHT}x${WIDTH}x${FRAMES}f steps=${NUM_STEPS} sde_steps=${LTX_NUM_SDE_STEPS} candidates=${LTX_SDE_STEP_CANDIDATES} clip=${CLIP_RANGE}" -echo "[run] batch: rollout=${ROLLOUT_BATCH_SIZE} n_samples=${N_SAMPLES_PER_PROMPT} samples/rollout=${SAMPLES_PER_ROLLOUT} optim_steps/rollout=${NUM_STEPS_PER_ROLLOUT} gbs=${DERIVED_GLOBAL_BATCH_SIZE} save_interval=${SAVE_INTERVAL}" - -DEBUG_ARGS=() -if [[ "${MILES_DIFFUSION_DEBUG:-0}" == "1" ]]; then - DEBUG_ARGS+=(--diffusion-debug-mode) -fi - -DUMP_ARGS=() -if [[ -n "${LTX_FORWARD_DUMP_ROOT:-}" ]]; then - mkdir -p "${LTX_FORWARD_DUMP_ROOT}" - DUMP_ARGS+=(--dump-details "${LTX_FORWARD_DUMP_ROOT}") -fi - -LTX_AV_ARGS=() -if [[ "${LTX_DISABLE_AV_CROSS_ATTN:-0}" == "1" ]]; then - LTX_AV_ARGS+=(--ltx-disable-av-cross-attn) - export MILES_LTX_DISABLE_AV_CROSS=1 +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=() @@ -171,78 +57,54 @@ if [[ "${USE_LORA:-1}" == "1" ]]; then ) fi -SKIP_OPT_ARGS=() -if [[ "${SKIP_OPTIMIZER:-0}" == "1" ]]; then - SKIP_OPT_ARGS+=(--debug-skip-optimizer-step) -fi - -# Resume: point LOAD_CKPT at a previously saved --save dir (LoRA adapter). -LOAD_ARGS=() -if [[ -n "${LOAD_CKPT:-}" ]]; then - LOAD_ARGS+=(--load "${LOAD_CKPT}") -fi - -# WandB: enabled when WANDB_API_KEY is set. Mirrors the legacy reward run so the -# reward curve is directly comparable. -WANDB_ARGS=() -if [[ -n "${WANDB_API_KEY:-}" ]]; then - WANDB_ARGS+=( - --use-wandb - --wandb-dir "${WANDB_DIR}" - --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:-5}" - --disable-wandb-random-suffix - ) +LTX_AV_ARGS=() +if [[ "${LTX_DISABLE_AV_CROSS_ATTN:-0}" == "1" ]]; then + LTX_AV_ARGS+=(--ltx-disable-av-cross-attn) fi -python -u "${ROOT_DIR}/train_diffusion.py" \ +"${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 "${PROMPT_DATA}" \ + --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}" \ - --micro-batch-size-tstep "${MICRO_BATCH_SIZE_TSTEP}" \ + --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}" \ + --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" \ --use-miles-router \ - --rollout-health-check-interval 120 \ - --miles-router-health-check-failure-threshold 30 \ - --sglang-server-concurrency 1 \ + --sglang-server-concurrency "${SGLANG_SERVER_CONCURRENCY:-1}" \ --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ "${LORA_ARGS[@]}" \ "${LTX_AV_ARGS[@]}" \ --lr 2e-4 \ --adam-beta2 0.999 \ --weight-decay 1e-4 \ - --diffusion-clip-range "${CLIP_RANGE}" \ + --diffusion-clip-range "${CLIP_RANGE:-1e-4}" \ --diffusion-kl-beta 0.0 \ - --diffusion-num-steps "${NUM_STEPS}" \ + --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}" \ - --ltx-sde-step-candidates "${LTX_SDE_STEP_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}" \ - --diffusion-width "${WIDTH}" \ - --ltx-frames "${FRAMES}" \ - --ltx-fps 24 \ + --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 \ @@ -252,17 +114,13 @@ python -u "${ROOT_DIR}/train_diffusion.py" \ --rm-type pickscore \ --diffusion-reward "pickscore:1.0" \ --reward-key avg \ - --pickscore-processor-path "${PICKSCORE_PROCESSOR:-/data/wenhao/hf_home/pickscore}" \ - --pickscore-model-path "${PICKSCORE_MODEL:-/data/wenhao/hf_home/pickscore}" \ - --pickscore-num-frames 3 \ + --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 \ - --pickscore-num-gpus-per-worker 0 \ --update-weight-buffer-size 2147483648 \ --save "${SAVE_DIR}" \ --save-interval "${SAVE_INTERVAL}" \ - "${LOAD_ARGS[@]}" \ - "${DEBUG_ARGS[@]}" \ - "${DUMP_ARGS[@]}" \ - "${SKIP_OPT_ARGS[@]}" \ "${WANDB_ARGS[@]}" \ "$@" From fcdf664131dd64f01c65afee2ac5802b9f509e62 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 23 Jun 2026 03:37:17 +0000 Subject: [PATCH 25/31] refactor(ltx): unify CPS SDE into sde_log_prob and drop ltx_sde.py LTX train CPS now calls the shared sde_step_with_logprob with explicit sigma/sigma_prev, matching rollout math without a separate dynamics module. --- miles/backends/fsdp_utils/configs/ltx.py | 42 ++++--- miles/backends/fsdp_utils/ltx_sde.py | 144 ----------------------- miles/utils/sde_log_prob.py | 78 ++++++++---- 3 files changed, 81 insertions(+), 183 deletions(-) delete mode 100644 miles/backends/fsdp_utils/ltx_sde.py diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 06b41f6a..7231dd0c 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -316,6 +316,14 @@ def server_kwargs_extras(args) -> dict: 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, @@ -333,9 +341,7 @@ def patch_rollout_sampling_params( if evaluation: return - from miles.backends.fsdp_utils.ltx_sde import normalize_dynamics_type - - dynamics = normalize_dynamics_type(getattr(args, "ltx_dynamics_type", "cps")) + 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." @@ -720,7 +726,7 @@ def sde_step( noise_level: float, extra: dict | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - from miles.backends.fsdp_utils.ltx_sde import sde_step_with_logprob_dynamics + 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( @@ -731,21 +737,23 @@ def sde_step( sigma_view = timesteps.float() sigma_next = sigmas[torch.clamp(step_indices + 1, max=len(sigmas) - 1)] - dynamics_type = extra.get("dynamics_type", "cps") - sigma_min_override = extra.get("sigma_min_override", None) - if sigma_min_override == 0.0: - sigma_min_override = None + 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, _dt_sqrt = sde_step_with_logprob_dynamics( - model_output=noise_pred.float(), - sigma=sigma_view, - sigma_next=sigma_next, - sample=sample.float(), - sigmas=sigmas, - prev_sample=prev_sample.float(), - sigma_min_override=sigma_min_override, + 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, - dynamics_type=dynamics_type, + 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))) diff --git a/miles/backends/fsdp_utils/ltx_sde.py b/miles/backends/fsdp_utils/ltx_sde.py deleted file mode 100644 index cadc9b1c..00000000 --- a/miles/backends/fsdp_utils/ltx_sde.py +++ /dev/null @@ -1,144 +0,0 @@ -"""LTX-specific SDE dynamics (schedule-decoupled CPS / Flow-SDE / ODE paths).""" - -from __future__ import annotations - -import math -from typing import Optional - -import torch - -CANONICAL_DYNAMICS_TYPES = ("sde", "flow_sde", "cps", "ode", "dance_sde") - - -def normalize_dynamics_type(name: str) -> str: - """Map a dynamics-type alias (CLI / legacy casing) to its canonical name.""" - key = str(name).strip().lower().replace("-", "_") - if key not in CANONICAL_DYNAMICS_TYPES: - raise ValueError( - f"Unknown dynamics_type {name!r}; expected one of {CANONICAL_DYNAMICS_TYPES}" - ) - return key - - -def sde_step_with_logprob_dynamics( - model_output: torch.FloatTensor, - sigma: torch.FloatTensor, - sigma_next: torch.FloatTensor, - sample: torch.FloatTensor, - sigmas: torch.FloatTensor, - prev_sample: Optional[torch.FloatTensor] = None, - generator: Optional[torch.Generator] = None, - deterministic: bool = False, - sigma_min_override: Optional[float] = None, - noise_level: float = 0.8, - dynamics_type: str = "flow_sde", -): - """Schedule-decoupled SDE step with log-prob for LTX-2.3 and similar models.""" - dynamics_type = normalize_dynamics_type(dynamics_type) - model_output = model_output.float() - sample = sample.float() - if prev_sample is not None: - prev_sample = prev_sample.float() - - ndim = sample.ndim - sigma_view = sigma.float() - sigma_next_view = sigma_next.float() - while sigma_view.ndim < ndim: - sigma_view = sigma_view.unsqueeze(-1) - while sigma_next_view.ndim < ndim: - sigma_next_view = sigma_next_view.unsqueeze(-1) - - dt = sigma_next_view - sigma_view - - sigma_max = sigmas[0].float().item() - if sigma_min_override is not None: - sigma_min = sigma_min_override - else: - sigma_min = max(sigmas[-2].float().item(), 1e-4) if len(sigmas) > 1 else 1e-4 - - if dynamics_type == "ode": - prev_sample_mean = sample + dt * model_output - std_dev_t = torch.zeros_like(sigma_view) - if prev_sample is None: - prev_sample = prev_sample_mean - log_prob = torch.zeros(sample.shape[0], dtype=sample.dtype, device=sample.device) - - elif dynamics_type == "flow_sde": - std_dev_t = (sigma_min + (sigma_max - sigma_min) * sigma_view) * noise_level - sigma_safe = torch.clamp(sigma_view, min=1e-8) - - drift_sample = 1.0 + std_dev_t**2 / (2.0 * sigma_safe) * dt - drift_model = (1.0 + std_dev_t**2 * (1.0 - sigma_view) / (2.0 * sigma_safe)) * dt - prev_sample_mean = sample * drift_sample + model_output * drift_model - - noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) - - if prev_sample is None: - if deterministic: - prev_sample = sample + dt * model_output - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + noise_scale * variance_noise - - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) - - torch.log(noise_scale + 1e-12) - - 0.5 * math.log(2.0 * math.pi) - ) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - elif dynamics_type == "cps": - angle = torch.tensor(noise_level, dtype=sigma_next_view.dtype, device=sigma_next_view.device) * torch.pi / 2 - std_dev_t = sigma_next_view * torch.sin(angle) - - x0 = sample - sigma_view * model_output - x1 = sample + model_output * (1.0 - sigma_view) - sqrt_term = torch.sqrt(torch.clamp(sigma_next_view**2 - std_dev_t**2, min=1e-12)) - prev_sample_mean = x0 * (1.0 - sigma_next_view) + x1 * sqrt_term - - if prev_sample is None: - if deterministic: - prev_sample = prev_sample_mean - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + std_dev_t * variance_noise - - log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - elif dynamics_type == "dance_sde": - sigma_safe = torch.clamp(sigma_view, min=1e-8) - x0_pred = sample - sigma_safe * model_output - std_dev_t = torch.as_tensor(noise_level, dtype=sample.dtype, device=sample.device) - log_term = 0.5 * noise_level**2 * (sample - x0_pred * (1.0 - sigma_view)) / (sigma_safe**2) - prev_sample_mean = sample + (model_output + log_term) * dt - noise_scale = std_dev_t * torch.sqrt(torch.clamp(-dt, min=1e-12)) - - if prev_sample is None: - if deterministic: - prev_sample = sample + dt * model_output - else: - variance_noise = torch.randn( - sample.shape, generator=generator, device=sample.device, dtype=sample.dtype, - ) - prev_sample = prev_sample_mean + noise_scale * variance_noise - - log_prob = ( - -((prev_sample.detach() - prev_sample_mean) ** 2) / (2.0 * noise_scale**2 + 1e-12) - - torch.log(noise_scale + 1e-12) - - 0.5 * math.log(2.0 * math.pi) - ) - log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) - - else: - raise ValueError( - f"dynamics_type {dynamics_type!r} is not supported by the " - "schedule-decoupled path; use flow_sde / cps / ode / dance_sde." - ) - - dt_sqrt = torch.sqrt(torch.clamp(-dt, min=1e-12)) - return prev_sample, log_prob, prev_sample_mean, std_dev_t, dt_sqrt 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))) From 7f7a590071311d9b3e80d3623d37991915d721ca Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 23 Jun 2026 03:50:08 +0000 Subject: [PATCH 26/31] refactor(ltx): extract TrainStepBackend and drop actor is_ltx branches Move load/forward/SDE orchestration into DiffusersTrainStepBackend and LTXTrainStepBackend so actor stays model-agnostic; TrainPipelineConfig keeps cond/trajectory/CFG hooks only. --- miles/backends/fsdp_utils/actor.py | 256 ++++--------- miles/backends/fsdp_utils/configs/ltx.py | 59 +-- .../configs/train_pipeline_config.py | 26 ++ .../backends/fsdp_utils/train_step_backend.py | 361 ++++++++++++++++++ 4 files changed, 468 insertions(+), 234 deletions(-) create mode 100644 miles/backends/fsdp_utils/train_step_backend.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 1a34cc67..8230c3bb 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -1,17 +1,14 @@ 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.qwen_image # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.sd3 # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config -from miles.backends.fsdp_utils.configs.ltx import is_ltx_model from miles.ray.train_actor import TrainRayActor from miles.utils import tracking_utils, train_metric_utils from miles.utils.context_utils import with_defer @@ -19,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 @@ -69,39 +65,25 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty diffusion_model_id = args.diffusion_model or args.hf_checkpoint self.train_pipeline_config = get_train_pipeline_config(diffusion_model_id) - self._is_ltx = is_ltx_model(args) + self.train_step_backend = self.train_pipeline_config.get_train_step_backend() - if self._is_ltx: - model, self.scheduler = self.train_pipeline_config.load_model_and_scheduler( - self.args, init_context_factory=self._get_init_weight_context_manager, - ) - else: - 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 + 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 and not self._is_ltx: - 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 = getattr(self.train_pipeline_config, "fsdp_wrap_classes", None) if self._is_ltx else None + fsdp_wrap = self.train_step_backend.get_fsdp_wrap_classes() model = apply_fsdp2( model, mesh=self.parallel_state.dp_mesh, @@ -309,10 +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 - if self._is_ltx: - use_cfg = False + use_cfg = self.train_step_backend.should_use_cfg(self.args) # ------------- KL loss ------------- kl_beta = float(self.args.diffusion_kl_beta) @@ -344,16 +323,12 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: num_train_timesteps = ( int(sched_config.num_train_timesteps) if sched_config is not None else 1000 ) - if self._is_ltx: - sigmas_ref = self.train_pipeline_config.resolve_sigmas_ref( - timesteps_ref, sigmas_snapshot, self.scheduler, - ) - else: - 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)]) + 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 @@ -465,19 +440,14 @@ def _build_train_grids( # prepare cond kwargs (denoising_env) denoising_env = denoising_envs[traj_idx] - if self._is_ltx: - positive_cond_kwargs_list.append( - train_pipeline_config.build_train_cond_kwargs( - denoising_env.pos_cond_kwargs, - latents=latents, - args=self.args, - device=device, - ) - ) - else: - positive_cond_kwargs_list.append( - train_pipeline_config.prepare_cond_kwargs(denoising_env.pos_cond_kwargs, device) + positive_cond_kwargs_list.append( + 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( train_pipeline_config.prepare_cond_kwargs(denoising_env.neg_cond_kwargs, device) @@ -652,6 +622,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"] @@ -664,12 +635,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) - if self._is_ltx: - timesteps_for_sde = timesteps_flat / float( - getattr(train_pipeline_config, "sde_timestep_divisor", 1000.0) - ) - else: - timesteps_for_sde = timesteps_flat + 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. @@ -713,81 +679,34 @@ 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:] + ) - if self._is_ltx: - def _forward(cond: dict) -> torch.Tensor: - return train_pipeline_config.forward_velocity( - self.model, latents_input, timesteps_input, cond, - ) - - def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: - adapter_ctx = self.model.disable_adapter() if disable_adapter else nullcontext() - with adapter_ctx: - return _forward(pos_cond_tile) - - noise_pred_flat = _compute_noise_pred() - sde_extra = train_pipeline_config.build_sde_extra( - self.scheduler, grids, sample_indices, tstep_indices, self.args, - ) - prev_sample_dummy, log_prob_new_flat, prev_sample_mean, std_dev_t = train_pipeline_config.sde_step( - self.scheduler, - noise_pred_flat, - timesteps_for_sde, - latents_flat, - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ), - noise_level=noise_level, - extra=sde_extra, - ) - del prev_sample_dummy - else: - 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_level=noise_level, - ) + 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 log_prob_new = log_prob_new_flat.reshape(tile_sample_count, tile_tstep_count) @@ -799,31 +718,31 @@ 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. - if self._is_ltx: - _, _, prev_sample_mean_ref, _ = train_pipeline_config.sde_step( - self.scheduler, - ref_noise_pred_flat, - timesteps_for_sde, - latents_flat, - prev_sample=next_latents_tile.reshape( - tile_sample_count * tile_tstep_count, *next_latents_tile.shape[2:] - ), - noise_level=noise_level, - extra=sde_extra, - ) - else: - _, _, 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(), - noise_level=noise_level, - ) + 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)), keepdim=True, @@ -857,12 +776,9 @@ 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()) - if self._is_ltx: - flat_train = noise_pred_flat.float().reshape(noise_pred_flat.shape[0], -1) - flat_rollout = rollout_mo_flat.float().reshape(rollout_mo_flat.shape[0], -1) - log_stats["model_output_cosine_sim"].append( - torch.nn.functional.cosine_similarity(flat_train, flat_rollout, dim=1).mean().detach() - ) + train_step_backend.append_model_output_compare_stats( + log_stats, noise_pred_flat, rollout_mo_flat + ) return loss @@ -910,20 +826,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 diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 7231dd0c..958ad60f 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -10,6 +10,7 @@ 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 @@ -452,60 +453,16 @@ def validate_args(args: Namespace) -> None: 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 - fsdp_wrap_classes = ["BasicAVTransformerBlock"] - lora_target_modules = [ "to_q", "to_k", "to_v", "to_out.0", "net.0.proj", "net.2", ] - def load_model_and_scheduler(self, args, init_context_factory): - from dataclasses import dataclass, field - - from ltx_core.components.schedulers import LTX2Scheduler - # load_ltx_transformer_for_train / resolve_transformer_checkpoint defined above - - @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") - master_dtype = {"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=master_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 prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: if cond is None: return {} @@ -532,18 +489,6 @@ def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> kwargs["audio_context_mask"] = audio_mask return kwargs - def resolve_sigmas_ref( - self, - timesteps_ref: torch.Tensor, - sigmas_snapshot: torch.Tensor | None, - scheduler, - ) -> torch.Tensor: - device = timesteps_ref.device - if sigmas_snapshot is not None: - return sigmas_snapshot.to(device).float() - sigmas_ref = timesteps_ref / float(self.sde_timestep_divisor) - return torch.cat([sigmas_ref, sigmas_ref.new_zeros(1)]) - def build_train_cond_kwargs( self, cond: CondKwargs | None, 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/train_step_backend.py b/miles/backends/fsdp_utils/train_step_backend.py new file mode 100644 index 00000000..67a42d05 --- /dev/null +++ b/miles/backends/fsdp_utils/train_step_backend.py @@ -0,0 +1,361 @@ +"""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: + pass + + +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() + ) From 776a0474391c7d5909ab8e063b8243954b74f942 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Tue, 23 Jun 2026 07:30:53 +0000 Subject: [PATCH 27/31] fix: rollout timeout --- scripts/run-diffusion-grpo-ltx23-sglang.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index 3b8a2913..ac074dae 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -84,6 +84,8 @@ fi --rollout-num-gpus "${NUM_GPUS}" \ --rollout-num-gpus-per-engine "${ROLLOUT_NUM_GPUS_PER_ENGINE:-1}" \ --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:-1}" \ --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ "${LORA_ARGS[@]}" \ From f0bffd205367dccfcaf635d5198df95f527a3098 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 24 Jun 2026 09:46:26 +0000 Subject: [PATCH 28/31] clean .gitignore --- .gitignore | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/.gitignore b/.gitignore index b0fb9d67..a017d4c6 100644 --- a/.gitignore +++ b/.gitignore @@ -177,6 +177,8 @@ settings.json .pypirc wandb/ +logs/ +dist/ outputs/ data/ local/ @@ -193,15 +195,3 @@ glm/ _examples_synced/ .env .DS_Store - -# crash dumps / runtime artifacts -core -core.* - -logs/ -wandb/ - -# local debug scripts / alignment tools (not shipped with the library) -dist/ -scripts/run-ltx23-grpo-local.sh -.claude/skills/align-ltx-train-rollout/ \ No newline at end of file From 019b606ca00dc50b6a54572bbcec09a9f3bb26f9 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 24 Jun 2026 11:57:27 +0000 Subject: [PATCH 29/31] style: fix lint --- miles/backends/fsdp_utils/actor.py | 19 ++---- miles/backends/fsdp_utils/configs/ltx.py | 64 ++++++++----------- miles/backends/fsdp_utils/ltx_geometry.py | 12 +--- .../backends/fsdp_utils/train_step_backend.py | 26 ++++---- .../sglang_diffusion_utils/configs/ltx.py | 3 +- .../monkey_patches/__init__.py | 4 +- .../patch_ltx2_ltxcore_parity.py | 31 +++------ .../patch_ltx2_rollout_cond_kwargs.py | 4 +- .../sglang_diffusion_engine.py | 11 +--- miles/ray/rollout.py | 1 - miles/rollout/rm_hub/pickscore.py | 5 +- miles/rollout/sglang_diffusion_rollout.py | 9 +-- miles/rollout/step_strategy_hub.py | 9 +-- miles/utils/arguments.py | 3 +- miles/utils/diffusion_rollout_response.py | 4 +- 15 files changed, 74 insertions(+), 131 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 8230c3bb..49bda3bc 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -6,9 +6,9 @@ import torch import torch.distributed as dist +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 -import miles.backends.fsdp_utils.configs.ltx # noqa: F401 — register pipeline config from miles.ray.train_actor import TrainRayActor from miles.utils import tracking_utils, train_metric_utils from miles.utils.context_utils import with_defer @@ -320,9 +320,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: timesteps_ref = dit_trajectories[0].timesteps.to(device).float() sigmas_snapshot = getattr(dit_trajectories[0], "sigmas", None) sched_config = getattr(self.scheduler, "config", None) - num_train_timesteps = ( - int(sched_config.num_train_timesteps) if sched_config is not None else 1000 - ) + 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, @@ -541,8 +539,7 @@ def _build_train_grids( "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) + if sde_indices_per_sample_list and len(sde_indices_per_sample_list) == (traj_end - traj_start) else None ), } @@ -776,9 +773,7 @@ def _forward_tile( 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 - ) + train_step_backend.append_model_output_compare_stats(log_stats, noise_pred_flat, rollout_mo_flat) return loss @@ -897,9 +892,9 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, fsdp_wrap_classe 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" - ) + 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 index 958ad60f..d19de0e0 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -70,9 +70,7 @@ def resolve_model_id(args) -> str: def _diffusion_cache_root() -> Path: - return Path( - os.environ.get("SGLANG_DIFFUSION_CACHE_ROOT", "/data/wenhao/sgl_diffusion_cache") - ) + 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: @@ -97,8 +95,7 @@ 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" + f"Materialized LTX model at {materialized_dir} is missing " f"transformer/model.safetensors" ) return checkpoint @@ -118,9 +115,7 @@ def _read_materialized_transformer_config(checkpoint: Path) -> dict: config_json = _materialized_config_path(checkpoint) if config_json is None: - raise FileNotFoundError( - f"Materialized LTX checkpoint {checkpoint} is missing sibling config.json" - ) + raise FileNotFoundError(f"Materialized LTX checkpoint {checkpoint} is missing sibling config.json") transformer_cfg = json.loads(config_json.read_text()) return {"transformer": transformer_cfg} @@ -142,10 +137,7 @@ def load_ltx_transformer_for_train( 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 ( - LTXModelConfigurator, - LTXV_MODEL_COMFY_RENAMING_MAP, - ) + 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(): @@ -160,7 +152,11 @@ def load_ltx_transformer_for_train( meta_model = create_meta_model(LTXModelConfigurator, config, ()) loader = SafetensorsModelStateDictLoader() sd = load_state_dict( - str(checkpoint), loader, DummyRegistry(), torch.device("cpu"), None, + str(checkpoint), + loader, + DummyRegistry(), + torch.device("cpu"), + None, ) state_dict = sd.sd if dtype is not None: @@ -252,7 +248,8 @@ def resolve_transformer_checkpoint( if _is_hf_model_id(str(diffusion_model)): materialized_dir = resolve_materialized_model_dir( - str(diffusion_model), materialize=materialize, + str(diffusion_model), + materialize=materialize, ) if materialized_dir is not None: checkpoint = _transformer_checkpoint_in_dir(materialized_dir) @@ -264,7 +261,8 @@ def resolve_transformer_checkpoint( return str(checkpoint) materialized_dir = resolve_materialized_model_dir( - LTX_DEFAULT_HF_MODEL, materialize=materialize, + LTX_DEFAULT_HF_MODEL, + materialize=materialize, ) if materialized_dir is not None: checkpoint = _transformer_checkpoint_in_dir(materialized_dir) @@ -273,7 +271,7 @@ def resolve_transformer_checkpoint( raise FileNotFoundError( "Could not resolve LTX transformer checkpoint. Pass --diffusion-model " - f"Lightricks/LTX-2.3 (recommended) or a .safetensors override." + "Lightricks/LTX-2.3 (recommended) or a .safetensors override." ) @@ -288,11 +286,10 @@ def server_kwargs_extras(args) -> dict: weights_path = None if explicit: weights_path = resolve_transformer_checkpoint( - getattr(args, "diffusion_model", None), explicit_path=explicit, + getattr(args, "diffusion_model", None), + explicit_path=explicit, ) - elif getattr(args, "diffusion_model", None) and str(args.diffusion_model).endswith( - ".safetensors" - ): + elif getattr(args, "diffusion_model", None) and str(args.diffusion_model).endswith(".safetensors"): weights_path = resolve_transformer_checkpoint(args.diffusion_model) if weights_path: @@ -344,9 +341,7 @@ def patch_rollout_sampling_params( 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." - ) + 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 @@ -459,8 +454,12 @@ class LTXTrainPipelineConfig(TrainPipelineConfig): sde_timestep_divisor = 1000.0 lora_target_modules = [ - "to_q", "to_k", "to_v", "to_out.0", - "net.0.proj", "net.2", + "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: @@ -502,9 +501,7 @@ def build_train_cond_kwargs( 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" - ) + 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: @@ -656,9 +653,7 @@ def forward_velocity_cfg_joint( timesteps_input: torch.Tensor, joint_cond: dict, ) -> torch.Tensor: - raise NotImplementedError( - "LTX trains with guidance_scale=1.0; --fsdp-cfg-batching is not supported." - ) + raise NotImplementedError("LTX trains with guidance_scale=1.0; --fsdp-cfg-batching is not supported.") def sde_step( self, @@ -674,9 +669,7 @@ def sde_step( 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',...}" - ) + 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() @@ -685,8 +678,7 @@ def sde_step( 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})." + 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( diff --git a/miles/backends/fsdp_utils/ltx_geometry.py b/miles/backends/fsdp_utils/ltx_geometry.py index 788fd352..9aef9c8a 100644 --- a/miles/backends/fsdp_utils/ltx_geometry.py +++ b/miles/backends/fsdp_utils/ltx_geometry.py @@ -66,9 +66,7 @@ def prepare_ltx_video_positions( 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 - ) + 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) @@ -76,9 +74,7 @@ def prepare_ltx_video_positions( 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, ...] + _LTX_CAUSAL_OFFSET - _LTX_SCALE_FACTORS[0]).clamp(min=0) pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) return pixel_coords @@ -96,9 +92,7 @@ def build_ltx_t2v_geometry( 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 - ) + 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( diff --git a/miles/backends/fsdp_utils/train_step_backend.py b/miles/backends/fsdp_utils/train_step_backend.py index 67a42d05..5a085fd5 100644 --- a/miles/backends/fsdp_utils/train_step_backend.py +++ b/miles/backends/fsdp_utils/train_step_backend.py @@ -49,8 +49,7 @@ def load_model_and_scheduler( init_context_factory, *, master_dtype: torch.dtype, - ) -> tuple[torch.nn.Module, object]: - ... + ) -> tuple[torch.nn.Module, object]: ... def apply_gradient_checkpointing(self, model: torch.nn.Module, args) -> None: if args.gradient_checkpointing: @@ -97,8 +96,7 @@ def compute_noise_pred( true_cfg_scale: float | None, fsdp_cfg_batching: bool, disable_adapter: bool = False, - ) -> torch.Tensor: - ... + ) -> torch.Tensor: ... @abc.abstractmethod def sde_step_logprob( @@ -124,7 +122,8 @@ def append_model_output_compare_stats( noise_pred: torch.Tensor, rollout_mo_flat: torch.Tensor, ) -> None: - pass + """Optional hook for comparing rollout vs train noise predictions.""" + return class DiffusersTrainStepBackend(TrainStepBackend): @@ -258,11 +257,14 @@ def to(self, 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] + resolved_dtype = ( + master_dtype + or { + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32, + }[master_dtype_name] + ) checkpoint = resolve_transformer_checkpoint( args.diffusion_model, @@ -274,7 +276,9 @@ def to(self, device): ltx_sched = LTX2Scheduler() sigmas = ltx_sched.execute(steps=num_steps).float() scheduler = _LTXSchedulerHolder( - sigmas=sigmas, timesteps=sigmas[:num_steps], num_inference_steps=num_steps, + sigmas=sigmas, + timesteps=sigmas[:num_steps], + num_inference_steps=num_steps, ) if getattr(args, "gradient_checkpointing", False): diff --git a/miles/backends/sglang_diffusion_utils/configs/ltx.py b/miles/backends/sglang_diffusion_utils/configs/ltx.py index 7f6ee79d..da0fed31 100644 --- a/miles/backends/sglang_diffusion_utils/configs/ltx.py +++ b/miles/backends/sglang_diffusion_utils/configs/ltx.py @@ -45,7 +45,8 @@ def resolve_ltx_transformer_weights_path( ) -> str | None: try: return resolve_transformer_checkpoint( - diffusion_model, explicit_path=explicit_path, + 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 e9127350..67c8ea4f 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -63,9 +63,7 @@ def apply_sgld_monkey_patches(*, include_ltx2_ltxcore: bool | None = None) -> No if include_ltx2_ltxcore is None: include_ltx2_ltxcore = os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" if include_ltx2_ltxcore: - from miles.backends.sglang_diffusion_utils.monkey_patches import ( - patch_ltx2_ltxcore_parity, - ) + from miles.backends.sglang_diffusion_utils.monkey_patches import patch_ltx2_ltxcore_parity patch_ltx2_ltxcore_parity.apply() diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py index fa02747a..bff31835 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py @@ -6,7 +6,8 @@ from __future__ import annotations -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import torch import torch.nn.functional as F @@ -17,12 +18,7 @@ def expand_temb_for_hidden(temb: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: """Broadcast batch-level temb ``[B, 1, D]`` to ``[B, T, D]`` when uniform.""" - if ( - temb.ndim == 3 - and temb.shape[1] == 1 - and hidden_states.ndim == 3 - and hidden_states.shape[1] > 1 - ): + if temb.ndim == 3 and temb.shape[1] == 1 and hidden_states.ndim == 3 and hidden_states.shape[1] > 1: return temb.expand(-1, hidden_states.shape[1], -1) return temb @@ -43,9 +39,7 @@ def _ltx_pytorch_sdpa( mask = mask.unsqueeze(0) if mask.ndim == 3: mask = mask.unsqueeze(1) - out = F.scaled_dot_product_attention( - q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False - ) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) return out.transpose(1, 2).reshape(b, -1, heads * dim_head) @@ -117,13 +111,8 @@ def _patched_get_ada_values( ) -> tuple[torch.Tensor, ...]: num_ada_params = int(scale_shift_table.shape[0]) ada_values = ( - scale_shift_table[indices] - .unsqueeze(0) - .unsqueeze(0) - .to(device=timestep.device, dtype=timestep.dtype) - + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[ - :, :, indices, : - ] + scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) + + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] ).unbind(dim=2) return ada_values @@ -176,9 +165,7 @@ def _patched_forward( gather_context_kv_for_sp: bool = False, ) -> torch.Tensor: from sglang.multimodal_gen.runtime.distributed import get_tp_world_size - from sglang.multimodal_gen.runtime.models.dits.ltx_2 import ( - apply_interleaved_rotary_emb, - ) + from sglang.multimodal_gen.runtime.models.dits.ltx_2 import apply_interleaved_rotary_emb if get_tp_world_size() > 1 or gather_context_kv_for_sp or self.use_local_attention: return orig_forward( @@ -259,9 +246,7 @@ def _patched_forward(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, tor audio_hidden_states = args[1] if isinstance(audio_hidden_states, torch.Tensor): kwargs = dict(kwargs) - kwargs["temb_audio"] = expand_temb_for_hidden( - kwargs["temb_audio"], audio_hidden_states - ) + kwargs["temb_audio"] = expand_temb_for_hidden(kwargs["temb_audio"], audio_hidden_states) return orig_forward(self, *args, **kwargs) return _patched_forward 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 index 1e32a580..99265003 100644 --- 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 @@ -26,9 +26,7 @@ def apply() -> None: if _APPLIED: return - from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( - LTX2DenoisingStage, - ) + from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import LTX2DenoisingStage if not hasattr(LTX2DenoisingStage, "_prepare_denoising_loop"): logger.warning( diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 502abd36..7cdb856a 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -68,7 +68,6 @@ def _scheduler_process_with_sgld_monkey_patches(*args, **kwargs): 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": @@ -190,17 +189,11 @@ 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 = bool(getattr(self.args, "apply_sgld_monkey_patches", False)) - apply_ltx = ( - ltx_config.is_ltx_model(self.args) - and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" - ) + apply_ltx = ltx_config.is_ltx_model(self.args) and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" 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)" - ) + 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") diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index b69e8c77..9a0e54a1 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -6,7 +6,6 @@ from pathlib import Path from typing import Any -import numpy as np import ray import torch from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy diff --git a/miles/rollout/rm_hub/pickscore.py b/miles/rollout/rm_hub/pickscore.py index c8b4245e..872fde39 100644 --- a/miles/rollout/rm_hub/pickscore.py +++ b/miles/rollout/rm_hub/pickscore.py @@ -74,10 +74,7 @@ def first_frame_for_wandb(t: torch.Tensor) -> np.ndarray | None: 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 - ] + 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: diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 51af6012..63398bb2 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -83,10 +83,7 @@ def build_rollout_generate_payload( ) -> dict[str, Any]: """Build full JSON payload for ``POST /rollout/generate`` (``RolloutImageRequest``).""" sampling_params["prompt"] = prompt - if ( - sampling_params.get("negative_prompt") is None - and float(sampling_params.get("guidance_scale", 1.0)) != 1.0 - ): + 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 @@ -451,9 +448,7 @@ async def eval_rollout_single_dataset( reward_key = args.eval_reward_key return { dataset_config.name: { - "rewards": [ - sample.get_reward_value(args, reward_key=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 862a9c3f..cf6846c3 100644 --- a/miles/rollout/step_strategy_hub.py +++ b/miles/rollout/step_strategy_hub.py @@ -51,14 +51,9 @@ def ltx_sde_candidates( 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 - ) + 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)" - ) + 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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 8b1e7adf..044ec060 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1303,8 +1303,7 @@ def miles_validate_args(args): args.train_memory_margin_bytes = 0 assert not (args.debug_rollout_only and args.debug_train_only), ( - "debug_rollout_only and debug_train_only cannot be set at the same time, " - "please set only one of them." + "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() diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 3eb82921..f014dd4c 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -86,9 +86,7 @@ def _parse_cond_kwargs( 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") - ), + 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), ) From 18921dcee8afb29bc8ae66ed581b28fb9bcc6f96 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Thu, 25 Jun 2026 12:00:40 +0000 Subject: [PATCH 30/31] refactor(ltx): simplify rollout patches and always disable AV cross-attn Remove ltxcore parity patch and optional toggles; LTX rollout now always applies cond kwargs and AV cross-off for train/rollout alignment. --- miles/backends/fsdp_utils/configs/ltx.py | 13 +- .../monkey_patches/__init__.py | 25 +- .../patch_ltx2_disable_av_cross.py | 14 +- .../patch_ltx2_ltxcore_parity.py | 300 ------------------ .../sglang_diffusion_engine.py | 6 +- miles/utils/arguments.py | 4 +- scripts/run-diffusion-grpo-ltx23-sglang.sh | 6 - 7 files changed, 17 insertions(+), 351 deletions(-) delete mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index d19de0e0..156e6228 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -358,11 +358,8 @@ def patch_rollout_engine_env_vars(env_vars: dict[str, str], args) -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV - if getattr(args, "ltx_disable_av_cross_attn", False): - env_vars["MILES_LTX_DISABLE_AV_CROSS"] = "1" - for name in (LTX_ROLLOUT_PATCHES_ENV, "MILES_LTX_DISABLE_AV_CROSS"): - if os.environ.get(name): - env_vars[name] = os.environ[name] + 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: @@ -403,12 +400,6 @@ def register_args(parser: ArgumentParser) -> None: default=None, help="Override σ_min for LTX SDE step.", ) - parser.add_argument( - "--ltx-disable-av-cross-attn", - action="store_true", - default=False, - help="Disable LTX A2V/V2A cross-attn in sglang rollout (align with ltx_core video-only train).", - ) parser.add_argument( "--ltx-gemma-path", type=str, diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 67c8ea4f..29b4e2f9 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -5,7 +5,7 @@ ``apply_rollout_patch_group``. - ``sgld``: diffusers / SD3 op parity (RMSNorm, RoPE, attention, …). -- ``ltx``: LTX-2 ltx_core parity + AV-off (rollout uses official gs=1 path). +- ``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. @@ -13,14 +13,12 @@ from __future__ import annotations -import os - 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_LTX2_LTXCORE_PARITY" +LTX_ROLLOUT_PATCHES_ENV = "MILES_APPLY_LTX_ROLLOUT_PATCHES" def resolve_rollout_patch_group(args) -> str | None: @@ -30,7 +28,7 @@ def resolve_rollout_patch_group(args) -> str | None: from miles.backends.sglang_diffusion_utils.configs.ltx import is_ltx_model - if is_ltx_model(args) and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1": + if is_ltx_model(args): return PATCH_GROUP_LTX return None @@ -38,12 +36,12 @@ def resolve_rollout_patch_group(args) -> str | None: def apply_rollout_patch_group(group: str | None) -> None: if group == PATCH_GROUP_SGLD: - apply_sgld_monkey_patches(include_ltx2_ltxcore=False) + apply_sgld_monkey_patches() elif group == PATCH_GROUP_LTX: apply_ltx2_rollout_patches() -def apply_sgld_monkey_patches(*, include_ltx2_ltxcore: bool | None = None) -> None: +def apply_sgld_monkey_patches() -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import ( patch_layernorm_scale_shift, patch_mul_add, @@ -60,22 +58,13 @@ def apply_sgld_monkey_patches(*, include_ltx2_ltxcore: bool | None = None) -> No patch_usp_attention.apply() patch_qk_norm_rope.apply() - if include_ltx2_ltxcore is None: - include_ltx2_ltxcore = os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" - if include_ltx2_ltxcore: - from miles.backends.sglang_diffusion_utils.monkey_patches import patch_ltx2_ltxcore_parity - - patch_ltx2_ltxcore_parity.apply() - def apply_ltx2_rollout_patches() -> None: - """LTX-2 ltx_core parity + video-only train alignment.""" + """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_ltxcore_parity, patch_ltx2_rollout_cond_kwargs, ) - patch_ltx2_ltxcore_parity.apply() - patch_ltx2_disable_av_cross.apply() 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 index f1172a57..2cb073d5 100644 --- 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 @@ -1,26 +1,18 @@ """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. Enabled via ``MILES_LTX_DISABLE_AV_CROSS=1`` -(set by miles rollout when ``--ltx-disable-av-cross-attn``), this injects the -disable flags into the DiT forward so the rollout video branch matches train. +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 -import os - _APPLIED = False -def _enabled() -> bool: - return os.environ.get("MILES_LTX_DISABLE_AV_CROSS", "0") == "1" - - def apply() -> None: global _APPLIED - if _APPLIED or not _enabled(): - _APPLIED = True + if _APPLIED: return from sglang.multimodal_gen.runtime.models.dits import ltx_2 as ltx2_mod diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py deleted file mode 100644 index bff31835..00000000 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_ltxcore_parity.py +++ /dev/null @@ -1,300 +0,0 @@ -"""LTX-2 DiT parity patches: align sglang ltx_2.py with miles/ltx_core. - -TODO(upstream): remove once sgl-d LTX rollout matches ltx_core AdaLN / temb / -velocity-to-x0 paths natively (train/rollout alignment checks pass without patch). -""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import Any - -import torch -import torch.nn.functional as F - -_ORIGINALS: dict[str, Any] = {} -_APPLIED = False - - -def expand_temb_for_hidden(temb: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: - """Broadcast batch-level temb ``[B, 1, D]`` to ``[B, T, D]`` when uniform.""" - if temb.ndim == 3 and temb.shape[1] == 1 and hidden_states.ndim == 3 and hidden_states.shape[1] > 1: - return temb.expand(-1, hidden_states.shape[1], -1) - return temb - - -def _ltx_pytorch_sdpa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - heads: int, - attn_mask: torch.Tensor | None = None, -) -> torch.Tensor: - b, _, dim_head = q.shape - dim_head //= heads - q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v)) - mask = attn_mask - if mask is not None: - if mask.ndim == 2: - mask = mask.unsqueeze(0) - if mask.ndim == 3: - mask = mask.unsqueeze(1) - out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) - return out.transpose(1, 2).reshape(b, -1, heads * dim_head) - - -def _linear_out(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: - return F.linear(x, module.weight, module.bias) - - -def _ltxcore_rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: - try: - from ltx_core.utils import rms_norm as ltx_rms_norm - - return ltx_rms_norm(x, eps=eps) - except ImportError: - return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps) - - -def _ltxcore_apply_split_rotary_emb( - x: torch.Tensor, - freqs: tuple[torch.Tensor, torch.Tensor], -) -> torch.Tensor: - cos, sin = freqs - try: - from ltx_core.model.transformer.rope import apply_split_rotary_emb as ltx_apply - - return ltx_apply(x, cos, sin) - except ImportError: - return _pytorch_apply_split_rotary_emb(x, cos, sin) - - -def _pytorch_apply_split_rotary_emb( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, -) -> torch.Tensor: - x_dtype = x.dtype - needs_reshape = False - if x.ndim != 4 and cos.ndim == 4: - b = x.shape[0] - _, h, t, _ = cos.shape - x = x.reshape(b, t, h, -1).swapaxes(1, 2) - needs_reshape = True - - last = x.shape[-1] - split_x = x.reshape(*x.shape[:-1], 2, last // 2) - first_x = split_x[..., :1, :] - second_x = split_x[..., 1:, :] - - cos_u = cos.unsqueeze(-2) - sin_u = sin.unsqueeze(-2) - - out = split_x * cos_u - first_out = out[..., :1, :] - second_out = out[..., 1:, :] - first_out.addcmul_(-sin_u, second_x) - second_out.addcmul_(sin_u, first_x) - - out = out.reshape(*out.shape[:-2], last) - if needs_reshape: - out = out.swapaxes(1, 2).reshape(b, t, -1) - return out.to(dtype=x_dtype) - - -def _patched_get_ada_values( - self, - scale_shift_table: torch.Tensor, - batch_size: int, - timestep: torch.Tensor, - indices: slice, -) -> tuple[torch.Tensor, ...]: - num_ada_params = int(scale_shift_table.shape[0]) - ada_values = ( - scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) - + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] - ).unbind(dim=2) - return ada_values - - -def _patched_ltx2_adaln_single_forward( - self, - timestep: torch.Tensor, - hidden_dtype: torch.dtype | None = None, -): - """Match ltx_core AdaLayerNormSingle embedding path.""" - from ltx_core.model.transformer.timestep_embedding import get_timestep_embedding - - t = timestep.reshape(-1).to(dtype=torch.float32) - t_freq = get_timestep_embedding( - t, - 256, - flip_sin_to_cos=True, - downscale_freq_shift=0, - ) - if hidden_dtype is not None: - t_freq = t_freq.to(dtype=hidden_dtype) - - te = self.emb.timestep_embedder - x = F.silu(_linear_out(te.linear_1, t_freq)) - embedded_timestep = _linear_out(te.linear_2, x).to(dtype=self.linear.weight.dtype) - out = _linear_out(self.linear, F.silu(embedded_timestep)) - - if timestep.ndim == 0: - batch = 1 - elif timestep.ndim == 1: - batch = 1 - else: - batch = timestep.shape[0] - out = out.view(batch, -1, out.shape[-1]) - embedded_timestep = embedded_timestep.view(batch, -1, embedded_timestep.shape[-1]) - return out, embedded_timestep - - -def _make_patched_ltx2_attention_forward(orig_forward: Callable[..., torch.Tensor]): - def _patched_forward( - self, - x: torch.Tensor, - context: torch.Tensor | None = None, - mask: torch.Tensor | None = None, - pe: tuple[torch.Tensor, torch.Tensor] | None = None, - k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, - perturbation_mask: torch.Tensor | None = None, - all_perturbed: bool = False, - skip_sequence_parallel_override: bool = False, - gather_context_kv_for_sp: bool = False, - ) -> torch.Tensor: - from sglang.multimodal_gen.runtime.distributed import get_tp_world_size - from sglang.multimodal_gen.runtime.models.dits.ltx_2 import apply_interleaved_rotary_emb - - if get_tp_world_size() > 1 or gather_context_kv_for_sp or self.use_local_attention: - return orig_forward( - self, - x, - context=context, - mask=mask, - pe=pe, - k_pe=k_pe, - perturbation_mask=perturbation_mask, - all_perturbed=all_perturbed, - skip_sequence_parallel_override=skip_sequence_parallel_override, - gather_context_kv_for_sp=gather_context_kv_for_sp, - ) - - gate_input = x - context_ = x if context is None else context - v = _linear_out(self.to_v, context_) - use_attention = not all_perturbed - - if use_attention: - q = _linear_out(self.to_q, x) - k = _linear_out(self.to_k, context_) - - if self.qk_norm: - assert self.q_norm is not None and self.k_norm is not None - q = self.q_norm(q) - k = self.k_norm(k) - - if pe is not None: - cos, sin = pe - k_cos, k_sin = pe if k_pe is None else k_pe - if cos.dim() == 3: - q = apply_interleaved_rotary_emb(q, (cos, sin)) - k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) - else: - q = _ltxcore_apply_split_rotary_emb(q, (cos, sin)) - k = _ltxcore_apply_split_rotary_emb(k, (k_cos, k_sin)) - - out = _ltx_pytorch_sdpa(q, k, v, self.local_heads, mask) - - if perturbation_mask is not None: - if perturbation_mask.ndim == out.ndim - 1: - perturbation_mask = perturbation_mask.unsqueeze(-1) - out = out * perturbation_mask + v * (1 - perturbation_mask) - else: - out = v - - if self.to_gate_logits is not None: - gate_logits = _linear_out(self.to_gate_logits, gate_input) - b, t = out.shape[:2] - out = out.view(b, t, self.local_heads, self.dim_head) - out = out * (2.0 * torch.sigmoid(gate_logits).unsqueeze(-1)) - out = out.view(b, t, self.local_heads * self.dim_head) - - return _linear_out(self.to_out[0], out) - - return _patched_forward - - -def _make_patched_ltx2_block_forward(orig_forward: Callable[..., tuple[torch.Tensor, torch.Tensor]]): - def _patched_forward(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: - args = list(args) - if len(args) >= 6 and isinstance(args[0], torch.Tensor) and isinstance(args[4], torch.Tensor): - args[4] = expand_temb_for_hidden(args[4], args[0]) - if len(args) >= 7 and isinstance(args[1], torch.Tensor) and isinstance(args[5], torch.Tensor): - args[5] = expand_temb_for_hidden(args[5], args[1]) - if "temb" in kwargs: - hidden_states = kwargs.get("hidden_states") - if hidden_states is None and args: - hidden_states = args[0] - if isinstance(hidden_states, torch.Tensor): - kwargs = dict(kwargs) - kwargs["temb"] = expand_temb_for_hidden(kwargs["temb"], hidden_states) - if "temb_audio" in kwargs: - audio_hidden_states = kwargs.get("audio_hidden_states") - if audio_hidden_states is None and len(args) >= 2: - audio_hidden_states = args[1] - if isinstance(audio_hidden_states, torch.Tensor): - kwargs = dict(kwargs) - kwargs["temb_audio"] = expand_temb_for_hidden(kwargs["temb_audio"], audio_hidden_states) - return orig_forward(self, *args, **kwargs) - - return _patched_forward - - -def apply() -> None: - global _APPLIED - if _APPLIED: - return - - from sglang.multimodal_gen.runtime.models.dits import ltx_2 as ltx2_mod - - if "rms_norm" not in _ORIGINALS: - _ORIGINALS["rms_norm"] = ltx2_mod.rms_norm - - def _patched_rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: - return _ltxcore_rms_norm(x, eps=eps) - - ltx2_mod.rms_norm = _patched_rms_norm - - if "apply_split_rotary_emb" not in _ORIGINALS: - _ORIGINALS["apply_split_rotary_emb"] = ltx2_mod.apply_split_rotary_emb - - def _patched_apply_split_rotary_emb( - x: torch.Tensor, - freqs: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - return _ltxcore_apply_split_rotary_emb(x, freqs) - - ltx2_mod.apply_split_rotary_emb = _patched_apply_split_rotary_emb - - adaln_cls = ltx2_mod.LTX2AdaLayerNormSingle - if "LTX2AdaLayerNormSingle.forward" not in _ORIGINALS: - _ORIGINALS["LTX2AdaLayerNormSingle.forward"] = adaln_cls.forward - adaln_cls.forward = _patched_ltx2_adaln_single_forward - - block_cls = ltx2_mod.LTX2TransformerBlock - if "LTX2TransformerBlock.get_ada_values" not in _ORIGINALS: - _ORIGINALS["LTX2TransformerBlock.get_ada_values"] = block_cls.get_ada_values - block_cls.get_ada_values = _patched_get_ada_values - - if "LTX2TransformerBlock.forward" not in _ORIGINALS: - _ORIGINALS["LTX2TransformerBlock.forward"] = block_cls.forward - block_cls.forward = _make_patched_ltx2_block_forward(block_cls.forward) - - attn_cls = ltx2_mod.LTX2Attention - if "LTX2Attention.forward" not in _ORIGINALS: - _ORIGINALS["LTX2Attention.forward"] = attn_cls.forward - attn_cls.forward = _make_patched_ltx2_attention_forward(attn_cls.forward) - - _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 7cdb856a..217e4881 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -10,7 +10,6 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs from miles.backends.sglang_diffusion_utils.configs import ltx as ltx_config -from miles.backends.sglang_diffusion_utils.monkey_patches import LTX_ROLLOUT_PATCHES_ENV from miles.ray.ray_actor import RayActor from miles.utils.http_utils import get_host_info @@ -64,6 +63,7 @@ def _scheduler_process_with_sgld_monkey_patches(*args, **kwargs): # 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 ( + LTX_ROLLOUT_PATCHES_ENV, apply_ltx2_rollout_patches, apply_sgld_monkey_patches, ) @@ -188,8 +188,10 @@ 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() + 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) and os.environ.get(LTX_ROLLOUT_PATCHES_ENV, "1") == "1" + 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" diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 044ec060..ab70bddb 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -395,9 +395,7 @@ def add_rollout_arguments(parser): default=False, help=( "Apply miles.backends.sglang_diffusion_utils.monkey_patches at " - "sglang-d startup. For SD3: diffusers bf16 parity. For LTX-2: also " - "applies ltx_core AdaLN/RoPE/attention parity patches " - "(disable via MILES_APPLY_LTX2_LTXCORE_PARITY=0). Small perf hit." + "sglang-d startup for SD3 diffusers bf16 parity. Small perf hit." ), ) parser.add_argument( diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index ac074dae..2b816413 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -57,11 +57,6 @@ if [[ "${USE_LORA:-1}" == "1" ]]; then ) fi -LTX_AV_ARGS=() -if [[ "${LTX_DISABLE_AV_CROSS_ATTN:-0}" == "1" ]]; then - LTX_AV_ARGS+=(--ltx-disable-av-cross-attn) -fi - "${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ --train-backend fsdp \ --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ @@ -89,7 +84,6 @@ fi --sglang-server-concurrency "${SGLANG_SERVER_CONCURRENCY:-1}" \ --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ "${LORA_ARGS[@]}" \ - "${LTX_AV_ARGS[@]}" \ --lr 2e-4 \ --adam-beta2 0.999 \ --weight-decay 1e-4 \ From c283eef36775a97a54f7a1d1589d546d096bde5c Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Sat, 11 Jul 2026 02:12:56 +0000 Subject: [PATCH 31/31] update script multi gpu train --- scripts/run-diffusion-grpo-ltx23-sglang.sh | 42 +++++++++++++++++----- 1 file changed, 33 insertions(+), 9 deletions(-) diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index 2b816413..925e8ea7 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -1,18 +1,36 @@ #!/usr/bin/env bash # LTX-2.3 video PickScore GRPO: sglang rollout + FSDP train (colocate). # -# Default: 1-GPU colocate (train FSDP + sglang rollout). Override NUM_GPUS for -# multi-GPU colocate. CPS dynamics, 3 SDE steps from candidates 0–9, clip 1e-4. +# 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. # -# Layout mirrors other scripts/run-diffusion-grpo-*.sh recipes: -# train+rollout share the first NUM_GPUS in CUDA_VISIBLE_DEVICES; -# optional pickscore worker uses additional GPUs when configured. +# 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:-0}" +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}" @@ -27,12 +45,18 @@ if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then fi DIFFUSION_MODEL="${DIFFUSION_MODEL:-Lightricks/LTX-2.3}" -NUM_GPUS="${NUM_GPUS:-1}" 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 @@ -77,11 +101,11 @@ fi --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:-1}" \ + --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:-1}" \ + --sglang-server-concurrency "${SGLANG_SERVER_CONCURRENCY}" \ --sglang-attention-backend "${SGLANG_ATTENTION_BACKEND:-torch_sdpa}" \ "${LORA_ARGS[@]}" \ --lr 2e-4 \