hiyouga/LlamaFactory · error · NotImplementedError

Stage does not supported: {stage}.

Error message

Stage does not supported: {stage}.

What it means

In KTO's concatenated_forward (trainer.py:192), per-sequence logps must align 1:1 with the batch's kto_tags boolean mask used to split chosen/rejected. If len(target_logps) != len(batch['kto_tags']) the split would silently misindex, so the trainer raises ValueError. A mismatch means the tokenized inputs and the kto_tag labels were produced from inconsistent batch shapes — almost always a data/collation problem, not model logic.

Source

Thrown at scripts/stat_utils/cal_ppl.py:104

        )
    )
    tokenizer_module = load_tokenizer(model_args)
    tokenizer = tokenizer_module["tokenizer"]
    template = get_template_and_fix_tokenizer(tokenizer, data_args)
    trainset = get_dataset(template, model_args, data_args, training_args, stage, **tokenizer_module)["train_dataset"]
    model = load_model(tokenizer, model_args, finetuning_args, is_trainable=False)
    if stage == "pt":
        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    elif stage == "sft":
        data_collator = MultiModalDataCollatorForSeq2Seq(
            template=template, tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX
        )
    elif stage == "rm":
        data_collator = PairwiseDataCollatorWithPadding(
            template=template, tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX, train_on_prompt=train_on_prompt
        )
    else:
        raise NotImplementedError(f"Stage does not supported: {stage}.")

    dataloader = DataLoader(trainset, batch_size, shuffle=False, collate_fn=data_collator, pin_memory=True)
    criterion = torch.nn.CrossEntropyLoss(reduction="none")
    total_ppl = 0
    perplexities = []
    batch: dict[str, torch.Tensor]
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Computing perplexities"):
            batch = batch.to(model.device)
            outputs = model(**batch)
            shift_logits: torch.Tensor = outputs["logits"][..., :-1, :]
            shift_labels: torch.Tensor = batch["labels"][..., 1:]
            loss_mask = shift_labels != IGNORE_INDEX
            flatten_logits = shift_logits.contiguous().view(shift_labels.size(0) * shift_labels.size(1), -1)
            flatten_labels = shift_labels.contiguous().view(-1)
            token_logps: torch.Tensor = criterion(flatten_logits, flatten_labels)
            token_logps = token_logps.contiguous().view(shift_logits.size(0), -1)
            sentence_logps = (token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)

View on GitHub (pinned to f28afaf635)

Solutions

1. Validate the dataset offline: every KTO example must have exactly one kto_tag (bool) matching one input sequence — check `len(tokenizer(ex['prompt']+ex['response']).input_ids)` rows against tags
2. Regenerate the dataset with the documented KTO format (prompt/response/kto_tag columns) via dataset_info.json rather than a custom script
3. If using a custom collator, ensure it never changes batch size between inputs and kto_tags
4. Re-run tokenization after changing cutoff_len or template so filtering is applied consistently

Example fix

# before: custom KTO rows with drifting tags
[
  {"prompt": "...", "response": "...", "kto_tag": true},
  {"prompt": "..."}  # missing tag -> misaligned batch
]

# after: every row tagged, standard dataset_info entry
[
  {"prompt": "...", "response": "...", "kto_tag": true},
  {"prompt": "...", "response": "...", "kto_tag": false}
]

# Offline KTO alignment check before training
bad = [
    i for i, ex in enumerate(dataset)
    if not isinstance(ex.get('kto_tag'), bool) or not ex.get('prompt') or not ex.get('response')
]
assert not bad, f'Rows with missing/invalid kto_tag: {bad[:5]}'
assert all(len(dataset[i]['prompt']) + len(dataset[i]['response']) > 0 for i in range(len(dataset)))

def is_valid_kto_example(example: dict) -> bool:
    return (
        isinstance(example.get('kto_tag'), bool)
        and isinstance(example.get('prompt'), str)
        and isinstance(example.get('response'), str)
    )

try:
    trainer.train()
except ValueError as e:
    if 'Mismatched shape' in str(e):
        raise SystemExit('KTO dataset misaligned: verify one kto_tag per sequence and collator batch handling') from e
    raise

When it happens

Trigger: KTO dataset rows missing or misaligned kto_tag fields, a custom data collator that pads/duplicates inputs but not kto_tags, or preprocessing that drops sequences (e.g. over-long samples filtered on one side only), so the collated batch has N logps but M tags.

Common situations: Hand-built KTO datasets where kto_tag length differs from input length; mixing a non-KTO collator into the KTO stage; cutoff_len filtering applied after tagging; corrupted/sharegpt-format conversions where some rows lack the tag.

Related errors

• Expected a string, got {value}
• The model does not have a submodule named '{submodule_name}'
• Unsupported model type: {getattr(config, 'model_type')}.
• Device not supported: {device_name}.
• Not allowed

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