hiyouga/LlamaFactory · error · ValueError

Mismatched shape of inputs and labels.

Error message

Mismatched shape of inputs and labels.

What it means

Error "Mismatched shape of inputs and labels." thrown in hiyouga/LlamaFactory.

Source

Thrown at src/llamafactory/train/kto/trainer.py:192

            model_inputs["aspect_ratio_mask"] = batch["aspect_ratio_mask"]

        if f"{prefix}cross_attention_mask" in batch:
            model_inputs["cross_attention_mask"] = batch[f"{prefix}cross_attention_mask"]

        logits = model(**model_inputs, return_dict=True, use_cache=False).logits.to(torch.float32)
        logps, valid_length = get_batch_logps(logits=logits, labels=batch[f"{prefix}labels"])
        return logits, logps, logps / valid_length

    @override
    def concatenated_forward(
        self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
    ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        target_logits, target_logps, target_logps_avg = self.forward(model, batch)
        with torch.no_grad():
            _, kl_logps, _ = self.forward(model, batch, prefix="kl_")

        if len(target_logps) != len(batch["kto_tags"]):
            raise ValueError("Mismatched shape of inputs and labels.")

        chosen_logits = target_logits[batch["kto_tags"]]
        chosen_logps = target_logps[batch["kto_tags"]]
        rejected_logits = target_logits[~batch["kto_tags"]]
        rejected_logps = target_logps[~batch["kto_tags"]]
        chosen_logps_avg = target_logps_avg[batch["kto_tags"]]
        return chosen_logps, rejected_logps, chosen_logits, rejected_logits, kl_logps, chosen_logps_avg

    @override
    def compute_reference_log_probs(
        self, model: "PreTrainedModel", batch: dict[str, "torch.Tensor"]
    ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        r"""Compute log probabilities of the reference model."""
        if self.ref_model is None:
            ref_model = model
            ref_context = self.accelerator.unwrap_model(model).disable_adapter()
        else:
            ref_model = self.ref_model

View on GitHub (pinned to f28afaf635)

Solutions

1. Ensure each KTO batch contains exactly one kto_tag label per input sample; check that the dataset uses the KTO format (single response with a boolean kto_tag) and is not truncated or padded inconsistently.

Example fix

# dataset sample
{"messages": [...], "kto_tag": true}

When it happens

Trigger: Thrown at src/llamafactory/train/kto/trainer.py:192 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of hiyouga/LlamaFactory@f28afaf635 (2026-08-14). Data as JSON: /api/errors/f0dc7c952d773b36. Report an issue: GitHub.