Fix SFT mask to avoid calculating loss on generation prompt tokens

[ChingTsai]

Description

• Modern models like Llama 3 or Qwen add generation prompt tokens to indicate the start of an assistant response. These are appened during inference, therefore, are not part of the completion tokens.

• Many LLM fine-tuning frameworks, such as unsloth+trl, and ms-swift, do not calculate loss on the generation prompt.

  • Unsloth splits tokens after generation_prompt as the 'response_part'
  • Visualise Unslosh masking
    

• This PR aligns the MaxText SFT masking logic with these other frameworks

  • Include generation_prompt tokens as part of prompt tokens here, hence, the loss would not be calculated during SFT.

FIXES: b/475383130

Tests

Unit Test

python3 -m pytest  tests/unit/sft_data_processing_test.py

Compare sft qwen3-4b initial loss

• Drop from 5.868 to 5.1988
  

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!