| """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 |
| |
| 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() |
|
|
