"""Train PRIMA on M4SEG with native single-GPU or DeepSpeed distributed execution.""" import argparse from functools import partial import json import math import os from pathlib import Path import random import numpy as np import torch from torch.utils.data import DataLoader, DistributedSampler from tqdm import tqdm from dataset.dataset import custom_collate_fn_prima from eval.tasks import M4SEG_SOURCES, build_dataset from model.llava import conversation as conversation_lib from tools.batch import move_batch from tools.checkpoint import export_checkpoint, load_checkpoint from tools.config import parse_with_config, validate_paths from tools.model_loader import (add_model_arguments, normalize_model_args, setup_tokenizer_and_special_tokens, setup_qformer_tokenizer, initialize_model, prepare_model_for_inference, setup_lora_config) def parse_args(argv=None): parser = argparse.ArgumentParser(description=__doc__) add_model_arguments(parser) parser.set_defaults(ckpt_path='') parser.add_argument('--dataset_dir', default='/path/to/datasets') parser.add_argument('--m4seg_annpath', default='/path/to/training_annotations.json') parser.add_argument('--m4seg_data', default=M4SEG_SOURCES) parser.add_argument('--batch_size', type=int, default=4) parser.add_argument('--grad_accumulation_steps', type=int, default=10) parser.add_argument('--epochs', type=int, default=1) parser.add_argument('--workers', type=int, default=4) parser.add_argument('--lr', type=float, default=3e-4) parser.add_argument('--beta1', type=float, default=0.9) parser.add_argument('--beta2', type=float, default=0.95) parser.add_argument('--weight_decay', type=float, default=0.0) parser.add_argument('--warmup_steps', type=int, default=100) parser.add_argument('--max_steps', type=int, default=0, help='Optional optimizer-step limit.') parser.add_argument('--save_every', type=int, default=100) parser.add_argument('--gradient_checkpointing', action=argparse.BooleanOptionalAction, default=True) parser.add_argument('--unfreeze_mm_projector', action=argparse.BooleanOptionalAction, default=True) parser.add_argument('--ce_loss_weight', type=float, default=1.0) parser.add_argument('--dice_loss_weight', type=float, default=0.5) parser.add_argument('--bce_loss_weight', type=float, default=2.0) parser.add_argument('--seed', type=int, default=27) parser.add_argument('--log_base_dir', default='outputs/training') parser.add_argument('--exp_name', default='prima') parser.add_argument('--resume', default='', help='Recoverable checkpoint directory.') parser.add_argument('--auto_resume', action='store_true') parser.set_defaults(train_mask_decoder=True) args = normalize_model_args(parse_with_config(parser, argv)) if min(args.batch_size, args.grad_accumulation_steps, args.epochs) < 1: parser.error('Batch size, accumulation, and epochs must be positive.') if args.lora_r < 1: parser.error('PRIMA training requires a positive LoRA rank.') return args def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def prepare_training_model(args, tokenizer, qformer_tokenizer): model = initialize_model(args, tokenizer) model = prepare_model_for_inference(model, args, qformer_tokenizer) if args.merged_ckpt_path: from peft import PeftModel, get_peft_model if not isinstance(model, PeftModel): model = get_peft_model(model, setup_lora_config(model, args)) if args.ckpt_path and not args.ckpt_path.startswith('/path/to/'): load_checkpoint(model, args.ckpt_path) model.to(device=torch.device(f'cuda:{args.local_rank}'), dtype=torch.bfloat16) trainable = ('lora_', 'grounding_encoder.mask_decoder.', 'text_hidden_fcs.', 'query_tokens', 'residual_extractor.') for name, parameter in model.named_parameters(): parameter.requires_grad_(any(group in name for group in trainable) or (args.unfreeze_mm_projector and 'mm_projector.' in name)) model.enable_input_require_grads() if args.gradient_checkpointing: model.gradient_checkpointing_enable() model.config.use_cache = False return model def train_mode(model): model.train() unwrapped = model.module if hasattr(model, 'module') else model core = unwrapped.get_model() for frozen in (core.visual_encoder, core.ln_vision, core.Qformer, core.grounding_encoder.image_encoder, core.grounding_encoder.prompt_encoder): frozen.eval() def learning_rate_factor(step, *, warmup, total): if warmup and step < warmup: return float(step) / warmup progress = (step - warmup) / max(1, total - warmup) return 0.5 * (1 + math.cos(math.pi * min(1.0, max(0.0, progress)))) def save_checkpoint(model, optimizer, scheduler, output, state, *, distributed): checkpoint = output / 'checkpoint' if distributed: model.save_checkpoint(str(checkpoint), tag=f"step{state['global_step']}", client_state=state, exclude_frozen_parameters=True) elif int(os.environ.get('RANK', '0')) == 0: checkpoint.mkdir(parents=True, exist_ok=True) export_checkpoint(model, checkpoint / 'weights.bin') payload = {**state, 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(), 'python_rng': random.getstate(), 'numpy_rng': np.random.get_state()} temporary = checkpoint / 'training.tmp' torch.save(payload, temporary) temporary.replace(checkpoint / 'training.pt') if int(os.environ.get('RANK', '0')) == 0: core = model.module if distributed else model export_checkpoint(core, output / 'exports' / f"step{state['global_step']}.bin") def resume_checkpoint(model, optimizer, scheduler, path, *, distributed): if distributed: loaded, state = model.load_checkpoint(str(path), load_module_strict=False) if not loaded: raise ValueError('No distributed training checkpoint was found.') else: load_checkpoint(model, Path(path) / 'weights.bin') state = torch.load(Path(path) / 'training.pt', map_location='cpu') optimizer.load_state_dict(state['optimizer']) scheduler.load_state_dict(state['scheduler']) torch.set_rng_state(state['torch_rng']) torch.cuda.set_rng_state_all(state['cuda_rng']) random.setstate(state['python_rng']) np.random.set_state(state['numpy_rng']) return state def main(args): validate_paths(args, checkpoint=bool(args.ckpt_path), dataset=True) if not Path(args.m4seg_annpath).is_file(): raise ValueError('Set m4seg_annpath to an existing M4SEG training annotation path.') world_size = int(os.environ.get('WORLD_SIZE', '1')) rank = int(os.environ.get('RANK', '0')) args.local_rank = int(os.environ.get('LOCAL_RANK', args.local_rank)) device = torch.device(f'cuda:{args.local_rank}') torch.cuda.set_device(device) distributed = world_size > 1 if distributed: import deepspeed deepspeed.init_distributed(dist_backend='nccl') set_seed(args.seed) tokenizer = setup_tokenizer_and_special_tokens(args) tokenizer.padding_side = 'right' qformer_tokenizer = setup_qformer_tokenizer(args) conversation_lib.default_conversation = conversation_lib.conv_templates[args.conv_type] model = prepare_training_model(args, tokenizer, qformer_tokenizer) dataset = build_dataset(args, mode='train') if not len(dataset): raise ValueError('M4SEG training annotations contain no supported records.') sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=args.seed) if distributed else None # Use an epoch-specific generator so skipping resumed batches preserves shuffle order. generator = torch.Generator() collate = partial(custom_collate_fn_prima, tokenizer=tokenizer, qformer_tokenizer=qformer_tokenizer, inference=False, img_emb_len=args.img_emb_len, local_rank=args.local_rank) loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=sampler is None, sampler=sampler, num_workers=args.workers, collate_fn=collate, generator=generator, pin_memory=True) steps_per_epoch = math.ceil(len(loader) / args.grad_accumulation_steps) total_steps = steps_per_epoch * args.epochs optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr, betas=(args.beta1, args.beta2), weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, partial(learning_rate_factor, warmup=args.warmup_steps, total=total_steps)) if distributed: model, optimizer, _, scheduler = deepspeed.initialize( model=model, optimizer=optimizer, lr_scheduler=scheduler, config={'train_micro_batch_size_per_gpu': args.batch_size, 'gradient_accumulation_steps': args.grad_accumulation_steps, 'bf16': {'enabled': True}, 'gradient_clipping': 1.0, 'zero_optimization': {'stage': 2}}, ) output = Path(args.log_base_dir) / args.exp_name output.mkdir(parents=True, exist_ok=True) if rank == 0: (output / 'configuration.json').write_text(json.dumps(vars(args), default=str, indent=2) + '\n') start_epoch = start_batch = global_step = 0 resume = args.resume or (str(output / 'checkpoint') if args.auto_resume and (output / 'checkpoint').exists() else '') if resume: state = resume_checkpoint(model, optimizer, scheduler, resume, distributed=distributed) start_epoch, start_batch, global_step = state['epoch'], state['next_batch'], state['global_step'] if state['batch_size'] != args.batch_size or state['accumulation'] != args.grad_accumulation_steps: raise ValueError('Resume batch size and accumulation must match the saved run.') train_mode(model) optimizer.zero_grad() stop = False for epoch in range(start_epoch, args.epochs): generator.manual_seed(args.seed + epoch) if sampler is not None: sampler.set_epoch(epoch) for batch_index, batch in enumerate(tqdm(loader, desc=f'Epoch {epoch + 1}', disable=rank != 0)): if epoch == start_epoch and batch_index < start_batch: continue batch = move_batch(batch, device) result = model(**batch) loss = result['loss'] if not torch.isfinite(loss): raise RuntimeError(f'Non-finite M4SEG loss at epoch {epoch}, batch {batch_index}.') window_start = (batch_index // args.grad_accumulation_steps) * args.grad_accumulation_steps window_size = min(args.grad_accumulation_steps, len(loader) - window_start) boundary = (batch_index + 1) % args.grad_accumulation_steps == 0 or batch_index + 1 == len(loader) if distributed: model.set_gradient_accumulation_boundary(boundary) model.backward(loss * args.grad_accumulation_steps / window_size) model.step() else: (loss / window_size).backward() if boundary: torch.nn.utils.clip_grad_norm_((p for p in model.parameters() if p.requires_grad), 1.0, error_if_nonfinite=True) optimizer.step() scheduler.step() optimizer.zero_grad() if not boundary: continue global_step += 1 if rank == 0: values = {key: float(value.detach()) for key, value in result.items() if torch.is_tensor(value) and value.numel() == 1} print(json.dumps({'step': global_step, **values}), flush=True) state = {'epoch': epoch, 'next_batch': batch_index + 1, 'global_step': global_step, 'batch_size': args.batch_size, 'accumulation': args.grad_accumulation_steps} stop = bool(args.max_steps and global_step >= args.max_steps) if stop or (args.save_every > 0 and global_step % args.save_every == 0): save_checkpoint(model, optimizer, scheduler, output, state, distributed=distributed) if stop: break if stop: break state = {'epoch': epoch + 1, 'next_batch': 0, 'global_step': global_step, 'batch_size': args.batch_size, 'accumulation': args.grad_accumulation_steps} save_checkpoint(model, optimizer, scheduler, output, state, distributed=distributed) if distributed: torch.distributed.barrier() torch.distributed.destroy_process_group()