hpsv3/model/qwen2vl_trainer.py

import os
import pdb
import warnings
import time
import math
import json
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms as transforms

from typing import List, Optional, Dict, Union, Any
import pandas as pd
import safetensors
import numpy as np
import torch
import torch.nn as nn
import datasets
from torch.utils.data import Dataset, DataLoader
from peft import PeftModel
from transformers import Qwen2VLForConditionalGeneration
from transformers import AutoConfig
from transformers.modeling_utils import PreTrainedModel
from transformers.trainer import TrainerCallback
from transformers.trainer import (
    is_sagemaker_mp_enabled,
    is_peft_available,
    is_datasets_available,
    WEIGHTS_NAME,
    TRAINING_ARGS_NAME,
    SAFE_WEIGHTS_NAME,
    TRAINER_STATE_NAME,
    PREFIX_CHECKPOINT_DIR,
    logger,
    speed_metrics,
    deepspeed_init,
    speed_metrics,
    has_length,
    EvalPrediction,
    EvalLoopContainer,
    PredictionOutput,
    is_torch_xla_available,
    denumpify_detensorize,
    PredictionOutput,
    EvalLoopOutput,
    DistributedTensorGatherer,
    SequentialDistributedSampler,
    nested_concat,
)
from transformers.trainer_pt_utils import IterableDatasetShard
from transformers.trainer_callback import TrainerControl, TrainerState

from transformers.trainer_pt_utils import nested_detach, find_batch_size
from transformers.training_args import TrainingArguments
from trl import RewardTrainer
from hpsv3.utils.training_utils import get_peft_state_non_lora_maybe_zero_3

class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration):
    def __init__(
        self,
        config,
        output_dim=4,
        reward_token="last",
        special_token_ids=None,
        rm_head_type="default",
        rm_head_kwargs=None,
    ):
        super().__init__(config)
        # pdb.set_trace()
        self.output_dim = output_dim
        if rm_head_type == "default":
            self.rm_head = nn.Linear(config.hidden_size, output_dim, bias=False)
        elif rm_head_type == "ranknet":
            if rm_head_kwargs is not None:
                for layer in range(rm_head_kwargs.get("num_layers", 3)):
                    if layer == 0:
                        self.rm_head = nn.Sequential(
                            nn.Linear(config.hidden_size, rm_head_kwargs["hidden_size"]),
                            nn.ReLU(),
                            nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
                        )
                    elif layer < rm_head_kwargs.get("num_layers", 3) - 1:
                        self.rm_head.add_module(
                            f"layer_{layer}",
                            nn.Sequential(
                                nn.Linear(rm_head_kwargs["hidden_size"], rm_head_kwargs["hidden_size"]),
                                nn.ReLU(),
                                nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
                            ),
                        )
                    else:
                        self.rm_head.add_module(
                            f"output_layer",
                            nn.Linear(rm_head_kwargs["hidden_size"], output_dim, bias=rm_head_kwargs.get("bias", False)),
                        )

            else:
                self.rm_head = nn.Sequential(
                    nn.Linear(config.hidden_size, 1024),
                    nn.ReLU(),
                    nn.Dropout(0.05),
                    nn.Linear(1024, 16),
                    nn.ReLU(),
                    nn.Linear(16, output_dim),
                )

        self.rm_head.to(torch.float32)
        self.reward_token = reward_token

        self.special_token_ids = special_token_ids
        if self.special_token_ids is not None:
            self.reward_token = "special"

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
    ):
        ## modified from the origin class Qwen2VLForConditionalGeneration
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        # pdb.set_trace()
        if inputs_embeds is None:
            inputs_embeds = self.model.embed_tokens(input_ids)
            if pixel_values is not None:
                pixel_values = pixel_values.type(self.visual.get_dtype())
                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
                image_mask = (
                    (input_ids == self.config.image_token_id)
                    .unsqueeze(-1)
                    .expand_as(inputs_embeds)
                )
                image_embeds = image_embeds.to(
                    inputs_embeds.device, inputs_embeds.dtype
                )
                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

            if pixel_values_videos is not None:
                pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
                video_mask = (
                    (input_ids == self.config.video_token_id)
                    .unsqueeze(-1)
                    .expand_as(inputs_embeds)
                )
                video_embeds = video_embeds.to(
                    inputs_embeds.device, inputs_embeds.dtype
                )
                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

            if attention_mask is not None:
                attention_mask = attention_mask.to(inputs_embeds.device)

        outputs = self.model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]  # [B, L, D]
        with torch.autocast(device_type='cuda', dtype=torch.float32):
            logits = self.rm_head(hidden_states)  # [B, L, N]

        if input_ids is not None:
            batch_size = input_ids.shape[0]
        else:
            batch_size = inputs_embeds.shape[0]

        ## get sequence length
        if self.config.pad_token_id is None and batch_size != 1:
            raise ValueError(
                "Cannot handle batch sizes > 1 if no padding token is defined."
            )
        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
                sequence_lengths = (
                    torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
                )
                sequence_lengths = sequence_lengths % input_ids.shape[-1]
                sequence_lengths = sequence_lengths.to(logits.device)
            else:
                sequence_lengths = -1

        ## get the last token's logits
        if self.reward_token == "last":
            pooled_logits = logits[
                torch.arange(batch_size, device=logits.device), sequence_lengths
            ]
        elif self.reward_token == "mean":
            ## get the mean of all valid tokens' logits
            valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
            pooled_logits = torch.stack(
                [logits[i, : valid_lengths[i]].mean(dim=0) for i in range(batch_size)]
            )
        elif self.reward_token == "special":
            # special_token_ids = self.tokenizer.convert_tokens_to_ids(self.special_tokens)
            # create a mask for special tokens
            special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
            for special_token_id in self.special_token_ids:
                special_token_mask = special_token_mask | (
                    input_ids == special_token_id
                )
            pooled_logits = logits[special_token_mask, ...]
            pooled_logits = pooled_logits.view(
                batch_size, 1, -1
            )  # [B, 3, N] assert 3 attributes
            pooled_logits = pooled_logits.view(batch_size, -1)

            # pdb.set_trace()
        else:
            raise ValueError("Invalid reward_token")

        return {"logits": pooled_logits}


def _convert_A_B_to_chosen_rejected(
    rewards_A,
    rewards_B,
    tied_threshold=None,
    choice_dist=None,
):
    """
    Inputs:
        rewards_A: [B, 1]
        rewards_B: [B, 1]
    Outputs:
        rewards_chosen: [B, 1]
        rewards_rejected: [B, 1]
        nontied_mask: [B, 1] (preference labels that is not tied)
    """
    chosen_label = torch.ones_like(rewards_A, dtype=torch.int64).to(
        rewards_A.device
    )  # [B, 1]
    chosen_mask = chosen_label == 1
    rejected_mask = chosen_label != 1

    rewards_chosen = rewards_A 
    rewards_rejected = rewards_B

    if tied_threshold is None:
        nontied_mask = torch.ones_like(chosen_label, dtype=torch.float32).to(
            rewards_A.device
        )
    else:
        nontied_mask = (
            torch.abs(
                (choice_dist[:, 0] - choice_dist[:, 1]) / torch.sum(choice_dist, dim=-1)
            )
            > tied_threshold
        )
        print(nontied_mask)
    return (
        rewards_chosen,
        rewards_rejected,
        nontied_mask,
    )


class PartialEmbeddingUpdateCallback(TrainerCallback):
    """
    Callback to update the embedding of special tokens
    Only the special tokens are updated, the rest of the embeddings are kept fixed
    """

    def __init__(self, special_token_ids):
        super().__init__()
        self.special_token_ids = special_token_ids
        self.orig_embeds_params = None

    def on_train_begin(self, args, state, control, **kwargs):
        model = kwargs.get("model")
        self.orig_embeds_params = model.get_input_embeddings().weight.clone().detach()

    def on_step_end(self, args, state, control, **kwargs):
        # pdb.set_trace()
        model = kwargs.get("model")
        tokenizer = kwargs.get("tokenizer")

        index_no_updates = torch.ones((len(tokenizer),), dtype=torch.bool)
        index_no_updates[self.special_token_ids] = False
        with torch.no_grad():
            model.get_input_embeddings().weight[index_no_updates] = (
                self.orig_embeds_params[index_no_updates]
            )


class VLMRewardTrainer(RewardTrainer):
    def __init__(self, loss_type="regular", loss_hyperparameters={}, tied_threshold=None, 
                 visualization_steps=500, max_viz_samples=4, *args, **kwargs):
        super(VLMRewardTrainer, self).__init__(*args, **kwargs)
        self.loss_type = loss_type
        self.tied_threshold = tied_threshold
        self.rewards_chosen_accumulated = []
        self.rewards_rejected_accumulated = []
        self.loss_hyperparameters = loss_hyperparameters
        self.visualization_steps = visualization_steps
        self.max_viz_samples = max_viz_samples

    def get_eval_dataloader(
        self, eval_dataset: Optional[Union[str, Dataset]] = None
    ) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        # If we have persistent workers, don't do a fork bomb especially as eval datasets
        # don't change during training
        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
        if (
            hasattr(self, "_eval_dataloaders")
            and dataloader_key in self._eval_dataloaders
            and self.args.dataloader_persistent_workers
        ):
            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])

        eval_dataset = (
            self.eval_dataset[eval_dataset]
            if isinstance(eval_dataset, str)
            else eval_dataset if eval_dataset is not None else self.eval_dataset
        )

        data_collator = self.data_collator

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(
                eval_dataset, description="evaluation"
            )
        else:
            data_collator = self._get_collator_with_removed_columns(
                data_collator, description="evaluation"
            )

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        # accelerator.free_memory() will destroy the references, so
        # we need to store the non-prepared version
        eval_dataloader = DataLoader(eval_dataset, **dataloader_params)
        if self.args.dataloader_persistent_workers:
            if hasattr(self, "_eval_dataloaders"):
                self._eval_dataloaders[dataloader_key] = eval_dataloader
            else:
                self._eval_dataloaders = {dataloader_key: eval_dataloader}

        return self.accelerator.prepare(eval_dataloader)

    def create_optimizer(self):
        """
        Setup the optimizer.
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
        """
        if is_sagemaker_mp_enabled():
            return super().create_optimizer()

        opt_model = self.model

        if self.optimizer is None:
            decay_parameters = self.get_decay_parameter_names(opt_model)
            decay_parameters = [name for name in decay_parameters if "bias" not in name]
            lr_mapper = {}
            visual_parameters = []
            merger_parameters = []
            rm_head_parameters = []

            if self.args.vision_lr is not None:
                lr_mapper["visual"] = self.args.vision_lr
                visual_parameters = [
                    name
                    for name, _ in opt_model.named_parameters()
                    if "visual" in name and "merger" not in name
                ]
            if self.args.merger_lr is not None:
                lr_mapper["merger"] = self.args.merger_lr
                merger_parameters = [
                    name for name, _ in opt_model.named_parameters() if "merger" in name
                ]
                
            if self.args.rm_head_lr is not None:
                lr_mapper["rm_head"] = self.args.rm_head_lr
                rm_head_parameters = [
                    name for name, _ in opt_model.named_parameters() if "rm_head" in name
                ]

            if len(lr_mapper) > 0:
                special_lr_parameters = merger_parameters + visual_parameters + rm_head_parameters

                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p
                            for n, p in opt_model.named_parameters()
                            if (
                                n in decay_parameters
                                and n not in special_lr_parameters
                                and p.requires_grad
                            )
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p
                            for n, p in opt_model.named_parameters()
                            if (
                                n not in decay_parameters
                                and n not in special_lr_parameters
                                and p.requires_grad
                            )
                        ],
                        "weight_decay": 0.0,
                    },
                ]

                if visual_parameters:
                    optimizer_grouped_parameters.extend(
                        [
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n in decay_parameters
                                        and n in visual_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": self.args.weight_decay,
                                "lr": self.args.vision_lr,
                            },
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n not in decay_parameters
                                        and n in visual_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": 0.0,
                                "lr": self.args.vision_lr,
                            },
                        ]
                    )

                if merger_parameters:
                    optimizer_grouped_parameters.extend(
                        [
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n in decay_parameters
                                        and n in merger_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": self.args.weight_decay,
                                "lr": self.args.merger_lr,
                            },
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n not in decay_parameters
                                        and n in merger_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": 0.0,
                                "lr": self.args.merger_lr,
                            },
                        ]
                    )
                
                if rm_head_parameters:
                    optimizer_grouped_parameters.extend(
                        [
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n in decay_parameters
                                        and n in rm_head_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": self.args.weight_decay,
                                "lr": self.args.rm_head_lr,
                            },
                            {
                                "params": [
                                    p
                                    for n, p in opt_model.named_parameters()
                                    if (
                                        n not in decay_parameters
                                        and n in rm_head_parameters
                                        and p.requires_grad
                                    )
                                ],
                                "weight_decay": 0.0,
                                "lr": self.args.rm_head_lr,
                            },
                        ]
                    )

            else:
                optimizer_grouped_parameters = [
                    {
                        "params": [
                            p
                            for n, p in opt_model.named_parameters()
                            if (n in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": self.args.weight_decay,
                    },
                    {
                        "params": [
                            p
                            for n, p in opt_model.named_parameters()
                            if (n not in decay_parameters and p.requires_grad)
                        ],
                        "weight_decay": 0.0,
                    },
                ]

            if self.model.special_token_ids:
                special_token_embeddings = opt_model.get_input_embeddings().weight

                special_token_embeddings.requires_grad = True

                optimizer_grouped_parameters.extend(
                    [
                        {
                            # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)],
                            "params": [special_token_embeddings],
                            "lr": self.args.special_token_lr,
                            "weight_decay": 0.0,
                        },
                    ]
                )

            optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(
                self.args, opt_model
            )

            self.optimizer = optimizer_cls(
                optimizer_grouped_parameters, **optimizer_kwargs
            )

        return self.optimizer

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        rewards_A = model(return_dict=True, **inputs["batch_1"])["logits"]
        rewards_B = model(return_dict=True, **inputs["batch_2"])["logits"]
        
        # Log to TensorBoard for visualization
        if (hasattr(self.state, 'global_step') and 
            self.state.global_step % self.visualization_steps == 0 and
            self.state.global_step > 0):
            # Pass the original inputs which should contain the text prompts
            self._log_training_visualization(inputs, rewards_A, rewards_B)
        
        # calculate loss, optionally modulate with margin
        # get chosen and rejected rewards from the chosen label
        (
            rewards_chosen,
            rewards_rejected,
            nontied_mask,
        ) = _convert_A_B_to_chosen_rejected(
            rewards_A,
            rewards_B,
            tied_threshold=self.tied_threshold,
            choice_dist=inputs["choice_dist"],
        )

        loss_dict = {}

        if self.loss_type == "bt":
            # Bradley-Terry model
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected)
            out_mask = nontied_mask
            loss = loss * out_mask
            loss = loss.mean()
        elif self.loss_type == "likelihood_displacement":
            # Bradley-Terry model
            loss = -nn.functional.logsigmoid(rewards_chosen - self.loss_hyperparameters['tau'] * rewards_rejected)
            out_mask = nontied_mask
            loss = loss * out_mask
            loss = loss.mean()

        elif self.loss_type == "constant_margin":
            # Bradley-Terry model with constant margin
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - 0.57)
            out_mask = nontied_mask
            loss = loss * out_mask
            loss = loss.mean()
        elif self.loss_type == "btt":
            # Bradley-Terry-With-Ties model
            k = 5.0
            log_k = math.log(k)
            log_k2_sub_1 = math.log(k**2 - 1)
            bt_loss = -nn.functional.logsigmoid(
                rewards_chosen - rewards_rejected - log_k
            )
            same_loss = (
                -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k)
                - nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k)
                - log_k2_sub_1
            )
            loss = bt_loss * nontied_mask.float() + same_loss * (
                1 - nontied_mask.float()
            )
            out_mask = torch.ones_like(nontied_mask, dtype=torch.float32).to(
                rewards_A.device
            )  # [B, 1]
            loss = loss * out_mask

            loss = loss.mean()
        elif self.loss_type == "hpsv2":
            device = rewards_A.device
            rewards = torch.nn.functional.softmax(
                torch.cat([rewards_A, rewards_B], dim=-1), dim=-1
            )
            text_0_logits, text_1_logits = rewards[:, 0], rewards[:, 1]
            label_0, label_1 = torch.ones_like(text_0_logits), torch.zeros_like(
                text_0_logits
            )

            text_logits = torch.stack([text_0_logits, text_1_logits], dim=-1)
            text_0_labels = torch.zeros(
                text_logits.shape[0], device=device, dtype=torch.long
            )
            text_1_labels = text_0_labels + 1

            text_0_loss = torch.nn.functional.cross_entropy(
                text_logits, text_0_labels, reduction="none"
            )
            text_1_loss = torch.nn.functional.cross_entropy(
                text_logits, text_1_labels, reduction="none"
            )

            loss = label_0 * text_0_loss + label_1 * text_1_loss

            # absolute_example_weight = 1 / num_per_prompt
            # denominator = absolute_example_weight.sum()
            # weight_per_example = absolute_example_weight / denominator
            # text_loss *= weight_per_example
            loss = loss.sum()
        elif self.loss_type == "uncertainty":
            batch_size = rewards_A.shape[0]
            mean_chosen = rewards_A[:, 0]
            mean_rejected = rewards_B[:, 0]
            sigma_chosen = torch.exp(rewards_A[:, 1])
            sigma_rejected = torch.exp(rewards_B[:, 1])

            mean_z = mean_chosen - mean_rejected
            sigma_z = torch.sqrt(sigma_chosen**2 + sigma_rejected**2)

            z_samples = torch.randn(batch_size, 1000).to(sigma_z.device).to(
                torch.float16
            ) * sigma_z.unsqueeze(1).repeat(1, 1000) + mean_z.unsqueeze(1).repeat(
                1, 1000
            )
            loss = -torch.nn.functional.logsigmoid(z_samples).mean()
        else:
            raise NotImplementedError(f"Loss type {self.loss_type} not implemented.")

        loss_dict.update({"loss": loss.item()})

        if return_outputs:
            ## return rewards_A/B instead of chosen/rejected
            ## easier to calculate metrics for multi-attribute
            return loss, {
                "rewards_A": rewards_A,
                "rewards_B": rewards_B,
            }
        return loss

    def prediction_step(
        self,
        model,
        inputs,
        prediction_loss_only,
        ignore_keys=None,
    ):
        model.eval()
        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(
                    self.model.config, "keys_to_ignore_at_inference", []
                )
            else:
                ignore_keys = []

        with torch.no_grad():
            loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True)

        if prediction_loss_only:
            return (loss, None, None)
        loss = loss.detach()
        logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys)
        logits = nested_detach(logits)
        if self.loss_type != "uncertainty":
            logits = torch.cat(logits, dim=1)  # [B, 2]
        else:
            logits = torch.cat([p[:, [0]] for p in logits], dim=1)

        labels = torch.ones((logits.shape[0], 1)).to(logits.device)

        return loss, logits, labels

    def _log_training_visualization(self, inputs, rewards_A, rewards_B):
        """Log training samples and predictions to TensorBoard"""
        try:
            # Get tensorboard writer from trainer
            writer = None
            if hasattr(self, 'log_metrics'):
                # Try to get the writer from the logger
                if hasattr(self.args, 'report_to') and 'tensorboard' in self.args.report_to:
                    from torch.utils.tensorboard import SummaryWriter
                    if not hasattr(self, '_tb_writer'):
                        self._tb_writer = SummaryWriter(log_dir=self.args.logging_dir)
                    writer = self._tb_writer
            
            if writer is None:
                return
            
            step = self.state.global_step
            batch_size = min(len(rewards_A), self.max_viz_samples)
            
            # Log scalar metrics
            for i in range(batch_size):
                score_A = rewards_A[i].float().detach().cpu().numpy()
                score_B = rewards_B[i].float().detach().cpu().numpy()
                
                # Convert to float for logging
                score_A_val = float(score_A.mean()) if score_A.ndim > 0 else float(score_A)
                score_B_val = float(score_B.mean()) if score_B.ndim > 0 else float(score_B)
                score_diff = score_A_val - score_B_val
                
                writer.add_scalar(f'train_viz/sample_{i}/score_A', score_A_val, step)
                writer.add_scalar(f'train_viz/sample_{i}/score_B', score_B_val, step)
                writer.add_scalar(f'train_viz/sample_{i}/score_diff', score_diff, step)
                
                try:
                    # Get image data from inputs
                    image_A = inputs['image_1'][i] if 'image_1' in inputs else None
                    image_B = inputs['image_2'][i] if 'image_2' in inputs else None
                    
                    # Get prompt text from the original batch (now properly stored)
                    prompt_A = inputs.get('text_1', ['Unknown prompt'])[i] if 'text_1' in inputs else 'Unknown prompt'
                    
                    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 8))
                    fig.text(0.05, 0.05, f'Prompt:\n{prompt_A[:200]}{"..." if len(prompt_A) > 200 else ""}', 
                                 ha='left', va='bottom', fontsize=8, wrap=True,
                                 bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.7))
                    img_A_np = np.array(image_A)
                    if img_A_np.ndim == 3 and img_A_np.shape[0] == 3:  # CHW format
                        img_A_np = np.transpose(img_A_np, (1, 2, 0))
                    img_A_np = np.clip(img_A_np, 0, 1)  # Ensure values are in [0,1]
                    axes[0].imshow(img_A_np)
                    axes[0].set_title(f'Image A - Score: {score_A_val:.3f}')
                    axes[0].axis('off')
                    
                    img_B_np = np.array(image_B)
                    if img_B_np.ndim == 3 and img_B_np.shape[0] == 3:  # CHW format
                        img_B_np = np.transpose(img_B_np, (1, 2, 0))
                    img_B_np = np.clip(img_B_np, 0, 1)  # Ensure values are in [0,1]
                    axes[1].imshow(img_B_np)
                
                    axes[1].set_title(f'Image B - Score: {score_B_val:.3f}')
                    axes[1].axis('off')
                    
                    # Add prediction info
                    winner = "A" if score_diff > 0 else "B"
                    plt.suptitle(f'Step {step} - Sample {i} | Predicted Winner: Image {winner} | Diff: {score_diff:.3f}', fontsize=14)
                    plt.tight_layout()

                    # Log figure to tensorboard
                    writer.add_figure(f'train_viz/sample_{i}_comparison', fig, step)
                    plt.close(fig)
                except Exception as viz_error:
                    print(f"Warning: Could not extract images for visualization: {viz_error}")
                    continue

            # Log aggregate statistics
            all_scores_A = rewards_A.float().detach().cpu().numpy()
            all_scores_B = rewards_B.float().detach().cpu().numpy()
            
            writer.add_histogram('train_viz/all_scores_A', all_scores_A, step)
            writer.add_histogram('train_viz/all_scores_B', all_scores_B, step)
            writer.add_scalar('train_viz/mean_score_A', float(all_scores_A.mean()), step)
            writer.add_scalar('train_viz/mean_score_B', float(all_scores_B.mean()), step)
            writer.add_scalar('train_viz/mean_score_diff', float((all_scores_A - all_scores_B).mean()), step)
            
        except Exception as e:
            print(f"Error in training visualization: {e}")

    def _save_checkpoint(self, model, trial, metrics=None):

        if isinstance(self.model, PeftModel):
            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

            if self.hp_search_backend is None and trial is None:
                self.store_flos()

            run_dir = self._get_output_dir(trial=trial)
            output_dir = os.path.join(run_dir, checkpoint_folder)
            os.makedirs(output_dir, exist_ok=True)

            # TODO: Just Temp
            self.save_model(output_dir, _internal_call=True)
            # pdb.set_trace()

            if not self.args.save_full_model:
                non_lora_weights = get_peft_state_non_lora_maybe_zero_3(
                    self.model.named_parameters(), require_grad_only=True
                )
                torch.save(
                    non_lora_weights,
                    os.path.join(output_dir, "non_lora_state_dict.pth"),
                )
            # safetensors.torch.save(non_lora_weights, os.path.join(output_dir, "non_lora_model.safetensors"))

            if not self.args.save_only_model:
                # Save optimizer and scheduler
                self._save_optimizer_and_scheduler(output_dir)
                # Save RNG state
                self._save_rng_state(output_dir)

        else:
            super(RewardTrainer, self)._save_checkpoint(model, trial, metrics)

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        # If we are executing this function, we are the process zero, so we don't check for that.
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")
        # pdb.set_trace()

        supported_classes = (
            (PreTrainedModel,)
            if not is_peft_available()
            else (PreTrainedModel, PeftModel)
        )
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, supported_classes):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(self.accelerator.unwrap_model(self.model), supported_classes):
                self.accelerator.unwrap_model(self.model).save_pretrained(
                    output_dir,
                    state_dict=state_dict,
                    safe_serialization=self.args.save_safetensors,
                )
            else:
                logger.info(
                    "Trainer.model is not a `PreTrainedModel`, only saving its state dict."
                )
                if self.args.save_safetensors:
                    safetensors.torch.save_file(
                        state_dict,
                        os.path.join(output_dir, SAFE_WEIGHTS_NAME),
                        metadata={"format": "pt"},
                    )
                else:
                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            if not self.args.save_full_model:
                state_dict = {k: v for k, v in state_dict.items() if "wte" not in k}
                self.model.save_pretrained(
                    output_dir,
                    state_dict=state_dict,
                    safe_serialization=self.args.save_safetensors,
                )
            else:
                torch.save(state_dict, os.path.join(output_dir, "model.pth"))

        if self.tokenizer is not None:
            os.makedirs(os.path.join(output_dir, "tokenizer"), exist_ok=True)
            self.tokenizer.save_pretrained(os.path.join(output_dir, "tokenizer"))

        # Good practice: save your training arguments together with the trained model
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        # pdb.set_trace()


def compute_multi_attr_accuracy(eval_pred, metainfo_idxs=None) -> Dict[str, float]:
    predictions, labels = eval_pred
    metrics = {}

    pred_curr = predictions
    label_curr = labels.squeeze(1)
    total_count = np.sum(label_curr != 0)

    rewards_chosen = pred_curr[:, 0]
    rewards_rejected = pred_curr[:, 1]

    rewards_chosen_avg = np.sum(rewards_chosen) / total_count
    rewards_rejected_avg = np.sum(rewards_rejected) / total_count

    accuracy = np.sum(rewards_chosen > rewards_rejected) / total_count

    metrics.update(
        {
            f"Acc": accuracy,
            f"R_chosen_avg": rewards_chosen_avg,
            f"R_rejected_avg": rewards_rejected_avg,
        }
    )
    return metrics