diff --git a/.gitignore b/.gitignore index 4ad6f92ff..7fe542dc5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,8 @@ build/ +/build-cpu/ +/build-cuda/ +/build-ddp/ +/runs/ .cache/ .vscode/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 709bc30f2..8f3a806da 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,6 +104,11 @@ endif() if(USE_CUDA) add_compile_definitions(USE_CUDA=1) + set(CMAKE_CUDA_STANDARD 20) + set(CMAKE_CUDA_STANDARD_REQUIRED ON) + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "75;80;90") + endif() enable_language(CUDA) find_package(CUDAToolkit REQUIRED) include_directories(${CUDAToolkit_INCLUDE_DIRS}) @@ -115,7 +120,7 @@ if(USE_CUDA) file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu) add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS}) - set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90") + set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") target_link_libraries(infini_train_cuda_kernels PUBLIC @@ -240,6 +245,12 @@ add_executable(mnist ) link_infini_train_exe(mnist) +add_executable(mnist_align + example/mnist/align.cc + example/mnist/net.cc +) +link_infini_train_exe(mnist_align) + add_executable(gpt2 example/gpt2/main.cc example/common/tiny_shakespeare_dataset.cc diff --git a/README.md b/README.md index abd8070b2..4d8083b23 100644 --- a/README.md +++ b/README.md @@ -139,10 +139,23 @@ The generated files can be passed directly to the corresponding executables: ##### MNIST +The MNIST example uses a two-layer CNN: +`Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten → Linear(18432,10)`. +It accepts flattened DataLoader batches or NCHW images and returns raw logits. +See [MNIST CNN architecture and tests](docs/mnist_cnn.md) for details. +The reproducible CPU/CUDA training configuration and measured results are in +[MNIST end-to-end training](docs/mnist_training.md). +For one-process-per-GPU NCCL training and real two-GPU integration tests, see +[MNIST distributed training](docs/mnist_ddp.md). + ```bash ./build/mnist \ --device cpu \ - --dataset data/mnist + --dataset data/mnist \ + --bs 64 \ + --num_epoch 3 \ + --lr 0.05 \ + --output_dir runs/mnist-cpu ``` ##### GPT-2 124M diff --git a/cmake/FindNCCL.cmake b/cmake/FindNCCL.cmake index 046d305da..f259e732c 100644 --- a/cmake/FindNCCL.cmake +++ b/cmake/FindNCCL.cmake @@ -40,4 +40,12 @@ find_package_handle_standard_args(NCCL if (NCCL_FOUND) set(NCCL_INCLUDE_DIRS ${NCCL_INCLUDE_DIR}) set(NCCL_LIBRARIES ${NCCL_LIBRARY}) + # Honor NCCL_ROOT for both headers and linkage instead of resolving bare + # -lnccl against an unrelated system installation. + if(NOT TARGET nccl) + add_library(nccl UNKNOWN IMPORTED) + set_target_properties(nccl PROPERTIES + IMPORTED_LOCATION "${NCCL_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${NCCL_INCLUDE_DIR}") + endif() endif() diff --git a/docs/cnn_basic_support.md b/docs/cnn_basic_support.md new file mode 100644 index 000000000..560159523 --- /dev/null +++ b/docs/cnn_basic_support.md @@ -0,0 +1,96 @@ +# CNN 基础训练能力（任务第一步） + +依据：[项目要求：小模型训练支持](https://gxtctab8no8.feishu.cn/docx/QiMHdE5w4omePTx9KBicX6gZnSd)，任务拆解第 1 项。 + +第二步的网络搭建与 MNIST 示例接入见 [MNIST CNN 网络说明](mnist_cnn.md)。 + +## 已实现接口 + +| 能力 | 接口 | 范围 | +| --- | --- | --- | +| 卷积模块 | `nn::Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=true, device=Device())` | 方形卷积核，整数步长、对称零填充，可选偏置 | +| 卷积函数 | `nn::function::Conv2d(input, weight, bias=nullptr, stride=1, padding=0)` | 权重为 OIHW；函数也支持矩形卷积核 | +| 激活 | `nn::ReLU`、`nn::function::ReLU(input)` | 非原地操作，零点导数为 0 | +| 展平模块 | `nn::Flatten(start_dim=1, end_dim=-1)` | 支持负维度，复用 `Tensor::Flatten` 的变形与自动求导 | + +卷积与 ReLU 支持 CPU/CUDA 上的 FP32。卷积输入为连续的 `[N,C,H,W]`，权重为 `[C_out,C_in,K_h,K_w]`，执行与 PyTorch Conv2d 相同的互相关计算。当前不提供 groups、dilation、其他 padding 模式、混合精度或非连续张量接口。卷积要求正的 batch、通道与空间尺寸；Flatten 要求合法的非标量维度区间。 + +输出高度为 `(H + 2 * padding - K_h) / stride + 1`，采用整数向下取整，宽度同理。输入 rank、通道、dtype、设备、bias 形状、stride/padding 和可用空间均在计算前检查。 + +## 框架接入 + +- `nn::Conv2d` 注册可训练的 `weight` 和可选 `bias`，复用 Kaiming/Uniform 初始化、参数枚举、设备迁移和优化器。 +- `autograd::Conv2d` 保存反向所需张量，通过 Dispatcher 调用后端，计算输入、权重和偏置梯度，并跳过不需要的梯度输出。 +- CPU 使用直接卷积，前向、输入梯度和权重梯度可使用 OpenMP；CUDA 每个线程负责一个输出或梯度元素，使用框架当前 stream、DeviceGuard 和 launch 错误检查。无 CPU 回退和浮点原子累加。 +- CPU/CUDA 共享标量索引计算，单元测试使用独立的双精度参考实现验证；该实现优先保证正确性，尚未做 im2col/GEMM 或 cuDNN 性能优化。 +- 新源码和测试接入项目原有的 CMake 源码收集规则。新增文件后需要重新运行 CMake 配置。 +- CUDA 使用 C++20，允许通过 `CMAKE_CUDA_ARCHITECTURES` 指定 GPU 架构；未指定时保留原默认值 `75;80;90`。 + +必要头文件： + +```cpp +#include "infini_train/include/nn/modules/convolution.h" +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/modules/flatten.h" +#include "infini_train/include/nn/functional.h" +``` + +```cpp +auto conv = std::make_shared(1, 16, 3); +auto relu = std::make_shared(); +auto flatten = std::make_shared(); // 保留 batch 维 +auto features = (*flatten)((*relu)((*conv)({input})))[0]; +// [N,1,28,28] -> [N,16,26,26] -> [N,10816] +``` + +与框架现有模块保持一致，参与设备迁移或参数枚举的模块使用 `std::shared_ptr` 管理。 + +## 测试与复现 + +测试源文件：`tests/autograd/test_autograd_cnn.cc`，测试组：`CnnTest`。 + +覆盖内容： + +1. 手算卷积前向、输入/权重/bias 梯度及非均匀上游梯度。 +2. 多 batch、多输入/输出通道、矩形输入/卷积核、stride/padding、有无 bias。前向与独立双精度参考计算比较；三类梯度逐元素做中心有限差分检查。 +3. 全部七种可训练输入/权重/bias 组合，无 bias 模块、1×1 卷积。 +4. ReLU 的正负数、零点导数、非原地行为和 NaN/无穷值前向。 +5. Flatten 的默认/负维度/局部/全维展平、数值顺序与反向形状恢复。 +6. 两层 Conv2d + ReLU + Flatten + Linear + CrossEntropyLoss 的四步 SGD。检查六组参数的梯度非零、有限，每个元素满足 SGD 更新公式，loss 下降且 ZeroGrad 生效。 +7. NoGrad 推理、重复反向的梯度累积，以及非法参数失败检查。 + +前向及参数更新比较使用绝对误差 `1e-5`，有限差分步长 `1e-4`、梯度绝对误差 `2e-5`。测试数据采用可精确表示的二进制小数以降低参考输入转换误差。 + +在 Linux/WSL 中，从项目根目录运行： + +```bash +# 若使用 Git checkout，先补齐仓库声明的 third_party 子模块。 +git submodule update --init --recursive + +cmake -S . -B build-cpu -DUSE_CUDA=OFF -DUSE_NCCL=OFF \ + -DUSE_OMP=ON -DBUILD_TEST=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build-cpu --target test_autograd_cpu -j 6 +OMP_NUM_THREADS=2 ./build-cpu/tests/autograd/test_autograd_cpu --gtest_filter='*CnnTest*' + +cmake -S . -B build-cuda -DUSE_CUDA=ON -DUSE_NCCL=OFF \ + -DUSE_OMP=ON -DBUILD_TEST=ON -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=native +cmake --build build-cuda --target test_autograd_cuda -j 6 +OMP_NUM_THREADS=2 ./build-cuda/tests/autograd/test_autograd_cuda --gtest_filter='CUDA/CnnTest.*' +``` + +CMake 4.x 配合较早版本的第三方依赖时，可增加 `-DCMAKE_POLICY_VERSION_MINIMUM=3.5`。未在 PATH 中的 nvcc 可用 `-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc` 指定。CUDA 参数错误测试复用 CPU 检查并在 CUDA 测试组中跳过，避免 fork 已初始化的 CUDA context。 + +本次工作止于基础算子、模块与训练链路单测；MNIST 网络改造、完整数据集训练和 PyTorch 端到端对齐属于文档后续任务。 + +## 本机验证记录 + +验证环境：WSL Ubuntu 26.04、GCC 15.2、CMake 4.2.3、CUDA 13.3，GPU 为 NVIDIA GeForce RTX 5060 Laptop GPU（编译参数 `CMAKE_CUDA_ARCHITECTURES=120`），启用 OpenMP。 + +原压缩包中的第三方目录为空，本机从原依赖上游补齐了 glog 0.7.1、gflags 2.2.2、GoogleTest 1.17.0 和 InfiniTensor/eigen-mirror（5.0.1-dev），未修改这些依赖的源码。 + +- CPU 独立构建成功，新增 CNN 测试 8/8 通过。 +- CPU 自动求导回归：141 项中 140 项通过，1 项原有 BF16 测试按现有条件跳过，无失败。 +- CUDA 构建和链接成功，包含新增 Conv2d/ReLU 内核及测试程序。 +- 本机曾因 Windows 驱动与 CUDA 13.3 运行库不兼容而无法执行 GPU 测试。最终验收改在驱动 570.124.06、CUDA 12.8、RTX 4090 服务器上完成。 +- 服务器实测：CPU 的 8 项 `CnnTest` 全部通过；CUDA 的 7 项计算/训练测试全部通过，1 项非法参数测试因与 CPU 共用校验逻辑按设计跳过。没有失败项。 diff --git a/docs/mnist_cnn.md b/docs/mnist_cnn.md new file mode 100644 index 000000000..ef733817d --- /dev/null +++ b/docs/mnist_cnn.md @@ -0,0 +1,81 @@ +# MNIST CNN 网络（任务第二步） + +依据：[小模型训练支持项目要求](https://gxtctab8no8.feishu.cn/docx/QiMHdE5w4omePTx9KBicX6gZnSd)，任务拆解第 2 项。网络直接采用文档中的两层卷积结构，并复用第一步的基础能力。 + +## 网络结构 + +| 模块名 | 配置 | 输出形状 | +| --- | --- | --- | +| 输入 | FP32，像素除以 255 | `[N,1,28,28]` | +| `conv1` | Conv2d(1,16,3)，stride=1，padding=0，bias=true | `[N,16,26,26]` | +| `relu1` | ReLU | `[N,16,26,26]` | +| `conv2` | Conv2d(16,32,3)，stride=1，padding=0，bias=true | `[N,32,24,24]` | +| `relu2` | ReLU | `[N,32,24,24]` | +| `flatten` | Flatten(1,-1) | `[N,18432]` | +| `fc` | Linear(18432,10)，bias=true | `[N,10]` | + +共 **189,130** 个可训练参数，分为 `conv1.weight/bias`、`conv2.weight/bias`、`fc.weight/bias` 六组。输出为原始 logits，直接交给现有 CrossEntropyLoss；网络末尾不额外使用 Softmax。预测时取 logits 最大值对应的类别。 + +`example/mnist/net.h` / `net.cc` 中的 `MNIST` 类保留原名称，替换原 MLP。网络接受两种批量输入： + +- `[N,784]`：适配现有 DataLoader，在模型入口通过可求导的 View 恢复为 NCHW。 +- `[N,1,28,28]`：直接用于图像输入和后续参考实现对齐。 + +其他形状、非 FP32 输入和空 batch 会明确报错。batch 大小不固定，支持末批数量不足的情况。参数初始化沿用 Conv2d/Linear 的默认实现。 + +## 示例接入与必要修复 + +- `example/mnist/main.cc` 使用 `shared_ptr` 管理网络，使参数枚举、设备迁移与优化器正常工作；通过模块调用接口执行前向。 +- 测试集前向使用 `NoGradGuard`。读取设备回传的指标前同步，避免异步拷贝后提前访问数据。 +- 修复 `MNISTDataset` 的图片切片步长：IDX 中每张图片占 784 字节，但归一化成 FP32 后占 3136 字节。以前仍按 784 字节偏移读取，会错误读取第二张及后续图片；现在按实际 FP32 大小切片。 +- 保留通用 DataLoader 的原有行为，将图像形状适配局限于 MNIST 网络入口。 + +使用示例： + +```cpp +auto network = std::make_shared(); +network->To(device); +infini_train::optimizers::SGD optimizer(network->Parameters(), 0.001f); +infini_train::nn::CrossEntropyLoss loss_fn; +optimizer.ZeroGrad(); +auto logits = (*network)({images})[0]; +auto loss = loss_fn({logits, labels})[0]; +loss->Backward(); +optimizer.Step(); +``` + +这里的 `images` 是已归一化的 FP32 Tensor，`labels` 是类别索引 Tensor，两者已位于所选设备；构建网络之后先迁移设备，再创建优化器。 + +## 验证与运行 + +新增 `tests/mnist/test_mnist.cc`，CMake 注册 `test_mnist_cpu` / `test_mnist_cuda`。测试覆盖逐层输出尺寸、六组参数名称及总量、batch=1/2/3、两种输入布局的结果一致性、原始 logits、无梯度推理、完整模型的 Loss/Backward/SGD、非法输入，以及真实 IDX 格式的合成文件与末批数据接入。 + +```bash +cmake -S . -B build-cpu -DUSE_CUDA=OFF -DUSE_NCCL=OFF \ + -DBUILD_TEST=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build-cpu --target mnist test_mnist_cpu -j 6 +OMP_NUM_THREADS=2 ctest --test-dir build-cpu -R MnistCnnTest --output-on-failure + +# 准备好真实 MNIST 数据后，可继续使用原示例入口。 +OMP_NUM_THREADS=2 ./build-cpu/mnist --device cpu --dataset data/mnist --bs 64 --num_epoch 1 --lr 0.01 +``` + +完整数据集训练采用 `batch_size=64`、`epochs=3`、`lr=0.05`、`seed=42`，配置与结果见 [MNIST 端到端训练](mnist_training.md)。CMake 4.x 配合较旧依赖时需额外传入 `-DCMAKE_POLICY_VERSION_MINIMUM=3.5`。 + +CUDA 构建及环境要求见 [CNN 基础能力说明](cnn_basic_support.md)。已有 CUDA 构建目录可执行： + +```bash +cmake -S . -B build-cuda +cmake --build build-cuda --target mnist test_mnist_cuda -j 6 +OMP_NUM_THREADS=2 ./build-cuda/tests/mnist/test_mnist_cuda --gtest_filter='CUDA/*' +``` + +真实 MNIST 已在 CPU 和 CUDA 上完成端到端训练；PyTorch 数值对齐也已在同一 CUDA 环境完成。第二步的合成数据检查仍只用于验证网络接入，不作为准确率结果。 + +## 本机验证记录 + +验证环境沿用第一步的 WSL / GCC / CUDA 环境。CPU 版 `mnist` 已成功构建，并用三张合成 IDX 图片、batch=2、1 epoch、lr=0.001 跑完训练和评估，覆盖末批大小为 1 的情况，loss 有限、程序正常退出。该检查只证明入口可运行，不代表真实手写数字识别精度。 + +CPU 的 5 项 `MnistCnnTest` 均经 CTest 运行通过：逐层尺寸/参数、输入布局/批量大小、完整网络反向与 SGD、IDX 读取/末批接入、非法输入检查。反向测试确认全部六组参数梯度有限且非零，更新逐元素符合 SGD 公式，更新后的同批次 loss 下降。 + +最终在驱动 570.124.06、CUDA 12.8、RTX 4090 服务器上完成 GPU 验收。CPU 的 7 项 `MnistCnnTest` 全部通过；CUDA 的 6 项计算/训练测试通过，1 项非法输入测试因与 CPU 共用校验逻辑按设计跳过。相同 seed 和超参数下，CPU/CUDA 都完成 3 个 epoch 的真实 MNIST 训练并达到 97.88% 测试准确率，详见 [MNIST 端到端训练](mnist_training.md)。 diff --git a/docs/mnist_ddp.md b/docs/mnist_ddp.md new file mode 100644 index 000000000..6b09b4c2c --- /dev/null +++ b/docs/mnist_ddp.md @@ -0,0 +1,87 @@ +# MNIST CNN 单机多卡 DDP + +在已有 CPU/CUDA CNN 上复用框架的 `DistributedDataParallel`、`ProcessGroup`、 +NCCL 和 `infini_run`，采用一进程一卡，支持 bucketed 和逐参数两条梯度归约路径。 +本 Demo 支持单机多卡；不支持多节点、ZeRO、混合精度或训练中断续训。 +checkpoint 可在单卡或双卡上仅评估。 + +## 构建与启动 + +需要 CUDA、NCCL 开发库和至少两张可见 GPU。RTX 5090 使用 CUDA 12.8 以上， +`CMAKE_CUDA_ARCHITECTURES=120`；其他 GPU 可设置为 `native`。 + +```bash +cmake -S . -B build-ddp -DUSE_CUDA=ON -DUSE_NCCL=ON \ + -DBUILD_TEST=ON -DBUILD_MNIST_DDP_TESTS=ON -DUSE_OMP=ON -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=120 +cmake --build build-ddp --target mnist mnist_align infini_run \ + test_mnist_ddp test_mnist_cpu test_mnist_cuda test_autograd_cpu test_autograd_cuda -j 8 + +# --bs 是每个 rank 的 batch；此处 global batch = 32 * 2 = 64。 +OMP_NUM_THREADS=2 ./build-ddp/infini_run --nproc_per_node=2 ./build-ddp/mnist \ + --ddp=true --device=cuda --dataset=data/mnist --bs=32 \ + --num_epoch=3 --lr=0.05 --seed=42 --threads=8 --output_dir=runs/mnist-ddp + +# 与上述训练使用相同的全局 batch、初始权重、样本集合和学习率。 +./build-ddp/mnist --device=cuda --dataset=data/mnist --bs=64 \ + --num_epoch=3 --lr=0.05 --seed=42 --threads=8 --output_dir=runs/mnist-single +``` + +使用独立 NCCL 安装时设置 `-DNCCL_ROOT=/path/to/nccl`，该目录包含 `include/nccl.h` +和 `lib/libnccl.so`。切换已有 build 的 NCCL 版本时先清除 CMake 缓存中的旧路径： +`cmake -S . -B build-ddp -U NCCL_INCLUDE_DIR -U NCCL_LIBRARY -DNCCL_ROOT=/path/to/nccl`。 +运行 PyTorch 参考脚本时也将对应 `lib` 目录放在 `LD_LIBRARY_PATH` 首位。 + +## 训练语义 + +- `--ddp=false` 默认保持原有 CPU/单卡行为。多进程启动必须显式传 `--ddp=true`。 +- `--ddp_buckets=true` 默认复用框架 bucket reducer；设为 `false` 使用逐参数 AllReduce。 +- `infini_run` 提供 `RANK/WORLD_SIZE/LOCAL_RANK/LOCAL_WORLD_SIZE` 和独立 rendezvous ID。 + `LOCAL_RANK` 映射可见 GPU，进程组使用 DP 轴，不创建 TP/PP 组。 +- 先在 CPU 确定性初始化，迁移到本卡，再将 rank 0 的全部参数广播到所有 rank。 +- 每个 epoch 使用同一 `seed+epoch` 生成全局排列，rank r 读取位置 r、r+world_size 等样本。 + 训练只丢弃总样本数对 world_size 的余数，每个 rank 样本数相等。 + 不足 `--bs` 的本地尾批仍执行；各卡尾批大小相同，NCCL 平均本地 mean-loss 梯度， + 因而等价于全局 batch 的 mean-loss 梯度，学习率不额外乘除卡数。 +- MNIST 60,000 个训练样本在双卡上不丢弃也不重复；每卡 30,000 个，最后本地 batch=16， + 全局尾批=32。每个 epoch 938 次 SGD，3 epochs 共 2,814 次。 +- 测试集不截断、不补齐；允许不等长或空评估分片。FP64 AllReduce 汇总 loss_sum、correct、 + samples，再计算全局均值，避免平均每卡 accuracy/loss 引入偏差。 +- 只由 rank 0 输出全局 JSONL 与 checkpoint；每个 rank 输出独立配置。结束前广播参数副本， + 校验所有 189,130 个参数有限且与 rank 0 逐字节相等，写入 `rank*/replica_check.txt`。 + 所有 rank 在保存完成后同步退出。非空输出目录会报错。 + +## 测试 + +```bash +OMP_NUM_THREADS=2 ctest --test-dir build-ddp -R '^mnist_ddp_' --output-on-failure -V +./build-ddp/tests/autograd/test_autograd_cpu --gtest_filter='CPU/CnnTest.*' +./build-ddp/tests/autograd/test_autograd_cuda --gtest_filter='CUDA/CnnTest.*' +./build-ddp/tests/mnist/test_mnist_cpu --gtest_filter='CPU/MnistCnnTest.*:MnistDistributedSampler.*' +./build-ddp/tests/mnist/test_mnist_cuda --gtest_filter='CUDA/MnistCnnTest.*' +``` + +两个 DDP CTest 均实际启动两个进程/两张 GPU，分别验证 bucket 与逐参数路径。 +测试故意用不同 seed 初始化两个 rank，验证参数广播；使用 batch=6、2、6 连续三步， +检查 logits、全局 loss、输入梯度、六组参数梯度、SGD 后参数，并覆盖不均匀/空评估分片。 +数值判据为 `abs(candidate-reference) <= atol + rtol*abs(reference)`：logits 为 +`1e-5/1e-4`，梯度为 `1e-6/1e-3`，更新参数为 `1e-6/1e-5`。 +采样器单测验证确定性、分片互斥、截断策略、尾批及全评估集覆盖。 +这些测试仅在 `BUILD_MNIST_DDP_TESTS=ON` 时注册（默认 OFF，避免影响单 GPU 的普通 CTest）。 +无双卡环境应只运行 CPU/单卡测试，不将其当作已验证的 DDP。 + +项目报告随附的 `validate_ddp.py`（PR 外）可一键运行上述测试、3 类固定 fixture 的 +PyTorch 单卡/DDP 数值对齐、MNIST 全量训练、双卡重复训练与 checkpoint 重新评估。 +输入梯度没有 DDP 参数 hook，因此导出时除以 world_size，转换为全局 mean-loss 的导数。 +完整训练不要求与单卡逐位相同：不同矩阵 batch 形状和归约次序存在 FP32 舍入差异。 +固定短轨迹的阈值对齐与相同配置的独立重复运行分别检验正确性与可复现性。 + +```bash +./build-ddp/infini_run --nproc_per_node=2 ./build-ddp/mnist \ + --device=cuda --ddp=true --dataset=data/mnist --bs=32 --eval_only=true \ + --checkpoint=runs/mnist-ddp/checkpoint --output_dir=runs/mnist-ddp-eval +``` + +常见错误：缺少 `USE_NCCL` 时重建；`LOCAL_RANK` 超出 GPU 数量时检查 `CUDA_VISIBLE_DEVICES`； +训练样本少于卡数时减少卡数；出现 NCCL 错误时开启 `NCCL_DEBUG=INFO` 并核对动态库实际版本。 +不要用跳过多卡测试或禁用梯度同步替代修复通信环境。 diff --git a/docs/mnist_training.md b/docs/mnist_training.md new file mode 100644 index 000000000..fdc4c74d2 --- /dev/null +++ b/docs/mnist_training.md @@ -0,0 +1,113 @@ +# MNIST CNN 端到端训练 + +本文记录任务第三步的可复现配置与实测结果。训练入口完整执行 +`Dataset → DataLoader → Forward → CrossEntropyLoss → Backward → SGD → Evaluation`， +并输出 epoch、step、training loss、test loss 与 test accuracy。 + +## 数据与模型 + +使用官方 MNIST 的 60,000 张训练图片和 10,000 张测试图片。程序读取解压后的 +IDX 文件，将像素转换为 FP32 并除以 255。服务器实测文件的 SHA256 如下： + +| 文件 | SHA256 | +| --- | --- | +| `train-images-idx3-ubyte` | `ba891046e6505d7aadcbbe25680a0738ad16aec93bde7f9b65e87a2fc25776db` | +| `train-labels-idx1-ubyte` | `65a50cbbf4e906d70832878ad85ccda5333a97f0f4c3dd2ef09a8a9eef7101c5` | +| `t10k-images-idx3-ubyte` | `0fa7898d509279e482958e8ce81c8e77db3f2f8254e26661ceb7762c4d494ce7` | +| `t10k-labels-idx1-ubyte` | `ff7bcfd416de33731a308c3f266cc351222c34898ecbeaf847f06e48f7ec33f2` | + +网络为 `Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten → +Linear(18432,10)`，共 189,130 个可训练参数。详细结构见 +[MNIST CNN 网络](mnist_cnn.md)。 + +## 训练配置 + +| 参数 | 值 | +| --- | --- | +| dtype | FP32 | +| optimizer | SGD，momentum=0，weight_decay=0 | +| learning rate | 0.05 | +| batch size | 64 | +| epochs | 3 | +| seed | 42 | +| shuffle | 每个 epoch 使用 `seed + epoch` 确定性打乱训练集 | +| evaluation | 训练前一次，每个 epoch 后一次 | +| checkpoint | 第 3 个 epoch 后保存，并重新加载复验 | + +实测环境为 NVIDIA GeForce RTX 4090（24564 MiB）、驱动 570.124.06、 +CUDA 12.8、GCC 13.3.0 和 CMake 3.31.4。CUDA 与 CPU 使用相同初始化、样本顺序和超参数。 + +## 训练结果 + +### CUDA + +| epoch | train loss | train accuracy | test loss | test accuracy | train seconds | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | — | — | 2.302524 | 10.21% | — | +| 1 | 0.297060 | 91.173% | 0.111381 | 96.73% | 8.331 | +| 2 | 0.099873 | 97.100% | 0.069132 | 97.88% | 8.306 | +| 3 | 0.069905 | 97.882% | 0.069654 | 97.88% | 8.253 | + +### CPU + +| epoch | train loss | train accuracy | test loss | test accuracy | train seconds | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 0 | — | — | 2.302522 | 10.21% | — | +| 1 | 0.297042 | 91.175% | 0.111398 | 96.73% | 103.777 | +| 2 | 0.099882 | 97.100% | 0.069135 | 97.88% | 103.192 | +| 3 | 0.069932 | 97.885% | 0.069621 | 97.88% | 103.341 | + +两种设备上，测试 loss 均由约 2.3025 降至约 0.0696，测试 accuracy 由 +10.21% 提升至 97.88%。CUDA 与 CPU 的最终 accuracy 相同；最终 test loss +相差约 3.3e-5，符合不同 kernel 浮点归约顺序带来的预期差异。两个最终 +checkpoint 重新加载后均复现各自的最终 test loss 和 97.88% accuracy。 + +## 单元测试 + +| 测试组 | 通过 | 跳过 | 失败 | +| --- | ---: | ---: | ---: | +| CNN 基础能力 CPU | 8 | 0 | 0 | +| CNN 基础能力 CUDA | 7 | 1 | 0 | +| MNIST CNN CPU | 7 | 0 | 0 | +| MNIST CNN CUDA | 6 | 1 | 0 | + +两个 CUDA 跳过项只检查非法参数和非法输入；校验逻辑与 CPU 共用，已由对应 CPU +测试覆盖。CUDA 的 Forward、Backward、SGD、数据读取、末批处理、指标统计和确定性 +初始化路径均实际执行并通过。 + +## 构建和运行 + +CPU 构建： + +```bash +cmake -S . -B build-cpu -DUSE_CUDA=OFF -DUSE_NCCL=OFF \ + -DUSE_OMP=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build-cpu --target mnist -j 4 +./build-cpu/mnist --dataset=data/mnist --device=cpu --bs=64 \ + --num_epoch=3 --lr=0.05 --seed=42 --shuffle=true --threads=16 \ + --log_interval=100 --output_dir=runs/mnist-cpu-seed42 +``` + +CUDA 构建： + +```bash +cmake -S . -B build-cuda -DUSE_CUDA=ON -DUSE_NCCL=OFF \ + -DUSE_OMP=ON -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=native +cmake --build build-cuda --target mnist -j 4 +./build-cuda/mnist --dataset=data/mnist --device=cuda --bs=64 \ + --num_epoch=3 --lr=0.05 --seed=42 --shuffle=true --threads=8 \ + --log_interval=100 --output_dir=runs/mnist-cuda-seed42 +``` + +每个输出目录包含 `config.txt`、`metrics.jsonl` 和 `checkpoint/`。程序拒绝写入 +非空输出目录，避免覆盖已有结果。用保存的 checkpoint 仅评估： + +```bash +./build-cuda/mnist --dataset=data/mnist --device=cuda --bs=64 \ + --eval_only=true --checkpoint=runs/mnist-cuda-seed42/checkpoint \ + --output_dir=runs/mnist-cuda-checkpoint-eval +``` + +`--dataset` 必须包含四个解压后的 IDX 文件；`--device=cuda` 要求以 +`USE_CUDA=ON` 构建。batch size、epoch、learning rate、线程数和日志间隔必须为正值。 diff --git a/example/mnist/align.cc b/example/mnist/align.cc new file mode 100644 index 000000000..1dced7997 --- /dev/null +++ b/example/mnist/align.cc @@ -0,0 +1,138 @@ +// Deterministic numerical-alignment entry point. PyTorch fixture generation +// and comparison scripts are distributed separately from the framework PR. +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "example/mnist/distributed.h" +#include "example/mnist/net.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/optimizer.h" + +DEFINE_string(fixture_dir, "", "Fixture directory exported by the PyTorch reference"); +DEFINE_string(output_dir, "", "New directory for logits, loss, gradients and updated weights"); +DEFINE_string(device, "cpu", "cpu or cuda"); +DEFINE_bool(ddp, false, "Launch one process per GPU using infini_run"); +DEFINE_bool(ddp_buckets, true, "Enable gradient bucketing"); + +using namespace infini_train; +namespace { +void Read(const std::filesystem::path &path, const std::shared_ptr &tensor) { + CHECK(tensor->GetDevice().IsCPU()); + CHECK_EQ(std::filesystem::file_size(path), tensor->SizeInBytes()) << path; + std::ifstream file(path, std::ios::binary); + CHECK(file.is_open()) << path; + file.read(static_cast(tensor->DataPtr()), tensor->SizeInBytes()); + CHECK(file.good()) << "Incomplete fixture: " << path; +} + +void Write(const std::filesystem::path &path, const std::shared_ptr &tensor) { + CHECK(tensor) << "Missing tensor: " << path; + CHECK(tensor->Dtype() == DataType::kFLOAT32); + auto cpu = tensor->To(Device()); + const auto device = tensor->GetDevice(); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *data = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < cpu.NumElements(); ++i) { CHECK(std::isfinite(data[i])) << path << ": " << i; } + std::ofstream file(path, std::ios::binary); + CHECK(file.is_open()) << path; + file.write(reinterpret_cast(data), cpu.SizeInBytes()); + CHECK(file.good()) << "Failed to write " << path; +} +} // namespace + +int main(int argc, char **argv) { + static_assert(std::endian::native == std::endian::little, "Alignment fixtures use little-endian arrays"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + CHECK(!FLAGS_fixture_dir.empty() && !FLAGS_output_dir.empty()); + CHECK(FLAGS_device == "cpu" || FLAGS_device == "cuda"); +#ifndef USE_CUDA + CHECK_EQ(FLAGS_device, "cpu") << "Rebuild with USE_CUDA=ON"; +#endif + const mnist::DistributedContext distributed(FLAGS_ddp, FLAGS_device == "cuda"); + const std::filesystem::path fixture(FLAGS_fixture_dir); + const auto output = FLAGS_ddp + ? std::filesystem::path(FLAGS_output_dir) / ("rank" + std::to_string(distributed.rank)) + : std::filesystem::path(FLAGS_output_dir); + CHECK(!std::filesystem::exists(output) || std::filesystem::is_empty(output)) + << "Use a new output directory; existing results are preserved"; + std::ifstream config(fixture / "case.txt"); + std::string version; + int64_t batch = 0, steps = 0; + float lr = 0; + CHECK(static_cast(config >> version >> batch >> steps >> lr)); + CHECK_EQ(version, "mnist-alignment-v1"); + CHECK(batch > 0 && batch <= 1024); + CHECK(steps > 0 && steps <= 100); + CHECK(std::isfinite(lr) && lr > 0); + auto network = std::make_shared(); + for (const auto &[name, param] : network->NamedParameters()) { Read(fixture / (name + ".bin"), param); } + auto input = std::make_shared(std::vector{batch, 784}, DataType::kFLOAT32); + auto labels = std::make_shared(std::vector{batch}, DataType::kINT64); + Read(fixture / "input.bin", input); + Read(fixture / "labels.bin", labels); + const auto *target = static_cast(labels->DataPtr()); + for (int64_t i = 0; i < batch; ++i) { CHECK(target[i] >= 0 && target[i] < MNIST::kNumClasses); } + CHECK_EQ(batch % distributed.world_size, 0) << "Alignment batch must divide evenly across ranks"; + const int64_t local_batch = batch / distributed.world_size; + const int64_t offset = distributed.rank * local_batch; + input = std::make_shared(*input, offset * 784 * sizeof(float), std::vector{local_batch, 784}); + labels = std::make_shared(*labels, offset * sizeof(int64_t), std::vector{local_batch}); + const auto device = distributed.device; + network->To(device); + auto model = distributed.Wrap(network, FLAGS_ddp_buckets); + input = std::make_shared(input->To(device))->RequiresGrad(); + labels = std::make_shared(labels->To(device)); + optimizers::SGD optimizer(network->NamedParameters(), lr); + nn::CrossEntropyLoss loss_fn; + std::filesystem::create_directories(output); + // Initial snapshots also prove the candidate loaded the fixture unchanged. + for (const auto &[name, param] : network->NamedParameters()) { + Write(output / ("initial." + name + ".bin"), param); + } + Write(output / "input.bin", input); + for (int64_t step = 0; step < steps; ++step) { + optimizer.ZeroGrad(); + input->ZeroGrad(); + auto logits = (*model)({input})[0]; + auto loss = loss_fn({logits, labels})[0]; + loss->Backward(); + const auto prefix = "step" + std::to_string(step) + "."; + Write(output / (prefix + "logits.bin"), logits); + Write(output / (prefix + "loss.bin"), loss); + { + autograd::NoGradGuard no_grad; + // Input gradients are local (DDP only reduces parameter grads). + // Export derivatives of the GLOBAL mean loss for concatenation. + Write(output / (prefix + "input_grad.bin"), input->grad()->Mul(1.0f / distributed.world_size)); + } + for (const auto &[name, param] : network->NamedParameters()) { + Write(output / (prefix + "grad." + name + ".bin"), param->grad()); + } + optimizer.Step(); + for (const auto &[name, param] : network->NamedParameters()) { + Write(output / (prefix + "updated." + name + ".bin"), param); + } + std::cout << "Completed alignment step " << step << " on " << FLAGS_device << std::endl; + } + std::ofstream complete(output / "complete.txt"); + complete << version << '\n' + << batch << ' ' << steps << ' ' << std::setprecision(9) << lr << '\n' + << FLAGS_device << '\n'; + CHECK(complete.good()); + complete.close(); + std::ofstream rank_info(output / "rank.txt"); + rank_info << distributed.rank << ' ' << distributed.world_size << ' ' << local_batch << '\n'; + rank_info.close(); + distributed.Barrier(); + return 0; +} diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d4..4dbfa13bb 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -110,6 +110,9 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) } } image_file_.tensor = std::move(transposed_tensor); + // Tensor views use byte offsets. Images now contain FP32 values, not the + // UINT8 values in the IDX file; advance by a full FP32 image per sample. + image_size_in_bytes_ = 28 * 28 * sizeof(float); } std::pair, std::shared_ptr> diff --git a/example/mnist/dataset.h b/example/mnist/dataset.h index 8d6193fb7..0c7dfd2fe 100644 --- a/example/mnist/dataset.h +++ b/example/mnist/dataset.h @@ -40,6 +40,6 @@ class MNISTDataset : public infini_train::Dataset { SN3PascalVincentFile label_file_; std::vector image_dims_; std::vector label_dims_; - const size_t image_size_in_bytes_ = 0; + size_t image_size_in_bytes_ = 0; const size_t label_size_in_bytes_ = 0; }; diff --git a/example/mnist/distributed.h b/example/mnist/distributed.h new file mode 100644 index 000000000..7efd22a9f --- /dev/null +++ b/example/mnist/distributed.h @@ -0,0 +1,102 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "example/mnist/training.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/utils.h" + +namespace mnist { +// One process per GPU, launched with infini_run. Reuse the framework's NCCL +// process group, parameter broadcast and autograd-driven DDP gradient reducer. +class DistributedContext { +public: + DistributedContext(bool ddp, bool cuda) { + using namespace infini_train; + using namespace nn::parallel; + global::InitAllEnv(1, 1, false, 1, 1); + world_size = global::GetWorldSize(); + rank = global::GetGlobalProcRank(); + global::thread_global_rank = rank; + CHECK(ddp || world_size == 1) << "Multi-process launch requires --ddp=true"; + CHECK(!ddp || cuda) << "DDP requires --device=cuda"; +#if !defined(USE_CUDA) || !defined(USE_NCCL) + CHECK(!ddp) << "DDP requires a build with USE_CUDA=ON and USE_NCCL=ON"; +#endif + device = cuda ? Device(Device::DeviceType::kCUDA, global::GetLocalProcRank()) : Device(); + if (cuda) { + auto *runtime = core::GetDeviceGuardImpl(device.type()); + CHECK_LT(device.index(), runtime->DeviceCount()) << "Not enough visible GPUs for LOCAL_RANK"; + runtime->SetDevice(device); + } + if (ddp && world_size > 1) { + CHECK_EQ(global::GetNnodes(), 1) << "This MNIST demo supports single-node DDP"; + group = ProcessGroupFactory::Instance(device.type()) + ->GetOrCreate(GetDataParallelProcessGroupName(rank), global::GetGroupRanks(global::DP, rank)); + } + } + + std::shared_ptr Wrap(const std::shared_ptr &net, bool buckets) const { + using namespace infini_train::nn::parallel; + if (!group) { + return net; + } + // Do not rely on coincidentally identical per-rank RNG state. + group->Broadcast(net->Parameters(), 0); + DistributedDataParallelConfig config; + config.gradient_bucketing_enabled = buckets; + config.average_in_collective = true; + config.zero_stage = 0; + return std::make_shared(net, Rank(rank, 0, global::GetNprocPerNode(), 1), config); + } + + Metrics Sum(const Metrics &local) const { + using namespace infini_train; + if (!group) { + return local; + } + // Sum sufficient statistics, never unweighted per-rank averages. + auto cpu = std::make_shared(std::vector{3}, DataType::kFLOAT64); + const double values[] + = {local.loss_sum, static_cast(local.correct), static_cast(local.samples)}; + std::memcpy(cpu->DataPtr(), values, sizeof(values)); + auto tensor = std::make_shared(cpu->To(device)); + group->AllReduce(tensor, nn::parallel::function::ReduceOpType::kSum); + auto result = tensor->To(Device()); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *v = static_cast(result.DataPtr()); + return {v[0], static_cast(v[1]), static_cast(v[2])}; + } + + void Barrier() const { (void)Sum(Metrics{}); } + + void VerifyReplicas(const MNIST &net) const { + using namespace infini_train; + if (!group) { + return; + } + for (const auto &[name, parameter] : net.NamedParameters()) { + auto root = std::make_shared(parameter->Dims(), parameter->Dtype(), device); + root->CopyFrom(parameter); + group->Broadcast({root}, 0); + auto actual = parameter->To(Device()), expected = root->To(Device()); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *values = static_cast(actual.DataPtr()); + for (size_t i = 0; i < actual.NumElements(); ++i) { CHECK(std::isfinite(values[i])) << name; } + CHECK_EQ(std::memcmp(actual.DataPtr(), expected.DataPtr(), actual.SizeInBytes()), 0) + << "DDP replica differs from rank 0: " << name; + } + } + + int rank = 0; + int world_size = 1; + infini_train::Device device; + const infini_train::nn::parallel::ProcessGroup *group = nullptr; +}; +} // namespace mnist diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..5933e49a2 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -1,132 +1,188 @@ #include -#include -#include +#include +#include +#include +#include +#include #include #include -#include -#include - +#include +#ifdef USE_OMP +#include +#endif +#include "example/mnist/dataset.h" +#include "example/mnist/distributed.h" +#include "example/mnist/training.h" #include "gflags/gflags.h" #include "glog/logging.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/checkpoint/checkpoint.h" #include "infini_train/include/optimizer.h" -#include "example/mnist/dataset.h" -#include "example/mnist/net.h" - -DEFINE_string(dataset, "", "mnist dataset path"); -DEFINE_int32(bs, 64, "batch size"); -DEFINE_int32(num_epoch, 1, "num epochs"); -DEFINE_double(lr, 0.01, "learning rate"); -DEFINE_string(device, "cpu", "device type (cpu/cuda)"); +DEFINE_string(dataset, "", "Directory containing uncompressed MNIST IDX files"); +DEFINE_int32(bs, 64, "Batch size"); +DEFINE_int32(num_epoch, 3, "Number of complete training epochs"); +DEFINE_double(lr, 0.05, "SGD learning rate, without momentum or weight decay"); +DEFINE_string(device, "cpu", "Device: cpu or cuda"); +DEFINE_uint32(seed, 42, "Initialization and per-epoch shuffle seed"); +DEFINE_bool(shuffle, true, "Shuffle only the training split each epoch"); +DEFINE_int32(threads, 4, "OpenMP thread count"); +DEFINE_int32(log_interval, 100, "Steps between training progress records"); +DEFINE_string(output_dir, "runs/mnist", "New run directory for config, metrics and checkpoint"); +DEFINE_bool(eval_only, false, "Evaluate --checkpoint without training"); +DEFINE_string(checkpoint, "", "Checkpoint directory, used only with --eval_only"); +DEFINE_bool(ddp, false, "Single-node NCCL DDP; launch one process per GPU with infini_run"); +DEFINE_bool(ddp_buckets, true, "Use the framework DDP gradient buckets (false: per-parameter all-reduce)"); using namespace infini_train; - namespace { -constexpr int kNumItersOfOutputDuration = 10; -constexpr int kNumClasses = 10; - -constexpr char kDeviceCPU[] = "cpu"; -constexpr char kDeviceCUDA[] = "cuda"; -}; // namespace - -DEFINE_validator(device, - [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); - -int main(int argc, char *argv[]) { +using Clock = std::chrono::steady_clock; +void WriteRecord(std::ofstream &out, const char *split, int epoch, int64_t step, const mnist::Metrics &m, + double seconds) { + out << std::setprecision(12) << "{\"split\":\"" << split << "\",\"epoch\":" << epoch << ",\"step\":" << step + << ",\"samples\":" << m.samples << ",\"correct\":" << m.correct << ",\"loss\":" << m.Loss() + << ",\"accuracy\":" << m.Accuracy() << ",\"seconds\":" << seconds << "}\n"; + out.flush(); + CHECK(out.good()) << "Failed to write metrics"; + std::cout << std::fixed << std::setprecision(6) << split << " epoch=" << epoch << " step=" << step + << " samples=" << m.samples << " loss=" << m.Loss() << " accuracy=" << m.Accuracy() + << " seconds=" << seconds << std::endl; +} +} // namespace +int main(int argc, char **argv) { gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); - - auto train_dataset = std::make_shared(FLAGS_dataset, true); - DataLoader train_dataloader(train_dataset, FLAGS_bs); - - // TODO(dcj): Add sampler & eval dataloader later. + CHECK(!FLAGS_dataset.empty()) << "--dataset is required"; + CHECK(FLAGS_device == "cpu" || FLAGS_device == "cuda"); + CHECK_GT(FLAGS_bs, 0); + CHECK_GT(FLAGS_threads, 0); + CHECK_GT(FLAGS_log_interval, 0); + CHECK_GT(FLAGS_num_epoch, 0); + CHECK(std::isfinite(FLAGS_lr) && FLAGS_lr > 0); + CHECK(FLAGS_eval_only == !FLAGS_checkpoint.empty()) << "Use --checkpoint together with --eval_only"; +#ifndef USE_CUDA + CHECK_EQ(FLAGS_device, "cpu") << "Rebuild with USE_CUDA=ON"; +#endif +#ifdef USE_OMP + omp_set_dynamic(0); + omp_set_num_threads(FLAGS_threads); +#endif + const mnist::DistributedContext distributed(FLAGS_ddp, FLAGS_device == "cuda"); + const auto device = distributed.device; + const bool main_rank = distributed.rank == 0; + const std::filesystem::path output_dir(FLAGS_output_dir); + CHECK(!FLAGS_output_dir.empty()); + if (main_rank) { + CHECK(!std::filesystem::exists(output_dir) || std::filesystem::is_empty(output_dir)) + << "Choose a new --output_dir to preserve previous results"; + std::filesystem::create_directories(output_dir); + } + distributed.Barrier(); + const auto rank_dir = FLAGS_ddp ? output_dir / ("rank" + std::to_string(distributed.rank)) : output_dir; + std::filesystem::create_directories(rank_dir); + std::ofstream config(rank_dir / "config.txt"); + CHECK(config.is_open()); + config << "model=Conv2d(1,16,3)-ReLU-Conv2d(16,32,3)-ReLU-Flatten-Linear(18432,10)\n" + << "dtype=float32\nnormalization=pixel/255\noptimizer=SGD\nmomentum=0\nweight_decay=0\n" + << "dataset=" << std::filesystem::absolute(FLAGS_dataset) << "\ndevice=" << FLAGS_device + << "\nbatch_size=" << FLAGS_bs << "\nepochs=" << FLAGS_num_epoch << "\nlearning_rate=" << FLAGS_lr + << "\nseed=" << FLAGS_seed << "\nshuffle=" << FLAGS_shuffle << "\nthreads=" << FLAGS_threads + << "\neval_only=" << FLAGS_eval_only << "\ncheckpoint=" << FLAGS_checkpoint << "\nddp=" << FLAGS_ddp + << "\nworld_size=" << distributed.world_size << "\nrank=" << distributed.rank + << "\nlocal_device=" << static_cast(device.index()) + << "\nglobal_batch_size=" << static_cast(FLAGS_bs) * distributed.world_size + << "\nddp_buckets=" << FLAGS_ddp_buckets << '\n'; + config.close(); + std::ofstream metrics; + if (main_rank) { + metrics.open(output_dir / "metrics.jsonl"); + CHECK(metrics.is_open()); + } + auto net = std::make_shared(); + mnist::Initialize(*net, FLAGS_seed); + TrainerState state; + if (FLAGS_eval_only) { + Checkpoint::Load(FLAGS_checkpoint, *net, nullptr, state, nullptr); + } + net->To(device); + auto model = distributed.Wrap(net, FLAGS_ddp_buckets); auto test_dataset = std::make_shared(FLAGS_dataset, false); - DataLoader test_dataloader(test_dataset, FLAGS_bs); - - auto network = MNIST(); - Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); - Device cpu_device = Device(); - network.To(device); - - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); - - for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { - int train_idx = 0; - float total_loss = 0.0; - - const auto epoch_start = std::chrono::high_resolution_clock::now(); - - for (const auto &[image, label] : train_dataloader) { - auto new_image = std::make_shared(image->To(device)); - auto new_label = std::make_shared(label->To(device)); - - auto outputs = network.Forward({new_image}); - optimizer.ZeroGrad(); - - auto loss = loss_fn.Forward({outputs[0], new_label}); - loss[0]->Backward(); - - // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA - // between forward and backward. - auto loss_cpu = loss[0]->To(cpu_device); - float current_loss = static_cast(loss_cpu.DataPtr())[0]; - total_loss += current_loss; - if (train_idx % kNumItersOfOutputDuration == 0) { - LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() - << "] " - << " loss: " << current_loss; - } - - optimizer.Step(); - train_idx += 1; + CHECK_GT(test_dataset->Size(), 0); + auto test_shard = std::make_shared(test_dataset, distributed.rank, distributed.world_size); + DataLoader test_loader(test_shard, FLAGS_bs); + auto evaluate = [&](int epoch) { + const auto start = Clock::now(); + const auto result = distributed.Sum(mnist::Evaluate(*net, test_loader, device)); + CHECK_EQ(result.samples, test_dataset->Size()); + if (main_rank) { + WriteRecord(metrics, "test", epoch, state.global_step, result, + std::chrono::duration(Clock::now() - start).count()); } - - const auto epoch_end = std::chrono::high_resolution_clock::now(); - const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); - - LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", - epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, - train_dataset->Size() / (duration_us / 1e6)); - } - - // TODO(dcj): Add no_grad() context manager later. - std::vector test_losses; - int correct = 0; - int total = 0; - for (const auto &[image, label] : test_dataloader) { - auto new_image = std::make_shared(image->To(device)); - auto new_label = std::make_shared(label->To(device)); - - auto label_cpu = label->To(cpu_device); - auto outputs = network.Forward({new_image}); - auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn.Forward({outputs[0], new_label}); - auto loss_cpu = loss[0]->To(cpu_device); - - const int batch_size = output_cpu.Dims()[0]; - for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; - const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; - const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; - if (output_index == label_index) { - ++correct; + }; + evaluate(0); + if (!FLAGS_eval_only) { + auto train_data = std::make_shared(FLAGS_dataset, true); + CHECK_GT(train_data->Size(), 0); + auto sampler + = std::make_shared(train_data, distributed.rank, distributed.world_size, true); + CHECK_GT(sampler->Size(), 0) << "Training dataset must have at least world_size samples"; + DataLoader train_loader(sampler, FLAGS_bs); + config.open(rank_dir / "config.txt", std::ios::app); + config << "train_samples=" << train_data->Size() << "\ntest_samples=" << test_dataset->Size() + << "\nlocal_train_samples=" << sampler->Size() + << "\ndropped_train_samples=" << train_data->Size() % distributed.world_size << '\n'; + config.close(); + nn::CrossEntropyLoss loss_fn; + optimizers::SGD optimizer(net->Parameters(), FLAGS_lr); + for (int epoch = 1; epoch <= FLAGS_num_epoch; ++epoch) { + sampler->Reset(FLAGS_seed + static_cast(epoch), FLAGS_shuffle); + mnist::Metrics train_metrics; + const auto start = Clock::now(); + for (const auto &[image, label] : train_loader) { + auto x = std::make_shared(image->To(device)); + auto target = std::make_shared(label->To(device)); + optimizer.ZeroGrad(); + auto logits = (*model)({x})[0]; + auto loss = loss_fn({logits, target})[0]; + loss->Backward(); + optimizer.Step(); + mnist::Accumulate(train_metrics, logits, loss, label); + ++state.global_step; + state.consumed_train_samples += image->Dims()[0] * distributed.world_size; + if (state.global_step % FLAGS_log_interval == 0) { + const auto global_metrics = distributed.Sum(train_metrics); + if (main_rank) { + WriteRecord(metrics, "train_progress", epoch, state.global_step, global_metrics, + std::chrono::duration(Clock::now() - start).count()); + } + } } + CHECK_EQ(train_metrics.samples, sampler->Size()); + const auto global_metrics = distributed.Sum(train_metrics); + CHECK_EQ(global_metrics.samples, train_data->Size() / distributed.world_size * distributed.world_size); + if (main_rank) { + WriteRecord(metrics, "train", epoch, state.global_step, global_metrics, + std::chrono::duration(Clock::now() - start).count()); + } + evaluate(epoch); + } + optimizer.ZeroGrad(); + distributed.VerifyReplicas(*net); + if (FLAGS_ddp) { + std::ofstream replica_check(rank_dir / "replica_check.txt"); + replica_check << "All 189130 parameters finite and bitwise equal to rank 0\n"; + CHECK(replica_check.good()); + } + if (main_rank) { + net->To(Device()); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + Checkpoint::Save(output_dir / "checkpoint", *net, nullptr, state, nullptr); } - total += batch_size; - test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); } - const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); - LOG(ERROR) << "Total: " << total << ", Correct: " << correct - << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; - + distributed.Barrier(); + if (main_rank) { + std::cout << "Results: " << std::filesystem::absolute(output_dir) << std::endl; + } gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); - return 0; } diff --git a/example/mnist/net.cc b/example/mnist/net.cc index 501fee7ef..8015ad0cb 100644 --- a/example/mnist/net.cc +++ b/example/mnist/net.cc @@ -7,25 +7,43 @@ #include "glog/logging.h" #include "infini_train/include/nn/modules/activations.h" -#include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/convolution.h" +#include "infini_train/include/nn/modules/flatten.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" namespace nn = infini_train::nn; -MNIST::MNIST() { - std::vector> layers; - layers.push_back(std::make_shared(784, 30)); - layers.push_back(std::make_shared()); - modules_["sequential"] = std::make_shared(std::move(layers)); - modules_["linear2"] = std::make_shared(30, 10); +MNIST::MNIST() : CloneableModule(kType) { + modules_["conv1"] = std::make_shared(1, 16, 3); + modules_["relu1"] = std::make_shared(); + modules_["conv2"] = std::make_shared(16, 32, 3); + modules_["relu2"] = std::make_shared(); + modules_["flatten"] = std::make_shared(); + modules_["fc"] = std::make_shared(kFlattenFeatures, kNumClasses); } std::vector> MNIST::Forward(const std::vector> &x) { CHECK_EQ(x.size(), 1); - auto x1 = (*modules_["sequential"])(x); - auto x2 = (*modules_["linear2"])(x1); - return x2; + CHECK(x[0]); + CHECK(x[0]->Dtype() == infini_train::DataType::kFLOAT32) << "MNIST CNN expects FP32 images"; + const auto &dims = x[0]->Dims(); + CHECK(dims.size() == 2 || dims.size() == 4) << "MNIST CNN expects [N,784] or [N,1,28,28]"; + CHECK_GT(dims[0], 0); + auto image = x[0]; + if (dims.size() == 2) { + CHECK_EQ(dims[1], kImageSize * kImageSize); + image = image->View({dims[0], 1, kImageSize, kImageSize}); + } else { + CHECK_EQ(dims[1], 1); + CHECK_EQ(dims[2], kImageSize); + CHECK_EQ(dims[3], kImageSize); + } + std::vector> outputs{image}; + for (const auto *name : {"conv1", "relu1", "conv2", "relu2", "flatten", "fc"}) { + outputs = (*modules_.at(name))(outputs); + } + return outputs; } diff --git a/example/mnist/net.h b/example/mnist/net.h index 5f4cfa33b..c78eaf139 100644 --- a/example/mnist/net.h +++ b/example/mnist/net.h @@ -9,8 +9,14 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" -class MNIST : public infini_train::nn::Module { +// FP32 CNN for 28x28 grayscale images. Accepts [N,784] (DataLoader) or +// [N,1,28,28] and returns unnormalized [N,10] logits for CrossEntropyLoss. +class MNIST : public infini_train::nn::CloneableModule { public: + static constexpr char kType[] = "MNIST"; + static constexpr int64_t kImageSize = 28; + static constexpr int64_t kNumClasses = 10; + static constexpr int64_t kFlattenFeatures = 32 * 24 * 24; MNIST(); std::vector> diff --git a/example/mnist/training.h b/example/mnist/training.h new file mode 100644 index 000000000..382cd5d9a --- /dev/null +++ b/example/mnist/training.h @@ -0,0 +1,146 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "example/mnist/net.h" +#include "infini_train/include/autograd/grad_mode.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dataloader.h" +#include "infini_train/include/dataset.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/nn/modules/loss.h" + +namespace mnist { +// Keep the shared DataLoader API unchanged. Reset from identity so each epoch's +// sample order depends only on its seed, not the previous epoch's permutation. +class ShuffledDataset : public infini_train::Dataset { +public: + explicit ShuffledDataset(std::shared_ptr source, int rank = 0, int world_size = 1, + bool equal_shards = false) + : source_(std::move(source)), rank_(rank), world_size_(world_size), equal_shards_(equal_shards) { + CHECK(source_); + CHECK_GT(world_size_, 0); + CHECK_GE(rank_, 0); + CHECK_LT(rank_, world_size_); + Reset(0, false); + } + void Reset(uint32_t seed, bool shuffle = true) { + std::vector global_order(source_->Size()); + std::iota(global_order.begin(), global_order.end(), size_t{0}); + if (shuffle) { + std::mt19937 rng(seed); + std::shuffle(global_order.begin(), global_order.end(), rng); + } + // Training drops at most world_size-1 samples, so every rank performs + // the same number of collectives and local mean losses have equal weight. + // Evaluation keeps every sample exactly once, including uneven shards. + const size_t size = equal_shards_ ? global_order.size() / world_size_ * world_size_ : global_order.size(); + order_.clear(); + for (size_t i = rank_; i < size; i += world_size_) { order_.push_back(global_order[i]); } + } + size_t Size() const override { return order_.size(); } + std::pair, std::shared_ptr> + operator[](size_t i) const override { + return (*source_)[order_.at(i)]; + } + const std::vector &order() const { return order_; } + +private: + std::shared_ptr source_; + std::vector order_; + int rank_; + int world_size_; + bool equal_shards_; +}; + +// Explicit CPU initialization avoids dependence on the global RNG or OpenMP +// worker count. The distribution matches the existing Conv2d/Linear defaults. +inline void Initialize(MNIST &net, uint32_t seed) { + using namespace infini_train; + std::mt19937 rng(seed); + for (const auto *layer_name : {"conv1", "conv2", "fc"}) { + auto &layer = net.mutable_module(layer_name); + const auto [fan_in, fan_out] = nn::init::CalculateFanInAndFanOut(layer->parameter("weight")); + const float bound = 1.0f / std::sqrt(static_cast(fan_in)); + for (const auto *name : {"weight", "bias"}) { + const auto ¶m = layer->parameter(name); + CHECK(param->GetDevice().IsCPU()) << "Initialize MNIST before moving it to CUDA"; + auto *p = static_cast(param->DataPtr()); + for (size_t i = 0; i < param->NumElements(); ++i) { + const float uniform = static_cast(rng() >> 8) / 16777216.0f; + p[i] = (2.0f * uniform - 1.0f) * bound; + } + } + } +} + +struct Metrics { + double loss_sum = 0; + int64_t correct = 0; + int64_t samples = 0; + void Add(float mean_loss, int64_t batch_correct, int64_t batch_size) { + CHECK(std::isfinite(mean_loss)) << "Non-finite loss"; + CHECK_GT(batch_size, 0); + CHECK_GE(batch_correct, 0); + CHECK_LE(batch_correct, batch_size); + loss_sum += static_cast(mean_loss) * batch_size; + correct += batch_correct; + samples += batch_size; + } + double Loss() const { + CHECK_GT(samples, 0); + return loss_sum / samples; + } + double Accuracy() const { + CHECK_GT(samples, 0); + return static_cast(correct) / samples; + } +}; + +inline void Accumulate(Metrics &metrics, const std::shared_ptr &logits, + const std::shared_ptr &loss, + const std::shared_ptr &cpu_labels) { + using namespace infini_train; + CHECK(cpu_labels->GetDevice().IsCPU()); + CHECK(cpu_labels->Dtype() == DataType::kUINT8); + auto y = logits->To(Device()); + auto l = loss->To(Device()); + const auto device = logits->GetDevice(); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto n = y.Dims()[0]; + const auto classes = y.Dims()[1]; + const auto *values = static_cast(y.DataPtr()); + const auto *labels = static_cast(cpu_labels->DataPtr()); + CHECK_EQ(cpu_labels->NumElements(), n); + int64_t correct = 0; + for (int64_t i = 0; i < n; ++i) { + const float *row = values + i * classes; + for (int64_t j = 0; j < classes; ++j) { CHECK(std::isfinite(row[j])) << "Non-finite logits"; } + CHECK_LT(labels[i], classes); + correct += (std::max_element(row, row + classes) - row) == labels[i]; + } + metrics.Add(*static_cast(l.DataPtr()), correct, n); +} + +inline Metrics Evaluate(MNIST &net, const infini_train::DataLoader &loader, infini_train::Device device) { + using namespace infini_train; + autograd::NoGradGuard no_grad; + nn::CrossEntropyLoss loss_fn; + Metrics result; + for (const auto &[image, label] : loader) { + auto x = std::make_shared(image->To(device)); + auto target = std::make_shared(label->To(device)); + auto logits = net({x})[0]; + Accumulate(result, logits, loss_fn({logits, target})[0], label); + } + return result; +} +} // namespace mnist diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..0771fddec 100644 --- a/infini_train/include/autograd/activations.h +++ b/infini_train/include/autograd/activations.h @@ -10,6 +10,16 @@ class Tensor; } namespace infini_train::autograd { +class ReLU : public Function { +public: + static constexpr char kType[] = "ReLUFunction"; + ReLU() : Function(kType) {} + std::vector> Forward(const std::vector> &inputs) override; + void SetupContext(const std::vector> &inputs, + const std::vector> &outputs) override; + std::vector> Backward(const std::vector> &grads) override; +}; + class Sigmoid : public Function { public: static constexpr char kType[] = "SigmoidFunction"; diff --git a/infini_train/include/autograd/convolution.h b/infini_train/include/autograd/convolution.h new file mode 100644 index 000000000..463c1cbfd --- /dev/null +++ b/infini_train/include/autograd/convolution.h @@ -0,0 +1,19 @@ +#pragma once + +#include "infini_train/include/autograd/function.h" + +namespace infini_train::autograd { +class Conv2d : public Function { +public: + static constexpr char kType[] = "Conv2dFunction"; + explicit Conv2d(int64_t stride = 1, int64_t padding = 0) : Function(kType), stride_(stride), padding_(padding) {} + std::vector> Forward(const std::vector> &inputs) override; + void SetupContext(const std::vector> &inputs, + const std::vector> &outputs) override; + std::vector> Backward(const std::vector> &grads) override; + +private: + int64_t stride_, padding_; + bool bias_ = false; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index e4354fd10..05b38ca7e 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -10,6 +10,13 @@ class Tensor; namespace infini_train::nn::function { +// FP32 NCHW cross-correlation, OIHW weights, scalar stride / zero padding, +// dilation=groups=1. Bias is optional. CPU and CUDA support autograd. +std::shared_ptr Conv2d(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias = nullptr, int64_t stride = 1, int64_t padding = 0); +// Out-of-place FP32 ReLU; the derivative at zero is zero. +std::shared_ptr ReLU(const std::shared_ptr &input); + // Returns the lower triangular part of a 2D tensor or a batch of matrices. // // The lower triangular part includes elements on and below the specified diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..35b26aa18 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -10,6 +10,13 @@ class Tensor; } namespace infini_train::nn { +class ReLU : public CloneableModule { +public: + static constexpr char kType[] = "ReLU"; + ReLU() : CloneableModule(kType) {} + std::vector> Forward(const std::vector> &inputs) override; +}; + class Sigmoid : public CloneableModule { public: static constexpr char kType[] = "Sigmoid"; diff --git a/infini_train/include/nn/modules/convolution.h b/infini_train/include/nn/modules/convolution.h new file mode 100644 index 000000000..dcba57214 --- /dev/null +++ b/infini_train/include/nn/modules/convolution.h @@ -0,0 +1,22 @@ +#pragma once + +#include "infini_train/include/nn/modules/module.h" + +namespace infini_train::nn { +// Contiguous FP32 NCHW input; square kernel, scalar stride and zero padding. +// Dilation and groups are fixed to 1. Weight layout: [out_channels, in_channels, k, k]. +class Conv2d : public CloneableModule { +public: + static constexpr char kType[] = "Conv2d"; + static constexpr char kParamWeightName[] = "weight"; + static constexpr char kParamBiasName[] = "bias"; + Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, int64_t stride = 1, int64_t padding = 0, + bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &inputs) override; + bool has_bias() const { return bias_; } + +private: + int64_t stride_, padding_; + bool bias_; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/flatten.h b/infini_train/include/nn/modules/flatten.h new file mode 100644 index 000000000..11744bd06 --- /dev/null +++ b/infini_train/include/nn/modules/flatten.h @@ -0,0 +1,16 @@ +#pragma once + +#include "infini_train/include/nn/modules/module.h" + +namespace infini_train::nn { +class Flatten : public CloneableModule { +public: + static constexpr char kType[] = "Flatten"; + explicit Flatten(int64_t start_dim = 1, int64_t end_dim = -1) + : CloneableModule(kType), start_dim_(start_dim), end_dim_(end_dim) {} + std::vector> Forward(const std::vector> &inputs) override; + +private: + int64_t start_dim_, end_dim_; +}; +} // namespace infini_train::nn diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5ea..d48a8f817 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -6,6 +6,30 @@ #include "infini_train/include/tensor.h" namespace infini_train::autograd { +std::vector> ReLU::Forward(const std::vector> &inputs) { + CHECK_EQ(inputs.size(), 1); + CHECK(inputs[0]); + CHECK(inputs[0]->Dtype() == DataType::kFLOAT32) << "ReLU supports FP32 only"; + return {Dispatcher::Instance().Call>({inputs[0]->GetDevice().type(), "ReLUForward"}, + inputs[0])}; +} + +void ReLU::SetupContext(const std::vector> &inputs, + const std::vector> &) { + ctx_.SaveForBackward({inputs[0]}); +} + +std::vector> ReLU::Backward(const std::vector> &grads) { + CHECK_EQ(grads.size(), 1); + const auto input = ctx_.GetSavedTensors()[0]; + CHECK(grads[0]); + CHECK(grads[0]->Dims() == input->Dims()); + CHECK(grads[0]->Dtype() == input->Dtype()); + CHECK(grads[0]->GetDevice() == input->GetDevice()); + return {Dispatcher::Instance().Call>({input->GetDevice().type(), "ReLUBackward"}, input, + grads[0])}; +} + std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { CHECK_EQ(input_tensors.size(), 1); const auto &input = input_tensors[0]; diff --git a/infini_train/src/autograd/convolution.cc b/infini_train/src/autograd/convolution.cc new file mode 100644 index 000000000..be2d0a738 --- /dev/null +++ b/infini_train/src/autograd/convolution.cc @@ -0,0 +1,62 @@ +#include "infini_train/include/autograd/convolution.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Conv2d::Forward(const std::vector> &inputs) { + CHECK(inputs.size() == 2 || inputs.size() == 3); + CHECK_GT(stride_, 0); + CHECK_GE(padding_, 0); + for (const auto &t : inputs) { + CHECK(t); + CHECK(t->Dtype() == DataType::kFLOAT32) << "Conv2d supports FP32 only"; + CHECK(t->GetDevice() == inputs[0]->GetDevice()) << "Conv2d device mismatch"; + for (auto d : t->Dims()) { CHECK_GT(d, 0); } + } + const auto &x = inputs[0]; + const auto &w = inputs[1]; + CHECK_EQ(x->Dims().size(), 4) << "Conv2d expects NCHW input"; + CHECK_EQ(w->Dims().size(), 4) << "Conv2d expects OIHW weight"; + CHECK_EQ(x->Dims()[1], w->Dims()[1]); + CHECK_GE(x->Dims()[2] + 2 * padding_, w->Dims()[2]); + CHECK_GE(x->Dims()[3] + 2 * padding_, w->Dims()[3]); + const auto bias = inputs.size() == 3 ? inputs[2] : nullptr; + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], w->Dims()[0]); + } + return {Dispatcher::Instance().Call>({x->GetDevice().type(), "Conv2dForward"}, x, w, bias, + stride_, padding_)}; +} + +void Conv2d::SetupContext(const std::vector> &inputs, + const std::vector> &) { + ctx_.SaveForBackward({inputs[0], inputs[1]}); + bias_ = inputs.size() == 3; +} + +std::vector> Conv2d::Backward(const std::vector> &grads) { + CHECK_EQ(grads.size(), 1); + CHECK(grads[0]); + const auto saved = ctx_.GetSavedTensors(); + const auto &x = saved[0]; + const auto &w = saved[1]; + const auto &dy = grads[0]; + const std::vector shape + = {x->Dims()[0], w->Dims()[0], (x->Dims()[2] + 2 * padding_ - w->Dims()[2]) / stride_ + 1, + (x->Dims()[3] + 2 * padding_ - w->Dims()[3]) / stride_ + 1}; + CHECK(dy->Dims() == shape); + CHECK(dy->Dtype() == DataType::kFLOAT32); + CHECK(dy->GetDevice() == x->GetDevice()); + const auto &needs = ctx_.needs_input_grad(); + CHECK_EQ(needs.size(), bias_ ? 3 : 2); + auto result = Dispatcher::Instance().Call>>( + {x->GetDevice().type(), "Conv2dBackward"}, x, w, dy, stride_, padding_, static_cast(needs[0]), + static_cast(needs[1]), bias_ && needs[2]); + if (!bias_) { + result.resize(2); + } + return result; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc index 60f60b619..f56713a08 100644 --- a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc +++ b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc @@ -1,6 +1,7 @@ #include "infini_train/src/core/runtime/cuda/cuda_guard_impl.h" #include +#include #include #include @@ -83,7 +84,8 @@ void CudaGuardImpl::SetDevice(Device device) const { int CudaGuardImpl::DeviceCount() const { int device_count = 0; - CUDA_DRIVER_CHECK(cuDeviceGetCount(&device_count)); + // Runtime API initializes the driver even when this is the first CUDA call. + CUDA_CHECK(cudaGetDeviceCount(&device_count)); return device_count; } diff --git a/infini_train/src/kernels/common/conv2d.h b/infini_train/src/kernels/common/conv2d.h new file mode 100644 index 000000000..c102fd9a7 --- /dev/null +++ b/infini_train/src/kernels/common/conv2d.h @@ -0,0 +1,93 @@ +#pragma once + +#include + +#ifdef __CUDACC__ +#define INFINI_CONV_HD __host__ __device__ +#else +#define INFINI_CONV_HD +#endif + +namespace infini_train::kernels { +struct Conv2dShape { + int64_t n, ci, h, w, co, kh, kw, oh, ow, stride, padding; +}; + +// Each output/gradient element has a single writer, including on CUDA. +// This direct implementation favors a small, deterministic baseline over GEMM tuning. +INFINI_CONV_HD inline float Conv2dForwardAt(int64_t i, const float *x, const float *weight, const float *bias, + Conv2dShape s) { + const int64_t ow = i % s.ow, oh = i / s.ow % s.oh, oc = i / (s.ow * s.oh) % s.co; + const int64_t n = i / (s.ow * s.oh * s.co); + float sum = bias ? bias[oc] : 0.0f; + for (int64_t ic = 0; ic < s.ci; ++ic) { + for (int64_t kh = 0; kh < s.kh; ++kh) { + const int64_t ih = oh * s.stride - s.padding + kh; + if (ih < 0 || ih >= s.h) { + continue; + } + for (int64_t kw = 0; kw < s.kw; ++kw) { + const int64_t iw = ow * s.stride - s.padding + kw; + if (iw >= 0 && iw < s.w) { + sum += x[((n * s.ci + ic) * s.h + ih) * s.w + iw] + * weight[((oc * s.ci + ic) * s.kh + kh) * s.kw + kw]; + } + } + } + } + return sum; +} + +INFINI_CONV_HD inline float Conv2dInputGradAt(int64_t i, const float *weight, const float *dy, Conv2dShape s) { + const int64_t iw = i % s.w, ih = i / s.w % s.h, ic = i / (s.w * s.h) % s.ci; + const int64_t n = i / (s.w * s.h * s.ci); + float sum = 0.0f; + for (int64_t oc = 0; oc < s.co; ++oc) { + for (int64_t kh = 0; kh < s.kh; ++kh) { + const int64_t ph = ih + s.padding - kh; + if (ph < 0 || ph % s.stride != 0 || ph / s.stride >= s.oh) { + continue; + } + for (int64_t kw = 0; kw < s.kw; ++kw) { + const int64_t pw = iw + s.padding - kw; + if (pw >= 0 && pw % s.stride == 0 && pw / s.stride < s.ow) { + sum += dy[((n * s.co + oc) * s.oh + ph / s.stride) * s.ow + pw / s.stride] + * weight[((oc * s.ci + ic) * s.kh + kh) * s.kw + kw]; + } + } + } + } + return sum; +} + +INFINI_CONV_HD inline float Conv2dWeightGradAt(int64_t i, const float *x, const float *dy, Conv2dShape s) { + const int64_t kw = i % s.kw, kh = i / s.kw % s.kh, ic = i / (s.kw * s.kh) % s.ci; + const int64_t oc = i / (s.kw * s.kh * s.ci); + float sum = 0.0f; + for (int64_t n = 0; n < s.n; ++n) { + for (int64_t oh = 0; oh < s.oh; ++oh) { + const int64_t ih = oh * s.stride - s.padding + kh; + if (ih < 0 || ih >= s.h) { + continue; + } + for (int64_t ow = 0; ow < s.ow; ++ow) { + const int64_t iw = ow * s.stride - s.padding + kw; + if (iw >= 0 && iw < s.w) { + sum += x[((n * s.ci + ic) * s.h + ih) * s.w + iw] * dy[((n * s.co + oc) * s.oh + oh) * s.ow + ow]; + } + } + } + } + return sum; +} + +INFINI_CONV_HD inline float Conv2dBiasGradAt(int64_t oc, const float *dy, Conv2dShape s) { + float sum = 0.0f; + for (int64_t n = 0; n < s.n; ++n) { + for (int64_t p = 0; p < s.oh * s.ow; ++p) { sum += dy[(n * s.co + oc) * s.oh * s.ow + p]; } + } + return sum; +} +} // namespace infini_train::kernels + +#undef INFINI_CONV_HD diff --git a/infini_train/src/kernels/cpu/convolution.cc b/infini_train/src/kernels/cpu/convolution.cc new file mode 100644 index 000000000..15a5f3cf3 --- /dev/null +++ b/infini_train/src/kernels/cpu/convolution.cc @@ -0,0 +1,76 @@ +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/kernels/common/conv2d.h" + +namespace infini_train::kernels::cpu { +namespace { +Conv2dShape Shape(const Tensor &x, const Tensor &w, int64_t stride, int64_t padding) { + const auto &a = x.Dims(); + const auto &b = w.Dims(); + return {a[0], + a[1], + a[2], + a[3], + b[0], + b[2], + b[3], + (a[2] + 2 * padding - b[2]) / stride + 1, + (a[3] + 2 * padding - b[3]) / stride + 1, + stride, + padding}; +} +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &x, const std::shared_ptr &w, + const std::shared_ptr &b, int64_t stride, int64_t padding) { + const auto s = Shape(*x, *w, stride, padding); + auto y = std::make_shared(std::vector{s.n, s.co, s.oh, s.ow}, x->Dtype(), x->GetDevice()); + const auto xp = static_cast(x->DataPtr()); + const auto wp = static_cast(w->DataPtr()); + const auto bp = b ? static_cast(b->DataPtr()) : nullptr; + auto yp = static_cast(y->DataPtr()); + const int64_t count = y->NumElements(); +#ifdef USE_OMP +#pragma omp parallel for +#endif + for (int64_t i = 0; i < count; ++i) { yp[i] = Conv2dForwardAt(i, xp, wp, bp, s); } + return y; +} + +std::vector> Conv2dBackward(const std::shared_ptr &x, const std::shared_ptr &w, + const std::shared_ptr &dy, int64_t stride, int64_t padding, + bool need_x, bool need_w, bool need_b) { + const auto s = Shape(*x, *w, stride, padding); + const auto xp = static_cast(x->DataPtr()); + const auto wp = static_cast(w->DataPtr()); + const auto gp = static_cast(dy->DataPtr()); + std::vector> grads(3); + if (need_x) { + grads[0] = std::make_shared(x->Dims(), x->Dtype(), x->GetDevice()); + auto p = static_cast(grads[0]->DataPtr()); + const int64_t count = x->NumElements(); +#ifdef USE_OMP +#pragma omp parallel for +#endif + for (int64_t i = 0; i < count; ++i) { p[i] = Conv2dInputGradAt(i, wp, gp, s); } + } + if (need_w) { + grads[1] = std::make_shared(w->Dims(), w->Dtype(), w->GetDevice()); + auto p = static_cast(grads[1]->DataPtr()); + const int64_t count = w->NumElements(); +#ifdef USE_OMP +#pragma omp parallel for +#endif + for (int64_t i = 0; i < count; ++i) { p[i] = Conv2dWeightGradAt(i, xp, gp, s); } + } + if (need_b) { + grads[2] = std::make_shared(std::vector{s.co}, x->Dtype(), x->GetDevice()); + auto p = static_cast(grads[2]->DataPtr()); + for (int64_t i = 0; i < s.co; ++i) { p[i] = Conv2dBiasGradAt(i, gp, s); } + } + return grads; +} +} // namespace infini_train::kernels::cpu + +REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, Conv2dForward, infini_train::kernels::cpu::Conv2dForward) +REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, Conv2dBackward, infini_train::kernels::cpu::Conv2dBackward) diff --git a/infini_train/src/kernels/cpu/relu.cc b/infini_train/src/kernels/cpu/relu.cc new file mode 100644 index 000000000..9a46d19a3 --- /dev/null +++ b/infini_train/src/kernels/cpu/relu.cc @@ -0,0 +1,24 @@ +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr ReLUForward(const std::shared_ptr &x) { + auto y = std::make_shared(x->Dims(), x->Dtype(), x->GetDevice()); + const auto xp = static_cast(x->DataPtr()); + auto yp = static_cast(y->DataPtr()); + for (size_t i = 0; i < x->NumElements(); ++i) { yp[i] = xp[i] <= 0.0f ? 0.0f : xp[i]; } + return y; +} + +std::shared_ptr ReLUBackward(const std::shared_ptr &x, const std::shared_ptr &dy) { + auto dx = std::make_shared(x->Dims(), x->Dtype(), x->GetDevice()); + const auto xp = static_cast(x->DataPtr()); + const auto gp = static_cast(dy->DataPtr()); + auto p = static_cast(dx->DataPtr()); + for (size_t i = 0; i < x->NumElements(); ++i) { p[i] = xp[i] <= 0.0f ? 0.0f : gp[i]; } + return dx; +} +} // namespace infini_train::kernels::cpu + +REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, ReLUForward, infini_train::kernels::cpu::ReLUForward) +REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, ReLUBackward, infini_train::kernels::cpu::ReLUBackward) diff --git a/infini_train/src/kernels/cuda/convolution.cu b/infini_train/src/kernels/cuda/convolution.cu new file mode 100644 index 000000000..1c85b09c8 --- /dev/null +++ b/infini_train/src/kernels/cuda/convolution.cu @@ -0,0 +1,100 @@ +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" +#include "infini_train/src/kernels/common/conv2d.h" + +namespace infini_train::kernels::cuda { +namespace { +// One thread reduces one output element. No atomics or host staging are needed. +template +__global__ void ConvKernel(float *out, int64_t count, const float *x, const float *w, const float *b, const float *dy, + Conv2dShape s) { + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < count; + i += static_cast(blockDim.x) * gridDim.x) { + if constexpr (Mode == 0) { + out[i] = Conv2dForwardAt(i, x, w, b, s); + } + if constexpr (Mode == 1) { + out[i] = Conv2dInputGradAt(i, w, dy, s); + } + if constexpr (Mode == 2) { + out[i] = Conv2dWeightGradAt(i, x, dy, s); + } + if constexpr (Mode == 3) { + out[i] = Conv2dBiasGradAt(i, dy, s); + } + } +} + +Conv2dShape Shape(const Tensor &x, const Tensor &w, int64_t stride, int64_t padding) { + const auto &a = x.Dims(); + const auto &b = w.Dims(); + return {a[0], + a[1], + a[2], + a[3], + b[0], + b[2], + b[3], + (a[2] + 2 * padding - b[2]) / stride + 1, + (a[3] + 2 * padding - b[3]) / stride + 1, + stride, + padding}; +} + +template +void Launch(const std::shared_ptr &out, const float *x, const float *w, const float *b, const float *dy, + Conv2dShape s) { + const auto device = out->GetDevice(); + core::DeviceGuard guard(device); + auto stream = dynamic_cast(core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + const int64_t count = out->NumElements(); + if (!count) { + return; + } + const int blocks = static_cast(std::min((count + 255) / 256, 65535)); + ConvKernel<<>>(static_cast(out->DataPtr()), count, x, w, b, dy, s); + CUDA_CHECK(cudaGetLastError()); +} +} // namespace + +std::shared_ptr Conv2dForward(const std::shared_ptr &x, const std::shared_ptr &w, + const std::shared_ptr &b, int64_t stride, int64_t padding) { + const auto s = Shape(*x, *w, stride, padding); + auto y = std::make_shared(std::vector{s.n, s.co, s.oh, s.ow}, x->Dtype(), x->GetDevice()); + Launch<0>(y, static_cast(x->DataPtr()), static_cast(w->DataPtr()), + b ? static_cast(b->DataPtr()) : nullptr, nullptr, s); + return y; +} + +std::vector> Conv2dBackward(const std::shared_ptr &x, const std::shared_ptr &w, + const std::shared_ptr &dy, int64_t stride, int64_t padding, + bool need_x, bool need_w, bool need_b) { + const auto s = Shape(*x, *w, stride, padding); + const auto xp = static_cast(x->DataPtr()); + const auto wp = static_cast(w->DataPtr()); + const auto gp = static_cast(dy->DataPtr()); + std::vector> grads(3); + if (need_x) { + grads[0] = std::make_shared(x->Dims(), x->Dtype(), x->GetDevice()); + Launch<1>(grads[0], nullptr, wp, nullptr, gp, s); + } + if (need_w) { + grads[1] = std::make_shared(w->Dims(), w->Dtype(), w->GetDevice()); + Launch<2>(grads[1], xp, nullptr, nullptr, gp, s); + } + if (need_b) { + grads[2] = std::make_shared(std::vector{s.co}, x->Dtype(), x->GetDevice()); + Launch<3>(grads[2], nullptr, nullptr, nullptr, gp, s); + } + return grads; +} +} // namespace infini_train::kernels::cuda + +REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, Conv2dForward, infini_train::kernels::cuda::Conv2dForward) +REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, Conv2dBackward, infini_train::kernels::cuda::Conv2dBackward) diff --git a/infini_train/src/kernels/cuda/relu.cu b/infini_train/src/kernels/cuda/relu.cu new file mode 100644 index 000000000..7e2b72e44 --- /dev/null +++ b/infini_train/src/kernels/cuda/relu.cu @@ -0,0 +1,44 @@ +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { +__global__ void ReLUKernel(const float *x, const float *dy, float *out, int64_t count) { + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < count; + i += static_cast(blockDim.x) * gridDim.x) { + out[i] = x[i] <= 0.0f ? 0.0f : (dy ? dy[i] : x[i]); + } +} + +std::shared_ptr ApplyReLU(const std::shared_ptr &x, const std::shared_ptr &dy) { + const auto device = x->GetDevice(); + core::DeviceGuard guard(device); + auto out = std::make_shared(x->Dims(), x->Dtype(), device); + const int64_t count = x->NumElements(); + if (!count) { + return out; + } + auto stream = dynamic_cast(core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + const int blocks = static_cast(std::min((count + 255) / 256, 65535)); + ReLUKernel<<>>(static_cast(x->DataPtr()), + dy ? static_cast(dy->DataPtr()) : nullptr, + static_cast(out->DataPtr()), count); + CUDA_CHECK(cudaGetLastError()); + return out; +} +} // namespace + +std::shared_ptr ReLUForward(const std::shared_ptr &x) { return ApplyReLU(x, nullptr); } +std::shared_ptr ReLUBackward(const std::shared_ptr &x, const std::shared_ptr &dy) { + return ApplyReLU(x, dy); +} +} // namespace infini_train::kernels::cuda + +REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, ReLUForward, infini_train::kernels::cuda::ReLUForward) +REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, ReLUBackward, infini_train::kernels::cuda::ReLUBackward) diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index c33e23684..2f1c9913d 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -5,6 +5,7 @@ #include #include "infini_train/include/autograd/activations.h" +#include "infini_train/include/autograd/convolution.h" #include "infini_train/include/autograd/elementwise.h" #include "infini_train/include/autograd/reduction.h" #include "infini_train/include/autograd/softmax.h" @@ -13,6 +14,17 @@ #include "infini_train/include/tensor.h" namespace infini_train::nn::function { +std::shared_ptr Conv2d(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, int64_t stride, int64_t padding) { + auto fn = std::make_shared(stride, padding); + return fn->Apply(bias ? std::vector>{input, weight, bias} + : std::vector>{input, weight})[0]; +} + +std::shared_ptr ReLU(const std::shared_ptr &input) { + return std::make_shared()->Apply({input})[0]; +} + std::shared_ptr Tril(const std::shared_ptr &input, int64_t diagonal) { return std::make_shared(diagonal)->Apply({input})[0]; } diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..a47af5e7b 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -8,6 +8,10 @@ #include "infini_train/include/tensor.h" namespace infini_train::nn { +std::vector> ReLU::Forward(const std::vector> &inputs) { + return std::make_shared()->Apply(inputs); +} + std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { return std::make_shared()->Apply(input_tensors); } diff --git a/infini_train/src/nn/modules/convolution.cc b/infini_train/src/nn/modules/convolution.cc new file mode 100644 index 000000000..8b1d5eb96 --- /dev/null +++ b/infini_train/src/nn/modules/convolution.cc @@ -0,0 +1,36 @@ +#include "infini_train/include/nn/modules/convolution.h" + +#include + +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +Conv2d::Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, int64_t stride, int64_t padding, + bool bias, Device device) + : CloneableModule(kType), stride_(stride), padding_(padding), bias_(bias) { + CHECK_GT(in_channels, 0); + CHECK_GT(out_channels, 0); + CHECK_GT(kernel_size, 0); + CHECK_GT(stride, 0); + CHECK_GE(padding, 0); + device_ = device; + auto weight = std::make_shared(std::vector{out_channels, in_channels, kernel_size, kernel_size}, + DataType::kFLOAT32, device, true); + parameters_[kParamWeightName] = weight; + init::KaimingUniform(weight, std::sqrt(5.0f)); + if (bias) { + auto b = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device, true); + parameters_[kParamBiasName] = b; + const float bound = 1.0f / std::sqrt(static_cast(in_channels * kernel_size * kernel_size)); + init::Uniform(b, -bound, bound); + } +} + +std::vector> Conv2d::Forward(const std::vector> &inputs) { + CHECK_EQ(inputs.size(), 1); + return {function::Conv2d(inputs[0], parameters_.at(kParamWeightName), + bias_ ? parameters_.at(kParamBiasName) : nullptr, stride_, padding_)}; +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/flatten.cc b/infini_train/src/nn/modules/flatten.cc new file mode 100644 index 000000000..9a2df56f3 --- /dev/null +++ b/infini_train/src/nn/modules/flatten.cc @@ -0,0 +1,17 @@ +#include "infini_train/include/nn/modules/flatten.h" + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +std::vector> Flatten::Forward(const std::vector> &inputs) { + CHECK_EQ(inputs.size(), 1); + CHECK(inputs[0]); + const int64_t rank = inputs[0]->Dims().size(); + const int64_t start = start_dim_ < 0 ? start_dim_ + rank : start_dim_; + const int64_t end = end_dim_ < 0 ? end_dim_ + rank : end_dim_; + CHECK_GE(start, 0); + CHECK_GE(end, start); + CHECK_LT(end, rank); + return {inputs[0]->Flatten(start, end)}; +} +} // namespace infini_train::nn diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5b9838781..05914459f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -9,11 +9,16 @@ add_subdirectory(common) # Distributed tests add_subdirectory(distributed) + # Module tests add_subdirectory(module) + # Backend extension contract tests add_subdirectory(backend) +# MNIST tests +add_subdirectory(mnist) + # DataLoader tests add_subdirectory(dataloader) diff --git a/tests/autograd/test_autograd_cnn.cc b/tests/autograd/test_autograd_cnn.cc new file mode 100644 index 000000000..e3eef590c --- /dev/null +++ b/tests/autograd/test_autograd_cnn.cc @@ -0,0 +1,286 @@ +#include +#include +#include + +#include "infini_train/include/autograd/convolution.h" +#include "infini_train/include/autograd/grad_mode.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/modules/activations.h" +#include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/convolution.h" +#include "infini_train/include/nn/modules/flatten.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/optimizer.h" +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class CnnTest : public test::InfiniTrainTest { +protected: + void SetUp() override { + if (GetDevice().IsCUDA()) { + REQUIRE_MIN_DEVICES(1); + } + } + + std::shared_ptr Make(const std::vector &v, const std::vector &shape, bool grad = true) { + auto t = std::make_shared(v.data(), shape, DataType::kFLOAT32, GetDevice()); + t->set_requires_grad(grad); + return t; + } + + std::vector Values(const std::shared_ptr &t) { + auto cpu = t->To(Device()); + core::GetDeviceGuardImpl(GetDevice().type())->SynchronizeDevice(GetDevice()); + const auto p = static_cast(cpu.DataPtr()); + return {p, p + cpu.NumElements()}; + } + + void Near(const std::shared_ptr &t, const std::vector &expected, float tol = 1e-5f) { + ASSERT_NE(t, nullptr); + const auto actual = Values(t); + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { EXPECT_NEAR(actual[i], expected[i], tol) << "index " << i; } + } +}; + +TEST_P(CnnTest, ConvKnownForwardAndBackward) { + auto x = Make({1, 2, 3, 4, 5, 6, 7, 8, 9}, {1, 1, 3, 3}); + auto w = Make({1, 0, 0, -1}, {1, 1, 2, 2}); + auto b = Make({2}, {1}); + auto y = nn::function::Conv2d(x, w, b); + EXPECT_EQ(y->Dims(), (std::vector{1, 1, 2, 2})); + Near(y, {-2, -2, -2, -2}); + y->Backward(Make({1, 2, 3, 4}, y->Dims(), false)); + Near(x->grad(), {1, 2, 0, 3, 3, -2, 0, -3, -4}); + Near(w->grad(), {37, 47, 67, 77}); + Near(b->grad(), {10}); +} + +// Independent double-precision reference, including virtual zero padding. +// Finite differences below validate all three analytical derivatives. +namespace { +std::vector Reference(const std::vector &x, const std::vector &w, const std::vector &b, + int stride, int pad) { + constexpr int n = 2, ci = 2, h = 4, width = 5, co = 2, kh = 2, kw = 3; + const int oh = (h + 2 * pad - kh) / stride + 1; + const int ow = (width + 2 * pad - kw) / stride + 1; + std::vector y(n * co * oh * ow); + for (int batch = 0; batch < n; ++batch) { + for (int oc = 0; oc < co; ++oc) { + for (int row = 0; row < oh; ++row) { + for (int col = 0; col < ow; ++col) { + double sum = b.empty() ? 0.0 : b[oc]; + for (int ic = 0; ic < ci; ++ic) { + for (int r = 0; r < kh; ++r) { + for (int c = 0; c < kw; ++c) { + int ir = row * stride + r - pad, jc = col * stride + c - pad; + if (ir >= 0 && ir < h && jc >= 0 && jc < width) { + sum += x[((batch * ci + ic) * h + ir) * width + jc] + * w[((oc * ci + ic) * kh + r) * kw + c]; + } + } + } + } + y[((batch * co + oc) * oh + row) * ow + col] = sum; + } + } + } + } + return y; +} +} // namespace + +TEST_P(CnnTest, ConvMultiChannelFiniteDifferences) { + for (const auto [stride, pad] : std::vector>{{1, 0}, {2, 1}}) { + for (bool bias : {false, true}) { + SCOPED_TRACE(::testing::Message() << "stride=" << stride << " padding=" << pad << " bias=" << bias); + std::vector xv(80), wv(24), bv(bias ? 2 : 0); + for (int i = 0; i < 80; ++i) { xv[i] = (i % 13 - 6) * 0.125; } + for (int i = 0; i < 24; ++i) { wv[i] = (i % 7 - 3) * 0.0625; } + if (bias) { + bv = {0.125, -0.25}; + } + auto x = Make(std::vector(xv.begin(), xv.end()), {2, 2, 4, 5}); + auto w = Make(std::vector(wv.begin(), wv.end()), {2, 2, 2, 3}); + auto b = bias ? Make(std::vector(bv.begin(), bv.end()), {2}) : nullptr; + auto y = nn::function::Conv2d(x, w, b, stride, pad); + auto expected = Reference(xv, wv, bv, stride, pad); + Near(y, std::vector(expected.begin(), expected.end())); + std::vector upstream(expected.size()); + for (size_t i = 0; i < upstream.size(); ++i) { upstream[i] = (static_cast(i % 9) - 4) * 0.125f; } + y->Backward(Make(upstream, y->Dims(), false)); + auto objective = [&]() { + const auto out = Reference(xv, wv, bv, stride, pad); + return std::inner_product(out.begin(), out.end(), upstream.begin(), 0.0); + }; + auto check = [&](std::vector &v, const std::shared_ptr &grad) { + const auto actual = Values(grad); + ASSERT_EQ(v.size(), actual.size()); + for (size_t i = 0; i < v.size(); ++i) { + const double old = v[i], eps = 1e-4; + v[i] = old + eps; + const double plus = objective(); + v[i] = old - eps; + const double minus = objective(); + v[i] = old; + EXPECT_NEAR(actual[i], (plus - minus) / (2 * eps), 2e-5) << "index " << i; + } + }; + check(xv, x->grad()); + check(wv, w->grad()); + if (bias) { + check(bv, b->grad()); + } + } + } +} + +TEST_P(CnnTest, ConvSelectiveGradientsAndNoBias) { + for (int mask = 1; mask < 8; ++mask) { + auto x = Make(std::vector(9, 1), {1, 1, 3, 3}, mask & 1); + auto w = Make(std::vector(4, 1), {1, 1, 2, 2}, mask & 2); + auto b = Make({0}, {1}, mask & 4); + nn::function::Conv2d(x, w, b)->Backward(Make(std::vector(4, 1), {1, 1, 2, 2}, false)); + EXPECT_EQ(x->grad() != nullptr, static_cast(mask & 1)); + EXPECT_EQ(w->grad() != nullptr, static_cast(mask & 2)); + EXPECT_EQ(b->grad() != nullptr, static_cast(mask & 4)); + } + auto conv = std::make_shared(1, 2, 1, 1, 0, false, GetDevice()); + EXPECT_EQ(conv->Parameters().size(), 1); + EXPECT_FALSE(conv->has_bias()); + conv->parameter("weight")->Fill(1.0f); + auto y = (*conv)({Make({1, 2, 3, 4}, {1, 1, 2, 2})})[0]; + Near(y, {1, 2, 3, 4, 1, 2, 3, 4}); + y->Backward(Make(std::vector(8, 1), y->Dims(), false)); + Near(conv->parameter("weight")->grad(), {10, 10}); +} + +TEST_P(CnnTest, ReLUForwardBackwardAndZero) { + auto x = Make({-2, -1, 0, 1, 2, 3}, {2, 3}); + nn::ReLU relu; + auto y = relu({x})[0]; + Near(y, {0, 0, 0, 1, 2, 3}); + Near(x, {-2, -1, 0, 1, 2, 3}); + y->Backward(Make({1, 2, 3, 4, 5, 6}, {2, 3}, false)); + Near(x->grad(), {0, 0, 0, 4, 5, 6}); + auto special + = nn::function::ReLU(Make({-std::numeric_limits::infinity(), std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}, + {3}, false)); + auto values = Values(special); + EXPECT_EQ(values[0], 0); + EXPECT_TRUE(std::isinf(values[1])); + EXPECT_TRUE(std::isnan(values[2])); +} + +TEST_P(CnnTest, FlattenShapeValuesAndGradient) { + std::vector values(24); + std::iota(values.begin(), values.end(), 0.0f); + for (const auto [start, end] : std::vector>{{1, -1}, {-3, -1}, {1, 2}, {0, -1}}) { + auto x = Make(values, {2, 2, 2, 3}); + nn::Flatten flatten(start, end); + auto y = flatten({x})[0]; + const std::vector expected = start == 0 ? std::vector{24} + : end == 2 ? std::vector{2, 4, 3} + : std::vector{2, 12}; + EXPECT_EQ(y->Dims(), expected); + Near(y, values); + y->Backward(Make(values, y->Dims(), false)); + ASSERT_NE(x->grad(), nullptr); + EXPECT_EQ(x->grad()->Dims(), x->Dims()); + Near(x->grad(), values); + } +} + +TEST_P(CnnTest, SequentialLossBackwardAndOptimizer) { + auto conv = std::make_shared(1, 2, 2); + auto conv2 = std::make_shared(2, 2, 2); + auto linear = std::make_shared(8, 2); + auto net = std::make_shared(std::vector>{ + conv, std::make_shared(), conv2, std::make_shared(), std::make_shared(), + linear}); + net->To(GetDevice()); + conv->parameter("weight")->Fill(0.1f); + conv->parameter("bias")->Fill(0.1f); + conv2->parameter("weight")->Fill(0.1f); + conv2->parameter("bias")->Fill(0.1f); + std::vector weights(16); + for (int i = 0; i < 16; ++i) { weights[i] = (i < 8 ? 1 : -1) * 0.05f; } + linear->parameter("weight")->CopyFrom(Make(weights, {2, 8}, false)); + linear->parameter("bias")->Fill(0.0f); + EXPECT_EQ(net->NamedParameters().size(), 6); + auto x = Make(std::vector(32, 0.5f), {2, 1, 4, 4}, false); + auto labels = Make({1, 1}, {2}, false)->To(DataType::kINT64); + auto target = std::make_shared(labels); + nn::CrossEntropyLoss loss; + optimizers::SGD optimizer(net->Parameters(), 0.02f); + float initial = 0, final = 0; + for (int step = 0; step < 4; ++step) { + optimizer.ZeroGrad(); + auto l = loss({(*net)({x})[0], target})[0]; + final = Values(l)[0]; + if (!step) { + initial = final; + } + l->Backward(); + std::vector> expected; + for (const auto &p : net->Parameters()) { + ASSERT_NE(p->grad(), nullptr); + auto before = Values(p), grad = Values(p->grad()); + bool nonzero = false; + for (size_t i = 0; i < before.size(); ++i) { + EXPECT_TRUE(std::isfinite(grad[i])); + nonzero |= grad[i] != 0; + before[i] -= 0.02f * grad[i]; + } + EXPECT_TRUE(nonzero); + expected.push_back(before); + } + optimizer.Step(); + const auto params = net->Parameters(); + for (size_t i = 0; i < params.size(); ++i) { Near(params[i], expected[i]); } + } + EXPECT_LT(final, initial); + optimizer.ZeroGrad(); + for (const auto &p : net->Parameters()) { EXPECT_EQ(p->grad(), nullptr); } +} + +TEST_P(CnnTest, NoGradInferenceAndAccumulation) { + auto x = Make(std::vector(9, 1), {1, 1, 3, 3}); + auto w = Make(std::vector(4, 1), {1, 1, 2, 2}); + { + autograd::NoGradGuard guard; + auto y = nn::function::ReLU(nn::function::Conv2d(x, w)); + EXPECT_FALSE(y->requires_grad()); + EXPECT_EQ(y->grad_fn(), nullptr); + } + for (int i = 0; i < 2; ++i) { + nn::function::Conv2d(x, w)->Backward(Make(std::vector(4, 1), {1, 1, 2, 2}, false)); + } + Near(w->grad(), {8, 8, 8, 8}); +} + +TEST_P(CnnTest, InvalidArgumentsFailClearly) { + GTEST_FLAG_SET(death_test_style, "threadsafe"); + // CUDA death tests must not fork after a CUDA context has been initialized. + if (GetDevice().IsCUDA()) { + GTEST_SKIP() << "Argument validation is shared with CPU"; + } + auto x = Make(std::vector(9, 1), {1, 1, 3, 3}); + auto w = Make(std::vector(4, 1), {1, 1, 2, 2}); + EXPECT_DEATH(nn::function::Conv2d(x, w, nullptr, 0), "stride_"); + EXPECT_DEATH(nn::function::Conv2d(x, w, nullptr, 1, -1), "padding_"); + EXPECT_DEATH(nn::function::Conv2d(x->View({1, 9}), w), "NCHW"); + EXPECT_DEATH(nn::function::Conv2d(x, Make(std::vector(8, 1), {1, 2, 2, 2})), "Check failed"); + EXPECT_DEATH(nn::function::Conv2d(x, Make(std::vector(16, 1), {1, 1, 4, 4})), "Check failed"); + EXPECT_DEATH(nn::function::Conv2d(x, w, Make({1, 2}, {2})), "Check failed"); + EXPECT_DEATH(nn::function::Conv2d(std::make_shared(x->To(DataType::kFLOAT64)), w), "FP32"); + EXPECT_DEATH(nn::Conv2d(1, 1, 0), "kernel_size"); + EXPECT_DEATH(nn::Flatten(1, 4)({x}), "Check failed"); + EXPECT_DEATH(nn::Flatten(-5)({x}), "Check failed"); +} + +INFINI_TRAIN_REGISTER_TEST(CnnTest); diff --git a/tests/mnist/CMakeLists.txt b/tests/mnist/CMakeLists.txt new file mode 100644 index 000000000..a2f44e7b5 --- /dev/null +++ b/tests/mnist/CMakeLists.txt @@ -0,0 +1,21 @@ +infini_train_add_test_suite(test_mnist + SOURCES + test_mnist.cc + ${PROJECT_SOURCE_DIR}/example/mnist/net.cc + ${PROJECT_SOURCE_DIR}/example/mnist/dataset.cc + TEST_TIMEOUT 60 +) + +if(USE_CUDA AND USE_NCCL) + add_executable(test_mnist_ddp test_mnist_ddp.cc ${PROJECT_SOURCE_DIR}/example/mnist/net.cc) + link_infini_train_exe(test_mnist_ddp) + option(BUILD_MNIST_DDP_TESTS "Register MNIST integration tests requiring two visible CUDA GPUs" OFF) + if(BUILD_MNIST_DDP_TESTS) + foreach(buckets IN ITEMS true false) + add_test(NAME mnist_ddp_${buckets} + COMMAND $ --nproc_per_node=2 $ --ddp_buckets=${buckets}) + set_tests_properties(mnist_ddp_${buckets} PROPERTIES TIMEOUT 120 PROCESSORS 2 + RESOURCE_LOCK mnist_gpu LABELS "mnist;ddp;cuda;requires_2_gpus") + endforeach() + endif() +endif() diff --git a/tests/mnist/test_mnist.cc b/tests/mnist/test_mnist.cc new file mode 100644 index 000000000..e53e8b080 --- /dev/null +++ b/tests/mnist/test_mnist.cc @@ -0,0 +1,336 @@ +#include +#include +#include +#include +#include +#include + +#include "example/mnist/dataset.h" +#include "example/mnist/net.h" +#include "example/mnist/training.h" +#include "infini_train/include/autograd/grad_mode.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dataloader.h" +#include "infini_train/include/nn/modules/loss.h" +#include "infini_train/include/optimizer.h" +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +// Real IDX encoding, deliberately distinct images to catch byte-offset errors +// after the dataset converts its storage from UINT8 to FP32. +class IdxFixture { +public: + IdxFixture() { + path = std::filesystem::temp_directory_path() + / ("infinitrain-mnist-" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count())); + CHECK(std::filesystem::create_directory(path)); + for (const auto *prefix : {"train", "t10k"}) { + std::ofstream images(path / (std::string(prefix) + "-images-idx3-ubyte"), std::ios::binary); + for (uint32_t v : {2051u, 3u, 28u, 28u}) { WriteInt(images, v); } + for (unsigned char v : {0, 127, 255}) { + const std::string image(784, static_cast(v)); + images.write(image.data(), image.size()); + } + std::ofstream labels(path / (std::string(prefix) + "-labels-idx1-ubyte"), std::ios::binary); + for (uint32_t v : {2049u, 3u}) { WriteInt(labels, v); } + const char values[] = {1, 5, 9}; + labels.write(values, sizeof(values)); + CHECK(images.good()); + CHECK(labels.good()); + } + } + ~IdxFixture() { + std::error_code ec; + std::filesystem::remove_all(path, ec); + } + std::filesystem::path path; + +private: + static void WriteInt(std::ofstream &out, uint32_t v) { + for (int shift : {24, 16, 8, 0}) { out.put(static_cast((v >> shift) & 255)); } + } +}; +} // namespace + +class MnistCnnTest : public test::InfiniTrainTest { +protected: + void SetUp() override { +#ifdef USE_CUDA + if (GetDevice().IsCUDA()) { + int count = 0; + const auto status = cudaGetDeviceCount(&count); + if (status != cudaSuccess) { + GTEST_SKIP() << cudaGetErrorString(status); + } + if (!count) { + GTEST_SKIP() << "No CUDA device available"; + } + } +#endif + } + + std::shared_ptr Network() { + auto net = std::make_shared(); + net->To(GetDevice()); + return net; + } + + std::shared_ptr Input(int64_t batch) { + std::vector v(batch * 784); + for (size_t i = 0; i < v.size(); ++i) { v[i] = static_cast(i % 251) / 255.0f; } + return std::make_shared(v.data(), std::vector{batch, 784}, DataType::kFLOAT32, GetDevice()); + } + + std::vector Values(const std::shared_ptr &t) { + auto cpu = t->To(Device()); + core::GetDeviceGuardImpl(GetDevice().type())->SynchronizeDevice(GetDevice()); + auto p = static_cast(cpu.DataPtr()); + return {p, p + cpu.NumElements()}; + } +}; + +TEST_P(MnistCnnTest, LayerShapesAndNamedParameters) { + auto net = Network(); + const std::map> expected_params = { + {"conv1.weight", {16, 1, 3, 3}}, {"conv1.bias", {16}}, {"conv2.weight", {32, 16, 3, 3}}, {"conv2.bias", {32}}, + {"fc.weight", {10, 18432}}, {"fc.bias", {10}}}; + size_t numel = 0; + ASSERT_EQ(net->NamedParameters().size(), expected_params.size()); + for (const auto &[name, p] : net->NamedParameters()) { + EXPECT_EQ(p->Dims(), expected_params.at(name)); + EXPECT_TRUE(p->requires_grad()); + EXPECT_EQ(p->GetDevice(), GetDevice()); + numel += p->NumElements(); + } + EXPECT_EQ(numel, 189130); + autograd::NoGradGuard guard; + auto x = Input(2)->View({2, 1, 28, 28}); + for (const auto &[name, shape] : + std::vector>>{{"conv1", {2, 16, 26, 26}}, + {"relu1", {2, 16, 26, 26}}, + {"conv2", {2, 32, 24, 24}}, + {"relu2", {2, 32, 24, 24}}, + {"flatten", {2, 18432}}, + {"fc", {2, 10}}}) { + x = (*net->mutable_module(name))({x})[0]; + EXPECT_EQ(x->Dims(), shape) << name; + } +} + +TEST_P(MnistCnnTest, FlatAndNchwInputsAgreeForVariableBatches) { + auto net = Network(); + // A constant classifier bias confirms the network returns raw logits, + // without accidentally applying sigmoid or softmax before the loss. + net->mutable_module("fc")->parameter("weight")->Fill(0.0f); + net->mutable_module("fc")->parameter("bias")->Fill(2.0f); + autograd::NoGradGuard guard; + for (int64_t batch : {1, 3}) { + auto flat = Input(batch); + auto a = (*net)({flat})[0]; + auto b = (*net)({flat->View({batch, 1, 28, 28})})[0]; + EXPECT_EQ(a->Dims(), (std::vector{batch, 10})); + EXPECT_EQ(Values(a), Values(b)); + for (float v : Values(a)) { EXPECT_EQ(v, 2.0f); } + EXPECT_FALSE(a->requires_grad()); + EXPECT_EQ(a->grad_fn(), nullptr); + } + // Also compare nonconstant logits so layout errors cannot hide behind a zero head. + net = Network(); + auto flat = Input(2); + EXPECT_EQ(Values((*net)({flat})[0]), Values((*net)({flat->View({2, 1, 28, 28})})[0])); +} + +TEST_P(MnistCnnTest, LossBackwardAndSgdReachEveryLayer) { + auto net = Network(); + for (const auto *name : {"conv1", "conv2"}) { + net->mutable_module(name)->parameter("weight")->Fill(0.01f); + net->mutable_module(name)->parameter("bias")->Fill(0.02f); + } + std::vector weights(10 * 18432); + for (size_t i = 0; i < weights.size(); ++i) { weights[i] = (static_cast(i / 18432) - 4) * 0.0001f; } + auto w = std::make_shared(weights.data(), std::vector{10, 18432}, DataType::kFLOAT32, GetDevice()); + net->mutable_module("fc")->parameter("weight")->CopyFrom(w); + net->mutable_module("fc")->parameter("bias")->Fill(0.0f); + auto x = Input(2)->RequiresGrad(); + const float label_values[] = {3, 7}; + auto labels = std::make_shared(label_values, std::vector{2}, DataType::kFLOAT32, GetDevice()); + labels = std::make_shared(labels->To(DataType::kINT64)); + nn::CrossEntropyLoss loss_fn; + optimizers::SGD optimizer(net->Parameters(), 0.001f); + auto loss = loss_fn({(*net)({x})[0], labels})[0]; + const float before_loss = Values(loss)[0]; + ASSERT_TRUE(std::isfinite(before_loss)); + loss->Backward(); + ASSERT_NE(x->grad(), nullptr); + EXPECT_EQ(x->grad()->Dims(), x->Dims()); + std::map> expected; + for (const auto &[name, p] : net->NamedParameters()) { + ASSERT_NE(p->grad(), nullptr) << name; + EXPECT_EQ(p->grad()->Dims(), p->Dims()); + auto values = Values(p), grad = Values(p->grad()); + bool nonzero = false; + for (size_t i = 0; i < values.size(); ++i) { + ASSERT_TRUE(std::isfinite(grad[i])) << name; + nonzero |= grad[i] != 0.0f; + values[i] -= 0.001f * grad[i]; + } + EXPECT_TRUE(nonzero) << name; + expected[name] = std::move(values); + } + optimizer.Step(); + for (const auto &[name, p] : net->NamedParameters()) { + auto values = Values(p); + for (size_t i = 0; i < values.size(); ++i) { ASSERT_NEAR(values[i], expected.at(name)[i], 1e-6f) << name; } + } + optimizer.ZeroGrad(); + for (const auto &p : net->Parameters()) { EXPECT_EQ(p->grad(), nullptr); } + autograd::NoGradGuard guard; + EXPECT_LT(Values(loss_fn({(*net)({x})[0], labels})[0])[0], before_loss); +} + +TEST_P(MnistCnnTest, IdxImagesAndPartialBatchesReachNetwork) { + IdxFixture files; + auto net = Network(); + autograd::NoGradGuard guard; + for (bool train : {true, false}) { + auto dataset = std::make_shared(files.path.string(), train); + ASSERT_EQ(dataset->Size(), 3); + const std::vector pixel = {0, 127.0f / 255.0f, 1}; + const std::vector expected_labels = {1, 5, 9}; + for (int i = 0; i < 3; ++i) { + auto [image, label] = (*dataset)[i]; + EXPECT_EQ(image->Dims(), (std::vector{28, 28})); + EXPECT_EQ(image->Dtype(), DataType::kFLOAT32); + for (size_t j = 0; j < 784; ++j) { ASSERT_FLOAT_EQ(static_cast(image->DataPtr())[j], pixel[i]); } + EXPECT_EQ(*static_cast(label->DataPtr()), expected_labels[i]); + } + DataLoader loader(dataset, 2); + int consumed = 0; + for (const auto &[images, labels] : loader) { + const int64_t n = consumed == 0 ? 2 : 1; + EXPECT_EQ(images->Dims(), (std::vector{n, 784})); + auto x = std::make_shared(images->To(GetDevice())); + auto logits = (*net)({x})[0]; + EXPECT_EQ(logits->Dims(), (std::vector{n, 10})); + for (float v : Values(logits)) { EXPECT_TRUE(std::isfinite(v)); } + for (int j = 0; j < n; ++j) { + EXPECT_EQ(static_cast(labels->DataPtr())[j], expected_labels[consumed + j]); + } + consumed += n; + } + EXPECT_EQ(consumed, 3); + } +} + +TEST_P(MnistCnnTest, RejectsInvalidImageShapes) { + if (GetDevice().IsCUDA()) { + GTEST_SKIP() << "Validation is shared with CPU"; + } + GTEST_FLAG_SET(death_test_style, "threadsafe"); + auto net = Network(); + auto input = Input(1); + EXPECT_DEATH((*net)({}), "x.size"); + EXPECT_DEATH((*net)({input->View({1, 28, 28})}), "expects"); + EXPECT_DEATH((*net)({input->View({1, 1, 14, 56})}), "dims"); + EXPECT_DEATH((*net)({input->View({1, 2, 14, 28})}), "dims"); + EXPECT_DEATH((*net)({input->View({2, 392})}), "dims"); + auto integer = std::make_shared(input->To(DataType::kINT64)); + EXPECT_DEATH((*net)({integer}), "FP32"); +} + +TEST_P(MnistCnnTest, TrainingMetricsWeightPartialBatchesAndEvaluationDoesNotBuildGraph) { + mnist::Metrics metrics; + metrics.Add(2.0f, 1, 2); + metrics.Add(5.0f, 1, 1); + EXPECT_DOUBLE_EQ(metrics.Loss(), 3.0); + EXPECT_DOUBLE_EQ(metrics.Accuracy(), 2.0 / 3.0); + IdxFixture files; + auto dataset = std::make_shared(files.path.string(), false); + auto net = Network(); + auto &fc = net->mutable_module("fc"); + fc->parameter("weight")->Fill(0.0f); + std::vector bias(10); + std::iota(bias.begin(), bias.end(), 0.0f); + fc->parameter("bias")->CopyFrom( + std::make_shared(bias.data(), std::vector{10}, DataType::kFLOAT32, GetDevice())); + auto a = mnist::Evaluate(*net, DataLoader(dataset, 2), GetDevice()); + auto b = mnist::Evaluate(*net, DataLoader(dataset, 3), GetDevice()); + EXPECT_EQ(a.samples, 3); + EXPECT_EQ(a.correct, 1); + EXPECT_NEAR(a.Loss(), b.Loss(), 1e-6); + double exp_sum = 0; + for (int i = 0; i < 10; ++i) { exp_sum += std::exp(static_cast(i)); } + EXPECT_NEAR(a.Loss(), std::log(exp_sum) - 5.0, 1e-6); + EXPECT_TRUE(autograd::GradMode::IsEnabled()); + for (auto &p : net->Parameters()) { EXPECT_EQ(p->grad(), nullptr); } +} + +TEST_P(MnistCnnTest, InitializationAndShufflingAreSeeded) { + auto a = std::make_shared(); + auto b = std::make_shared(); + mnist::Initialize(*a, 42); + mnist::Initialize(*b, 42); + const auto pa = a->NamedParameters(), pb = b->NamedParameters(); + for (size_t i = 0; i < pa.size(); ++i) { + EXPECT_EQ(pa[i].first, pb[i].first); + const auto *x = static_cast(pa[i].second->DataPtr()); + const auto *y = static_cast(pb[i].second->DataPtr()); + EXPECT_TRUE(std::equal(x, x + pa[i].second->NumElements(), y)); + } + mnist::Initialize(*b, 43); + EXPECT_NE(*static_cast(pa[0].second->DataPtr()), *static_cast(pb[0].second->DataPtr())); + IdxFixture files; + auto dataset = std::make_shared(files.path.string(), true); + mnist::ShuffledDataset sampler(dataset); + sampler.Reset(42); + auto order = sampler.order(); + sampler.Reset(97); + sampler.Reset(42); + EXPECT_EQ(sampler.order(), order); + std::sort(order.begin(), order.end()); + EXPECT_EQ(order, (std::vector{0, 1, 2})); + for (size_t i = 0; i < sampler.Size(); ++i) { + EXPECT_EQ(*static_cast(sampler[i].second->DataPtr()), + *static_cast((*dataset)[sampler.order()[i]].second->DataPtr())); + } + sampler.Reset(42, false); + EXPECT_EQ(sampler.order(), order); +} + +INFINI_TRAIN_REGISTER_TEST(MnistCnnTest); + +// One representative sampler test covers uneven sizes, deterministic reshuffle, +// disjoint shards, full evaluation coverage and an empty evaluation rank. +TEST(MnistDistributedSampler, DisjointDeterministicShardsAndTailPolicy) { + class IndexDataset : public Dataset { + public: + size_t Size() const override { return 11; } + std::pair, std::shared_ptr> operator[](size_t) const override { return {}; } + }; + auto source = std::make_shared(); + mnist::ShuffledDataset full(source), a(source, 0, 2, true), b(source, 1, 2, true); + full.Reset(42); + a.Reset(42); + b.Reset(42); + ASSERT_EQ(a.Size(), 5); + ASSERT_EQ(b.Size(), 5); + for (size_t i = 0; i < 5; ++i) { + EXPECT_EQ(a.order()[i], full.order()[2 * i]); + EXPECT_EQ(b.order()[i], full.order()[2 * i + 1]); + } + auto previous = a.order(); + a.Reset(43); + EXPECT_NE(previous, a.order()); + a.Reset(42); + EXPECT_EQ(previous, a.order()); + // local bs=3 gives batches 3+2 on both ranks, so the tail is retained. + EXPECT_EQ((a.Size() + 2) / 3, (b.Size() + 2) / 3); + mnist::ShuffledDataset eval_a(source, 0, 2), eval_b(source, 1, 2), empty(source, 11, 12); + auto combined = eval_a.order(); + combined.insert(combined.end(), eval_b.order().begin(), eval_b.order().end()); + std::sort(combined.begin(), combined.end()); + EXPECT_EQ(combined, (std::vector{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10})); + EXPECT_EQ(empty.Size(), 0); +} diff --git a/tests/mnist/test_mnist_ddp.cc b/tests/mnist/test_mnist_ddp.cc new file mode 100644 index 000000000..fbc525f95 --- /dev/null +++ b/tests/mnist/test_mnist_ddp.cc @@ -0,0 +1,116 @@ +// Real two-process NCCL integration test, including a partial final batch. +#include +#include +#include +#include +#include +#include +#include + +#include "example/mnist/distributed.h" +#include "gflags/gflags.h" +#include "infini_train/include/optimizer.h" + +DEFINE_bool(ddp_buckets, true, "Test bucketed or per-parameter DDP reduction"); +using namespace infini_train; + +namespace { +std::vector Values(const std::shared_ptr &tensor) { + CHECK(tensor); + auto cpu = tensor->To(Device()); + const auto device = tensor->GetDevice(); + core::GetDeviceGuardImpl(device.type())->SynchronizeDevice(device); + const auto *p = static_cast(cpu.DataPtr()); + return {p, p + cpu.NumElements()}; +} + +void Compare(const std::vector &candidate, const std::vector &reference, float atol, float rtol, + const std::string &name) { + CHECK_EQ(candidate.size(), reference.size()); + for (size_t i = 0; i < candidate.size(); ++i) { + CHECK(std::isfinite(candidate[i]) && std::isfinite(reference[i])) << name; + CHECK_LE(std::abs(candidate[i] - reference[i]), atol + rtol * std::abs(reference[i])) + << name << " element=" << i << " candidate=" << candidate[i] << " reference=" << reference[i]; + } +} +} // namespace + +int main(int argc, char **argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + mnist::DistributedContext context(true, true); + CHECK_EQ(context.world_size, 2); + auto net = std::make_shared(); + auto reference = std::make_shared(); + mnist::Initialize(*net, 42 + context.rank); // Deliberately different before broadcast. + mnist::Initialize(*reference, 42); + net->To(context.device); + reference->To(context.device); + auto model = context.Wrap(net, FLAGS_ddp_buckets); + const auto params = net->NamedParameters(), ref_params = reference->NamedParameters(); + for (size_t i = 0; i < params.size(); ++i) { + CHECK_EQ(params[i].first, ref_params[i].first); + Compare(Values(params[i].second), Values(ref_params[i].second), 0, 0, "broadcast"); + } + optimizers::SGD optimizer(net->Parameters(), 0.01f), ref_optimizer(reference->Parameters(), 0.01f); + nn::CrossEntropyLoss criterion; + for (int batch : {6, 2, 6}) { + const int local = batch / 2; + std::vector data(batch * 784), targets(batch); + for (size_t i = 0; i < data.size(); ++i) { data[i] = static_cast((i * 17 + 5) % 251) / 255; } + for (int i = 0; i < batch; ++i) { targets[i] = (i * 3 + 1) % 10; } + auto x = std::make_shared(data.data() + context.rank * local * 784, std::vector{local, 784}, + DataType::kFLOAT32, context.device) + ->RequiresGrad(); + auto all_x = std::make_shared(data.data(), std::vector{batch, 784}, DataType::kFLOAT32, + context.device) + ->RequiresGrad(); + auto labels_float + = std::make_shared(targets.data(), std::vector{batch}, DataType::kFLOAT32, context.device); + auto all_labels = std::make_shared(labels_float->To(DataType::kINT64)); + auto labels = std::make_shared(*all_labels, context.rank * local * sizeof(int64_t), + std::vector{local}); + optimizer.ZeroGrad(); + ref_optimizer.ZeroGrad(); + auto logits = (*model)({x})[0]; + auto ref_logits = (*reference)({all_x})[0]; + auto ref_values = Values(ref_logits); + std::vector slice(ref_values.begin() + context.rank * local * 10, + ref_values.begin() + (context.rank + 1) * local * 10); + Compare(Values(logits), slice, 1e-5f, 1e-4f, "logits"); + auto loss = criterion({logits, labels})[0], ref_loss = criterion({ref_logits, all_labels})[0]; + mnist::Metrics local_metrics; + local_metrics.Add(Values(loss)[0], 0, local); + auto global_metrics = context.Sum(local_metrics); + CHECK_EQ(global_metrics.samples, batch); + CHECK_LE(std::abs(global_metrics.Loss() - Values(ref_loss)[0]), 1e-6); + loss->Backward(); + ref_loss->Backward(); + auto local_grad = Values(x->grad()), global_grad = Values(all_x->grad()); + for (auto &v : local_grad) { v /= 2; } + slice.assign(global_grad.begin() + context.rank * local * 784, + global_grad.begin() + (context.rank + 1) * local * 784); + Compare(local_grad, slice, 1e-6f, 1e-3f, "input gradient"); + for (size_t i = 0; i < params.size(); ++i) { + Compare(Values(params[i].second->grad()), Values(ref_params[i].second->grad()), 1e-6f, 1e-3f, + params[i].first + " gradient"); + } + optimizer.Step(); + ref_optimizer.Step(); + for (size_t i = 0; i < params.size(); ++i) { + Compare(Values(params[i].second), Values(ref_params[i].second), 1e-6f, 1e-5f, params[i].first + " updated"); + } + } + // Weighted reduction also handles unequal and empty evaluation shards. + auto uneven = context.Sum(context.rank == 0 ? mnist::Metrics{6.0, 2, 3} : mnist::Metrics{4.0, 1, 1}); + CHECK_EQ(uneven.samples, 4); + CHECK_EQ(uneven.correct, 3); + CHECK_EQ(uneven.Loss(), 2.5); + auto empty = context.Sum(context.rank == 0 ? mnist::Metrics{6.0, 2, 3} : mnist::Metrics{}); + CHECK_EQ(empty.samples, 3); + CHECK_EQ(empty.Loss(), 2.0); + context.Barrier(); + std::cout << "PASS rank=" << context.rank << " device=" << static_cast(context.device.index()) + << " buckets=" << FLAGS_ddp_buckets << " broadcast/logits/loss/gradients/SGD/tails/metrics" << std::endl; + return 0; +}