axolotl 预训练与持续预训练实战指南：completion 非流式与 pretraining_dataset 流式方案全解析

axolotl 预训练与持续预训练实战指南：completion 非流式与 pretraining_dataset 流式方案全解析

【免费下载链接】axolotlGo ahead and axolotl questions项目地址: https://gitcode.com/GitHub_Trending/ax/axolotl

导读

本文聚焦 axolotl 中“基于纯文本、无输入掩码”的预训练 / 持续预训练（Pretraining / Continual Pretraining）两条技术路线：适用于内存可容纳数据集的非流式type: completion方案，以及面向海量语料的流式pretraining_dataset方案。通过本文，你将掌握两种方案的适用场景选型、完整 YAML 配置、流式模式下的必需参数（如max_steps、streaming_multipack_buffer_size、pretrain_multipack_attn）与底层数据编码原理，并能直接对照仓库中的真实示例配置开启一次预训练任务。

适用场景：什么时候需要预训练 / 持续预训练

axolotl 的预训练路径针对“整段文本都是训练信号”的裸文本数据，即不使用输入掩码（no input masking）——模型对输入序列中的每一个 token 都进行预测，这与 SFT 中掩码掉 prompt 部分、只对回答计算损失的做法有本质区别。

根据文档 docs/agents/pretraining.md，典型的应用场景包括：

• 领域语料持续预训练（Continual Pretraining）：在某个垂直领域（法律、医疗、代码、金融等）的大规模语料上继续训练已有基座模型，注入领域知识；
• 语言 / 领域适配：在下游微调之前，先把基座模型适配到一种新语言或新领域；
• 预训练式数据：整个文本内容本身就是唯一的训练信号，不需要构造指令-回答对。

上述场景中，文档明确指出存在两条技术路线，取舍的核心在于数据集规模是否能放进内存。

方案选型：非流式 vs 流式

两条路线在代码层面也有清晰的分流：在 src/axolotl/utils/data/sft.py#L63-L64 中，数据集准备函数会首先判断cfg.streaming or cfg.pretraining_dataset，命中则走_prepare_streaming_dataset流式路径，否则走_prepare_standard_dataset非流式路径。

补充说明：axolotl 还提供了axolotl preprocess命令，将数据先离线分词、缓存后再训练。对于较小的数据集或需要反复使用的数据集，非流式 + 预处理的组合通常效率更高；而流式则适合“数据大得没必要/没可能先全量预处理”的场景（详见 docs/streaming.qmd）。

非流式方案：type: completion

最小配置

对于能放进内存的中小型语料，直接在datasets:下声明数据集并指定type: completion即可：

datasets: - path: my_corpus type: completion # field: text # 文本列名（默认: "text"）

底层实现：无掩码的 token 化策略

completion类型的核心实现在 src/axolotl/prompt_strategies/completion.py。其要点如下：

• CompletionPromptTokenizingStrategy.parse_instruction_fields将prompt[field]作为 instruction，而 input 与 response 均为空字符串，因此整个文本都被当作训练目标，不做任何掩码；
• 默认读取的列名为text（_field: str = "text"），可通过数据集配置中的field:覆盖（见load()中strat.field = ds_cfg["field"]）；
• 在tokenize_prompt中，对每个样本按self.sequence_len步长切分：for i in range(0, len(val), self.sequence_len)，这正是“长文本自动切分为多段”的实现机制；
• load()中传入的max_length=cfg.sequence_len * 64，即单个样本的最大 token 数上限为sequence_len的 64 倍，超出部分会通过_tokenize的截断逻辑处理。

也就是说，非流式路径在训练开始前就对全量数据完成了分词与切分，训练时直接读取 token id，因此可以在 CPU 上完成全部预处理工作，再把结果搬到 GPU。

数据格式

completion类型要求每条样本是“整段文档文本”的 JSON 对象：

{"text": "The complete document text goes here."}

流式方案：pretraining_dataset

为什么流式必须显式设置max_steps

当语料规模大到无法放入内存时，使用pretraining_dataset让数据按需流入训练循环。由于流式数据集没有确定的大小，axolotl 无法自行推断总步数，因此max_steps是必填项：

pretraining_dataset: - path: HuggingFaceFW/fineweb-edu type: pretrain text_column: text split: train max_steps: 1000 # 必填——axolotl 无法推断数据集大小 streaming_multipack_buffer_size: 10000 # 样本打包（sample packing）缓冲区 pretrain_multipack_attn: true # 阻止打包样本之间的交叉注意力

max_steps与 token 消耗量的换算关系（文档原文）：一个 step =sequence_len × micro_batch_size × gradient_accumulation_steps × num_gpus个 token。据此可以估算总训练 token 数并校准学习率与调度器。

流式模式的底层原理

流式数据路径实现在 src/axolotl/utils/data/streaming.py：

• encode_streaming()（src/axolotl/utils/data/streaming.py#L20）：对每个样本先按max_tokens - 2截断分词，随后在序列尾部追加eos_token_id与pad_token_id（labels 中对应追加eos_token_id与-100，attention_mask 对应追加1与0），再按固定长度把多个样本**拼接（concatenate）**为定长序列，不足max_tokens的部分以pad_token_id（labels 为-100）填充。这正是文档表格中“拼接为固定长度序列”的实现细节；
• wrap_streaming_dataset()（src/axolotl/utils/data/streaming.py#L179）：当开启sample_packing时，构建PretrainingBatchSamplerDataCollatorForSeq2Seq并使用encode_packed_streaming进行多样本打包；其中multipack_attn的取值逻辑是——SFT 数据强制为True，预训练数据则取决于pretrain_multipack_attn配置（src/axolotl/utils/data/streaming.py#L186-L190）。

配置字段的权威定义位于 src/axolotl/utils/schemas/config.py：

• pretraining_dataset（src/axolotl/utils/schemas/config.py#L410）：设置为 HF 数据集以启用流式预训练；
• pretrain_multipack_attn（src/axolotl/utils/schemas/config.py#L761）：默认True，用于“阻止预训练打包序列之间的交叉注意力”；
• streaming_multipack_buffer_size（src/axolotl/utils/schemas/config.py#L778）：默认10_000，控制流式多包（multipack）样本的缓冲区大小，缓冲区越大打包效率越高但内存占用也越大；
• 旧字段pretrain_multipack_buffer_size已在 v0.13.0 废弃，v0.14.0 将移除，请统一使用streaming_multipack_buffer_size。

从流式数据到训练：完整示例配置

仓库提供了可直接运行的完整流式预训练配置 examples/streaming/pretrain.yaml，涵盖数据集、流式缓冲、序列打包、批量大小、优化器、精度、日志与检查点等全部环节：

base_model: HuggingFaceTB/SmolLM2-135M # Streaming pretraining configuration pretraining_dataset: - path: HuggingFaceFW/fineweb-edu name: sample-10BT type: pretrain text_column: text split: train # Streaming-specific settings streaming_multipack_buffer_size: 10000 shuffle_merged_datasets: true # Training configuration max_steps: 1000 output_dir: ./outputs/smollm2-135m-pretrain-streaming # Sequence and packing settings sequence_len: 1024 sample_packing: true pretrain_multipack_attn: true # Prevent cross-attention between packed sequences attn_implementation: flash_attention_2 # Batch size settings gradient_accumulation_steps: 8 micro_batch_size: 1 # Optimizer and scheduler optimizer: adamw_torch lr_scheduler: cosine learning_rate: 5e-4 warmup_ratio: 0.1 weight_decay: 0.01 # Precision and performance bf16: auto tf32: true # Logging and checkpointing logging_steps: 10 save_strategy: steps save_steps: 250 save_total_limit: 3 # Weights & Biases (optional) wandb_project: wandb_entity: wandb_watch: wandb_name: wandb_log_model: # Special tokens special_tokens: pad_token: "<|endoftext|>" # save_first_step: true # uncomment this to validate checkpoint saving works with your config

对其中关键项做要点解读：

• pretraining_dataset各子字段：path为 HF 数据集标识；name为子集（如 fineweb-edu 的sample-10BT）；type: pretrain标记预训练流式加载；text_column: text指定文本列；split: train指定划分；
• shuffle_merged_datasets: true：使用缓冲区对流式数据进行洗牌（会额外占用 shuffle 缓冲内存），关闭则按原始顺序读取；
• sample_packing: true+pretrain_multipack_attn: true：把多条样本打包进定长序列以最大化 GPU 利用率，同时用注意力隔离防止打包样本间串扰；
• attn_implementation: flash_attention_2：打包模式下必须启用（下文“关键设置”详述）；
• save_first_step：注释中的提示说明，可用它来快速验证检查点保存链路是否正常。

流式 SFT（对照参考）

流式模式不仅适用于预训练，也可以用于 SFT（配置streaming: true+datasets:），对照示例见 examples/streaming/sft.yaml。SFT 流式打包时注意力隔离是强制开启的（不依赖pretrain_multipack_attn），这一点与预训练路径不同（见 src/axolotl/utils/data/streaming.py#L185-L190）。

关键设置速查

文档 docs/agents/pretraining.md 给出的核心设置如下：

• sample_packing: true+pad_to_sequence_len: true：把多个文档打包进定长序列，减少 padding 浪费、提升 GPU 利用率；
• flash_attention: true（等价地，attn_implementation: flash_attention_2）：样本打包的前提条件，必须启用 flash attention 才能配合打包；
• 不挂 adapter（No adapter）：预训练通常是全参微调（full fine-tune），一般不需要 LoRA 等适配器；
• train_on_inputs: true：completion 类型的默认行为——对所有 token 计算损失。补充说明：在 src/axolotl/utils/schemas/training.py#L103-L108 中train_on_inputs的全局默认值为False（用于 SFT 时掩码掉 prompt），但预训练路径下因为整段文本都是输入，语义上等价于始终对全部 token 训练，无需掩码。

流式模式的注意事项

综合 docs/streaming.qmd 与源码实现，流式预训练还需关注以下几点：

内存与缓冲

• 通过streaming_multipack_buffer_size调节内存占用：缓冲区越大，打包效率越高，但内存开销越大；
• sample_packing需要缓冲多个样本才能打包；
• 开启shuffle_merged_datasets需要额外的洗牌缓冲内存；
• 计算好max_steps与 token 预算：一个 step 消耗sequence_len × micro_batch_size × gradient_accumulation_steps × num_gpus个 token。

性能

• 流式按需实时分词，相比预分词数据集有更高的训练前延迟；
• 从远端或本地流式读取时，网络带宽 / 磁盘读取速度会成为关键瓶颈；
• 对较小或高频复用的数据集，建议改用axolotl preprocess离线预处理。

评估数据集

评估数据集不会走流式——为了保证评估指标的一致性，即使训练使用流式，评估集仍按常规方式完整加载（这也是配置评估集时需要保证其能放入内存的原因）。

文件地图：快速定位相关实现

以下是预训练 / 持续预训练相关代码与示例的速查清单（来源：文档 docs/agents/pretraining.md 的 File Map，路径已转换为仓库根目录相对路径）：

src/axolotl/ prompt_strategies/completion.py # 非流式: completion prompt strategy（无掩码） utils/data/sft.py # 非流式: 数据集加载与处理（含流式/非流式分流） utils/data/streaming.py # 流式: encode_streaming(), wrap_streaming_dataset() utils/schemas/config.py # 配置字段: pretraining_dataset, pretrain_multipack_attn 等 examples/streaming/pretrain.yaml # 完整的流式预训练示例配置 examples/streaming/sft.yaml # 完整的流式 SFT 示例配置

小结

axolotl 为预训练 / 持续预训练提供了清晰的两级方案：数据能进内存就用datasets: + type: completion（CPU 预分词、长文本自动切分、无掩码全 token 训练）；数据量巨大则用pretraining_dataset（按需实时分词、定长拼接、可配合 sample packing 与注意力隔离）。流式路径务必显式声明max_steps，并按“一个 step 等于sequence_len × micro_batch_size × gradient_accumulation_steps × num_gpus个 token”的公式规划训练预算。实际落地时，直接以 examples/streaming/pretrain.yaml 为模板，替换数据集路径与训练超参即可快速启动。

关于流式数据集的完整说明（包括基本流式配置、内存/性能注意事项、评估集行为），可进一步阅读 docs/streaming.qmd 与 docs/agents/pretraining.md。

【免费下载链接】axolotlGo ahead and axolotl questions项目地址: https://gitcode.com/GitHub_Trending/ax/axolotl