import dotenv dotenv.load_dotenv(override=True) from typing import Union, List, Optional, Tuple import time from contextlib import contextmanager from copy import deepcopy import argparse from collections import defaultdict import logging import math import os import json import random import shutil from functools import partial from pathlib import Path from omegaconf import OmegaConf from tqdm.auto import tqdm import numpy as np import matplotlib.pyplot as plt import torch import torch.nn.functional as F import torch.utils.checkpoint from torchvision.transforms.functional import crop, to_pil_image, to_tensor from einops import repeat, rearrange import accelerate from accelerate import Accelerator from accelerate.state import AcceleratorState from accelerate.logging import get_logger from accelerate.utils import ProjectConfiguration, set_seed, DataLoaderConfiguration from accelerate import init_empty_weights from accelerate.utils import gather_object import transformers from transformers import AutoTokenizer, AutoProcessor from transformers import Qwen2_5_VLModel as TextEncoder import diffusers from diffusers.optimization import get_scheduler from diffusers.utils.torch_utils import is_compiled_module from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from peft import LoraConfig from omnigen2.training_utils import EMAModel from omnigen2.utils.logging_utils import TqdmToLogger from omnigen2.utils.tensor_util import pad_to_length, expand_as from omnigen2.dataset.omnigen2_train_dataset import OmniGen2TrainDataset, OmniGen2Collator, RepeatedDistributedBatchSampler from omnigen2.models.transformers.transformer_omnigen2 import OmniGen2Transformer2DModel from omnigen2.models.transformers.repo import OmniGen2RotaryPosEmbed from omnigen2.grpo.reward_client_edit import evaluate_images from omnigen2.grpo.utils import forward_logprob, process_grpo_rewards, compute_single_step_ppo_loss from omnigen2.pipelines.omnigen2.pipeline_omnigen2 import FMPipelineOutput logger = get_logger(__name__) def parse_args(root_path) -> OmegaConf: parser = argparse.ArgumentParser(description="OmniGen2 training script") parser.add_argument( "--config", type=str, required=True, help="Path to configuration file (YAML format)", ) parser.add_argument( "--global_batch_size", type=int, default=None, help="Global batch size.", ) parser.add_argument( "--data_path", type=str, default=None, help="Data path.", ) args = parser.parse_args() conf = OmegaConf.load(args.config) output_dir = os.path.join(root_path, 'experiments', conf.name) conf.root_dir = root_path conf.output_dir = output_dir conf.config_file = args.config # Override config with command line arguments if args.global_batch_size is not None: conf.train.global_batch_size = args.global_batch_size if args.data_path is not None: conf.data.data_path = args.data_path return conf def setup_logging(args: OmegaConf, accelerator: Accelerator) -> None: """ Set up logging configuration for training. Args: accelerator: Accelerator instance args: Configuration object logging_dir: Directory for log files """ logging_dir = Path(args.output_dir, "logs") if accelerator.is_main_process: if args.output_dir is not None: os.makedirs(args.output_dir, exist_ok=True) shutil.copy(args.config_file, args.output_dir) # Create logging directory and file handler os.makedirs(logging_dir, exist_ok=True) log_file = Path(logging_dir, f'{time.strftime("%Y%m%d-%H%M%S")}.log') formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(name)s - %(message)s') file_handler = logging.FileHandler(log_file, 'w') file_handler.setFormatter(formatter) file_handler.setLevel(logging.INFO) logger.logger.addHandler(file_handler) # Configure basic logging logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) # Set verbosity for different processes log_level = logging.INFO if accelerator.is_local_main_process else logging.ERROR transformers.utils.logging.set_verbosity(log_level) diffusers.utils.logging.set_verbosity(log_level) def log_model_info(name: str, model: torch.nn.Module): """Logs parameter counts for a given model.""" total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"--- {name} ---") logger.info(model) logger.info(f"Total parameters (M): {total_params / 1e6:.2f}") logger.info(f"Trainable parameters (M): {trainable_params / 1e6:.2f}") def get_qwen2_prompt_embeds( text_encoder, tokenizer, prompt: Union[str, List[str]], device: Optional[torch.device] = None, max_sequence_length: int = 256, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Get prompt embeddings from the Qwen2 text encoder. Args: prompt: The prompt or list of prompts to encode. device: The device to place the embeddings on. If None, uses the pipeline's device. max_sequence_length: Maximum sequence length for tokenization. Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - The prompt embeddings tensor - The attention mask tensor Raises: Warning: If the input text is truncated due to sequence length limitations. """ prompt = [prompt] if isinstance(prompt, str) else prompt text_inputs = tokenizer( prompt, padding="longest", max_length=max_sequence_length, truncation=True, return_tensors="pt", ) text_input_ids = text_inputs.input_ids.to(device) untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device) if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): removed_text = tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1]) logger.warning( "The following part of your input was truncated because Gemma can only handle sequences up to" f" {max_sequence_length} tokens: {removed_text}" ) prompt_attention_mask = text_inputs.attention_mask.to(device) prompt_embeds = text_encoder( input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=False, ).last_hidden_state return prompt_embeds, prompt_attention_mask @contextmanager def disabled_adapters(model): try: # model.disable_adapters() from peft.tuners.tuners_utils import BaseTunerLayer for _, module in model.named_modules(): if isinstance(module, BaseTunerLayer): if hasattr(module, "enable_adapters"): module._disable_adapters = True else: module.disable_adapters = True yield finally: # model.enable_adapters() from peft.tuners.tuners_utils import BaseTunerLayer for _, module in model.named_modules(): if isinstance(module, BaseTunerLayer): if hasattr(module, "enable_adapters"): module._disable_adapters = False else: module.disable_adapters = False def main(args): accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=Path(args.output_dir, 'logs')) accelerator = Accelerator( gradient_accumulation_steps=args.train.gradient_accumulation_steps * math.ceil( args.train.rl.num_inference_step * args.train.rl.get("train_timesteps_fraction", 1.0) ), mixed_precision=args.train.mixed_precision, log_with=OmegaConf.to_object(args.logger.log_with), project_config=accelerator_project_config, dataloader_config=DataLoaderConfiguration(split_batches=True), ) setup_logging(args, accelerator) # Reproducibility if args.seed is not None: set_seed(args.seed, device_specific=args.get('device_specific_seed', False)) # Set performance flags if args.train.allow_tf32: torch.backends.cuda.matmul.allow_tf32 = True if args.train.get('benchmark_cudnn', False): torch.backends.cudnn.benchmark = True weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 def unwrap_model(model): model = accelerator.unwrap_model(model) model = model._orig_mod if is_compiled_module(model) else model return model ema_decay = args.train.get('ema_decay', 0) if args.model.pretrained_model_path: with init_empty_weights(): model = OmniGen2Transformer2DModel(**args.model.arch_opt) state_dict = torch.load(args.model.pretrained_model_path, mmap=True, weights_only=True) missing, unexpect = model.load_state_dict(state_dict, assign=True, strict=False) else: model = OmniGen2Transformer2DModel(**args.model.arch_opt) model.train() freqs_cis = OmniGen2RotaryPosEmbed.get_freqs_cis( model.config.axes_dim_rope, model.config.axes_lens, theta=10000, ) if ema_decay != 0: model_ema = deepcopy(model) model_ema._requires_grad = False processor = AutoProcessor.from_pretrained(args.model.pretrained_text_encoder_model_name_or_path) text_tokenizer = processor.tokenizer text_tokenizer.padding_side = "right" if accelerator.is_main_process: text_tokenizer.save_pretrained(os.path.join(args.output_dir, 'tokenizer')) text_encoder = TextEncoder.from_pretrained( args.model.pretrained_text_encoder_model_name_or_path, torch_dtype=weight_dtype, ) if args.model.get('resize_token_embeddings', False): text_encoder.resize_token_embeddings(len(text_tokenizer)) if accelerator.is_main_process: text_encoder.save_pretrained(os.path.join(args.output_dir, 'text_encoder')) log_model_info("text_encoder", text_encoder) vae = AutoencoderKL.from_pretrained( args.model.pretrained_vae_model_name_or_path, subfolder=args.model.get("vae_subfolder", "vae"), local_files_only=True, ) logger.info(vae) logger.info("***** Move vae, text_encoder to device and cast to weight_dtype *****") # Move vae, unet, text_encoder and controlnet_ema to device and cast to weight_dtype vae = vae.to(accelerator.device, dtype=weight_dtype) text_encoder = text_encoder.to(accelerator.device, dtype=weight_dtype) args.train.lora_ft = args.train.get('lora_ft', False) if args.train.lora_ft: model.requires_grad_(False) target_modules = ["to_k", "to_q", "to_v", "to_out.0"] lora_config = LoraConfig( r=args.train.lora_rank, lora_alpha=args.train.lora_alpha, lora_dropout=args.train.lora_dropout, init_lora_weights="gaussian", target_modules=target_modules, ) model.add_adapter(lora_config) if args.train.gradient_checkpointing: model.enable_gradient_checkpointing() if args.train.scale_lr: args.train.learning_rate = ( args.train.learning_rate * args.train.gradient_accumulation_steps * args.train.batch_size * accelerator.num_processes ) # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs if args.train.use_8bit_adam: try: import bitsandbytes as bnb except ImportError: raise ImportError( "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." ) optimizer_class = bnb.optim.AdamW8bit else: optimizer_class = torch.optim.AdamW log_model_info("transformer", model) # Optimizer creation trainable_params = list(filter(lambda p: p.requires_grad, model.parameters())) optimizer = optimizer_class( trainable_params, lr=args.train.learning_rate, betas=(args.train.adam_beta1, args.train.adam_beta2), weight_decay=args.train.adam_weight_decay, eps=args.train.adam_epsilon, ) logger.info("***** Prepare dataset *****") with accelerator.main_process_first(): train_dataset = OmniGen2TrainDataset( args.data.data_path, tokenizer=text_tokenizer, num_workers=args.train.dataloader_num_workers, use_chat_template=args.data.use_chat_template, prompt_dropout_prob=args.data.get('prompt_dropout_prob', 0.0), ref_img_dropout_prob=args.data.get('ref_img_dropout_prob', 0.0), max_input_pixels=OmegaConf.to_object(args.data.get('max_input_pixels', 1024 * 1024)), max_output_pixels=args.data.get('max_output_pixels', 1024 * 1024), max_side_length=args.data.get('max_side_length', 2048), ) logger.info(f"Number of training samples: {len(train_dataset)}") if args.seed is not None and args.get("workder_specific_seed", False): from omnigen2.utils.reproducibility import worker_init_fn worker_init_fn = partial( worker_init_fn, num_processes=AcceleratorState().num_processes, num_workers=args.train.dataloader_num_workers, process_index=AcceleratorState().process_index, seed=args.seed, same_seed_per_epoch=args.get("same_seed_per_epoch", False), ) else: worker_init_fn = None logger.info("***** Prepare dataLoader *****") train_dataloader = torch.utils.data.DataLoader( train_dataset, num_workers=args.train.dataloader_num_workers, batch_sampler=RepeatedDistributedBatchSampler( dataset=train_dataset, batch_size=args.train.batch_size, num_repeats=args.train.rl.num_images_per_prompt, num_replicas=AcceleratorState().num_processes, rank=AcceleratorState().process_index, shuffle=True, seed=args.seed, drop_last=True, ), worker_init_fn=worker_init_fn, collate_fn=OmniGen2Collator(tokenizer=text_tokenizer, max_token_len=args.data.maximum_text_tokens) ) logger.info(f"{args.train.batch_size=} {args.train.gradient_accumulation_steps=} {accelerator.num_processes=} {args.train.global_batch_size=}") assert args.train.batch_size % (args.train.rl.batch_size_per_forward * args.train.gradient_accumulation_steps) == 0, f"{args.train.batch_size=} % ({args.train.rl.batch_size_per_forward=} * {args.train.rl.gradient_accumulation_steps=}) != 0" assert args.train.batch_size // (args.train.rl.batch_size_per_forward * args.train.gradient_accumulation_steps) == args.train.rl.num_update_steps_per_sampling assert args.train.global_batch_size // args.train.rl.num_images_per_prompt == args.train.rl.num_unique_prompts_per_sampling, f"{args.train.global_batch_size=} // {args.train.rl.num_images_per_prompt=} != {args.train.rl.num_unique_prompts_per_sampling=}" # Scheduler and math around the number of training steps. overrode_max_train_steps = False num_update_steps_per_epoch = math.ceil(len(train_dataloader) * args.train.rl.num_update_steps_per_sampling) if 'max_train_steps' not in args.train: args.train.max_train_steps = args.train.num_train_epochs * num_update_steps_per_epoch overrode_max_train_steps = True if args.train.lr_scheduler == 'timm_cosine': from omnigen2.optim.scheduler.cosine_lr import CosineLRScheduler lr_scheduler = CosineLRScheduler(optimizer=optimizer, t_initial=args.train.t_initial, lr_min=args.train.lr_min, cycle_decay=args.train.cycle_decay, warmup_t=args.train.warmup_t, warmup_lr_init=args.train.warmup_lr_init, warmup_prefix=args.train.warmup_prefix, t_in_epochs=args.train.t_in_epochs) elif args.train.lr_scheduler == 'timm_constant_with_warmup': from omnigen2.optim.scheduler.step_lr import StepLRScheduler lr_scheduler = StepLRScheduler( optimizer=optimizer, decay_t=1, decay_rate=1, warmup_t=args.train.warmup_t, warmup_lr_init=args.train.warmup_lr_init, warmup_prefix=args.train.warmup_prefix, t_in_epochs=args.train.t_in_epochs, ) else: lr_scheduler = get_scheduler( args.train.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.train.lr_warmup_steps, num_training_steps=args.train.max_train_steps, num_cycles=args.train.lr_num_cycles, power=args.train.lr_power, ) logger.info("***** Prepare everything with our accelerator *****") if args.train.ema_decay != 0: model, model_ema, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( model, model_ema, optimizer, train_dataloader, lr_scheduler ) model_ema = EMAModel(model_ema.parameters(), decay=ema_decay, model_cls=type(unwrap_model(model)), model_config=model_ema.config) else: model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( model, optimizer, train_dataloader, lr_scheduler ) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) * args.train.rl.num_update_steps_per_sampling) if overrode_max_train_steps: args.train.max_train_steps = args.train.num_train_epochs * num_update_steps_per_epoch # Afterwards we recalculate our number of training epochs args.train.num_train_epochs = math.ceil(args.train.max_train_steps / num_update_steps_per_epoch) # We need to initialize the trackers we use, and also store our configuration. # The trackers initializes automatically on the main process. if accelerator.is_main_process: accelerator.init_trackers("OmniGen2-RL", init_kwargs={"wandb": {"name": args.name}}) # Train! total_batch_size = args.train.batch_size * accelerator.num_processes * args.train.gradient_accumulation_steps logger.info("***** Running training *****") logger.info(f" Num examples = {len(train_dataset)}") logger.info(f" Num batches each epoch = {len(train_dataloader)}") logger.info(f" Num Epochs = {args.train.num_train_epochs}") logger.info(f" Instantaneous batch size per device = {args.train.batch_size}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logger.info(f" Gradient Accumulation steps = {args.train.gradient_accumulation_steps}") logger.info(f" Total optimization steps = {args.train.max_train_steps}") global_step = 0 first_epoch = 0 # Potentially load in the weights and states from a previous save if args.resume_from_checkpoint: if args.resume_from_checkpoint != "latest": path = os.path.basename(args.resume_from_checkpoint) else: # Get the most recent checkpoint dirs = os.listdir(args.output_dir) dirs = [d for d in dirs if d.startswith("checkpoint")] dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) path = dirs[-1] if len(dirs) > 0 else None if path is None: accelerator.print( f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." ) args.resume_from_checkpoint = None initial_global_step = 0 else: accelerator.print(f"Resuming from checkpoint {path}") accelerator.load_state(os.path.join(args.output_dir, path)) global_step = int(path.split("-")[1]) initial_global_step = global_step first_epoch = global_step // num_update_steps_per_epoch else: initial_global_step = 0 progress_bar = tqdm( range(0, args.train.max_train_steps), initial=initial_global_step, desc="Steps", # Only show the progress bar once on each machine. disable=not accelerator.is_local_main_process, file=TqdmToLogger(logger, level=logging.INFO) ) if accelerator.is_main_process: for tracker in accelerator.trackers: if tracker.name == "wandb": logger.info(f"***** Wandb log dir: {tracker.run.dir} *****") from omnigen2.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline from omnigen2.schedulers.scheduling_flow_match_euler_maruyama_discrete import FlowMatchEulerMaruyamaDiscreteScheduler pipeline = OmniGen2Pipeline( transformer=model, vae=vae, scheduler=FlowMatchEulerMaruyamaDiscreteScheduler( sigma_coef=args.train.rl.get('sigma_coef', 0.7), time_shift_base_res=args.train.rl.get('time_shift_base_res', 320) ), mllm=None, processor=processor, ) pipeline.set_progress_bar_config(disable=True) with torch.no_grad(): if args.train.rl.get('use_ori_neg_prompt_template', False): negative_prompt = [ { "role": "system", "content": "You are a helpful assistant.", }, {"role": "user", "content": args.train.rl.negative_prompt}, ] negative_prompt = pipeline.processor.tokenizer.apply_chat_template( negative_prompt, tokenize=False, add_generation_prompt=False ) else: negative_prompt = pipeline._apply_chat_template(args.train.rl.negative_prompt) negative_prompt_embeds, negative_prompt_attention_mask = get_qwen2_prompt_embeds( text_encoder=text_encoder, tokenizer=text_tokenizer, prompt=negative_prompt, device=accelerator.device, max_sequence_length=1024, ) seq_len = negative_prompt_embeds.shape[1] negative_prompt_embeds = negative_prompt_embeds.repeat(1, args.train.rl.batch_size_per_forward, 1) negative_prompt_embeds = negative_prompt_embeds.view(args.train.rl.batch_size_per_forward, seq_len, -1) negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(args.train.rl.batch_size_per_forward, 1) negative_prompt_attention_mask = negative_prompt_attention_mask.view( args.train.rl.batch_size_per_forward, -1 ) ref_latents_N, ref_img_mask_N, l_effective_ref_img_len_N, ref_img_sizes_N = pipeline.transformer.flat_and_pad_to_seq_ref_img(None, args.train.rl.batch_size_per_forward, weight_dtype, accelerator.device) reward_server_config = OmegaConf.load(args.reward_server_config) for epoch in range(first_epoch, args.train.num_train_epochs): if 'max_train_steps' in args.train and global_step >= args.train.max_train_steps: break train_dataloader.batch_sampler.batch_sampler.set_epoch(epoch) for step, batch in enumerate(train_dataloader): instruction = batch['instruction'] input_images = batch['input_images'] input_images_pil = batch['input_images_pil'] target_img_size = batch['target_img_size'] text_mask = batch['text_mask'] text_input_ids = batch['text_ids'] total_results = None total_text_feats = None batch_size_per_forward = args.train.rl.batch_size_per_forward for i in range(args.train.batch_size // batch_size_per_forward): with torch.no_grad(): text_feats = text_encoder( input_ids=text_input_ids[i*batch_size_per_forward:(i+1)*batch_size_per_forward], attention_mask=text_mask[i*batch_size_per_forward:(i+1)*batch_size_per_forward], output_hidden_states=False, ).last_hidden_state results = pipeline( prompt_embeds=text_feats, prompt_attention_mask=text_mask[i*batch_size_per_forward:(i+1)*batch_size_per_forward], negative_prompt_embeds=negative_prompt_embeds, negative_prompt_attention_mask=negative_prompt_attention_mask, input_images=input_images[i*batch_size_per_forward:(i+1)*batch_size_per_forward], size=target_img_size[i*batch_size_per_forward:(i+1)*batch_size_per_forward], num_inference_steps=args.train.rl.num_inference_step, max_sequence_length=1024, text_guidance_scale=args.train.rl.text_guidance_scale, image_guidance_scale=args.train.rl.image_guidance_scale, cfg_range=(args.train.rl.cfg_range_start, args.train.rl.cfg_range_end), num_images_per_prompt=1, output_type="pil", enable_parallel_cfg=True, return_middle_statistics=True, mixed_precision=True, do_normalize=False ) if i == 0: total_text_feats = text_feats total_results = results for k in ['img_mask', 'ref_latents', 'ref_img_mask', 'middle_latents']: total_results.__dict__[k] = [total_results.__dict__[k]] else: total_text_feats = torch.cat([total_text_feats, text_feats], dim=0) for k in ['images', 'l_effective_img_len', 'img_sizes', 'l_effective_ref_img_len', 'ref_img_sizes']: total_results.__dict__[k].extend(results.__dict__[k]) for k in ['img_mask', 'ref_latents', 'ref_img_mask', 'middle_latents']: total_results.__dict__[k].append(results.__dict__[k]) for i in range(len(results.log_probs)): total_results.log_probs[i] = torch.cat([total_results.log_probs[i], results.log_probs[i]], dim=0) for i in range(len(batch['meta_data'])): json_data = json.loads(batch['meta_data'][i]) json_data['id'] = f"{global_step * args.train.global_batch_size + accelerator.process_index * args.train.batch_size + i}" batch['meta_data'][i] = json.dumps(json_data) local_batch_size = len(input_images_pil) gathered_input_images_pil = gather_object(input_images_pil) gathered_output_images = gather_object(total_results.images) gathered_meta_data = gather_object(batch['meta_data']) if accelerator.is_main_process: scores, rewards, reasoning, meta_data = evaluate_images( input_images=gathered_input_images_pil, output_image=gathered_output_images, meta_datas=gathered_meta_data, proxy_host=reward_server_config.server.hosts[0], proxy_port=reward_server_config.server.proxy_port, server_type=args.train.rl.get('server_type', 'vlm') ) rewards_to_scatter = [rewards[i:i + local_batch_size] for i in range(0, len(rewards), local_batch_size)] reasoning_to_scatter = [reasoning[i:i + local_batch_size] for i in range(0, len(reasoning), local_batch_size)] meta_data_to_scatter = [meta_data[i:i + local_batch_size] for i in range(0, len(meta_data), local_batch_size)] else: rewards_to_scatter = [None for _ in range(accelerator.num_processes)] reasoning_to_scatter = [None for _ in range(accelerator.num_processes)] meta_data_to_scatter = [None for _ in range(accelerator.num_processes)] accelerator.wait_for_everyone() # Extract the current process’s own rewards, reasoning, and meta_data. rewards = [None] reasoning = [None] meta_data = [None] torch.distributed.scatter_object_list(rewards, rewards_to_scatter) torch.distributed.scatter_object_list(reasoning, reasoning_to_scatter) torch.distributed.scatter_object_list(meta_data, meta_data_to_scatter) rewards = rewards[0] reasoning = reasoning[0] meta_data = meta_data[0] assert len(rewards) == len(reasoning) == len(meta_data) == local_batch_size rewards = torch.tensor(rewards, dtype=torch.float32, device=accelerator.device) advantages, prompt_stats = process_grpo_rewards( rewards=rewards, prompts=instruction, accelerator=accelerator, std_level=args.train.rl.get('std_level', 'group'), ) reuse_samples_nums = args.train.rl.reuse_samples_nums # reuse times of samples clip_range = args.train.rl.clip_range # PPO clip range # prepare data for GRPO timesteps = pipeline.scheduler._timesteps # [batch_size, num_timesteps+1] assert reuse_samples_nums == 1 for reuse_step in range(reuse_samples_nums): logs = defaultdict(list) for forward_step in range(args.train.batch_size // batch_size_per_forward): results = FMPipelineOutput( images=[], middle_latents=[], log_probs=[], img_mask=[], l_effective_img_len=[], img_sizes=[], ref_latents=[], ref_img_mask=[], l_effective_ref_img_len=[], ref_img_sizes=[], ) for k in ['images', 'l_effective_img_len', 'img_sizes', 'l_effective_ref_img_len', 'ref_img_sizes']: results.__dict__[k] = total_results.__dict__[k][forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward] for k in ['img_mask', 'ref_latents', 'ref_img_mask', 'middle_latents']: results.__dict__[k] = total_results.__dict__[k][forward_step] results.log_probs = [total_results.log_probs[i][forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward] for i in range(len(total_results.log_probs))] text_feats = total_text_feats[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward] old_log_probs = [total_results.log_probs[i][forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward] for i in range(len(total_results.log_probs))] train_timesteps = list(range(args.train.rl.num_inference_step)) sample_steps = math.ceil(args.train.rl.num_inference_step * args.train.rl.get('train_timesteps_fraction', 1.0)) train_timesteps = sorted(random.sample(train_timesteps, k=sample_steps)) if args.train.rl.policy_loss_reweighting: sigma_ts = [] for i in train_timesteps: t = timesteps[0, i] t_next = timesteps[0, i+1] sigma_t = pipeline.scheduler.get_sigma_t(t, t_next if i == 0 else None) # [batch_size] dt = t_next - t sigma_ts.append(sigma_t * math.sqrt(dt)) sigma_ts = torch.stack(sigma_ts) normalize_factor = sigma_ts.mean() for idx, i in enumerate(train_timesteps): with accelerator.accumulate(model): text_guidance_scale = args.train.rl.text_guidance_scale if args.train.rl.cfg_range_start <= i / args.train.rl.num_inference_step <= args.train.rl.cfg_range_end else 1.0 image_guidance_scale = args.train.rl.image_guidance_scale if args.train.rl.cfg_range_start <= i / args.train.rl.num_inference_step <= args.train.rl.cfg_range_end else 1.0 latents = results.middle_latents[i] latents_next = results.middle_latents[i+1] t = timesteps[:, i] t_next = timesteps[:, i+1] model_kwargs = dict( hidden_states=latents, timestep=t, freqs_cis=freqs_cis, flat_and_pad=False, img_mask=results.img_mask, l_effective_img_len=results.l_effective_img_len, img_sizes=results.img_sizes, ) model_pred_kwargs = dict( text_hidden_states=text_feats, text_attention_mask=text_mask[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward], ref_image_hidden_states=results.ref_latents, ref_img_mask=results.ref_img_mask, l_effective_ref_img_len=results.l_effective_ref_img_len, ref_img_sizes=results.ref_img_sizes, ) if image_guidance_scale > 1 and text_guidance_scale > 1: model_kwargs['hidden_states'] = torch.cat([latents, latents, latents], dim=0) model_kwargs['timestep'] = torch.cat([t, t, t], dim=0) model_kwargs['img_mask'] = torch.cat([results.img_mask, results.img_mask, results.img_mask], dim=0) model_kwargs['l_effective_img_len'] = results.l_effective_img_len * 3 model_kwargs['img_sizes'] = results.img_sizes * 3 model_pred_kwargs['text_hidden_states'] = torch.cat([text_feats, pad_to_length(negative_prompt_embeds, len=text_feats.shape[1]), pad_to_length(negative_prompt_embeds, len=text_feats.shape[1])], dim=0) model_pred_kwargs['text_attention_mask'] = torch.cat([text_mask[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward], pad_to_length(negative_prompt_attention_mask, len=text_mask.shape[1]), pad_to_length(negative_prompt_attention_mask, len=text_mask.shape[1])], dim=0) model_pred_kwargs['ref_image_hidden_states'] = torch.cat([results.ref_latents, results.ref_latents, pad_to_length(ref_latents_N, len=results.ref_latents.shape[1])], dim=0) model_pred_kwargs['ref_img_mask'] = torch.cat([results.ref_img_mask, results.ref_img_mask, pad_to_length(ref_img_mask_N, len=results.ref_img_mask.shape[1])], dim=0) model_pred_kwargs['l_effective_ref_img_len'] = results.l_effective_ref_img_len * 2 + l_effective_ref_img_len_N model_pred_kwargs['ref_img_sizes'] = results.ref_img_sizes * 2 + ref_img_sizes_N elif text_guidance_scale > 1: model_kwargs['hidden_states'] = torch.cat([latents, latents], dim=0) model_kwargs['timestep'] = torch.cat([t, t], dim=0) model_kwargs['img_mask'] = torch.cat([results.img_mask, results.img_mask], dim=0) model_kwargs['l_effective_img_len'] = results.l_effective_img_len * 2 model_kwargs['img_sizes'] = results.img_sizes * 2 model_pred_kwargs['text_hidden_states'] = torch.cat([text_feats, pad_to_length(negative_prompt_embeds, len=text_feats.shape[1])], dim=0) model_pred_kwargs['text_attention_mask'] = torch.cat([text_mask[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward], pad_to_length(negative_prompt_attention_mask, len=text_mask.shape[1])], dim=0) model_pred_kwargs['ref_image_hidden_states'] = torch.cat([results.ref_latents, pad_to_length(ref_latents_N, len=results.ref_latents.shape[1])], dim=0) model_pred_kwargs['ref_img_mask'] = torch.cat([results.ref_img_mask, pad_to_length(ref_img_mask_N, len=results.ref_img_mask.shape[1])], dim=0) model_pred_kwargs['l_effective_ref_img_len'] = results.l_effective_ref_img_len + l_effective_ref_img_len_N model_pred_kwargs['ref_img_sizes'] = results.ref_img_sizes + ref_img_sizes_N step_log_probs, mean_t, sigma_t = forward_logprob( latents=latents, latents_next=latents_next, t=t, t_next=t_next, step_index=i, img_mask=results.img_mask, model=model, model_kwargs=model_kwargs, model_pred_kwargs=model_pred_kwargs, scheduler=pipeline.scheduler, apply_cfg=args.train.rl.apply_cfg_in_training, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, ) if args.train.rl.kl_loss_weight > 0: with torch.no_grad(): with disabled_adapters(unwrap_model(model)): _, mean_t_ref, _ = forward_logprob( latents=latents, latents_next=latents_next, t=t, t_next=t_next, step_index=i, img_mask=results.img_mask, model=model, model_kwargs=model_kwargs, model_pred_kwargs=model_pred_kwargs, scheduler=pipeline.scheduler, apply_cfg=args.train.rl.apply_cfg_in_training, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, ) loss = 0 ( policy_loss, policy_clip_frac, approx_kl, unclipped_loss, clipped_loss, ratio, ratio_positive, ratio_negative, num_positive, num_negative, ratio_large_than_1, ratio_small_than_1, ) = compute_single_step_ppo_loss( step_log_probs=step_log_probs, old_step_log_probs=old_log_probs[i], advantages=advantages[ forward_step * batch_size_per_forward : ( forward_step + 1 ) * batch_size_per_forward ], clip_range=clip_range, adv_clip_max=args.train.rl.adv_clip_max, ) logs['policy_loss'].append(policy_loss.detach()) logs['policy_clip_frac'].append(policy_clip_frac.detach()) logs['approx_kl'].append(approx_kl.detach()) logs['advantages'].append(advantages[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward].mean().detach()) logs['advantages_std'].append(advantages[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward].std().detach()) logs['advantages_min'].append(advantages[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward].min().detach()) logs['advantages_max'].append(advantages[forward_step*batch_size_per_forward:(forward_step+1)*batch_size_per_forward].max().detach()) logs['policy_loss_unclipped'].append(unclipped_loss.mean().detach()) logs['policy_loss_clipped'].append(clipped_loss.mean().detach()) logs['ratio'].append(ratio.mean().detach()) logs['ratio_positive'].append(ratio_positive.detach()) logs['ratio_negative'].append(ratio_negative.detach()) logs['num_positive'].append(num_positive.detach()) logs['num_negative'].append(num_negative.detach()) logs['ratio_large_than_1'].append(ratio_large_than_1.detach()) logs['ratio_small_than_1'].append(ratio_small_than_1.detach()) if args.train.rl.policy_loss_reweighting: policy_loss = policy_loss * sample_steps * (sigma_ts[idx] / normalize_factor) loss += policy_loss if args.train.rl.kl_loss_weight > 0: kl_loss = torch.mean( ((mean_t - mean_t_ref.detach()) ** 2) .flatten(start_dim=1) .mean(dim=1) / (2 * sigma_t) ) logs['kl_loss'].append(kl_loss.detach()) loss += args.train.rl.kl_loss_weight * kl_loss logs['loss'].append(loss.detach()) accelerator.backward(loss) if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(model.parameters(), args.train.max_grad_norm) logs['grad_norm'].append(grad_norm.to(accelerator.device).detach()) optimizer.step() if 'timm' in args.train.lr_scheduler: lr_scheduler.step(global_step) else: lr_scheduler.step() optimizer.zero_grad(set_to_none=args.train.set_grads_to_none) # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: logs = {k: torch.mean(torch.stack(v)) for k, v in logs.items()} logs = accelerator.reduce(logs, reduction="mean") logs = {k: v.item() for k, v in logs.items()} logs.update( { "lr": lr_scheduler.get_last_lr()[0], "rewards_min": np.mean( [v["min"] for k, v in prompt_stats.items()] ), "rewards_max": np.mean( [v["max"] for k, v in prompt_stats.items()] ), "rewards_mean": np.mean( [v["mean"] for k, v in prompt_stats.items()] ), "rewards_std": np.mean( [v["std"] for k, v in prompt_stats.items()] ), "zero_std_ratio": np.mean( [ v["std"] == 0 for k, v in prompt_stats.items() ] ) } ) if ema_decay != 0: model_ema.step(model.parameters()) global_step += 1 if global_step % args.logger.checkpointing_steps == 0: if accelerator.is_main_process: if args.logger.checkpoints_total_limit is not None: checkpoints = os.listdir(args.output_dir) checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) if len(checkpoints) >= args.logger.checkpoints_total_limit: num_to_remove = len(checkpoints) - args.logger.checkpoints_total_limit + 1 removing_checkpoints = checkpoints[0:num_to_remove] logger.info( f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" ) logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") for removing_checkpoint in removing_checkpoints: removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) shutil.rmtree(removing_checkpoint) accelerator.wait_for_everyone() save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path) logger.info(f"Saved state to {save_path}") if 'train_visualization_interval' in args.val and (global_step - 1) % args.val.train_visualization_interval == 0: num_samples = min(args.val.get('num_train_visualization_samples', 2), args.train.batch_size) if accelerator.is_main_process: target_instruction = instruction[:num_samples] else: target_instruction = [None] * num_samples torch.distributed.broadcast_object_list(target_instruction) rewards_per_instruction = [[] for _ in range(num_samples)] for i in range(len(target_instruction)): for j in range(len(instruction)): if target_instruction[i] == instruction[j]: rewards_per_instruction[i].append((rewards[j].item(), accelerator.process_index, j)) gathered_rewards_per_instruction = [None for _ in range(accelerator.num_processes)] torch.distributed.all_gather_object(gathered_rewards_per_instruction, rewards_per_instruction) gathered_rewards_per_instruction_flat = [[] for _ in range(num_samples)] for i in range(num_samples): for j in range(accelerator.num_processes): gathered_rewards_per_instruction_flat[i].extend(gathered_rewards_per_instruction[j][i]) p = [{} for _ in range(num_samples)] for i in range(len(target_instruction)): gathered_rewards_per_instruction_flat[i].sort(key=lambda x: x[0]) for j in range(len(gathered_rewards_per_instruction_flat[i])): if gathered_rewards_per_instruction_flat[i][j][1] == accelerator.process_index: p[i][gathered_rewards_per_instruction_flat[i][j][2]] = j with torch.no_grad(): for i in range(len(target_instruction)): cnt = 0 for j in range(len(instruction)): if instruction[j] == target_instruction[i]: if cnt == 0: for k in range(len(input_images_pil[j])): input_images_pil[j][k].save(os.path.join(args.output_dir, f"input_visualization_{global_step}_{i}_input_{k}.png")) total_results.images[j].save(os.path.join(args.output_dir, f"input_visualization_{global_step}_{i}_{p[i][j]}.png")) with open(os.path.join(args.output_dir, f"instruction_{global_step}_{i}_{p[i][j]}.txt"), "w", encoding='utf-8') as f: f.write(f"instruction: {instruction[j]}\nreward: {rewards[j].item()}\nadvantages: {advantages[j].item()}\nreasoning: {reasoning[j]}\nmeta_data: {batch['meta_data'][j]}\ncur_receive_meta_data: {meta_data[j]}") cnt += 1 progress_bar.set_postfix(**logs) progress_bar.update(1) accelerator.log(logs, step=global_step) logs = defaultdict(list) if 'max_train_steps' in args.train and global_step >= args.train.max_train_steps: break checkpoints = os.listdir(args.output_dir) checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) if len(checkpoints) > 0 and int(checkpoints[-1].split("-")[1]) < global_step: if accelerator.is_main_process: if args.logger.checkpoints_total_limit is not None: if len(checkpoints) >= args.logger.checkpoints_total_limit: num_to_remove = len(checkpoints) - args.logger.checkpoints_total_limit + 1 removing_checkpoints = checkpoints[0:num_to_remove] logger.info( f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" ) logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") for removing_checkpoint in removing_checkpoints: removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) shutil.rmtree(removing_checkpoint) accelerator.wait_for_everyone() save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path) logger.info(f"Saved state to {save_path}") accelerator.end_training() if __name__ == "__main__": root_path = os.path.abspath(os.path.join(__file__, os.path.pardir)) args = parse_args(root_path) main(args)