def concatenated_forward(
        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
        """
        Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
        We do this to avoid doing two forward passes, because it's faster for FSDP.
        """
        concatenated_batch = self.concatenated_inputs(
            batch,
            is_encoder_decoder=self.is_encoder_decoder,
            label_pad_token_id=self.label_pad_token_id,
            padding_value=self.padding_value,
            device=self.accelerator.device,
        )

        model_kwargs = (
            {
                "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]),
            }
            if self.is_encoder_decoder
            else {}
        )

        if self.aux_loss_enabled:
            model_kwargs["output_router_logits"] = True

        if self.is_encoder_decoder:
            labels = concatenated_batch["concatenated_labels"].clone()
        else:
            labels = concatenated_batch["concatenated_input_ids"].clone()
            attention_mask = concatenated_batch["concatenated_attention_mask"]
            labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id)

        if isinstance(model, FullyShardedDataParallel):
            outputs = _FSDPForwardRedirection()(
                model,
                model._fsdp_wrapped_module.model,
                concatenated_batch["concatenated_input_ids"],
                attention_mask=concatenated_batch["concatenated_attention_mask"],
                use_cache=False,
                **model_kwargs,
            )
        else:
            if isinstance(model, torch.nn.DataParallel):
                model = model.module
            outputs = model.model(
                concatenated_batch["concatenated_input_ids"],
                attention_mask=concatenated_batch["concatenated_attention_mask"],
                use_cache=False,
                **model_kwargs,
            )

        orpo_loss_fn = LigerFusedLinearORPOLoss(ignore_index=self.label_pad_token_id, beta=self.beta)

        def orpo_partial(lm_head, last_hidden_state, concatenated_labels, nll_target):
            return orpo_loss_fn(
                lm_head.weight, last_hidden_state, concatenated_labels, lm_head.bias, nll_target=nll_target
            )

        orpo_loss, aux_outputs = _FSDPForwardRedirection()(
            model,
            orpo_partial,
            model.lm_head,
            outputs.last_hidden_state[:, :-1] if not self.is_encoder_decoder else outputs.last_hidden_state,
            concatenated_batch["concatenated_labels"][:, 1:]
            if not self.is_encoder_decoder
            else concatenated_batch["concatenated_labels"],
            labels[:, 1:] if not self.is_encoder_decoder else labels,
        )
        # if aux_loss_enabled, add the aux_loss to the orpo_loss
        if self.aux_loss_enabled:
            orpo_loss += self.aux_loss_coef * outputs.aux_loss

        return orpo_loss, aux_outputs

        Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
        We do this to avoid doing two forward passes, because it's faster for FSDP.
        
src/liger_kernel/transformers/trainer/orpo_trainer.py