feat(xtoken): replace DP token alignment with offset-based cluster alignment#3286
Open
avenkateshha wants to merge 3 commits into
Open
feat(xtoken): replace DP token alignment with offset-based cluster alignment#3286avenkateshha wants to merge 3 commits into
avenkateshha wants to merge 3 commits into
Conversation
…ignment Cross-tokenizer alignment now pairs student and teacher tokens by the character spans they cover (return_offsets_mapping) instead of a Needleman-Wunsch DP over canonicalized token strings. Consecutive same-span tokens collapse into a cluster and a strict char-end walker pairs clusters covering the same character range; special tokens (offset (0,0)) pair by role. is_correct is still recomputed from canonicalized text, so pair_is_correct and the gold-loss exact partition keep their meaning. The collator emits per-token char offsets and threads them into TokenAligner.align. The DP kernel, anchor segmentation, post-processing, and byte/mojibake canonicalization helpers are removed; canonical_token (still used by tools/x_token and the is_correct mask) stays. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
Adds the chat/instruct alignment path to TokenAligner (Part 2 of the offset-alignment PR): align_one_offset_per_asst aligns each assistant message independently (isolate content, rebase offsets to a shared coordinate system, run the offset-cluster kernel on the slice, map positions back to the full sequence), with native reasoning/close/answer regions, a synthetic EOT pair, drop_first_content_pair, and included_regions. align_one_offset adds the single-sample offset path with an NFC decode-fix is_correct mask, and the constructor gains alignment_method (offset_cluster_decode_fix | same_tokenizer_identity). Returns raw 7-tuples; the chat collator (next commit) converts to AlignmentPair. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
Adds a chat mode to CrossTokenizerCollator (Part 2 of the offset-alignment PR). In chat mode the collator renders each side's chat template, tokenizes with char offsets, derives the per-token assistant mask + each assistant message's char span, and aligns each assistant message independently via TokenAligner.align_one_offset_per_asst (multi-teacher). Loss fires only on assistant content (token_mask = attention_mask * assistant_mask). find_rendered_message_content_span (chat_templates.py) locates message content in the rendered string. Config knobs (collator_mode, drop_first_assistant_chunk_kl, include_thinking_in_loss, native_thinking_alignment, kd_alignment_regions, num_packed_rows) are wired from data_config and default to the raw-text path. Scope: whole-message chat alignment. native_thinking_alignment (native reasoning/close/answer parsing) and num_packed_rows>1 (lockstep packing) raise NotImplementedError as clearly-marked follow-ups. Validated by unit tests (char-level fake chat tokenizer); a GPU chat smoke on real instruct data is a follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Adithya Hanasoge <avenkateshha@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Moves cross-tokenizer alignment from Needleman-Wunsch DP to an offset-based cluster aligner, and extends it to chat / instruct (and "thinking") data.
Alignment pairs tokens by their character-span offsets (
return_offsets_mapping): consecutive same-span tokens collapse into a cluster, a strict char-end walker pairs clusters covering the same character range, and special tokens (offset(0,0)) pair by role.is_correctis recomputed from canonicalized text, sopair_is_correctand the gold-loss exact partition keep their meaning.Scope / status
Part 1 — offset alignment (done)
nemo_rl/algorithms/x_token/token_aligner.py— DP kernel, anchor segmentation, post-processing, and byte/mojibake canonicalization removed; the offset-cluster kernel (align_by_offsets_cluster) added.canonical_token(still used bytools/x_tokenand theis_correctmask) kept.nemo_rl/data/cross_tokenizer_collate.py—_tokenize_batchemitsoffset_mapping;align()receives student/teacher char offsets.test_token_aligner.pyrewritten for the offset path; the collator's fake tokenizer emits char offsets.Part 2 — chat / native-template alignment (this PR)
Each model renders a conversation with its own chat template, so the two rendered texts — and their offsets — differ, and the full sequence can't be aligned directly. This part aligns per assistant message: isolate each side's assistant content, rebase offsets to a shared coordinate system, align the slice with the same
align_by_offsets_clusterkernel, then map positions back to the full sequence. "Thinking" data is split intoreasoning/close/answerregions (+ an EOT marker).Adds:
token_aligner.py—align_one_offset_per_asst()(per-turn slice → rebase → align → map-back) andalign_one_offset()(canonical + NFC decode-fixis_correct).cross_tokenizer_collate.py— achatmode (__call__→_call_text/_call_chat) with render / assistant-mask / native-region / lockstep-packing helpers and new constructor knobs.chat_templates.py—find_rendered_message_content_span().processors.py—chat_kd_processor()(readsmessages/conversation).collator_mode,include_thinking_in_loss,native_thinking_alignment,drop_first_assistant_chunk_kl,kd_alignment_regions,num_packed_rows.Testing
🤖 Generated with Claude Code