feat(dataset): pluggable conversation renderers, pretokenized dataset loading, and tool-call preservation - #163
Merged
Merged
Conversation
… conversations
`_normalize_conversation` rebuilt each ShareGPT message as a fresh
`{"role", "content"}` dict plus an optional `reasoning_content`, silently
dropping every other key. Tool-calling corpora lost `tool_calls`,
`tool_call_id`, and `name`, so an assistant turn that only emitted a tool
call normalized to an empty-content message with no recoverable call, and
`tool` responses lost the id that binds them to their request.
Carry all unrecognized keys through instead, with `role`/`content` written
last so a stale duplicate in the source row cannot shadow the normalized
values. The reasoning aliases still collapse into `reasoning_content` so
callers keep a single field to read.
Existing chat templates only read `role` and `content`, so the extra keys
are inert on the template path and become available to tokenizers whose
`apply_chat_template` understands them. Also widen the `Conversation` alias
to `List[Dict[str, Any]]`, which was already inaccurate for multimodal
content lists and is now inaccurate for `tool_calls` too.
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
…late cannot express `ChatTemplate` describes a conversation as a flat header/content/end-of-turn string triple. That is enough for most chat models but cannot express nested message structures, tool-call serialization, or a checkpoint whose own `apply_chat_template` is the only correct source of truth — and it derives the loss mask by string matching on a header the template itself defines. Add a `ConversationRenderer` protocol plus a `RENDERER_REGISTRY` alongside the existing `TEMPLATE_REGISTRY`. A renderer is constructed with the target model's tokenizer and owns both halves of the job: it returns `(input_ids, loss_mask)` directly from the structured messages, so supervision is derived from the token stream it just built rather than re-discovered from text. Selection is `dataset.renderer`, which takes precedence over `dataset.chat_template`; the template path is untouched when it is unset. Renderer-specific plumbing this adds to the dataset path: - Per-row `tools` and `generation_config` are read from the source dataset and forwarded to `render`. `generation_config` is opaque to the framework — it carries whatever per-sample rendering options a renderer needs (a reasoning-effort setting, a serving-time sampling config) and renderers that do not need it ignore it. Both columns are added to `load_hf_dataset`'s keep-list, which previously dropped them for Hub datasets, and `data_id` is accepted as an alias for `id`. - `multimodal_inputs` are left unflattened for renderers, since a renderer's tokenizer consumes structured content directly instead of a placeholder string. - `CACHE_VERSION` is folded into the tokenization cache key so a renderer whose semantics change cannot silently reuse a stale cache. - `train_with_decode` is rejected: appending a generation prompt is a template-path concept, and a renderer emitting one would supervise tokens the model never produces. - A renderer owns its supervision, so a row with zero supervised tokens is dropped rather than trained on, independent of `min_loss_tokens`. Renderers also force `dynamic_loss_mask` on outside offline replay, because a renderer emits its mask *before* the inference engine expands media placeholders. Any multimodal row's pre-expansion mask is therefore stale and is dropped — both on the fresh path and when loading a cache written by an earlier run — so the training-time matcher recomputes it against the engine's real token IDs. `get_assistant_token_ids` grows a renderer branch to supply that matcher, and the controller now accepts `input_ids` without a `packed_loss_mask` when dynamic masking is on. The one behavior change on the existing template path: a multimodal row tokenized with `defer_tokenization=False` and no dynamic mask used to be given a mask rendered against unexpanded placeholders, which cannot align with what the engine feeds the model. That now fails closed with an explanation instead of training on a silently misaligned mask. No renderer is registered here; the registry ships empty. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
… re-rendering Rendering a large corpus at training time is wasted work when the tokens are already fixed, and it makes supervision non-reproducible: a per-row mask policy decided by an offline pipeline cannot be recovered from a single global `last_turn_loss_only` flag, and re-rendering under a different tokenizer or template version silently shifts the mask. Detect a dataset that already carries `input_ids` plus a loss mask and load it straight through — no tokenizer, no template, no renderer, and no tokenization cache, since there is nothing to cache. Either mask encoding is accepted: a packed segment-length string, or an explicit binary mask that gets packed on load. The loader validates rather than repairs, because a bad row here trains on the wrong tokens with no downstream symptom: - `input_ids` must be a non-empty list of ints, and `seq_len`/`loss_tokens`, if present, must agree with them. - An over-length row raises instead of truncating. Truncation would drop supervised tokens the producer counted, so it is a corpus/config mismatch to fix at the source. - A row with both mask encodings must have them agree position-by-position, not merely in supervised-token count. - Every row needs at least one supervised token, and `data_id`s must be unique. Columns the loader does not consume are carried through to `metadata` untouched, so a producer's own provenance bookkeeping survives into training without this code having to know its schema. Non-scalar values are skipped so a stray array column is not pinned in memory for the dataset's lifetime. Detection needs a typed schema, so this covers Parquet/Arrow files and Hub repos; a JSONL file reports no columns when streamed and is still treated as raw conversations. `load_hf_dataset` also no longer strips unknown columns from a Hub repo that exposes `input_ids` — that filter would have dropped the token columns and silently demoted a pretokenized repo to raw text. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c8501fef67
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+335
to
+336
| if not renderer_name and not chat_template_name: | ||
| raise ValueError("Either renderer or chat_template must be set for dataset tokenization") |
There was a problem hiding this comment.
This validation runs before checking whether the dataset already has input_ids plus a loss mask, so a Parquet/Arrow pretokenized corpus fails when the config explicitly sets dataset.chat_template=null and no renderer, even though the pretokenized branch below never renders and does not need either setting. Load the schema and short-circuit pretokenized datasets before enforcing raw-conversation renderer/template requirements.
Useful? React with 👍 / 👎.
… masks With the default `last_turn_loss_only: "auto"`, only the deferred-formatting worker recorded its per-sample decision. `_tokenize_single` — the path a renderer takes — resolved the decision, baked it into the mask, and discarded it. A multimodal row's mask is then dropped as stale, so the controller had no `has_thinking` to forward and the training fetcher fell back to the global `"auto"` string, which `compute_assistant_loss_mask` reads as truthy. Every non-thinking multi-turn multimodal sample was silently trained on its final assistant turn alone. Record the resolved decision as `has_thinking` metadata whenever the flag is `"auto"`, on both the renderer and the template branch, so it survives the mask drop and reaches the training-time matcher. `_format_single` now shares the same helper rather than repeating the logic. Also stop `resolve_loss_mask` from reading an unresolved `"auto"` sentinel as last-turn-only: it is a per-sample question that token IDs cannot answer, so supervise every assistant turn instead. And fail closed when a cache written before this change would lose a multimodal mask with no recorded decision — the source messages are gone, so it cannot be recomputed. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
…corpora The "either renderer or chat_template must be set" check ran before the pretokenized branch, so a Parquet/Arrow corpus carrying input_ids plus a loss mask failed whenever the config explicitly set `dataset.chat_template=null` with no renderer — even though that branch never renders and needs neither setting. Load the schema and short-circuit pretokenized detection before enforcing it. The two checks that remain above the load are decidable from the config alone and stay there: `renderer` with `defer_tokenization=True` is a contradiction, and a renderer name absent from the registry is wrong whatever the corpus turns out to be. `test_dataset_loading_requires_a_renderer_or_a_chat_template` now writes the raw JSONL file it points at, since the requirement it asserts can no longer be decided before the schema is known. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 601a44cffd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+365
to
+367
| if renderer_name and renderer_name not in RENDERER_REGISTRY.get_all_renderer_names(): | ||
| available = ", ".join(RENDERER_REGISTRY.get_all_renderer_names()) or "<none>" | ||
| raise ValueError(f"Unknown dataset renderer {renderer_name!r}; available: {available}") |
There was a problem hiding this comment.
When train_data_path points at a pretokenized Parquet/Arrow/Hub dataset but the config still has dataset.renderer set to a renderer that is not registered in this process, this check raises before the schema is inspected. The remaining pre-schema check is the renderer registry lookup; the pretokenized branch below never calls a renderer, so reusing a renderer-based config for offline-tokenized data is blocked even though input_ids plus a loss mask are already present. Move the unknown-renderer validation until after _is_pretokenized_dataset short-circuits, or skip it for pretokenized inputs.
Useful? React with 👍 / 👎.
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Three independent dataset-path changes, one commit each. Together they let a model whose conversation format a
ChatTemplatecannot express be trained without a bespoke fork ofdataset.py, and let a corpus that was tokenized offline be trained on as-is.1.
fix(data)— preserve extended message fields when normalizing ShareGPT conversations_normalize_conversationrebuilt each message as a fresh{"role", "content"}dict, dropping every other key. Tool-calling corpora losttool_calls,tool_call_id, andname, so an assistant turn that only emitted a tool call normalized to empty content with no recoverable call. Unrecognized keys now pass through, withrole/contentwritten last so a stale duplicate in the source row cannot shadow them.Existing chat templates only read
roleandcontent, so the extra keys are inert on the template path.2.
feat(dataset)— pluggable conversation renderersChatTemplatedescribes a conversation as a flat header/content/end-of-turn string triple, and derives the loss mask by string-matching the header it defines. That cannot express nested message structures, tool-call serialization, or a checkpoint whose ownapply_chat_templateis the only correct source of truth.This adds a
ConversationRendererprotocol and aRENDERER_REGISTRYalongsideTEMPLATE_REGISTRY. A renderer is constructed with the target model's tokenizer and returns(input_ids, loss_mask)directly, so supervision comes from the token stream it just built rather than being re-discovered from text. Selection isdataset.renderer, which takes precedence overdataset.chat_template; the template path is untouched when unset.Supporting plumbing: per-row
toolsandgeneration_configare forwarded torender(generation_configis opaque to the framework — whatever per-sample options a renderer needs);multimodal_inputsstay unflattened, since a renderer's tokenizer consumes structured content directly;CACHE_VERSIONis folded into the tokenization cache key;train_with_decodeis rejected, because appending a generation prompt is a template-path concept.Renderers also force
dynamic_loss_maskon outside offline replay, because a renderer emits its mask before the engine expands media placeholders. Stale pre-expansion masks are dropped for multimodal rows — on the fresh path and when loading an older cache — so the training-time matcher recomputes them against the engine's real token IDs.get_assistant_token_idsgrows a renderer branch to supply that matcher.No renderer is registered here; the registry ships empty. A concrete implementation follows in a separate PR.
3.
feat(dataset)— load pretokenizedinput_ids/loss-mask datasets without re-renderingRendering a large corpus at training time is wasted work when the tokens are already fixed, and it makes supervision non-reproducible: a per-row mask policy chosen by an offline pipeline cannot be recovered from one global
last_turn_loss_onlyflag.A dataset carrying
input_idsplus a loss mask (packed segment string or explicit binary) is now loaded straight through — no tokenizer, no template, no renderer, no tokenization cache. The loader validates rather than repairs: an over-length row raises instead of truncating, both mask encodings must agree position-by-position, every row needs a supervised token, anddata_ids must be unique. Columns the loader does not consume are carried through tometadatauntouched, so a producer's provenance bookkeeping survives without this code knowing its schema.Behavior changes to be aware of
defer_tokenization=Falseand no dynamic mask used to receive a mask rendered against unexpanded placeholders, which cannot align with what the engine feeds the model. That now fails closed with an explanation.load_hf_datasetno longer strips unknown columns from a Hub repo that exposesinput_ids; that filter would have dropped the token columns and silently demoted a pretokenized repo to raw text.dataset.chat_templateis nowOptional[str]so it can be unset in favor ofdataset.renderer. Its default is unchanged.Test plan
tests/test_conversation_normalization.py(new, 4 tests) — field preservation, no input mutation, reasoning-alias collapsing, identity on already-normalized input.tests/test_renderer_registry.py(new, 17 tests) — registry semantics andCACHE_VERSIONfallback;tools/generation_config/last_turn_onlyforwarding;train_with_decodeand unknown-renderer rejection; end-to-end ShareGPT load through a stub renderer; multimodal mask deferral and the fail-closed path;dynamic_loss_maskderivation including the offline-replay exemption; the renderer-supplied assistant matcher.tests/test_pretokenized_dataset.py(new, 17 tests) — round-trip load, generic metadata passthrough, both mask encodings, every validation branch, and an end-to-end load of a real Parquet file asserting no tokenizer is loaded and no cache is written.ruff check .andruff format --check .clean.upstream/maincheckout in the same sandbox: identical failure set (93 pre-existing failures both before and after, zero new). Those are CUDA/mooncake/ray.dag/flash_attnsandbox limitations.mooncakeinstalled before merging — this sandbox is CPU-only, so trainer and inference-engine tests could not execute.