Add SFT workload support to Slime recipes - #517
Conversation
|
Documentation preview: https://modal-labs-training-gym--training-gym-previews-preview-r-78da35.modal.run/517/docs |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
There was a problem hiding this comment.
🔍 Masking semantics lack coverage
Tests validate conversation shape but never exercise tokenizer masking. Multi-turn assistant spans, end tokens, and truncation behavior remain unverified.
Was this helpful? React with 👍 or 👎 to provide feedback.
| class SFTDataset(HuggingFaceDataset): | ||
| """A Hugging Face dataset formatted as complete conversations for SFT. | ||
|
|
||
| Args: | ||
| messages_column: Source column containing OpenAI-style messages. Use this | ||
| or `input_column` plus `output_column`, not both. | ||
| input_column: Source prompt column for constructing single-turn examples. | ||
| output_column: Source assistant-response column for constructing examples. | ||
| system_prompt: Optional system message for constructed examples. | ||
| prompt_template: Template applied to constructed user prompts. | ||
|
|
||
| Assistant turns are supervised targets. In-training evaluation is not | ||
| materialized because Slime's SFT rollout supports training data only. | ||
| """ | ||
|
|
||
| messages_column: str = "" | ||
| input_key: str = "messages" | ||
| label_key: str = "label" | ||
| apply_chat_template: bool = False | ||
| writes_eval_paths: bool = False | ||
| supports_sft: ClassVar[bool] = True |
There was a problem hiding this comment.
🔴 Custom SFT datasets reuse RL data
A custom dataset accepted for SFT keeps its RL cache path because only SFTDataset receives the SFT suffix. Existing RL data then bypasses preparation, so SFT can train on the wrong format.
Learn more
Custom DatasetConfig implementations are explicitly accepted for SFT, but BaseTrainRecipe._resolve_data_paths only separates caches when the object is an SFTDataset. The resolver does not know the recipe workload, so a custom dataset used for both RL and SFT resolves to one file. The Slime launcher skips dataset.prepare when that file exists. Make the workload part of Slime data-path resolution for every SFT run, while preserving Miles and existing RL paths. Update both launcher preparation paths and CLI serialization consistently, then cover a custom dataset used by both workloads.
Was this helpful? React with 👍 or 👎 to provide feedback.
| def _from_pair(row: dict) -> dict: | ||
| for column in (in_col, out_col): | ||
| if not isinstance(row[column], str) or not row[column].strip(): | ||
| raise TrainingGymConfigError( | ||
| f"SFT column {column!r} requires non-empty text" | ||
| ) | ||
| messages = [] | ||
| if sys_prompt: | ||
| messages.append({"role": "system", "content": sys_prompt}) | ||
| messages.extend( | ||
| [ | ||
| {"role": "user", "content": template.format(input=row[in_col])}, | ||
| {"role": "assistant", "content": row[out_col]}, |
There was a problem hiding this comment.
🟡 Rendered prompts bypass validation
With a blank prompt_template, pair inputs pass validation before rendering. The resulting conversation has an empty user turn, so training learns targets without their intended context.
Was this helpful? React with 👍 or 👎 to provide feedback.
| SFT starts only the Megatron trainer: no SGLang server or inference GPUs are | ||
| allocated. The model recipe's topology and memory settings are preserved, so | ||
| the same switch works for every existing model recipe. `num_rollout` is the | ||
| number of optimizer steps; choose an SFT learning rate for the model's scale. |
There was a problem hiding this comment.
🔍 Model-wide support remains unverified
Documentation promises SFT for every existing model recipe. Live evidence covers Qwen only, while specialized audio, vision, hybrid, and MoE presets remain untested.
Was this helpful? React with 👍 or 👎 to provide feedback.
| if self.training_type == "sft" and dataset.apply_chat_template: | ||
| raise TrainingGymConfigError( | ||
| f"{type(ds).__name__}: input_key and label_key are both " | ||
| f"{ds.input_key!r}; they must name distinct columns." | ||
| "SFT requires apply_chat_template=False so Slime can build its " | ||
| "loss mask from raw messages" | ||
| ) |
There was a problem hiding this comment.
🔴 SFT datasets run as RL
Pairing SFTDataset with the default RL recipe passes validate_dataset. Its complete conversations and blank labels then enter reward training instead of supervised training.
Prompt for agents
SlimeRecipe.validate_dataset enforces apply_chat_template=False when training_type is sft, but it does not enforce the inverse relationship. SFTDataset always writes complete conversations with an empty label, so using it with the default training_type='rl' sends SFT-shaped data into the reward-learning path. Reject SFTDataset unless training_type is 'sft'. Also consider rejecting SFTDataset in Miles, where SFT is unsupported, and add mismatch tests for both directions.
Was this helpful? React with 👍 or 👎 to provide feedback.
| for message in messages: | ||
| if not isinstance(message, dict): | ||
| raise TrainingGymConfigError("Every SFT message must be a dictionary") | ||
| if not isinstance(message.get("role"), str) or not message["role"].strip(): | ||
| raise TrainingGymConfigError( | ||
| "Every SFT message requires a non-empty role string" | ||
| ) | ||
| if not any(message["role"] == "assistant" for message in messages): | ||
| raise TrainingGymConfigError( | ||
| "Each SFT conversation requires at least one assistant turn" | ||
| ) |
There was a problem hiding this comment.
🟡 Empty assistant targets pass validation
A conversation containing {"role": "assistant"} passes _validate_conversation. It has no supervised content, so tokenization can fail or produce an empty loss mask.
Learn more
SFTDataset._validate_conversation validates role presence but not whether assistant turns contain a usable target. Preserve supported text conversations and any intentionally supported structured message forms, but reject conversations where every assistant turn lacks supervised content. Validate other message content needed by the selected chat templates and add tests for missing, null, empty, and whitespace-only assistant content without tool calls.
Was this helpful? React with 👍 or 👎 to provide feedback.
| rollout_batch_size: int = 8 | ||
|
|
||
| # ── App identity ───────────────────────────────────────────────────────── | ||
| training_type: Literal["rl", "sft"] = "rl" |
There was a problem hiding this comment.
🔍 SFT option name is inconsistent
The PR description advertises workload_type="sft", but the public recipe and tutorial use training_type="sft". Reconcile the release-facing contract.
Was this helpful? React with 👍 or 👎 to provide feedback.
| if self.training_type == "rl": | ||
| return self |
There was a problem hiding this comment.
🟡 RL variants retain SFT behavior
Changing an SFT recipe to training_type="rl" leaves its normalized fields intact. TrainingGroup can label train-only SFT execution as RL.
Prompt for agents
Avoid destructive, irreversible normalization in SlimeRecipe's validator. Preserve the user-supplied or model-preset values separately, or derive effective SFT fields only during CLI and allocation resolution, so rebuilding a recipe after changing training_type restores RL behavior. Add a TrainingGroup regression test whose base recipe is SFT and whose grid selects training_type=rl.
Was this helpful? React with 👍 or 👎 to provide feedback.
| for message in messages: | ||
| if not isinstance(message, dict): | ||
| raise TrainingGymConfigError("Every SFT message must be a dictionary") | ||
| if not isinstance(message.get("role"), str) or not message["role"].strip(): | ||
| raise TrainingGymConfigError( | ||
| "Every SFT message requires a non-empty role string" | ||
| ) | ||
| if not any(message["role"] == "assistant" for message in messages): | ||
| raise TrainingGymConfigError( | ||
| "Each SFT conversation requires at least one assistant turn" | ||
| ) |
There was a problem hiding this comment.
🔍 Conversation contract exceeds stated scope
_validate_conversation accepts arbitrary roles and tool-call messages. Confirm Slime supports these shapes, or restrict validation to text SFT.
Was this helpful? React with 👍 or 👎 to provide feedback.
61b9f3d to
38bec4b
Compare
38bec4b to
863acce
Compare
andrewhinh
left a comment
andrewhinh
left a comment
There was a problem hiding this comment.
thanks for taking this on! it's badly needed.
plz address my and Devin's comments and lmk if you have any questions.
| ds.to_json(path, orient="records", lines=True) | ||
|
|
||
|
|
||
| class SFTDataset(HuggingFaceDataset): |
There was a problem hiding this comment.
wait are we only supporting SFT for HuggingFace datasets?
also, im not sure i understand what the difference between this and HuggingFaceDataset is.
| - [HuggingFaceDataset](https://gym.modal.dev/reference/huggingfacedataset) | ||
| - [MultimodalDataset](https://gym.modal.dev/reference/multimodaldataset) | ||
| - [OnlineRollout](https://gym.modal.dev/reference/onlinerollout) | ||
| - [SFTDataset](https://gym.modal.dev/reference/sftdataset) |
There was a problem hiding this comment.
hmm, i wonder if creating all these use-case specific dataset classes (i.e., this, OnlineRollout, MultimodalDataset) is good practice. as a user, i'd expect these to either be automatically inferred or exposed as flags.
btw this pertains to a latent issue rather than something this PR introduces.
| SlimeRecipe._validate_custom_model_architecture(model) | ||
| SlimeRecipe._validate_datasets(dataset, eval_dataset) | ||
|
|
||
| if slime.training_type == "sft" and dataset.apply_chat_template(): |
There was a problem hiding this comment.
to my point above, this seems like a sharp edge for a user to have to experience.
| Megatron actor nodes. | ||
| actor_num_gpus_per_node: | ||
| GPUs per actor node. | ||
| num_gpus_per_node: |
There was a problem hiding this comment.
hmm, im not a big fan of having separate parameters like this. even more confusingly, although this is for slime, there is already a miles flag with this name. i wonder if we should instead make a SFTRecipe.
| @@ -0,0 +1,37 @@ | |||
| # --- | |||
There was a problem hiding this comment.
i think this is better suited as a new section in the Dataset guide.
2 participants
Adds text SFT through the existing Slime model recipes: set
workload_type="sft"and useSFTDatasetwith conversations or prompt/answer columns. RL remains the default. SFT reserves only actor GPUs and disables inference, critic, KL, and reward-based training work.Includes data validation, workload identity in run summaries, tests, a one-GPU tutorial, and the three-PR release plan.
colocate=Falseselects Slime's train-only execution path; it does not reserve an inference node.Stack: #517 → #518 → #519. Review this first against
main; #518 adds reliable SFT observability and #519 adds UI/UX. Each PR is based on its immediate parent. Rebased onto main13f9e3efbefore implementation. Review model/dataset semantics before generated documentation.Validation: focused SFT, recipe identity, dashboard schema, and CLI tests pass; lint/compile and documentation generation pass. The combined stack passes 980 Python tests (one skipped) and six frontend tests/build. A Qwen3-0.6B one-H100 proof and ten-step smoke completed with actor=1/inference=0/critic=0, saved checkpoints, and all ten optimizer metrics. Full integration evidence and remaining release checks are in #519 and
plan.md.Text SFT is the initial supported scope; multimodal/tool-call SFT and Miles SFT are excluded. Broad GPU model-validation CI was deliberately cancelled before execution to honor the one-H100 budget; snapshot CI awaits an authorized environment reviewer. Do not treat the cancelled matrix as a passing all-model validation.