from collections import defaultdict import contextlib import os import datetime from concurrent import futures import time import json import hashlib from absl import app, flags from accelerate import Accelerator from ml_collections import config_flags from accelerate.utils import set_seed, ProjectConfiguration from accelerate.logging import get_logger from diffusers import StableDiffusion3Pipeline from diffusers.utils.torch_utils import is_compiled_module import numpy as np import flow_grpo.prompts import flow_grpo.rewards from flow_grpo.stat_tracking import PerPromptStatTracker from flow_grpo.diffusers_patch.sd3_pipeline_with_logprob import pipeline_with_logprob from flow_grpo.diffusers_patch.sd3_sde_with_logprob import sde_step_with_logprob from flow_grpo.diffusers_patch.train_dreambooth_lora_sd3 import encode_prompt import torch import wandb from functools import partial import tqdm import tempfile from PIL import Image from peft import LoraConfig, get_peft_model, set_peft_model_state_dict, PeftModel import random from torch.utils.data import Dataset, DataLoader, Sampler from flow_grpo.ema import EMAModuleWrapper import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from copd_loss import opd_positive_loss, copd_loss as copd_loss_fn, regen_velocity_per_step, sde_time_weight tqdm = partial(tqdm.tqdm, dynamic_ncols=True) FLAGS = flags.FLAGS config_flags.DEFINE_config_file("config", "config/base.py", "Training configuration.") logger = get_logger(__name__) class TextPromptDataset(Dataset): def __init__(self, dataset, split='train'): self.file_path = os.path.join(dataset, f'{split}.txt') with open(self.file_path, 'r') as f: self.prompts = [line.strip() for line in f.readlines()] def __len__(self): return len(self.prompts) def __getitem__(self, idx): return {"prompt": self.prompts[idx], "metadata": {}} @staticmethod def collate_fn(examples): prompts = [example["prompt"] for example in examples] metadatas = [example["metadata"] for example in examples] return prompts, metadatas class GenevalPromptDataset(Dataset): def __init__(self, dataset, split='train'): self.file_path = os.path.join(dataset, f'{split}_metadata.jsonl') with open(self.file_path, 'r', encoding='utf-8') as f: self.metadatas = [json.loads(line) for line in f] self.prompts = [item['prompt'] for item in self.metadatas] def __len__(self): return len(self.prompts) def __getitem__(self, idx): return {"prompt": self.prompts[idx], "metadata": self.metadatas[idx]} @staticmethod def collate_fn(examples): prompts = [example["prompt"] for example in examples] metadatas = [example["metadata"] for example in examples] return prompts, metadatas class DistributedKRepeatSampler(Sampler): def __init__(self, dataset, batch_size, k, num_replicas, rank, seed=0): self.dataset = dataset self.batch_size = batch_size # Batch size per replica self.k = k # Number of repetitions per sample self.num_replicas = num_replicas # Total number of replicas self.rank = rank # Current replica rank self.seed = seed # Random seed for synchronization # Compute the number of unique samples needed per iteration self.total_samples = self.num_replicas * self.batch_size assert self.total_samples % self.k == 0, f"k can not divide n*b, k{k}-num_replicas{num_replicas}-batch_size{batch_size}" self.m = self.total_samples // self.k # Number of unique samples self.epoch = 0 def __iter__(self): while True: # Generate a deterministic random sequence to ensure all replicas are synchronized g = torch.Generator() g.manual_seed(self.seed + self.epoch) # Randomly select m unique samples indices = torch.randperm(len(self.dataset), generator=g)[:self.m].tolist() # Repeat each sample k times to generate n*b total samples repeated_indices = [idx for idx in indices for _ in range(self.k)] # Shuffle to ensure uniform distribution shuffled_indices = torch.randperm(len(repeated_indices), generator=g).tolist() shuffled_samples = [repeated_indices[i] for i in shuffled_indices] # Split samples to each replica per_card_samples = [] for i in range(self.num_replicas): start = i * self.batch_size end = start + self.batch_size per_card_samples.append(shuffled_samples[start:end]) # Return current replica's sample indices yield per_card_samples[self.rank] def set_epoch(self, epoch): self.epoch = epoch # Used to synchronize random state across epochs def compute_text_embeddings(prompt, text_encoders, tokenizers, max_sequence_length, device): with torch.no_grad(): prompt_embeds, pooled_prompt_embeds = encode_prompt( text_encoders, tokenizers, prompt, max_sequence_length ) prompt_embeds = prompt_embeds.to(device) pooled_prompt_embeds = pooled_prompt_embeds.to(device) return prompt_embeds, pooled_prompt_embeds def calculate_zero_std_ratio(prompts, gathered_rewards): """ Calculate the proportion of unique prompts whose reward standard deviation is zero. Args: prompts: List of prompts. gathered_rewards: Dictionary containing rewards, must include the key 'ori_avg'. Returns: zero_std_ratio: Proportion of prompts with zero standard deviation. prompt_std_devs: Mean standard deviation across all unique prompts. """ # Convert prompt list to NumPy array prompt_array = np.array(prompts) # Get unique prompts and their group information unique_prompts, inverse_indices, counts = np.unique( prompt_array, return_inverse=True, return_counts=True ) # Group rewards for each prompt grouped_rewards = gathered_rewards['ori_avg'][np.argsort(inverse_indices)] split_indices = np.cumsum(counts)[:-1] reward_groups = np.split(grouped_rewards, split_indices) # Calculate standard deviation for each group prompt_std_devs = np.array([np.std(group) for group in reward_groups]) # Calculate the ratio of zero standard deviation zero_std_count = np.count_nonzero(prompt_std_devs == 0) zero_std_ratio = zero_std_count / len(prompt_std_devs) return zero_std_ratio, prompt_std_devs.mean() def create_generator(prompts, base_seed): generators = [] for prompt in prompts: # Use a stable hash (SHA256), then convert it to an integer seed hash_digest = hashlib.sha256(prompt.encode()).digest() prompt_hash_int = int.from_bytes(hash_digest[:4], 'big') # Take the first 4 bytes as part of the seed seed = (base_seed + prompt_hash_int) % (2**31) # Ensure the number is within a valid range gen = torch.Generator().manual_seed(seed) generators.append(gen) return generators def compute_log_prob(transformer, pipeline, sample, j, embeds, pooled_embeds, config): if config.train.cfg: noise_pred = transformer( hidden_states=torch.cat([sample["latents"][:, j]] * 2), timestep=torch.cat([sample["timesteps"][:, j]] * 2), encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0] noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = ( noise_pred_uncond + config.sample.guidance_scale * (noise_pred_text - noise_pred_uncond) ) else: noise_pred = transformer( hidden_states=sample["latents"][:, j], timestep=sample["timesteps"][:, j], encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0] # compute the log prob of next_latents given latents under the current model prev_sample, log_prob, prev_sample_mean, std_dev_t = sde_step_with_logprob( pipeline.scheduler, noise_pred.float(), sample["timesteps"][:, j], sample["latents"][:, j].float(), prev_sample=sample["next_latents"][:, j].float(), noise_level=config.sample.noise_level, ) return prev_sample, log_prob, prev_sample_mean, std_dev_t def compute_velocity(transformer, sample, j, embeds, pooled_embeds): """Flow-CoPD: raw velocity-field prediction v(x_j, t_j) under the currently ACTIVE adapter. No CFG (config.train.cfg must be False for OPD distillation), so embeds are the conditional embeds only.""" return transformer( hidden_states=sample["latents"][:, j], timestep=sample["timesteps"][:, j], encoder_hidden_states=embeds, pooled_projections=pooled_embeds, return_dict=False, )[0] def eval(pipeline, test_dataloader, text_encoders, tokenizers, config, accelerator, global_step, reward_fn, executor, autocast, num_train_timesteps, ema, transformer_trainable_parameters): if config.train.ema: ema.copy_ema_to(transformer_trainable_parameters, store_temp=True) neg_prompt_embed, neg_pooled_prompt_embed = compute_text_embeddings([""], text_encoders, tokenizers, max_sequence_length=128, device=accelerator.device) sample_neg_prompt_embeds = neg_prompt_embed.repeat(config.sample.test_batch_size, 1, 1) sample_neg_pooled_prompt_embeds = neg_pooled_prompt_embed.repeat(config.sample.test_batch_size, 1) # test_dataloader = itertools.islice(test_dataloader, 2) all_rewards = defaultdict(list) for test_batch in tqdm( test_dataloader, desc="Eval: ", disable=not accelerator.is_local_main_process, position=0, ): prompts, prompt_metadata = test_batch prompt_embeds, pooled_prompt_embeds = compute_text_embeddings( prompts, text_encoders, tokenizers, max_sequence_length=128, device=accelerator.device ) # The last batch may not be full batch_size if len(prompt_embeds)= config.train.batch_size # assert config.sample.train_batch_size % config.train.batch_size == 0 # assert samples_per_epoch % total_train_batch_size == 0 epoch = 0 global_step = 0 train_iter = iter(train_dataloader) while True: #################### EVAL #################### pipeline.transformer.eval() if epoch % config.eval_freq == 0 and (epoch > 0 or getattr(config, "eval_at_start", True)): eval(pipeline, test_dataloader, text_encoders, tokenizers, config, accelerator, global_step, eval_reward_fn, executor, autocast, num_train_timesteps, ema, transformer_trainable_parameters) if epoch % config.save_freq == 0 and epoch > 0 and accelerator.is_main_process: save_ckpt(config.save_dir, transformer, global_step, accelerator, ema, transformer_trainable_parameters, config) #################### SAMPLING #################### pipeline.transformer.eval() samples = [] prompts = [] for i in tqdm( range(config.sample.num_batches_per_epoch), desc=f"Epoch {epoch}: sampling", disable=not accelerator.is_local_main_process, position=0, ): train_sampler.set_epoch(epoch * config.sample.num_batches_per_epoch + i) prompts, prompt_metadata = next(train_iter) prompt_embeds, pooled_prompt_embeds = compute_text_embeddings( prompts, text_encoders, tokenizers, max_sequence_length=128, device=accelerator.device ) prompt_ids = tokenizers[0]( prompts, padding="max_length", max_length=256, truncation=True, return_tensors="pt", ).input_ids.to(accelerator.device) # sample if config.sample.same_latent: generator = create_generator(prompts, base_seed=epoch*10000+i) else: generator = None with autocast(): with torch.no_grad(): images, latents, log_probs = pipeline_with_logprob( pipeline, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=sample_neg_prompt_embeds, negative_pooled_prompt_embeds=sample_neg_pooled_prompt_embeds, num_inference_steps=config.sample.num_steps, guidance_scale=config.sample.guidance_scale, output_type="pt", height=config.resolution, width=config.resolution, noise_level=config.sample.noise_level, generator=generator ) latents = torch.stack( latents, dim=1 ) # (batch_size, num_steps + 1, 16, 96, 96) log_probs = torch.stack(log_probs, dim=1) # shape after stack (batch_size, num_steps) timesteps = pipeline.scheduler.timesteps.repeat( config.sample.train_batch_size, 1 ) # (batch_size, num_steps) # compute rewards asynchronously rewards = executor.submit(reward_fn, images, prompts, prompt_metadata, only_strict=True) # yield to to make sure reward computation starts time.sleep(0) samples.append( { "prompt_ids": prompt_ids, "prompt_embeds": prompt_embeds, "pooled_prompt_embeds": pooled_prompt_embeds, "timesteps": timesteps, "latents": latents[ :, :-1 ], # each entry is the latent before timestep t "next_latents": latents[ :, 1: ], # each entry is the latent after timestep t "log_probs": log_probs, "rewards": rewards, } ) # wait for all rewards to be computed for sample in tqdm( samples, desc="Waiting for rewards", disable=not accelerator.is_local_main_process, position=0, ): rewards, reward_metadata = sample["rewards"].result() # accelerator.print(reward_metadata) sample["rewards"] = { key: torch.as_tensor(value, device=accelerator.device).float() for key, value in rewards.items() } # collate samples into dict where each entry has shape (num_batches_per_epoch * sample.batch_size, ...) samples = { k: torch.cat([s[k] for s in samples], dim=0) if not isinstance(samples[0][k], dict) else { sub_key: torch.cat([s[k][sub_key] for s in samples], dim=0) for sub_key in samples[0][k] } for k in samples[0].keys() } if epoch % 10 == 0 and accelerator.is_main_process: # this is a hack to force wandb to log the images as JPEGs instead of PNGs with tempfile.TemporaryDirectory() as tmpdir: num_samples = min(15, len(images)) sample_indices = random.sample(range(len(images)), num_samples) for idx, i in enumerate(sample_indices): image = images[i] pil = Image.fromarray( (image.cpu().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 = [prompts[i] for i in sample_indices] sampled_rewards = [rewards['avg'][i] for i in sample_indices] wandb.log( { "images": [ wandb.Image( os.path.join(tmpdir, f"{idx}.jpg"), caption=f"{prompt:.100} | avg: {avg_reward:.2f}", ) for idx, (prompt, avg_reward) in enumerate(zip(sampled_prompts, sampled_rewards)) ], }, step=global_step, ) samples["rewards"]["ori_avg"] = samples["rewards"]["avg"] # The purpose of repeating `adv` along the timestep dimension here is to make it easier to introduce timestep-dependent advantages later, such as adding a KL reward. samples["rewards"]["avg"] = samples["rewards"]["avg"].unsqueeze(1).repeat(1, num_train_timesteps) # gather rewards across processes gathered_rewards = {key: accelerator.gather(value) for key, value in samples["rewards"].items()} gathered_rewards = {key: value.cpu().numpy() for key, value in gathered_rewards.items()} # log rewards and images if accelerator.is_main_process: # Flow-CoPD monitoring: print proxy/gold reward means to stdout each epoch # (reward-hacking = reward_pickscore up while reward_aesthetic/gold down). _rk = {k: float(v.mean()) for k, v in gathered_rewards.items() if '_accuracy' not in k and k not in ('avg',)} print(f"[REWARDS][E{epoch}] " + " ".join(f"{k}={v:.4f}" for k, v in _rk.items()), flush=True) wandb.log( { "epoch": epoch, **{f"reward_{key}": value.mean() for key, value in gathered_rewards.items() if '_strict_accuracy' not in key and '_accuracy' not in key}, }, step=global_step, ) # per-prompt mean/std tracking if config.per_prompt_stat_tracking: # gather the prompts across processes prompt_ids = accelerator.gather(samples["prompt_ids"]).cpu().numpy() prompts = pipeline.tokenizer.batch_decode( prompt_ids, skip_special_tokens=True ) advantages = stat_tracker.update(prompts, gathered_rewards['avg']) if accelerator.is_local_main_process: print("len(prompts)", len(prompts)) print("len unique prompts", len(set(prompts))) group_size, trained_prompt_num = stat_tracker.get_stats() zero_std_ratio, reward_std_mean = calculate_zero_std_ratio(prompts, gathered_rewards) if accelerator.is_main_process: wandb.log( { "group_size": group_size, "trained_prompt_num": trained_prompt_num, "zero_std_ratio": zero_std_ratio, "reward_std_mean": reward_std_mean, }, step=global_step, ) stat_tracker.clear() else: advantages = (gathered_rewards['avg'] - gathered_rewards['avg'].mean()) / (gathered_rewards['avg'].std() + 1e-4) # ungather advantages; we only need to keep the entries corresponding to the samples on this process advantages = torch.as_tensor(advantages) samples["advantages"] = ( advantages.reshape(accelerator.num_processes, -1, advantages.shape[-1])[accelerator.process_index] .to(accelerator.device) ) if accelerator.is_local_main_process: print("advantages: ", samples["advantages"].abs().mean()) del samples["rewards"] del samples["prompt_ids"] # Get the mask for samples where all advantages are zero across the time dimension mask = (samples["advantages"].abs().sum(dim=1) != 0) # If the number of True values in mask is not divisible by config.sample.num_batches_per_epoch, # randomly change some False values to True to make it divisible num_batches = config.sample.num_batches_per_epoch true_count = mask.sum() if true_count % num_batches != 0: false_indices = torch.where(~mask)[0] num_to_change = num_batches - (true_count % num_batches) if len(false_indices) >= num_to_change: random_indices = torch.randperm(len(false_indices))[:num_to_change] mask[false_indices[random_indices]] = True if accelerator.is_main_process: wandb.log( { "actual_batch_size": mask.sum().item()//config.sample.num_batches_per_epoch, }, step=global_step, ) # Filter out samples where the entire time dimension of advantages is zero samples = {k: v[mask] for k, v in samples.items()} total_batch_size, num_timesteps = samples["timesteps"].shape # assert ( # total_batch_size # == config.sample.train_batch_size * config.sample.num_batches_per_epoch # ) assert num_timesteps == config.sample.num_steps #################### TRAINING #################### for inner_epoch in range(config.train.num_inner_epochs): # shuffle samples along batch dimension perm = torch.randperm(total_batch_size, device=accelerator.device) samples = {k: v[perm] for k, v in samples.items()} # rebatch for training samples_batched = { k: v.reshape(-1, total_batch_size//config.sample.num_batches_per_epoch, *v.shape[1:]) for k, v in samples.items() } # dict of lists -> list of dicts for easier iteration samples_batched = [ dict(zip(samples_batched, x)) for x in zip(*samples_batched.values()) ] # train pipeline.transformer.train() info = defaultdict(list) for i, sample in tqdm( list(enumerate(samples_batched)), desc=f"Epoch {epoch}.{inner_epoch}: training", position=0, disable=not accelerator.is_local_main_process, ): if config.train.cfg: # concat negative prompts to sample prompts to avoid two forward passes embeds = torch.cat( [train_neg_prompt_embeds[:len(sample["prompt_embeds"])], sample["prompt_embeds"]] ) pooled_embeds = torch.cat( [train_neg_pooled_prompt_embeds[:len(sample["pooled_prompt_embeds"])], sample["pooled_prompt_embeds"]] ) else: embeds = sample["prompt_embeds"] pooled_embeds = sample["pooled_prompt_embeds"] train_timesteps = [step_index for step_index in range(num_train_timesteps)] for j in tqdm( train_timesteps, desc="Timestep", position=1, leave=False, disable=not accelerator.is_local_main_process, ): with accelerator.accumulate(transformer): # ===== Flow-CoPD: on-policy DISTILLATION loss (no policy gradient) ===== # unwrap WITHOUT accelerator.unwrap_model (which imports deepspeed and # crashes when CUDA_HOME is unset); mirror flow_grpo's direct .module access. _peft = transformer while hasattr(_peft, "module"): _peft = _peft.module _tname = config.copd.teacher_adapter_name t_gs = float(getattr(config.copd, "teacher_guidance_scale", 1.0)) # --- TEACHER velocity FIRST (frozen): eval() kills dropout (MAJOR#4), # no_grad, freed before the student graph to cut peak mem (#7) --- with torch.no_grad(): _was_training = _peft.training _peft.eval() _peft.set_adapter(_tname) with autocast(): v_t_cond = compute_velocity(transformer, sample, j, embeds, pooled_embeds) if t_gs != 1.0: # CFG-composed teacher (CRITICAL#2): distill the GUIDED teacher _ne = train_neg_prompt_embeds[:len(sample["prompt_embeds"])] _npe = train_neg_pooled_prompt_embeds[:len(sample["pooled_prompt_embeds"])] v_t_uncond = transformer( hidden_states=sample["latents"][:, j], timestep=sample["timesteps"][:, j], encoder_hidden_states=_ne, pooled_projections=_npe, return_dict=False, )[0] v_teacher = v_t_uncond + t_gs * (v_t_cond - v_t_uncond) else: v_teacher = v_t_cond _peft.set_adapter("default") _peft.train(_was_training) # --- STUDENT velocity (trainable "default" adapter) --- with autocast(): v_student = compute_velocity(transformer, sample, j, embeds, pooled_embeds) # per-step noise level (review #1/#5): scheduler.sigmas aligned to last (sampling) call sigma_j = pipeline.scheduler.sigmas[j].to(v_student.device) weight_t = sde_time_weight(sigma_j, scheme=config.copd.weight_scheme) if config.copd.mode == "positive": # A0: positive-only OPD (Flow-OPD / DiffusionOPD reproduction) loss = opd_positive_loss(v_student, v_teacher, weight_t) info["loss_pos"].append(loss.detach()) else: # A2: contrastive negative — repel low-reward trajectories from the # PER-STEP velocity (x_t - x0_bad)/sigma_t that regenerates them. v_neg = regen_velocity_per_step( sample["latents"][:, j], sample["next_latents"][:, -1], sigma_j ) loss, comp = copd_loss_fn( v_student, v_teacher, v_neg, sample["advantages"][:, j], weight_t=weight_t, lambda_neg=config.copd.lambda_neg, neg_clamp=config.copd.neg_clamp, adv_thresh=config.copd.adv_thresh, ) for _k, _v in comp.items(): info[_k].append(_v) info["loss"].append(loss.detach()) # backward pass accelerator.backward(loss) if accelerator.sync_gradients: accelerator.clip_grad_norm_( transformer.parameters(), config.train.max_grad_norm ) optimizer.step() optimizer.zero_grad() # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: # assert (j == train_timesteps[-1]) and ( # i + 1 # ) % config.train.gradient_accumulation_steps == 0 # log training-related stuff info = {k: torch.mean(torch.stack(v)) for k, v in info.items()} info = accelerator.reduce(info, reduction="mean") info.update({"epoch": epoch, "inner_epoch": inner_epoch}) if accelerator.is_main_process: wandb.log(info, step=global_step) global_step += 1 info = defaultdict(list) if config.train.ema: ema.step(transformer_trainable_parameters, global_step) # make sure we did an optimization step at the end of the inner epoch # assert accelerator.sync_gradients epoch+=1 if __name__ == "__main__": app.run(main)