diff --git a/CMakeLists.txt b/CMakeLists.txt index 88911e84..9a557877 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,52 +231,20 @@ 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 # ------------------------------------------------------------------------------ if(PROJECT_IS_TOP_LEVEL) -add_executable(mnist - example/mnist/main.cc - example/mnist/dataset.cc - example/mnist/net.cc -) -link_infini_train_exe(mnist) - -add_executable(gpt2 - example/gpt2/main.cc - example/common/tiny_shakespeare_dataset.cc - example/common/utils.cc - example/common/tokenizer.cc - example/gpt2/checkpoint_loader.cc -) -link_infini_train_exe(gpt2) - -add_executable(mixtral - example/mixtral/main.cc - example/common/tiny_shakespeare_dataset.cc - example/common/utils.cc - example/mixtral/checkpoint_loader.cc -) -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) + add_subdirectory(example) endif() # Tools diff --git a/README.md b/README.md index 84cd4d9b..8a3d140c 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,11 @@ HF_TOKEN=hf_xxx \ #### Model Examples +GPT-2, LLaMA3, Qwen3, and Mixtral share the `pretrain_gpt` training program. +Model presets are selected with `--flagfile=example//train.flags`; +subsequent CLI options override preset values. The named executables load +`configs//train.flags` beside the binary automatically. + The generated files can be passed directly to the corresponding executables: ##### MNIST diff --git a/cmake/gpt_examples.cmake b/cmake/gpt_examples.cmake new file mode 100644 index 00000000..dff48e67 --- /dev/null +++ b/cmake/gpt_examples.cmake @@ -0,0 +1,51 @@ +include_guard(GLOBAL) + +function(infini_train_add_gpt_executable target_name) + add_executable(${target_name}) + target_link_libraries(${target_name} PRIVATE infini_train_gpt_sources InfiniTrain::training_cli) + # Keep the framework archives as direct link entries in their whole-archive + # group; transitive de-duplication can otherwise move the core outside it. + link_infini_train_exe(${target_name}) +endfunction() + +# Called by example//CMakeLists.txt. Each model contributes source files +# to the shared GPT program and declares a preset-loading executable. +function(infini_train_add_gpt_model model_name) + cmake_parse_arguments(MODEL "" "FLAGFILE" "SOURCES" ${ARGN}) + if(MODEL_UNPARSED_ARGUMENTS OR MODEL_KEYWORDS_MISSING_VALUES) + message(FATAL_ERROR + "infini_train_add_gpt_model(${model_name}): invalid or incomplete arguments") + endif() + if(NOT MODEL_SOURCES OR NOT MODEL_FLAGFILE) + message(FATAL_ERROR + "infini_train_add_gpt_model(${model_name}): SOURCES and FLAGFILE are required") + endif() + if(NOT TARGET infini_train_gpt_sources) + message(FATAL_ERROR "GPT models must be declared from the example build") + endif() + + # Absolute source paths retain the declaring directory's meaning when the + # interface sources are consumed by executables in other model directories. + foreach(source IN LISTS MODEL_SOURCES) + get_filename_component(source_path "${source}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") + target_sources(infini_train_gpt_sources INTERFACE "${source_path}") + endforeach() + get_filename_component(flagfile "${MODEL_FLAGFILE}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}") + if(NOT EXISTS "${flagfile}" OR IS_DIRECTORY "${flagfile}") + message(FATAL_ERROR "Training preset not found: ${flagfile}") + endif() + + infini_train_add_gpt_executable(${model_name}) + target_compile_definitions(${model_name} PRIVATE + INFINITRAIN_GPT_DEFAULT_FLAGFILE="configs/${model_name}/train.flags") + + # Target-file directories also cover multi-config generators and custom output + # locations. Preset edits are copied on the next build without recompilation. + add_custom_target(${model_name}_training_config + COMMAND ${CMAKE_COMMAND} -E make_directory "$/configs/${model_name}" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${flagfile}" "$/configs/${model_name}/train.flags" + VERBATIM + ) + add_dependencies(${model_name} ${model_name}_training_config) +endfunction() diff --git a/docs/unified_pretraining.md b/docs/unified_pretraining.md new file mode 100644 index 00000000..b7aaa586 --- /dev/null +++ b/docs/unified_pretraining.md @@ -0,0 +1,453 @@ +# 统一 GPT 预训练入口设计 + +统一训练入口将 GPT-2、LLaMA3、Qwen3、Mixtral 的训练过程组织为一套框架生命周期。任务入口负责提供模型、数据和计算逻辑，`infini_train::training::Pretrain()` 负责执行这些逻辑，并管理并行包装、优化器、学习率调度和训练状态。 + +四个模型都基于 `nn::TransformerModel`，架构由 `TransformerConfig` 描述。这一共同基础使模型特例可以通过 provider 接入统一训练流程。开发者增加模型时维护模型目录，增加训练能力时维护框架训练层。 + +本文以当前实现为准。结构调整的对照基线是 `master` 提交 `5a7be78`。 + +## 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 recompute = nn::ParseActivationRecomputeOptions( + options.activation_recompute, options.recompute_granularity, + options.recompute_method, options.recompute_num_layers); +const auto model_provider = models.Resolve(FLAGS_model, FLAGS_llmc_filepath, + recompute); + +training::Pretrain(options, dataset_provider, model_provider, + ForwardStepGPT, after_step); +``` + +`Resolve()` 查找名称、复制 factory 并绑定权重路径与 activation recompute 配置，返回一个零参数的 `ModelProvider`。模型实例在训练流程调用该 provider 时创建。这样，模型构建可以读取已经设置好的 TP/PP/EP 状态，并生成当前 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。 | +| Mixtral | [model_provider.cc](../example/mixtral/model_provider.cc) | 从 LLMC 文件头读取实际 MoE 结构，或使用 `TinyMixtralConfig()` 随机初始化；校验 TP、ETP、SP 与专家切分约束。 | + +每个 provider 文件通过文件内的静态初始化 lambda 调用 `Register()`。登记发生在进入 `main()` 前，登记内容是名称和函数；模型构建属于后续的 rank 初始化流程。训练阶段使用已经解析、绑定的回调，注册表的修改阶段位于程序启动期。 + +注册表要求名称唯一、factory 为可调用对象。名称冲突和非法注册在登记时报告；`--help` 和未知模型错误中的可用名称列表都由同一张注册表生成。 + +模型目录通过 `infini_train_add_gpt_model()` 声明 provider 和 loader 源码。构建层将这些源码加入共享的 `infini_train_gpt_sources` 接口目标，由各 executable 直接编译，使注册代码参与程序初始化。 + +### 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、设备 guard 和 DP/TP/PP 通信组，计算梯度累积次数。 +3. **构建模型。** 调用 provider，读取实际 `Config()`，执行设备迁移、精度检查名称映射和 LoRA 注入。 +4. **包装并行执行。** 根据拓扑构建 `PipelineParallel` 和 DDP；组合使用时，先建立流水线，再为本地 chunk 添加 DDP 包装。 +5. **建立数据与优化组件。** 调用 dataset provider，创建加载器，通过 `SetupOptimizer()` 和 `SetupLRScheduler()` 创建优化器及调度器。 +6. **恢复训练状态。** 通过 `ResumeFromCheckpoint()` 恢复状态，并根据消费样本数推进数据迭代器。 +7. **执行训练步。** 运行 forward/loss、梯度累积、backward，通过 `FinalizeModelGrads()` 完成梯度同步，再执行优化器更新，随后更新学习率调度器。 +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) +``` + +CLI 同时接受序列单位的 `micro_batch_size`、`global_batch_size`：非零 `micro_batch_size` 覆盖 `batch_size`，非零 `global_batch_size` 乘以最终的 `sequence_length` 后覆盖 `total_batch_size`。参数文件使用统一的 `batch_size` / `total_batch_size`，原有 Mixtral 命令可以继续传入序列单位的参数。同一调用中建议选择一套 batch 参数；两套同时出现时，非零序列参数优先。 + +配置要求 `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/VPP 元数据。保存目录由 checkpoint manager 组织为 `iter_0000001/` 等迭代目录。 + +恢复训练时 provider 建立匹配的逻辑模型结构，checkpoint 层通过各模块的 `BuildShardedStateDict()` 描述本地张量到逻辑参数的映射；PP 包装器提供本地 chunk 到全局层名的映射。优化器状态恢复遵循 checkpoint 模块对拓扑与优化器类型的约束。`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 后，普通循环调用 `FinalizeModelGrads({model})`，PP 调度器对本地 chunks 调用同一接口，再执行 optimizer step。该接口统一收尾 DDP 梯度通信及 SP / QK norm 的 TP 梯度归并。Dense 模型的 DDP `no_sync` 与 microbatch 累积配合；MoE 按下述配置在最终同步阶段处理稀疏路由参数。 + +### 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` 执行。 + +### 5.3 MoE 与专家并行 + +`TrainingOptions` 的 `expert_parallel`、`expert_tensor_parallel` 传给全局拓扑初始化。Dense rank generator 描述 TP/DP/PP，expert rank generator 描述 ETP/EP/EDP/PP。数据分发与全局 batch 计算使用 dense DP；专家参数的同步范围由 DDP 根据 expert EDP 拓扑和参数标记确定。 + +训练层从 `TransformerConfig::ffn_type` 判断 MoE，并为整模型或每个 PP chunk 设置 `handle_unused_parameters=true`、`overlap_grad_reduce=false`。这样，路由未命中的专家按零梯度处理，各 rank 在固定顺序完成梯度同步。MoE 使用 ZeRO 0/1/2。 + +Mixtral provider 封装当前模型能力边界：TP=1、ETP=1、SP 关闭，专家数可被 EP 大小整除。随机初始化支持 PP/VPP，当前 LLMC loader 使用 PP=1。EP=1 可使用公共 checkpoint 路径；EP>1 的训练 checkpoint 需要专家逻辑编号与分片元数据，入口在请求保存或恢复时明确校验这一条件。 + +### 5.4 Activation recompute + +公共 CLI 将 `activation_recompute`、`recompute_granularity`、`recompute_method`、`recompute_num_layers` 转成 `ActivationRecomputeOptions`，注册表按值绑定到 model provider。Provider 在模型构造前把这些选项应用到实际配置；LLMC 加载与随机初始化共用这一时机。Transformer 执行层据此安排 full-layer recompute，支持 `none`、`uniform`、`block` 方法。 + +`log_interval` 控制公共训练日志频率，零表示关闭步骤日志；`print_timing` 在迭代边界同步设备并汇总平均耗时。 + +## 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) | [Mixtral](../example/mixtral/train.flags) | +| --- | --- | --- | --- | +| optimizer | `sgd` | `adam` | `adam` | +| learning_rate | `0.0001` | `0.00001` | `0.0001` | +| device | `cuda` | `cuda` | `cpu` | +| lora_target_modules | `c_attn,c_proj` | `c_attn,c_proj,c_fc,c_fc2` | 框架默认值 | +| load_optimizer_state / save_optimizer_state | `false` | `true` | 框架默认值 `true` | +| profile_name | `gpt2` | 对应模型名 | `mixtral` | + +预设还包含 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 共用 `infini_train_gpt_sources`。当前模型目录生成以下产物： + +| executable | 配置选择 | +| --- | --- | +| `pretrain_gpt` | CLI 显式提供模型与配置。 | +| `gpt2` | 先加载 `configs/gpt2/train.flags`，再解析 CLI。 | +| `llama3` | 先加载 `configs/llama3/train.flags`，再解析 CLI。 | +| `qwen3` | 先加载 `configs/qwen3/train.flags`，再解析 CLI。 | +| `mixtral` | 先加载 `configs/mixtral/train.flags`，再解析 CLI。 | + +构建组织分为三个层次： + +| 层次 | 入口 | 职责 | +| --- | --- | --- | +| 项目 | [CMakeLists.txt](../CMakeLists.txt) | 定义框架及公共链接组件，通过 `add_subdirectory(example)` 加载示例。 | +| 示例聚合 | [example/CMakeLists.txt](../example/CMakeLists.txt) | 定义 GPT 公共任务源码和统一入口，发现 `example/*/CMakeLists.txt`。 | +| 模型目录 | `example//CMakeLists.txt` | 声明该模型的 provider、loader、参数文件及 executable 名称。 | + +目录发现使用 `file(GLOB ... CONFIGURE_DEPENDS)`，下一次构建会检测模型清单文件的新增或移除。每个目录通过自己的 `CMakeLists.txt` 显式列出源码。GPT 模型调用 [gpt_examples.cmake](../cmake/gpt_examples.cmake) 中的 `infini_train_add_gpt_model()`；MNIST 的目录使用常规 executable 定义。 + +`infini_train_gpt_sources` 是一个接口目标，携带公共任务源码和各模型登记的源码。每个 executable 编译这组接口源码，并使用自己的编译定义，因此后端适配与默认配置设置作用于对应程序。公共 helper 为目标链接训练 CLI，并通过 `link_infini_train_exe()` 组织框架静态库的 whole-archive 链接组。 + +构建 helper 通过 `INFINITRAIN_GPT_DEFAULT_FLAGFILE` 为模型 executable 指定相对配置路径，并同步声明的参数文件。默认输出位置保持在构建根目录；配置同步使用实际目标文件目录，适配多配置生成器和自定义输出位置： + +```text +build/ + pretrain_gpt + gpt2 + llama3 + qwen3 + mixtral + configs/ + gpt2/train.flags + llama3/train.flags + qwen3/train.flags + mixtral/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 根据权重路径选择加载或随机初始化，并在构造模型前应用绑定的 recompute 配置。 +4. 在 `train.flags` 中提供模型选择和训练预设。 +5. 在模型目录添加 `CMakeLists.txt`，声明源码和参数文件： + +```cmake +infini_train_add_gpt_model(my_model + SOURCES model_provider.cc checkpoint_loader.cc + FLAGFILE train.flags +) +``` + +例如，新目录 `example/my_model/` 包含这些文件后，执行： + +```bash +cmake --build build --target pretrain_gpt my_model +``` + +构建系统自动加载该目录的声明，将 provider 编入统一入口，并生成 `my_model` executable 和随包配置 `configs/my_model/train.flags`。已有模型的构建规则同样位于各自目录。 + +CMake 声明决定编译哪些实现，C++ provider 注册决定 `--model` 接受哪些名称。注册表为入口提供名称解析、帮助列表和诊断信息；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 | 通过模型目录声明 provider/loader，链接 GPT 源码接口目标及所需后端组件，并分发相应参数文件。 | + +## 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) | 默认值一致性、参数覆盖顺序、序列 batch 换算及溢出、并行与重计算选项、文件语法、嵌套包含和路径处理。 | + +```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 检查提供构建产物与启动流程验证。 + +### 10.1 基于 master `5a7be78` 的接入验证 + +本轮验证使用四层固定 LLMC 权重和固定 token 数据，每组训练两步，对照输出 loss；Mixtral 的随机初始化 PP 用例使用默认模型配置训练一步。 + +| 验证范围 | 结果 | +| --- | --- | +| 四模型的 CPU / CUDA、命名 executable / `pretrain_gpt` | loss 对照通过；命名 executable 自动读取随包参数文件。 | +| 四模型的 full / uniform recompute 与普通训练、保存一轮后恢复第二轮 | loss 对照通过。 | +| GPT-2、LLaMA3、Qwen3 的两卡 DP、TP、SP、PP、VPP | 与同一模型单卡训练的 loss 对照通过。 | +| Mixtral 两卡 DP、EP，四卡 EP+EDP | 与单卡训练的 loss 对照通过。 | +| Mixtral EP+ZeRO 0/1/2 | 与 master 原始 Mixtral 训练逻辑在同配置下的 loss 一致。对照程序使用文件头建立同一小型模型，训练循环沿用 master。 | +| Mixtral 随机初始化 PP、VPP、EP+PP | 均完成训练，loss 有限。 | +| C++ / CTest | 选中的 60 项中，50 项通过、10 项按现有测试条件跳过，零失败；覆盖训练参数、注册表、checkpoint、Transformer 和并行拓扑。 | +| 代码规范 | 全量 clang-format 16、Black、clang-tidy 18、IWYU 检查通过。 | + +Mixtral 的 ZeRO 1/2 相对 ZeRO 0 仍存在 master 同样具有的小幅数值偏差：本轮第二步 loss 分别为 `4.120690`、`4.120387`，ZeRO 0 为 `4.120744`。统一入口与 master 在各自相同 ZeRO 配置下的结果一致；该对照验证的是入口迁移的数值保持。 + +本地构建与检查日志位于 `build/rebase-*.log`，训练与对照记录位于 `build/rebase-validation/`。这些是本地验证产物；持续回归测试由上述 C++ 测试维护。 diff --git a/example/CMakeLists.txt b/example/CMakeLists.txt new file mode 100644 index 00000000..d85bea15 --- /dev/null +++ b/example/CMakeLists.txt @@ -0,0 +1,29 @@ +# Keep the established build/ paths when examples use subdirectories. +if(NOT DEFINED CMAKE_RUNTIME_OUTPUT_DIRECTORY) + set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}") +endif() + +include(${PROJECT_SOURCE_DIR}/cmake/gpt_examples.cmake) + +# Interface sources are compiled directly into each executable, preserving model +# registration objects and each consumer's external-backend compile definitions. +add_library(infini_train_gpt_sources INTERFACE) +target_sources(infini_train_gpt_sources INTERFACE + "${CMAKE_CURRENT_SOURCE_DIR}/pretrain_gpt.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/common/tiny_shakespeare_dataset.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/common/utils.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/common/tokenizer.cc" +) +infini_train_add_gpt_executable(pretrain_gpt) + +# Discover only directory manifests; each manifest explicitly declares its files. +# CONFIGURE_DEPENDS notices newly added/removed model directories on the next build. +file(GLOB example_manifests CONFIGURE_DEPENDS + RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/*/CMakeLists.txt" +) +list(SORT example_manifests) +foreach(manifest IN LISTS example_manifests) + get_filename_component(example_dir "${manifest}" DIRECTORY) + add_subdirectory("${example_dir}") +endforeach() diff --git a/example/gpt2/CMakeLists.txt b/example/gpt2/CMakeLists.txt new file mode 100644 index 00000000..acb0da4f --- /dev/null +++ b/example/gpt2/CMakeLists.txt @@ -0,0 +1,4 @@ +infini_train_add_gpt_model(gpt2 + SOURCES model_provider.cc checkpoint_loader.cc + FLAGFILE train.flags +) diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index b140edbd..d8098be9 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -96,6 +96,7 @@ std::shared_ptr LoadFromLLMC( 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); nn::ApplyActivationRecomputeOptions(&gpt2_config, recompute_options); diff --git a/example/gpt2/config.h b/example/gpt2/config.h index 30ae14b8..67f693a0 100644 --- a/example/gpt2/config.h +++ b/example/gpt2/config.h @@ -1,11 +1,26 @@ #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, @@ -24,6 +39,20 @@ 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 e0d35588..00000000 --- a/example/gpt2/main.cc +++ /dev/null @@ -1,780 +0,0 @@ -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -#include "infini_train/include/datatype.h" -#include "infini_train/include/nn/lora/lora_config.h" -#include "infini_train/include/nn/modules/transformer/transformer_config.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" -#include "infini_train/include/nn/parallel/process_group.h" -#include "infini_train/include/tensor.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/comm.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/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/nn/parallel/utils.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/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."); - -// activation recompute -DEFINE_bool(activation_recompute, false, - "Enable activation recompute to trade compute for memory."); -DEFINE_string(recompute_granularity, "full", - "Activation recompute granularity: none|full|selective"); -DEFINE_string(recompute_method, "none", - "Activation recompute method: none|uniform|block"); -DEFINE_uint32(recompute_num_layers, 0, - "Number of transformer layers per recompute region for " - "uniform/block methods."); - -// 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; - - const auto recompute_options = nn::ParseActivationRecomputeOptions( - FLAGS_activation_recompute, FLAGS_recompute_granularity, - FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); - - if (!FLAGS_llmc_filepath.empty()) { - model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); - } else if (kModelToConfigs.count(FLAGS_model)) { - model_config = kModelToConfigs.at(FLAGS_model); - gpt2::SanitizeGPT2Config(model_config); - nn::ApplyActivationRecomputeOptions(&model_config, recompute_options); - model = std::make_shared(model_config); - } - - CHECK(model) << "GPT2 example expects GPT2 model."; - - model->To(device); - - 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()); - } - - // 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, - .original_vocab_size = model_config.original_vocab_size, - .padded_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, - .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), - .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"; - } - - nn::parallel::FinalizeModelGrads({model}); - 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); - comm::AllReduce(lossf_tensor, comm::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::kDeviceTimePercentage); - 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, - /*expert_parallel_size=*/1, - /*expert_tensor_parallel_size=*/1); - 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 00000000..b2fac20c --- /dev/null +++ b/example/gpt2/model_provider.cc @@ -0,0 +1,40 @@ +#include +#include +#include +#include + +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/training/model_provider_registry.h" + +#include "example/gpt2/checkpoint_loader.h" +#include "example/gpt2/config.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, + const infini_train::nn::ActivationRecomputeOptions &recompute) { + if (!weights.empty()) { + return LoadFromLLMC(weights, recompute); + } + auto config = GPT2Config( + name, + infini_train::nn::parallel::global::GetTensorParallelSize()); + SanitizeGPT2Config(config); + infini_train::nn::ApplyActivationRecomputeOptions(&config, recompute); + 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 00000000..980faea7 --- /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/CMakeLists.txt b/example/llama3/CMakeLists.txt new file mode 100644 index 00000000..fa4423c4 --- /dev/null +++ b/example/llama3/CMakeLists.txt @@ -0,0 +1,4 @@ +infini_train_add_gpt_model(llama3 + SOURCES model_provider.cc checkpoint_loader.cc + FLAGFILE train.flags +) diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index bbd6a996..50e0ea8f 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -97,6 +97,7 @@ std::shared_ptr LoadFromLLMC( 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/main.cc b/example/llama3/main.cc deleted file mode 100644 index f8188cef..00000000 --- a/example/llama3/main.cc +++ /dev/null @@ -1,755 +0,0 @@ -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -#include "infini_train/include/datatype.h" -#include "infini_train/include/nn/lora/lora_config.h" -#include "infini_train/include/nn/modules/transformer/transformer_config.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" -#include "infini_train/include/tensor.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/comm.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/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/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/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."); - -// activation recompute -DEFINE_bool(activation_recompute, false, - "Enable activation recompute to trade compute for memory."); -DEFINE_string(recompute_granularity, "full", - "Activation recompute granularity: none|full|selective"); -DEFINE_string(recompute_method, "none", - "Activation recompute method: none|uniform|block"); -DEFINE_uint32(recompute_num_layers, 0, - "Number of transformer layers per recompute region for " - "uniform/block methods."); -// 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 = {"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 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 = llama3::LLaMA3Config(); - std::shared_ptr model = nullptr; - const auto recompute_options = nn::ParseActivationRecomputeOptions( - FLAGS_activation_recompute, FLAGS_recompute_granularity, - FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); - if (!FLAGS_llmc_filepath.empty()) { - model = llama3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); - } else { - nn::ApplyActivationRecomputeOptions(&model_config, recompute_options); - llama3::SanitizeLLaMA3Config(model_config); - model = std::make_shared(model_config); - } - - CHECK(model) << "LLaMA3 example expects LLaMA3 model."; - - 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): 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, - .original_vocab_size = model_config.original_vocab_size, - .padded_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, - .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), - .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"; - } - - nn::parallel::FinalizeModelGrads({model}); - 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); - comm::AllReduce(lossf_tensor, comm::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("llama3.report", - Profiler::SortBy::kDeviceTimePercentage); - Profiler::Instance().PrintRecords("llama3.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, - /*expert_parallel_size=*/1, - /*expert_tensor_parallel_size=*/1); - 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/llama3/model_provider.cc b/example/llama3/model_provider.cc new file mode 100644 index 00000000..c6428e07 --- /dev/null +++ b/example/llama3/model_provider.cc @@ -0,0 +1,31 @@ +#include +#include + +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/training/model_provider_registry.h" + +#include "example/llama3/checkpoint_loader.h" +#include "example/llama3/config.h" + +namespace llama3 { +namespace { + +[[maybe_unused]] const bool registered = [] { + infini_train::training::ModelProviderRegistry::Instance().Register( + "llama3", + [](const std::string &weights, + const infini_train::nn::ActivationRecomputeOptions &recompute) { + if (!weights.empty()) { + return LoadFromLLMC(weights, recompute); + } + auto config = LLaMA3Config(); + SanitizeLLaMA3Config(config); + infini_train::nn::ApplyActivationRecomputeOptions(&config, recompute); + 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 00000000..a7b36d2f --- /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/mixtral/CMakeLists.txt b/example/mixtral/CMakeLists.txt new file mode 100644 index 00000000..e65e2906 --- /dev/null +++ b/example/mixtral/CMakeLists.txt @@ -0,0 +1,4 @@ +infini_train_add_gpt_model(mixtral + SOURCES model_provider.cc checkpoint_loader.cc + FLAGFILE train.flags +) diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc deleted file mode 100644 index 82ee2b39..00000000 --- a/example/mixtral/main.cc +++ /dev/null @@ -1,420 +0,0 @@ -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -#include "infini_train/include/datatype.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/nn/modules/transformer/transformer_config.h" - -// Out-of-tree builds inject the selected provider 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/core/runtime/device_guard.h" -#include "infini_train/include/dataloader.h" -#include "infini_train/include/device.h" -#include "infini_train/include/nn/modules/loss.h" -#include "infini_train/include/nn/modules/transformer/transformer.h" -#include "infini_train/include/nn/parallel/comm.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.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/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/tensor.h" -#ifdef PROFILE_MODE -#include "infini_train/include/profiler.h" -#endif - -#include "example/common/tiny_shakespeare_dataset.h" -#include "example/mixtral/checkpoint_loader.h" -#include "example/mixtral/config.h" - -// I/O -DEFINE_string(input_bin, "", "input .bin to train on"); -DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); -// token layout for each step of the optimization -DEFINE_uint32(micro_batch_size, 4, "batch size, in units of #batch dimensions"); -DEFINE_uint32(global_batch_size, 4, - "total desired batch size, in units of #batch dimensions"); -DEFINE_uint32(sequence_length, 64, "sequence length"); -// workload (number of steps) -DEFINE_uint32(num_iteration, 10, "number of iterations to run"); -// optimization -DEFINE_double(learning_rate, 1e-4, "Peak learning rate."); -DEFINE_int32(zero_stage, 0, - "ZeRO stage (0/1/2); 0 disables DistributedOptimizer"); -// memory management -DEFINE_string(device, "cpu", "Training device."); -// 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_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."); -DEFINE_uint32(expert_parallel, 1, "Expert Parallel world size"); -DEFINE_uint32(expert_tensor_parallel, 1, - "Expert Tensor Parallel world size (currently must be 1)"); -// precision -DEFINE_string(dtype, "float32", - "precision used in training (float32/bfloat16)"); -// logging -DEFINE_uint32(log_interval, 1, - "Print train loss every N steps. 0 disables step loss logging."); -DEFINE_bool(print_timing, false, - "Print training-loop elapsed time and token throughput."); - -using namespace infini_train; - -namespace { - -constexpr char kDtypeFP32[] = "float32"; -constexpr char kDtypeBF16[] = "bfloat16"; - -void ValidateRuntimeFlags(const nn::TransformerConfig &config) { - CHECK(!FLAGS_input_bin.empty()) << "Mixtral training requires --input_bin"; - CHECK_GT(FLAGS_micro_batch_size, 0); - CHECK_GT(FLAGS_global_batch_size, 0); - CHECK_GT(FLAGS_sequence_length, 0); - CHECK_LE(FLAGS_sequence_length, config.block_size) - << "sequence_length must be <= model max positions (block_size)"; - CHECK(config.moe_config.has_value()); - CHECK_EQ(config.moe_config->num_experts % FLAGS_expert_parallel, 0) - << "The number of Mixtral experts must be divisible by expert_parallel"; - CHECK_LE(FLAGS_zero_stage, 2) << "Mixtral supports ZeRO-0/1/2"; - - const uint64_t batch_per_micro_step = - static_cast(FLAGS_micro_batch_size) * - nn::parallel::global::GetDataParallelSize(); - CHECK_EQ(FLAGS_global_batch_size % batch_per_micro_step, 0) - << "global_batch_size must be divisible by micro_batch_size * " - "data_parallel_size"; -} - -} // namespace - -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 <= 2; -}); - -void Train(const nn::parallel::Rank &rank) { - using namespace nn::parallel; - const auto device_type = Device::ParseType(FLAGS_device).value(); - - global::thread_global_rank = rank.GlobalRank(); - - const int ddp_world_size = global::GetDataParallelSize(); - const int pp_world_size = global::GetPipelineParallelSize(); - int ddp_rank = 0; - int pp_rank = 0; - Device device; - - if (rank.IsParallel()) { - device = Device(device_type, global::GetDeviceIndex(rank.thread_rank())); - auto *pg_factory = ProcessGroupFactory::Instance(device.type()); - if (ddp_world_size > 1) { - const auto *ddp_pg = pg_factory->GetOrCreate( - GetDataParallelProcessGroupName(rank.GlobalRank()), - GetDataParallelGroupRanks(rank.GlobalRank())); - CHECK_NOTNULL(ddp_pg); - ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); - } - if (pp_world_size > 1) { - const auto *pp_pg = pg_factory->GetOrCreate( - GetPipelineParallelProcessGroupName(rank.GlobalRank()), - GetPipelineParallelGroupRanks(rank.GlobalRank())); - CHECK_NOTNULL(pp_pg); - pp_rank = pp_pg->GetGroupRank(rank.GlobalRank()); - nn::parallel::pp_rank = pp_rank; - } - nn::parallel::tp_rank = 0; - } else { - device = Device(device_type, 0); - } - - // Profiling runs before kernel-local guards; keep this rank's device current - // throughout training so its events and synchronization use the correct - // stream. - core::DeviceGuard device_guard(device); - nn::TransformerConfig model_config = mixtral::TinyMixtralConfig(); - mixtral::SanitizeTinyMixtralConfig(model_config); - ValidateRuntimeFlags(model_config); - - std::shared_ptr model; - if (!FLAGS_llmc_filepath.empty()) { - model = mixtral::LoadFromLLMC(FLAGS_llmc_filepath, model_config); - } else { - model = std::make_shared(model_config); - } - model->To(device); - - const uint32_t grad_accum_steps = - FLAGS_global_batch_size / - (FLAGS_micro_batch_size * static_cast(ddp_world_size)); - if (rank.IsMainRank()) { - LOG(INFO) << "global batch size: " << FLAGS_global_batch_size - << ", data parallel size: " << ddp_world_size - << ", gradient accumulation steps: " << grad_accum_steps; - } - - const uint32_t loader_batch_size = - pp_world_size > 1 ? FLAGS_micro_batch_size * grad_accum_steps - : FLAGS_micro_batch_size; - DistributedDataLoader train_loader( - std::make_shared(FLAGS_input_bin, - FLAGS_sequence_length), - loader_batch_size, ddp_rank, ddp_world_size); - auto train_iter = train_loader.begin(); - - auto loss_fn = std::make_shared(); - loss_fn->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."; - } - - if (pp_world_size > 1) { - const std::vector> recv_shapes{ - {FLAGS_micro_batch_size, FLAGS_sequence_length, model_config.n_embd}}; - model = std::make_shared( - model, pp_world_size, grad_accum_steps, recv_shapes, pp_rank, device, - model_config.GetChunkSize()); - if (ddp_world_size > 1) { - auto ddp_config = DistributedDataParallelConfig{}; - ddp_config.handle_unused_parameters = true; - // Unused experts may never signal readiness; finalize all buckets in a - // fixed order. - ddp_config.overlap_grad_reduce = false; - ddp_config.zero_stage = FLAGS_zero_stage; - auto *mutable_chunks = - dynamic_cast(model.get())->mutable_chunks(); - for (size_t chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) { - mutable_chunks->at(chunk_id) = - std::make_shared( - mutable_chunks->at(chunk_id), rank, ddp_config); - } - } - } else if (ddp_world_size > 1) { - auto ddp_config = DistributedDataParallelConfig{}; - ddp_config.handle_unused_parameters = true; - // Unused experts may never signal readiness; finalize all buckets in a - // fixed order. - ddp_config.overlap_grad_reduce = false; - ddp_config.zero_stage = FLAGS_zero_stage; - model = std::make_shared(model, rank, ddp_config); - } - - auto optimizer_creator = - optimizers::Adam::CreateNamed(static_cast(FLAGS_learning_rate)); - const auto params_to_optimize = model->NamedParameters(); - std::shared_ptr optimizer; - 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, params_to_optimize, model_chunks, ddp_world_size, - ddp_rank); - } else { - optimizer = optimizer_creator(params_to_optimize); - } - auto device_impl = core::GetDeviceGuardImpl(device.type()); - std::vector step_duration_ms; - step_duration_ms.reserve(FLAGS_num_iteration); - const double tokens_per_step = - static_cast(FLAGS_global_batch_size) * FLAGS_sequence_length; - - LOG(INFO) << "Rank " << rank.GlobalRank() - << ": start tiny Mixtral training on " << device.ToString(); - for (uint32_t step = 0; step < FLAGS_num_iteration; ++step) { -#ifdef PROFILE_MODE - Profiler::Instance().SetTag("Step_" + std::to_string(step)); -#endif - device_impl->ResetMemPoolHighWatermarks(device); - device_impl->SynchronizeDevice(device); - const auto step_start_time = std::chrono::steady_clock::now(); - - float lossf = 0.0f; - if (pp_world_size > 1) { - if (train_iter == train_loader.end()) { - train_iter = train_loader.begin(); - } - auto [x_cpu, y_cpu] = *train_iter; - ++train_iter; - auto x = std::make_shared(x_cpu->To(device)); - auto y = std::make_shared(y_cpu->To(device)); - lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); - } else { - optimizer->ZeroGrad(); - for (uint32_t micro_step = 0; micro_step < grad_accum_steps; - ++micro_step) { - AutocastGuard autocast_guard(device.type(), dtype); - if (train_iter == train_loader.end()) { - train_iter = train_loader.begin(); - } - auto [x_cpu, y_cpu] = *train_iter; - ++train_iter; - auto x = std::make_shared(x_cpu->To(device)); - auto y = std::make_shared(y_cpu->To(device)); - auto logits = (*model)({x})[0]; - auto loss = - (*loss_fn)({logits, y})[0] / static_cast(grad_accum_steps); - autocast_guard.Disable(); - loss->Backward(); - // Defer the loss D2H copy until after backward; reading it earlier - // would synchronize the device between forward and backward. - auto loss_cpu = loss->To(Device()); - lossf += static_cast(loss_cpu.DataPtr())[0]; - } - nn::parallel::FinalizeModelGrads({model}); - optimizer->Step(); - } - - if (ddp_world_size > 1) { - auto loss_tensor = std::make_shared( - &lossf, std::vector{}, DataType::kFLOAT32, device); - const auto *dp_group = - ProcessGroupFactory::Instance(device.type()) - ->Get(GetDataParallelProcessGroupName(rank.GlobalRank())); - comm::AllReduce(loss_tensor, comm::ReduceOpType::kAvg, dp_group); - lossf = - static_cast(loss_tensor->To(Device()).DataPtr())[0]; - } - - device_impl->SynchronizeDevice(device); - const auto step_end_time = std::chrono::steady_clock::now(); - const double duration_ms = std::chrono::duration( - step_end_time - step_start_time) - .count(); - step_duration_ms.push_back(duration_ms); - - if (rank.IsLastRank() && FLAGS_log_interval > 0 && - ((step + 1) % FLAGS_log_interval == 0 || - step + 1 == FLAGS_num_iteration)) { - size_t used_mb = 0; - size_t reserved_mb = 0; - std::tie(used_mb, reserved_mb) = device_impl->GetMemPoolPeakMB(device); - std::cout << std::format( - "step {:4d}/{} | train loss {:.6f} | norm -1.0000 | lr " - "{:.2e} | ({:.2f} ms | {:.0f} tok/s | " - "peak used: {:5d} MB | peak reserved: {:5d} MB)", - step + 1, FLAGS_num_iteration, lossf, - FLAGS_learning_rate, duration_ms, - tokens_per_step / (duration_ms / 1e3), used_mb, - reserved_mb) - << std::endl; - } - } - - if (rank.IsLastRank() && FLAGS_print_timing && !step_duration_ms.empty()) { - double duration_sum_ms = 0.0; - for (size_t idx = step_duration_ms.size() > 1 ? 1 : 0; - idx < step_duration_ms.size(); ++idx) { - duration_sum_ms += step_duration_ms[idx]; - } - const size_t averaged_steps = step_duration_ms.size() > 1 - ? step_duration_ms.size() - 1 - : step_duration_ms.size(); - std::cout << std::format("final {} iters avg: {:.3f}ms", averaged_steps, - duration_sum_ms / averaged_steps) - << std::endl; - } - -#ifdef PROFILE_MODE - if (rank.IsLastRank()) { - Profiler::Instance().Report("mixtral.report", - Profiler::SortBy::kDeviceTimePercentage); - Profiler::Instance().PrintRecords("mixtral.records.log"); - } -#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]); - - CHECK_GT(FLAGS_nthread_per_process, 0); - CHECK_EQ(FLAGS_tensor_parallel, 1) - << "The Mixtral EP launcher currently requires --tensor_parallel=1"; - CHECK_EQ(FLAGS_expert_tensor_parallel, 1) - << "The Mixtral EP launcher currently requires " - "--expert_tensor_parallel=1"; - CHECK_GT(FLAGS_pipeline_parallel, 0); - CHECK_GT(FLAGS_virtual_pipeline_parallel, 0); - CHECK_GT(FLAGS_expert_parallel, 0); - - nn::parallel::global::InitAllEnv( - FLAGS_nthread_per_process, FLAGS_tensor_parallel, - /*sequence_parallel_enabled=*/false, FLAGS_pipeline_parallel, - FLAGS_virtual_pipeline_parallel, FLAGS_expert_parallel, - FLAGS_expert_tensor_parallel); - - LOG(INFO) << nn::parallel::global::ProcessGroupOverview( - nn::parallel::global::GetDenseRankGenerator(), - nn::parallel::global::GetExpertRankGenerator()); - // Use one worker thread for each logical rank owned by this process. - 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(), /*thread_rank=*/0, - nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); - Train(rank); - } - - gflags::ShutDownCommandLineFlags(); - google::ShutdownGoogleLogging(); - - return 0; -} diff --git a/example/mixtral/model_provider.cc b/example/mixtral/model_provider.cc new file mode 100644 index 00000000..22f79524 --- /dev/null +++ b/example/mixtral/model_provider.cc @@ -0,0 +1,47 @@ +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/training/model_provider_registry.h" + +#include "example/mixtral/checkpoint_loader.h" +#include "example/mixtral/config.h" + +namespace mixtral { +namespace { + +[[maybe_unused]] const bool registered = [] { + infini_train::training::ModelProviderRegistry::Instance().Register( + "mixtral", + [](const std::string &weights, + const infini_train::nn::ActivationRecomputeOptions &recompute) { + namespace global = infini_train::nn::parallel::global; + CHECK_EQ(global::GetTensorParallelSize(), 1) + << "Mixtral currently requires TP=1"; + CHECK_EQ(global::GetExpertTensorParallelSize(), 1) + << "Mixtral currently requires ETP=1"; + CHECK(!global::GetSequenceParallelEnabled()) + << "Mixtral currently requires SP disabled"; + auto config = + weights.empty() ? TinyMixtralConfig() : ConfigFromLLMC(weights); + SanitizeTinyMixtralConfig(config); + CHECK_EQ( + config.moe_config->num_experts % global::GetExpertParallelSize(), + 0); + infini_train::nn::ApplyActivationRecomputeOptions(&config, recompute); + if (!weights.empty()) { + return LoadFromLLMC(weights, config); + } + return std::make_shared(config); + }); + return true; +}(); + +} // namespace +} // namespace mixtral diff --git a/example/mixtral/train.flags b/example/mixtral/train.flags new file mode 100644 index 00000000..04cd35a3 --- /dev/null +++ b/example/mixtral/train.flags @@ -0,0 +1,12 @@ +# Mixtral training defaults. Global batch size is expressed in tokens. +--model=mixtral +--device=cpu +--optimizer=adam +--learning_rate=0.0001 +--batch_size=4 +--sequence_length=64 +--total_batch_size=256 +--num_iteration=10 +--dtype=float32 +--zero_stage=0 +--profile_name=mixtral diff --git a/example/mnist/CMakeLists.txt b/example/mnist/CMakeLists.txt new file mode 100644 index 00000000..41dfc436 --- /dev/null +++ b/example/mnist/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(mnist main.cc dataset.cc net.cc) +target_link_libraries(mnist PRIVATE InfiniTrain::infini_train_executable) diff --git a/example/pretrain_gpt.cc b/example/pretrain_gpt.cc new file mode 100644 index 00000000..2bcd1bf9 --- /dev/null +++ b/example/pretrain_gpt.cc @@ -0,0 +1,153 @@ +// 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 +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS +#include "infini_train/include/device.h" +#endif +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/modules/module.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" + +#include "example/common/tiny_shakespeare_dataset.h" +#include "example/common/tokenizer.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"); + +namespace infini_train { +class Device; +class Tensor; +} // namespace infini_train + +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"; + CHECK(!FLAGS_input_bin.empty()) << "--input_bin is required"; + auto options = training::TrainingOptionsFromFlags(); + const auto recompute = nn::ParseActivationRecomputeOptions( + options.activation_recompute, options.recompute_granularity, + options.recompute_method, options.recompute_num_layers); + const auto model_provider = + models.Resolve(model_name, FLAGS_llmc_filepath, recompute); +#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/CMakeLists.txt b/example/qwen3/CMakeLists.txt new file mode 100644 index 00000000..cc227bd8 --- /dev/null +++ b/example/qwen3/CMakeLists.txt @@ -0,0 +1,4 @@ +infini_train_add_gpt_model(qwen3 + SOURCES model_provider.cc checkpoint_loader.cc + FLAGFILE train.flags +) diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc index 7270f7bb..fb66c3ee 100644 --- a/example/qwen3/checkpoint_loader.cc +++ b/example/qwen3/checkpoint_loader.cc @@ -115,6 +115,7 @@ std::shared_ptr LoadFromLLMC( 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 a9f925fe..00000000 --- a/example/qwen3/main.cc +++ /dev/null @@ -1,753 +0,0 @@ -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "gflags/gflags.h" -#include "glog/logging.h" - -#include "infini_train/include/datatype.h" -#include "infini_train/include/nn/lora/lora_config.h" -#include "infini_train/include/nn/modules/transformer/transformer_config.h" -#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" -#include "infini_train/include/tensor.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/comm.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/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/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."); - -// activation recompute -DEFINE_bool(activation_recompute, false, - "Enable activation recompute to trade compute for memory."); -DEFINE_string(recompute_granularity, "full", - "Activation recompute granularity: none|full|selective"); -DEFINE_string(recompute_method, "none", - "Activation recompute method: none|uniform|block"); -DEFINE_uint32(recompute_num_layers, 0, - "Number of transformer layers per recompute region for " - "uniform/block methods."); -// 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; - const auto recompute_options = nn::ParseActivationRecomputeOptions( - FLAGS_activation_recompute, FLAGS_recompute_granularity, - FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); - if (!FLAGS_llmc_filepath.empty()) { - model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); - } else { - nn::ApplyActivationRecomputeOptions(&model_config, recompute_options); - 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, - .original_vocab_size = model_config.original_vocab_size, - .padded_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, - .vpp_size = static_cast(FLAGS_virtual_pipeline_parallel), - .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"; - } - - nn::parallel::FinalizeModelGrads({model}); - 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); - comm::AllReduce(lossf_tensor, comm::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::kDeviceTimePercentage); - 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, - /*expert_parallel_size=*/1, - /*expert_tensor_parallel_size=*/1); - 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 00000000..4819a606 --- /dev/null +++ b/example/qwen3/model_provider.cc @@ -0,0 +1,31 @@ +#include +#include + +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/training/model_provider_registry.h" + +#include "example/qwen3/checkpoint_loader.h" +#include "example/qwen3/config.h" + +namespace qwen3 { +namespace { + +[[maybe_unused]] const bool registered = [] { + infini_train::training::ModelProviderRegistry::Instance().Register( + "qwen3", + [](const std::string &weights, + const infini_train::nn::ActivationRecomputeOptions &recompute) { + if (!weights.empty()) { + return LoadFromLLMC(weights, recompute); + } + auto config = Qwen3Config(); + SanitizeQwen3Config(config); + infini_train::nn::ApplyActivationRecomputeOptions(&config, recompute); + 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 00000000..bc70fd59 --- /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 19a96a78..d2d9bd28 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h +++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h @@ -9,6 +9,7 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/shard_spec.h" +#include "infini_train/include/training/training.h" namespace infini_train { class Tensor; @@ -38,7 +39,8 @@ class PipelineParallel : public Module { PipelineParallel(const std::shared_ptr module, int num_stages, int num_micro_batches, const std::vector> &recv_shape, - int rank, Device device, int vpp); + int rank, Device device, int vpp, + training::ForwardStepFunction forward_step = {}); float TrainStep(const std::vector> &input, const std::vector> &target, @@ -59,7 +61,8 @@ class PipelineParallel : public Module { 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 8e845695..bb6563eb 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_schedule.h +++ b/infini_train/include/nn/parallel/pp/pipeline_schedule.h @@ -6,6 +6,8 @@ #include #include +#include "infini_train/include/training/training.h" + namespace infini_train { class Tensor; class Optimizer; @@ -23,8 +25,11 @@ 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) {} + 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; @@ -44,6 +49,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 00000000..4a3cd0fc --- /dev/null +++ b/infini_train/include/training/arguments.h @@ -0,0 +1,21 @@ +#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 00000000..1a25c536 --- /dev/null +++ b/infini_train/include/training/model_provider_registry.h @@ -0,0 +1,39 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#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, + const nn::ActivationRecomputeOptions &recompute)>; + +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 nn::ActivationRecomputeOptions &recompute = {}) 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 00000000..0b04df21 --- /dev/null +++ b/infini_train/include/training/training.h @@ -0,0 +1,125 @@ +#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; + uint32_t expert_parallel = 1; + uint32_t expert_tensor_parallel = 1; + bool activation_recompute = false; + std::string recompute_granularity = "full"; + std::string recompute_method = "none"; + int64_t recompute_num_layers = 0; + uint32_t log_interval = 1; + bool print_timing = false; + 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 d91115f4..2430c3ac 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -30,9 +30,10 @@ void PipelineParallel::BuildPipelineStage( 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( @@ -92,7 +93,7 @@ StageInfo PipelineParallel::GetStageInfo(int total_layers, int pp_size, 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) + Device device, int chunk_size, training::ForwardStepFunction forward_step) : num_stages_(num_stages), rank_(pp_rank) { modules_[kModuleName] = std::move(module); @@ -115,7 +116,7 @@ PipelineParallel::PipelineParallel( BuildPipelineStage(recv_shape, device, std::move(chunks)); - SetupSchedule(num_micro_batches); + SetupSchedule(num_micro_batches, std::move(forward_step)); } std::vector> *PipelineParallel::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 b40a18de..30eead70 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -243,6 +244,8 @@ float PipelineSchedule::StepMicroBatches( 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()); for (const auto &chunk : stage_->chunks()) { @@ -273,8 +276,20 @@ float PipelineSchedule::StepMicroBatches( } } - activations[task.local_chunk_idx][mb] = - stage_->ForwardOneChunk(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()) { @@ -297,10 +312,15 @@ float PipelineSchedule::StepMicroBatches( infini_train::AutocastGuard autocast_guard(stage_->device().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 00000000..3cd48ba2 --- /dev/null +++ b/infini_train/src/training/arguments.cc @@ -0,0 +1,423 @@ +#include "infini_train/include/training/arguments.h" + +#include +#include +#include +#include +#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_uint32(expert_parallel, kTrainingDefaults.expert_parallel, + "Expert parallel world size"); +DEFINE_uint32(expert_tensor_parallel, kTrainingDefaults.expert_tensor_parallel, + "Expert tensor parallel world size"); +DEFINE_bool(activation_recompute, kTrainingDefaults.activation_recompute, + "Enable activation recomputation"); +DEFINE_string(recompute_granularity, kTrainingDefaults.recompute_granularity, + "Recompute granularity: full|selective"); +DEFINE_string(recompute_method, kTrainingDefaults.recompute_method, + "Recompute method: none|uniform|block"); +DEFINE_int64(recompute_num_layers, kTrainingDefaults.recompute_num_layers, + "Number of layers per recompute group"); +DEFINE_uint32(log_interval, kTrainingDefaults.log_interval, + "Report loss every N steps"); +DEFINE_bool(print_timing, kTrainingDefaults.print_timing, + "Synchronize each iteration and report average training time"); +// Sequence-based aliases are normalized to the framework's token-based options. +DEFINE_uint32(micro_batch_size, 0, + "Alias for batch_size; nonzero overrides batch_size"); +DEFINE_uint32(global_batch_size, 0, + "Global batch in sequences; nonzero overrides total_batch_size"); +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; + options.expert_parallel = FLAGS_expert_parallel; + options.expert_tensor_parallel = FLAGS_expert_tensor_parallel; + options.activation_recompute = FLAGS_activation_recompute; + options.recompute_granularity = FLAGS_recompute_granularity; + options.recompute_method = FLAGS_recompute_method; + options.recompute_num_layers = FLAGS_recompute_num_layers; + options.log_interval = FLAGS_log_interval; + options.print_timing = FLAGS_print_timing; + if (FLAGS_micro_batch_size != 0) { + options.batch_size = FLAGS_micro_batch_size; + } + if (FLAGS_global_batch_size != 0) { + const auto tokens = static_cast(FLAGS_global_batch_size) * + options.sequence_length; + if (tokens > std::numeric_limits::max()) { + throw std::invalid_argument( + "global_batch_size * sequence_length exceeds uint32 range"); + } + options.total_batch_size = static_cast(tokens); + } + 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 00000000..6d441e4b --- /dev/null +++ b/infini_train/src/training/model_provider_registry.cc @@ -0,0 +1,46 @@ +#include "infini_train/include/training/model_provider_registry.h" + +#include +#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 nn::ActivationRecomputeOptions &recompute) const { + const auto it = factories_.find(name); + CHECK(it != factories_.end()) + << "Unknown model: " << name << ". Available models: " << Names(); + return [factory = it->second, weights, recompute] { + return factory(weights, recompute); + }; +} + +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/infini_train/src/training/training.cc b/infini_train/src/training/training.cc new file mode 100644 index 00000000..4cb1cd64 --- /dev/null +++ b/infini_train/src/training/training.cc @@ -0,0 +1,640 @@ +#include "infini_train/include/training/training.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#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/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/lr_scheduler.h" +#include "infini_train/include/nn/lora/lora_config.h" +#include "infini_train/include/nn/lora/lora_utils.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/modules/transformer/transformer.h" +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/nn/parallel/comm.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.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/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/tensor.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 + +namespace infini_train::training { +namespace { + +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_GT(options.expert_parallel, 0); + CHECK_GT(options.expert_tensor_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"; +} + +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; + + 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()); + } + + 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(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 (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)."; + } + + 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); + } + + core::DeviceGuard device_guard(device); + + // calculate gradient accumulation from the desired total batch size and the + // current run configuration + 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: " << options.total_batch_size + << " => calculated gradient accumulation steps: " + << grad_accum_steps; + } + + auto base_model = model_provider(); + CHECK(base_model) << "model_provider returned a null model"; + const auto model_config = base_model->Config(); + const bool is_moe = model_config.ffn_type == nn::FFNType::kMoE; + if (is_moe) { + CHECK_LE(options.zero_stage, 2) << "MoE supports ZeRO stages 0, 1 and 2"; + if (options.expert_parallel > 1) { + CHECK(options.load.empty() && options.save.empty()) + << "MoE training checkpoints currently require EP=1; expert shard " + "metadata is not yet available"; + } + } else { + CHECK_EQ(options.expert_parallel, 1) + << "Expert parallelism requires an MoE model"; + } + 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); + + if (options.synchronize_device) { + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + } + + utils::PrecisionChecker::BuildNameMap(model.get()); + + // Apply LoRA using GetLoRAModel (in-place injection) + bool lora_enabled = options.lora_rank > 0; + if (lora_enabled) { + 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 (!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 + nn::lora::PrintLoRASummary(model, rank.GlobalRank()); + } + + LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device."; + + DataType dtype; + if (options.dtype == "float32") { + dtype = DataType::kFLOAT32; + } else if (options.dtype == "bfloat16") { + dtype = DataType::kBFLOAT16; + } else { + LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << options.dtype + << " not supported."; + } + + 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>{ + {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(), rank_forward_step); + if (ddp_world_size > 1) { + auto ddp_config = + DistributedDataParallelConfig{.handle_unused_parameters = is_moe, + .overlap_grad_reduce = !is_moe, + .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] = 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{.handle_unused_parameters = is_moe, + .overlap_grad_reduce = !is_moe, + .zero_stage = options.zero_stage}; + model = std::make_shared(model, rank, ddp_config); + } + + 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); + } + + auto optimizer = + SetupOptimizer(options, model, ddp_world_size, ddp_rank, pp_world_size); + auto scheduler = SetupLRScheduler(options, optimizer); + + auto train_iter = train_loader.begin(); + 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 = options.load, + .rank = rank, + .model = model, + .optimizer = options.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): 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, + .original_vocab_size = model_config.original_vocab_size, + .padded_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, + .vpp_size = static_cast(options.virtual_pipeline_parallel), + .checkpoint_root_dir = options.save, + .max_checkpoint_keep = options.max_checkpoint_keep, + .rank = rank, + .model = *model, + .optimizer = options.save_optimizer_state ? optimizer.get() : nullptr, + .lr_scheduler = scheduler.get(), + }); + }; + + double training_time_us = 0.0; + int64_t completed_steps = 0; + 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 == options.num_iteration; + + impl->ResetMemPoolHighWatermarks(device); + + if (options.print_timing) { + impl->SynchronizeDevice(device); + } + const auto iter_start = std::chrono::high_resolution_clock::now(); + + // once in a while evaluate the validation dataset + 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 (options.sample_every > 0 && + (step % options.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(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 (options.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 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 = result.loss_func(result.output); + // 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 && !is_moe && + 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"; + } + + nn::parallel::FinalizeModelGrads({model}); + 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, nullptr, dtype); + if (scheduler) { + scheduler->Step(); + } + } + + if (ddp_world_size > 1) { + auto lossf_tensor = std::make_shared( + &lossf, std::vector{}, DataType::kFLOAT32, device); + comm::AllReduce(lossf_tensor, comm::ReduceOpType::kAvg, ddp_pg); + lossf = + static_cast(lossf_tensor->To(Device()).DataPtr())[0]; + } + + if (options.print_timing) { + impl->SynchronizeDevice(device); + } + const auto iter_end = std::chrono::high_resolution_clock::now(); + const double duration_us = + std::chrono::duration(iter_end - iter_start) + .count(); + training_time_us += duration_us; + ++completed_steps; + const double tps = options.total_batch_size / (duration_us / 1e6); + + if (rank.IsLastRank()) { + if (options.log_interval > 0 && ((step + 1) % options.log_interval == 0 || + step + 1 == options.num_iteration)) { + 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, 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 (after_step) { + after_step(*model, device, step + 1); + } + } + + 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(options.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); + } + } + } + + if (options.print_timing && rank.IsLastRank() && completed_steps > 0) { + LOG(ERROR) << "Average training iteration: " + << training_time_us / completed_steps / 1000.0 << " ms"; + } + + // Save LoRA weights if enabled and path specified + 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(options.profile_name + ".report", + Profiler::SortBy::kDeviceTimePercentage); + Profiler::Instance().PrintRecords(options.profile_name + ".records.log"); +#endif + + if (options.synchronize_device) { + impl->SynchronizeDevice(device); + } +} + +} // namespace + +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, options.expert_parallel, + options.expert_tensor_parallel); + utils::PrecisionCheckEnv::Instance().Init(precision_config); + LOG(INFO) << nn::parallel::global::ProcessGroupOverview( + nn::parallel::global::GetDenseRankGenerator(), + nn::parallel::global::GetExpertRankGenerator()); + + 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 < options.nthread_per_process; ++idx) { + threads.emplace_back(run_rank, idx); + } + for (auto &thread : threads) { + thread.join(); + } + } else { + run_rank(0); + } +} + +} // namespace infini_train::training diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e2029460..d986f767 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -44,6 +44,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 00000000..cee64a8b --- /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 00000000..d87e4b89 --- /dev/null +++ b/tests/pretrain/test_model_provider_registry.cc @@ -0,0 +1,88 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +#include "infini_train/include/training/model_provider_registry.h" +#include "infini_train/include/training/training.h" + +namespace infini_train::nn { +class TransformerModel; +} + +using infini_train::training::ModelFactory; +using infini_train::training::ModelProvider; +using infini_train::training::ModelProviderRegistry; + +TEST(ModelProviderRegistryTest, ResolvesWithoutBuildingAndOwnsBoundArguments) { + int calls = 0; + std::string received; + int64_t received_layers = 0; + ModelProvider provider; + { + ModelProviderRegistry registry; + registry.Register( + "unused", [](const std::string &, + const infini_train::nn::ActivationRecomputeOptions &) { + ADD_FAILURE() << "Resolved the wrong model factory"; + return std::shared_ptr{}; + }); + registry.Register( + "selected", + [&](const std::string &weights, + const infini_train::nn::ActivationRecomputeOptions &recompute) { + ++calls; + received = weights; + received_layers = recompute.num_layers; + return std::shared_ptr{}; + }); + std::string weights = "initial.bin"; + infini_train::nn::ActivationRecomputeOptions recompute; + recompute.num_layers = 3; + provider = registry.Resolve("selected", weights, recompute); + recompute.num_layers = 9; + 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"); + EXPECT_EQ(received_layers, 3); + provider(); + EXPECT_EQ(calls, 2); // Each rank invokes its provider independently. +} + +TEST(ModelProviderRegistryTest, RejectsDuplicateAndInvalidRegistrations) { + ModelProviderRegistry registry; + ModelFactory factory = + [](const std::string &, + const infini_train::nn::ActivationRecomputeOptions &) { + 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 &, + const infini_train::nn::ActivationRecomputeOptions &) { + 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 00000000..9ecb1cb9 --- /dev/null +++ b/tests/pretrain/test_training_arguments.cc @@ -0,0 +1,172 @@ +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "gtest/gtest.h" + +#include "infini_train/include/training/arguments.h" +#include "infini_train/include/training/training.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, + NormalizesSequenceBatchAliasesAndParallelOptions) { + auto defaults = + Write("mixtral.flags", + "--batch_size=4\n--sequence_length=64\n--total_batch_size=256\n"); + Parse({"--micro_batch_size=2", "--global_batch_size=16", + "--sequence_length=8", "--expert_parallel=2", + "--expert_tensor_parallel=1", "--activation_recompute", + "--recompute_method=uniform", "--recompute_num_layers=2"}, + defaults); + const auto actual = TrainingOptionsFromFlags(); + EXPECT_EQ(actual.batch_size, 2); + EXPECT_EQ(actual.total_batch_size, 128); + EXPECT_EQ(actual.expert_parallel, 2); + EXPECT_EQ(actual.expert_tensor_parallel, 1); + EXPECT_TRUE(actual.activation_recompute); + EXPECT_EQ(actual.recompute_method, "uniform"); + EXPECT_EQ(actual.recompute_num_layers, 2); +} + +TEST_F(TrainingArgumentsTest, RejectsOverflowInSequenceBatchConversion) { + Parse({"--global_batch_size=4294967295", "--sequence_length=2"}); + EXPECT_THROW(TrainingOptionsFromFlags(), std::invalid_argument); +} + +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