Skip to content

feat: add MultiturnJSON datamodule + full multi-turn SFT support - #2317

Open
ysjprojects wants to merge 2 commits into
sj/multiturnsftdatasetfrom
sj/multiturnjson-datamodule
Open

feat: add MultiturnJSON datamodule + full multi-turn SFT support#2317
ysjprojects wants to merge 2 commits into
sj/multiturnsftdatasetfrom
sj/multiturnjson-datamodule

Conversation

@ysjprojects

@ysjprojects ysjprojects commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds MultiturnJSON, the DataModule that wires MultiturnSFTDataset (#<PR for sj/multiturnsftdataset>)
up to litgpt finetune, plus a fix to a downstream assumption that would otherwise
crash any real training run using it. Together these make multi-turn JSON finetuning
work end-to-end: litgpt finetune lora --data MultiturnJSON --data.json_path ... --data.prompt_style llama3 ....

Stacked on #2315, which itself stacks on #2314

What's included

MultiturnJSON DataModule (litgpt/data/multiturn_json_data.py)

  • Loads a JSON/JSONL file or a train.json/val.json directory, same file/dir/split-loading
    behavior as JSON (litgpt/data/json_data.py) — reuses its pattern rather than duplicating it.
  • prompt_style is required, no default (unlike JSON's prompt_style="alpaca") — there's
    no single safe default chat template across checkpoints, and silently applying the wrong one is
    a real correctness risk.
  • Validates prompt_style.supports_multiturn at construction, raising immediately with a clear
    message instead of failing deep inside a DataLoader worker.
  • mask_prompt defaults to True (vs. JSON's False) — masking is the whole point of
    multi-turn SFT: only assistant turns should contribute to the loss.
  • Two input schemas, auto-detected per example (to_messages()):
    • OpenAI-style: {"messages": [{"role": "system"|"user"|"assistant", "content": str}]}
    • ShareGPT-style: {"conversations": [{"from": "system"|"human"|"gpt", "value": str}]}
      (role-mapped via sharegpt_to_messages(): humanuser, gptassistant)
    • Detection happens per record, so a single file can freely mix both formats.
    • Unrecognized shape raises ValueError naming the offending example's keys.
  • Exported from litgpt.data alongside JSON/MultiturnSFTDataset.

Fix: select_sft_generate_example crashed on multi-turn data (litgpt/utils.py)

This function powers the periodic and end-of-training "generate a sample" step in every
finetune script (litgpt/finetune/lora.py etc.) — it runs unconditionally at the end of
every training run and periodically every eval.interval steps, so it's not something a config
flag can route around. It assumed every dataset example has an "instruction" key, which isn't
true for multi-turn examples ({"messages": [...]} / {"conversations": [...]}"), so any real
MultiturnJSON-backed finetune run would KeyError mid-training.

Added _instruction_of(example, transform): applies the dataset's own transform first (reusing
the same mechanism MultiturnSFTDataset/SFTDataset already carry, rather than hardcoding
schema knowledge into utils.py), then returns example["instruction"] for single-turn examples
or the last user turn's content for multi-turn ones. This is fully backward compatible — for
every existing single-turn DataModule (with or without a transform), behavior is unchanged.

This is now ready for real multi-turn finetuning end-to-end

Verified beyond unit tests — actually ran the full path with a real (non-mocked) MultiturnJSON
instance: DataModule.setup()MultiturnSFTDatasetDataLoader batching → collate →
select_sft_generate_example(), for both OpenAI-style and ShareGPT-style records, mixed in the
same file. litgpt finetune lora --data MultiturnJSON ... should now run start to finish without
crashing at the sample-generation step.

Tests

  • tests/data/test_multiturn_json.py (new, mirrors tests/data/test_json.py's structure):
    file/JSONL loading, directory splits, path/warning validation, the supports_multiturn guard,
    ShareGPT format, mixed-format files, and direct unit tests of to_messages/sharegpt_to_messages.
  • tests/test_utils.py::test_select_sft_generate_example — updated so its mocks accurately
    reflect a real dataset's default transform=None (this test regressed against the fix until
    fixed, since MagicMock().transform auto-vends a truthy mock instead of None).
  • tests/test_utils.py::test_select_sft_generate_example_multiturn (new) — direct multi-turn
    data with no transform, MultiturnJSON-style data with a transform, and multi-user-turn
    conversations (confirms the last user turn is used).

Out of scope (follow-up)

  • litgpt/prompts.py's prompt_styles name-lookup dict doesn't register the ChatML/R1Base
    base classes by name — only ChatML's specific subclasses (qwen2.5, qwen3, smollm2,
    salamandra) and Llama3 are. So --data.prompt_style chatml currently KeyErrors;
    --data.prompt_style llama3 (or a ChatML subclass name) works today. Worth a small follow-up
    since chatml is the most likely first thing people try.
  • Inference-side multi-turn support (litgpt chat conversation history, LLM.generate(messages=...))
    is unrelated to this data-loading PR.
  • deita.py/lima.py still flatten multi-turn source data into independent single-turn pairs
    rather than routing through MultiturnSFTDataset — not touched here.

Test plan

  • pytest tests/data/test_multiturn_json.py tests/data/test_json.py tests/data/test_base.py tests/test_prompts.py tests/test_utils.py — 125 passing, 6 skipped (unrelated/environment-gated).
  • ruff check / ruff format --check clean.
  • Manually exercised the real (non-mocked) end-to-end path for both MultiturnJSON and plain
    JSON through select_sft_generate_example, confirming no crash and correct output.

@ysjprojects ysjprojects changed the title feat: add MultiturnJSON datamodule feat: add MultiturnJSON datamodule + full multi-turn sf support Aug 31, 2026
@ysjprojects ysjprojects changed the title feat: add MultiturnJSON datamodule + full multi-turn sf support feat: add MultiturnJSON datamodule + full multi-turn SFT support Aug 31, 2026
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