import copy 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 tempfile 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 StableDiffusion3Pipeline from ml_collections import config_flags from peft import LoraConfig, PeftModel, get_peft_model from PIL import Image from torch.cuda.amp import GradScaler from torch.cuda.amp import autocast as torch_autocast 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__))))) 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_debug_image_subset, 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.post_training.rewards from diffusion.post_training.diffusers_patch.pipeline_with_logprob import pipeline_with_logprob_sd3 from diffusion.post_training.diffusers_patch.text_encode import encode_sd3_prompt from diffusion.post_training.ema import EMAModuleWrapper from diffusion.post_training.stat_tracking import PerPromptStatTracker tqdm = tqdm.tqdm FLAGS = flags.FLAGS config_flags.DEFINE_config_file( "config", "configs/sol_rl/sd3.py", "Training configuration.", ) logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") TEXT_ENCODER_MAX_SEQ_LEN = 128 TOKENIZER_MAX_LENGTH = 256 WANDB_MAX_LOG_IMAGES = 12 def compute_text_embeddings(prompts, text_encoders, tokenizers, max_sequence_length, device): with torch.no_grad(): prompt_embeds, pooled_prompt_embeds = encode_sd3_prompt( text_encoders, tokenizers, prompts, max_sequence_length, device=device, ) return prompt_embeds, pooled_prompt_embeds def _build_sd3_latents_from_seeds(seed_list, latent_shape, device, dtype): latents = [] channels, latent_h, latent_w = latent_shape for seed in seed_list: generator = torch.Generator(device=device).manual_seed(int(seed)) latents.append( torch.randn( 1, channels, latent_h, latent_w, device=device, dtype=dtype, generator=generator, ) ) return torch.cat(latents, dim=0) def eval_fn( pipeline, test_dataloader, text_encoders, tokenizers, 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) sequential_decode = bool(getattr(config, "sequential_decode", True)) if config.train.ema and ema is not None: ema.copy_ema_to(transformer_trainable_parameters, store_temp=True) pipeline.transformer.eval() neg_prompt_embed, neg_pooled_prompt_embed = compute_text_embeddings( [""], text_encoders, tokenizers, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device ) 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=test_dataloader.num_workers, ) for prompts, prompt_metadata in tqdm( eval_loader, desc="Eval", disable=not is_main_process(rank), dynamic_ncols=True, ): prompt_embeds, pooled_prompt_embeds = compute_text_embeddings( prompts, text_encoders, tokenizers, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device ) bs = len(prompts) with torch_autocast(enabled=(config.mixed_precision in ["fp16", "bf16"]), dtype=mixed_precision_dtype): with torch.no_grad(): images, _, _ = pipeline_with_logprob_sd3( pipeline, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=neg_prompt_embed.repeat(bs, 1, 1), negative_pooled_prompt_embeds=neg_pooled_prompt_embed.repeat(bs, 1), num_inference_steps=config.sample.eval_num_steps, guidance_scale=config.eval_sample_guidance_scale, output_type="pt", height=config.resolution, width=config.resolution, noise_level=config.sample.noise_level, deterministic=True, solver=config.sample.solver, sequential_decode=sequential_decode, ) 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(): rewards_tensor = torch.as_tensor(value, device=device).float() all_rewards[key].append(gather_tensor_to_all(rewards_tensor, world_size).numpy()) enable_debug_image_save = bool(getattr(config, "enable_debug_image_save", True)) if is_main_process(rank): final_rewards = {key: np.concatenate(value_list) for key, value_list in all_rewards.items()} images_to_log = images.cpu() prompts_to_log = prompts if enable_debug_image_save: eval_debug_dir = os.path.join(config.save_dir, "debug_images", "eval", f"step_{global_step}") save_debug_image_subset( images=images_to_log, prompts=prompts_to_log, save_root=eval_debug_dir, prefix="eval", resolution=config.resolution, rewards=final_rewards.get("avg", None), max_images=getattr(config, "debug_image_subset_size", 6), ) with tempfile.TemporaryDirectory() as tmpdir: num_to_log = min(WANDB_MAX_LOG_IMAGES, len(images_to_log)) for idx in range(num_to_log): image = images_to_log[idx].float() pil = Image.fromarray((image.numpy().transpose(1, 2, 0) * 255).astype(np.uint8)) pil = pil.resize((config.resolution, config.resolution)) pil.save(os.path.join(tmpdir, f"{idx}.jpg")) sampled_prompts_log = [prompts_to_log[i] for i in range(num_to_log)] sampled_rewards_log = [{k: final_rewards[k][i] for k in final_rewards} for i in range(num_to_log)] wandb.log( { "eval_images": [ wandb.Image( os.path.join(tmpdir, f"{idx}.jpg"), caption=f"{prompt:.1000} | " + " | ".join(f"{k}: {v:.2f}" for k, v in reward.items() if v != -10), ) for idx, (prompt, reward) in enumerate(zip(sampled_prompts_log, sampled_rewards_log)) ], **{f"eval_reward_{k}": np.mean(v[v != -10]) for k, v in final_rewards.items()}, }, commit=False, ) if config.train.ema and ema is not None: ema.copy_temp_to(transformer_trainable_parameters) if world_size > 1: dist.barrier() def _swap_pipeline_model(pipeline, mode, inference_models, transformer_ddp, original_transformer): """Swap pipeline.transformer to the model specified by *mode*. mode: "compile_nvfp4" | "compile" | "peft" """ if mode == "peft": pipeline.transformer = original_transformer transformer_ddp.module.set_adapter("old") else: pipeline.transformer = inference_models[mode] def _rollout_for_one_prompt( pipeline, reward_fn, executor, prompt_text, prompt_meta, prompt_embed_single, pooled_embed_single, neg_prompt_embed_single, neg_pooled_prompt_embed_single, prompt_token_ids_single, config, device, inference_models=None, transformer_ddp=None, original_transformer=None, ): sequential_decode = bool(getattr(config, "sequential_decode", True)) amp_dtype = ( torch.bfloat16 if config.mixed_precision == "bf16" else (torch.float16 if config.mixed_precision == "fp16" else None) ) enable_amp = amp_dtype is not 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 = int(getattr(config.sample, "full_rollout_num", config.sample.best_of_n)) full_rollout_num = max(1, min(full_rollout_num, draft_total)) latent_h = config.resolution // 8 latent_w = config.resolution // 8 latent_shape = (16, latent_h, latent_w) seed_pool = [] draft_reward_pool = [] prompt_samples = [] final_images = None final_prompts = None full_chunks = int(config.sample.rollout_batch_size) 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 original_transformer is not None if preview_step > 0: # --- Stage 1: draft preview (fast screening) --- if _can_swap: _swap_pipeline_model(pipeline, preview_model_key, inference_models, transformer_ddp, original_transformer) with torch.no_grad(): for iter_idx in range(config.sample.per_prompt_iter_num): batch_size = int(config.sample.rollout_batch_size) prompt_embeds = prompt_embed_single.repeat(batch_size, 1, 1) pooled_prompt_embeds = pooled_embed_single.repeat(batch_size, 1) neg_prompt_embeds = neg_prompt_embed_single.repeat(batch_size, 1, 1) neg_pooled_prompt_embeds = neg_pooled_prompt_embed_single.repeat(batch_size, 1) seed_list = torch.randint( low=0, high=2**31 - 1, size=(batch_size,), device="cpu", ).tolist() init_latents = _build_sd3_latents_from_seeds( seed_list, latent_shape=latent_shape, device=device, dtype=prompt_embeds.dtype, ) with torch_autocast(enabled=enable_amp, dtype=amp_dtype): images, _, _ = pipeline_with_logprob_sd3( pipeline, latents=init_latents, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, negative_pooled_prompt_embeds=neg_pooled_prompt_embeds, num_inference_steps=preview_step, guidance_scale=config.rollout_sample_guidance_scale, output_type="pt", height=config.resolution, width=config.resolution, noise_level=config.sample.noise_level, deterministic=True, solver=config.sample.solver, sequential_decode=sequential_decode, ) rewards, _ = reward_fn( images, [prompt_text] * batch_size, [prompt_meta] * batch_size, only_strict=True, ) draft_avg = torch.as_tensor(rewards["avg"], device=device).float() seed_pool.extend(int(s) for s in seed_list) draft_reward_pool.extend(draft_avg.detach().cpu().tolist()) draft_rewards = torch.as_tensor(draft_reward_pool, device=device).float() stage1_indices = select_indices_by_mode( draft_rewards, target_count=full_rollout_num, mode=getattr(config.sample, "stage1_select_mode", "best_worst"), ) selected_seeds = [seed_pool[int(i)] for i in stage1_indices.detach().cpu().tolist()] # --- Stage 2: full rollout (may use a different model) --- if _can_swap and fullrollout_model_key != preview_model_key: _swap_pipeline_model( pipeline, fullrollout_model_key, inference_models, transformer_ddp, original_transformer ) for start in range(0, len(selected_seeds), full_chunks): seed_chunk = selected_seeds[start : start + full_chunks] bs = len(seed_chunk) prompt_embeds = prompt_embed_single.repeat(bs, 1, 1) pooled_prompt_embeds = pooled_embed_single.repeat(bs, 1) neg_prompt_embeds = neg_prompt_embed_single.repeat(bs, 1, 1) neg_pooled_prompt_embeds = neg_pooled_prompt_embed_single.repeat(bs, 1) init_latents = _build_sd3_latents_from_seeds( seed_chunk, latent_shape=latent_shape, device=device, dtype=prompt_embeds.dtype, ) with torch.no_grad(): with torch_autocast(enabled=enable_amp, dtype=amp_dtype): images, latents, _ = pipeline_with_logprob_sd3( pipeline, latents=init_latents, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, negative_pooled_prompt_embeds=neg_pooled_prompt_embeds, num_inference_steps=full_steps, guidance_scale=config.rollout_sample_guidance_scale, output_type="pt", height=config.resolution, width=config.resolution, noise_level=config.sample.noise_level, deterministic=True, solver=config.sample.solver, sequential_decode=sequential_decode, ) timesteps = pipeline.scheduler.timesteps.repeat(bs, 1).to(device) latents = torch.stack(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": prompt_embeds, "pooled_prompt_embeds": pooled_prompt_embeds, "timesteps": timesteps, "next_timesteps": torch.concatenate([timesteps[:, 1:], torch.zeros_like(timesteps[:, :1])], dim=1), "latents_clean": latents[:, -1], "rewards_future": rewards_future, } ) final_images = images final_prompts = [prompt_text] * bs else: for iter_idx in range(config.sample.per_prompt_iter_num): batch_size = int(config.sample.rollout_batch_size) prompt_embeds = prompt_embed_single.repeat(batch_size, 1, 1) pooled_prompt_embeds = pooled_embed_single.repeat(batch_size, 1) neg_prompt_embeds = neg_prompt_embed_single.repeat(batch_size, 1, 1) neg_pooled_prompt_embeds = neg_pooled_prompt_embed_single.repeat(batch_size, 1) seed_list = torch.randint( low=0, high=2**31 - 1, size=(batch_size,), device="cpu", ).tolist() init_latents = _build_sd3_latents_from_seeds( seed_list, latent_shape=latent_shape, device=device, dtype=prompt_embeds.dtype, ) with torch.no_grad(): with torch_autocast(enabled=enable_amp, dtype=amp_dtype): images, latents, _ = pipeline_with_logprob_sd3( pipeline, latents=init_latents, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, negative_pooled_prompt_embeds=neg_pooled_prompt_embeds, num_inference_steps=full_steps, guidance_scale=config.rollout_sample_guidance_scale, output_type="pt", height=config.resolution, width=config.resolution, noise_level=config.sample.noise_level, deterministic=True, solver=config.sample.solver, sequential_decode=sequential_decode, ) timesteps = pipeline.scheduler.timesteps.repeat(batch_size, 1).to(device) latents = torch.stack(latents, dim=1) rewards_future = executor.submit( reward_fn, images, [prompt_text] * batch_size, [prompt_meta] * batch_size, True, ) time.sleep(0) prompt_samples.append( { "prompt_ids": prompt_token_ids_single.repeat(batch_size, 1), "prompt_embeds": prompt_embeds, "pooled_prompt_embeds": pooled_prompt_embeds, "timesteps": timesteps, "next_timesteps": torch.concatenate([timesteps[:, 1:], torch.zeros_like(timesteps[:, :1])], dim=1), "latents_clean": latents[:, -1], "rewards_future": rewards_future, } ) final_images = images final_prompts = [prompt_text] * batch_size 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) final_rewards = collated["rewards"]["avg"] keep_indices = select_indices_by_mode( final_rewards, target_count=config.sample.best_of_n, mode=getattr(config.sample, "stage2_select_mode", "best_worst"), ) collated = filter_by_indices(collated, keep_indices) return collated, final_images, final_prompts def main(_): config = FLAGS.config 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("\n%s", config) set_seed(config.seed, rank) mixed_precision_dtype = None if config.mixed_precision == "fp16": mixed_precision_dtype = torch.float16 elif config.mixed_precision == "bf16": mixed_precision_dtype = torch.bfloat16 enable_amp = mixed_precision_dtype is not None scaler = GradScaler(enabled=enable_amp) pipeline = StableDiffusion3Pipeline.from_pretrained(config.pretrained.model) pipeline.vae.requires_grad_(False) pipeline.text_encoder.requires_grad_(False) pipeline.text_encoder_2.requires_grad_(False) pipeline.text_encoder_3.requires_grad_(False) pipeline.transformer.requires_grad_(not config.use_lora) pipeline.safety_checker = None pipeline.set_progress_bar_config(disable=True) text_encoders = [pipeline.text_encoder, pipeline.text_encoder_2, pipeline.text_encoder_3] tokenizers = [pipeline.tokenizer, pipeline.tokenizer_2, pipeline.tokenizer_3] text_encoder_dtype = mixed_precision_dtype if enable_amp else torch.float32 pipeline.vae.to(device, dtype=torch.float32) pipeline.text_encoder.to(device, dtype=text_encoder_dtype) pipeline.text_encoder_2.to(device, dtype=text_encoder_dtype) pipeline.text_encoder_3.to(device, dtype=text_encoder_dtype) transformer = pipeline.transformer.to(device) # --- Inference models: clean copies (no PEFT), optionally nvfp4 + torch.compile --- 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 = copy.deepcopy(transformer) m.eval() m.requires_grad_(False) m.to(dtype=torch.bfloat16) if "nvfp4" in mtype: if not _HAS_TE: raise RuntimeError(f"model type {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} nn.Linear -> te.Linear, skipped {n_skip}") wrap_forward_with_fp8(m) logger.info(f"[NVFP4] {mtype}: wrap_forward_with_fp8 applied") if is_main_process(rank): report_path = os.path.join(config.save_dir, f"nvfp4_quant_report_{mtype}.txt") os.makedirs(os.path.dirname(report_path), exist_ok=True) with open(report_path, "w") as f: f.write(f"NVFP4 Quantization Report ({mtype})\n{'=' * 60}\n") f.write(f"skip_modules: {nvfp4_skip_modules}\n") f.write(f"min_dim: {nvfp4_min_dim}\n") f.write(f"replaced: {n_rep} skipped: {n_skip}\n\n") f.write(f"Replaced (te.Linear + NVFP4):\n{'-' * 60}\n") for fqn, inf, outf, bias, _ in rep_d: f.write(f" {fqn:60s} in={inf:6d} out={outf:6d} bias={bias}\n") f.write(f"\nSkipped (kept as nn.Linear):\n{'-' * 60}\n") for fqn, inf, outf, bias, reason in skip_d: f.write(f" {fqn:60s} in={inf:6d} out={outf:6d} bias={bias} reason={reason}\n") logger.info(f"[NVFP4] Report saved to {report_path}") 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) transformer_lora_config = 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, transformer_lora_config) transformer.add_adapter("old", transformer_lora_config) transformer.set_adapter("default") 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, ) _, train_dataloader, train_sampler, _, test_dataloader = build_datasets_and_loaders(config, world_size, rank) neg_prompt_embed, neg_pooled_prompt_embed = compute_text_embeddings( [""], text_encoders, tokenizers, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device ) if config.sample.best_of_n == 1: config.per_prompt_stat_tracking = False if config.per_prompt_stat_tracking: stat_tracker = PerPromptStatTracker(config.global_std) else: raise ValueError("per_prompt_stat_tracking must be enabled for this recipe") 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) executor = futures.ThreadPoolExecutor(max_workers=8) ema = None if config.train.ema: ema = EMAModuleWrapper(transformer_trainable_parameters, decay=0.9, update_step_interval=1, device=device) num_train_timesteps = int(config.rollout_sample_num_steps * config.train.timestep_fraction) train_iter = iter(train_dataloader) optimizer.zero_grad() # --- Resume from checkpoint --- first_epoch = 0 global_step = 0 candidates = find_resume_candidates(config) global_step, resume_parameters = resume_from_checkpoint( candidates, transformer_ddp.module, ema, optimizer, scaler, device, ) first_epoch = global_step if not resume_parameters: for src_param, tgt_param in zip( transformer_trainable_parameters, old_transformer_trainable_parameters, strict=True ): tgt_param.data.copy_(src_param.detach().data) if global_step != 0: for i in range(global_step): prompts, prompt_metadata = next(train_iter) if world_size > 1: dist.barrier() # Sync old adapter weights → all inference models (after resume or fresh init) for mtype, inf_model in inference_models.items(): n_synced = sync_lora_to_inference( transformer_ddp.module, unwrap_compiled(inf_model), adapter_name="old", ) logger.info(f"[SYNC] Initial sync: merged {n_synced} LoRA layers → {mtype!r}") time_logger = DistributedTimeLogger(device) start_time = time.time() 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: py_rng_state = random.getstate() np_rng_state = np.random.get_state() torch_rng_state = torch.random.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state_all() eval_fn( pipeline, test_dataloader, text_encoders, tokenizers, config, device, rank, world_size, global_step, eval_reward_fn, executor, mixed_precision_dtype, ema, transformer_trainable_parameters, ) random.setstate(py_rng_state) np.random.set_state(np_rng_state) torch.random.set_rng_state(torch_rng_state) torch.cuda.set_rng_state_all(cuda_rng_state) time_logger.end("eval_time") time_logger.start("rollout_time") pipeline.transformer.eval() prompts, prompt_metadata = next(train_iter) prompt_embeds_all, pooled_prompt_embeds_all = compute_text_embeddings( prompts, text_encoders, tokenizers, max_sequence_length=TEXT_ENCODER_MAX_SEQ_LEN, device=device ) prompt_ids_all = tokenizers[0]( prompts, padding="max_length", max_length=TOKENIZER_MAX_LENGTH, truncation=True, return_tensors="pt", ).input_ids.to(device) prompt_wise_samples = [] step_prompt_reward_groups = [] images_for_log = None prompts_for_log = None rewards_for_log = None _saved_pipeline_transformer = pipeline.transformer # For non-preview path, set up the model once (same as before) if preview_step == 0: if fullrollout_model_key != "peft" and fullrollout_model_key in inference_models: pipeline.transformer = inference_models[fullrollout_model_key] else: transformer_ddp.module.set_adapter("old") 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( pipeline=pipeline, reward_fn=reward_fn, executor=executor, prompt_text=prompts[prompt_idx], prompt_meta=prompt_metadata[prompt_idx], prompt_embed_single=prompt_embeds_all[prompt_idx : prompt_idx + 1], pooled_embed_single=pooled_prompt_embeds_all[prompt_idx : prompt_idx + 1], neg_prompt_embed_single=neg_prompt_embed, neg_pooled_prompt_embed_single=neg_pooled_prompt_embed, prompt_token_ids_single=prompt_ids_all[prompt_idx : prompt_idx + 1], config=config, device=device, inference_models=inference_models, transformer_ddp=transformer_ddp, original_transformer=_saved_pipeline_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=prompt_idx, prompt_text=prompts[prompt_idx], prompt_meta=prompt_meta_i, intra_prompt_data_list=[collated_prompt_samples], ) ) images_for_log = final_images prompts_for_log = final_prompts rewards_for_log = collated_prompt_samples["rewards"]["avg"] # Restore original transformer and switch back to "default" adapter pipeline.transformer = _saved_pipeline_transformer 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 = { key: gather_tensor_to_all(value, world_size).numpy() for key, value in collated_samples["rewards"].items() } if is_main_process(rank): rewards_to_log = gathered_rewards_dict["avg"] rewards_to_log = rewards_to_log.reshape( world_size * config.sample.per_gpu_to_process_prompts, -1, num_train_timesteps ) rewards_to_log = rewards_to_log.mean(axis=-1) wandb.log( { "epoch": epoch, "reward/mean": rewards_to_log.mean(), "reward/max": rewards_to_log.max(axis=1).mean(), "reward/min": rewards_to_log.min(axis=1).mean(), "reward/range": rewards_to_log.max(axis=1).mean() - rewards_to_log.min(axis=1).mean(), }, commit=False, ) prompt_ids_all_global = gather_tensor_to_all(collated_samples["prompt_ids"], world_size) prompts_all_decoded = tokenizers[0].batch_decode(prompt_ids_all_global.cpu().numpy(), skip_special_tokens=True) advantages = stat_tracker.update(prompts_all_decoded, gathered_rewards_dict["avg"]) if is_main_process(rank): group_size, trained_prompt_num = stat_tracker.get_stats() zero_std_ratio, reward_std_mean = calculate_zero_std_ratio(prompts_all_decoded, gathered_rewards_dict) wandb.log( { "group_size": group_size, "trained_prompt_num": trained_prompt_num, "zero_std_ratio": zero_std_ratio, "reward_std_mean": reward_std_mean, "mean_reward_100": stat_tracker.get_mean_of_top_rewards(100), "mean_reward_75": stat_tracker.get_mean_of_top_rewards(75), "mean_reward_50": stat_tracker.get_mean_of_top_rewards(50), "mean_reward_25": stat_tracker.get_mean_of_top_rewards(25), "mean_reward_10": stat_tracker.get_mean_of_top_rewards(10), }, commit=False, ) stat_tracker.clear() samples_per_gpu = collated_samples["timesteps"].shape[0] if advantages.ndim == 1: advantages = advantages[:, None] if advantages.shape[0] != world_size * samples_per_gpu: raise RuntimeError("Unexpected advantage shape after all-gather") 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 assert total_batch_size_filtered == config.sample.per_gpu_total_samples_to_train time_logger.start("train_time") transformer_ddp.train() 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_samples = {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_samples[key] = shuffled_samples[key][ torch.arange(total_batch_size_filtered, device=device)[:, None], perms_time, ] training_batch_size = config.train.batch_size batches = [] for batch_idx in range(config.train.n_batch_per_epoch): start = batch_idx * training_batch_size end = (batch_idx + 1) * training_batch_size batches.append({k: v[start:end] for k, v in shuffled_samples.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, ): current_bs = len(train_batch["prompt_embeds"]) if config.train_sample_guidance_scale > 1.0: embeds = torch.cat( [ neg_prompt_embed.repeat(current_bs, 1, 1), train_batch["prompt_embeds"], ] ) pooled_embeds = torch.cat( [ neg_pooled_prompt_embed.repeat(current_bs, 1), train_batch["pooled_prompt_embeds"], ] ) else: embeds = train_batch["prompt_embeds"] pooled_embeds = train_batch["pooled_prompt_embeds"] for j_idx in range(num_train_timesteps): x0 = train_batch["latents_clean"] t = train_batch["timesteps"][:, j_idx] / 1000.0 t_expanded = t.view(-1, *([1] * (len(x0.shape) - 1))) noise = torch.randn_like(x0.float()) xt = (1 - t_expanded) * x0 + t_expanded * noise with torch_autocast(enabled=enable_amp, dtype=mixed_precision_dtype): transformer_ddp.module.set_adapter("old") with torch.no_grad(): old_prediction = transformer_ddp( hidden_states=xt, timestep=train_batch["timesteps"][:, j_idx], encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0].detach() transformer_ddp.module.set_adapter("default") forward_prediction = transformer_ddp( hidden_states=xt, timestep=train_batch["timesteps"][:, j_idx], encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0] with torch.no_grad(): with transformer_ddp.module.disable_adapter(): ref_forward_prediction = transformer_ddp( hidden_states=xt, timestep=train_batch["timesteps"][:, j_idx], encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0] transformer_ddp.module.set_adapter("default") 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) normalized_adv = (advantages_clip / config.train.adv_clip_max) / 2.0 + 0.5 r = torch.clamp(normalized_adv, 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 - t_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 - t_expanded * implicit_negative_prediction with torch.no_grad(): negative_weight_factor = ( 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 / negative_weight_factor).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 = {} 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) kl_div_loss = torch.mean(kl_div_loss) loss_terms["kl_div_loss"] = torch.mean(kl_div_loss).detach() loss_terms["kl_div"] = torch.mean( ((forward_prediction - ref_forward_prediction) ** 2).mean(dim=tuple(range(1, x0.ndim))) ).detach() 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 k_info, v_info in loss_terms.items(): info_accumulated[k_info].append(v_info) 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()} if torch.is_tensor(grad_norm): log_info["grad_norm"] = grad_norm.detach().float().item() else: log_info["grad_norm"] = 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_log = {k: info_tensor[i].item() for i, k in enumerate(sorted(log_info))} if is_main_process(rank): wandb.log( { "global_step": global_step, "gradient_update_times": gradient_update_times, "epoch": epoch, "inner_epoch": inner_epoch, "current_time": time.time() - start_time, **reduced_log, }, commit=False, ) 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_param, tgt_param in zip( transformer_trainable_parameters, old_transformer_trainable_parameters, strict=True ): tgt_param.data.copy_(tgt_param.detach().data * decay + src_param.detach().clone().data * (1.0 - decay)) # Sync updated old adapter → all inference models for next rollout 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): time_logs = {f"time/{k}": v for k, v in stats.items()} wandb.log(time_logs, commit=True) logger.info("Step %d Time Report: %s", global_step, time_logs) time_logger.empty_cache() if is_main_process(rank): wandb.finish() cleanup_distributed() if __name__ == "__main__": app.run(main)