code/training/trainer.py

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