Skip to content

Add semantic SFT loss masks#743

Open
vivekkalyan wants to merge 4 commits into
mainfrom
vivek/sft-loss-mask
Open

Add semantic SFT loss masks#743
vivekkalyan wants to merge 4 commits into
mainfrom
vivek/sft-loss-mask

Conversation

@vivekkalyan

@vivekkalyan vivekkalyan commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds semantic SFT loss masks to ART trajectories and compiles them into the existing token-level label tensors during SFT preprocessing. This lets local/Megatron SFT train only selected assistant turns, such as the last GPT-4.1 continuation in a multi-turn production trace, while preserving the full prompt context.

Changes

  • Add typed SFT loss mask specs on Trajectory: all, none, last_assistant, message_mask, and roles.
  • Compile semantic masks in tokenizer preprocessing by rendering message-prefix spans and writing selected tokens into labels while keeping ignored positions at -100.
  • Preserve loss_mask through JSONL SFT input, trajectory logging, and migration paths.
  • Reject trajectory.loss_mask on serverless SFT for now, since this support is intended for local/Megatron SFT.
  • Add Qwen disable-thinking response-part configs so empty <think> scaffolds are not trained as target content.

Why

Production-trace SFT needs to keep full multi-turn context but apply loss only to verified GPT-4.1 assistant outputs. Message-index masks in the dataset are fragile; this keeps the dataset representation semantic and lets ART compile it at tokenization time.

Validation

  • uv run --no-sync pytest tests/unit/test_trajectory_parquet.py tests/unit/test_preprocessing_tokenize.py tests/unit/test_sft.py tests/unit/test_serverless_pipeline_trainer_compat.py -q
  • uv run --no-sync prek run --all-files

* origin/main:
  feat: unify shared-prefix Megatron execution (#739)
  Add DeepSeek V4 Flash support for ART Megatron (#745)
  Add ART Megatron support for GPT OSS (#744)
  ci: bound GPU prewarm cleanup
  ci: prewarm GPU image on all GPU clusters
  ci: dispatch Caladan GPU image builds
  Fix Gemma4 MoE Triton flex stage config (#742)
  fix: split GPU image dependency warmup

# Conflicts:
#	src/art/dev/model.py
#	src/art/serverless/backend.py
@vivekkalyan vivekkalyan marked this pull request as ready for review July 9, 2026 00:52
@vivekkalyan vivekkalyan requested a review from Kovbo July 9, 2026 00:52
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