Skip to content

v-code01/maskbreak

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

5 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

maskbreak

apply_chat_template(..., return_assistant_tokens_mask=True) returns the assistant loss mask as all zeros when the conversation is truncated from the left. Left truncation is the standard setting for decoder-only chat models, where an over-long conversation drops its oldest turns to keep the most recent one. The mask builder walks the assistant spans in document order and breaks out of the loop at the first span whose start no longer maps to a token. Under left truncation that first span is the one that was dropped, so the break fires immediately and the surviving final assistant turn is never marked. The training example then contributes zero loss on exactly the tokens the model is supposed to learn to produce.

The code

src/transformers/tokenization_utils_base.py, apply_chat_template, current main (lines 3151-3157; identical in the 5.12.1 release):

for assistant_start_char, assistant_end_char in generation_indices[i]:
    start_token = out.char_to_token(i, assistant_start_char)
    end_token = out.char_to_token(i, assistant_end_char - 1)
    if start_token is None:
        # start_token is out of bounds maybe due to truncation.
        break
    for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])):
        current_mask[token_id] = 1

The break assumes a span whose start maps to no token was truncated from the end, so every later span is gone too. That holds only when truncation removes the tail. When truncation removes the head, the early span is missing while later spans survive, and the break throws them away with it. The correct control-flow is continue.

What it does

The mask must be 1 exactly on the surviving assistant tokens. With a two-turn conversation whose oldest turn is dropped from the left:

left truncation (drop oldest turn):
  transformers masks 0 assistant tokens, correct masks 3
  the surviving assistant turn trains with zero loss

The base case and the right-truncation control confirm the harness and isolate the bug: with no truncation both turns are masked, and with right truncation the two paths agree because the surviving span comes first and the break never fires. Only left truncation diverges. This reproduces directly on transformers 5.12.1: a Qwen2.5 tokenizer with a {% generation %} chat template, truncation_side="left", and a max_length that drops the first assistant turn returns an empty assistant mask, where the surviving FINAL RESPONSE keep me tokens should be marked.

The same two lines carry a second, narrower defect. The end-token guard end_token + 1 if end_token else len(input_ids[i]) uses truthiness, but char_to_token can legitimately return 0. When an assistant span ends at token 0 the guard treats it as missing and extends the mask to the end of the whole sequence, masking the user tokens too:

assistant span ending at token 0 (truthiness guard):
  transformers masks 6 of 6 tokens, correct masks 1
  the whole sequence, including user tokens, is trained on

This one needs assistant content to occupy token index 0, which a normal chat template avoids by emitting a role header first, so it is a latent fragility rather than a mainstream trigger. The correct guard is end_token is not None.

Scope

Both defects are on current main, in a supported combination: the char_to_token is None branch exists specifically to handle truncation, and truncation_side is a first-class documented setting. Nothing states that assistant masks require right truncation. Any instruction-tuning pipeline that left-truncates long conversations (the common case) silently trains with zero assistant loss on every truncated example. The fix is one line each: break to continue, and if end_token to if end_token is not None.

Layout

  • excerpt.py: the mask loop quoted with line numbers, and the two flags for the break and the truthiness guard (Apache-2.0).
  • chatmask.py: the current-main loop and a correct version, differing only in break vs continue and the end-token guard, over a character-to-token map that models truncation.
  • consequence.py: the zero-loss left-truncation case and the token-0 over-masking case.
  • test_maskbreak.py: the no-truncation base case matches, right truncation matches, left truncation zeroes the surviving span, and the token-0 guard over-masks the whole sequence.

Reproduce

python chatmask.py
python consequence.py
python test_maskbreak.py

The loop is quoted from current main. The character-to-token model reproduces the branch behaviour without a model download; the finding was confirmed against transformers 5.12.1 with a real Qwen2.5 tokenizer and a left-truncated conversation.

About

transformers apply_chat_template returns an all-zero assistant loss mask under left truncation: break aborts the loop at the dropped early span before masking the surviving assistant turn

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages