# SPDX-License-Identifier: Apache-2.0 """DiffusionNFT multi-reward policy optimization method.""" from __future__ import annotations import contextlib from collections import defaultdict from collections.abc import Iterator from typing import Any import torch import torch.distributed as dist from tqdm.auto import tqdm from fastvideo.dataset.parquet_dataset_map_style import ( get_parquet_files_and_length, read_row_from_parquet_file, ) from fastvideo.dataset.utils import collate_rows_from_parquet_schema from fastvideo.logger import init_logger from fastvideo.pipelines import TrainingBatch from fastvideo.train.methods.base import LogScalar, TrainingMethod from fastvideo.train.models.base import ModelBase from fastvideo.train.methods.rl.rewards import build_multi_reward_scorer from fastvideo.train.methods.rl.common import ( DiffusionSampler, RLValidationConfig, SamplingConfig, distributed_k_repeat_indices, media_to_video_array, validation_caption, validation_shard_indices, ) from fastvideo.train.utils.checkpoint import _FullModelState from fastvideo.train.utils.config import ( get_optional_float, get_optional_int, parse_betas, ) from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler from fastvideo.training.training_utils import ( EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases, ) logger = init_logger(__name__) class _DiffusionNFTEMAState: """DCP state wrapper for method-owned DiffusionNFT EMA.""" def __init__(self, method: DiffusionNFTMethod) -> None: self._method = method def state_dict(self) -> dict[str, Any]: ema = self._method._student_ema return { "shadow": ema.state_dict() if ema is not None else {}, "update_count": torch.tensor(self._method._ema_update_count, dtype=torch.long), } def load_state_dict( self, state_dict: dict[str, Any], ) -> None: ema = self._method._student_ema shadow = state_dict.get("shadow", {}) if ema is not None and isinstance(shadow, dict): ema.load_state_dict(shadow) update_count = state_dict.get("update_count", 0) if torch.is_tensor(update_count): update_count = int(update_count.item()) self._method._ema_update_count = int(update_count) class DiffusionNFTMethod(TrainingMethod): """DiffusionNFT-style RL for diffusion models. This method owns the algorithm's sample-then-inner-train loop. One ``Trainer`` step corresponds to one DiffusionNFT outer epoch. """ def __init__( self, *, cfg: Any, role_models: dict[str, ModelBase], ) -> None: super().__init__(cfg=cfg, role_models=role_models) if "old" not in role_models: raise ValueError("DiffusionNFTMethod requires role 'old'") if "reference" not in role_models: raise ValueError("DiffusionNFTMethod requires role 'reference'") if not self.student._trainable: raise ValueError("DiffusionNFTMethod requires a trainable student") self.old = role_models["old"] self.reference = role_models["reference"] self.student.init_preprocessors(self.training_config) self._sampling_config = self._parse_sampling_config() self._sampler = DiffusionSampler(self._sampling_config) self._validation_config = RLValidationConfig.from_mapping(self.method_config.get("validation")) self._validation_sampling_config = self._parse_validation_sampling_config() self._validation_sampler = DiffusionSampler(self._validation_sampling_config) self._validation_items: list[tuple[int, bool, dict[str, Any]]] | None = None self._sample_steps = int(self._sampling_config.num_steps) self._sample_train_batch_size = self._read_int( "sample_train_batch_size", max(1, int(self.training_config.data.train_batch_size or 1)), ) self._train_batch_size = self._read_int("train_batch_size", self._sample_train_batch_size) self._num_batches_per_epoch = self._read_int("num_batches_per_epoch", 48) self._num_inner_epochs = self._read_int("num_inner_epochs", 1) self._num_video_per_prompt = self._read_int("num_video_per_prompt", 24) self._adv_clip_max = self._read_float("adv_clip_max", 5.0) self._timestep_fraction = self._read_float("timestep_fraction", 0.99) self._kl_beta = self._read_float("kl_beta", 0.0001) self._nft_beta = self._read_float("beta", 0.1) self._max_grad_norm = self._read_float("max_grad_norm", 1.0) self._decay_type = self._read_int("decay_type", 1) self._adv_mode = str(self.method_config.get("adv_mode", "all") or "all").strip().lower() self._terminal_progress = bool(self.method_config.get("terminal_progress", True)) ema_config = self._parse_ema_config() self._ema_enabled = bool(ema_config["enabled"]) self._ema_decay = float(ema_config["decay"]) self._ema_update_after_step = int(ema_config["update_after_step"]) self._validation_use_ema = bool(ema_config["validation"]) self._student_ema: EMA_FSDP | None = None self._ema_update_count = 0 self._trained_prompt_hashes: set[int] = set() if self._adv_mode not in {"all", "positive_only", "negative_only", "one_only", "binary"}: raise ValueError("method.adv_mode must be one of " "{all, positive_only, negative_only, one_only, binary}") reward_fn = self.method_config.get("reward_fn", None) if not isinstance(reward_fn, dict) or not reward_fn: raise ValueError("method.reward_fn must be a non-empty mapping, " "for example {pickscore: 1.0, clipscore: 1.0}") self._reward_fn_config = {str(k): float(v) for k, v in reward_fn.items()} unsupported = sorted(set(self._reward_fn_config) - {"pickscore", "clipscore"}) if unsupported: raise ValueError(f"Unsupported DiffusionNFT reward(s): {unsupported}. " "Only pickscore and clipscore are currently ported.") self._reward_scorer: Any | None = None self._init_optimizer_and_scheduler() @property def _optimizer_dict(self, ) -> dict[str, torch.optim.Optimizer]: return {"student": self._student_optimizer} @property def _lr_scheduler_dict(self) -> dict[str, Any]: return {"student": self._student_lr_scheduler} def manages_optimization(self) -> bool: return True def single_train_step( self, batch: dict[str, Any], iteration: int, ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]: del batch raise RuntimeError("DiffusionNFTMethod uses managed_train_step()") def managed_train_step( self, data_stream: Iterator[dict[str, Any]], iteration: int, ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]: self._log_progress(f"[DiffusionNFT] outer step {iteration}: start sampling " f"{self._num_batches_per_epoch} batches") sample_items = self._sample_epoch(data_stream, iteration) self._log_progress(f"[DiffusionNFT] outer step {iteration}: scoring rewards") rewards = self._score_samples(sample_items) self._log_progress(f"[DiffusionNFT] outer step {iteration}: computing advantages") advantages = self._compute_advantages(sample_items, rewards) self._log_progress(f"[DiffusionNFT] outer step {iteration}: start inner training") loss_map, metrics = self._inner_train(sample_items, advantages, iteration) self._update_old_model(iteration) metrics.update(self._reward_metrics(rewards)) metrics.update(self._reward_diagnostic_metrics(sample_items, rewards)) metrics["nft/num_sampled"] = float(sum(item["latents_clean"].shape[0] for item in sample_items)) return loss_map, {}, metrics def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]: return self._maybe_run_validation(iteration) def get_optimizers( self, iteration: int, ) -> list[torch.optim.Optimizer]: del iteration return [self._student_optimizer] def get_lr_schedulers( self, iteration: int, ) -> list[Any]: del iteration return [self._student_lr_scheduler] def checkpoint_state(self) -> dict[str, Any]: states = super().checkpoint_state() states["roles.old.transformer"] = _FullModelState(self.old.transformer) if self._ema_enabled: states["diffusion_nft.ema"] = _DiffusionNFTEMAState(self) return states def seed_optimizer_state_for_resume(self) -> None: super().seed_optimizer_state_for_resume() def on_train_start(self) -> None: super().on_train_start() self._sync_old_from_student() if self._ema_enabled: self._student_ema = EMA_FSDP( self.student.transformer, decay=self._ema_decay, mode="local_shard", ) self._log_progress(f"[DiffusionNFT] EMA enabled (decay={self._ema_decay})") self._reward_scorer = build_multi_reward_scorer( self._reward_fn_config, device=self.student.device, ) def _init_optimizer_and_scheduler(self) -> None: lr = float(self.training_config.optimizer.learning_rate) betas = self.training_config.optimizer.betas betas_raw = self.method_config.get("betas", None) if betas_raw is not None: betas = parse_betas(betas_raw, where="method.betas") params = [p for p in self.student.transformer.parameters() if p.requires_grad] self._student_optimizer, self._student_lr_scheduler = build_optimizer_and_scheduler( params=params, optimizer_config=self.training_config.optimizer, loop_config=self.training_config.loop, learning_rate=lr, betas=betas, scheduler_name=str(self.training_config.optimizer.lr_scheduler), ) def _sample_epoch( self, data_stream: Iterator[dict[str, Any]], iteration: int, ) -> list[dict[str, Any]]: self.student.transformer.eval() self.old.transformer.eval() sample_items: list[dict[str, Any]] = [] with torch.no_grad(): for batch_idx in tqdm( range(self._num_batches_per_epoch), desc=f"DiffusionNFT step {iteration}: sampling", position=1, leave=False, disable=not self._show_terminal_progress(), ): raw_batch = self._sample_prompt_batch(data_stream, iteration, batch_idx) prompts = self._extract_prompts(raw_batch) batch = self.student.prepare_batch( raw_batch, generator=self.cuda_generator, latents_source="zeros", ) sampling_result = self._sampler.sample( self.old, batch, generator=self.cuda_generator, ) latents_clean = sampling_result.latents train_timesteps = sampling_result.timesteps[:self._num_train_timesteps()] media = self.student.decode_latents(latents_clean) sample_items.append({ "encoder_hidden_states": batch.encoder_hidden_states.detach(), "encoder_attention_mask": batch.encoder_attention_mask.detach(), "latents_clean": latents_clean.detach(), "timesteps": train_timesteps.detach().unsqueeze(0).repeat(latents_clean.shape[0], 1), "media": media.detach().cpu(), "prompts": prompts, }) return sample_items def _sample_prompt_batch( self, data_stream: Iterator[dict[str, Any]], iteration: int, batch_idx: int, ) -> dict[str, Any]: dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) if dataset is None or not all( hasattr(dataset, attr) for attr in ("parquet_files", "lengths", "parquet_schema")): return self._repeat_first_prompt(next(data_stream), self._sample_train_batch_size) rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 total_rows = int(sum(dataset.lengths)) sample = distributed_k_repeat_indices( dataset_length=total_rows, batch_size=self._sample_train_batch_size, repeats_per_prompt=self._num_video_per_prompt, world_size=world_size, rank=rank, seed=int(self.training_config.data.seed) + (int(iteration) - 1) * self._num_batches_per_epoch + batch_idx, ) rows = [] for prompt_idx in sample.local_indices: row = read_row_from_parquet_file(dataset.parquet_files, prompt_idx, dataset.lengths) row["_sample_index"] = prompt_idx rows.append(row) return self._collate_rows(rows) def _score_samples(self, sample_items: list[dict[str, Any]]) -> dict[str, torch.Tensor]: if self._reward_scorer is None: raise RuntimeError("Reward scorer has not been initialized") media = torch.cat([item["media"] for item in sample_items], dim=0) prompts = [prompt for item in sample_items for prompt in item["prompts"]] reward_dict = self._reward_scorer(media, prompts) device = self.student.device return {k: v.to(device=device, dtype=torch.float32) for k, v in reward_dict.items()} @torch.no_grad() def _maybe_run_validation(self, iteration: int) -> dict[str, LogScalar]: config = self._validation_config if config.every_steps <= 0 or iteration % config.every_steps != 0: return {} if self._reward_scorer is None: raise RuntimeError("Reward scorer has not been initialized") use_ema = self._validation_use_ema and self._student_ema is not None suffix = " with EMA" if use_ema else "" self._log_progress(f"[DiffusionNFT] validation step {iteration}: start " f"{config.num_prompts} prompts, {config.num_steps} sampling steps{suffix}") with self._ema_context(): return self._run_validation(iteration) @torch.no_grad() def _run_validation(self, iteration: int) -> dict[str, LogScalar]: config = self._validation_config if self._reward_scorer is None: raise RuntimeError("Reward scorer has not been initialized") self.student.transformer.eval() rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 prepare_generator = torch.Generator(device=self.student.device).manual_seed(config.seed + 100_000 + rank) validation_generator = torch.Generator(device=self.student.device).manual_seed(config.seed + rank) local_rewards: dict[str, list[torch.Tensor]] = defaultdict(list) local_masks: list[torch.Tensor] = [] local_logs: list[dict[str, Any]] = [] items = self._get_validation_items() total_batches = max(1, (len(items) + config.batch_size - 1) // config.batch_size) for start in tqdm( range(0, len(items), config.batch_size), total=total_batches, desc=f"DiffusionNFT validation {iteration}", position=1, leave=False, disable=not self._show_terminal_progress(), ): batch_items = items[start:start + config.batch_size] raw_batch = self._collate_validation_rows([item[2] for item in batch_items]) prompts = self._extract_prompts(raw_batch) batch = self.student.prepare_batch( raw_batch, generator=prepare_generator, latents_source="zeros", ) sampling_result = self._validation_sampler.sample( self.student, batch, generator=validation_generator, ) media = self.student.decode_latents(sampling_result.latents).detach().cpu() # DiffusionNFT's original ``train_nft_sd3.py::eval_fn`` scores # validation samples and puts reward values in sample captions. rewards = self._reward_scorer(media, prompts) valid_mask = torch.tensor([item[1] for item in batch_items], device=self.student.device, dtype=torch.float32) local_masks.append(valid_mask) for key, value in rewards.items(): local_rewards[key].append(value.to(device=self.student.device, dtype=torch.float32)) if config.log_samples: for sample_idx, (global_idx, valid, _) in enumerate(batch_items): if not valid: continue local_logs.append({ "index": int(global_idx), "prompt": prompts[sample_idx], "media": media[sample_idx], "rewards": { key: float(value[sample_idx].detach().float().cpu()) for key, value in rewards.items() }, }) if not local_rewards or not local_masks: return {} local_mask = torch.cat(local_masks, dim=0) gathered_mask = self._gather_tensor(local_mask).bool() metrics: dict[str, LogScalar] = {} for key, chunks in local_rewards.items(): values = torch.cat(chunks, dim=0) gathered_values = self._gather_tensor(values.detach().float()) valid_values = gathered_values[gathered_mask] if valid_values.numel() > 0: metrics[f"validation/reward/{key}"] = valid_values.mean() metrics["validation/num_prompts"] = gathered_mask.float().sum() if config.log_samples: self._log_progress(f"[DiffusionNFT] validation step {iteration}: logging samples") self._log_validation_samples(local_logs, iteration) self._log_progress(f"[DiffusionNFT] validation step {iteration}: finished") return metrics def _get_validation_items(self) -> list[tuple[int, bool, dict[str, Any]]]: if self._validation_items is not None: return self._validation_items dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) if dataset is None: raise RuntimeError("DiffusionNFT validation requires the student dataloader dataset") data_path = self._validation_config.data_path or getattr(dataset, "path", None) if not data_path: data_path = self.training_config.data.data_path if self._validation_config.data_path is None or data_path == getattr(dataset, "path", None): parquet_files = list(dataset.parquet_files) lengths = list(dataset.lengths) else: parquet_files, lengths = get_parquet_files_and_length(data_path) parquet_files = list(parquet_files) lengths = list(lengths) total_rows = int(sum(lengths)) if total_rows <= 0: raise RuntimeError(f"Validation data_path {data_path!r} has no rows") num_prompts = min(self._validation_config.num_prompts, total_rows) rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 items: list[tuple[int, bool, dict[str, Any]]] = [] for prompt_idx, valid in validation_shard_indices(num_prompts, rank=rank, world_size=world_size): row = read_row_from_parquet_file(parquet_files, prompt_idx, lengths) row["_sample_index"] = prompt_idx items.append((prompt_idx, valid, row)) self._validation_items = items return items def _collate_validation_rows(self, rows: list[dict[str, Any]]) -> dict[str, Any]: return self._collate_rows(rows) def _collate_rows(self, rows: list[dict[str, Any]]) -> dict[str, Any]: dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) if dataset is None or not hasattr(dataset, "parquet_schema"): raise RuntimeError("DiffusionNFT requires a parquet-backed student dataset") return collate_rows_from_parquet_schema( rows, dataset.parquet_schema, int(getattr(dataset, "text_padding_length", 512)), cfg_rate=0.0, seed=self._validation_config.seed, ) def _log_validation_samples( self, local_logs: list[dict[str, Any]], iteration: int, ) -> None: logs = self._gather_objects(local_logs) rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 if rank != 0 or not logs: return tracker = getattr(self, "tracker", None) if tracker is None: return artifacts = [] for item in sorted(logs, key=lambda x: int(x["index"])): artifact = tracker.video( media_to_video_array(item["media"]), caption=validation_caption(str(item["prompt"]), item["rewards"]), fps=1, ) if artifact is not None: artifacts.append(artifact) if artifacts: tracker.log_artifacts({"validation/videos": artifacts}, step=iteration) def _compute_advantages( self, sample_items: list[dict[str, Any]], rewards: dict[str, torch.Tensor], ) -> torch.Tensor: avg = rewards["avg"].detach().float() local_prompts = [prompt for item in sample_items for prompt in item["prompts"]] gathered_rewards = self._gather_tensor(avg) gathered_prompts = self._gather_prompts(local_prompts) rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 local_count = int(avg.shape[0]) if len(gathered_prompts) != int(gathered_rewards.shape[0]): raise RuntimeError("Gathered prompt count does not match gathered rewards") global_advantages = torch.empty_like(gathered_rewards) prompt_to_indices: dict[str, list[int]] = defaultdict(list) for idx, prompt in enumerate(gathered_prompts): prompt_to_indices[prompt].append(idx) for indices in prompt_to_indices.values(): index_tensor = torch.tensor(indices, device=gathered_rewards.device, dtype=torch.long) group_rewards = gathered_rewards[index_tensor] group_std = group_rewards.std(unbiased=False) global_advantages[index_tensor] = (group_rewards - group_rewards.mean()) / (group_std + 1e-4) start = rank * local_count end = start + local_count advantages = global_advantages[start:end].to(avg.device) num_train_timesteps = self._num_train_timesteps() return advantages.unsqueeze(1).repeat(1, num_train_timesteps) def _inner_train( self, sample_items: list[dict[str, Any]], advantages: torch.Tensor, iteration: int, ) -> tuple[dict[str, torch.Tensor], dict[str, LogScalar]]: self.student.transformer.train() self.old.transformer.eval() self.reference.transformer.eval() self._student_optimizer.zero_grad(set_to_none=True) samples = self._collate_samples(sample_items) total_samples = int(samples["latents_clean"].shape[0]) num_train_timesteps = self._num_train_timesteps() if total_samples != int(advantages.shape[0]): raise RuntimeError("advantages and samples have mismatched batch sizes") effective_grad_accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1)) effective_grad_accum *= max(1, num_train_timesteps) current_accum = 0 optimizer_steps = 0 loss_terms: dict[str, list[torch.Tensor]] = defaultdict(list) num_batches = max(1, total_samples // max(1, self._train_batch_size)) training_batch_size = max(1, total_samples // num_batches) total_train_batches = self._num_inner_epochs * num_batches with tqdm( total=total_train_batches, desc=f"DiffusionNFT step {iteration}: training", position=1, leave=False, disable=not self._show_terminal_progress(), ) as progress: for _ in range(self._num_inner_epochs): perm = torch.randperm( total_samples, device=self.student.device, generator=self.cuda_generator, ) shuffled = {k: v[perm] for k, v in samples.items()} shuffled_adv = advantages[perm] perms_time = torch.stack([ torch.randperm(num_train_timesteps, device=self.student.device, generator=self.cuda_generator) for _ in range(total_samples) ]) shuffled["timesteps"] = shuffled["timesteps"][ torch.arange(total_samples, device=self.student.device)[:, None], perms_time, ] shuffled_adv = shuffled_adv[ torch.arange(total_samples, device=self.student.device)[:, None], perms_time, ] for batch_idx in range(num_batches): start = batch_idx * training_batch_size end = total_samples if batch_idx == num_batches - 1 else (batch_idx + 1) * training_batch_size train_sample = {k: v[start:end] for k, v in shuffled.items()} train_adv = shuffled_adv[start:end] for timestep_idx in range(num_train_timesteps): losses, backward_ctx = self._training_timestep_loss( train_sample, train_adv[:, timestep_idx], timestep_idx, ) self.student.backward( losses["total_loss"], backward_ctx, grad_accum_rounds=effective_grad_accum, ) current_accum += 1 for key, value in losses.items(): loss_terms[key].append(value.detach()) if current_accum % effective_grad_accum == 0: self._clip_student_grads() self._student_optimizer.step() self._student_lr_scheduler.step() self._update_ema() self._student_optimizer.zero_grad(set_to_none=True) optimizer_steps += 1 progress.update(1) if current_accum % effective_grad_accum != 0: self._clip_student_grads() self._student_optimizer.step() self._student_lr_scheduler.step() self._update_ema() self._student_optimizer.zero_grad(set_to_none=True) optimizer_steps += 1 self._log_progress(f"[DiffusionNFT] outer step {iteration}: finished inner training " f"({current_accum} micro-steps, {optimizer_steps} optimizer steps)") reduced_local = { key: torch.stack(values).mean() if values else torch.zeros((), device=self.student.device) for key, values in loss_terms.items() } reduced = {key: self._mean_scalar_across_ranks(value) for key, value in reduced_local.items()} reduced.setdefault("total_loss", torch.zeros((), device=self.student.device)) metrics: dict[str, LogScalar] = { "nft/iteration": float(iteration), "nft/num_inner_epochs": float(self._num_inner_epochs), "nft/inner_micro_steps": float(current_accum), "nft/optimizer_steps": float(optimizer_steps), "ema/update_count": float(self._ema_update_count), } return reduced, metrics @contextlib.contextmanager def _ema_context(self) -> Iterator[None]: if self._validation_use_ema and self._student_ema is not None: with self._student_ema.apply_to_model(self.student.transformer): yield return yield def _update_ema(self) -> None: if not self._ema_enabled or self._student_ema is None: return if self._ema_update_count >= self._ema_update_after_step: self._student_ema.update(self.student.transformer) self._ema_update_count += 1 def _log_progress(self, message: str) -> None: if not self._terminal_progress: return rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 if rank == 0: logger.info(message) def _show_terminal_progress(self) -> bool: rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 return bool(self._terminal_progress and rank == 0) def _training_timestep_loss( self, sample: dict[str, torch.Tensor], advantages: torch.Tensor, timestep_idx: int, ) -> tuple[dict[str, torch.Tensor], tuple[torch.Tensor, Any]]: # Ported from DiffusionNFT's ``scripts/train_nft_sd3.py`` inner # timestep loss, adapted to WanModel's predict_noise interface. x0 = sample["latents_clean"] timestep = sample["timesteps"][:, timestep_idx].to(device=x0.device) t = timestep.float() / float(self.student.num_train_timesteps) t_expanded = t.view(-1, *([1] * (x0.ndim - 1))) noise = torch.randn( x0.shape, device=x0.device, dtype=x0.dtype, generator=self.cuda_generator, ) xt = ((1 - t_expanded) * x0 + t_expanded * noise).to(dtype=x0.dtype) batch = self._make_training_batch(sample, timestep) with torch.no_grad(): old_prediction = self.old.predict_noise( xt, timestep, batch, conditional=True, attn_kind="dense", ).detach() ref_forward_prediction = self.reference.predict_noise( xt, timestep, batch, conditional=True, attn_kind="dense", ).detach() forward_prediction = self.student.predict_noise( xt, timestep, batch, conditional=True, attn_kind="dense", ) advantages_clip = torch.clamp(advantages, -self._adv_clip_max, self._adv_clip_max) if self._adv_mode == "positive_only": advantages_clip = torch.clamp(advantages_clip, 0, self._adv_clip_max) elif self._adv_mode == "negative_only": advantages_clip = torch.clamp(advantages_clip, -self._adv_clip_max, 0) elif self._adv_mode == "one_only": advantages_clip = torch.where(advantages_clip > 0, torch.ones_like(advantages_clip), torch.zeros_like(advantages_clip)) elif self._adv_mode == "binary": advantages_clip = torch.sign(advantages_clip) normalized_advantages_clip = (advantages_clip / self._adv_clip_max) / 2.0 + 0.5 r = torch.clamp(normalized_advantages_clip, 0, 1) positive_prediction = self._nft_beta * forward_prediction + (1 - self._nft_beta) * old_prediction.detach() implicit_negative_prediction = ((1.0 + self._nft_beta) * old_prediction.detach() - self._nft_beta * forward_prediction) x0_prediction = xt - t_expanded * positive_prediction with torch.no_grad(): weight_factor = torch.abs(x0_prediction.double() - x0.double()) weight_factor = weight_factor.mean(dim=tuple(range(1, x0.ndim)), keepdim=True).clip(min=0.00001) positive_loss = ((x0_prediction - x0)**2 / weight_factor).mean(dim=tuple(range(1, x0.ndim))) negative_x0_prediction = xt - t_expanded * implicit_negative_prediction with torch.no_grad(): negative_weight_factor = torch.abs(negative_x0_prediction.double() - x0.double()) negative_weight_factor = negative_weight_factor.mean(dim=tuple(range(1, x0.ndim)), keepdim=True).clip(min=0.00001) negative_loss = ((negative_x0_prediction - x0)**2 / negative_weight_factor).mean(dim=tuple(range(1, x0.ndim))) ori_policy_loss = r * positive_loss / self._nft_beta + (1.0 - r) * negative_loss / self._nft_beta policy_loss = (ori_policy_loss * self._adv_clip_max).mean() kl_div_loss = ((forward_prediction - ref_forward_prediction)**2).mean(dim=tuple(range(1, x0.ndim))).mean() total_loss = policy_loss + self._kl_beta * kl_div_loss losses = { "total_loss": total_loss, "policy_loss": policy_loss, "unweighted_policy_loss": ori_policy_loss.mean(), "kl_div_loss": kl_div_loss, "old_deviate": ((forward_prediction - old_prediction)**2).mean(), "old_kl_div": ((old_prediction - ref_forward_prediction)**2).mean(), "x0_norm": torch.mean(x0**2), } return losses, (batch.timesteps, batch.attn_metadata) def _repeat_first_prompt( self, raw_batch: dict[str, Any], repeat: int, ) -> dict[str, Any]: out: dict[str, Any] = {} for key, value in raw_batch.items(): if torch.is_tensor(value) and value.shape[0] > 0: out[key] = value[:1].repeat((repeat, ) + (1, ) * (value.ndim - 1)) elif key == "info_list" and isinstance(value, list) and value: out[key] = [dict(value[0]) for _ in range(repeat)] elif isinstance(value, list) and value: out[key] = [value[0] for _ in range(repeat)] else: out[key] = value return out @staticmethod def _extract_prompts(raw_batch: dict[str, Any]) -> list[str]: infos = raw_batch.get("info_list") if isinstance(infos, list) and infos: prompts: list[str] = [] for info in infos: if isinstance(info, dict): prompts.append(str(info.get("prompt") or info.get("caption") or "")) else: prompts.append("") return prompts captions = raw_batch.get("caption_text") if isinstance(captions, list): return [str(c) for c in captions] raise ValueError("Could not find prompts in batch info_list or caption_text") def _collate_samples(self, sample_items: list[dict[str, Any]]) -> dict[str, torch.Tensor]: keys = ["encoder_hidden_states", "encoder_attention_mask", "latents_clean", "timesteps"] return {key: torch.cat([item[key] for item in sample_items], dim=0).to(self.student.device) for key in keys} @staticmethod def _make_training_batch( sample: dict[str, torch.Tensor], timestep: torch.Tensor, ) -> TrainingBatch: batch = TrainingBatch() batch.encoder_hidden_states = sample["encoder_hidden_states"] batch.encoder_attention_mask = sample["encoder_attention_mask"] batch.conditional_dict = { "encoder_hidden_states": batch.encoder_hidden_states, "encoder_attention_mask": batch.encoder_attention_mask, } batch.timesteps = timestep batch.raw_latent_shape = tuple(sample["latents_clean"].shape) return batch def _num_train_timesteps(self) -> int: schedule_len = self._sample_steps if self._sampling_config.timesteps is not None: schedule_len = len(self._sampling_config.timesteps) elif self._sampling_config.sigmas is not None: schedule_len = len(self._sampling_config.sigmas) return max(1, min(schedule_len, int(schedule_len * self._timestep_fraction))) def _clip_student_grads(self) -> None: if self._max_grad_norm <= 0.0: return clip_grad_norm_while_handling_failing_dtensor_cases( [p for p in self.student.transformer.parameters()], self._max_grad_norm, foreach=None, ) def _sync_old_from_student(self) -> None: with torch.no_grad(): for src, tgt in zip(self.student.transformer.parameters(), self.old.transformer.parameters(), strict=True): tgt.data.copy_(src.detach().data) def _update_old_model(self, iteration: int) -> None: decay = self._return_decay(iteration, self._decay_type) with torch.no_grad(): for src, tgt in zip(self.student.transformer.parameters(), self.old.transformer.parameters(), strict=True): tgt.data.copy_(tgt.detach().data * decay + src.detach().data * (1.0 - decay)) @staticmethod def _return_decay(step: int, decay_type: int) -> float: # Ported from DiffusionNFT's ``scripts/train_nft_sd3.py::return_decay``. if decay_type == 0: flat, uprate, uphold = 0, 0.0, 0.0 elif decay_type == 1: flat, uprate, uphold = 0, 0.001, 0.5 elif decay_type == 2: flat, uprate, uphold = 75, 0.0075, 0.999 else: raise ValueError(f"Unsupported decay_type: {decay_type}") if step < flat: return 0.0 return min((step - flat) * uprate, uphold) def _reward_metrics(self, rewards: dict[str, torch.Tensor]) -> dict[str, LogScalar]: metrics: dict[str, LogScalar] = {} for key, value in rewards.items(): gathered = self._gather_tensor(value.detach().float()) metrics[f"reward/{key}"] = gathered.mean() return metrics def _reward_diagnostic_metrics( self, sample_items: list[dict[str, Any]], rewards: dict[str, torch.Tensor], ) -> dict[str, LogScalar]: avg = rewards["avg"].detach().float() local_prompts = [prompt for item in sample_items for prompt in item["prompts"]] gathered_rewards = self._gather_tensor(avg) gathered_prompts = self._gather_prompts(local_prompts) if len(gathered_prompts) != int(gathered_rewards.shape[0]): raise RuntimeError("Gathered prompt count does not match gathered rewards") prompt_to_indices: dict[str, list[int]] = defaultdict(list) for idx, prompt in enumerate(gathered_prompts): prompt_to_indices[prompt].append(idx) self._trained_prompt_hashes.add(hash(prompt)) if not prompt_to_indices: return {} group_sizes = [] group_stds = [] grouped_rewards = [] for indices in prompt_to_indices.values(): index_tensor = torch.tensor(indices, device=gathered_rewards.device, dtype=torch.long) group_rewards = gathered_rewards[index_tensor].float() grouped_rewards.append(group_rewards) group_sizes.append(float(group_rewards.numel())) group_stds.append(group_rewards.std(unbiased=False)) std_tensor = torch.stack(group_stds) metrics: dict[str, LogScalar] = { "group_size": sum(group_sizes) / max(1, len(group_sizes)), "trained_prompt_num": float(len(self._trained_prompt_hashes)), "zero_std_ratio": (std_tensor == 0).float().mean(), "reward_std_mean": std_tensor.mean(), } for top_percentage in (100, 75, 50, 25, 10): metrics[f"mean_reward_{top_percentage}"] = self._mean_of_top_group_rewards( grouped_rewards, top_percentage, ) return metrics @staticmethod def _mean_of_top_group_rewards( grouped_rewards: list[torch.Tensor], top_percentage: int, ) -> torch.Tensor: per_prompt_means = [] for rewards in grouped_rewards: if rewards.numel() == 0: continue if top_percentage == 100: per_prompt_means.append(rewards.mean()) continue threshold = torch.quantile( rewards.float(), (100 - float(top_percentage)) / 100.0, ) top_rewards = rewards[rewards >= threshold] if top_rewards.numel() > 0: per_prompt_means.append(top_rewards.mean()) if not per_prompt_means: device = grouped_rewards[0].device if grouped_rewards else torch.device("cpu") return torch.zeros((), device=device) return torch.stack(per_prompt_means).mean() @staticmethod def _gather_tensor(tensor: torch.Tensor) -> torch.Tensor: if not dist.is_available() or not dist.is_initialized(): return tensor.detach() gathered = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())] dist.all_gather(gathered, tensor.contiguous()) return torch.cat(gathered, dim=0).detach() @staticmethod def _mean_scalar_across_ranks(value: torch.Tensor) -> torch.Tensor: if not dist.is_available() or not dist.is_initialized(): return value.detach() reduced = value.detach().clone() dist.all_reduce(reduced, op=dist.ReduceOp.AVG) return reduced @staticmethod def _gather_prompts(prompts: list[str]) -> list[str]: if not dist.is_available() or not dist.is_initialized(): return list(prompts) gathered: list[list[str] | None] = [None for _ in range(dist.get_world_size())] dist.all_gather_object(gathered, list(prompts)) flat: list[str] = [] for rank_prompts in gathered: if rank_prompts is None: continue flat.extend(rank_prompts) return flat @staticmethod def _gather_objects(items: list[Any]) -> list[Any]: if not dist.is_available() or not dist.is_initialized(): return list(items) gathered: list[list[Any] | None] = [None for _ in range(dist.get_world_size())] dist.all_gather_object(gathered, list(items)) flat: list[Any] = [] for rank_items in gathered: if rank_items is None: continue flat.extend(rank_items) return flat def _read_int(self, key: str, default: int) -> int: value = get_optional_int(self.method_config, key, where=f"method.{key}") return int(default if value is None else value) def _read_float(self, key: str, default: float) -> float: value = get_optional_float(self.method_config, key, where=f"method.{key}") return float(default if value is None else value) def _parse_sampling_config(self) -> SamplingConfig: raw = self.method_config.get("sampling", None) if isinstance(raw, dict): raw = dict(raw) if raw.get("flow_shift", None) in (None, "inherit"): pipeline_config = getattr(self.training_config, "pipeline_config", None) flow_shift = getattr(pipeline_config, "flow_shift", None) if flow_shift is not None: raw["flow_shift"] = float(flow_shift) return SamplingConfig.from_mapping(raw) def _parse_validation_sampling_config(self) -> SamplingConfig: raw = dict(self.method_config.get("sampling", {}) or {}) validation_sampling = self._validation_config.sampling if validation_sampling is not None: raw.update(validation_sampling) if validation_sampling is None or "num_steps" not in validation_sampling: raw["num_steps"] = self._validation_config.num_steps if raw.get("flow_shift") in (None, "inherit"): pipeline_config = getattr(self.training_config, "pipeline_config", None) flow_shift = getattr(pipeline_config, "flow_shift", None) if flow_shift is not None: raw["flow_shift"] = float(flow_shift) return SamplingConfig.from_mapping(raw) def _parse_ema_config(self) -> dict[str, bool | float | int]: raw = self.method_config.get("ema", None) config: dict[str, bool | float | int] = { "enabled": True, "decay": 0.9, "update_after_step": 0, "validation": True, } if raw is None: return config if isinstance(raw, bool | str | int): config["enabled"] = self._coerce_bool(raw, where="method.ema") return config if not isinstance(raw, dict): raise ValueError("method.ema must be a bool or mapping") allowed = set(config) unknown = sorted(set(raw) - allowed) if unknown: raise ValueError(f"Unsupported method.ema key(s): {unknown}") if "enabled" in raw: config["enabled"] = self._coerce_bool(raw["enabled"], where="method.ema.enabled") if "decay" in raw: decay = float(raw["decay"]) if not 0.0 <= decay <= 1.0: raise ValueError("method.ema.decay must be in [0, 1]") config["decay"] = decay if "update_after_step" in raw: update_after_step = int(raw["update_after_step"]) if update_after_step < 0: raise ValueError("method.ema.update_after_step must be >= 0") config["update_after_step"] = update_after_step if "validation" in raw: config["validation"] = self._coerce_bool(raw["validation"], where="method.ema.validation") return config @staticmethod def _coerce_bool(value: Any, *, where: str) -> bool: if isinstance(value, bool): return value if isinstance(value, int): return bool(value) if isinstance(value, str): lowered = value.strip().lower() if lowered in {"1", "true", "yes", "on"}: return True if lowered in {"0", "false", "no", "off"}: return False raise ValueError(f"{where} must be a boolean")