# Adapted from https://github.com/andimarafioti/florence2-finetuning/blob/main/distributed_train.py import argparse import os from functools import partial import logging import torch import math import transformers import diffusers import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from pathlib import Path from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoProcessor, get_scheduler import yaml from types import SimpleNamespace from peft import LoraConfig, get_peft_model from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, DistributedType from data import collate_fn, get_dataset logger = get_logger(__name__) def load_config(config_file): with open(config_file, "r") as f: config = yaml.safe_load(f) config = convert_config_types(config) return SimpleNamespace(**config) def convert_config_types(obj): if isinstance(obj, dict): return {k: convert_config_types(v) for k, v in obj.items()} elif isinstance(obj, list): return [convert_config_types(item) for item in obj] elif isinstance(obj, str): # Try to convert to int, then float; if both fail, keep as string. try: return int(obj) except ValueError: try: return float(obj) except ValueError: return obj else: return obj def create_data_loaders( train_dataset, val_dataset, batch_size, num_workers, processor, ): train_loader = DataLoader( train_dataset, batch_size=batch_size, pin_memory=True, shuffle=True, collate_fn=partial(collate_fn, processor=processor), num_workers=num_workers, ) val_loader = DataLoader( val_dataset, batch_size=batch_size // 2, pin_memory=True, collate_fn=partial(collate_fn, processor=processor), num_workers=num_workers, ) return train_loader, val_loader def forward_with_model(model, inputs, labels, weight_dtype=torch.float16): input_ids = inputs.input_ids pixel_values = inputs.pixel_values.to(weight_dtype) labels = labels return model(input_ids=input_ids, pixel_values=pixel_values, labels=labels) def evaluate_model(model, val_loader, device, global_step, max_val_item_count, weight_dtype, disable_pbar): # Evaluation phase model.eval() val_loss = 0 with torch.no_grad(): val_item_count = 0 for batch in tqdm(val_loader, desc=f"Evaluation at step {global_step}", disable=disable_pbar): val_item_count += len(batch) # Prepare the input and target tensors color_inputs, colors = batch["color_inputs"], batch["colors"] lighting_inputs, lightings = batch["lighting_inputs"], batch["lightings"] lighting_type_inputs, lighting_types = batch["lighting_type_inputs"], batch["lighting_types"] composition_inputs, compositions = batch["composition_inputs"], batch["compositions"] losses = [] for inputs, labels in [ (color_inputs, colors), (lighting_inputs, lightings), (lighting_type_inputs, lighting_types), (composition_inputs, compositions), ]: losses.append(forward_with_model(model, inputs, labels, weight_dtype=weight_dtype).loss) loss = torch.stack(losses).mean() val_loss += loss.item() if val_item_count > max_val_item_count: break avg_val_loss = val_loss / val_item_count model.train() return avg_val_loss def train_model(accelerator, args): # Load the dataset based on the dataset_name argument dataset = get_dataset( accelerator=accelerator, dataset_id=args.dataset_id, num_proc=args.num_proc, cache_dir=args.cache_dir ) with accelerator.main_process_first(): splits = dataset.train_test_split(0.1, seed=2025) train_dataset, val_dataset = splits["train"], splits["test"] with accelerator.main_process_first(): further_splits = val_dataset.train_test_split(0.1, seed=2025) val_dataset = further_splits["train"] # Load the model and processor ft_model = AutoModelForCausalLM.from_pretrained(args.model_id, trust_remote_code=True) processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True) # Freeze the vision tower if needed if args.freeze_vision_tower: for param in ft_model.vision_tower.parameters(): param.requires_grad = False # LoRA config. if args.use_lora: TARGET_MODULES = ["q_proj", "o_proj", "k_proj", "v_proj", "linear", "Conv2d", "lm_head", "fc2"] config = LoraConfig( r=8, lora_alpha=8, target_modules=TARGET_MODULES, task_type="CAUSAL_LM", lora_dropout=0.05, bias="none", inference_mode=False, use_rslora=True, init_lora_weights="gaussian", ) ft_model = get_peft_model(ft_model, config) # Saving and loading hooks. def save_model_hook(models, weights, output_dir): if accelerator.is_main_process: for model in models: if isinstance(accelerator.unwrap_model(model), type(accelerator.unwrap_model(ft_model))): model = accelerator.unwrap_model(model) model.save_pretrained(output_dir) else: raise ValueError(f"unexpected save model: {model.__class__}") # make sure to pop weight so that corresponding model is not saved again if weights: weights.pop() def load_model_hook(models, input_dir): transformer_ = None if not accelerator.distributed_type == DistributedType.DEEPSPEED: while len(models) > 0: model = models.pop() if isinstance(accelerator.unwrap_model(model), type(accelerator.unwrap_model(ft_model))): transformer_ = model # noqa: F841 else: raise ValueError(f"unexpected save model: {accelerator.unwrap_model(model).__class__}") else: transformer_ = AutoModelForCausalLM.from_pretrained(input_dir) # noqa: F841 accelerator.register_save_state_pre_hook(save_model_hook) accelerator.register_load_state_pre_hook(load_model_hook) if args.gradient_checkpointing: ft_model.gradient_checkpointing_enable() # Create DataLoaders train_loader, val_loader = create_data_loaders( train_dataset, val_dataset, args.batch_size, args.num_proc, processor, ) # Optimizer and scheduler optimizer_cls = torch.optim.AdamW if args.use_8bit_adam: import bitsandbytes as bnb optimizer_cls = bnb.optim.AdamW8bit trainable_params = list(filter(lambda p: p.requires_grad, ft_model.parameters())) optimizer = optimizer_cls(trainable_params, lr=args.lr) # Math around scheduler steps and training steps. len_train_dataloader_after_sharding = math.ceil(len(train_loader) / accelerator.num_processes) num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps) lr_scheduler = get_scheduler( name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=math.ceil(num_update_steps_per_epoch * args.epochs) * accelerator.num_processes, ) ft_model, train_loader, val_loader, optimizer, lr_scheduler = accelerator.prepare( ft_model, train_loader, val_loader, optimizer, lr_scheduler ) # Update again if needed. num_update_steps_per_epoch = math.ceil(len(train_loader) / args.gradient_accumulation_steps) max_train_steps = args.epochs * num_update_steps_per_epoch args.epochs = math.ceil(max_train_steps / num_update_steps_per_epoch) # 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: logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.") args.resume_from_checkpoint = None initial_global_step = 0 else: logger.info(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 # Start training! progress_bar = tqdm( range(0, 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, ) weight_dtype = torch.float16 if args.mixed_precision == "fp16" else torch.float32 global_step = 0 first_epoch = 0 for epoch in range(first_epoch, args.epochs): # Training phase ft_model.train() for batch in train_loader: with accelerator.accumulate(ft_model): # Prepare the input and target tensors color_inputs, colors = batch["color_inputs"], batch["colors"] lighting_inputs, lightings = batch["lighting_inputs"], batch["lightings"] lighting_type_inputs, lighting_types = batch["lighting_type_inputs"], batch["lighting_types"] composition_inputs, compositions = batch["composition_inputs"], batch["compositions"] losses = [] for inputs, labels in [ (color_inputs, colors), (lighting_inputs, lightings), (lighting_type_inputs, lighting_types), (composition_inputs, compositions), ]: losses.append(forward_with_model(ft_model, inputs, labels, weight_dtype=weight_dtype).loss) loss = torch.stack(losses).mean() accelerator.backward(loss) if accelerator.sync_gradients: params_to_clip = ft_model.parameters() accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) optimizer.step() lr_scheduler.step() optimizer.zero_grad() # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process: if global_step % args.save_steps == 0: 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}") val_loss = None if global_step % args.eval_steps == 0: val_loss = evaluate_model( ft_model, val_loader, accelerator.device, global_step, args.max_val_item_count, weight_dtype=weight_dtype, disable_pbar=not accelerator.is_local_main_process, ) logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} if val_loss: logs.update({"val_loss": val_loss}) progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) if global_step >= max_train_steps: break evaluate_model( ft_model, val_loader, accelerator.device, global_step, args.max_val_item_count, weight_dtype=weight_dtype, disable_pbar=not accelerator.is_local_main_process, ) # Finish run. accelerator.wait_for_everyone() if accelerator.is_main_process: accelerator.unwrap_model(ft_model).save_pretrained(args.output_dir) processor.save_pretrained(args.output_dir) accelerator.end_training() def main(args): model_slug = args.model_id.split("/")[-1] ds_slug = args.dataset_id.split("/")[-1] run_name = f"model@{model_slug}-ds@{ds_slug}-bs@{args.batch_size}-8bit@{args.use_8bit_adam}-lora@{args.use_lora}-lr@{args.lr}-mp@{args.mixed_precision}-fve@{args.freeze_vision_tower}" output_dir = Path("./model_checkpoints_accelerate") / run_name args.output_dir = output_dir accelerator_project_config = ProjectConfiguration(project_dir=output_dir, logging_dir=output_dir / "logs") kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, project_config=accelerator_project_config, kwargs_handlers=[kwargs], ) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() if accelerator.is_main_process: if args.output_dir is not None: os.makedirs(args.output_dir, exist_ok=True) tracker_name = "shot-categorizer" accelerator.init_trackers(tracker_name, config=vars(args), init_kwargs={"wandb": {"name": run_name}}) train_model(accelerator, args) if __name__ == "__main__": parser = argparse.ArgumentParser( description="Train Florence-2 model on specified dataset using YAML configuration" ) parser.add_argument("--config", type=str, default="config.yaml", help="Path to the YAML configuration file.") args = parser.parse_args() config = load_config(args.config) main(config)