diff --git a/CMakeLists.txt b/CMakeLists.txt index 035492ed7..d58a7256b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -72,6 +72,9 @@ endif() # Framework core sources (*.cc), excluding cpu kernels (they are built separately) file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc) list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*") +# CLI flags are opt-in: linking the core must not register names such as +# --device or --batch_size that applications may define themselves. +list(FILTER SRC EXCLUDE REGEX ".*/training/arguments\\.cc$") # Exclude backend-specific runtime/ccl translation units when the corresponding # backend is disabled. This keeps each build self-contained and avoids pulling @@ -228,6 +231,14 @@ link_infini_train_exe(infini_train_executable) add_library(InfiniTrain::infini_train_executable ALIAS infini_train_executable) unset(_infini_train_executable_archives) +# ------------------------------------------------------------------------------ +# Training CLI (opt-in for applications) +# ------------------------------------------------------------------------------ + +add_library(infini_train_training_cli STATIC infini_train/src/training/arguments.cc) +add_library(InfiniTrain::training_cli ALIAS infini_train_training_cli) +target_link_libraries(infini_train_training_cli PUBLIC InfiniTrain::infini_train) + # ------------------------------------------------------------------------------ # Examples # ------------------------------------------------------------------------------ @@ -240,14 +251,42 @@ add_executable(mnist ) link_infini_train_exe(mnist) -add_executable(gpt2 - example/gpt2/main.cc +# Only GPT task/data/weight-format code belongs to the example executable. +# The training runtime comes from infini_train; its CLI adapter is opt-in. +set(INFINITRAIN_GPT_EXAMPLE_SOURCES + example/pretrain_gpt.cc example/common/tiny_shakespeare_dataset.cc example/common/utils.cc example/common/tokenizer.cc + # Provider TUs contain startup registration and must be linked into the + # executable directly (not discarded as unreferenced archive members). + example/gpt2/model_provider.cc + example/llama3/model_provider.cc + example/qwen3/model_provider.cc example/gpt2/checkpoint_loader.cc + example/llama3/checkpoint_loader.cc + example/qwen3/checkpoint_loader.cc ) -link_infini_train_exe(gpt2) + +foreach(model IN ITEMS pretrain_gpt gpt2 llama3 qwen3) + add_executable(${model} ${INFINITRAIN_GPT_EXAMPLE_SOURCES}) + if(NOT model STREQUAL "pretrain_gpt") + target_compile_definitions(${model} PRIVATE + INFINITRAIN_GPT_DEFAULT_FLAGFILE="configs/${model}/train.flags") + # Keep runtime presets beside each executable, including multi-config builds. + # Copy on every requested build so editing a preset needs no C++ recompilation. + add_custom_target(${model}_training_config + COMMAND ${CMAKE_COMMAND} -E make_directory "$/configs/${model}" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${CMAKE_CURRENT_SOURCE_DIR}/example/${model}/train.flags" + "$/configs/${model}/train.flags" + VERBATIM + ) + add_dependencies(${model} ${model}_training_config) + endif() + target_link_libraries(${model} PRIVATE InfiniTrain::training_cli) + link_infini_train_exe(${model}) +endforeach() add_executable(mixtral example/mixtral/main.cc @@ -257,23 +296,6 @@ add_executable(mixtral ) link_infini_train_exe(mixtral) -add_executable(llama3 - example/llama3/main.cc - example/common/tiny_shakespeare_dataset.cc - example/common/utils.cc - example/common/tokenizer.cc - example/llama3/checkpoint_loader.cc -) -link_infini_train_exe(llama3) - -add_executable(qwen3 - example/qwen3/main.cc - example/common/tiny_shakespeare_dataset.cc - example/common/utils.cc - example/qwen3/checkpoint_loader.cc - example/common/tokenizer.cc -) -link_infini_train_exe(qwen3) endif() # Tools diff --git a/README.md b/README.md index 8f856700a..912771c53 100644 --- a/README.md +++ b/README.md @@ -91,8 +91,10 @@ Build Options: ## 🏋️ Training -Each model in the `example/` directory is compiled into an independent executable. -For example, the `llama3` example produces a binary named `llama3`. +GPT-2, LLaMA3, and Qwen3 use the shared `pretrain_gpt` training program. +Select a model and its training preset with `--flagfile=example//train.flags`. +The `gpt2`, `llama3`, and `qwen3` executables also run this program, loading their +respective presets automatically. To view available runtime options: @@ -135,7 +137,17 @@ HF_TOKEN=hf_xxx \ #### Model Examples -The generated files can be passed directly to the corresponding executables: +Each model's `train.flags` file defines its model selection and training settings, +including the optimizer, learning rate, batch size, and checkpoint options. Put +`--flagfile` before any command-line overrides; later values take precedence. +For example, append `--learning_rate=0.00002` to override the preset learning rate. + +The named executables (`gpt2`, `llama3`, and `qwen3`) load +`configs//train.flags` relative to the executable. Building a named target +copies its preset there. Include the adjacent `configs/` directory when moving +these executables, or pass an explicit `--flagfile` path. + +The following commands use the prepared datasets and weights: ##### MNIST diff --git a/docs/unified_pretraining.md b/docs/unified_pretraining.md new file mode 100644 index 000000000..e428446e0 --- /dev/null +++ b/docs/unified_pretraining.md @@ -0,0 +1,382 @@ +# 统一 GPT 预训练入口设计 + +统一训练入口将 GPT-2、LLaMA3、Qwen3 的训练过程组织为一套框架生命周期。任务入口负责提供模型、数据和计算逻辑，`infini_train::training::Pretrain()` 负责执行这些逻辑，并管理并行包装、优化器、学习率调度和训练状态。 + +三个模型都基于 `nn::TransformerModel`，架构由 `TransformerConfig` 描述。这一共同基础使模型特例可以通过 provider 接入统一训练流程。开发者增加模型时维护模型目录，增加训练能力时维护框架训练层。 + +本文以当前实现为准。结构调整的对照基线是 `master` 提交 `a70c996`（`test: align Qwen3 validation case configurations`）。 + +## 1. 架构与职责 + +设计分为任务、训练接口和执行机制三个层次： + +```mermaid +flowchart TD + P["各模型 train.flags"] -->|"--flagfile"| A["arguments.cc
解析配置，生成 TrainingOptions"] + A --> E["example/pretrain_gpt.cc
组装 GPT 任务"] + M["各模型 model_provider.cc
配置与权重加载"] --> R["ModelProviderRegistry"] + R -->|"ModelProvider"| E + E -->|"options / providers / callbacks"| T["training::Pretrain"] + T --> D["DistributedDataLoader"] + T --> S["普通执行 / PipelineSchedule"] + T --> O["DDP / Optimizer / LRScheduler"] + T --> C["Checkpoint"] +``` + +任务层决定训练什么模型、使用什么数据、如何计算输出和损失。框架层把这些任务定义组合成可执行的训练过程。并行组件、优化器和 checkpoint 模块承担各自的执行职责。 + +| 组件 | 位置 | 职责 | +| --- | --- | --- | +| GPT 任务入口 | [example/pretrain_gpt.cc](../example/pretrain_gpt.cc) | 解析任务参数，选择模型 provider，构造数据、forward/loss 和采样回调，调用 `Pretrain()`。 | +| 训练接口 | [training.h](../infini_train/include/training/training.h) | 定义 `TrainingOptions`、provider、训练步骤类型与 `Pretrain()` 声明。 | +| 训练执行 | [training.cc](../infini_train/src/training/training.cc) | 初始化 rank，创建训练组件，调度训练并管理恢复与保存。 | +| 命令行适配 | [arguments.cc](../infini_train/src/training/arguments.cc)、[arguments.h](../infini_train/include/training/arguments.h) | 定义公共参数，校验参数文件，定位默认配置，将参数转换为 `TrainingOptions`。 | +| 模型注册表 | [model_provider_registry.h](../infini_train/include/training/model_provider_registry.h)、[实现](../infini_train/src/training/model_provider_registry.cc) | 管理模型名称到构建函数的映射，生成名称列表和诊断信息。 | +| 模型适配 | `example//` | 提供 `config.h`、`model_provider.cc`、`checkpoint_loader.cc` 和 `train.flags`。 | +| GPT 数据适配 | `example/common/` | 提供 token 数据读取、tokenizer 和 LLMC 文件格式辅助功能。 | + +相对基线，原来分布在三个 `main.cc` 中的职责按下面的方式归集： + +| 基线中的职责 | 当前归属 | +| --- | --- | +| 三个模型各自组织完整训练程序 | 一份 `pretrain_gpt.cc` 组装任务，一份 `Pretrain()` 执行训练。 | +| 入口内部选择模型配置与 loader | 各模型的 provider 向注册表登记构建逻辑。 | +| 入口中的模型训练默认值 | 各模型目录的 `train.flags`。 | +| 通用参数、优化器、调度器、数据位置和 checkpoint 管理 | 框架 `training/` 层。 | +| 普通路径与 PP 路径中的任务计算 | 统一的 forward/loss 回调协议。 | + +## 2. 公共训练接口 + +公开接口位于 `infini_train::training`： + +```cpp +void Pretrain(const TrainingOptions &options, + const DatasetProvider &dataset_provider, + const ModelProvider &model_provider, + const ForwardStep &forward_step, + const AfterStep &after_step = {}); +``` + +`TrainingOptions` 描述一次训练运行的公共配置。Provider 和回调描述任务的可替换部分，框架负责它们的调用时机。当前接口以 `TransformerModel`、`TransformerConfig` 和定长语言模型 batch 为模型与数据约定。 + +| 接口 | 输入与返回值 | 执行契约 | +| --- | --- | --- | +| `ModelProvider` | 返回 `shared_ptr` | 每个 rank 在并行上下文就绪后调用一次，创建该 rank 对应的 CPU 模型。设备迁移、LoRA 注入和并行包装由训练流程完成。 | +| `DatasetProvider` | 接收目标全局训练样本数，返回 `DatasetSplits{train, valid}` | 每个 rank 创建自己的数据集对象，由框架建立分布式加载器并维护迭代位置。 | +| `ForwardStep` | 接收模型或本地 chunk、输入、标签和实际模型配置，返回 `ForwardStepResult` | 定义一个 microbatch 的前向计算和损失计算方式。 | +| `AfterStep` | 接收包装后的模型、设备和已完成步数 | 在优化器与调度器更新后，由最后一个全局 rank 调用；适合日志、结果展示和单 rank 采样。 | + +配置在启动工作线程前完成解析，训练期间按只读方式使用。模型和数据 provider 为各 rank 创建独立对象；回调捕获的共享状态应满足相应的并发访问约定。 + +### 2.1 训练步骤类型 + +普通训练和流水线执行共用 `training.h` 中的步骤协议： + +```cpp +using TensorList = std::vector>; +using LossFunction = std::function(const TensorList &)>; + +struct ForwardStepResult { + TensorList output; + LossFunction loss_func; +}; +``` + +`output` 保存模型输出或中间激活，`loss_func` 定义如何从输出计算当前 microbatch 的平均损失。闭包可以捕获该 microbatch 的标签等数据，使输出与损失计算保持对应关系。 + +接口中的两个 forward 类型承担不同角色： + +| 类型 | 使用位置 | 配置信息 | +| --- | --- | --- | +| `ForwardStep` | 任务入口传给 `Pretrain()` | 显式接收实际模型的 `TransformerConfig`。 | +| `ForwardStepFunction` | 普通训练循环与 PP 调度器 | 训练流程已绑定当前 rank 的模型配置，执行器直接传入模型、输入和标签。 | + +这些类型统一归属 `infini_train::training`。`training.h` 使用标准库类型和前向声明表达接口；具体类的定义由实现文件显式包含。PP 通过这个轻量公共头文件使用训练步骤协议。 + +## 3. 模型接入与注册 + +模型注册表把命令行名称映射为构建函数。入口的模型选择过程是： + +```cpp +const auto &models = training::ModelProviderRegistry::Instance(); +const auto model_provider = models.Resolve(FLAGS_model, FLAGS_llmc_filepath); + +training::Pretrain(options, dataset_provider, model_provider, + ForwardStepGPT, after_step); +``` + +`Resolve()` 查找名称、复制 factory 并绑定权重路径，返回一个零参数的 `ModelProvider`。模型实例在训练流程调用该 provider 时创建。这样，模型构建可以读取已经设置好的 TP/PP 状态，并生成当前 rank 的本地层和 VPP chunk。 + +各模型的实现集中在自己的目录： + +| 模型 | Provider | 配置与加载策略 | +| --- | --- | --- | +| GPT-2 | [model_provider.cc](../example/gpt2/model_provider.cc) | 从尺寸表登记标准名称及 `d*` 别名；选择 LLMC 加载或指定尺寸的随机初始化。 | +| LLaMA3 | [model_provider.cc](../example/llama3/model_provider.cc) | 使用 LLaMA3 配置校验和 LLMC loader。 | +| Qwen3 | [model_provider.cc](../example/qwen3/model_provider.cc) | 使用 Qwen3 配置校验和 LLMC loader。 | + +每个 provider 文件通过文件内的静态初始化 lambda 调用 `Register()`。登记发生在进入 `main()` 前，登记内容是名称和函数；模型构建属于后续的 rank 初始化流程。训练阶段使用已经解析、绑定的回调，注册表的修改阶段位于程序启动期。 + +注册表要求名称唯一、factory 为可调用对象。名称冲突和非法注册在登记时报告；`--help` 和未知模型错误中的可用名称列表都由同一张注册表生成。 + +CMake 将 provider 源文件直接编入 executable，使注册代码参与程序初始化。采用静态库封装 provider 时，链接规则需要显式保留这些注册对象。 + +### 3.1 实际模型配置 + +训练流程以已创建模型的配置作为结构信息来源： + +```cpp +auto base_model = model_provider(); +const auto model_config = base_model->Config(); +``` + +这份配置用于序列长度校验、PP 接收 shape、chunk 数量、TP loss 的有效词表范围和 checkpoint 元数据。模型的配置、参数布局和训练组件据此采用一致的尺寸信息。 + +模型适配层负责完整填充配置： + +- GPT-2 的 `n_kv_head` 与实际 `n_head` 相等，表达 MHA 结构。 +- LLaMA3 和 Qwen3 的 `original_vocab_size` 取自 LLMC 文件中的实际词表大小。 +- GPT-2 随机初始化先使用 `GPT2Config()`，再按尺寸表设置层数、head 数、KV-head 数和 hidden size，并按 TP 大小补齐词表。 + +GPT-2 的尺寸表位于 `example/gpt2/config.h`，包含 `gpt2/gpt2-medium/gpt2-large/gpt2-xl` 及对应的 `d12/d24/d36/d48` 别名。LLMC loader 根据文件头和权重布局构建实际模型，配置预设用于随机初始化路径。 + +## 4. 一次训练运行的生命周期 + +入口解析参数、取得 model provider，并组装 dataset、forward/loss、采样回调。随后由 `Pretrain()` 组织以下过程： + +1. **建立运行环境。** 校验公共选项，初始化全局并行环境和精度检查环境，在当前进程按线程配置运行本地 rank。 +2. **建立 rank 上下文。** 设置线程局部 rank、设备和 DP/TP/PP 通信组，计算梯度累积次数。 +3. **构建模型。** 调用 provider，读取实际 `Config()`，执行设备迁移、精度检查名称映射和 LoRA 注入。 +4. **包装并行执行。** 根据拓扑构建 `PipelineParallel` 和 DDP；组合使用时，先建立流水线，再为本地 chunk 添加 DDP 包装。 +5. **建立数据与优化组件。** 调用 dataset provider，创建加载器，通过 `SetupOptimizer()` 和 `SetupLRScheduler()` 创建优化器及调度器。 +6. **恢复训练状态。** 通过 `ResumeFromCheckpoint()` 恢复状态，并根据消费样本数推进数据迭代器。 +7. **执行训练步。** 运行 forward/loss、梯度累积、backward 和优化器更新，随后更新学习率调度器。 +8. **处理步后工作。** 汇总 loss、记录性能与内存指标，调用 `AfterStep`，按配置保存 checkpoint。 + +模型构建和并行包装有明确的先后关系：provider 使用 rank 上下文确定模型分片，设备迁移和 LoRA 注入准备待训练的参数集合，DDP/PP 包装建立执行与同步关系。`SetupOptimizer()` 最后从包装后的模型中选取普通参数或 LoRA 参数，并与 `NamedParameters()` 对齐。 + +### 4.1 Batch 与数据位置 + +`batch_size` 表示每个 microbatch 的序列数，`total_batch_size` 表示每次优化器更新的全局 token 数。设 DP 大小为 `D`、序列长度为 `S`、microbatch 大小为 `B`，梯度累积次数为： + +```text +M = total_batch_size / (B * S * D) +``` + +配置要求 `M` 为正整数。普通执行每次读取一个 microbatch，累计 `M` 次；PP 执行每次读取 `B * M` 条序列，再由调度器拆分为 microbatch。 + +Dataset provider 接收的目标全局训练样本数为： + +```text +num_iteration * total_batch_size / sequence_length +``` + +GPT provider 使用 `TinyShakespeareDataset` 读取有限的 token 文件，训练循环通过循环迭代达到目标步数。数据消费计数以全局样本数表示，因此不同执行路径可以使用同一套恢复位置逻辑。 + +### 4.2 精度与训练状态 + +前向计算和损失计算在 `AutocastGuard` 管理的精度环境中执行，backward 按训练流程设定的上下文运行。精度选择来自 `TrainingOptions::dtype`。 + +`llmc_filepath` 提供权重初始化来源，`load` 提供训练状态恢复来源。恢复前，model provider 先建立匹配的模型结构。Checkpoint 保存模型参数、按选项启用的 optimizer 状态、scheduler 状态、已完成步数和消费样本数，并记录模型尺寸与 DP/TP/SP/PP 元数据。 + +恢复训练时使用相同的模型结构、并行切分和调度配置。`num_iteration` 是包含已完成步数的目标总步数；`lr_decay_iters` 等参数定义整段训练的调度时间轴。 + +## 5. 普通执行与 PP 的共同计算协议 + +`ForwardStepGPT` 负责模型前向，并根据 TP 配置选择交叉熵损失。其返回的损失闭包捕获当前标签，执行器据此完成当前 microbatch 的训练。 + +普通路径的核心过程是： + +```text +读取 microbatch，迁移到设备 + → forward_step(model, inputs, target) + → result.loss_func(result.output) + → loss / M + → backward +``` + +完成 `M` 个 microbatch 后执行一次 optimizer step。DDP 的 `no_sync` 控制与梯度累积次数配合，最后一个 microbatch 触发相应同步。 + +### 5.1 流水线中的前向与损失 + +`PipelineParallel` 接收可选的 `training::ForwardStepFunction`，通过 `SetupSchedule()` 传给 `PipelineSchedule`。调度器按任务表执行每个本地 chunk： + +| 执行位置 | 输入 | 标签 | 回调结果的用途 | +| --- | --- | --- | --- | +| 首个 chunk | token 张量 | 最终 chunk 才接收标签 | 输出作为后续 chunk 的激活。 | +| 中间 chunk | 通信接收的激活 | 空指针 | 输出继续参与流水线传输。 | +| 最终 chunk | 前序激活 | 当前 microbatch 的标签 | 保存模型输出和损失闭包，供对应反向任务使用。 | + +调度器按 microbatch 索引保存最终损失闭包。在该 microbatch 的反向任务中，执行闭包、释放闭包引用、将损失除以 `M`，再调用 backward。输出和标签的生命周期由张量引用与闭包共同管理。 + +这份协议把任务计算和执行时序连接起来：任务负责模型输出与 microbatch 平均损失，训练执行器负责累积缩放、反向、梯度同步和参数更新。任务回调应遵守这一职责分配。 + +PP 接收 shape 由实际模型配置生成，采用 `[batch_size, sequence_length / SP, n_embd]`，其中 `SP` 为序列并行分片数。当前接收缓冲区使用 `float32`，跨阶段激活的形状与 dtype 应满足这一通信约定。 + +PP 使用现有的调度表生成、通信、microbatch 拆分和 `no_sync` 机制。`PipelineSchedule::Step()` 负责 zero-grad 和 optimizer step，外层训练循环负责 scheduler step。采用损失模块接口的调用方可以使用默认构造参数，经由 `ForwardOneChunk()` 和 `loss_fn` 执行。 + +## 6. 配置模型与命令行 + +配置由三层组成：框架默认值、模型运行预设、调用时覆盖。 + +### 6.1 默认值来源 + +`TrainingOptions` 定义框架公共选项及默认值。`arguments.cc` 的 `DEFINE_*` 从同一份默认配置读取初始值，使 C++ 直接调用与 CLI 启动使用一致的基础配置。 + +每个模型目录的 `train.flags` 表达该模型的运行预设： + +| 配置 | [GPT-2](../example/gpt2/train.flags) | [LLaMA3](../example/llama3/train.flags) / [Qwen3](../example/qwen3/train.flags) | +| --- | --- | --- | +| optimizer | `sgd` | `adam` | +| learning_rate | `0.0001` | `0.00001` | +| lora_target_modules | `c_attn,c_proj` | `c_attn,c_proj,c_fc,c_fc2` | +| load_optimizer_state / save_optimizer_state | `false` | `true` | +| profile_name | `gpt2` | 对应模型名 | + +预设还包含 batch、序列长度、精度和训练步数等参数。`--model` 的职责是选择注册的模型构建函数；优化器等训练设置由上述配置层共同确定。 + +### 6.2 参数文件与覆盖顺序 + +参数文件采用以下格式： + +```text +# LLaMA3 training preset +--model=llama3 +--optimizer=adam +--learning_rate=0.00001 +``` + +每行一个 `--name=value`，布尔参数也支持 `--flag` / `--noflag`。空行和以 `#` 开头的注释行用于组织内容。多个文件可以通过 `--flagfile=a.flags,b.flags` 加载，文件包含关系通过直接的 `--flagfile=...` 表达。显式文件路径及嵌套文件中的相对路径按进程工作目录解释。 + +参数按顺序生效，后出现的同名参数覆盖前面的值： + +```bash +./build/pretrain_gpt --flagfile=example/llama3/train.flags \ + --input_bin=/path/to/train.bin --llmc_filepath=/path/to/llama3.bin \ + --learning_rate=0.00002 --num_iteration=100 +``` + +这个调用使用 LLaMA3 预设，并将学习率和目标步数设为命令行给出的值。直接以 CLI 表达配置也是等价的使用方式： + +```bash +./build/pretrain_gpt --model=gpt2 --optimizer=sgd --learning_rate=0.0001 \ + --load_optimizer_state=false --save_optimizer_state=false \ + --input_bin=/path/to/train.bin --llmc_filepath=/path/to/gpt2.bin +``` + +通用入口要求通过 CLI 或参数文件提供模型名。其余公共选项按照框架默认值补齐；例如单独指定 `--model=gpt2` 时，optimizer 和 learning rate 的基础值仍是 `TrainingOptions` 中的 Adam、`1e-5`。加载 GPT-2 预设会将其设为 SGD、`1e-4`。 + +### 6.3 严格解析 + +`ParseTrainingFlags()` 将参数文件校验和 gflags 解析组织为同一启动步骤。文件内的参数名应已注册，非布尔参数应显式赋值，文件包含关系应可解析且无环。参数名错误或非布尔参数缺少赋值时，诊断包含文件和行号；文件读取和循环包含错误会指出相关路径。数值与布尔值的类型转换及其错误诊断由 gflags 完成。 + +gflags 执行 validator 时持有注册表锁，因此解析流程先采集参数类型快照，文件校验器再基于快照工作。文件校验器、路径定位和其他内部 helper 都位于 `arguments.cc` 的匿名 namespace 中。公共接口是 `ParseTrainingFlags()` 与 `TrainingOptionsFromFlags()`。 + +## 7. 构建产物与部署 + +训练能力由两个链接组件提供： + +| 组件 | 内容 | 使用方式 | +| --- | --- | --- | +| `InfiniTrain::infini_train` | 公共训练接口及执行实现、模型注册表 | C++ 调用者构造 `TrainingOptions` 与回调，调用 `Pretrain()`。 | +| `InfiniTrain::training_cli` | 参数定义、文件校验、路径定位和配置转换 | 采用标准训练命令行的应用按需链接。 | + +CLI 作为单独静态库，使训练参数的注册范围与应用的链接选择保持一致。框架 executable 链接接口使用 whole-archive 保留运行时注册对象；GPT 的 provider 源文件直接列入 executable 源码列表，参与启动期注册。 + +四个 executable 共用 `INFINITRAIN_GPT_EXAMPLE_SOURCES`： + +| executable | 配置选择 | +| --- | --- | +| `pretrain_gpt` | CLI 显式提供模型与配置。 | +| `gpt2` | 先加载 `configs/gpt2/train.flags`，再解析 CLI。 | +| `llama3` | 先加载 `configs/llama3/train.flags`，再解析 CLI。 | +| `qwen3` | 先加载 `configs/qwen3/train.flags`，再解析 CLI。 | + +CMake 通过 `INFINITRAIN_GPT_DEFAULT_FLAGFILE` 为三个兼容 executable 指定相对配置路径，并在构建时同步模型目录中的文件： + +```text +build/ + pretrain_gpt + gpt2 + llama3 + qwen3 + configs/ + gpt2/train.flags + llama3/train.flags + qwen3/train.flags +``` + +默认文件相对实际 executable 定位。部署时将程序与相邻的 `configs/` 一起分发，运行目录可以独立选择；Linux 通过 `/proc/self/exe` 定位，支持从 `PATH` 启动。 + +默认文件存在时，取值顺序为“默认文件 → 显式参数文件 → 后续 CLI”。默认文件缺失时，用户可通过显式 `--flagfile` 提供完整的运行配置；缺少配置来源时，入口会报告定位结果和处理建议。 + +编辑 `example//train.flags` 后，执行对应构建目标即可同步随包副本。直接编辑随包文件，或显式读取源码中的文件，更新会在下次启动时生效。 + +外部后端在入口解析参数前完成注册。后端通过设备抽象提供运行时操作，并可设置 `TrainingOptions::synchronize_device`，要求在模型上传后及训练资源释放前执行设备同步。 + +## 8. 开发与扩展 + +### 8.1 接入模型 + +新增同类 Transformer 模型时，围绕模型目录组织实现： + +1. 在 `config.h` 中定义架构、尺寸及其校验规则。 +2. 在 `checkpoint_loader.cc` 中实现权重格式映射，并填充实际模型配置。 +3. 在 `model_provider.cc` 中登记名称和 factory；factory 根据权重路径选择加载或随机初始化。 +4. 在 `train.flags` 中提供模型选择和训练预设。 +5. 将 provider、loader 源文件加入 GPT executable 的 CMake 源码列表。 + +注册表随后为入口提供名称解析、帮助列表和诊断信息。GPT-2 的尺寸别名通过该模型自己的尺寸表统一登记。 + +### 8.2 扩展数据与任务计算 + +新的数据来源通过实现 `Dataset` 并提供 `DatasetProvider` 接入。数据层负责文件格式、样本构造和可用序列长度，模型配置提供上下文长度约束；训练参数应同时满足二者。 + +新的目标函数通过 `ForwardStep` 和损失闭包表达。回调返回当前 microbatch 的平均损失，执行器按梯度累积次数缩放并完成更新。需要跨 microbatch 加权的目标，应同时定义清晰的损失归一化契约。 + +`AfterStep` 的执行位置是最后一个全局 rank。单 rank 文本生成通过 `freq_generate_txt` 控制频率，零表示关闭。涉及 TP collective 的采样流程需要组织相应 rank 协同执行。`val_loss_every`、`sample_every` 和 `overfit_single_batch` 当前是预留选项，对应分支为占位逻辑；扩展这些能力时，可分别接入评估步骤、采样步骤和可恢复的 sampler。 + +### 8.3 扩展框架训练能力 + +| 需求 | 主要修改位置 | +| --- | --- | +| 新增公共训练参数 | `TrainingOptions` 定义字段及默认值，`arguments.cc` 提供 CLI 映射和转换。 | +| 新增优化器或调度策略 | `SetupOptimizer()`、`SetupLRScheduler()` 及对应实现模块。 | +| 新增训练状态 | checkpoint 的保存、恢复和数据消费位置处理。 | +| 调整流水线执行 | `nn/parallel/pp` 中的调度和通信实现，并验证训练回调与损失模块两条调用路径。 | +| 构建外部 GPT executable | 使用 GPT 任务源码、provider/loader、框架 executable 链接接口和 CLI 组件，并分发相应参数文件。 | + +## 9. 与 Megatron 的语义对应 + +设计参考了 Megatron-LM `core_v0.12.0` 的 [pretrain_gpt.py](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/pretrain_gpt.py)、[training.py](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/megatron/training/training.py) 和 [pipeline schedules](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/megatron/core/pipeline_parallel/schedules.py)，并核对了提交 [4e1df1b6](https://github.com/NVIDIA/Megatron-LM/blob/4e1df1b62ccb799269b131046362f1f948ad5308/megatron/training/training.py) 中的配置与 builder 路径。 + +两者共享的设计原则是由任务提供模型、数据及 forward/loss，由训练基础设施决定模型包装、状态管理和执行时序。 + +| 语义 | Megatron | InfiniTrain 中的映射 | +| --- | --- | --- | +| 模型供应 | model provider 根据阶段信息创建模型 | provider 在 rank 上下文中创建 `TransformerModel`，由模型自身组织本地层和 VPP chunk。 | +| 数据供应 | dataset provider 提供训练所需数据集 | `DatasetProvider` 返回 `DatasetSplits`，加载器和消费位置由训练流程管理。 | +| 任务计算 | forward step 返回输出和 loss 回调 | `ForwardStepResult` 表达输出与延迟损失，执行器完成归一化和反向。 | +| 执行调度 | 公共训练流程选择 forward/backward schedule | 普通循环和 `PipelineSchedule` 消费同一套任务步骤协议。 | + +InfiniTrain 的 DataLoader 已负责组成 batch，PP 调度器负责拆分 microbatch，因此 forward 接口接收准备好的张量。模型阶段信息由已有并行上下文和 `TransformerModel` 管理。这些映射使 provider 语义与框架现有类型、线程模型和执行组件相衔接。 + +## 10. 验证方式 + +仓库中的自动化检查采用 C++ / GoogleTest，并通过 CTest 执行： + +| 测试 | 关注点 | +| --- | --- | +| [test_model_provider_registry.cc](../tests/pretrain/test_model_provider_registry.cc) | 延迟构建、回调与绑定参数的生命周期、名称冲突及诊断信息。 | +| [test_training_arguments.cc](../tests/pretrain/test_training_arguments.cc) | 默认值一致性、参数覆盖顺序、文件语法、嵌套包含、路径处理和显式配置接入。 | + +```bash +cmake --build build --target test_model_provider_registry test_training_arguments +ctest --test-dir build -R 'ModelProviderRegistryTest|TrainingArgumentsTest' --output-on-failure +``` + +训练行为验证使用固定模型权重和 token 数据，分别检查普通执行、并行执行和中断恢复。关键观察量包括 loss、参数更新、optimizer/scheduler 状态及消费样本数。中断恢复应在相同模型结构、拓扑和调度时间轴下，与连续训练进行对照。 + +已有本地记录覆盖 CPU-only/CUDA 构建、三模型数值对照、LoRA、断点恢复、两卡并行以及 LLaMA3 四卡组合。头文件检查验证了 `training.h` 的独立包含，运行检查还覆盖普通 CPU 训练和两卡 PP/VPP 冒烟检查。训练记录主要使用固定的小层数、小 hidden size 模型；GoogleTest 的持续覆盖范围如上表所列。 + +配置部署验证采用复制和移动 executable/configs 目录、切换工作目录、通过 `PATH` 启动等方式，并检查最终生效的配置。默认值和文件校验单测提供参数语义验证，实际 executable 检查提供构建产物与启动流程验证。 diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index 95e54730b..aabe16b7c 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -87,6 +87,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) gpt2_config.original_vocab_size = vocab_size; gpt2_config.n_layer = n_layer; gpt2_config.n_head = n_head; + gpt2_config.n_kv_head = n_head; gpt2_config.n_embd = n_embd; gpt2::SanitizeGPT2Config(gpt2_config); auto local_gpt2 = std::make_shared(gpt2_config); diff --git a/example/gpt2/config.h b/example/gpt2/config.h index c50680086..2a6dffc5c 100644 --- a/example/gpt2/config.h +++ b/example/gpt2/config.h @@ -1,11 +1,24 @@ #pragma once +#include +#include + #include "glog/logging.h" #include "infini_train/include/nn/modules/transformer/transformer_config.h" namespace nn = infini_train::nn; namespace gpt2 { +struct GPT2Size { + int64_t layers; + int64_t heads; + int64_t hidden; +}; +inline const std::unordered_map kModelSizes = { + {"gpt2", {12, 12, 768}}, {"d12", {12, 12, 768}}, {"gpt2-medium", {24, 16, 1024}}, {"d24", {24, 16, 1024}}, + {"gpt2-large", {36, 20, 1280}}, {"d36", {36, 20, 1280}}, {"gpt2-xl", {48, 25, 1600}}, {"d48", {48, 25, 1600}}, +}; + inline nn::TransformerConfig GPT2Config() { return {.block_size = 1024, .vocab_size = 50304, @@ -26,6 +39,19 @@ inline nn::TransformerConfig GPT2Config() { .rotary_interleaved = true}; } +inline nn::TransformerConfig GPT2Config(const std::string &name, int tensor_parallel_size) { + CHECK_GT(tensor_parallel_size, 0); + const auto &size = kModelSizes.at(name); + auto config = GPT2Config(); + config.n_layer = size.layers; + config.n_head = config.n_kv_head = size.heads; + config.n_embd = size.hidden; + // Scratch models pad vocabulary only as required by TP. + config.vocab_size + = (config.original_vocab_size + tensor_parallel_size - 1) / tensor_parallel_size * tensor_parallel_size; + return config; +} + inline void SanitizeGPT2Config(const nn::TransformerConfig &c) { CHECK_GT(c.block_size, 0); CHECK_GT(c.vocab_size, 0); diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc deleted file mode 100644 index 7697ad213..000000000 --- a/example/gpt2/main.cc +++ /dev/null @@ -1,635 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -// Out-of-tree builds inject the selected provider's declaration without adding -// a vendor dependency to the upstream example. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#endif - -#include "infini_train/include/autocast.h" -#include "infini_train/include/checkpoint/checkpoint.h" -#include "infini_train/include/core/runtime/device_guard.h" -#include "infini_train/include/dataloader.h" -#include "infini_train/include/device.h" -#include "infini_train/include/lr_scheduler.h" -#include "infini_train/include/nn/lora/lora_utils.h" -#include "infini_train/include/nn/modules/loss.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/nn/modules/transformer/transformer.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" -#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" -#include "infini_train/include/nn/parallel/global.h" -#include "infini_train/include/nn/parallel/parallel_functional.h" -#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" -#include "infini_train/include/nn/parallel/rank.h" -#include "infini_train/include/nn/parallel/reduce_op_type.h" -#include "infini_train/include/nn/parallel/tensor_parallel.h" -#include "infini_train/include/optimizer.h" -#ifdef PROFILE_MODE -#include "infini_train/include/profiler.h" -#endif -#include "infini_train/include/checkpoint/checkpoint_manager.h" -#include "infini_train/include/nn/parallel/utils.h" -#include "infini_train/include/utils/global_module_hook_registry.h" -#include "infini_train/include/utils/precision_check_config.h" -#include "infini_train/include/utils/precision_checker.h" - -#include "example/common/tiny_shakespeare_dataset.h" -#include "example/common/tokenizer.h" -#include "example/gpt2/checkpoint_loader.h" -#include "example/gpt2/config.h" - -// TODO(jym): Reorganize CLI flags into categories for better readability and maintainability. -// I/O -DEFINE_string(input_bin, "", "input .bin to train on"); -DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); -DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); -// model bin file is downloaded and processed using the script at -// https://github.com/karpathy/llm.c/blob/master/train_gpt2.py -DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); -DEFINE_string(model, "gpt2", "gpt2|gpt2-medium|gpt2-large|gpt2-xl|d12|d24|d36|d48"); -// token layout for each step of the optimization -DEFINE_uint32(batch_size, 4, "batch size, in units of #batch dimensions"); -DEFINE_uint32(sequence_length, 64, "sequence length"); -DEFINE_uint32(total_batch_size, 256, "total desired batch size, in units of #tokens"); -// workload (number of steps) -DEFINE_uint32(num_iteration, 10, "number of iterations to run"); -DEFINE_uint32(freq_generate_txt, 10, "frequency of text generation"); -DEFINE_uint32(text_length, 64, "the length of the generated text"); -// optimization -DEFINE_double(learning_rate, 1e-4, "Peak learning rate."); -DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); -// lr scheduler -DEFINE_double(min_lr, 0.0, "Minimum learning rate."); -DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); -DEFINE_int64(lr_warmup_iters, 0, "Number of linear warmup iterations."); -DEFINE_double(lr_warmup_init, 0.0, "Initial learning rate at the start of warmup."); -DEFINE_int64(lr_decay_iters, 0, "Number of iterations to decay LR over (0 = num_iteration)."); -// evaluation -DEFINE_uint32(val_loss_every, 0, "every how many steps to evaluate val loss?"); -DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); -// debugging -DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); -// memory management -DEFINE_string(device, "cuda", - "device type (cpu/cuda/privateuse1/), useless if using parallel " - "training mode"); -// parallel -DEFINE_int32(nthread_per_process, 1, - "Number of threads to use for each process. " - "When set > 1, enables data parallelism on the specified accelerator devices."); -DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); -DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); -DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); -DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); - -// precision -DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); -DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving"); -DEFINE_string(load, "", "checkpoint directory to resume from"); -DEFINE_string(save, "", "root directory used to store checkpoints"); -DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); -// precision check -DEFINE_string( - precision_check, "", - "precision check config: level=N,format=simple|table,output_md5=true|false,output_path=PATH,baseline=PATH"); - -// LoRA parameters -DEFINE_int32(lora_rank, 0, "LoRA rank (0 = disabled)"); -DEFINE_double(lora_alpha, 16.0, "LoRA alpha scaling factor"); -DEFINE_string(lora_target_modules, "c_attn,c_proj", - "LoRA target modules (comma-separated: c_attn,c_proj,c_fc,c_fc2,mlp.c_proj)"); -DEFINE_string(lora_save_path, "", "Path to save LoRA weights after training"); -DEFINE_string(lora_load_path, "", "Path to load LoRA weights from"); - -using namespace infini_train; - -namespace { -// validation -const std::unordered_set kSupportedModels - = {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"}; -constexpr char kDtypeFP32[] = "float32"; -constexpr char kDtypeBF16[] = "bfloat16"; -const std::unordered_set kSupportedLRDecayStyles - = {"none", "constant", "linear", "cosine", "inverse-square-root"}; - -// -const std::unordered_map kModelToConfigs = { - {"d12", {.block_size = 1024, .vocab_size = 50257, .n_layer = 12, .n_head = 12, .n_embd = 768}}, - {"d24", {.block_size = 1024, .vocab_size = 50257, .n_layer = 24, .n_head = 16, .n_embd = 1024}}, - {"d36", {.block_size = 1024, .vocab_size = 50257, .n_layer = 36, .n_head = 20, .n_embd = 1280}}, - {"d48", {.block_size = 1024, .vocab_size = 50257, .n_layer = 48, .n_head = 25, .n_embd = 1600}}, -}; - -} // namespace - -DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); -DEFINE_validator(device, [](const char *, const std::string &value) { return Device::ParseType(value).has_value(); }); -DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); -DEFINE_validator(lr_decay_style, - [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); - -void Train(const nn::parallel::Rank &rank) { - using namespace nn::parallel; - - { - if (rank.IsLastRank()) { - if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { - LOG(FATAL) << "Invalid configuration: --save is set ('" << FLAGS_save - << "'), but --save_interval is 0. " - << "They must be set together."; - } - - if (FLAGS_save.empty() && FLAGS_save_interval > 0) { - LOG(FATAL) << "Invalid configuration: --save_interval is set to " << FLAGS_save_interval - << ", but --save is empty. " - << "They must be set together."; - } - } - } - - // select the device - Device device; - const auto device_type = Device::ParseType(FLAGS_device).value(); - - int ddp_world_size = global::GetDataParallelSize(); - int tp_world_size = global::GetTensorParallelSize(); - int sp_world_size = global::GetSequenceParallelEnabled() ? tp_world_size : 1; - int pp_world_size = global::GetPipelineParallelSize(); - - if (FLAGS_sequence_parallel) { - CHECK_EQ(FLAGS_sequence_length % tp_world_size, 0) - << "sequence_length must be divisible by tp_world_size when SP is enabled (pad later if needed)."; - } - - int ddp_rank = 0; - int tp_rank = 0; - int pp_rank = 0; - - // Set thread-local global rank - // TODO(dcj): Use DeviceGuardImpl to get GlobalRank later. - nn::parallel::global::thread_global_rank = rank.GlobalRank(); - - const ProcessGroup *ddp_pg = nullptr; - const ProcessGroup *tp_pg = nullptr; - const ProcessGroup *pp_pg = nullptr; - - if (rank.IsParallel()) { - CHECK(device_type != Device::DeviceType::kCPU) << "Parallel training requires an accelerator backend"; - device = Device(device_type, global::GetDeviceIndex(rank.thread_rank())); - auto *pg_factory = ProcessGroupFactory::Instance(device.type()); - - if (ddp_world_size > 1) { - ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()), - GetDataParallelGroupRanks(rank.GlobalRank())); - ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); - } - - if (tp_world_size > 1) { - tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()), - GetTensorParallelGroupRanks(rank.GlobalRank())); - tp_rank = tp_pg->GetGroupRank(rank.GlobalRank()); - // NOTE(zbl): Reserved for VocabParallelEmbedding - nn::parallel::tp_rank = tp_rank; - } - - if (pp_world_size > 1) { - pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()), - GetPipelineParallelGroupRanks(rank.GlobalRank())); - pp_rank = pp_pg->GetGroupRank(rank.GlobalRank()); - - nn::parallel::pp_rank = pp_rank; - } - } else { - device = Device(device_type, 0); - } - - // calculate gradient accumulation from the desired total batch size and the current run configuration - const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size; - CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0); - const auto grad_accum_steps = FLAGS_total_batch_size / tokens_per_fwdbwd; - LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size - << " => calculated gradient accumulation steps: " << grad_accum_steps; - - // rng / reproducibility - // ManualSeed(42); - - // init the model, either from scratch or from OpenAI pretrained checkpoint - nn::TransformerConfig model_config = gpt2::GPT2Config(); - std::shared_ptr model = nullptr; - - if (!FLAGS_llmc_filepath.empty()) { - model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath); - } else if (kModelToConfigs.count(FLAGS_model)) { - model_config = kModelToConfigs.at(FLAGS_model); - gpt2::SanitizeGPT2Config(model_config); - model = std::make_shared(model_config); - } - - model->To(device); - - utils::PrecisionChecker::BuildNameMap(model.get()); - - // Get chunk size before wrapping with LoRA (needed for PipelineParallel) - auto gpt2_model = std::dynamic_pointer_cast(model); - CHECK(gpt2_model) << "GPT2 example expects GPT2 model."; - - // Apply LoRA using GetLoRAModel (in-place injection) - bool lora_enabled = FLAGS_lora_rank > 0; - if (lora_enabled) { - nn::lora::LoRAConfig lora_config{FLAGS_lora_rank, static_cast(FLAGS_lora_alpha), 0.0f, - nn::lora::ParseLoRATargetModules(FLAGS_lora_target_modules)}; - - // GetLoRAModel: in-place injection, modifies module tree directly - model = nn::lora::GetLoRAModel(model, lora_config); - - // Load LoRA weights if specified - if (!FLAGS_lora_load_path.empty()) { - LOG(INFO) << "Loading LoRA weights from: " << FLAGS_lora_load_path; - nn::lora::LoadLoRAWeights(model, FLAGS_lora_load_path); - } - - // Print LoRA summary - nn::lora::PrintLoRASummary(model, rank.GlobalRank()); - } - - // select the data type - // TODO(lzm): change to solely rely on the weight file info for determining the dtype when autocast is supported - DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { - dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { - dtype = DataType::kBFLOAT16; - } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported."; - } - - auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size); - - // Create optimizer - use GetLoRAParameters if LoRA is enabled - std::vector> params_to_optimize; - if (lora_enabled) { - params_to_optimize = nn::lora::GetLoRAParameters(model); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters"; - } else { - params_to_optimize = model->Parameters(); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters"; - } - - if (pp_world_size > 1) { - // NOTE(dcj): To ensure that the tensor shapes at the pipeline stage boundaries remain correct - // when sequence parallelism (SP) is enabled, we need to divide by sp_world_size. - auto shapes = std::vector>{ - {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}}; - - model = std::make_shared(model, pp_world_size, num_micro_batches, shapes, - pp_rank, device, model_config.GetChunkSize()); - if (ddp_world_size > 1) { - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; - auto *mutable_chunks = dynamic_cast(model.get())->mutable_chunks(); - for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) { - (*mutable_chunks)[chunk_id] - = std::make_shared(mutable_chunks->at(chunk_id), rank, ddp_config); - } - } - } else if (ddp_world_size > 1) { - // NOTE(dcj): Complete all device (.to(device)) and dtype (.to(dtype)) conversions - // before wrapping the model with DistributedDataParallel (DDP). - // Otherwise, DDP’s gradient hooks may be lost because new parameter tensors - // are created during the conversion. - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; - model = std::make_shared(model, rank, ddp_config); - } - - const size_t train_loader_batch_size = pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size; - DistributedDataLoader train_loader(std::make_shared(FLAGS_input_bin, FLAGS_sequence_length), - train_loader_batch_size, ddp_rank, ddp_world_size); - - std::optional val_loader = std::nullopt; - if (!FLAGS_input_val_bin.empty()) { - val_loader = DistributedDataLoader( - std::make_shared(FLAGS_input_val_bin, FLAGS_sequence_length), FLAGS_batch_size, - ddp_rank, ddp_world_size); - } - - // - // main training loop - // - - std::unique_ptr tokenizer = nullptr; - if (!FLAGS_tokenizer_bin.empty()) { - tokenizer = std::make_unique(FLAGS_tokenizer_bin); - } - - // TODO(dcj): support more complex optimizer later - // auto optimizer = optimizers::SGD(model->Parameters(), FLAGS_learning_rate); - auto optimizer_creator = optimizers::SGD::CreateNamed(FLAGS_learning_rate); - std::shared_ptr optimizer = nullptr; - std::unordered_set params_to_optimize_set; - params_to_optimize_set.reserve(params_to_optimize.size()); - for (const auto ¶m : params_to_optimize) { params_to_optimize_set.insert(param.get()); } - - NamedParameterList named_parameters; - for (const auto &[name, param] : model->NamedParameters()) { - if (params_to_optimize_set.contains(param.get())) { - named_parameters.emplace_back(name, param); - } - } - CHECK_EQ(named_parameters.size(), params_to_optimize.size()); - - if (FLAGS_zero_stage >= 1) { - auto model_chunks = (pp_world_size > 1) - ? *(dynamic_cast(model.get())->mutable_chunks()) - : std::vector>{model}; - optimizer = std::make_shared(optimizer_creator, named_parameters, - model_chunks, ddp_world_size, ddp_rank); - } else { - optimizer = optimizer_creator(named_parameters); - } - - const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; - TrainingLRSchedulerConfig sched_config; - sched_config.lr = static_cast(FLAGS_learning_rate); - sched_config.min_lr = static_cast(FLAGS_min_lr); - sched_config.lr_decay_style = FLAGS_lr_decay_style; - sched_config.lr_decay_iters = lr_decay_iters; - sched_config.lr_warmup_iters = FLAGS_lr_warmup_iters; - sched_config.lr_warmup_init = static_cast(FLAGS_lr_warmup_init); - auto scheduler = CreateLRScheduler(optimizer, sched_config); - - auto train_iter = train_loader.begin(); - std::shared_ptr loss_fn - = (tp_world_size > 1) ? std::static_pointer_cast( - std::make_shared(model_config.original_vocab_size)) - : std::static_pointer_cast(std::make_shared()); - loss_fn->To(device); - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start training"; - - auto impl = core::GetDeviceGuardImpl(device.type()); - - int start_step = 0; - TrainerState state; - const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, - .rank = rank, - .model = model, - .optimizer = nullptr, - .model_config = model_config, - .state = state, - .lr_scheduler = scheduler}); - start_step = resume_result.global_step; - size_t consumed_train_samples = resume_result.consumed_train_samples; - - auto advance_train_iter = [&]() { - ++train_iter; - if (train_iter == train_loader.end()) { - train_iter = train_loader.begin(); - } - }; - - // TODO(jym): Move resume position handling into a Sampler abstraction when available. - if (consumed_train_samples > 0) { - const size_t num_skips - = DataLoaderBatchesToSkip(consumed_train_samples, train_loader_batch_size, ddp_world_size); - for (size_t i = 0; i < num_skips; ++i) { advance_train_iter(); } - } - - auto next_train_batch = [&]() { - auto batch = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - advance_train_iter(); - consumed_train_samples += train_loader_batch_size * ddp_world_size; - return batch; - }; - - auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) { - SaveCheckpoint({ - .save_dir = save_dir, - .global_step = global_step, - .consumed_train_samples = consumed_train_samples, - .n_layer = model_config.n_layer, - .n_head = model_config.n_head, - .n_kv_head = model_config.n_kv_head, - .n_embd = model_config.n_embd, - .vocab_size = model_config.vocab_size, - .ddp_size = ddp_world_size, - .tp_size = tp_world_size, - .sp_size = sp_world_size, - .pp_size = pp_world_size, - .checkpoint_root_dir = FLAGS_save, - .max_checkpoint_keep = FLAGS_max_checkpoint_keep, - .rank = rank, - .model = *model, - .optimizer = nullptr, - .lr_scheduler = scheduler.get(), - }); - }; - - LOG(INFO) << "start training"; - - for (int step = start_step; step < FLAGS_num_iteration + 1; ++step) { - // Reset precision check counters at start of each iteration for file overwrite - utils::PrecisionChecker::ResetCounters(); - - const bool last_step = step == FLAGS_num_iteration; - - impl->ResetMemPoolHighWatermarks(device); - - const auto iter_start = std::chrono::high_resolution_clock::now(); - - // once in a while evaluate the validation dataset - if (FLAGS_val_loss_every > 0 && (step % FLAGS_val_loss_every == 0 || last_step) && val_loader.has_value()) { - // TODO(dcj): implement this after model.eval() is supported - } - // once in a while perform model inference on the master process - if (FLAGS_sample_every > 0 && (step % FLAGS_sample_every == 0 || last_step)) { - // TODO(dcj): implement this after model.eval() is supported - } - - // bit confusing: we want to make sure to eval and sample on 0th iteration - // but also after the very last iteration. so we loop for step <= num_iterations - // instead of just < num_iterations (one extra due to <=), only to do - // the validation/sampling one last time, and then we break right here as we're done. - if (last_step) { - break; - } - -#ifdef PROFILE_MODE - Profiler::Instance().SetTag("Step_" + std::to_string(step)); -#endif - - const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); - float lossf = 0.0f; - // model->Train(); - if (pp_world_size == 1) { - optimizer->ZeroGrad(); - - // if we are trying to overfit a single batch, we reset the loader here - if (FLAGS_overfit_single_batch) { - // train_loader.Reset(); - } - - for (int micro_step = 0; micro_step < grad_accum_steps; ++micro_step) { - // enable autocast for the current step - infini_train::AutocastGuard autocast_guard(device.type(), dtype); - - // (bs, seq_len), (bs, seq_len) - auto [x, y] = next_train_batch(); - x = std::make_shared(x->To(device)); - y = std::make_shared(y->To(device)); - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; - - // (bs, seq_len, vocab_size) - auto logits = (*model)({x, y})[0]; - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; - auto loss = (*loss_fn)({logits, y})[0]; - // FIXME(jym): verify gradient accumulation precision - loss = loss / grad_accum_steps; - - // disable autocast for the current step (backward is not under autocast) - autocast_guard.Disable(); - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward"; - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward"; - std::unique_ptr no_sync_guard; - if (ddp_world_size > 1 && micro_step != grad_accum_steps - 1) { - no_sync_guard = model->no_sync(); - } - loss->Backward(); - // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA - // between forward and backward. - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0]; - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; - } - - optimizer->Step(); - if (scheduler) { - scheduler->Step(); - } - } else { - auto [x, y] = next_train_batch(); - x = std::make_shared(x->To(device)); - y = std::make_shared(y->To(device)); - - lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); - if (scheduler) { - scheduler->Step(); - } - } - - if (ddp_world_size > 1) { - auto lossf_tensor = std::make_shared(&lossf, std::vector{}, DataType::kFLOAT32, device); - function::AllReduce(lossf_tensor, function::ReduceOpType::kAvg, ddp_pg); - lossf = static_cast(lossf_tensor->To(Device()).DataPtr())[0]; - } - - const auto iter_end = std::chrono::high_resolution_clock::now(); - const double duration_us = std::chrono::duration(iter_end - iter_start).count(); - const double tps = FLAGS_total_batch_size / (duration_us / 1e6); - - if (rank.IsLastRank()) { - size_t used_mb = 0, reserved_mb = 0; - std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); - LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " - "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", - step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, - used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, - pp_world_size); - - if ((step + 1) % FLAGS_freq_generate_txt == 0) { - if (tokenizer) { - // FIXME(jym): to support PP - CHECK_EQ(pp_world_size, 1); - tokenizer->GenerateText(*model, FLAGS_batch_size, FLAGS_sequence_length, FLAGS_text_length, device); - } - } - } - - if (!FLAGS_save.empty() && FLAGS_save_interval > 0) { - if ((step + 1) % FLAGS_save_interval == 0 || (step + 1) == FLAGS_num_iteration) { - std::filesystem::path step_dir - = std::filesystem::path(FLAGS_save) / std::format("checkpoint_step_{:06d}", step + 1); - if (rank.IsParallel()) { - step_dir /= std::format("rank_{:06d}", rank.GlobalRank()); - } - save_checkpoint(step_dir, step + 1); - } - } - } - - // Save LoRA weights if enabled and path specified - if (lora_enabled && !FLAGS_lora_save_path.empty()) { - LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; - nn::lora::SaveLoRAWeights(model, FLAGS_lora_save_path); - } - -#ifdef PROFILE_MODE - Profiler::Instance().Report("gpt2.report", Profiler::SortBy::DeviceTimePercentage); - Profiler::Instance().PrintRecords("gpt2.records.log"); -#endif - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA needs synchronization before training teardown. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { - impl->SynchronizeDevice(device); - } -#endif -} - -int main(int argc, char *argv[]) { - // Register provider metadata and implementations before gflags validates - // --device. The device runtime initializes lazily on first DeviceGuard use. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR - INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); -#endif - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); - - auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); - nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, - FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); - utils::PrecisionCheckEnv::Instance().Init(precision_config); - - LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); - - if (FLAGS_nthread_per_process > 1) { - std::vector threads; - for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), idx, - nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); - threads.emplace_back(Train, rank); - } - - for (auto &thread : threads) { thread.join(); } - } else { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), - FLAGS_nthread_per_process); - Train(rank); - } - - gflags::ShutDownCommandLineFlags(); - google::ShutdownGoogleLogging(); - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA parallel execution bypasses static destruction to avoid teardown failures. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - && Device::ParseType(FLAGS_device).value() == Device::DeviceType::kPrivateUse1) { - std::_Exit(0); - } -#endif - return 0; -} diff --git a/example/gpt2/model_provider.cc b/example/gpt2/model_provider.cc new file mode 100644 index 000000000..e9433b807 --- /dev/null +++ b/example/gpt2/model_provider.cc @@ -0,0 +1,26 @@ +#include "example/gpt2/checkpoint_loader.h" +#include "example/gpt2/config.h" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/training/model_provider_registry.h" + +namespace gpt2 { +namespace { + +[[maybe_unused]] const bool registered = [] { + auto ®istry = infini_train::training::ModelProviderRegistry::Instance(); + for (const auto &[name, size] : kModelSizes) { + registry.Register(name, [name = name](const std::string &weights) { + if (!weights.empty()) { + return LoadFromLLMC(weights); + } + auto config = GPT2Config(name, infini_train::nn::parallel::global::GetTensorParallelSize()); + SanitizeGPT2Config(config); + return std::make_shared(config); + }); + } + return true; +}(); + +} // namespace +} // namespace gpt2 diff --git a/example/gpt2/train.flags b/example/gpt2/train.flags new file mode 100644 index 000000000..980faea7a --- /dev/null +++ b/example/gpt2/train.flags @@ -0,0 +1,15 @@ +# GPT-2 training defaults. CLI options after --flagfile override these values. +--model=gpt2 +--optimizer=sgd +--learning_rate=0.0001 +--batch_size=4 +--sequence_length=64 +--total_batch_size=256 +--num_iteration=10 +--dtype=float32 +--zero_stage=0 +--lr_decay_style=constant +--lora_target_modules=c_attn,c_proj +--load_optimizer_state=false +--save_optimizer_state=false +--profile_name=gpt2 diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index 32b11e430..e4c5022b5 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -71,6 +71,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) nn::TransformerConfig llama3_config = llama3::LLaMA3Config(); llama3_config.block_size = block_size; llama3_config.vocab_size = vocab_size; + llama3_config.original_vocab_size = vocab_size; llama3_config.n_layer = n_layer; llama3_config.n_head = n_head; llama3_config.n_kv_head = n_kv_head; diff --git a/example/llama3/model_provider.cc b/example/llama3/model_provider.cc new file mode 100644 index 000000000..96fda2285 --- /dev/null +++ b/example/llama3/model_provider.cc @@ -0,0 +1,22 @@ +#include "example/llama3/checkpoint_loader.h" +#include "example/llama3/config.h" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/training/model_provider_registry.h" + +namespace llama3 { +namespace { + +[[maybe_unused]] const bool registered = [] { + infini_train::training::ModelProviderRegistry::Instance().Register("llama3", [](const std::string &weights) { + if (!weights.empty()) { + return LoadFromLLMC(weights); + } + auto config = LLaMA3Config(); + SanitizeLLaMA3Config(config); + return std::make_shared(config); + }); + return true; +}(); + +} // namespace +} // namespace llama3 diff --git a/example/llama3/train.flags b/example/llama3/train.flags new file mode 100644 index 000000000..a7b36d2fc --- /dev/null +++ b/example/llama3/train.flags @@ -0,0 +1,15 @@ +# LLaMA3 training defaults. CLI options after --flagfile override these values. +--model=llama3 +--optimizer=adam +--learning_rate=0.00001 +--batch_size=4 +--sequence_length=64 +--total_batch_size=256 +--num_iteration=10 +--dtype=float32 +--zero_stage=0 +--lr_decay_style=constant +--lora_target_modules=c_attn,c_proj,c_fc,c_fc2 +--load_optimizer_state=true +--save_optimizer_state=true +--profile_name=llama3 diff --git a/example/pretrain_gpt.cc b/example/pretrain_gpt.cc new file mode 100644 index 000000000..9c983b68f --- /dev/null +++ b/example/pretrain_gpt.cc @@ -0,0 +1,115 @@ +// Out-of-tree providers register before parsing --device. +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER +#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER +#endif + +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "example/common/tiny_shakespeare_dataset.h" +#include "example/common/tokenizer.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/training/arguments.h" +#include "infini_train/include/training/model_provider_registry.h" +#include "infini_train/include/training/training.h" + +DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(input_val_bin, "", "input .bin for validation (evaluation is not implemented yet)"); +DEFINE_string(tokenizer_bin, "", "input tokenizer .bin"); +DEFINE_string(llmc_filepath, "", "LLMC weights to initialize the model; --load resumes training state"); +DEFINE_string(model, "", "Registered model name; see --help for available models"); +DEFINE_uint32(freq_generate_txt, 10, "text generation frequency; 0 disables generation"); +DEFINE_uint32(text_length, 64, "number of tokens to generate"); + +using namespace infini_train; + +namespace { + +training::ForwardStepResult ForwardStepGPT(nn::Module &model, const training::TensorList &inputs, + const std::shared_ptr &target, const nn::TransformerConfig &config) { + auto output = model(inputs); + if (!target) { + return {std::move(output), {}}; + } + std::shared_ptr loss_fn; + if (nn::parallel::global::GetTensorParallelSize() > 1) { + loss_fn = std::make_shared(config.original_vocab_size); + } else { + loss_fn = std::make_shared(); + } + return {std::move(output), [loss_fn, target](const training::TensorList &logits) { + return (*loss_fn)({logits.at(0), target}).at(0); + }}; +} + +} // namespace + +int main(int argc, char *argv[]) { +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); +#endif + const auto &models = training::ModelProviderRegistry::Instance(); + gflags::SetUsageMessage("GPT pretraining. Available models: " + models.Names() + + "\nUse --model= or --flagfile=example//train.flags."); +#ifdef INFINITRAIN_GPT_DEFAULT_FLAGFILE + training::ParseTrainingFlags(&argc, &argv, INFINITRAIN_GPT_DEFAULT_FLAGFILE); +#else + training::ParseTrainingFlags(&argc, &argv); +#endif + google::InitGoogleLogging(argv[0]); + + const std::string model_name = FLAGS_model; + CHECK(!model_name.empty()) << "--model is required; pass it directly or use --flagfile=example//train.flags"; + const auto model_provider = models.Resolve(model_name, FLAGS_llmc_filepath); + CHECK(!FLAGS_input_bin.empty()) << "--input_bin is required"; + auto options = training::TrainingOptionsFromFlags(); +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + options.synchronize_device = INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + && Device::ParseType(options.device) == Device::DeviceType::kPrivateUse1; +#endif + + const std::string train_path = FLAGS_input_bin; + const std::string valid_path = FLAGS_input_val_bin; + const auto dataset_provider = [train_path, valid_path, sequence_length = options.sequence_length](size_t) { + training::DatasetSplits datasets; + datasets.train = std::make_shared(train_path, sequence_length); + if (!valid_path.empty()) { + datasets.valid = std::make_shared(valid_path, sequence_length); + } + return datasets; + }; + + std::shared_ptr tokenizer; + if (!FLAGS_tokenizer_bin.empty() && FLAGS_freq_generate_txt > 0) { + CHECK_EQ(options.pipeline_parallel, 1) << "Text generation does not support PP"; + tokenizer = std::make_shared(FLAGS_tokenizer_bin); + } + const auto after_step = [tokenizer, frequency = FLAGS_freq_generate_txt, text_length = FLAGS_text_length, + &options](nn::Module &model, const Device &device, int64_t step) { + if (tokenizer && step % frequency == 0) { + tokenizer->GenerateText(model, options.batch_size, options.sequence_length, text_length, device); + } + }; + + training::Pretrain(options, dataset_provider, model_provider, ForwardStepGPT, after_step); + + gflags::ShutDownCommandLineFlags(); + google::ShutdownGoogleLogging(); +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + // Preserve the external backend's existing static-destruction workaround. + if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS + && Device::ParseType(options.device).value() == Device::DeviceType::kPrivateUse1) { + std::_Exit(0); + } +#endif + return 0; +} diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc index 181217a32..a9814de26 100644 --- a/example/qwen3/checkpoint_loader.cc +++ b/example/qwen3/checkpoint_loader.cc @@ -97,6 +97,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) nn::TransformerConfig qwen3_config = qwen3::Qwen3Config(); qwen3_config.block_size = block_size; qwen3_config.vocab_size = vocab_size; + qwen3_config.original_vocab_size = vocab_size; qwen3_config.n_layer = n_layer; qwen3_config.n_head = n_head; qwen3_config.n_kv_head = n_kv_head; diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc deleted file mode 100644 index ba969fd21..000000000 --- a/example/qwen3/main.cc +++ /dev/null @@ -1,621 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -// Out-of-tree builds inject the selected provider's declaration without adding -// a vendor dependency to the upstream example. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#endif - -#include "infini_train/include/autocast.h" -#include "infini_train/include/checkpoint/checkpoint.h" -#include "infini_train/include/checkpoint/checkpoint_manager.h" -#include "infini_train/include/core/runtime/device_guard.h" -#include "infini_train/include/dataloader.h" -#include "infini_train/include/device.h" -#include "infini_train/include/lr_scheduler.h" -#include "infini_train/include/nn/lora/lora_utils.h" -#include "infini_train/include/nn/modules/loss.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/nn/modules/transformer/transformer.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" -#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" -#include "infini_train/include/nn/parallel/global.h" -#include "infini_train/include/nn/parallel/parallel_functional.h" -#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" -#include "infini_train/include/nn/parallel/process_group.h" -#include "infini_train/include/nn/parallel/rank.h" -#include "infini_train/include/nn/parallel/reduce_op_type.h" -#include "infini_train/include/nn/parallel/tensor_parallel.h" -#include "infini_train/include/nn/parallel/utils.h" -#include "infini_train/include/optimizer.h" -#include "infini_train/include/utils/global_module_hook_registry.h" -#include "infini_train/include/utils/precision_check_config.h" -#include "infini_train/include/utils/precision_checker.h" -#ifdef PROFILE_MODE -#include "infini_train/include/profiler.h" -#endif - -#include "example/common/tiny_shakespeare_dataset.h" -#include "example/common/tokenizer.h" -#include "example/qwen3/checkpoint_loader.h" -#include "example/qwen3/config.h" - -// TODO(jym): Reorganize CLI flags into categories for better readability and maintainability. -// I/O -DEFINE_string(input_bin, "", "input .bin to train on"); -DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); -DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); -// model bin file is downloaded and processed using the script at -// Converted from the official Hugging Face checkpoint into the shared LLMC v4 format. -DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); -DEFINE_string(model, "qwen3", "Qwen/Qwen3-8B"); -// token layout for each step of the optimization -DEFINE_uint32(batch_size, 4, "batch size, in units of #batch dimensions"); -DEFINE_uint32(sequence_length, 64, "sequence length"); -DEFINE_uint32(total_batch_size, 256, "total desired batch size, in units of #tokens"); -// workload (number of steps) -DEFINE_uint32(num_iteration, 10, "number of iterations to run"); -DEFINE_uint32(freq_generate_txt, 10, "frequency of text generation"); -DEFINE_uint32(text_length, 64, "the length of the generated text"); -// optimization -DEFINE_double(learning_rate, 1e-5, "Peak learning rate."); -DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); -// lr scheduler -DEFINE_double(min_lr, 0.0, "Minimum learning rate."); -DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); -DEFINE_int64(lr_warmup_iters, 0, "Number of linear warmup iterations."); -DEFINE_double(lr_warmup_init, 0.0, "Initial learning rate at the start of warmup."); -DEFINE_int64(lr_decay_iters, 0, "Number of iterations to decay LR over (0 = num_iteration)."); -// evaluation -DEFINE_uint32(val_loss_every, 0, "every how many steps to evaluate val loss?"); -DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); -// debugging -DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); -// memory management -DEFINE_string(device, "cuda", - "device type (cpu/cuda/privateuse1/), useless if using parallel " - "training mode"); -// parallel -DEFINE_int32(nthread_per_process, 1, - "Number of threads to use for each process. " - "When set > 1, enables data parallelism on the specified accelerator devices."); -DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); -DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); -DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); -DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); -// precision -DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); -DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving"); -DEFINE_string(load, "", "checkpoint directory to resume from"); -DEFINE_string(save, "", "root directory used to store checkpoints"); -DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); -DEFINE_bool(load_optimizer_state, true, "whether optimizer state is restored from checkpoints"); -DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints"); - -// precision check -DEFINE_string( - precision_check, "", - "precision check config: level=N,format=simple|table,output_md5=true|false,output_path=PATH,baseline=PATH"); -// LoRA parameters -DEFINE_int32(lora_rank, 0, "LoRA rank (0 = disabled)"); -DEFINE_double(lora_alpha, 16.0, "LoRA alpha scaling factor"); -DEFINE_string(lora_target_modules, "c_attn,c_proj,c_fc,c_fc2", "LoRA target modules (comma-separated)"); -DEFINE_string(lora_save_path, "", "Path to save LoRA weights after training"); -DEFINE_string(lora_load_path, "", "Path to load LoRA weights from"); - -using namespace infini_train; - -namespace { -// validation -const std::unordered_set kSupportedModels = {"qwen3"}; -constexpr char kDtypeFP32[] = "float32"; -constexpr char kDtypeBF16[] = "bfloat16"; -const std::unordered_set kSupportedLRDecayStyles - = {"none", "constant", "linear", "cosine", "inverse-square-root"}; -} // namespace - -DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); -DEFINE_validator(device, [](const char *, const std::string &value) { return Device::ParseType(value).has_value(); }); -DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); -DEFINE_validator(lr_decay_style, - [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); - -void Train(const nn::parallel::Rank &rank) { - using namespace nn::parallel; - - { - if (rank.IsLastRank()) { - if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { - LOG(FATAL) << "Invalid configuration: --save is set ('" << FLAGS_save - << "'), but --save_interval is 0. " - << "They must be set together."; - } - - if (FLAGS_save.empty() && FLAGS_save_interval > 0) { - LOG(FATAL) << "Invalid configuration: --save_interval is set to " << FLAGS_save_interval - << ", but --save is empty. " - << "They must be set together."; - } - } - } - - // select the device - Device device; - const auto device_type = Device::ParseType(FLAGS_device).value(); - - int ddp_world_size = global::GetDataParallelSize(); - int tp_world_size = global::GetTensorParallelSize(); - int sp_world_size = global::GetSequenceParallelEnabled() ? tp_world_size : 1; - int pp_world_size = global::GetPipelineParallelSize(); - - if (FLAGS_sequence_parallel) { - CHECK_EQ(FLAGS_sequence_length % tp_world_size, 0) - << "sequence_length must be divisible by tp_world_size when SP is enabled (pad later if needed)."; - } - - int ddp_rank = 0; - int tp_rank = 0; - int pp_rank = 0; - - // Set thread-local global rank - nn::parallel::global::thread_global_rank = rank.GlobalRank(); - - const ProcessGroup *ddp_pg = nullptr; - const ProcessGroup *tp_pg = nullptr; - const ProcessGroup *pp_pg = nullptr; - - if (rank.IsParallel()) { - CHECK(device_type != Device::DeviceType::kCPU) << "Parallel training requires an accelerator backend"; - device = Device(device_type, global::GetDeviceIndex(rank.thread_rank())); - auto *pg_factory = ProcessGroupFactory::Instance(device.type()); - - if (ddp_world_size > 1) { - ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()), - GetDataParallelGroupRanks(rank.GlobalRank())); - ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); - } - - if (tp_world_size > 1) { - tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()), - GetTensorParallelGroupRanks(rank.GlobalRank())); - tp_rank = tp_pg->GetGroupRank(rank.GlobalRank()); - // NOTE(zbl): Reserved for VocabParallelEmbedding - nn::parallel::tp_rank = tp_rank; - } - - if (pp_world_size > 1) { - pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()), - GetPipelineParallelGroupRanks(rank.GlobalRank())); - pp_rank = pp_pg->GetGroupRank(rank.GlobalRank()); - - nn::parallel::pp_rank = pp_rank; - } - } else { - device = Device(device_type, 0); - } - - // calculate gradient accumulation from the desired total batch size and the current run configuration - const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size; - CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0); - const auto grad_accum_steps = FLAGS_total_batch_size / tokens_per_fwdbwd; - if (rank.IsMainRank()) { - LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size - << " => calculated gradient accumulation steps: " << grad_accum_steps; - } - - // rng / reproducibility - // ManualSeed(42); - - nn::TransformerConfig model_config = qwen3::Qwen3Config(); - std::shared_ptr model = nullptr; - if (!FLAGS_llmc_filepath.empty()) { - model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath); - } else { - qwen3::SanitizeQwen3Config(model_config); - model = std::make_shared(model_config); - } - - model->To(device); - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA requires model upload to finish before communication begins. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { - core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); - } -#endif - - utils::PrecisionChecker::BuildNameMap(model.get()); - - // Apply LoRA using GetLoRAModel (in-place injection) - bool lora_enabled = FLAGS_lora_rank > 0; - if (lora_enabled) { - nn::lora::LoRAConfig lora_config{FLAGS_lora_rank, static_cast(FLAGS_lora_alpha), 0.0f, - nn::lora::ParseLoRATargetModules(FLAGS_lora_target_modules)}; - - // GetLoRAModel: in-place injection, modifies module tree directly - model = nn::lora::GetLoRAModel(model, lora_config); - - // Load LoRA weights if specified - if (!FLAGS_lora_load_path.empty()) { - LOG(INFO) << "Loading LoRA weights from: " << FLAGS_lora_load_path; - nn::lora::LoadLoRAWeights(model, FLAGS_lora_load_path); - } - - // Print LoRA summary - nn::lora::PrintLoRASummary(model, rank.GlobalRank()); - } - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; - - DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { - dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { - dtype = DataType::kBFLOAT16; - } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported."; - } - - auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size); - - if (pp_world_size > 1) { - // NOTE(dcj): To ensure that the tensor shapes at the pipeline stage boundaries remain correct - // when sequence parallelism (SP) is enabled, we need to divide by sp_world_size. - auto shapes = std::vector>{ - {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}}; - - model = std::make_shared(model, pp_world_size, num_micro_batches, shapes, - pp_rank, device, model_config.GetChunkSize()); - if (ddp_world_size > 1) { - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; - auto *mutable_chunks = dynamic_cast(model.get())->mutable_chunks(); - for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) { - (*mutable_chunks)[chunk_id] - = std::make_shared(mutable_chunks->at(chunk_id), rank, ddp_config); - } - } - } else if (ddp_world_size > 1) { - // NOTE(dcj): Complete all device (.to(device)) and dtype (.to(dtype)) conversions - // before wrapping the model with DistributedDataParallel (DDP). - // Otherwise, DDP's gradient hooks may be lost because new parameter tensors - // are created during the conversion. - - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; - model = std::make_shared(model, rank, ddp_config); - } - - const size_t train_loader_batch_size = pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size; - DistributedDataLoader train_loader(std::make_shared(FLAGS_input_bin, FLAGS_sequence_length), - train_loader_batch_size, ddp_rank, ddp_world_size); - - std::optional val_loader = std::nullopt; - if (!FLAGS_input_val_bin.empty()) { - val_loader = DistributedDataLoader( - std::make_shared(FLAGS_input_val_bin, FLAGS_sequence_length), FLAGS_batch_size, - ddp_rank, ddp_world_size); - } - - // - // main training loop - // - std::unique_ptr tokenizer = nullptr; - if (!FLAGS_tokenizer_bin.empty()) { - tokenizer = std::make_unique(FLAGS_tokenizer_bin); - } - - // TODO(dcj): support more complex optimizer later - // auto optimizer = optimizers::Adam(model->Parameters(), FLAGS_learning_rate); - auto optimizer_creator = optimizers::Adam::CreateNamed(FLAGS_learning_rate); - std::shared_ptr optimizer = nullptr; - - std::vector> params_to_optimize; - if (lora_enabled) { - params_to_optimize = nn::lora::GetLoRAParameters(model); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters"; - } else { - params_to_optimize = model->Parameters(); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters"; - } - std::unordered_set params_to_optimize_set; - params_to_optimize_set.reserve(params_to_optimize.size()); - for (const auto ¶m : params_to_optimize) { params_to_optimize_set.insert(param.get()); } - - NamedParameterList named_parameters; - for (const auto &[name, param] : model->NamedParameters()) { - if (params_to_optimize_set.contains(param.get())) { - named_parameters.emplace_back(name, param); - } - } - CHECK_EQ(named_parameters.size(), params_to_optimize.size()); - - if (FLAGS_zero_stage >= 1) { - auto model_chunks = (pp_world_size > 1) - ? *(dynamic_cast(model.get())->mutable_chunks()) - : std::vector>{model}; - optimizer = std::make_shared(optimizer_creator, named_parameters, - model_chunks, ddp_world_size, ddp_rank); - } else { - optimizer = optimizer_creator(named_parameters); - } - - const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; - TrainingLRSchedulerConfig sched_config; - sched_config.lr = static_cast(FLAGS_learning_rate); - sched_config.min_lr = static_cast(FLAGS_min_lr); - sched_config.lr_decay_style = FLAGS_lr_decay_style; - sched_config.lr_decay_iters = lr_decay_iters; - sched_config.lr_warmup_iters = FLAGS_lr_warmup_iters; - sched_config.lr_warmup_init = static_cast(FLAGS_lr_warmup_init); - auto scheduler = CreateLRScheduler(optimizer, sched_config); - - auto train_iter = train_loader.begin(); - std::shared_ptr loss_fn - = (tp_world_size > 1) ? std::static_pointer_cast(std::make_shared()) - : std::static_pointer_cast(std::make_shared()); - loss_fn->To(device); - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start training"; - - auto impl = core::GetDeviceGuardImpl(device.type()); - - int start_step = 0; - TrainerState state; - const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, - .rank = rank, - .model = model, - .optimizer = FLAGS_load_optimizer_state ? optimizer : nullptr, - .model_config = model_config, - .state = state, - .lr_scheduler = scheduler}); - - start_step = resume_result.global_step; - size_t consumed_train_samples = resume_result.consumed_train_samples; - - auto advance_train_iter = [&]() { - ++train_iter; - if (train_iter == train_loader.end()) { - train_iter = train_loader.begin(); - } - }; - - // TODO(jym): Replace with Sampler abstraction when available. - if (consumed_train_samples > 0) { - const size_t num_skips - = DataLoaderBatchesToSkip(consumed_train_samples, train_loader_batch_size, ddp_world_size); - for (size_t i = 0; i < num_skips; ++i) { advance_train_iter(); } - } - - auto next_train_batch = [&]() { - auto batch = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - advance_train_iter(); - consumed_train_samples += train_loader_batch_size * ddp_world_size; - return batch; - }; - - auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) { - SaveCheckpoint({ - .save_dir = save_dir, - .global_step = global_step, - .consumed_train_samples = consumed_train_samples, - .n_layer = model_config.n_layer, - .n_head = model_config.n_head, - .n_kv_head = model_config.n_kv_head, - .n_embd = model_config.n_embd, - .vocab_size = model_config.vocab_size, - .ddp_size = ddp_world_size, - .tp_size = tp_world_size, - .sp_size = sp_world_size, - .pp_size = pp_world_size, - .checkpoint_root_dir = FLAGS_save, - .max_checkpoint_keep = FLAGS_max_checkpoint_keep, - .rank = rank, - .model = *model, - .optimizer = FLAGS_save_optimizer_state ? optimizer.get() : nullptr, - .lr_scheduler = scheduler.get(), - }); - }; - - for (int step = start_step; step < FLAGS_num_iteration + 1; ++step) { - // Reset precision check counters at start of each iteration for file overwrite - utils::PrecisionChecker::ResetCounters(); - - const bool last_step = step == FLAGS_num_iteration; - - impl->ResetMemPoolHighWatermarks(device); - - const auto iter_start = std::chrono::high_resolution_clock::now(); - - // once in a while evaluate the validation dataset - if (FLAGS_val_loss_every > 0 && (step % FLAGS_val_loss_every == 0 || last_step) && val_loader.has_value()) { - // TODO(dcj): implement this after model.eval() is supported - } - // once in a while perform model inference on the master process - if (FLAGS_sample_every > 0 && (step % FLAGS_sample_every == 0 || last_step)) { - // TODO(dcj): implement this after model.eval() is supported - } - - // bit confusing: we want to make sure to eval and sample on 0th iteration - // but also after the very last iteration. so we loop for step <= num_iterations - // instead of just < num_iterations (one extra due to <=), only to do - // the validation/sampling one last time, and then we break right here as we're done. - if (last_step) { - break; - } - -#ifdef PROFILE_MODE - Profiler::Instance().SetTag("Step_" + std::to_string(step)); -#endif - - const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); - float lossf = 0.0f; - if (pp_world_size == 1) { - // model->Train(); - optimizer->ZeroGrad(); - - // if we are trying to overfit a single batch, we reset the loader here - if (FLAGS_overfit_single_batch) { - // train_loader.Reset(); - } - - for (int micro_step = 0; micro_step < grad_accum_steps; ++micro_step) { - // enable autocast for the current step - infini_train::AutocastGuard autocast_guard(device.type(), dtype); - - // (bs, seq_len), (bs, seq_len) - auto [x, y] = next_train_batch(); - x = std::make_shared(x->To(device)); - y = std::make_shared(y->To(device)); - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; - // (bs, seq_len, vocab_size) - auto logits = (*model)({x, y})[0]; - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; - auto loss = (*loss_fn)({logits, y})[0]; - // FIXME(jym): verify gradient accumulation precision - loss = loss / grad_accum_steps; - - // disable autocast for the current step (backward is not under autocast) - autocast_guard.Disable(); - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward"; - - LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward"; - std::unique_ptr no_sync_guard; - if (ddp_world_size > 1 && micro_step != grad_accum_steps - 1) { - no_sync_guard = model->no_sync(); - } - loss->Backward(); - // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA - // between forward and backward. - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0]; - LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; - } - - optimizer->Step(); - if (scheduler) { - scheduler->Step(); - } - } else { - auto [x, y] = next_train_batch(); - x = std::make_shared(x->To(device)); - y = std::make_shared(y->To(device)); - - lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); - if (scheduler) { - scheduler->Step(); - } - } - - if (ddp_world_size > 1) { - auto lossf_tensor = std::make_shared(&lossf, std::vector{}, DataType::kFLOAT32, device); - function::AllReduce(lossf_tensor, function::ReduceOpType::kAvg, ddp_pg); - lossf = static_cast(lossf_tensor->To(Device()).DataPtr())[0]; - } - - const auto iter_end = std::chrono::high_resolution_clock::now(); - const double duration_us = std::chrono::duration(iter_end - iter_start).count(); - const double tps = FLAGS_total_batch_size / (duration_us / 1e6); - - if (rank.IsLastRank()) { - size_t used_mb = 0, reserved_mb = 0; - std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); - LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " - "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", - step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, - used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, - pp_world_size); - - if ((step + 1) % FLAGS_freq_generate_txt == 0) { - // FIXME(jym): to support PP - if (tokenizer) { - CHECK_EQ(pp_world_size, 1); - tokenizer->GenerateText(*model, FLAGS_batch_size, FLAGS_sequence_length, FLAGS_text_length, device); - } - } - } - - if (!FLAGS_save.empty() && FLAGS_save_interval > 0) { - if ((step + 1) % FLAGS_save_interval == 0 || (step + 1) == FLAGS_num_iteration) { - std::filesystem::path step_dir - = std::filesystem::path(FLAGS_save) / std::format("checkpoint_step_{:06d}", step + 1); - if (rank.IsParallel()) { - step_dir /= std::format("rank_{:06d}", rank.GlobalRank()); - } - save_checkpoint(step_dir, step + 1); - } - } - } - - // Save LoRA weights if enabled and path specified - if (lora_enabled && !FLAGS_lora_save_path.empty()) { - LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; - nn::lora::SaveLoRAWeights(model, FLAGS_lora_save_path); - } - -#ifdef PROFILE_MODE - Profiler::Instance().Report("qwen3.report", Profiler::SortBy::DeviceTimePercentage); - Profiler::Instance().PrintRecords("qwen3.records.log"); -#endif - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA needs synchronization before training teardown. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { - impl->SynchronizeDevice(device); - } -#endif -} - -int main(int argc, char *argv[]) { - // Register provider metadata and implementations before gflags validates - // --device. The device runtime initializes lazily on first DeviceGuard use. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR - INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); -#endif - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); - - auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); - nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, - FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); - utils::PrecisionCheckEnv::Instance().Init(precision_config); - - LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); - - if (FLAGS_nthread_per_process > 1) { - std::vector threads; - for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), idx, - nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); - threads.emplace_back(Train, rank); - } - - for (auto &thread : threads) { thread.join(); } - } else { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), - FLAGS_nthread_per_process); - Train(rank); - } - - gflags::ShutDownCommandLineFlags(); - google::ShutdownGoogleLogging(); - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA parallel execution bypasses static destruction to avoid teardown failures. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - && Device::ParseType(FLAGS_device).value() == Device::DeviceType::kPrivateUse1) { - std::_Exit(0); - } -#endif - return 0; -} diff --git a/example/qwen3/model_provider.cc b/example/qwen3/model_provider.cc new file mode 100644 index 000000000..8d4bfdca5 --- /dev/null +++ b/example/qwen3/model_provider.cc @@ -0,0 +1,22 @@ +#include "example/qwen3/checkpoint_loader.h" +#include "example/qwen3/config.h" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/training/model_provider_registry.h" + +namespace qwen3 { +namespace { + +[[maybe_unused]] const bool registered = [] { + infini_train::training::ModelProviderRegistry::Instance().Register("qwen3", [](const std::string &weights) { + if (!weights.empty()) { + return LoadFromLLMC(weights); + } + auto config = Qwen3Config(); + SanitizeQwen3Config(config); + return std::make_shared(config); + }); + return true; +}(); + +} // namespace +} // namespace qwen3 diff --git a/example/qwen3/train.flags b/example/qwen3/train.flags new file mode 100644 index 000000000..bc70fd599 --- /dev/null +++ b/example/qwen3/train.flags @@ -0,0 +1,15 @@ +# Qwen3 training defaults. CLI options after --flagfile override these values. +--model=qwen3 +--optimizer=adam +--learning_rate=0.00001 +--batch_size=4 +--sequence_length=64 +--total_batch_size=256 +--num_iteration=10 +--dtype=float32 +--zero_stage=0 +--lr_decay_style=constant +--lora_target_modules=c_attn,c_proj,c_fc,c_fc2 +--load_optimizer_state=true +--save_optimizer_state=true +--profile_name=qwen3 diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h index 25939bdc2..bf68a8901 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h +++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h @@ -5,6 +5,7 @@ #include #include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/training/training.h" namespace infini_train { class Tensor; @@ -30,7 +31,8 @@ struct StageInfo { class PipelineParallel : public Module { public: PipelineParallel(const std::shared_ptr module, int num_stages, int num_micro_batches, - const std::vector> &recv_shape, int rank, Device device, int vpp); + const std::vector> &recv_shape, int rank, Device device, int vpp, + training::ForwardStepFunction forward_step = {}); float TrainStep(const std::vector> &input, const std::vector> &target, const std::shared_ptr &optimizer, @@ -44,7 +46,7 @@ class PipelineParallel : public Module { void BuildPipelineStage(const std::vector> &recv_shape, Device device, std::vector> &&chunks); - void SetupSchedule(int num_micro_batches); + void SetupSchedule(int num_micro_batches, training::ForwardStepFunction forward_step); int num_stages_ = -1; int rank_ = -1; diff --git a/infini_train/include/nn/parallel/pp/pipeline_schedule.h b/infini_train/include/nn/parallel/pp/pipeline_schedule.h index cae190f82..c8252db2e 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_schedule.h +++ b/infini_train/include/nn/parallel/pp/pipeline_schedule.h @@ -4,6 +4,7 @@ #include #include "infini_train/include/datatype.h" +#include "infini_train/include/training/training.h" namespace infini_train { class Tensor; @@ -19,8 +20,9 @@ class PipelineStage; class PipelineSchedule { public: - PipelineSchedule(std::shared_ptr stage, int num_stages, int num_micro_batches) - : stage_(std::move(stage)), num_micro_batches_(num_micro_batches) {} + PipelineSchedule(std::shared_ptr stage, int num_stages, int num_micro_batches, + training::ForwardStepFunction forward_step = {}) + : stage_(std::move(stage)), num_micro_batches_(num_micro_batches), forward_step_(std::move(forward_step)) {} virtual ~PipelineSchedule() = default; @@ -37,6 +39,7 @@ class PipelineSchedule { protected: int num_micro_batches_ = -1; std::shared_ptr stage_ = nullptr; + training::ForwardStepFunction forward_step_; }; class PipelineParallelScheduler { diff --git a/infini_train/include/training/arguments.h b/infini_train/include/training/arguments.h new file mode 100644 index 000000000..f611d0515 --- /dev/null +++ b/infini_train/include/training/arguments.h @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infini_train/include/training/training.h" + +namespace infini_train::training { + +// Optional gflags adapter, provided by InfiniTrain::training_cli. The training +// runtime itself accepts typed options and does not register command-line flags. +TrainingOptions TrainingOptionsFromFlags(); + +// Parse gflags with strict flag-file validation. A relative default file is +// resolved beside the executable. Missing defaults may be replaced by an +// explicit --flagfile. Invalid input exits with status 1, like gflags itself. +// Call once during startup, before worker threads are created. +void ParseTrainingFlags(int *argc, char ***argv, const std::filesystem::path &default_flagfile = {}); + +} // namespace infini_train::training diff --git a/infini_train/include/training/model_provider_registry.h b/infini_train/include/training/model_provider_registry.h new file mode 100644 index 000000000..f5b598875 --- /dev/null +++ b/infini_train/include/training/model_provider_registry.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/training/training.h" + +namespace infini_train::training { + +// A factory owns model-specific configuration and checkpoint import. Binding +// a factory must not construct a model: Pretrain invokes it after rank setup. +using ModelFactory = std::function(const std::string &weights)>; + +class ModelProviderRegistry { +public: + static ModelProviderRegistry &Instance(); + + // Register during startup, before concurrent lookups/training begin. + // Each model/alias must have exactly one nonempty factory. + void Register(const std::string &name, ModelFactory factory); + + // Return an independently owned callback; invocation is deferred to Pretrain. + ModelProvider Resolve(const std::string &name, const std::string &weights) const; + + // Stable ordering for CLI help and unknown-model diagnostics. + std::string Names() const; + +private: + std::map factories_; +}; + +} // namespace infini_train::training diff --git a/infini_train/include/training/training.h b/infini_train/include/training/training.h new file mode 100644 index 000000000..f945d5a0c --- /dev/null +++ b/infini_train/include/training/training.h @@ -0,0 +1,106 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +// Public training contracts depend only on forward declarations. In particular, +// including them from PP must not pull TransformerModel back into the scheduler. +namespace infini_train { +class Dataset; +class Device; +class Tensor; +namespace nn { +class Module; +class TransformerModel; +struct TransformerConfig; +} // namespace nn +} // namespace infini_train + +namespace infini_train::training { + +using TensorList = std::vector>; +using LossFunction = std::function(const TensorList &)>; + +// Task forward returns activations and a deferred, microbatch-mean loss. +// The training executor owns loss scaling, backward and gradient synchronization. +struct ForwardStepResult { + TensorList output; + LossFunction loss_func; +}; + +// Configured step used by both ordinary training and pipeline execution. Under +// PP, inputs are tokens on the first chunk and activations on later chunks; +// target is supplied only to the last chunk, whose loss callback is evaluated. +using ForwardStepFunction + = std::function &target)>; + +// Parsed once before worker threads start; providers must not mutate CLI flags. +struct TrainingOptions { + uint32_t batch_size = 4; + uint32_t sequence_length = 64; + uint32_t total_batch_size = 256; + uint32_t num_iteration = 10; + double learning_rate = 1e-5; + int32_t zero_stage = 0; + double min_lr = 0.0; + std::string lr_decay_style = "constant"; + int64_t lr_warmup_iters = 0; + double lr_warmup_init = 0.0; + int64_t lr_decay_iters = 0; + uint32_t val_loss_every = 0; + uint32_t sample_every = 0; + bool overfit_single_batch = true; + std::string device = "cuda"; + int32_t nthread_per_process = 1; + uint32_t tensor_parallel = 1; + bool sequence_parallel = false; + uint32_t pipeline_parallel = 1; + uint32_t virtual_pipeline_parallel = 1; + std::string dtype = "float32"; + std::string precision_check; + uint32_t save_interval = 0; + std::string load = ""; + std::string save = ""; + uint32_t max_checkpoint_keep = 3; + bool load_optimizer_state = true; + bool save_optimizer_state = true; + int32_t lora_rank = 0; + double lora_alpha = 16.0; + std::string lora_target_modules = "c_attn,c_proj,c_fc,c_fc2"; + std::string lora_save_path = ""; + std::string lora_load_path = ""; + std::string optimizer = "adam"; + std::string profile_name = "pretrain"; + // Some backends require synchronization after model upload and before teardown. + bool synchronize_device = false; +}; + +struct DatasetSplits { + std::shared_ptr train; + std::shared_ptr valid; +}; + +// Called once per rank, after TP/PP thread-local state and process groups exist. +// The provider returns the local TransformerModel on CPU, without DDP/PP wrappers. +// TransformerModel constructs all local VPP chunks using the existing global layout. +using ModelProvider = std::function()>; +// The argument is the total requested number of global training samples (not tokens). +// A finite dataset may be returned: the trainer cycles it and restores its position. +using DatasetProvider = std::function; +// Called for every microbatch (and every local chunk under PP). Only the last +// chunk's loss callback is evaluated. The config is from the loaded model. +using ForwardStep = std::function &, + const nn::TransformerConfig &)>; +// Optional reporting/sampling hook, called on the last global rank after each step. +using AfterStep = std::function; + +// Owns initialization, model wrapping, optimizer/scheduler, data iterators, +// checkpoint resume/save and forward/backward scheduling. Providers own task logic. +void Pretrain(const TrainingOptions &options, const DatasetProvider &dataset_provider, + const ModelProvider &model_provider, const ForwardStep &forward_step, const AfterStep &after_step = {}); + +} // namespace infini_train::training diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc index c0369cdeb..0e1c7d6fd 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -22,8 +22,9 @@ void PipelineParallel::BuildPipelineStage(const std::vector pipeline_stage_ = std::make_shared(rank_, num_stages_, recv_shape, device, std::move(chunks)); } -void PipelineParallel::SetupSchedule(int num_micro_batches) { - schedule_ = std::make_shared(pipeline_stage_, num_stages_, num_micro_batches); +void PipelineParallel::SetupSchedule(int num_micro_batches, training::ForwardStepFunction forward_step) { + schedule_ + = std::make_shared(pipeline_stage_, num_stages_, num_micro_batches, std::move(forward_step)); } float PipelineParallel::TrainStep(const std::vector> &input, @@ -78,7 +79,7 @@ StageInfo PipelineParallel::GetStageInfo(int total_layers, int pp_size, int rank PipelineParallel::PipelineParallel(const std::shared_ptr module, int num_stages, int num_micro_batches, const std::vector> &recv_shape, int pp_rank, Device device, - int chunk_size) + int chunk_size, training::ForwardStepFunction forward_step) : num_stages_(num_stages), rank_(pp_rank) { modules_[kModuleName] = std::move(module); @@ -100,7 +101,7 @@ PipelineParallel::PipelineParallel(const std::shared_ptr module, int num BuildPipelineStage(recv_shape, device, std::move(chunks)); - SetupSchedule(num_micro_batches); + SetupSchedule(num_micro_batches, std::move(forward_step)); } std::vector> *PipelineParallel::mutable_chunks() { return pipeline_stage_->mutable_chunks(); } diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 0ea2a4710..608446d9c 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -211,6 +211,7 @@ float PipelineSchedule::StepMicroBatches(const std::vector>>> activations( vpp_size, std::vector>>(n)); + std::vector loss_functions(n); std::vector> no_sync_guards; no_sync_guards.reserve(stage_->chunks().size()); @@ -239,7 +240,17 @@ float PipelineSchedule::StepMicroBatches(const std::vectorForwardOneChunk(inputs, task.local_chunk_idx); + if (forward_step_) { + auto result = forward_step_(*stage_->chunks().at(task.local_chunk_idx), inputs, + task.is_last_chunk ? microbatch_targets[mb] : nullptr); + activations[task.local_chunk_idx][mb] = std::move(result.output); + if (task.is_last_chunk) { + CHECK(result.loss_func) << "forward_step must return a loss callback on the last chunk"; + loss_functions[mb] = std::move(result.loss_func); + } + } else { + activations[task.local_chunk_idx][mb] = stage_->ForwardOneChunk(inputs, task.local_chunk_idx); + } if (!task.is_last_chunk) { if (stage_->IsLastStage()) { @@ -259,9 +270,14 @@ float PipelineSchedule::StepMicroBatches(const std::vectordevice().type(), dtype); - auto target_on_device = target->To(activations[task.local_chunk_idx][mb][0]->GetDevice()); - loss = (*loss_fn)( - {activations[task.local_chunk_idx][mb][0], std::make_shared(target_on_device)})[0]; + if (forward_step_) { + loss = loss_functions[mb](activations[task.local_chunk_idx][mb]); + loss_functions[mb] = {}; + } else { + auto target_on_device = target->To(activations[task.local_chunk_idx][mb][0]->GetDevice()); + loss = (*loss_fn)( + {activations[task.local_chunk_idx][mb][0], std::make_shared(target_on_device)})[0]; + } loss = loss / n; } loss->Backward(); diff --git a/infini_train/src/training/arguments.cc b/infini_train/src/training/arguments.cc new file mode 100644 index 000000000..8831ea7f7 --- /dev/null +++ b/infini_train/src/training/arguments.cc @@ -0,0 +1,323 @@ +#include "infini_train/include/training/arguments.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" + +DECLARE_string(flagfile); + +namespace { +// TrainingOptions owns the framework defaults for both C++ and CLI callers. +const infini_train::training::TrainingOptions kTrainingDefaults; +} // namespace + +DEFINE_uint32(batch_size, kTrainingDefaults.batch_size, "batch size, in units of #batch dimensions"); +DEFINE_uint32(sequence_length, kTrainingDefaults.sequence_length, "sequence length"); +DEFINE_uint32(total_batch_size, kTrainingDefaults.total_batch_size, "total desired batch size, in units of #tokens"); +DEFINE_uint32(num_iteration, kTrainingDefaults.num_iteration, "number of iterations to run"); +DEFINE_double(learning_rate, kTrainingDefaults.learning_rate, "Peak learning rate."); +DEFINE_int32(zero_stage, kTrainingDefaults.zero_stage, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); +DEFINE_double(min_lr, kTrainingDefaults.min_lr, "Minimum learning rate."); +DEFINE_string(lr_decay_style, kTrainingDefaults.lr_decay_style, + "LR decay style: none|constant|linear|cosine|inverse-square-root"); +DEFINE_int64(lr_warmup_iters, kTrainingDefaults.lr_warmup_iters, "Number of linear warmup iterations."); +DEFINE_double(lr_warmup_init, kTrainingDefaults.lr_warmup_init, "Initial learning rate at the start of warmup."); +DEFINE_int64(lr_decay_iters, kTrainingDefaults.lr_decay_iters, + "Number of iterations to decay LR over (0 = num_iteration)."); +DEFINE_uint32(val_loss_every, kTrainingDefaults.val_loss_every, "every how many steps to evaluate val loss?"); +DEFINE_uint32(sample_every, kTrainingDefaults.sample_every, "how often to sample from the model?"); +DEFINE_bool(overfit_single_batch, kTrainingDefaults.overfit_single_batch, "overfit just one batch of data"); +DEFINE_string(device, kTrainingDefaults.device, + "device type (cpu/cuda/privateuse1/), useless if using parallel " + "training mode"); +DEFINE_int32(nthread_per_process, kTrainingDefaults.nthread_per_process, + "Number of threads to use for each process. " + "When set > 1, enables data parallelism on the specified accelerator devices."); +DEFINE_uint32(tensor_parallel, kTrainingDefaults.tensor_parallel, "Tensor Parallel world size"); +DEFINE_bool(sequence_parallel, kTrainingDefaults.sequence_parallel, "Whether to enable Sequence Parallel"); +DEFINE_uint32(pipeline_parallel, kTrainingDefaults.pipeline_parallel, + "Pipeline Parallel world size, specified the number of PP stages."); +DEFINE_uint32(virtual_pipeline_parallel, kTrainingDefaults.virtual_pipeline_parallel, "Number of chunks in PP stage."); +DEFINE_string(dtype, kTrainingDefaults.dtype, "precision used in training (float32/bfloat16)"); +DEFINE_uint32(save_interval, kTrainingDefaults.save_interval, "save checkpoint every N steps; 0 disables saving"); +DEFINE_string(load, kTrainingDefaults.load, "checkpoint directory to resume from"); +DEFINE_string(save, kTrainingDefaults.save, "root directory used to store checkpoints"); +DEFINE_uint32(max_checkpoint_keep, kTrainingDefaults.max_checkpoint_keep, "max number of checkpoint steps to keep"); +DEFINE_bool(load_optimizer_state, kTrainingDefaults.load_optimizer_state, + "whether optimizer state is restored from checkpoints"); +DEFINE_bool(save_optimizer_state, kTrainingDefaults.save_optimizer_state, + "whether optimizer state is persisted in checkpoints"); +DEFINE_int32(lora_rank, kTrainingDefaults.lora_rank, "LoRA rank (0 = disabled)"); +DEFINE_double(lora_alpha, kTrainingDefaults.lora_alpha, "LoRA alpha scaling factor"); +DEFINE_string(lora_target_modules, kTrainingDefaults.lora_target_modules, "LoRA target modules (comma-separated)"); +DEFINE_string(lora_save_path, kTrainingDefaults.lora_save_path, "Path to save LoRA weights after training"); +DEFINE_string(lora_load_path, kTrainingDefaults.lora_load_path, "Path to load LoRA weights from"); +DEFINE_string(optimizer, kTrainingDefaults.optimizer, "Optimizer: sgd|adam"); +DEFINE_string(profile_name, kTrainingDefaults.profile_name, "Prefix for profiler output files"); +DEFINE_string(precision_check, kTrainingDefaults.precision_check, + "precision check: level=N,format=simple|table,output_md5=true|false,output_path=PATH,baseline=PATH"); + +namespace infini_train::training { +namespace { + +namespace fs = std::filesystem; + +std::string Trim(const std::string &text) { + const auto begin = text.find_first_not_of(" \t\r\n"); + return begin == std::string::npos ? "" : text.substr(begin, text.find_last_not_of(" \t\r\n") - begin + 1); +} + +std::vector FileList(const std::string &text) { + std::vector files; + size_t begin = 0; + while (begin < text.size()) { + const auto end = text.find(',', begin); + auto file = Trim(text.substr(begin, end == std::string::npos ? end : end - begin)); + if (file.empty()) { + throw std::runtime_error("Empty path in --flagfile list"); + } + files.push_back(std::move(file)); + if (end == std::string::npos) { + break; + } + begin = end + 1; + if (begin == text.size()) { + throw std::runtime_error("Empty path in --flagfile list"); + } + } + return files; +} + +// gflags runs validators while holding its registry mutex. Snapshot flag types +// before parsing: querying gflags from inside the validator would deadlock. +std::map &FlagTypes() { + static std::map types; + return types; +} + +std::string FlagType(const std::string &name) { + const auto &types = FlagTypes(); + if (auto it = types.find(name); it != types.end()) { + return it->second; + } + if (name.starts_with("no")) { + auto it = types.find(name.substr(2)); + if (it != types.end() && it->second == "bool") { + return "bool"; + } + } + return ""; +} + +class FlagfileValidator { +public: + void Validate(const std::string &files) { + for (const auto &file : FileList(files)) { ValidateFile(file); } + } + + void ValidateFile(const fs::path &path) { + if (!fs::is_regular_file(path)) { + throw std::runtime_error("Cannot read flag file: " + path.string()); + } + const auto canonical = fs::canonical(path); + if (active_.size() >= 64 || !active_.insert(canonical).second) { + throw std::runtime_error("Cyclic or excessively nested flag file: " + path.string()); + } + std::ifstream input(path); + if (!input) { + throw std::runtime_error("Cannot read flag file: " + path.string()); + } + std::string line; + size_t number = 0; + while (std::getline(input, line)) { + ++number; + line = Trim(line); + if (line.empty() || line.starts_with('#')) { + continue; + } + const auto location = path.string() + ":" + std::to_string(number) + ": "; + if (!line.starts_with("--")) { + throw std::runtime_error(location + "Expected --name=value (filename sections are not supported)"); + } + const auto equal = line.find('='); + const auto name = line.substr(2, equal == std::string::npos ? equal : equal - 2); + const auto type = FlagType(name); + if (type.empty()) { + throw std::runtime_error(location + "Unknown flag: --" + name); + } + if (equal == std::string::npos && type != "bool") { + throw std::runtime_error(location + "Missing value for --" + name + "; use --name=value"); + } + if (name == "flagfile") { + Validate(line.substr(equal + 1)); + } else if (name == "fromenv" || name == "tryfromenv") { + for (const auto &flag : FileList(line.substr(equal + 1))) { + if (flag == "flagfile") { + throw std::runtime_error(location + "Use an explicit --flagfile instead of loading it fromenv"); + } + } + } + } + if (input.bad()) { + throw std::runtime_error("Failed reading flag file: " + path.string()); + } + active_.erase(canonical); + } + +private: + std::set active_; +}; + +bool ValidateFlagFiles(const char *, const std::string &files) { + try { + FlagfileValidator().Validate(files); + return true; + } catch (const std::exception &error) { + std::cerr << "ERROR: " << error.what() << '\n'; + return false; + } +} + +bool HasExplicitFlagfile(int argc, char *const argv[]) { + for (int i = 1; i < argc; ++i) { + std::string argument = argv[i]; + if (argument == "--") { + break; + } + if (!argument.starts_with('-')) { + continue; + } + argument.erase(0, argument.starts_with("--") ? 2 : 1); + const auto equal = argument.find('='); + const auto name = argument.substr(0, equal); + if (name == "flagfile") { + return true; + } + // A string argument's value may itself look like --flagfile; do not + // mistake that value for a separate command-line option. + const auto type = FlagType(name); + if (equal == std::string::npos && !type.empty() && type != "bool") { + ++i; + } + } + return false; +} + +fs::path ExecutableDirectory(const char *argv0) { +#ifdef __linux__ + std::error_code error; + const auto executable = fs::read_symlink("/proc/self/exe", error); + if (!error) { + return executable.parent_path(); + } +#endif + fs::path executable_path(argv0); + if (executable_path.has_parent_path()) { + return fs::canonical(executable_path).parent_path(); + } + const char *search_path = std::getenv("PATH"); + if (search_path) { + const std::string paths(search_path); + size_t begin = 0; + do { + const auto end = paths.find(':', begin); + const auto dir = paths.substr(begin, end == std::string::npos ? end : end - begin); + const auto candidate = fs::path(dir.empty() ? "." : dir) / executable_path; + if (fs::is_regular_file(candidate)) { + return fs::canonical(candidate).parent_path(); + } + if (end == std::string::npos) { + break; + } + begin = end + 1; + } while (begin <= paths.size()); + } + throw std::runtime_error("Cannot locate executable: " + executable_path.string()); +} + +} // namespace + +TrainingOptions TrainingOptionsFromFlags() { + TrainingOptions options; + options.profile_name = FLAGS_profile_name; + options.precision_check = FLAGS_precision_check; + options.batch_size = FLAGS_batch_size; + options.sequence_length = FLAGS_sequence_length; + options.total_batch_size = FLAGS_total_batch_size; + options.num_iteration = FLAGS_num_iteration; + options.learning_rate = FLAGS_learning_rate; + options.zero_stage = FLAGS_zero_stage; + options.min_lr = FLAGS_min_lr; + options.lr_decay_style = FLAGS_lr_decay_style; + options.lr_warmup_iters = FLAGS_lr_warmup_iters; + options.lr_warmup_init = FLAGS_lr_warmup_init; + options.lr_decay_iters = FLAGS_lr_decay_iters; + options.val_loss_every = FLAGS_val_loss_every; + options.sample_every = FLAGS_sample_every; + options.overfit_single_batch = FLAGS_overfit_single_batch; + options.device = FLAGS_device; + options.nthread_per_process = FLAGS_nthread_per_process; + options.tensor_parallel = FLAGS_tensor_parallel; + options.sequence_parallel = FLAGS_sequence_parallel; + options.pipeline_parallel = FLAGS_pipeline_parallel; + options.virtual_pipeline_parallel = FLAGS_virtual_pipeline_parallel; + options.dtype = FLAGS_dtype; + options.save_interval = FLAGS_save_interval; + options.load = FLAGS_load; + options.save = FLAGS_save; + options.max_checkpoint_keep = FLAGS_max_checkpoint_keep; + options.load_optimizer_state = FLAGS_load_optimizer_state; + options.save_optimizer_state = FLAGS_save_optimizer_state; + options.lora_rank = FLAGS_lora_rank; + options.lora_alpha = FLAGS_lora_alpha; + options.lora_target_modules = FLAGS_lora_target_modules; + options.lora_save_path = FLAGS_lora_save_path; + options.lora_load_path = FLAGS_lora_load_path; + options.optimizer = FLAGS_optimizer; + return options; +} + +void ParseTrainingFlags(int *argc, char ***argv, const fs::path &default_flagfile) { + try { + std::vector flags; + gflags::GetAllFlags(&flags); + auto &types = FlagTypes(); + types.clear(); + for (const auto &flag : flags) { types.emplace(flag.name, flag.type); } + if (!gflags::RegisterFlagValidator(&FLAGS_flagfile, &ValidateFlagFiles)) { + throw std::runtime_error("Another validator is already registered for --flagfile"); + } + if (!default_flagfile.empty()) { + const auto path = default_flagfile.is_absolute() ? default_flagfile + : ExecutableDirectory((*argv)[0]) / default_flagfile; + if (fs::exists(path)) { + // This is one resolved filesystem path, not a comma-separated + // --flagfile list; executable directories can contain commas. + FlagfileValidator().ValidateFile(path); + if (!gflags::ReadFromFlagsFile(path.string(), (*argv)[0], false)) { + std::exit(EXIT_FAILURE); + } + // Nested includes have already been applied in file order. + // Do not let ParseCommandLineFlags replay the final include. + FLAGS_flagfile.clear(); + } else if (!HasExplicitFlagfile(*argc, *argv)) { + throw std::runtime_error("Default training flag file not found: " + path.string() + + ". Copy configs/ beside the executable or pass --flagfile explicitly."); + } + } + gflags::ParseCommandLineFlags(argc, argv, true); + } catch (const std::exception &error) { + std::cerr << "ERROR: " << error.what() << '\n'; + std::exit(EXIT_FAILURE); + } +} + +} // namespace infini_train::training diff --git a/infini_train/src/training/model_provider_registry.cc b/infini_train/src/training/model_provider_registry.cc new file mode 100644 index 000000000..575d9dab1 --- /dev/null +++ b/infini_train/src/training/model_provider_registry.cc @@ -0,0 +1,38 @@ +#include "infini_train/include/training/model_provider_registry.h" + +#include + +#include "glog/logging.h" + +namespace infini_train::training { + +ModelProviderRegistry &ModelProviderRegistry::Instance() { + static ModelProviderRegistry registry; + return registry; +} + +void ModelProviderRegistry::Register(const std::string &name, ModelFactory factory) { + CHECK(!name.empty()) << "Model provider name must not be empty"; + CHECK(factory) << "Model factory must not be empty: " << name; + CHECK(!factories_.contains(name)) << "Model provider already registered: " << name; + factories_.emplace(name, std::move(factory)); +} + +ModelProvider ModelProviderRegistry::Resolve(const std::string &name, const std::string &weights) const { + const auto it = factories_.find(name); + CHECK(it != factories_.end()) << "Unknown model: " << name << ". Available models: " << Names(); + return [factory = it->second, weights] { return factory(weights); }; +} + +std::string ModelProviderRegistry::Names() const { + std::string names; + for (const auto &[name, factory] : factories_) { + if (!names.empty()) { + names += ", "; + } + names += name; + } + return names; +} + +} // namespace infini_train::training diff --git a/example/llama3/main.cc b/infini_train/src/training/training.cc similarity index 53% rename from example/llama3/main.cc rename to infini_train/src/training/training.cc index a9f405fa5..1d9fcd829 100644 --- a/example/llama3/main.cc +++ b/infini_train/src/training/training.cc @@ -1,3 +1,6 @@ +#include "infini_train/include/training/training.h" + +#include #include #include #include @@ -6,15 +9,8 @@ #include #include -#include "gflags/gflags.h" #include "glog/logging.h" -// Out-of-tree builds inject the selected provider's declaration without adding -// a vendor dependency to the upstream example. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#include INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER -#endif - #include "infini_train/include/autocast.h" #include "infini_train/include/checkpoint/checkpoint.h" #include "infini_train/include/checkpoint/checkpoint_manager.h" @@ -23,7 +19,6 @@ #include "infini_train/include/device.h" #include "infini_train/include/lr_scheduler.h" #include "infini_train/include/nn/lora/lora_utils.h" -#include "infini_train/include/nn/modules/loss.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/modules/transformer/transformer.h" #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" @@ -44,122 +39,93 @@ #include "infini_train/include/profiler.h" #endif -#include "example/common/tiny_shakespeare_dataset.h" -#include "example/common/tokenizer.h" -#include "example/llama3/checkpoint_loader.h" -#include "example/llama3/config.h" - -// TODO(jym): Reorganize CLI flags into categories for better readability and maintainability. -// I/O -DEFINE_string(input_bin, "", "input .bin to train on"); -DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); -DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); -// model bin file is downloaded and processed using the script at -// https://github.com/karpathy/llm.c/blob/master/train_llama3.py -DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); -DEFINE_string(model, "llama3", "meta-llama/Meta-Llama-3.1-8B"); -// token layout for each step of the optimization -DEFINE_uint32(batch_size, 4, "batch size, in units of #batch dimensions"); -DEFINE_uint32(sequence_length, 64, "sequence length"); -DEFINE_uint32(total_batch_size, 256, "total desired batch size, in units of #tokens"); -// workload (number of steps) -DEFINE_uint32(num_iteration, 10, "number of iterations to run"); -DEFINE_uint32(freq_generate_txt, 10, "frequency of text generation"); -DEFINE_uint32(text_length, 64, "the length of the generated text"); -// optimization -DEFINE_double(learning_rate, 1e-5, "Peak learning rate."); -DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); -// lr scheduler -DEFINE_double(min_lr, 0.0, "Minimum learning rate."); -DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); -DEFINE_int64(lr_warmup_iters, 0, "Number of linear warmup iterations."); -DEFINE_double(lr_warmup_init, 0.0, "Initial learning rate at the start of warmup."); -DEFINE_int64(lr_decay_iters, 0, "Number of iterations to decay LR over (0 = num_iteration)."); -// evaluation -DEFINE_uint32(val_loss_every, 0, "every how many steps to evaluate val loss?"); -DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); -// debugging -DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); -// memory management -DEFINE_string(device, "cuda", - "device type (cpu/cuda/privateuse1/), useless if using parallel " - "training mode"); -// parallel -DEFINE_int32(nthread_per_process, 1, - "Number of threads to use for each process. " - "When set > 1, enables data parallelism on the specified accelerator devices."); -DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size"); -DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel"); -DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages."); -DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage."); -// precision -DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)"); -DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving"); -DEFINE_string(load, "", "checkpoint directory to resume from"); -DEFINE_string(save, "", "root directory used to store checkpoints"); -DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); -DEFINE_bool(load_optimizer_state, true, "whether optimizer state is restored from checkpoints"); -DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints"); - -// precision check -DEFINE_string( - precision_check, "", - "precision check config: level=N,format=simple|table,output_md5=true|false,output_path=PATH,baseline=PATH"); -// LoRA parameters -DEFINE_int32(lora_rank, 0, "LoRA rank (0 = disabled)"); -DEFINE_double(lora_alpha, 16.0, "LoRA alpha scaling factor"); -DEFINE_string(lora_target_modules, "c_attn,c_proj,c_fc,c_fc2", "LoRA target modules (comma-separated)"); -DEFINE_string(lora_save_path, "", "Path to save LoRA weights after training"); -DEFINE_string(lora_load_path, "", "Path to load LoRA weights from"); - -using namespace infini_train; - +namespace infini_train::training { namespace { -// validation -const std::unordered_set kSupportedModels = {"llama3"}; -constexpr char kDtypeFP32[] = "float32"; -constexpr char kDtypeBF16[] = "bfloat16"; -const std::unordered_set kSupportedLRDecayStyles - = {"none", "constant", "linear", "cosine", "inverse-square-root"}; - -} // namespace -DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); -DEFINE_validator(device, [](const char *, const std::string &value) { return Device::ParseType(value).has_value(); }); -DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); -DEFINE_validator(lr_decay_style, - [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); +void ValidateOptions(const TrainingOptions &options) { + CHECK(Device::ParseType(options.device).has_value()) << "Unsupported device: " << options.device; + CHECK(options.dtype == "float32" || options.dtype == "bfloat16") << "Unsupported dtype: " << options.dtype; + CHECK(options.optimizer == "sgd" || options.optimizer == "adam") << "Unsupported optimizer: " << options.optimizer; + CHECK_GT(options.batch_size, 0); + CHECK_GT(options.sequence_length, 0); + CHECK_GT(options.total_batch_size, 0); + CHECK_GT(options.nthread_per_process, 0); + CHECK_GT(options.tensor_parallel, 0); + CHECK_GT(options.pipeline_parallel, 0); + CHECK_GT(options.virtual_pipeline_parallel, 0); + CHECK_GE(options.zero_stage, 0); + CHECK_LE(options.zero_stage, 3); + CHECK_EQ(options.save.empty(), options.save_interval == 0) << "--save and --save_interval must be set together"; +} -void Train(const nn::parallel::Rank &rank) { +std::shared_ptr SetupOptimizer(const TrainingOptions &options, const std::shared_ptr &model, + int ddp_world_size, int ddp_rank, int pp_world_size) { using namespace nn::parallel; + auto optimizer_creator = options.optimizer == "sgd" ? optimizers::SGD::CreateNamed(options.learning_rate) + : optimizers::Adam::CreateNamed(options.learning_rate); + std::shared_ptr optimizer = nullptr; - { - if (rank.IsLastRank()) { - if (!FLAGS_save.empty() && FLAGS_save_interval == 0) { - LOG(FATAL) << "Invalid configuration: --save is set ('" << FLAGS_save - << "'), but --save_interval is 0. " - << "They must be set together."; - } + std::vector> params_to_optimize; + if (options.lora_rank > 0) { + params_to_optimize = nn::lora::GetLoRAParameters(model); + LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters"; + } else { + params_to_optimize = model->Parameters(); + LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters"; + } + std::unordered_set params_to_optimize_set; + params_to_optimize_set.reserve(params_to_optimize.size()); + for (const auto ¶m : params_to_optimize) { params_to_optimize_set.insert(param.get()); } - if (FLAGS_save.empty() && FLAGS_save_interval > 0) { - LOG(FATAL) << "Invalid configuration: --save_interval is set to " << FLAGS_save_interval - << ", but --save is empty. " - << "They must be set together."; - } + NamedParameterList named_parameters; + for (const auto &[name, param] : model->NamedParameters()) { + if (params_to_optimize_set.contains(param.get())) { + named_parameters.emplace_back(name, param); } } + CHECK_EQ(named_parameters.size(), params_to_optimize.size()); + + if (options.zero_stage >= 1) { + auto model_chunks = (pp_world_size > 1) + ? *(dynamic_cast(model.get())->mutable_chunks()) + : std::vector>{model}; + optimizer = std::make_shared(optimizer_creator, named_parameters, + model_chunks, ddp_world_size, ddp_rank); + } else { + optimizer = optimizer_creator(named_parameters); + } + + return optimizer; +} + +std::shared_ptr SetupLRScheduler(const TrainingOptions &options, + const std::shared_ptr &optimizer) { + const int64_t lr_decay_iters = options.lr_decay_iters > 0 ? options.lr_decay_iters : options.num_iteration; + TrainingLRSchedulerConfig sched_config; + sched_config.lr = static_cast(options.learning_rate); + sched_config.min_lr = static_cast(options.min_lr); + sched_config.lr_decay_style = options.lr_decay_style; + sched_config.lr_decay_iters = lr_decay_iters; + sched_config.lr_warmup_iters = options.lr_warmup_iters; + sched_config.lr_warmup_init = static_cast(options.lr_warmup_init); + return CreateLRScheduler(optimizer, sched_config); +} + +void Train(const TrainingOptions &options, const nn::parallel::Rank &rank, const DatasetProvider &dataset_provider, + const ModelProvider &model_provider, const ForwardStep &forward_step, const AfterStep &after_step) { + using namespace nn::parallel; // select the device Device device; - const auto device_type = Device::ParseType(FLAGS_device).value(); + const auto device_type = Device::ParseType(options.device).value(); int ddp_world_size = global::GetDataParallelSize(); int tp_world_size = global::GetTensorParallelSize(); int sp_world_size = global::GetSequenceParallelEnabled() ? tp_world_size : 1; int pp_world_size = global::GetPipelineParallelSize(); - if (FLAGS_sequence_parallel) { - CHECK_EQ(FLAGS_sequence_length % tp_world_size, 0) + if (options.sequence_parallel) { + CHECK_EQ(options.sequence_length % tp_world_size, 0) << "sequence_length must be divisible by tp_world_size when SP is enabled (pad later if needed)."; } @@ -205,50 +171,46 @@ void Train(const nn::parallel::Rank &rank) { } // calculate gradient accumulation from the desired total batch size and the current run configuration - const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size; - CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0); - const auto grad_accum_steps = FLAGS_total_batch_size / tokens_per_fwdbwd; + const auto tokens_per_fwdbwd = static_cast(options.batch_size) * options.sequence_length * ddp_world_size; + CHECK_EQ(options.total_batch_size % tokens_per_fwdbwd, 0); + const auto grad_accum_steps = options.total_batch_size / tokens_per_fwdbwd; if (rank.IsMainRank()) { - LOG(INFO) << "total desired batch size: " << FLAGS_total_batch_size + LOG(INFO) << "total desired batch size: " << options.total_batch_size << " => calculated gradient accumulation steps: " << grad_accum_steps; } - // rng / reproducibility - // ManualSeed(42); - - nn::TransformerConfig model_config = llama3::LLaMA3Config(); - std::shared_ptr model = nullptr; - if (!FLAGS_llmc_filepath.empty()) { - model = llama3::LoadFromLLMC(FLAGS_llmc_filepath); - } else { - llama3::SanitizeLLaMA3Config(model_config); - model = std::make_shared(model_config); - } + auto base_model = model_provider(); + CHECK(base_model) << "model_provider returned a null model"; + const auto model_config = base_model->Config(); + CHECK_LE(options.sequence_length, model_config.block_size); + CHECK_GT(model_config.GetChunkSize(), 0) << "Every pipeline rank must own a transformer chunk"; + std::shared_ptr model = base_model; + ForwardStepFunction rank_forward_step + = [&](nn::Module &module, const TensorList &inputs, const std::shared_ptr &target) { + return forward_step(module, inputs, target, model_config); + }; model->To(device); -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA requires model upload to finish before communication begins. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { + if (options.synchronize_device) { core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); } -#endif utils::PrecisionChecker::BuildNameMap(model.get()); // Apply LoRA using GetLoRAModel (in-place injection) - bool lora_enabled = FLAGS_lora_rank > 0; + bool lora_enabled = options.lora_rank > 0; if (lora_enabled) { - nn::lora::LoRAConfig lora_config{FLAGS_lora_rank, static_cast(FLAGS_lora_alpha), 0.0f, - nn::lora::ParseLoRATargetModules(FLAGS_lora_target_modules)}; + nn::lora::LoRAConfig lora_config{options.lora_rank, static_cast(options.lora_alpha), 0.0f, + nn::lora::ParseLoRATargetModules(options.lora_target_modules)}; // GetLoRAModel: in-place injection, modifies module tree directly model = nn::lora::GetLoRAModel(model, lora_config); // Load LoRA weights if specified - if (!FLAGS_lora_load_path.empty()) { - LOG(INFO) << "Loading LoRA weights from: " << FLAGS_lora_load_path; - nn::lora::LoadLoRAWeights(model, FLAGS_lora_load_path); + if (!options.lora_load_path.empty()) { + LOG(INFO) << "Loading LoRA weights from: " << options.lora_load_path; + nn::lora::LoadLoRAWeights(model, options.lora_load_path); } // Print LoRA summary @@ -258,26 +220,27 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; DataType dtype; - if (FLAGS_dtype == kDtypeFP32) { + if (options.dtype == "float32") { dtype = DataType::kFLOAT32; - } else if (FLAGS_dtype == kDtypeBF16) { + } else if (options.dtype == "bfloat16") { dtype = DataType::kBFLOAT16; } else { - LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported."; + LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << options.dtype << " not supported."; } - auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size); + const auto num_micro_batches = grad_accum_steps; if (pp_world_size > 1) { // NOTE(dcj): To ensure that the tensor shapes at the pipeline stage boundaries remain correct // when sequence parallelism (SP) is enabled, we need to divide by sp_world_size. auto shapes = std::vector>{ - {FLAGS_batch_size, FLAGS_sequence_length / sp_world_size, model_config.n_embd}}; + {options.batch_size, options.sequence_length / sp_world_size, model_config.n_embd}}; - model = std::make_shared(model, pp_world_size, num_micro_batches, shapes, - pp_rank, device, model_config.GetChunkSize()); + model + = std::make_shared(model, pp_world_size, num_micro_batches, shapes, pp_rank, + device, model_config.GetChunkSize(), rank_forward_step); if (ddp_world_size > 1) { - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; + auto ddp_config = DistributedDataParallelConfig{.zero_stage = options.zero_stage}; auto *mutable_chunks = dynamic_cast(model.get())->mutable_chunks(); for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) { (*mutable_chunks)[chunk_id] @@ -290,89 +253,35 @@ void Train(const nn::parallel::Rank &rank) { // Otherwise, DDP’s gradient hooks may be lost because new parameter tensors // are created during the conversion. - auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage}; + auto ddp_config = DistributedDataParallelConfig{.zero_stage = options.zero_stage}; model = std::make_shared(model, rank, ddp_config); } - const size_t train_loader_batch_size = pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size; - DistributedDataLoader train_loader(std::make_shared(FLAGS_input_bin, FLAGS_sequence_length), - train_loader_batch_size, ddp_rank, ddp_world_size); - - std::optional val_loader = std::nullopt; - if (!FLAGS_input_val_bin.empty()) { - val_loader = DistributedDataLoader( - std::make_shared(FLAGS_input_val_bin, FLAGS_sequence_length), FLAGS_batch_size, - ddp_rank, ddp_world_size); - } - - // - // main training loop - // - std::unique_ptr tokenizer = nullptr; - if (!FLAGS_tokenizer_bin.empty()) { - tokenizer = std::make_unique(FLAGS_tokenizer_bin); - } - - // TODO(dcj): support more complex optimizer later - // auto optimizer = optimizers::Adam(model->Parameters(), FLAGS_learning_rate); - auto optimizer_creator = optimizers::Adam::CreateNamed(FLAGS_learning_rate); - std::shared_ptr optimizer = nullptr; - - std::vector> params_to_optimize; - if (lora_enabled) { - params_to_optimize = nn::lora::GetLoRAParameters(model); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " LoRA parameters"; - } else { - params_to_optimize = model->Parameters(); - LOG(INFO) << "Optimizing " << params_to_optimize.size() << " model parameters"; - } - std::unordered_set params_to_optimize_set; - params_to_optimize_set.reserve(params_to_optimize.size()); - for (const auto ¶m : params_to_optimize) { params_to_optimize_set.insert(param.get()); } - - NamedParameterList named_parameters; - for (const auto &[name, param] : model->NamedParameters()) { - if (params_to_optimize_set.contains(param.get())) { - named_parameters.emplace_back(name, param); - } - } - CHECK_EQ(named_parameters.size(), params_to_optimize.size()); - - if (FLAGS_zero_stage >= 1) { - auto model_chunks = (pp_world_size > 1) - ? *(dynamic_cast(model.get())->mutable_chunks()) - : std::vector>{model}; - optimizer = std::make_shared(optimizer_creator, named_parameters, - model_chunks, ddp_world_size, ddp_rank); - } else { - optimizer = optimizer_creator(named_parameters); + const size_t train_loader_batch_size + = pp_world_size > 1 ? options.batch_size * num_micro_batches : options.batch_size; + const auto datasets = dataset_provider(static_cast(options.num_iteration) + * (options.total_batch_size / options.sequence_length)); + CHECK(datasets.train) << "dataset_provider must return a training dataset"; + DistributedDataLoader train_loader(datasets.train, train_loader_batch_size, ddp_rank, ddp_world_size); + std::optional val_loader; + if (datasets.valid) { + val_loader.emplace(datasets.valid, options.batch_size, ddp_rank, ddp_world_size); } - const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; - TrainingLRSchedulerConfig sched_config; - sched_config.lr = static_cast(FLAGS_learning_rate); - sched_config.min_lr = static_cast(FLAGS_min_lr); - sched_config.lr_decay_style = FLAGS_lr_decay_style; - sched_config.lr_decay_iters = lr_decay_iters; - sched_config.lr_warmup_iters = FLAGS_lr_warmup_iters; - sched_config.lr_warmup_init = static_cast(FLAGS_lr_warmup_init); - auto scheduler = CreateLRScheduler(optimizer, sched_config); + auto optimizer = SetupOptimizer(options, model, ddp_world_size, ddp_rank, pp_world_size); + auto scheduler = SetupLRScheduler(options, optimizer); auto train_iter = train_loader.begin(); - std::shared_ptr loss_fn - = (tp_world_size > 1) ? std::static_pointer_cast(std::make_shared()) - : std::static_pointer_cast(std::make_shared()); - loss_fn->To(device); LOG(INFO) << "Rank " << rank.GlobalRank() << ": start training"; auto impl = core::GetDeviceGuardImpl(device.type()); int start_step = 0; TrainerState state; - const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, + const auto resume_result = ResumeFromCheckpoint({.resume_root = options.load, .rank = rank, .model = model, - .optimizer = FLAGS_load_optimizer_state ? optimizer : nullptr, + .optimizer = options.load_optimizer_state ? optimizer : nullptr, .model_config = model_config, .state = state, .lr_scheduler = scheduler}); @@ -417,31 +326,31 @@ void Train(const nn::parallel::Rank &rank) { .tp_size = tp_world_size, .sp_size = sp_world_size, .pp_size = pp_world_size, - .checkpoint_root_dir = FLAGS_save, - .max_checkpoint_keep = FLAGS_max_checkpoint_keep, + .checkpoint_root_dir = options.save, + .max_checkpoint_keep = options.max_checkpoint_keep, .rank = rank, .model = *model, - .optimizer = FLAGS_save_optimizer_state ? optimizer.get() : nullptr, + .optimizer = options.save_optimizer_state ? optimizer.get() : nullptr, .lr_scheduler = scheduler.get(), }); }; - for (int step = start_step; step < FLAGS_num_iteration + 1; ++step) { + for (int64_t step = start_step; step <= options.num_iteration; ++step) { // Reset precision check counters at start of each iteration for file overwrite utils::PrecisionChecker::ResetCounters(); - const bool last_step = step == FLAGS_num_iteration; + const bool last_step = step == options.num_iteration; impl->ResetMemPoolHighWatermarks(device); const auto iter_start = std::chrono::high_resolution_clock::now(); // once in a while evaluate the validation dataset - if (FLAGS_val_loss_every > 0 && (step % FLAGS_val_loss_every == 0 || last_step) && val_loader.has_value()) { + if (options.val_loss_every > 0 && (step % options.val_loss_every == 0 || last_step) && val_loader.has_value()) { // TODO(dcj): implement this after model.eval() is supported } // once in a while perform model inference on the master process - if (FLAGS_sample_every > 0 && (step % FLAGS_sample_every == 0 || last_step)) { + if (options.sample_every > 0 && (step % options.sample_every == 0 || last_step)) { // TODO(dcj): implement this after model.eval() is supported } @@ -457,14 +366,14 @@ void Train(const nn::parallel::Rank &rank) { Profiler::Instance().SetTag("Step_" + std::to_string(step)); #endif - const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); + const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(options.learning_rate); float lossf = 0.0f; if (pp_world_size == 1) { // model->Train(); optimizer->ZeroGrad(); // if we are trying to overfit a single batch, we reset the loader here - if (FLAGS_overfit_single_batch) { + if (options.overfit_single_batch) { // train_loader.Reset(); } @@ -479,9 +388,10 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; // (bs, seq_len, vocab_size) - auto logits = (*model)({x, y})[0]; + auto result = rank_forward_step(*model, {x}, y); + CHECK(result.loss_func) << "forward_step must return a loss callback"; LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; - auto loss = (*loss_fn)({logits, y})[0]; + auto loss = result.loss_func(result.output); // FIXME(jym): verify gradient accumulation precision loss = loss / grad_accum_steps; @@ -512,7 +422,7 @@ void Train(const nn::parallel::Rank &rank) { x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); - lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + lossf = model->TrainStep({x}, {y}, optimizer, nullptr, dtype); if (scheduler) { scheduler->Step(); } @@ -526,30 +436,26 @@ void Train(const nn::parallel::Rank &rank) { const auto iter_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(iter_end - iter_start).count(); - const double tps = FLAGS_total_batch_size / (duration_us / 1e6); + const double tps = options.total_batch_size / (duration_us / 1e6); if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", - step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, + step + 1, options.num_iteration, lossf, current_lr, duration_us / 1e3f, tps, used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, pp_world_size); - if ((step + 1) % FLAGS_freq_generate_txt == 0) { - // FIXME(jym): to support PP - if (tokenizer) { - CHECK_EQ(pp_world_size, 1); - tokenizer->GenerateText(*model, FLAGS_batch_size, FLAGS_sequence_length, FLAGS_text_length, device); - } + if (after_step) { + after_step(*model, device, step + 1); } } - if (!FLAGS_save.empty() && FLAGS_save_interval > 0) { - if ((step + 1) % FLAGS_save_interval == 0 || (step + 1) == FLAGS_num_iteration) { + if (!options.save.empty() && options.save_interval > 0) { + if ((step + 1) % options.save_interval == 0 || (step + 1) == options.num_iteration) { std::filesystem::path step_dir - = std::filesystem::path(FLAGS_save) / std::format("checkpoint_step_{:06d}", step + 1); + = std::filesystem::path(options.save) / std::format("checkpoint_step_{:06d}", step + 1); if (rank.IsParallel()) { step_dir /= std::format("rank_{:06d}", rank.GlobalRank()); } @@ -559,64 +465,47 @@ void Train(const nn::parallel::Rank &rank) { } // Save LoRA weights if enabled and path specified - if (lora_enabled && !FLAGS_lora_save_path.empty()) { - LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; - nn::lora::SaveLoRAWeights(model, FLAGS_lora_save_path); + if (lora_enabled && !options.lora_save_path.empty()) { + LOG(INFO) << "Saving LoRA weights to: " << options.lora_save_path; + nn::lora::SaveLoRAWeights(model, options.lora_save_path); } #ifdef PROFILE_MODE - Profiler::Instance().Report("llama3.report", Profiler::SortBy::DeviceTimePercentage); - Profiler::Instance().PrintRecords("llama3.records.log"); + Profiler::Instance().Report(options.profile_name + ".report", Profiler::SortBy::DeviceTimePercentage); + Profiler::Instance().PrintRecords(options.profile_name + ".records.log"); #endif -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA needs synchronization before training teardown. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { + if (options.synchronize_device) { impl->SynchronizeDevice(device); } -#endif } -int main(int argc, char *argv[]) { - // Register provider metadata and implementations before gflags validates - // --device. The device runtime initializes lazily on first DeviceGuard use. -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR - INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR(); -#endif - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); +} // namespace - auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); - nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, - FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); +void Pretrain(const TrainingOptions &options, const DatasetProvider &dataset_provider, + const ModelProvider &model_provider, const ForwardStep &forward_step, const AfterStep &after_step) { + ValidateOptions(options); + CHECK(dataset_provider); + CHECK(model_provider); + CHECK(forward_step); + auto precision_config = utils::PrecisionCheckConfig::Parse(options.precision_check); + nn::parallel::global::InitAllEnv(options.nthread_per_process, options.tensor_parallel, options.sequence_parallel, + options.pipeline_parallel, options.virtual_pipeline_parallel); utils::PrecisionCheckEnv::Instance().Init(precision_config); - LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); - if (FLAGS_nthread_per_process > 1) { + auto run_rank = [&](int thread_rank) { + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), thread_rank, + nn::parallel::global::GetNprocPerNode(), options.nthread_per_process); + Train(options, rank, dataset_provider, model_provider, forward_step, after_step); + }; + if (options.nthread_per_process > 1) { std::vector threads; - for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), idx, - nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); - threads.emplace_back(Train, rank); - } - + for (int idx = 0; idx < options.nthread_per_process; ++idx) { threads.emplace_back(run_rank, idx); } for (auto &thread : threads) { thread.join(); } } else { - nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), - FLAGS_nthread_per_process); - Train(rank); - } - - gflags::ShutDownCommandLineFlags(); - google::ShutdownGoogleLogging(); - -#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - // FIXME(cx): MACA parallel execution bypasses static destruction to avoid teardown failures. - if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS - && Device::ParseType(FLAGS_device).value() == Device::DeviceType::kPrivateUse1) { - std::_Exit(0); + run_rank(0); } -#endif - return 0; } + +} // namespace infini_train::training diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5b9838781..24b915582 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,6 +41,9 @@ add_subdirectory(transformer) # Checkpoint tests add_subdirectory(checkpoint) +# Pretraining provider contracts and model registry tests. +add_subdirectory(pretrain) + # Suite declaration above is device-independent. Instantiate the built-in # devices after every test subdirectory has registered its sources. _infini_train_instantiate_test_suites( diff --git a/tests/pretrain/CMakeLists.txt b/tests/pretrain/CMakeLists.txt new file mode 100644 index 000000000..cee64a8b0 --- /dev/null +++ b/tests/pretrain/CMakeLists.txt @@ -0,0 +1,10 @@ +infini_train_add_test(test_model_provider_registry + SOURCES test_model_provider_registry.cc + LABELS cpu pretrain +) + +infini_train_add_test(test_training_arguments + SOURCES test_training_arguments.cc + LINK_LIBRARIES InfiniTrain::training_cli + LABELS cpu pretrain +) diff --git a/tests/pretrain/test_model_provider_registry.cc b/tests/pretrain/test_model_provider_registry.cc new file mode 100644 index 000000000..c2836a0a3 --- /dev/null +++ b/tests/pretrain/test_model_provider_registry.cc @@ -0,0 +1,56 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/training/model_provider_registry.h" + +using infini_train::training::ModelFactory; +using infini_train::training::ModelProvider; +using infini_train::training::ModelProviderRegistry; + +TEST(ModelProviderRegistryTest, ResolvesWithoutBuildingAndOwnsBoundArguments) { + int calls = 0; + std::string received; + ModelProvider provider; + { + ModelProviderRegistry registry; + registry.Register("unused", [](const std::string &) { + ADD_FAILURE() << "Resolved the wrong model factory"; + return std::shared_ptr{}; + }); + registry.Register("selected", [&](const std::string &weights) { + ++calls; + received = weights; + return std::shared_ptr{}; + }); + std::string weights = "initial.bin"; + provider = registry.Resolve("selected", weights); + weights = "changed.bin"; + EXPECT_EQ(calls, 0); // Model construction must wait for rank initialization. + } + // A resolved provider must not depend on registry or argument lifetimes. + provider(); + EXPECT_EQ(calls, 1); + EXPECT_EQ(received, "initial.bin"); + provider(); + EXPECT_EQ(calls, 2); // Each rank invokes its provider independently. +} + +TEST(ModelProviderRegistryTest, RejectsDuplicateAndInvalidRegistrations) { + ModelProviderRegistry registry; + ModelFactory factory = [](const std::string &) { return std::shared_ptr{}; }; + registry.Register("existing", factory); + EXPECT_DEATH(registry.Register("existing", factory), "Model provider already registered: existing"); + EXPECT_DEATH(registry.Register("", factory), "Model provider name must not be empty"); + EXPECT_DEATH(registry.Register("empty", {}), "Model factory must not be empty: empty"); +} + +TEST(ModelProviderRegistryTest, ListsRegisteredNamesInUnknownModelError) { + ModelProviderRegistry registry; + ModelFactory factory = [](const std::string &) { return std::shared_ptr{}; }; + registry.Register("second", factory); + registry.Register("first", factory); + EXPECT_EQ(registry.Names(), "first, second"); + EXPECT_DEATH(registry.Resolve("missing", ""), "Unknown model: missing. Available models: first, second"); +} diff --git a/tests/pretrain/test_training_arguments.cc b/tests/pretrain/test_training_arguments.cc new file mode 100644 index 000000000..d58bb59c5 --- /dev/null +++ b/tests/pretrain/test_training_arguments.cc @@ -0,0 +1,129 @@ +#include +#include +#include +#include +#include + +#include + +#include "gflags/gflags.h" +#include "gtest/gtest.h" + +#include "infini_train/include/training/arguments.h" + +namespace { +namespace fs = std::filesystem; +using infini_train::training::ParseTrainingFlags; +using infini_train::training::TrainingOptions; +using infini_train::training::TrainingOptionsFromFlags; + +class TrainingArgumentsTest : public ::testing::Test { +protected: + void SetUp() override { + auto pattern = (fs::temp_directory_path() / "training arguments XXXXXX").string(); + ASSERT_NE(mkdtemp(pattern.data()), nullptr); + directory_ = pattern; + } + void TearDown() override { fs::remove_all(directory_); } + + fs::path Write(const std::string &name, const std::string &contents) { + const auto path = directory_ / name; + std::ofstream file(path); + file << contents; + return path; + } + + void Parse(std::vector flags, const fs::path &defaults = {}) { + flags.insert(flags.begin(), "test_training_arguments"); + std::vector argv; + for (auto &flag : flags) { argv.push_back(flag.data()); } + int argc = argv.size(); + argv.push_back(nullptr); + auto data = argv.data(); + ParseTrainingFlags(&argc, &data, defaults); + } + + gflags::FlagSaver flag_saver_; + fs::path directory_; +}; + +TEST_F(TrainingArgumentsTest, CppAndCliUseTheSameDefaults) { + const TrainingOptions expected; + Parse({}); + const auto actual = TrainingOptionsFromFlags(); + EXPECT_EQ(actual.batch_size, expected.batch_size); + EXPECT_EQ(actual.learning_rate, expected.learning_rate); + EXPECT_EQ(actual.nthread_per_process, expected.nthread_per_process); + EXPECT_EQ(actual.lr_decay_iters, expected.lr_decay_iters); + EXPECT_EQ(actual.overfit_single_batch, expected.overfit_single_batch); + EXPECT_EQ(actual.lora_target_modules, expected.lora_target_modules); + EXPECT_EQ(actual.optimizer, expected.optimizer); + EXPECT_EQ(actual.profile_name, expected.profile_name); +} + +TEST_F(TrainingArgumentsTest, FileAndCliOverridesPreserveOrder) { + auto defaults = Write("default.flags", "--optimizer=sgd\n--learning_rate=0.1\n--save_optimizer_state=false\n"); + auto first = Write("first.flags", "# comment\n\n--learning_rate=0.2\n--profile_name=a profile\n"); + auto nested = Write("nested.flags", "--learning_rate=0.3\n--nosave_optimizer_state\n"); + auto second = Write("second.flags", "--flagfile=" + nested.string() + "\n--learning_rate=0.4\n"); + Parse({"--flagfile=" + first.string(), "--flagfile", second.string(), "--learning_rate=0.5", + "--save_optimizer_state=true"}, + defaults); + const auto actual = TrainingOptionsFromFlags(); + EXPECT_EQ(actual.optimizer, "sgd"); + EXPECT_DOUBLE_EQ(actual.learning_rate, 0.5); + EXPECT_TRUE(actual.save_optimizer_state); + EXPECT_EQ(actual.profile_name, "a profile"); +} + +TEST_F(TrainingArgumentsTest, RejectsUnknownFileFlagsEvenWithUndefok) { + auto file = Write("typo.flags", "# header\n--learning_raet=0.2\n"); + EXPECT_EXIT(Parse({"--undefok=learning_raet", "--flagfile=" + file.string()}), ::testing::ExitedWithCode(1), + "typo.flags:2: Unknown flag: --learning_raet"); + EXPECT_EXIT(Parse({}, file), ::testing::ExitedWithCode(1), "typo.flags:2: Unknown flag"); +} + +TEST_F(TrainingArgumentsTest, DefaultFileDoesNotReplayItsNestedInclude) { + auto base = Write("base.flags", "--learning_rate=0.1\n"); + auto defaults = Write("defaults.flags", "--flagfile=" + base.string() + "\n--learning_rate=0.2\n"); + Parse({}, defaults); + EXPECT_DOUBLE_EQ(TrainingOptionsFromFlags().learning_rate, 0.2); +} + +TEST_F(TrainingArgumentsTest, DefaultFilePathMayContainCommas) { + auto defaults = Write("default, preset.flags", "--learning_rate=0.2\n"); + Parse({}, defaults); + EXPECT_DOUBLE_EQ(TrainingOptionsFromFlags().learning_rate, 0.2); +} + +TEST_F(TrainingArgumentsTest, RejectsMissingValuesAndUnsupportedSyntax) { + auto missing = Write("missing.flags", "--learning_rate\n"); + auto section = Write("section.flags", "some_program\n--learning_rate=0.2\n"); + auto type = Write("type.flags", "--batch_size=not-a-number\n"); + EXPECT_EXIT(Parse({"--flagfile=" + missing.string()}), ::testing::ExitedWithCode(1), + "missing.flags:1: Missing value"); + EXPECT_EXIT(Parse({"--flagfile=" + section.string()}), ::testing::ExitedWithCode(1), + "section.flags:1: Expected --name=value"); + EXPECT_EXIT(Parse({"--flagfile=" + type.string()}), ::testing::ExitedWithCode(1), "batch_size"); +} + +TEST_F(TrainingArgumentsTest, ValidatesNestedFilesAndRejectsCycles) { + auto bad = Write("nested bad.flags", "--does_not_exist=1\n"); + auto outer = Write("outer.flags", "--flagfile=" + bad.string() + "\n"); + EXPECT_EXIT(Parse({"--flagfile=" + outer.string()}), ::testing::ExitedWithCode(1), + "nested bad.flags:1: Unknown flag"); + auto cycle = Write("cycle.flags", "--flagfile=" + (directory_ / "cycle.flags").string() + "\n"); + EXPECT_EXIT(Parse({"--flagfile=" + cycle.string()}), ::testing::ExitedWithCode(1), "Cyclic.*cycle.flags"); +} + +TEST_F(TrainingArgumentsTest, ExplicitFileWorksWithoutPackagedDefaults) { + auto supplied = Write("custom.flags", "--learning_rate=0.25\n"); + Parse({"--flagfile", supplied.string()}, directory_ / "absent-default.flags"); + EXPECT_DOUBLE_EQ(TrainingOptionsFromFlags().learning_rate, 0.25); + EXPECT_EXIT(Parse({}, directory_ / "absent-default.flags"), ::testing::ExitedWithCode(1), + "Default training flag file not found"); + EXPECT_EXIT(Parse({"--flagfile=" + (directory_ / "absent.flags").string()}), ::testing::ExitedWithCode(1), + "Cannot read flag file"); +} + +} // namespace