hiyouga/LlamaFactory · error · ValueError

RM training requires pair data with token_type_ids. Ensure t

Error message

RM training requires pair data with token_type_ids. Ensure the dataset has chosen_messages/rejected_messages.

What it means

ValueError in RMTrainer.compute_loss (rm_trainer.py:107) when the batch carries no token_type_ids. RM packs chosen and rejected into one sequence and relies on token_type_ids (1=chosen, 2=rejected, 0=padding) to build the block-diagonal attention mask and locate each response's last token. Without it the loss is uncomputable, so the trainer refuses the batch.

Source

Thrown at src/llamafactory/v1/trainers/rm_trainer.py:107

                device_ids = None if self.device.type == "cpu" else [self.device.index]
                self.model = DDP(self.model, device_ids=device_ids, find_unused_parameters=True)
        else:
            super()._shard_model()

    @property
    def _unwrapped_model(self):
        """Access the underlying model, unwrapping DDP/FSDP wrappers if present."""
        model = self.model
        if hasattr(model, "module"):
            model = model.module
        return model

    def compute_loss(self, batch: BatchInput) -> Tensor:
        input_ids = batch["input_ids"].to(self.device, non_blocking=True)

        token_type_ids = batch.get("token_type_ids")
        if token_type_ids is None:
            raise ValueError(
                "RM training requires pair data with token_type_ids. "
                "Ensure the dataset has chosen_messages/rejected_messages."
            )
        token_type_ids = token_type_ids.to(self.device, non_blocking=True)

        # Use token_type_ids as document-index attention mask (values: 1=chosen, 2=rejected, 0=padding).
        # Transformers v5 models natively support this format in _update_causal_mask,
        # constructing the correct block-diagonal causal mask internally for all attention backends.
        model_attention_mask = token_type_ids

        # Build position_ids that reset at each document boundary.
        batch_size, seq_len = token_type_ids.shape
        arange = torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1)
        chosen_mask = token_type_ids == 1
        rejected_mask = token_type_ids == 2
        chosen_lens = chosen_mask.sum(dim=1, keepdim=True)
        position_ids = torch.zeros_like(token_type_ids)
        position_ids[chosen_mask] = arange[chosen_mask]

View on GitHub (pinned to f28afaf635)

Solutions

  1. Use the RM/pair data processor so batches contain token_type_ids alongside input_ids.
  2. If building batches manually, populate token_type_ids with 1 on chosen-response tokens, 2 on rejected-response tokens, 0 on padding before calling compute_loss.
  3. Ensure the dataset passed pair-format validation (chosen_messages/rejected_messages) so the processor marks token types.

Example fix

# before (custom batch)
batch = {"input_ids": ids, "attention_mask": mask}

# after
batch = {"input_ids": ids, "token_type_ids": token_types}  # 1=chosen, 2=rejected, 0=pad
Defensive patterns

Strategy: type-guard

Validate before calling

assert batch.get("token_type_ids") is not None, "RM batch missing token_type_ids (1=chosen, 2=rejected, 0=pad)"

Type guard

def is_rm_batch(batch) -> bool:
    tt = batch.get("token_type_ids")
    return tt is not None and bool(((tt == 1) | (tt == 2)).any())

Prevention

When it happens

Trigger: The data processor did not emit token_type_ids: non-pair dataset slipped past validation (e.g. a custom collator), a batch assembled without the pair processor, or a Manually constructed BatchInput missing the field. Note batch.get() returning None also covers an explicit None value.

Common situations: Custom collators or data pipelines that drop token_type_ids; using a processor stage that does not generate pair token types; mixing an SFT data pipeline with RMTrainer.

Related errors


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