""" Sana post-training (BON + preview rollout) using diffusers SanaPipeline. Structure mirrors train_sd3 (Sol-RL) with BON + preview rollout. """ import os import sys from collections import defaultdict _rank = int(os.environ.get("RANK", 0)) _cache_root = os.environ.get("CACHE_ROOT", os.path.expanduser("~/.cache/sol_rl")) os.environ.setdefault("TRITON_CACHE_DIR", f"{_cache_root}/triton/rank_{_rank}") os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{_cache_root}/torchinductor/rank_{_rank}") os.environ.setdefault("TORCHINDUCTOR_FX_GRAPH_CACHE", "1") import logging import random import time from concurrent import futures import numpy as np import torch import torch.distributed as dist import tqdm import wandb from absl import app, flags from diffusers import SanaPipeline from ml_collections import config_flags from peft import LoraConfig, PeftModel, get_peft_model from torch.cuda.amp import GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) import pyrallis from train_utils import ( _HAS_TE, DistributedTimeLogger, build_datasets_and_loaders, calculate_zero_std_ratio, cleanup_distributed, collate_dict_items, extract_prompt_reward_group, filter_by_indices, find_resume_candidates, gather_tensor_to_all, is_main_process, log_rollout_images, replace_linear_with_te, resume_from_checkpoint, return_decay, save_ckpt, save_step_reward_groups, select_indices_by_mode, set_seed, setup_distributed, slice_prompt_metadata, sync_lora_to_inference, unwrap_compiled, wrap_forward_with_fp8, ) import diffusion.model.nets.sana_multi_scale # noqa: F401 import diffusion.post_training.rewards from diffusion.model.builder import MODELS from diffusion.post_training.diffusers_patch.pipeline_with_logprob import pipeline_with_logprob_sana from diffusion.post_training.diffusers_patch.text_encode import encode_sana_prompt from diffusion.post_training.ema import EMAModuleWrapper from diffusion.post_training.stat_tracking import PerPromptStatTracker from diffusion.utils.config import SanaConfig, model_init_config from tools.download import find_model tqdm = tqdm.tqdm FLAGS = flags.FLAGS config_flags.DEFINE_config_file( "config", "configs/sol_rl/sana.py", "Training configuration.", ) flags.DEFINE_string("native_config", "", "Optional override for native model YAML config path") logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") TEXT_ENCODER_MAX_SEQ_LEN = 300 TOKENIZER_MAX_LENGTH = 300 WANDB_MAX_LOG_IMAGES = 12 def compute_text_embeddings(prompts, pipeline, max_sequence_length, device): with torch.no_grad(): return encode_sana_prompt( pipeline, prompts, max_sequence_length=max_sequence_length, device=device, negative_prompt="", do_classifier_free_guidance=True, ) def _resolve_native_checkpoint_source(config, native_cfg): native_model_path = str(getattr(config, "native_model_path", "") or "").strip() native_model_source = str(getattr(config, "native_model_source", "") or "").strip() if native_model_path: native_cfg.model.load_from = native_model_path if os.path.isfile(native_model_path): return native_model_path if native_model_source: logger.info( "[INIT] Native checkpoint missing at %s; falling back to %s", native_model_path, native_model_source ) return native_model_source return native_cfg.model.load_from def _prepare_latents_from_seeds(seed_list, num_channels, latent_h, latent_w, device, dtype): latents = [] for seed in seed_list: g = torch.Generator(device=device).manual_seed(int(seed)) latents.append(torch.randn(1, num_channels, latent_h, latent_w, device=device, dtype=dtype, generator=g)) return torch.cat(latents, dim=0) def _select_inference_transformer(mode, inference_models, transformer_ddp, peft_transformer): """Return the transformer to use for inference based on *mode*. mode: "compile_nvfp4" | "compile" | "peft" """ if mode == "peft": transformer_ddp.module.set_adapter("old") return peft_transformer return inference_models[mode] def _rollout_for_one_prompt( rollout_transformer, vae, num_channels, latent_size, reward_fn, executor, prompt_text, prompt_meta, prompt_embed_single, prompt_mask_single, neg_prompt_embed_single, neg_prompt_mask_single, prompt_token_ids_single, config, device, inference_models=None, transformer_ddp=None, peft_transformer=None, ): preview_step = int(getattr(config, "preview_step", 0)) full_steps = int(getattr(config, "rollout_sample_num_steps", config.sample.num_steps)) draft_total = int(config.sample.per_prompt_iter_num) * int(config.sample.rollout_batch_size) full_rollout_num = max( 1, min(int(getattr(config.sample, "full_rollout_num", config.sample.best_of_n)), draft_total) ) full_chunks = int(config.sample.rollout_batch_size) solver = str(getattr(config.sample, "solver", "flow")) noise_level = float(getattr(config.sample, "noise_level", 0.7)) deterministic = True seed_pool, draft_reward_pool = [], [] prompt_samples = [] final_images = final_prompts = None preview_model_key = str(getattr(config, "preview_model", "peft")) fullrollout_model_key = str(getattr(config, "fullrollout_model", "peft")) _can_swap = inference_models is not None and transformer_ddp is not None and peft_transformer is not None if preview_step > 0: if _can_swap: current_transformer = _select_inference_transformer( preview_model_key, inference_models, transformer_ddp, peft_transformer ) else: current_transformer = rollout_transformer with torch.no_grad(): for _ in range(config.sample.per_prompt_iter_num): bs = int(config.sample.rollout_batch_size) p_emb = prompt_embed_single.repeat(bs, 1, 1) p_mask = prompt_mask_single.repeat(bs, 1) neg_emb = neg_prompt_embed_single.repeat(bs, 1, 1) neg_mask = neg_prompt_mask_single.repeat(bs, 1) seed_list = torch.randint(0, 2**31 - 1, (bs,), device="cpu").tolist() init_latents = _prepare_latents_from_seeds( seed_list, num_channels, latent_size, latent_size, device, p_emb.dtype ) images, _, _ = pipeline_with_logprob_sana( current_transformer, vae, latents=init_latents, prompt_embeds=p_emb, prompt_attention_mask=p_mask, negative_prompt_embeds=neg_emb, negative_prompt_attention_mask=neg_mask, num_inference_steps=preview_step, guidance_scale=config.rollout_sample_guidance_scale, noise_level=noise_level, deterministic=deterministic, solver=solver, ) rewards, _ = reward_fn(images, [prompt_text] * bs, [prompt_meta] * bs, only_strict=True) seed_pool.extend(int(s) for s in seed_list) draft_reward_pool.extend(torch.as_tensor(rewards["avg"], device=device).float().detach().cpu().tolist()) draft_rewards = torch.as_tensor(draft_reward_pool, device=device).float() stage1_idx = select_indices_by_mode( draft_rewards, full_rollout_num, getattr(config.sample, "stage1_select_mode", "best_worst") ) selected_seeds = [seed_pool[int(i)] for i in stage1_idx.cpu().tolist()] if _can_swap and fullrollout_model_key != preview_model_key: current_transformer = _select_inference_transformer( fullrollout_model_key, inference_models, transformer_ddp, peft_transformer ) for start in range(0, len(selected_seeds), full_chunks): chunk_seeds = selected_seeds[start : start + full_chunks] bs = len(chunk_seeds) p_emb = prompt_embed_single.repeat(bs, 1, 1) p_mask = prompt_mask_single.repeat(bs, 1) neg_emb = neg_prompt_embed_single.repeat(bs, 1, 1) neg_mask = neg_prompt_mask_single.repeat(bs, 1) init_latents = _prepare_latents_from_seeds( chunk_seeds, num_channels, latent_size, latent_size, device, p_emb.dtype ) with torch.no_grad(): images, all_latents, step_sigmas = pipeline_with_logprob_sana( current_transformer, vae, latents=init_latents, prompt_embeds=p_emb, prompt_attention_mask=p_mask, negative_prompt_embeds=neg_emb, negative_prompt_attention_mask=neg_mask, num_inference_steps=full_steps, guidance_scale=config.rollout_sample_guidance_scale, noise_level=noise_level, deterministic=deterministic, solver=solver, ) timesteps = step_sigmas.unsqueeze(0).repeat(bs, 1) latents = torch.stack(all_latents, dim=1) rewards_future = executor.submit(reward_fn, images, [prompt_text] * bs, [prompt_meta] * bs, True) time.sleep(0) prompt_samples.append( { "prompt_ids": prompt_token_ids_single.repeat(bs, 1), "prompt_embeds": p_emb, "prompt_attention_mask": p_mask, "timesteps": timesteps, "next_timesteps": torch.cat([timesteps[:, 1:], torch.zeros_like(timesteps[:, :1])], dim=1), "latents_clean": latents[:, -1], "rewards_future": rewards_future, } ) final_images, final_prompts = images, [prompt_text] * bs else: for _ in range(config.sample.per_prompt_iter_num): bs = int(config.sample.rollout_batch_size) p_emb = prompt_embed_single.repeat(bs, 1, 1) p_mask = prompt_mask_single.repeat(bs, 1) neg_emb = neg_prompt_embed_single.repeat(bs, 1, 1) neg_mask = neg_prompt_mask_single.repeat(bs, 1) seed_list = torch.randint(0, 2**31 - 1, (bs,), device="cpu").tolist() init_latents = _prepare_latents_from_seeds( seed_list, num_channels, latent_size, latent_size, device, p_emb.dtype ) with torch.no_grad(): images, all_latents, step_sigmas = pipeline_with_logprob_sana( rollout_transformer, vae, latents=init_latents, prompt_embeds=p_emb, prompt_attention_mask=p_mask, negative_prompt_embeds=neg_emb, negative_prompt_attention_mask=neg_mask, num_inference_steps=full_steps, guidance_scale=config.rollout_sample_guidance_scale, noise_level=noise_level, deterministic=deterministic, solver=solver, ) timesteps = step_sigmas.unsqueeze(0).repeat(bs, 1) latents = torch.stack(all_latents, dim=1) rewards_future = executor.submit(reward_fn, images, [prompt_text] * bs, [prompt_meta] * bs, True) time.sleep(0) prompt_samples.append( { "prompt_ids": prompt_token_ids_single.repeat(bs, 1), "prompt_embeds": p_emb, "prompt_attention_mask": p_mask, "timesteps": timesteps, "next_timesteps": torch.cat([timesteps[:, 1:], torch.zeros_like(timesteps[:, :1])], dim=1), "latents_clean": latents[:, -1], "rewards_future": rewards_future, } ) final_images, final_prompts = images, [prompt_text] * bs for item in prompt_samples: rewards, _ = item["rewards_future"].result() item["rewards"] = {k: torch.as_tensor(v, device=device).float() for k, v in rewards.items()} del item["rewards_future"] collated = collate_dict_items(prompt_samples) keep = select_indices_by_mode( collated["rewards"]["avg"], config.sample.best_of_n, getattr(config.sample, "stage2_select_mode", "best_worst"), ) collated = filter_by_indices(collated, keep) return collated, final_images, final_prompts def eval_fn( pipeline, eval_transformer, vae, num_channels, latent_size, test_dataloader, config, device, rank, world_size, global_step, reward_fn, executor, mixed_precision_dtype, ema, transformer_trainable_parameters, ): set_seed(config.seed + 1_000_000, rank) if config.train.ema and ema is not None: ema.copy_ema_to(transformer_trainable_parameters, store_temp=True) eval_transformer.eval() all_rewards = defaultdict(list) test_sampler = ( DistributedSampler(test_dataloader.dataset, num_replicas=world_size, rank=rank, shuffle=False) if world_size > 1 else None ) eval_loader = DataLoader( test_dataloader.dataset, batch_size=config.sample.test_batch_size, sampler=test_sampler, collate_fn=test_dataloader.collate_fn, num_workers=0, ) solver = str(getattr(config.sample, "solver", "flow")) for test_batch in tqdm(eval_loader, desc="Eval:", disable=not is_main_process(rank)): prompts, prompt_metadata = test_batch with torch.no_grad(): prompt_embeds, prompt_mask, neg_embeds, neg_mask = compute_text_embeddings( prompts, pipeline, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device, ) images, _, _ = pipeline_with_logprob_sana( eval_transformer, vae, num_channels=num_channels, latent_size=latent_size, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt_mask, negative_prompt_embeds=neg_embeds, negative_prompt_attention_mask=neg_mask, num_inference_steps=config.sample.eval_num_steps, guidance_scale=config.eval_sample_guidance_scale, solver=solver, ) rewards_future = executor.submit(reward_fn, images, prompts, prompt_metadata, only_strict=False) time.sleep(0) rewards, _ = rewards_future.result() for key, value in rewards.items(): all_rewards[key].extend(value) final_rewards = {k: np.array([x.cpu() if torch.is_tensor(x) else x for x in v]) for k, v in all_rewards.items()} if is_main_process(rank) and final_rewards: wandb.log( {**{f"eval_reward_{k}": np.mean(v[v != -10]) for k, v in final_rewards.items()}}, commit=False, ) for k, v in final_rewards.items(): logger.info("eval_reward_%s: %.4f", k, np.mean(v[v != -10])) if config.train.ema and ema is not None: ema.copy_temp_to(transformer_trainable_parameters) if world_size > 1: dist.barrier() def main(_): config = FLAGS.config if FLAGS.native_config: config.native_config = FLAGS.native_config # cuDNN SDPA backward graph fails on Blackwell (sm_100); fall back to flash/math torch.backends.cuda.enable_cudnn_sdp(False) start_time = time.time() rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ["LOCAL_RANK"]) setup_distributed(rank, local_rank, world_size) device = torch.device(f"cuda:{local_rank}") if is_main_process(rank): log_dir = os.path.join(config.logdir, config.run_name) os.makedirs(log_dir, exist_ok=True) wandb.init( project="sol-rl", name=config.run_name, config=config.to_dict(), resume="allow", dir=log_dir, id=config.run_name, ) wandb.define_metric("global_step") wandb.define_metric("*", step_metric="global_step") logger.info(f"\n{config}") set_seed(config.seed, rank) mixed_precision_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16}.get(config.mixed_precision) enable_amp = mixed_precision_dtype is not None scaler = GradScaler(enabled=enable_amp and mixed_precision_dtype == torch.float16) # Build native transformer from the Sana YAML config. with open(config.native_config, encoding="utf-8") as _f: native_cfg = pyrallis.load(SanaConfig, _f) ckpt_source = _resolve_native_checkpoint_source(config, native_cfg) # Keep the diffusers text encoder + VAE, but replace the native transformer. pipeline = SanaPipeline.from_pretrained( "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", torch_dtype=torch.bfloat16 ) pipeline.vae.requires_grad_(False) pipeline.text_encoder.requires_grad_(False) pipeline.set_progress_bar_config(disable=True) text_encoder_dtype = mixed_precision_dtype if enable_amp else torch.float32 pipeline.vae.to(device, dtype=torch.bfloat16) pipeline.text_encoder.to(device, dtype=text_encoder_dtype) vae = pipeline.vae del pipeline.transformer torch.cuda.empty_cache() latent_size = config.resolution // 32 if getattr(native_cfg, "work_dir", None): os.makedirs(native_cfg.work_dir, exist_ok=True) model_kwargs = model_init_config(native_cfg, latent_size=latent_size) transformer = MODELS.build(dict(type=native_cfg.model.model), default_args=model_kwargs) transformer.to(device, dtype=torch.bfloat16) logger.info(f"[INIT] Loading native checkpoint from {ckpt_source}") ckpt = find_model(ckpt_source) if isinstance(ckpt, dict): if "state_dict" in ckpt: state = ckpt["state_dict"] elif "model" in ckpt: state = ckpt["model"] else: state = ckpt else: state = ckpt missing, unexpected = transformer.load_state_dict(state, strict=False) if missing: logger.warning(f"[weights] missing {len(missing)} keys (showing up to 10): {missing[:10]}") for blk in transformer.blocks: if hasattr(blk.attn, "eps"): blk.attn.eps = 1e-15 num_channels = native_cfg.vae.vae_latent_dim transformer.requires_grad_(not config.use_lora) # Create optional inference-only copies for compiled and/or NVFP4 rollout modes. compile_mode = str(getattr(config, "compile_mode", "max-autotune-no-cudagraphs")) preview_step = int(getattr(config, "preview_step", 0)) preview_model_key = str(getattr(config, "preview_model", "peft")) fullrollout_model_key = str(getattr(config, "fullrollout_model", "peft")) needed_model_types = set() if preview_step > 0: if preview_model_key != "peft": needed_model_types.add(preview_model_key) if fullrollout_model_key != "peft": needed_model_types.add(fullrollout_model_key) else: if fullrollout_model_key != "peft": needed_model_types.add(fullrollout_model_key) inference_models = {} nvfp4_skip_modules = list(getattr(config, "nvfp4_skip_modules", [])) nvfp4_min_dim = int(getattr(config, "nvfp4_min_dim", 0)) for mtype in sorted(needed_model_types): logger.info(f"[INIT] Creating inference model: {mtype!r} ...") m = MODELS.build(dict(type=native_cfg.model.model), default_args=model_kwargs) m.to(device, dtype=torch.bfloat16) m.load_state_dict(state, strict=False) for blk in m.blocks: if hasattr(blk.attn, "eps"): blk.attn.eps = 1e-15 m.eval() m.requires_grad_(False) if "nvfp4" in mtype: if not _HAS_TE: raise RuntimeError(f"{mtype!r} requires transformer_engine") n_rep, n_skip, rep_d, skip_d = replace_linear_with_te( m, skip_modules=nvfp4_skip_modules, min_dim=nvfp4_min_dim ) logger.info(f"[NVFP4] {mtype}: replaced {n_rep} -> te.Linear, skipped {n_skip}") if is_main_process(rank): os.makedirs(config.save_dir, exist_ok=True) report_path = os.path.join(config.save_dir, f"nvfp4_report_{mtype}.txt") with open(report_path, "w") as f: f.write(f"NVFP4 Quantization Report: {mtype}\n") f.write(f"skip_modules={nvfp4_skip_modules}, min_dim={nvfp4_min_dim}\n") f.write(f"replaced={n_rep}, skipped={n_skip}\n\n") f.write(f"REPLACED ({len(rep_d)}):\n") for fqn, inf, outf, b, _ in rep_d: f.write(f" {fqn:60s} {inf:>5d} -> {outf:>5d} bias={b}\n") f.write(f"\nSKIPPED ({len(skip_d)}):\n") for fqn, inf, outf, b, r in skip_d: f.write(f" {fqn:60s} {inf:>5d} -> {outf:>5d} bias={b} reason={r}\n") logger.info(f"[NVFP4] Report saved to {report_path}") wrap_forward_with_fp8(m) if world_size > 1: dist.barrier() logger.info(f"[COMPILE] torch.compile(mode={compile_mode!r}) on {mtype!r} ...") m = torch.compile(m, mode=compile_mode) inference_models[mtype] = m logger.info(f"[INIT] {mtype!r} ready") if inference_models: logger.info(f"[INIT] Inference models created: {list(inference_models.keys())}") else: logger.info("[INIT] No inference models needed, using PEFT model for all inference") if config.use_lora: init_lora_weights = getattr(config.train, "lora_init_mode", config.train.lora_init_weights) lora_cfg = LoraConfig( r=config.train.lora_rank, lora_alpha=config.train.lora_alpha, init_lora_weights=init_lora_weights, target_modules=list(config.train.lora_target_modules), ) if config.train.lora_path: transformer = PeftModel.from_pretrained(transformer, config.train.lora_path) transformer.set_adapter("default") else: transformer = get_peft_model(transformer, lora_cfg) transformer.add_adapter("old", lora_cfg) transformer.set_adapter("default") peft_transformer = transformer transformer_ddp = DDP(transformer, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False) transformer_ddp.module.set_adapter("default") transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, transformer_ddp.module.parameters())) transformer_ddp.module.set_adapter("old") old_transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, transformer_ddp.module.parameters())) transformer_ddp.module.set_adapter("default") if config.allow_tf32: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True optimizer = torch.optim.AdamW( transformer_trainable_parameters, lr=config.train.learning_rate, betas=(config.train.adam_beta1, config.train.adam_beta2), weight_decay=config.train.adam_weight_decay, eps=config.train.adam_epsilon, ) ema = None if config.train.ema: ema = EMAModuleWrapper( transformer_trainable_parameters, decay=getattr(config.train, "ema_decay", 0.9), update_step_interval=getattr(config.train, "ema_update_step_interval", 1), device=device, ) _, train_dataloader, train_sampler, _, test_dataloader = build_datasets_and_loaders(config, world_size, rank) train_iter = iter(train_dataloader) stat_tracker = PerPromptStatTracker(config.global_std) if config.per_prompt_stat_tracking else None executor = futures.ThreadPoolExecutor(max_workers=8) reward_fn = getattr(diffusion.post_training.rewards, "multi_score")(device, config.reward_fn) eval_reward_fn = getattr(diffusion.post_training.rewards, "multi_score")(device, config.reward_fn) timestep_clip = getattr(config, "timestep_clip", None) if timestep_clip is not None: _ts_start, _ts_end = int(timestep_clip[0]), int(timestep_clip[1]) num_train_timesteps = _ts_end - _ts_start else: _ts_start = _ts_end = None num_train_timesteps = int(config.rollout_sample_num_steps * config.train.timestep_fraction) first_epoch = 0 global_step = 0 candidates = find_resume_candidates(config) global_step, _resumed = resume_from_checkpoint( candidates, transformer_ddp.module, ema, optimizer, scaler, device, ) first_epoch = global_step if first_epoch == 0: for src_p, tgt_p in zip(transformer_trainable_parameters, old_transformer_trainable_parameters, strict=True): tgt_p.data.copy_(src_p.detach().data) # The rollout path reads from the "old" adapter weights. for mtype, inf_model in inference_models.items(): sync_lora_to_inference(transformer_ddp.module, unwrap_compiled(inf_model), adapter_name="old") optimizer.zero_grad() time_logger = DistributedTimeLogger(device) if global_step != 0: for _ in range(global_step): next(train_iter) if world_size > 1: dist.barrier() for epoch in range(first_epoch, config.num_epochs): time_logger.start("total_time") if hasattr(train_sampler, "set_epoch"): train_sampler.set_epoch(epoch) if epoch % config.save_freq == 0 and not config.debug: save_ckpt(config.save_dir, transformer_ddp, global_step, rank, ema, config, optimizer, scaler) time_logger.start("eval_time") if epoch % config.eval_freq == 0 and not config.debug: torch.cuda.empty_cache() py_rng = random.getstate() np_rng = np.random.get_state() torch_rng = torch.random.get_rng_state() cuda_rng = torch.cuda.get_rng_state_all() eval_fn( pipeline, peft_transformer, vae, num_channels, latent_size, test_dataloader, config, device, rank, world_size, global_step, eval_reward_fn, executor, mixed_precision_dtype, ema, transformer_trainable_parameters, ) random.setstate(py_rng) np.random.set_state(np_rng) torch.random.set_rng_state(torch_rng) torch.cuda.set_rng_state_all(cuda_rng) time_logger.end("eval_time") time_logger.start("rollout_time") peft_transformer.eval() prompts, prompt_metadata = next(train_iter) time_logger.start("text_tokenizer_time") with torch.no_grad(): prompt_embeds, prompt_masks, neg_embeds, neg_masks = compute_text_embeddings( prompts, pipeline, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device, ) txt_tokens = pipeline.tokenizer( prompts, max_length=TOKENIZER_MAX_LENGTH, padding="max_length", truncation=True, return_tensors="pt", ) time_logger.end("text_tokenizer_time") prompt_ids_all = txt_tokens.input_ids.to(device) prompt_embeds_list = [prompt_embeds[i : i + 1] for i in range(len(prompts))] prompt_masks_list = [prompt_masks[i : i + 1] for i in range(len(prompts))] neg_embeds_single = neg_embeds[0:1] neg_masks_single = neg_masks[0:1] prompt_wise_samples = [] step_prompt_reward_groups = [] images_for_log = prompts_for_log = rewards_for_log = None if preview_step == 0: if fullrollout_model_key != "peft" and fullrollout_model_key in inference_models: default_rollout_transformer = inference_models[fullrollout_model_key] else: transformer_ddp.module.set_adapter("old") default_rollout_transformer = peft_transformer else: default_rollout_transformer = peft_transformer for prompt_idx in tqdm( range(config.sample.per_gpu_to_process_prompts), desc=f"Epoch {epoch}: rollout", disable=not is_main_process(rank), dynamic_ncols=True, ): collated_prompt_samples, final_images, final_prompts = _rollout_for_one_prompt( rollout_transformer=default_rollout_transformer, vae=vae, num_channels=num_channels, latent_size=latent_size, reward_fn=reward_fn, executor=executor, prompt_text=prompts[prompt_idx], prompt_meta=prompt_metadata[prompt_idx], prompt_embed_single=prompt_embeds_list[prompt_idx], prompt_mask_single=prompt_masks_list[prompt_idx], neg_prompt_embed_single=neg_embeds_single, neg_prompt_mask_single=neg_masks_single, prompt_token_ids_single=prompt_ids_all[prompt_idx : prompt_idx + 1], config=config, device=device, inference_models=inference_models, transformer_ddp=transformer_ddp, peft_transformer=peft_transformer, ) prompt_wise_samples.append(collated_prompt_samples) prompt_meta_i = slice_prompt_metadata(prompt_metadata, prompt_idx) step_prompt_reward_groups.append( extract_prompt_reward_group(prompt_idx, prompts[prompt_idx], prompt_meta_i, [collated_prompt_samples]) ) images_for_log, prompts_for_log = final_images, final_prompts rewards_for_log = collated_prompt_samples["rewards"]["avg"] transformer_ddp.module.set_adapter("default") save_step_reward_groups( config=config, global_step=global_step, epoch=epoch, rank=rank, world_size=world_size, prompt_reward_groups=step_prompt_reward_groups, ) collated_samples = collate_dict_items(prompt_wise_samples) log_rollout_images(images_for_log, prompts_for_log, rewards_for_log, config, global_step, rank) collated_samples["rewards"]["avg"] = ( collated_samples["rewards"]["avg"].unsqueeze(1).repeat(1, num_train_timesteps) ) gathered_rewards_dict = { k: gather_tensor_to_all(v, world_size).numpy() for k, v in collated_samples["rewards"].items() } if is_main_process(rank): r2l = ( gathered_rewards_dict["avg"] .reshape(world_size * config.sample.per_gpu_to_process_prompts, -1, num_train_timesteps) .mean(axis=-1) ) wandb.log( { "epoch": epoch, "reward/mean": r2l.mean(), "reward/max": r2l.max(axis=1).mean(), "reward/min": r2l.min(axis=1).mean(), "reward/range": r2l.max(axis=1).mean() - r2l.min(axis=1).mean(), }, commit=False, ) prompt_ids_global = gather_tensor_to_all(collated_samples["prompt_ids"], world_size) prompts_decoded = pipeline.tokenizer.batch_decode(prompt_ids_global.cpu().numpy(), skip_special_tokens=True) if stat_tracker is not None: advantages = stat_tracker.update(prompts_decoded, gathered_rewards_dict["avg"]) if is_main_process(rank): gs, tp = stat_tracker.get_stats() zsr, rsm = calculate_zero_std_ratio(prompts_decoded, gathered_rewards_dict) wandb.log( { "group_size": gs, "trained_prompt_num": tp, "zero_std_ratio": zsr, "reward_std_mean": rsm, "mean_reward_100": stat_tracker.get_mean_of_top_rewards(100), "mean_reward_50": stat_tracker.get_mean_of_top_rewards(50), }, commit=False, ) stat_tracker.clear() else: avg = gathered_rewards_dict["avg"] advantages = (avg - avg.mean()) / (avg.std() + 1e-4) samples_per_gpu = collated_samples["timesteps"].shape[0] if advantages.ndim == 1: advantages = advantages[:, None] collated_samples["advantages"] = torch.from_numpy(advantages.reshape(world_size, samples_per_gpu, -1)[rank]).to( device ) del collated_samples["rewards"] del collated_samples["prompt_ids"] time_logger.end("rollout_time") total_batch_size_filtered, num_timesteps_filtered = collated_samples["timesteps"].shape time_logger.start("train_time") effective_grad_accum_steps = config.train.gradient_accumulation_steps * num_train_timesteps current_accumulated_steps = 0 gradient_update_times = 0 for inner_epoch in range(config.train.num_inner_epochs): perm = torch.randperm(total_batch_size_filtered, device=device) shuffled = {k: v[perm] for k, v in collated_samples.items()} perms_time = torch.stack( [torch.randperm(num_timesteps_filtered, device=device) for _ in range(total_batch_size_filtered)] ) for key in ["timesteps", "next_timesteps"]: shuffled[key] = shuffled[key][ torch.arange(total_batch_size_filtered, device=device)[:, None], perms_time ] batches = [] for bi in range(config.train.n_batch_per_epoch): s, e = bi * config.train.batch_size, (bi + 1) * config.train.batch_size batches.append({k: v[s:e] for k, v in shuffled.items()}) info_accumulated = defaultdict(list) for train_batch in tqdm( batches, desc=f"Epoch {epoch}.{inner_epoch}: train", disable=not is_main_process(rank), dynamic_ncols=True, ): embeds = train_batch["prompt_embeds"] embeds_mask = train_batch["prompt_attention_mask"] embeds_4d = embeds.unsqueeze(1) mask_4d = embeds_mask.unsqueeze(1).unsqueeze(1).to(torch.int16) if embeds_mask is not None else None for j_idx in range(num_train_timesteps): x0 = train_batch["latents_clean"] sigma = train_batch["timesteps"][:, j_idx] sigma_expanded = sigma.view(-1, *([1] * (len(x0.shape) - 1))) noise = torch.randn_like(x0.float()) xt = ((1 - sigma_expanded) * x0 + sigma_expanded * noise).to(torch.bfloat16) transformer_ddp.module.set_adapter("old") with torch.no_grad(): old_prediction = transformer_ddp(xt, sigma, embeds_4d, mask=mask_4d).detach() transformer_ddp.module.set_adapter("default") forward_prediction = transformer_ddp(xt, sigma, embeds_4d, mask=mask_4d) with torch.no_grad(): with transformer_ddp.module.disable_adapter(): ref_forward_prediction = transformer_ddp(xt, sigma, embeds_4d, mask=mask_4d) transformer_ddp.module.set_adapter("default") loss_terms = {} advantages_clip = torch.clamp( train_batch["advantages"][:, j_idx], -config.train.adv_clip_max, config.train.adv_clip_max ) if hasattr(config.train, "adv_mode"): if config.train.adv_mode == "positive_only": advantages_clip = torch.clamp(advantages_clip, 0, config.train.adv_clip_max) elif config.train.adv_mode == "negative_only": advantages_clip = torch.clamp(advantages_clip, -config.train.adv_clip_max, 0) elif config.train.adv_mode == "one_only": advantages_clip = torch.where( advantages_clip > 0, torch.ones_like(advantages_clip), torch.zeros_like(advantages_clip) ) elif config.train.adv_mode == "binary": advantages_clip = torch.sign(advantages_clip) r = torch.clamp((advantages_clip / config.train.adv_clip_max) / 2.0 + 0.5, 0, 1) positive_prediction = config.beta * forward_prediction + (1 - config.beta) * old_prediction.detach() implicit_negative_prediction = ( 1.0 + config.beta ) * old_prediction.detach() - config.beta * forward_prediction x0_prediction = xt - sigma_expanded * positive_prediction with torch.no_grad(): weight_factor = ( torch.abs(x0_prediction.double() - x0.double()) .mean(dim=tuple(range(1, x0.ndim)), keepdim=True) .clip(min=1e-5) ) positive_loss = ((x0_prediction - x0) ** 2 / weight_factor).mean(dim=tuple(range(1, x0.ndim))) negative_x0_prediction = xt - sigma_expanded * implicit_negative_prediction with torch.no_grad(): neg_wf = ( torch.abs(negative_x0_prediction.double() - x0.double()) .mean(dim=tuple(range(1, x0.ndim)), keepdim=True) .clip(min=1e-5) ) negative_loss = ((negative_x0_prediction - x0) ** 2 / neg_wf).mean(dim=tuple(range(1, x0.ndim))) ori_policy_loss = r * positive_loss / config.beta + (1.0 - r) * negative_loss / config.beta policy_loss = (ori_policy_loss * config.train.adv_clip_max).mean() loss = policy_loss loss_terms["policy_loss"] = policy_loss.detach() loss_terms["unweighted_policy_loss"] = ori_policy_loss.mean().detach() kl_div_loss = ((forward_prediction - ref_forward_prediction) ** 2).mean( dim=tuple(range(1, x0.ndim)) ) loss += config.train.beta * torch.mean(kl_div_loss) loss_terms["kl_div_loss"] = torch.mean(kl_div_loss).detach() loss_terms["kl_div"] = loss_terms["kl_div_loss"] loss_terms["old_kl_div"] = torch.mean( ((old_prediction - ref_forward_prediction) ** 2).mean(dim=tuple(range(1, x0.ndim))) ).detach() loss_terms["x0_norm"] = torch.mean(x0**2).detach() loss_terms["x0_norm_max"] = torch.max(x0**2).detach() loss_terms["old_deviate"] = torch.mean((forward_prediction - old_prediction) ** 2).detach() loss_terms["old_deviate_max"] = torch.max((forward_prediction - old_prediction) ** 2).detach() loss_terms["total_loss"] = loss.detach() scaled_loss = loss / effective_grad_accum_steps if torch.isnan(scaled_loss) or torch.isinf(scaled_loss): scaled_loss = scaled_loss * 0.0 if mixed_precision_dtype == torch.float16: scaler.scale(scaled_loss).backward() else: scaled_loss.backward() current_accumulated_steps += 1 for ki, vi in loss_terms.items(): info_accumulated[ki].append(vi) if current_accumulated_steps % effective_grad_accum_steps == 0: if mixed_precision_dtype == torch.float16: scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_( transformer_ddp.module.parameters(), config.train.max_grad_norm ) if mixed_precision_dtype == torch.float16: scaler.step(optimizer) scaler.update() else: optimizer.step() optimizer.zero_grad() gradient_update_times += 1 log_info = {k: torch.mean(torch.stack(v)).item() for k, v in info_accumulated.items()} log_info["grad_norm"] = ( grad_norm.detach().float().item() if torch.is_tensor(grad_norm) else float(grad_norm) ) info_tensor = torch.tensor([log_info[k] for k in sorted(log_info)], device=device) dist.all_reduce(info_tensor, op=dist.ReduceOp.AVG) reduced = {k: info_tensor[i].item() for i, k in enumerate(sorted(log_info))} if is_main_process(rank): _log_dict = { "global_step": global_step, "gradient_update_times": gradient_update_times, "epoch": epoch, "inner_epoch": inner_epoch, "current_time": time.time() - start_time, **reduced, } wandb.log(_log_dict, commit=False) logger.info( "[step %d] loss=%.6f policy=%.6f grad=%.6f kl=%.6f", global_step, reduced.get("total_loss", 0), reduced.get("policy_loss", 0), reduced.get("grad_norm", 0), reduced.get("kl_div_loss", 0), ) global_step += 1 info_accumulated = defaultdict(list) if ( config.train.ema and ema is not None and (current_accumulated_steps % effective_grad_accum_steps == 0) ): ema.step(transformer_trainable_parameters, global_step) time_logger.end("train_time") if world_size > 1: dist.barrier() with torch.no_grad(): decay = return_decay( global_step, config.decay_type, custom_decay_step=getattr(config, "custom_decay_step", 0), custom_decay_value=getattr(config, "custom_decay_value", 0.0), ) for src_p, tgt_p in zip( transformer_trainable_parameters, old_transformer_trainable_parameters, strict=True ): tgt_p.data.copy_(tgt_p.detach().data * decay + src_p.detach().clone().data * (1.0 - decay)) for mtype, inf_model in inference_models.items(): sync_lora_to_inference(transformer_ddp.module, unwrap_compiled(inf_model), adapter_name="old") time_logger.end("total_time") stats = time_logger.get_results() if is_main_process(rank): wandb.log({f"time/{k}": v for k, v in stats.items()}, commit=True) logger.info("Step %d Time: %s", global_step, stats) time_logger.empty_cache() if is_main_process(rank): wandb.finish() cleanup_distributed() if __name__ == "__main__": app.run(main)