utils.collators.dpo

utils.collators.dpo

DPO/ORPO/IPO/KTO data collator with pad_to_multiple_of support.

Extends TRL’s DPODataCollatorWithPadding to round padded sequence lengths up to a fixed multiple. This stabilizes Triton autotune caches for kernels that key on sequence length (e.g. fla’s linear attention kernels used by Qwen3.5), which otherwise re-autotune on every distinct batch length.

Classes

Name Description
AxolotlDPODataCollatorWithPadding DPO data collator that pads to a multiple of pad_to_multiple_of.

AxolotlDPODataCollatorWithPadding

utils.collators.dpo.AxolotlDPODataCollatorWithPadding(pad_to_multiple_of=None)

DPO data collator that pads to a multiple of pad_to_multiple_of.

Parameters

Name Type Description Default
pad_token_id Tokenizer pad token id (inherited). required
is_encoder_decoder Whether the model is encoder-decoder (inherited). required
pad_to_multiple_of int | None If set, padded lengths are rounded up to this multiple. Helps stabilize Triton autotune caches. None