Model-Optimizer 数据集混合（Dataset Blend）配置指南：面向 LLM QAT/QAD 训练的数据源混合、采样与缓存方案

人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载导读本指南聚焦 Model-Optimizer 仓库中examples/llm_qatLLM 量化感知训练 QAT 与量化感知蒸馏 QAD示例所采用的数据集混合配置方案。它通过 YAML 文件声明式地定义从哪些 HuggingFace 数据集取数、如何按权重采样、如何切分 train/eval/test、如何分词与打标签从而在有限样本预算下构造高质量的指令微调数据。读完本文你将掌握 blend YAML 的完整字段语义、Chat 模板标签掩码的三种策略、三种 split 模式、流式采样与磁盘缓存机制并能基于仓库源码理解底层实现从而为 QAT/QAD 训练自定义数据集混合配置。一、数据集混合在 LLM QAT/QAD 中的作用在量化感知训练QAT与量化感知蒸馏QAD流程中数据集是校准模型损失与保真度的基础。examples/llm_qat示例对应源码见 dataset_utils.py将数据准备与模型训练解耦为两个阶段混合Blend把多个领域的数据源代码、数学、科学问答、多语言对话、指令跟随等按比例拼成一个统一的训练集保证任务多样性Tokenize Cache预先完成分词并落盘缓存训练时直接复用避免重复下载与重复分词运行时子采样训练配置中的train_samples/eval_samples只是对已构建好数据集的运行时上限修改它们不会使缓存失效。数据集配置文件统一放在 examples/llm_qat/configs/dataset/ 目录下训练脚本通过--dataset_config指定默认值为configs/dataset/blend.yaml见 arguments.py。二、Blend YAML 结构总览Blend YAML 定义了三层信息数据集总规模blend_size、全局切分比例splits以及数据源列表sources。完整可运行示例见 blend_example.yaml基本结构如下blend_size: 100000 # total samples to download across all sources splits: # train/eval/test split ratios (must sum to 1.0) train: 0.80 eval: 0.10 test: 0.10 sources: - hf_path: nvidia/Nemotron-SWE-v1 split: r2e_gym ratio: 6000 category: codeblend_size表示从所有源下载的样本总数各源先按ratio归一化权重分配份额全部源的数据汇合后再由顶层splits全局切分为 train/eval/test。顶层字段FieldRequiredDefaultDescriptionblend_sizeno100000Total samples to download across all sourcessplitsno{train: 0.8, eval: 0.1, test: 0.1}Relative split weights; e.g.{train: 8, eval: 1, test: 1}gives 80/10/10blend_size与splits均非必填不填时使用上方默认值。注意splits并不强制各比例之和恰好为 1.0——BlendConfig.__post_init__会对splits值做归一化处理splits {k: v / total for k, v in splits.items()}因此{train: 8, eval: 1, test: 1}这类权重式写法同样有效见 dataset_utils.py。同样各 source 的ratio也是相对权重由_normalize_ratios归一化为总和 1.0见 dataset_utils.py。三、Per-Source 字段详解source列表中每个条目定义一个数据源字段如下FieldRequiredDefaultDescriptionhf_pathyes-HuggingFace dataset path or local pathratioyes-Relative weight (normalized across all sources)splityes-Split(s) to load (auto train/eval). See belowdataset_kwargsno{}Extra kwargs passed todatasets.load_dataset()(e.g.{name: 3.0.0})apply_chat_templatenotrueIf true, expects OpenAI messages formattrain_only_assistant_tokensnoautoLabel policy for chat datasets:auto,true, orfalse. See belowchat_keynomessagesKey containing conversationscategorynoLabel for logging其中hf_path、ratio、split为必填项若split缺失DatasetSourceConfig.__post_init__会抛出ValueError报错见 dataset_utils.py。这些字段与源码中的DatasetSourceConfigdataclass 一一对应dataset_utils.py字段名与 YAML 键完全一致。几个值得注意的细节hf_path可以是本地路径当hf_path指向本地存在的目录时代码会改用datasets.load_from_disk加载并对DatasetDict自动取对应 split见 dataset_utils.py 与 blend_example.yaml 中的注释示例category仅用于日志标注加载时会在日志中打印Source [i/n]: hf_path [category]方便观察各领域数据占比dataset_utils.py不影响采样逻辑apply_chat_template: false表示预训练风格纯文本此时使用make_pretrain_tokenize_fn从text/article/content字段取文本所有非 padding 的 token 都作为标签参与训练见 dataset_utils.py。四、Chat 标签掩码Chat Label Maskingapply_chat_template控制消息的格式化方式是否调用 tokenizer 的apply_chat_template而train_only_assistant_tokens控制哪些 chat-template token 会成为损失标签。三种取值语义如下auto当 tokenizer 原生支持{% generation %}掩码或属于已测试的 Qwen/Nemotron ChatML 启发式时仅用 assistant 输出作为标签否则训练所有非 padding 的 chat-template token并打印警告true强制要求 assistant-only 标签使用原生掩码或 ChatML 启发式若两者都不可用则直接抛出ValueError提示改为train_only_assistant_tokens: falsefalse对所有非 padding 的 chat-template token 都计算损失。在源码make_chat_tokenize_fn中dataset_utils.py掩码选择顺序为tokenizer 模板含{% generation %}→ 原生掩码return_assistant_tokens_maskTrue否则若 tokenizer 属于 Qwen/Nemotron 系列且识别出|im_start|/|im_end|等 ChatML 标记 → 使用启发式_chatml_assistant_mask手工构造掩码dataset_utils.py以上都不满足且取值为auto→ 退化到训练所有非 padding token并告警。已测试的模型家族由_TESTED_MODEL_FAMILIES (qwen, nemotron)限定dataset_utils.py若使用其他家族的模型且掩码不可用会在auto模式下告警、在true模式下直接失败。chat_key指定会话消息所在的字段名。默认messages对应 OpenAI 风格消息列表若数据集使用其他字段如 ShareGPT 格式的conversations通过chat_key覆盖。五、Split 模式的三种写法split指定从该数据源加载哪些 HuggingFace split。所有源的样本先汇合pooled再按顶层splits比例全局切分为 train/eval/test。三种写法等价于源码_parse_split_spec的解析结果dataset_utils.py# 单个 split split: train # 逗号分隔每个 split 等权重 split: code,math,stem # 字典按权重采样如 3:2:1 split: code: 3 math: 2 stem: 1单个 split 被解析为{train: 1.0}逗号分隔形式被解析为各 split 权重均为 1.0dict.fromkeys(parts, 1.0)字典形式按显式权重分配样本在_load_source_samples中每个 split 分配round(weight / total_weight * num_samples)个样本最后一个 split 取得剩余全部样本dataset_utils.py因此权重无需归一化。六、dataset_kwargs透传给 load_dataset 的额外参数dataset_kwargs中的键值对会原样合并进datasets.load_dataset()的调用load_kwargs.update(dataset_kwargs or {})见 dataset_utils.py典型用途包括指定 HF 数据集配置名如cnn_dailymail的3.0.0版本、信任远程代码、锁定 revision# HF config name (e.g. cnn_dailymail) dataset_kwargs: {name: 3.0.0} # 多个 kwargs dataset_kwargs: name: 3.0.0 trust_remote_code: true revision: main注意split与streamingTrue由框架内部注入无需也不应写在dataset_kwargs中。dataset_kwargs同时参与缓存键计算见下文缓存机制因此修改它会导致缓存重建。七、流式加载与洗牌Streaming and Shuffle所有 HuggingFace 数据集都以streamingTrue流式加载避免下载整个数据集dataset_utils.py。采样与洗牌由DataArguments中的shuffle与shuffle_buffer控制arguments.pyshuffle: true默认使用蓄水池采样Reservoir Sampling等价于dataset.shuffle(buffer_sizeN).take(n)。buffer 越大随机性越准确但更慢、内存占用更高shuffle: false直接取前 N 个样本等价于dataset.take(n)。速度快且结果确定便于复现。在分布式场景下每个 rank 只加载自己分片的数据流式数据通过skip(offset).take(per_rank)跳过而不存储本地数据集则利用select()随机访问dataset_utils.py。洗牌种子默认 42可通过--dataset_seed修改。八、预分词与磁盘缓存8.1 独立 CLI 预构建缓存为避免训练时重复下载与分词可先用dataset_utils.py独立预构建缓存对应 dataset_utils.py 的main入口python dataset_utils.py \ --dataset_config configs/dataset/blend.yaml \ --model_name_or_path Qwen/Qwen3-8B缓存默认写入dataset_cache_dir默认.dataset_cache/tokenized/之后使用相同数据集配置与 tokenizer 的任意训练运行都会复用该缓存。dataset_utils.py在构建时会打印各分区的样本数如Built dataset: train: 18000, eval: 1000, test: 1000便于核对比例。8.2 缓存键Cache Key机制缓存键的计算逻辑在_build_cache_path中dataset_utils.py由以下要素经 SHA-1 摘要取前 12 位生成顶层配置blend_size、splits按键排序拼接每个 sourcehf_path、ratio、split、dataset_kwargs、apply_chat_template、chat_key、train_only_assistant_tokens分词设置max_lengthtokenizer 指纹tokenizer 类名、vocab_size、eos/bos/pad/unk的 token id_tokenizer_fingerprintdataset_utils.py保证同类 tokenizer 不同词表得到不同缓存。关键结论训练配置中的train_samples与eval_samples不参与缓存键计算——它们只是对已构建数据集的运行时子采样上限。因此调整这两个值不会使缓存失效修改train_samples/eval_samples只需重新select(range(n))见 utils.py这也是 README 中训练配置与数据集配置解耦设计的核心原因见 llm_qat/README.md。8.3 分布式合并流程build_blend_datasetdataset_utils.py的完整流程为计算缓存路径并尝试从内存缓存_dataset_cache或磁盘加载每个 rank 按自己的分片流式加载各 source 的原始样本_load_all_source_samples每个 source 用make_chat_tokenize_fn或make_pretrain_tokenize_fn分词num_proc个 CPU worker 并行处理并过滤掉全部标签为IGNORE_INDEX的样本dataset_utils.py各 rank 将本地扁平数据存盘rank 0 合并所有分片、按dataset_seed确定性洗牌、再按splits比例切分并写回缓存_merge_distributed_shardsdataset_utils.py。并行参数由ParallelConfig管理dataset_utils.pynum_proc默认为 16在多卡节点上会根据LOCAL_WORLD_SIZE由 torchrun/SLURM 设置自动缩放为max(1, num_proc // local_world_size)避免 CPU 过度订阅。九、DataArguments 完整参数表处理参数cache_dir、shuffle、num_proc等在训练配置 YAML 或 CLI 中通过DataArguments设置定义见 arguments.pyArgumentDefaultDescription--dataset_configconfigs/dataset/blend.yamlPath to a dataset blend YAML config file--train_samples20000Number of training samples to use--eval_samples2000Number of evaluation samples to use--dataset_seed42Random seed for dataset shuffling--dataset_cache_dir.dataset_cache/tokenizedDirectory for caching tokenized datasets--shuffletrueWhether to shuffle dataset sources (reservoir sampling)--shuffle_buffer10000Buffer size for streaming shuffle--num_proc16Number of CPU workers for tokenizationdataset_config同时被quantize.py量化校准数据与train.pyQAT/QAD 训练数据使用见 arguments.py 中的参数分组因此一套数据集混合配置可同时服务于量化与训练两条流水线。十、仓库内置配置示例10.1 默认训练混合配置 blend.yamlconfigs/dataset/blend.yaml 是默认 SFT 混合配置使用 NVIDIA 后训练数据集总规模 20000、90/5/5 切分覆盖代码、数学、科学问答、指令跟随、多语言对话与竞赛编程等 7 个数据源各源按ratio相对权重分配6000/2500/1500/1500/5000/1500/1000默认启用 chat template 分词。10.2 全选项参考示例 blend_example.yamlblend_example.yaml 是展示所有混合选项的参考文件默认不被任何训练配置引用其中演示了单 split 自动 train/evalnvidia/Nemotron-SWE-v1、逗号分隔 splitcode,math,stem、字典权重 splitcode: 3, math: 2, stem: 1、预训练风格纯文本abisee/cnn_dailymailapply_chat_template: false、自定义chat_keyMagpie-Align/Magpie-Pro-MT-300K-v0.1chat_key: conversations以及本地数据集路径load_from_disk。10.3 快速测试配置 blend_test.yamlblend_test.yaml 是面向单元测试的微型快速配置仅cnn_dailymail一个源、blend_size: 200、关闭 shuffle 以加速构建可用于快速验证数据管线是否正确。十一、添加新数据集在 blend YAML 的sources下新增条目即可接入新数据集按数据形态分为三类示例见 blend_example.yamlsources: # Chat 数据集OpenAI messages 格式走 apply_chat_template - hf_path: your/dataset split: train ratio: 1000 # 会话字段不同的数据集自定义 chat_key全 token 参与训练 - hf_path: your/sharegpt-dataset split: train ratio: 500 chat_key: conversations train_only_assistant_tokens: false # 纯文本数据集预训练风格不做 chat template - hf_path: your/text-corpus split: train ratio: 500 apply_chat_template: false接入时注意三点其一ratio只决定相对权重实际样本数由ratio / Σratio × blend_size计算其二若希望单源内部再细分领域如同一数据集的不同 split 权重不同可对同一hf_path写多条 source 条目如blend.yaml中Nemotron-Science-v1的MCQ与RQA两条其三新增源后缓存键会随之变化首次运行时自动重建缓存。十二、常见问题与边界行为split 缺失或比例异常split是必填项缺失即报错splits与ratio的总和必须大于 0否则抛出ValueError见BlendConfig.__post_init__与_normalize_ratiosmask 不可用时的告警使用非 Qwen/Nemotron 家族模型且 chat template 不含{% generation %}时auto模式会告警并退化为全 token 训练true模式直接报错——若确实需要 assistant-only 标签应确认 tokenizer 是否支持原生 generation 掩码或 ChatML 格式空样本过滤分词后全部标签为IGNORE_INDEX的样本会被过滤如空消息日志中会报告丢弃数量dataset_utils.py缓存复用判定磁盘缓存需同时存在dataset_dict.json才视为有效_load_cached_datasetdataset_utils.py缓存目录为空或格式不完整时自动重建。综上Model-Optimizer 的 Dataset Blend 配置将多源数据混合、领域配比、Chat 标签策略与缓存管理统一收敛到一份 YAML 中配合dataset_utils.py的流式加载、分布式分片与确定性缓存设计为 QAT/QAD 训练提供了数据层面的一致性保障。赞分享人工智能大模型模型优化模型量化模型压缩【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址https://gitcode.com/GitHub_Trending/te/Model-Optimizer点击查看免费下载相关推荐MMPose 混合数据集训练实战CombinedDataset、KeypointConverter 与多源采样策略MMPose 混合数据集训练实战CombinedDataset、KeypointConverter 与多源采样策略 MMPose 通过 CombinedDat计算机视觉人工智能深度学习MMPose 混合数据集训练指南CombinedDataset 与多源采样策略深度解析MMPose 混合数据集训练指南CombinedDataset 与多源采样策略深度解析 MMPose 通过 CombinedDataset 封装器与 Keyp计算机视觉人工智能深度学习MMPose混合数据集训练指南高效融合COCO与AIC数据集MMPose混合数据集训练指南高效融合COCO与AIC数据集 混合数据集训练概述 在计算机视觉领域姿态估计模型的性能往往依赖于训练数据的多样性和规模。MMP计算机视觉人工智能深度学习创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考