Problems with Cross-Tokenizer Alignment in Correctness and Efficiency

[JqzChandler]

trl/trl/experimental/gold/gold_trainer.py

Lines 342 to 366 in 1eb561c

	def _build_alignment_groups_from_ids(self, student_token_ids, teacher_token_ids):
	"""
	Build alignment groups using a greedy substring-equality algorithm on decoded token pieces.
	Args:
	student_token_ids: List[int]
	teacher_token_ids: List[int]
	Returns:
	Tuple[List[List[int]], List[List[int]]]: student and teacher alignment groups
	"""
	def to_canonical_pieces(tok, ids):
	pieces = []
	prev = ""
	for k in range(len(ids)):
	# IMPORTANT: Do NOT skip special tokens - we need to align them too
	cur = tok.decode(ids[: k + 1], skip_special_tokens=False, clean_up_tokenization_spaces=False)
	# Extract the incremental addition (may include spaces/ZWJ/etc.)
	pieces.append(cur[len(prev) :])
	prev = cur
	return pieces
	s_pieces = to_canonical_pieces(self.student_tokenizer, student_token_ids)
	t_pieces = to_canonical_pieces(self.teacher_tokenizer, teacher_token_ids)

1. Current implementation produces incorrect decoding results in some cases
   BPE-based tokenizers always split some multi-byte characters, e.g. some Chinese characters. In the current implementation, the algorithm adds one token per loop iteration, calls the decode method, and takes the trailing incremental text as the corresponding piece. This causes decoding errors for these split bytes, with no subsequent validation or error correction mechanism. To reproduce this issue quickly, try preparing multilingual mixed text samples and load Qwen3 tokenizer, then assert ''.join(pieces) == orig_text to verify the splitting results for each sample.

2. Calling decode() repeatedly, causes significant computational overhead
   In our team's practice, we modified Python bindings in tokenizers (and correspondingly adjusted tokenization_utils_fast.py in transformers). We introduced parameter "offset_type" in both encode() and batch_encode() methods to directly obtain the offset sequences generated during the Rust code's single encoding execution. This approach has minimal computational overhead and ensures correctness.

This practice also provides an additional advantage: when aligning two families of BPE tokenizers, we can achieve more accurate token alignment results without relying on string as an intermediate modality.

def align_using_offsets(offsets_T: List, offsets_S: List):
    """
    offsets_T/_S be like:
    [(0, 1), (1, 3), (3, 8), ..., (start_i, end_i), ..., (start_I, TOTAL_BYTES_OF_ORIG_STRING)] # len = I == len(input_ids_T)
    [(0, 3), (3, 5), (5, 8), ..., (start_j, end_j), ..., (start_J, TOTAL_BYTES_OF_ORIG_STRING)] # len = J == len(input_ids_S)
    """
    
    constraint_ls = [{"T": 0, "S": 0}]

    cur_T, cur_S = 0, 0
    while True:
        if offsets_T[cur_T][1] == offsets_S[cur_S][1]:
            constraint_ls.append({"T": cur_T, "S": cur_S})
            cur_T, cur_S = cur_T + 1, cur_S + 1
        elif offsets_T[cur_T][1] < offsets_S[cur_S][1]:
            cur_T += 1
        else: # offsets_T[cur_T][1] > offsets_S[cur_S][1]
            cur_S += 1
        
        if cur_T == len(offsets_T)-1 or cur_S == len(offsets_S)-1:
            constraint_ls.append({"T": len(offsets_T)-1, "S": len(offsets_S)-1})
            break

    return constraint_ls

offsets_T = teacher_tokenizer.encode(gen_text, offset_type="byte")['offsets']
offsets_S = student_tokenizer.encode(gen_text, offset_type="byte")['offsets']

constraint_ls = align_using_offsets(offsets_T, offsets_S)

[JqzChandler]

corresponding patch in tokenizers
huggingface/tokenizers#1880

[kashif]

thanks @JqzChandler yes agree that there can be edge cases with regards to the different tokenizer and the current method is perhaps fragile. let me look at your solution in detail and report back

[JqzChandler]

We implemented a simple demo for speed testing and issue reproduction, using XNLI dataset.

• For speed, the decode-based implementation will show more significant speed disadvantages in longer texts, since this is O(N^2) while considering each call of decode() as O(N).
• For correctness, beside the already discussed problem in BPE by which CJK characters may be splitted, we observed likely problems for Arabic letters can cause piece-level decoding errors due to their different forms in word-initial, word-medial, and word-final positions. (We cannot read such languages; I visually observed and attribute the cause to its "positional variation" rule.) And this new phenomenon is not limited to BPE; it can be reproduced on other tokenizers as well.

import sys
from pathlib import Path

import json
from tqdm import tqdm
from transformers import AutoTokenizer

GPT_TOKENIZER_PATH = "(...)/SomeGPT2Tokenizer"
BPE_TOKENIZER_PATH = "(...)/Qwen3-8B"
XNIL_JSONL_PATH = "(...)/XNLI-1.0/xnli.test.jsonl"

def load_and_simplify_xnil_samples(jsonl_path):
    samples = []
    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            samples.append(obj)
    text_ls = [f"[{row['language']}]{row['sentence1']}{row['sentence2']}" for row in samples]
    return text_ls

def decode_based_alignment(text: str, tokenizer_pair: dict, ignore_errors: bool=True):
    # coplied from trl's current impl marked in issue: https://github.com/huggingface/trl/issues/4393

    def to_canonical_pieces(tok, ids):
        pieces = []
        prev = ""
        for k in range(len(ids)):
            # IMPORTANT: Do NOT skip special tokens - we need to align them too
            cur = tok.decode(ids[: k + 1], skip_special_tokens=False, clean_up_tokenization_spaces=False)
            # Extract the incremental addition (may include spaces/ZWJ/etc.)
            pieces.append(cur[len(prev) :])
            prev = cur
        return pieces

    s_pieces = to_canonical_pieces(tokenizer_pair["S"], tokenizer_pair["S"].encode(text, add_special_tokens=False))
    t_pieces = to_canonical_pieces(tokenizer_pair["T"], tokenizer_pair["T"].encode(text, add_special_tokens=False))

    if not ignore_errors:
        assert ''.join(s_pieces) == text, f"s_pieces: {s_pieces} != text: {text}"
        assert ''.join(t_pieces) == text, f"t_pieces: {t_pieces} != text: {text}"

    i = j = 0
    s_buf = t_buf = ""
    s_group = []
    t_group = []
    s_groups = []
    t_groups = []

    def flush():
        if s_group and t_group:
            s_groups.append(s_group.copy())
            t_groups.append(t_group.copy())

    # Greedily accumulate pieces until substrings match, then flush
    while i < len(s_pieces) or j < len(t_pieces):
        if s_buf == t_buf and s_buf != "":
            flush()
            s_buf = t_buf = ""
            s_group = []
            t_group = []
            continue

        if s_buf == "" and i < len(s_pieces):
            s_buf += s_pieces[i]
            s_group.append(i)
            i += 1
            continue
        if t_buf == "" and j < len(t_pieces):
            t_buf += t_pieces[j]
            t_group.append(j)
            j += 1
            continue

        if len(s_buf) <= len(t_buf):
            if i < len(s_pieces):
                s_buf += s_pieces[i]
                s_group.append(i)
                i += 1
            elif j < len(t_pieces):
                t_buf += t_pieces[j]
                t_group.append(j)
                j += 1
        else:
            if j < len(t_pieces):
                t_buf += t_pieces[j]
                t_group.append(j)
                j += 1
            elif i < len(s_pieces):
                s_buf += s_pieces[i]
                s_group.append(i)
                i += 1

    # Flush any remainder if both sides accumulated something
    if s_buf == t_buf and s_group and t_group:
        flush()
    elif s_group or t_group:
        # Handle remaining unmatched tokens by forcing a flush
        # This ensures both sides have the same number of alignment groups
        if s_group or t_group:
            # Ensure both groups have content (even if empty list)
            if not s_group:
                s_group = []
            if not t_group:
                t_group = []
            # Force flush even if buffers don't match
            if s_group or t_group:
                s_groups.append(s_group.copy() if s_group else [])
                t_groups.append(t_group.copy() if t_group else [])

    return s_groups, t_groups
    
    pass

def encode_based_alignment(text: str, tokenizer_pair: dict, ignore_errors: bool=True):
    # currently requires a local built version of tokenizers with our patch: https://github.com/huggingface/tokenizers/pull/1880
    # and expose the new arg "offset_type" in local installed transformers as well

    offset_T = tokenizer_pair["T"].encode_with_offsets(text, offset_type="byte")['offsets']
    offset_S = tokenizer_pair["S"].encode_with_offsets(text, offset_type="byte")['offsets']

    def align_using_offsets(offsets_T: list, offsets_S: list):
        """
        offsets_T/_S be like:
        [(0, 1), (1, 3), (3, 8), ..., (start_i, end_i), ..., (start_I, TOTAL_BYTES_OF_ORIG_STRING)] # len = I == len(input_ids_T)
        [(0, 3), (3, 5), (5, 8), ..., (start_j, end_j), ..., (start_J, TOTAL_BYTES_OF_ORIG_STRING)] # len = J == len(input_ids_S)
        """
        
        constraint_ls = [{"T": 0, "S": 0}]

        cur_T, cur_S = 0, 0
        while True:
            if offsets_T[cur_T][1] == offsets_S[cur_S][1]:
                constraint_ls.append({"T": cur_T, "S": cur_S})
                cur_T, cur_S = cur_T + 1, cur_S + 1
            elif offsets_T[cur_T][1] < offsets_S[cur_S][1]:
                cur_T += 1
            else: # offsets_T[cur_T][1] > offsets_S[cur_S][1]
                cur_S += 1
            
            if cur_T == len(offsets_T)-1 or cur_S == len(offsets_S)-1:
                constraint_ls.append({"T": len(offsets_T)-1, "S": len(offsets_S)-1})
                break

        return constraint_ls

    constraint_ls = align_using_offsets(offset_T, offset_S)

    return constraint_ls

if __name__ == "__main__":
    gpt_tokenizer = AutoTokenizer.from_pretrained(GPT_TOKENIZER_PATH)
    bpe_tokenizer = AutoTokenizer.from_pretrained(BPE_TOKENIZER_PATH)
    tokenizer_pair = {"T": gpt_tokenizer, "S": bpe_tokenizer}

    text_ls = load_and_simplify_xnil_samples(XNIL_JSONL_PATH)
    for text in tqdm(text_ls):
        _ = encode_based_alignment(text, tokenizer_pair)
    for text in tqdm(text_ls):
        _ = decode_based_alignment(text, tokenizer_pair, ignore_errors=True) # pass ignore_errors=False for speed test

[kashif]

thanks @JqzChandler that will help!

[kashif]

@JqzChandler i will open a PR and then share it with you if you want to contribute?

[JqzChandler]

@kashif happy to contribute!