[AusarYao28]

Description

DeepSeek-V4 checkpoints ship no Jinja chat_template, so input_pipeline_utils.apply_chat_template fails during SFT preprocessing. The official encoder (encoding_dsv4.py from the DeepSeek-V4 Hugging Face repo) was only vendored under benchmarks/api_server/encoding/, which src/maxtext cannot import. That copy had also lost the blank lines in TOOLS_TEMPLATE (and several docstrings), so the tool-calling prompts it renders no longer match the official encoder goldens.

• Move benchmarks/api_server/encoding/encoding_dsv4.py to src/maxtext/input_pipeline/encoding_dsv4.py, restore the stripped blank lines, and import it from benchmarks/api_server/{maxtext_server,server_utils}.py. One copy now serves SFT and the API server.
• input_pipeline_utils.apply_chat_template: when chat_template is None and the model or tokenizer name is DeepSeek-V4 (_is_deepseek_v4; the name check only, since the V3.2 tokenizer has the same special tokens), encode with encode_messages and split it into prompt and completion chunks. Messages after the final assistant turn are dropped, as in the template path, and any encoded text left after the final completion raises ValueError.
• In thinking mode, encode_messages drops the reasoning of assistant turns before the last user message (they are rendered after </think>). Those turns stay in the prompt; only turns that keep their reasoning are trained. Chat-mode turns are all trained.
• grain_data_processing.sft_preprocessing_pipeline: build the tokenizer with add_bos=False, add_eos=False. Chunks are encoded one by one, and with transformers 5.x add_bos_token=True makes the DeepSeek-V4 tokenizer prepend BOS to every chunk. The chat template or encoder already emits BOS/EOS. The HF SFT path and vision SFT already disable both.
• grain_data_processing and hf_data_processing: forward model_name to apply_chat_template.

Copybara note: //third_party/py/maxtext/src/maxtext/input_pipeline/BUILD lists srcs explicitly, so the import CL needs "encoding_dsv4.py" added to the :input_pipeline py_library. No new deps are needed: the module imports only the standard library. benchmarks/ has no internal BUILD target.

The pad_token == eos_token loss-masking fix is in #5853, and the DeepSeek-V4-Pro E2E SFT stage is in #5884.

FIXES: b/568401437

Tests

• pytest tests/unit/deepseek_v4_sft_masking_test.py tests/unit/instruction_data_processing_test.py tests/unit/hf_data_processing_test.py: 70 passed. Coverage:
  • golden parity of encode_messages with the 4 official encoder fixtures (tests/utils/deepseek4_reference/testdata);
  • Grain and HF SFT loss-mask checks, packed and unpacked, against labels derived independently from the encoder output;
  • _is_deepseek_v4, including a V3.2-style tokenizer with the same special tokens;
  • thinking-mode history turns staying in the prompt, and the Grain tokenizer built with add_bos=False, add_eos=False;
  • real-tokenizer loss masks with the tokenizer built the way the Grain pipeline builds it (checked locally with the cached DeepSeek-V4-Pro tokenizer);
  • model_name forwarding through hf_data_processing.preprocessing_pipeline and _format_chat_template_grain.
• Control: the unmodified benchmarks/api_server copy fails the golden test on fixtures 1 and 3 (the tool-calling prompts); with the blank lines restored, all 4 pass.
• benchmarks.api_server.server_utils imports encoding_dsv4 from maxtext.input_pipeline.
• End-to-end on TPU (deepseek4-1600b, 4 x v5p-128, 5-step train_sft_native on ultrachat_200k). The run used commit 8da2bf4ee, the pre-split branch that also carried the pad fix from Fix DeepSeek SFT pad_token == eos_token loss masking #5853.
  • Cloud Logging share link: https://cloudlogging.app.goo.gl/W8GtS5FMFi4hZ6Bt9
  • completed step: 0, seconds: 17.345, TFLOP/s/device: 17.845, Tokens/s/device: 59.037, total_weights: 169422, loss: 1.062, lm_loss: 1.062, perplexity: 2.893, moe_lb_loss: 0.000142, indexer_loss: 0.0
  • completed step: 4, seconds: 61.169, TFLOP/s/device: 5.060, Tokens/s/device: 16.740, total_weights: 173541, loss: 1.061, lm_loss: 1.061, perplexity: 2.889, moe_lb_loss: 0.000142, indexer_loss: 0.0

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 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.

[github-actions]

🤖 Hi @AusarYao28, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 Hi @AusarYao28, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[gemini-code-assist]

Code Review

This pull request integrates the official DeepSeek-V4 chat encoding and parsing logic into MaxText's input pipeline, enabling conversational formatting and SFT loss-masking support for DeepSeek-V4. It updates the Grain and Hugging Face data processing pipelines to pass the model name, implements DeepSeek-V4 detection and formatting utilities, and adds comprehensive unit tests. Feedback on the changes highlights a potential false positive in the DeepSeek-V4 detection logic when convert_tokens_to_ids returns an unknown token ID instead of None, and recommends replacing a bare except: clause with except Exception: to avoid catching system-exiting exceptions.

[gemini-code-assist]

For standard Hugging Face tokenizers, convert_tokens_to_ids returns the unk_token_id (which is an integer, e.g., 0) if a token is not found in the vocabulary. Therefore, checking convert_tokens_to_ids(...) is not None will evaluate to True for almost any tokenizer that has an unk_token_id defined, leading to false positives where non-DeepSeek models are incorrectly identified as DeepSeek-V4.\n\nTo prevent this, we should ensure that the returned token IDs are not equal to the tokenizer's unk_token_id.

    elif hasattr(tokenizer_model, "convert_tokens_to_ids"):\n      unk_id = getattr(tokenizer_model, "unk_token_id", None)\n      dsml_id = tokenizer_model.convert_tokens_to_ids("｜DSML｜")\n      user_id = tokenizer_model.convert_tokens_to_ids("<｜User｜>")\n      if dsml_id is not None and dsml_id != unk_id and user_id is not None and user_id != unk_id:\n        return True

[AusarYao28]

Fixed in 2500ee9fc: the convert_tokens_to_ids fallback now requires both ids to differ from unk_token_id. test_is_deepseek_v4_by_name covers a tokenizer that maps every token to unk_token_id.

[gemini-code-assist]

Avoid using bare except: clauses as they can catch system-exiting exceptions like SystemExit and KeyboardInterrupt, making it harder to interrupt the program. Use except Exception: instead.

Suggested change

	except:
	except Exception:

References

1. PEP 8 recommends avoiding bare except clauses. ^(link)

[AusarYao28]

The file is kept byte-for-byte with the official upstream encoding_dsv4.py (apart from restoring its blank lines), so it can be re-synced and diffed against upstream; it already carries # pylint: disable=bare-except. The try only wraps json.dumps, so leaving it as is has no practical effect.

[codecov]

Codecov Report

❌ Patch coverage is 73.80952% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/input_pipeline_utils.py	72.50%	8 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

📋 Review Summary

Verdict: Spec partially met with defects: conversation truncation risk, fragile tokenizer heuristic, duplicate vendored module bloat, and an unaddressed production EOS-loss bug.

Top issues:

1. [Bug] _format_conversations_dsv4 fails to verify cursor == len(text) after matching turns, silently dropping trailing conversation content (such as trailing user messages). → Validate cursor == len(text) or raise a ValueError when trailing content is detected.
2. [Bug] _is_deepseek_v4 fallback check convert_tokens_to_ids(...) is not None returns True for any tokenizer without get_vocab whose missing tokens map to unk_token_id. → Ensure retrieved token IDs are not equal to tokenizer_model.unk_token_id.
3. [Complexity] 752 lines of vendored DeepSeek-V4 chat encoder and DSML parser code were appended to instruction_data_processing.py, duplicating benchmarks/api_server/encoding/encoding_dsv4.py almost verbatim along with unused streaming/decoding functions. → Move the vendor encoder to a dedicated module (e.g., src/maxtext/input_pipeline/deepseek_v4_encoder.py) and share it with benchmarks/api_server/.
4. [Tests] deepseek_v4_sft_masking_test.py documents _EOS_DEFECT (where pad_id == eos_id == 1 drops all completion EOS targets from loss) but avoids it by testing only pad_id=2 with dead xfail helpers, while MaxText production tokenizer loading remains unconfigured for <｜▁pad▁｜>. → Configure pad_token = "<｜▁pad▁｜>" (id 2) for DeepSeek-V4 in MaxText tokenizer loading and restore end-to-end regression validation.
5. [Tests] hf_data_processing.py forwards model_name to apply_chat_template, but lacks unit tests verifying that HF pipeline preprocessing handles DeepSeek-V4 without a Jinja template. → Add a unit test for hf_data_processing.preprocessing_pipeline with model_name and chat_template=None.

Inferred spec & validation

• Sources: PR description; Buganizer ticket b/568401437 (referenced in description, not readable here); code before/after in src/maxtext/input_pipeline/ and tests/unit/deepseek_v4_sft_masking_test.py.
• Problem: DeepSeek-V4 checkpoints do not ship a Jinja chat_template in tokenizer_config.json, causing input_pipeline_utils.apply_chat_template to raise ValueError during SFT preprocessing.

Spec item	Implementation	Validating test
R1 (stated): Support DeepSeek-V4 conversational SFT preprocessing when tokenizer.chat_template is None	⚠️ partial (silently drops trailing conversation text if cursor != len(text))	deepseek_v4_sft_masking_test.py::test_apply_chat_template_deepseek_v4 ✅
R2 (stated): Correctly encode thinking mode, multi-turn chat, and DSML tool calls into prompt/completion turns	✅ preserved	deepseek_v4_sft_masking_test.py::test_vendored_encoder_matches_official_golden ✅
R3 (stated): Forward model_name from SFT config to apply_chat_template in Grain and HF pipelines	⚠️ partial (Grain and HF updated; Grain tested, HF untested)	deepseek_v4_sft_masking_test.py::test_format_chat_template_grain_deepseek_v4 ✅; HF pipeline ❌ none
R4 (inferred): Prevent false-positive DeepSeek-V4 detection for other tokenizers	⚠️ partial (convert_tokens_to_ids heuristic false-positives when unk_token_id is an int)	deepseek_v4_sft_masking_test.py::test_apply_chat_template_non_deepseek_without_template_raises ⚠️ superficial (mocks convert_tokens_to_ids to None)
C1 (stated): Preserve existing Jinja chat template behavior for tokenizers shipping chat_template	✅ preserved	tests/unit/chat_template_sft_test.py ✅
C2 (inferred): SFT training computes loss on assistant completion tokens (including EOS)	⚠️ partial (pad_id=2 tested in isolation, but production MaxText tokenizer loader still yields pad_id=1, triggering documented _EOS_DEFECT)	deepseek_v4_sft_masking_test.py::test_loss_mask_matches_official_format ⚠️ superficial (pad_id=1 dropped from test grid)

[github-actions]

🟡 While grain_data_processing has test_format_chat_template_grain_deepseek_v4, hf_data_processing lacks a unit test verifying that model_name is forwarded to apply_chat_template or that HuggingFace dataset preprocessing completes without errors for DeepSeek-V4. Add test coverage for hf_data_processing.preprocessing_pipeline with DeepSeek-V4 and missing chat_template.

[AusarYao28]

Added HfDataProcessingTest.test_sft_forwards_model_name_to_apply_chat_template in 2500ee9fc. It runs hf_data_processing.preprocessing_pipeline(use_sft=True) and asserts that apply_chat_template receives model_name.

[github-actions]

🟡 _param_grid drops xfail_reason and xfail_if (del xfail_reason, xfail_if) and removes pad_id=1 testing entirely, leaving dead helpers (_is_ds4_pad, _is_ds4_pad_unpacked) and directly contradicting the module docstring claiming pad_id=1 is tested as xfail(strict). More importantly, pad_token = "<｜▁pad▁｜>" (id 2) is only set in this unit test file; MaxText tokenizer loading (tokenizer.py, hf_data_processing._get_pad_id, data_processing_utils.get_tokenizer_and_pad_id) was not updated, leaving real DeepSeek-V4 SFT runs with pad_id == eos_id == 1 and triggering the documented _EOS_DEFECT where all completion EOS targets are dropped from the training loss. Ensure pad_token / pad_id is properly configured for DeepSeek-V4 in production code.

[AusarYao28]

Removed the unused _EOS_DEFECT, _SEGMENT_DEFECT, _is_ds4_pad* helpers and the stale docstring in 2500ee9fc. The production pad_token == eos_token remap (fix_pad_token_if_needed in the Grain and HF SFT pipelines) is in #5853; this PR was split from it at reviewer request.

[github-actions]

🟠 In _is_deepseek_v4, the fallback check using convert_tokens_to_ids(...) is not None returns True for any tokenizer without get_vocab where unseen tokens resolve to unk_token_id (an int, e.g., 0). Check that the token ID is not equal to unk_token_id.

Suggested change

	elif hasattr(tokenizer_model, "convert_tokens_to_ids"):
	if (
	tokenizer_model.convert_tokens_to_ids("｜DSML｜") is not None
	and tokenizer_model.convert_tokens_to_ids("<｜User｜>") is not None
	):
	return True
	elif hasattr(tokenizer_model, "convert_tokens_to_ids"):
	unk_id = getattr(tokenizer_model, "unk_token_id", None)
	dsml_id = tokenizer_model.convert_tokens_to_ids("｜DSML｜")
	user_id = tokenizer_model.convert_tokens_to_ids("<｜User｜>")
	if dsml_id is not None and dsml_id != unk_id and user_id is not None and user_id != unk_id:
	return True

[AusarYao28]

Fixed in 2500ee9fc (same change as the thread above): both ids must differ from unk_token_id. A unit test covers this.

[github-actions]

🟠 _format_conversations_dsv4 tracks cursor = match.end(2) across matches, but does not verify that cursor == len(text). Unlike _split_official in deepseek_v4_sft_masking_test.py which explicitly asserts this condition, any trailing conversation content after the final assistant turn (e.g., trailing user messages or incomplete conversations) is silently dropped from SFT preprocessing.

Suggested change

	cursor = 0
	for match in matches:
	completion = match.group(2)
	prompt = text[cursor : match.start(2)]
	messages.extend([prompt, completion])
	is_prompt.extend([True, False])
	cursor = match.end(2)
	cursor = 0
	for match in matches:
	completion = match.group(2)
	prompt = text[cursor : match.start(2)]
	messages.extend([prompt, completion])
	is_prompt.extend([True, False])
	cursor = match.end(2)
	if cursor != len(text):
	raise ValueError(
	f"Encoded text length ({len(text)}) does not match completion boundary ({cursor}). "
	f"Trailing conversation content after the final assistant turn was dropped: {text[cursor:]!r}"
	)

[AusarYao28]

Fixed in 2500ee9fc: _format_conversations_dsv4 drops messages after the final assistant turn (matching the existing template path) and raises ValueError if any encoded text remains after the last completion.

[github-actions]

🟠 Appending 752 lines of vendored DeepSeek-V4 chat encoding and DSML streaming parser logic directly into instruction_data_processing.py quadruples the size of a previously 212-line general utility file and introduces dead/unused functions (parse_message_from_completion_text, streaming formatters, task token handlers). Furthermore, this duplicates benchmarks/api_server/encoding/encoding_dsv4.py (736 lines) almost verbatim. Extract this vendor encoder into its own dedicated module (e.g. src/maxtext/input_pipeline/deepseek_v4_encoder.py) and share it across input_pipeline and benchmarks/api_server/.

[AusarYao28]

Agreed, thanks. In 2500ee9fc the encoder is no longer appended to instruction_data_processing.py (that file is back to main). benchmarks/api_server/encoding/encoding_dsv4.py moves to src/maxtext/input_pipeline/encoding_dsv4.py, and maxtext_server.py / server_utils.py now import it from there, so there is a single copy. The move also fixes a bug: the benchmarks copy had lost the blank lines in TOOLS_TEMPLATE, and it failed the official encoder goldens for fixtures 1 and 3 (tool calling). With the blank lines restored, all 4 goldens pass. parse_message_from_completion_text is still used by maxtext_server.py, so the file has no dead functions.

[darisoy]

Thanks! Moving the existing benchmarks copy is much cleaner than adding a second one, and nice catch on the stripped TOOLS_TEMPLATE blank lines. A few questions inline. Two general ones:

1. Now that encoding_dsv4.py lives under src/maxtext (and gets imported into google3), should its header include DeepSeek's MIT notice, with the LICENSE next to it rather than under tests/utils/?
2. Without #5853, DSv4 SFT would run but silently drop EOS from the loss. Should this land after #5853?

[darisoy]

It looks like the DeepSeek-V3.2 tokenizer also has <｜User｜>, <｜Assistant｜> and ｜DSML｜, and it also ships without a chat template. Would a V3.2 tokenizer (e.g. from a local path) end up with the V4 encoding here? Since model_name is now passed through from the config, could we rely on the name check only? That would also avoid calling get_vocab() for every example.

[AusarYao28]

You're right. The V3.2 tokenizer has the same special tokens and no chat template, so the vocab check would have sent it to the V4 encoder. _is_deepseek_v4 now only checks the config model_name and the tokenizer name, so it no longer calls get_vocab() for every example. I added a test case with a V3.2-style local tokenizer.

[darisoy]

For multi-turn thinking data, encode_messages drops reasoning from assistant turns before the last user message (drop_thinking=True), but those turns are still trained after </think>. Is that intended? Would it be worth either training only the final turn or adding a short comment about this behavior?

[AusarYao28]

Not intended: those turns were trained without the reasoning that produced them. In thinking mode, _format_conversations_dsv4 now keeps turns rendered after </think> (reasoning dropped by the encoder) in the prompt and trains only turns that keep their reasoning; there is a comment explaining this. Chat-mode data is unchanged, with all turns trained. When the encoder keeps history reasoning (e.g. tools are present), those turns render with <think> and are still trained. New test: test_thinking_history_turns_stay_in_prompt.

[darisoy]

If one row in a large dataset is malformed (e.g. two assistant messages in a row), this would stop the whole run. Is strict failure intended, or would logging and skipping the example be safer?

[AusarYao28]

Strict failure is intended, to match the existing template path. It raises ValueError for a misplaced system message or when the generation prompt does not line up. Skipping would silently change the training set, and a grain map cannot drop records, so malformed rows are better filtered before SFT.

[darisoy]

This checks the default AutoTokenizer load. The grain SFT path builds HFTokenizer with add_bos_token=config.add_bos (default true) and encodes each chunk separately. Could the grain path end up with an extra BOS in each chunk? Would it be worth checking with add_bos_token=True too?

[AusarYao28]

Yes, it does. With transformers 5.x (I checked 5.17.0 and 5.20.2), add_bos_token=True makes the DeepSeek-V4 tokenizer prepend BOS to every encode() call, so each prompt and completion chunk got an extra BOS. Transformers 4.57 ignores the flag, which is why the default-load test passed.

The grain SFT pipeline now builds its tokenizer with add_bos=False, add_eos=False, as the HF SFT path and vision SFT already do. _load_ds4_tokenizer now uses the same build_tokenizer(..., add_bos=False, add_eos=False) call, and test_grain_sft_tokenizer_disables_bos_and_eos checks the arguments the pipeline passes. For other models on the grain SFT path this matches what the HF path already does.

[darisoy]

nit: could this reuse input_pipeline_utils._DS4_TURN_RE so the two can't drift apart?

[AusarYao28]

Done. The test now uses input_pipeline_utils._DS4_TURN_RE.

[vipannalla]

Inference changes look good