Skip to content

Preserve group types and whitespace boundaries in trajectory tensorization - #904

Merged
bradhilton merged 2 commits into
mainfrom
hayek/text-equivalent-prompt-masks
Sep 16, 2026
Merged

bradhilton merged 2 commits into
mainfrom
hayek/text-equivalent-prompt-masks

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

Conventional multi-history training currently needs a cast after group.tensorize(), and some saved Qwen histories fail tokenization when the server and chat template encode identical whitespace differently.

  • Add overloads to TokenizedTrajectoryGroup.tensorize() that preserve the single-history or multi-history trajectory type. Runtime tensorization is unchanged; callers can access .histories without a cast.
  • Translate assistant/output masks across proven whitespace-equivalent replacements. A single merged token inherits participation from its characters; replacements with multiple target tokens require a uniform source mask. Served token IDs and log probabilities stay intact, and stop masks retain their existing strict token mapping.

The whitespace failure occurs when a length-stopped Qwen response forces rendering of a history whose prompt encoded two newlines as one token, while the renderer uses two tokens straddling the assistant-prefix boundary. Both fixes support Caladan experiment 056. This PR builds on the merged #903. Rebasing onto its squash merge preserved the incremental patch byte-for-byte.

Validation

  • 308 focused tokenization, compact-serialization, tokenized-model, and tensorized-model tests passed on the rebased commit with Preserve NaN log probabilities in trajectory JSON #903 included.
  • 282 tokenization/compact-serialization tests passed for the whitespace change, including the public regression, mask ownership cases, and rejection of changed text or ambiguous multi-token boundaries.
  • 12 tensorized-model tests passed; assert_type checks cover both group overloads, and Caladan's training loop type-checks without a cast.
  • Ruff, formatting, type checks, and git diff --check passed.
  • Both affected saved Qwen replays tokenize with the expected action and sampled-token masks.

@bradhilton bradhilton changed the title Handle merged whitespace tokens in exact prompt masks Preserve group types and whitespace boundaries in trajectory tensorization Sep 16, 2026
@bradhilton
bradhilton changed the base branch from hayek/preserve-nan-logprobs-json to main September 16, 2026 03:26
@bradhilton
bradhilton force-pushed the hayek/text-equivalent-prompt-masks branch from 76d327f to 2283a7c Compare September 16, 2026 03:26
@bradhilton
bradhilton marked this pull request as ready for review September 16, 2026 03:30
@bradhilton
bradhilton merged commit 53f2431 into main Sep 16, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant