MLX-LM training datasets and prompt masking
Parent: MLX-LM training and adapter lifecycle · Published reference · snapshot 2026-10-23 · viability 4.65 / 5 (research)
↓ Facts as markdownall context files
Source-anchored research for MLX-LM training datasets and prompt masking. Run status: COMPLETED.
These notes link each claim to its source. Some citations identify archived research reports that are no longer hosted on this site. A published reference means the content is available; it does not certify independent review or accuracy.Read the editorial policy and follow the sources before relying on a claim.
How it works
- For a local --data directory, MLX-LM reads JSONL records from train.jsonl for training, accepts valid.jsonl as an optional validation split, and reads test.jsonl for testing. The current loader returns an empty split when a file is missing, then raises for an empty training or test split and warns but continues when validation is empty. If the --data path does not exist, the loader treats it as a Hugging Face dataset ID. (high confidence) (verified as of 2026-10-23) [source]
- Prompt masking is disabled by default. Without --mask-prompt, MLX-LM computes loss for every sample token; with it, the documented supported formats are chat and completion, and the final chat message is treated as the completion. The current loader raises ValueError when prompt masking is requested for a text dataset. (high confidence) (verified as of 2026-10-23) [source]
- Chat masking computes one prefix offset from all messages except the final message. Earlier assistant actions in a multi-turn tool trace are therefore masked when a later tool result or answer follows them; a practitioner report documented this behavior and split rows so each assistant action became a final-message target. Keep the intended assistant target last in each row. (medium confidence) (verified as of 2026-10-23) [source]
- A CompletionsDataset prompt-masking crash was reported in issue #786 and addressed by PR #967, which merged into main on 2026-03-11 and changed the mask calculation to template a prompt-only message list. The current loader uses that prompt-only list when computing the completion offset. (high confidence) (verified as of 2026-10-23) [source]
How-to and procedures
- Use consistent JSONL row shapes: text records contain {"text":"..."}; completion records contain {"prompt":"...","completion":"..."}; chat records contain {"messages":[{"role":"user","content":"..."},{"role":"assistant","content":"..."}]}. prompt_feature and completion_feature rename completion columns; chat_feature and text_feature rename the other formats. (low confidence) (verified as of 2026-10-23) [source]
Problems, failure modes and limitations
- For local data, train.jsonl is the training split, valid.jsonl supplies validation, and test.jsonl supplies --test evaluation. A missing file becomes an empty split: the final loader rejects empty training/test splits when requested and warns for missing validation. A zero-line existing JSONL file instead reaches create_dataset(data[0]) without an empty-data guard and can fail before those final split checks; ensure every supplied split file has records. (low confidence) (verified as of 2026-10-23) [source]
Comparisons and alternatives
- Prompt masking is not a universal best choice. A 2026 TACL survey reports that including prompt-token loss can improve generalization in some long-instruction, low-data settings, while a prompt-loss-weighting study says the useful weight depends on the task and dataset. Evaluate masking choices on held-out task outputs instead of assuming completion-only loss always wins. (medium confidence) [source]
- A 2024 practitioner experiment using Llama-2-7b on a RACE subset reported that completion-only training reached 80% validation accuracy versus 75% with full-sequence loss and converged faster. The author cautioned that the result may not generalize beyond that dataset and setup. (low confidence) [source]
Facts and statements
- For Hugging Face datasets, YAML fields such as prompt_feature, completion_feature, chat_feature, and text_feature map dataset columns to MLX-LM inputs. MLX-LM applies the model's chat template to chat, tools, and completion records by default; Hugging Face documents that chat templates and control tokens vary by model, so the training serialization should match the tokenizer. (high confidence) (verified as of 2026-10-23) [source]
- The current loader chooses a row format from the first record: prompt and completion keys take precedence, then the chat key (messages by default), then the text key (text by default). Configurable feature names can rename those keys. Completion rows become a user/assistant message pair, and chat rows can pass a tools field to the tokenizer's chat template, so records within a split should use a consistent shape. (high confidence) (verified as of 2026-10-23) [source]
- MLX-LM calculates the prompt boundary in token space with tokenizer.apply_chat_template: chat data uses the template for all messages except the last, and completion data uses a prompt-only message with add_generation_prompt=True. Hugging Face documents that chat templates serialize messages into model-specific token sequences and that generation prompts affect this serialization, so the tokenizer and template used for training must match the intended format. (high confidence) (verified as of 2026-10-23) [source]
- Research does not establish completion-only loss as a universal optimum. The 2024 EMNLP study reports that useful prompt-token weights varied with generation length and benchmark, while the 2025 TACL study reports that the standard zero-prompt/full-response objective was often suboptimal across its tested models, datasets, and benchmarks; target-specific evaluation is needed before choosing masking. (medium confidence) [source]
- --mask-prompt removes loss on the prompt prefix. Completion masking templates the user prompt alone with add_generation_prompt=True to locate that prefix; chat masking templates all messages except the last. Chat and completion records support masking; text records raise ValueError("Prompt masking not supported for text dataset."). (low confidence) (verified as of 2026-10-23) [source]
Children
- Dataset format validation (frontier)
- Chat-template compatibility (frontier)
- Per-turn loss masking (frontier)
Frontier under this node: Chat-template compatibility, Dataset format validation, Per-turn loss masking