hiyouga/LlamaFactory · critical · ValueError

Logits (batchsize x seqlen) and labels must have the same sh

Error message

Logits (batchsize x seqlen) and labels must have the same shape.

What it means

`get_batch_logps` (src/llamafactory/train/trainer_utils.py:606) computes per-sequence log-probabilities used by DPO/SimPO-style losses. Before shifting labels, it asserts logits.shape[:-1] == labels.shape, i.e. logits must be (batch, seq_len, vocab) matching labels (batch, seq_len). A mismatch means the model emitted logits for a different sequence length than the labels (e.g. only last k positions via num_logits_to_keep, or truncated/extended sequences).

Source

Thrown at src/llamafactory/train/trainer_utils.py:606

        for param in optimizer_dict.keys():
            param.register_post_accumulate_grad_hook(scheduler_hook)


def get_batch_logps(
    logits: "torch.Tensor",
    labels: "torch.Tensor",
    label_pad_token_id: int = IGNORE_INDEX,
    ld_alpha: Optional[float] = None,
) -> tuple["torch.Tensor", "torch.Tensor"]:
    r"""Compute the log probabilities of the given labels under the given logits.

    Returns:
        logps: A tensor of shape (batch_size,) containing the sum of log probabilities.
        valid_length: A tensor of shape (batch_size,) containing the number of non-masked tokens.

    """
    if logits.shape[:-1] != labels.shape:
        raise ValueError("Logits (batchsize x seqlen) and labels must have the same shape.")

    labels = labels[:, 1:].clone()
    logits = logits[:, :-1, :]
    loss_mask = labels != label_pad_token_id
    labels[labels == label_pad_token_id] = 0  # dummy token
    per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)

    valid_length = loss_mask.sum(-1)
    if ld_alpha is not None:
        num_examples = labels.shape[0] // 2
        chosen_lengths = valid_length[:num_examples]
        rejected_lengths = valid_length[num_examples:]
        min_lengths = torch.min(chosen_lengths, rejected_lengths)
        start_positions = torch.argmax(loss_mask.int(), dim=1)
        public_lengths = start_positions + torch.cat([min_lengths, min_lengths], dim=0)

        seq_len = labels.shape[-1]
        position_ids = torch.arange(seq_len, device=per_token_logps.device).expand_as(per_token_logps)

View on GitHub (pinned to f28afaf635)

Solutions

  1. Pass full-sequence logits: run forward without num_logits_to_keep (or set it to 0/None) before calling get_batch_logps.
  2. Pass raw (batch, seq_len) labels — the function performs the [1:]/[:-1] shift itself; do not pre-shift.
  3. Check that cutoff_len/packing settings are applied identically to inputs and labels in your custom loop.
  4. Pin/align the transformers version with the one your LlamaFactory release expects (logits_to_keep behavior changed across versions).

Example fix

# before
outputs = model(**inputs, num_logits_to_keep=16)
logps, seqlens = get_batch_logps(logits=outputs.logits, labels=labels)  # logits shorter than labels

# after
outputs = model(**inputs)
logps, seqlens = get_batch_logps(logits=outputs.logits, labels=labels)
Defensive patterns

Strategy: type-guard

Validate before calling

def check_logps_inputs(logits, labels):
    assert logits.dim() == 3 and logits.shape[:-1] == labels.shape, (
        f"logits {tuple(logits.shape)} vs labels {tuple(labels.shape)}"
    )

Type guard

def is_full_seq_logits_pair(logits: "torch.Tensor", labels: "torch.Tensor") -> bool:
    return logits.dim() == 3 and logits.shape[:-1] == labels.shape

Try / catch

try:
    logps, len = get_batch_logps(logits, labels)
except ValueError as e:
    if "same shape" in str(e):
        raise RuntimeError("model returned truncated logits; disable num_logits_to_keep") from e
    raise

Prevention

When it happens

Trigger: Calling get_batch_logps with logits from a forward that used num_logits_to_keep/logits_to_keep (shorter first dim), labels sliced to a different length, packing/neat-packing mismatches, or a custom model that returns padded/truncated logits.

Common situations: Preference-training (DPO/KTO) code paths after a transformers version change that enables logits_to_keep by default; hand-written loops that pass labels[:, :-1] or logits before the function's own shift; multimodal setups where image tokens change expected lengths.

Related errors


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