src/trainer.py

from transformers.trainer import Trainer, TRAINING_ARGS_NAME
import torch.distributed as dist
from typing import Optional
import os
import torch
from src.loss import SimpleContrastiveLoss, DistributedContrastiveLoss, HardNegativeContrastiveLoss, DistributedHardNegativeContrastiveLoss
from itertools import repeat
from grad_cache.grad_cache import GradCache


MAX_INPUT_ID = int(1e9)
LLAVA_IMAGE_TOKEN_ID = 32000

class MMEBTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super(MMEBTrainer, self).__init__(*args, **kwargs)
        self.is_ddp = dist.is_initialized()
        self._dist_loss_scale_factor = dist.get_world_size() if self.is_ddp else 1

    def compute_loss(self, model, inputs, *args, **kwargs):
        if self.args.hard_neg:
            qry_inputs, tgt_inputs, neg_inputs = inputs
            return model(qry=qry_inputs, tgt=tgt_inputs, neg=neg_inputs)

        qry_inputs, tgt_inputs = inputs
        return model(qry=qry_inputs, tgt=tgt_inputs)

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        os.makedirs(output_dir, exist_ok=True)

        if state_dict is None:
            state_dict = self.model.state_dict()
        prefix = 'encoder.'
        assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
        state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
        self.model.encoder.save_pretrained(
            output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
        )

        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(output_dir)

        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))


def split_dense_inputs(model_input: dict, chunk_size: int):
    assert len(model_input) == 1
    arg_key = list(model_input.keys())[0]
    arg_val = model_input[arg_key]

    keys = list(arg_val.keys())
    chunked_tensors = [arg_val[k].split(chunk_size, dim=0) for k in keys]
    chunked_arg_val = [dict(zip(kk, tt)) for kk, tt in zip(repeat(keys), zip(*chunked_tensors))]

    return [{arg_key: c} for c in chunked_arg_val]


def split_vlm_inputs(model_input: dict, chunk_size: int):
    assert len(model_input) == 1
    arg_key = list(model_input.keys())[0]
    arg_val = model_input[arg_key]
    keys = list(arg_val.keys())

    # for input_ids and attention_mask, split directly
    chunked_tensors = [arg_val[k].split(chunk_size, dim=0) for k in ["input_ids", "attention_mask"]]

    # for pixel_values and image_sizes or any other image-related fields, need to split based on the position of images
    
    image_mask = "image_mask" if "image_mask" in keys else None

    if image_mask in keys:
        row_contain_image = torch.nonzero(arg_val[image_mask], as_tuple=False).squeeze()  # indicates which row in input_ids contain images
        if image_mask == "image_mask": 
            keys.remove(image_mask)
        num_chunks = len(chunked_tensors[0])
        chunk_image_count = []
        for chunk_idx in range(num_chunks):
            chunk_image_count.append(torch.sum(
                (row_contain_image >= chunk_idx * chunk_size) & (row_contain_image < (chunk_idx + 1) * chunk_size)).item())
        
        if "pixel_values" in keys:
            pixel_values = arg_val["pixel_values"]
            chunked_tensors.append(torch.split(pixel_values, chunk_image_count))
        if "image_sizes" in keys:
            image_sizes = arg_val["image_sizes"]
            chunked_tensors.append(torch.split(image_sizes, chunk_image_count))
        if "image_grid_thw" in keys:
            image_grid_thw = arg_val["image_grid_thw"]
            chunked_tensors.append(torch.split(image_grid_thw, chunk_image_count))
        
        if "image_flags" in keys:
            image_flags = arg_val["image_flags"]
            chunked_tensors.append(torch.split(image_flags, chunk_size))
            keys.remove("image_flags") 
        
    
    chunked_arg_val = []
    for kk, tt in zip(repeat(keys), zip(*chunked_tensors)):
        chunk_dict = {}
        # 先添加基本字段
        if "pixel_values" in keys and tt[2].numel() == 0:  # this chunk doesn't contain image
            chunk_dict.update(dict(zip(kk[:2], tt[:2])))
        else:
            chunk_dict.update(dict(zip(kk, tt)))
            
        # 如果有image_flags，添加对应的chunk
        if "image_flags" in arg_val:
            chunk_idx = len(chunked_arg_val)
            chunk_dict["image_flags"] = chunked_tensors[-1][chunk_idx]
            
        chunked_arg_val.append(chunk_dict)

    return [{arg_key: c} for c in chunked_arg_val]


def get_dense_rep(x):
    """
    Get either qry_reps or tgt_reps.
    """
    if x["qry_reps"] is None:
        return x["tgt_reps"]
    else:
        return x["qry_reps"]


class GradCacheTrainer(Trainer):
    """
    Adapted from gradcache repo.
    """
    def __init__(self, *args, **kwargs):
        super(GradCacheTrainer, self).__init__(*args, **kwargs)
        self.is_ddp = dist.is_initialized()
        self._dist_loss_scale_factor = dist.get_world_size() if self.is_ddp else 1
        # loss_fn_cls = DistributedContrastiveLoss if self.is_ddp else SimpleContrastiveLoss
        # 使用新的损失函数
        loss_fn_cls = DistributedHardNegativeContrastiveLoss if self.is_ddp else HardNegativeContrastiveLoss
        loss_fn = loss_fn_cls(temperature=self.model.temperature)

        self.gc = GradCache(
            models=[self.model, self.model],
            chunk_sizes=[self.args.gc_q_chunk_size, self.args.gc_p_chunk_size],
            loss_fn=loss_fn,
            split_input_fn=split_vlm_inputs,
            get_rep_fn=get_dense_rep,
            fp16=self.args.fp16,
            scaler=self.scaler if self.args.fp16 else None
        )

    def training_step(self, model, inputs, *args, **kwargs) -> torch.Tensor:
        model.train()
        
        if self.args.hard_neg:
            queries, passages, negatives = inputs
            queries, passages, negatives = {'qry': queries}, {'tgt': passages}, {'neg': negatives}
            
            if self.args.local_rank == 0:
                print(f"qry.shape={queries['qry']['input_ids'].shape}")
                print(f"tgt.shape={passages['tgt']['input_ids'].shape}")
                print(f"neg.shape={negatives['neg']['input_ids'].shape}")
                if 'pixel_values' in queries['qry']:
                    print(f"qry_img.shape={queries['qry']['pixel_values'].shape}")
                if 'pixel_values' in passages['tgt']:
                    print(f"tgt_img.shape={passages['tgt']['pixel_values'].shape}")
                if 'pixel_values' in negatives['neg']:
                    print(f"neg_img.shape={negatives['neg']['pixel_values'].shape}")
            
            _distributed = self.args.local_rank > -1
            self.gc.models = [model, model, model]  
            loss = self.gc(queries, passages, negatives, no_sync_except_last=_distributed)
        else:
            queries, passages = inputs
            queries, passages = {'qry': queries}, {'tgt': passages}
            
            if self.args.local_rank == 0:
                print(f"qry.shape={queries['qry']['input_ids'].shape}")
                print(f"tgt.shape={passages['tgt']['input_ids'].shape}")
                if 'pixel_values' in queries['qry']:
                    print(f"qry_img.shape={queries['qry']['pixel_values'].shape}")
                if 'pixel_values' in passages['tgt']:
                    print(f"tgt_img.shape={passages['tgt']['pixel_values'].shape}")
            
            _distributed = self.args.local_rank > -1
            self.gc.models = [model, model]
            loss = self.gc(queries, passages, no_sync_except_last=_distributed)

        return loss / self._dist_loss_scale_factor

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        print(f"Saving model to {output_dir}")
        os.makedirs(output_dir, exist_ok=True)

        if state_dict is None:
            state_dict = self.model.state_dict()
        prefix = 'encoder.'
        assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys())
        state_dict = {k[len(prefix):]: v for k, v in state_dict.items()}
        self.model.encoder.save_pretrained(
            output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
        )

        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(output_dir)

        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        self.model.encoder.config.to_json_file(os.path.join(output_dir, 'config.json'))