Extend DataCollatorForCompletionOnlyLM to support correct tool result masking

[shanghongsim]

Description

Context

Currently, the completions only collator does not consistently mask tool responses (Role.TOOL) correctly. For some chat templates (like Qwen), it accidentally gets it correct, while masking is consistently wrongly applied for llama chat templates. This PR enables the masking to be correctly applied for all chat templates, supporting all four roles (system, user, assistant and tool).

Without instruction_template, tool results are masked correctly. But with instruction_template, things get weird and some tool results are unmasked. Omitting instruction_template causes the collator to find the last response_template and masks everything before it (only the final assistant turn trains). With instruction_template, the collator searches for instruction_template and response_template pairs and masks everything in between.

Case 1: without instruction_template

RESPONSE_TEMPLATE = "<|im_start|>assistant\n"

old_no_inst = DataCollatorForCompletionOnlyLM(
    response_template=RESPONSE_TEMPLATE,
    instruction_template=None,
    tokenizer=tokenizer,
)
b = old_no_inst.torch_call([token_ids])
labels_A = b["labels"][0].tolist()
summarise("Case A", labels_A, N)
show_masking(b["input_ids"][0].tolist(), labels_A, tokenizer)

Without inst template, collator masks everything before the last resp template. So loss only sees

<|im_start|>assistant
The weather in Paris is sunny and 18°C.<|im_end|>

Case 2: with instruction_template

With instruction_template, it masks everything between a inst and resp template. Since there isn't an inst template before the tool result, it does not get masked properly.

INSTRUCTION_TEMPLATE = "<|im_start|>user\n"

old_with_inst = DataCollatorForCompletionOnlyLM(
    response_template=RESPONSE_TEMPLATE,
    instruction_template=INSTRUCTION_TEMPLATE,
    tokenizer=tokenizer,
)
b2 = old_with_inst.torch_call([token_ids])
labels_B = b2["labels"][0].tolist()
summarise("Case B", labels_B, N)
show_masking(b2["input_ids"][0].tolist(), labels_B, tokenizer)

In case 2, the loss sees tool result, which is incorrect.

...
<|im_start|>assistant
<tool_call>
{"name": "get_weather", "arguments": {"location": "Paris"}}
</tool_call><|im_end|>
...
<|im_start|>assistant
The weather in Paris is sunny and 18°C.<|im_end|>

To confirm my hypothesis of no inst template before the tool result being the cause, I experimented with adding a user turn before the tool call and masking works correctly in that case.

ToolAwareCompletionsCollator

With the new ToolAwareCompletionsCollator, tool results are masked properly

Without instruction_template

With instruction_template

Change 1: Extend DataCollatorForCompletionOnlyLM with span-based masking

Extend DataCollatorForCompletionOnlyLM with span-based masking (detecting assistant turn boundaries via response + end-of-turn tokens) beyond instruction based masking (matching instruction/response string pairs) supported currently.

Change 2: train_target enum for explicit control of which part of response to train on

• assistant_turn: Masks everything, then unmasks each assistant response bounded by response_template .. end_of_turn_template (inclusive of EOT).
• final_assistant_turn: Masks all tokens before the last response_template occurrence.

When train_target is set, no matter what templates have been provided, train_target has final authority on which combination is used. When train_target is not set, it is inferred from the templates provided. Eventually, we want users to specify train_target instead of relying on combinations of templates to control behavior. However, since we have many configs in production that do not use train_target and set templates explicitly, this shall be the interim solution to ensure old configs do not break. Existing configs using collator_kwargs with instruction_template + response_template continue to work via the legacy path (with a deprecation warning). New configs should use train_target explicitly. Users can set train_target and override the templates manually by setting collator_kwargs.

Before — users must provide model-specific token strings:

collator_name: "text_completions_only_with_padding"
collator_kwargs:
  response_template: "<|im_start|>assistant\n"
  end_of_turn_template: "<|im_end|>"

After — users express intent, templates auto-resolve from tokenizer:

collator_name: "text_completions_only_with_padding"
train_target: "assistant_turn"

Change 3: Auto detection of collator templates

Instead of fragile vocab and collator template matching (requires us to maintain collator templates for all popular models), shift to a more robust approach: apply the chat template to a known test conversation, then finds the assistant boundary strings in the rendered output.

Migration and deprecation considerations:

1. Configs that only specify collator_name -> removed support for this with the removal of the default llama template fallback. This should be ok as no configs in production specify collator_name without collator_kwargs. Only less than 10 configs in OSS fall in this edgecase and they are being updated in this PR.
2. Enterprise SFT configs -> they should still work with old legacy behavior. Will update them once OSS version is updated in API.

Open issues for next time

• Movetrain_target into collator_kwargs -> I considered putting train_target inside collator_kwargs instead of making it a separate field, but collator_kwargs is an opaque dict. Users wouldn't know train_target exists unless they read the docs. A top-level field makes it more visible. This will eventually be solved when we shift to CollatorParam.
• CollatorParams: This would mean changing every YAML from:

  collator_name: "text_completions_only_with_padding"
  collator_kwargs:
    train_target: "all_assistant_turns"

to

  collator:
    name: "text_completions_only_with_padding"
    train_target: "all_assistant_turns"

This is much cleaner but its a substantial refactor. In order to manage the scope of this PR, this shall be deferred to a future PR

• Vision specific concerns -> future PR to neaten the vision concerns that are now scattered throughout the LM specific things
• Per-split collator configuration -> build_collator_from_config builds a single collator from the training split's settings, which is reused for all splits (train, validation, test). Validation/test split-specific collator settings (e.g. a different train_target) are silently ignored. This is pre-existing behavior not introduced by this PR, but worth noting as a known limitation. Fixing it would require refactoring the training loop to build per-split collators.

Related issues

N/A

Before submitting

• This PR only changes documentation.
• Did you read the contributor guideline Pull Request guidelines?
• Did you link the issue(s) related to this PR in the section above?
• Did you add / update tests where needed?

Reviewers

At least one review from a member of oumi-ai/oumi-staff is required.

[gitar-bot]

Gitar is working

_Gitar

[sentry]

Bug: The TextCompletionsCollatorWithPadding constructor can receive ignore_index both as an explicit argument and within **kwargs, causing a TypeError if a user customizes it.
_{Severity: HIGH}

Suggested Fix

Before calling the TextCompletionsCollatorWithPadding constructor, pop ignore_index from the kwargs dictionary. Use the popped value to determine the final ignore_index value to be passed, preventing the argument duplication.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: src/oumi/builders/collators.py#L243-L249

Potential issue: When building the `text_completions_only_with_padding` collator in
`build_data_collator`, the `ignore_index` parameter is passed explicitly while also
being potentially present in the `**kwargs` passed to the
`TextCompletionsCollatorWithPadding` constructor. If a user specifies `ignore_index` in
their `collator_kwargs` configuration, this results in a `TypeError` because the
`__init__` method receives multiple values for the same keyword argument, causing a
runtime crash during training setup.

[jgreer013]

This is really clever, but we should maybe include a system instruction here?

[shanghongsim]

Like modify it to extract system instruction? We are current no longer relying on system instruction for masking

[jgreer013]

No I mean ensure that when we extract the different template portions, we do so using an example conversation that also includes a system instruction

Or rather, perhaps do it twice, once on a conversation with a system instruction, and once on a conversation without.

Why? Because I'm a bit concerned of the scenario where this logic doesn't use system instructions when extracting, but the user does in their data, and somehow the resulting templates we extract using this methodology don't wind up working when the data includes the presence of a system instruction for some reason.

[jgreer013]

nit: first_asst_start, and add _start to second turns

[jgreer013]

extract this to its own method

[jgreer013]

Same here, extract to its own method

[jgreer013]

This description scares me, I wonder if we could have a better workaround or a louder error

[jgreer013]

It's not clear to me what this collator does differently than the other one.

[shanghongsim]

Previously, it would instruction_template and response_template pairs then mask everything in between (see L324). However, there isn't an instruction_template before the tool result, so tool result content does not get masked properly. Some models (like Qwen) accidentally mask it correctly because their chat template stores tool result under role.user. But for most models (like Llama etc), tool result is stored under a role.tool with a separate tag so masking is incorrectly applied for the reasons above. The new approach basically masks everything then only unmask portions between the response_template and end_of_turn_template. This is more robust to the different roles and specifics of how tool results are formatted and stored by different chat templates.

[lefft]

Can we be sure to check whether template detection and collation happens properly when user or assistant turns begin with (one or more) leading \ns? We had issues with this breaking tokenization and leading to records being skipped in enterprise, would like to understand whether these changes will help, harm, or not impact that issue.

[lefft]

Has this been tested with gpt-oss models? They use a non-standard system for marking boundaries within responses, so wondering if boundary detection will work as it does for Qwen/Llama style templates.

[lefft]

Is there a plan/intention to eliminate support for instruction-based masking altogether, or will we just keep it in a "deprecated" state?

Want to make sure we retain compatibility with existing enterprise configs (or verify that moving to this style doesn't change training dynamics before adapting enterprise configs).