format: format files in examples and infini_train

ArcaLunar · ArcaLunar · commit 9c3e38c8ce30 · 2026-03-17T17:37:17.000+08:00
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
@@ -1,20 +1,20 @@
+#include <algorithm>
 #include <chrono>
 #include <cstdlib>
 #include <filesystem>
 #include <format>
 #include <limits>
 #include <memory>
 #include <optional>
-#include <algorithm>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "gflags/gflags.h"
 #include "glog/logging.h"
 
 #include "infini_train/include/autocast.h"
-#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/checkpoint.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
@@ -328,7 +328,7 @@ void Train(const nn::parallel::Rank &rank) {
     size_t saved_data_batch_idx = train_iter.BatchIndex();
     std::shared_ptr<nn::Module> loss_fn
         = (tp_world_size > 1) ? std::static_pointer_cast<nn::Module>(
-              std::make_shared<VocabParallelCrossEntropyLoss>(model_config.original_vocab_size))
+                                    std::make_shared<VocabParallelCrossEntropyLoss>(model_config.original_vocab_size))
                               : std::static_pointer_cast<nn::Module>(std::make_shared<nn::CrossEntropyLoss>());
     loss_fn->To(device);
     LOG(INFO) << "Rank " << rank.GlobalRank() << ": start training";
@@ -376,7 +376,8 @@ void Train(const nn::parallel::Rank &rank) {
 
     LOG(INFO) << "start training";
 
-    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
+    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step,
+                               bool prune_step_checkpoints) {
         const auto ckpt_start = std::chrono::high_resolution_clock::now();
 
         TrainerState state;
@@ -413,8 +414,7 @@ void Train(const nn::parallel::Rank &rank) {
                 const auto root = std::filesystem::path(FLAGS_checkpoint_dir);
                 if (std::filesystem::exists(root)) {
                     for (const auto &entry : std::filesystem::directory_iterator(root)) {
-                        if (entry.is_directory()
-                            && entry.path().filename().string().starts_with("checkpoint_step_")) {
+                        if (entry.is_directory() && entry.path().filename().string().starts_with("checkpoint_step_")) {
                             ckpts.push_back(entry.path());
                         }
                     }
diff --git a/example/gpt2/net.cc b/example/gpt2/net.cc
@@ -720,10 +720,8 @@ std::shared_ptr<GPT2> GPT2::FromLLMC(const std::string &filepath) {
 }
 
 void GPT2::SaveAsLLMC(const std::string &filepath) const {
-    CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1)
-        << "SaveAsLLMC currently supports TP=1 only.";
-    CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1)
-        << "SaveAsLLMC currently supports PP=1 only.";
+    CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1) << "SaveAsLLMC currently supports TP=1 only.";
+    CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1) << "SaveAsLLMC currently supports PP=1 only.";
 
     std::ofstream ofs(filepath, std::ios::binary);
     CHECK(ofs.is_open()) << "Failed to open model file for write: " << filepath;
@@ -737,7 +735,8 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
     header[5] = static_cast<int32_t>(config_.n_head);
     header[6] = static_cast<int32_t>(config_.n_embd);
     header[7] = static_cast<int32_t>(config_.vocab_size);
-    ofs.write(reinterpret_cast<const char *>(header.data()), static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
+    ofs.write(reinterpret_cast<const char *>(header.data()),
+              static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
 
     const auto state_dict = StateDict();
     auto get_tensor = [&](const std::string &name) -> std::shared_ptr<Tensor> {
@@ -763,12 +762,12 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
                                              nn::Embedding::kParamWeightName)));
 
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
-                                                 idx, Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
+                                                 Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
-                                                 idx, Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
+                                                 Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
         write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
@@ -791,12 +790,12 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
                                                  nn::parallel::RowParallelLinear::kParamBiasName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
-                                                 idx, Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
+                                                 Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
-                                                 idx, Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
+                                                 Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
         write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
@@ -819,10 +818,10 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
                                                  nn::parallel::RowParallelLinear::kParamBiasName)));
     }
 
-    write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName,
-                                             nn::LayerNorm::kParamWeightName)));
-    write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName,
-                                             nn::LayerNorm::kParamBiasName)));
+    write_tensor_fp32(get_tensor(
+        std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName, nn::LayerNorm::kParamWeightName)));
+    write_tensor_fp32(get_tensor(
+        std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName, nn::LayerNorm::kParamBiasName)));
 
     ofs.flush();
     CHECK(ofs.good()) << "Failed to flush model file: " << filepath;
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
@@ -1,18 +1,18 @@
+#include <algorithm>
 #include <cstdlib>
 #include <filesystem>
 #include <format>
 #include <limits>
 #include <memory>
 #include <optional>
-#include <algorithm>
 #include <unordered_set>
 
 #include "gflags/gflags.h"
 #include "glog/logging.h"
 
 #include "infini_train/include/autocast.h"
-#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/checkpoint.h"
+#include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/lora/lora_utils.h"
@@ -339,7 +339,8 @@ void Train(const nn::parallel::Rank &rank) {
         }
     }
 
-    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
+    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step,
+                               bool prune_step_checkpoints) {
         const auto ckpt_start = std::chrono::high_resolution_clock::now();
 
         TrainerState state;
@@ -374,8 +375,7 @@ void Train(const nn::parallel::Rank &rank) {
                 const auto root = std::filesystem::path(FLAGS_checkpoint_dir);
                 if (std::filesystem::exists(root)) {
                     for (const auto &entry : std::filesystem::directory_iterator(root)) {
-                        if (entry.is_directory()
-                            && entry.path().filename().string().starts_with("checkpoint_step_")) {
+                        if (entry.is_directory() && entry.path().filename().string().starts_with("checkpoint_step_")) {
                             ckpts.push_back(entry.path());
                         }
                     }
diff --git a/example/llama3/net.cc b/example/llama3/net.cc
@@ -1,8 +1,8 @@
 #include "example/llama3/net.h"
 
 #include <cmath>
-#include <cstring>
 #include <cstdlib>
+#include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <map>
@@ -758,10 +758,8 @@ std::shared_ptr<LLaMA3> LLaMA3::FromLLMC(const std::string &filepath) {
 }
 
 void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
-    CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1)
-        << "SaveAsLLMC currently supports TP=1 only.";
-    CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1)
-        << "SaveAsLLMC currently supports PP=1 only.";
+    CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1) << "SaveAsLLMC currently supports TP=1 only.";
+    CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1) << "SaveAsLLMC currently supports PP=1 only.";
 
     std::ofstream ofs(filepath, std::ios::binary);
     CHECK(ofs.is_open()) << "Failed to open model file for write: " << filepath;
@@ -789,7 +787,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
     header[13] = static_cast<int32_t>(config_.max_gen_batch_size);
     header[14] = 1; // version_major
     header[15] = 0; // version_minor
-    ofs.write(reinterpret_cast<const char *>(header.data()), static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
+    ofs.write(reinterpret_cast<const char *>(header.data()),
+              static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
 
     const auto state_dict = StateDict();
     auto get_tensor = [&](const std::string &name) -> std::shared_ptr<Tensor> {
@@ -810,8 +809,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
                                              nn::parallel::VocabParallelEmbedding::kParamWeightName)));
 
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
-                                                 idx, Block::kLn1LayerName, RMSNorm::kParamWeightName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName, idx,
+                                                 Block::kLn1LayerName, RMSNorm::kParamWeightName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
         write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
@@ -824,8 +823,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
                                                  nn::parallel::RowParallelLinear::kParamWeightName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
-        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
-                                                 idx, Block::kLn2LayerName, RMSNorm::kParamWeightName)));
+        write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName, idx,
+                                                 Block::kLn2LayerName, RMSNorm::kParamWeightName)));
     }
     for (int idx = 0; idx < config_.n_layer; ++idx) {
         write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
@@ -843,10 +842,10 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
                                                  nn::parallel::RowParallelLinear::kParamWeightName)));
     }
 
-    write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, LLaMA3LastStage::kLnFLayerName,
-                                             RMSNorm::kParamWeightName)));
-    write_tensor_fp32(get_tensor(std::format("{}.{}", LLaMA3LastStage::kLMHeadLayerName,
-                                             nn::parallel::ColumnParallelLinear::kParamWeightName)));
+    write_tensor_fp32(get_tensor(
+        std::format("{}.{}.{}", kTransformerLayerName, LLaMA3LastStage::kLnFLayerName, RMSNorm::kParamWeightName)));
+    write_tensor_fp32(get_tensor(
+        std::format("{}.{}", LLaMA3LastStage::kLMHeadLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)));
 
     ofs.flush();
     CHECK(ofs.good()) << "Failed to flush model file: " << filepath;
diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc
@@ -83,8 +83,8 @@ SN3PascalVincentFile ReadSN3PascalVincentFile(const std::string &path) {
 } // namespace
 
 MNISTDataset::MNISTDataset(const std::string &dataset, bool train)
-    : image_file_(
-        ReadSN3PascalVincentFile(std::format("{}/{}-images-idx3-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))),
+    : image_file_(ReadSN3PascalVincentFile(
+          std::format("{}/{}-images-idx3-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))),
       label_file_(ReadSN3PascalVincentFile(
           std::format("{}/{}-labels-idx1-ubyte", dataset, train ? kTrainPrefix : kTestPrefix))),
       image_dims_(image_file_.dims.begin() + 1, image_file_.dims.end()),
diff --git a/infini_train/include/autograd/softmax.h b/infini_train/include/autograd/softmax.h
@@ -15,7 +15,7 @@ class Softmax : public Function {
 public:
     static constexpr char kType[] = "SoftmaxFunction";
 
-    explicit Softmax(int64_t dim = -1) : Function(kType), dim_(dim){};
+    explicit Softmax(int64_t dim = -1) : Function(kType), dim_(dim) {};
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
     void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
diff --git a/infini_train/include/checkpoint.h b/infini_train/include/checkpoint.h
@@ -49,11 +49,11 @@ class Checkpoint {
                      TrainerState *state, const CheckpointLoadOptions &options = {});
 
 private:
-    static void SaveStateDictBinary(
-        const std::filesystem::path &path,
-        const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
+    static void SaveStateDictBinary(const std::filesystem::path &path,
+                                    const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
 
-    static std::unordered_map<std::string, std::shared_ptr<Tensor>> LoadStateDictBinary(const std::filesystem::path &path);
+    static std::unordered_map<std::string, std::shared_ptr<Tensor>>
+    LoadStateDictBinary(const std::filesystem::path &path);
 
     static void SaveTrainerState(const std::filesystem::path &path, const TrainerState &state);
     static TrainerState LoadTrainerState(const std::filesystem::path &path);
diff --git a/infini_train/include/common/common.h b/infini_train/include/common/common.h
@@ -7,7 +7,7 @@
 
 #include "infini_train/include/datatype.h"
 
-#define CEIL_DIV(x, y) (((x) + (y)-1) / (y))
+#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))
 #define LOG_LOC(LEVEL, MSG) LOG(LEVEL) << MSG << " at " << __FILE__ << ":" << __LINE__
 #define LOG_UNSUPPORTED_DTYPE(DTYPE, CONTEXT_IDENTIFIER)                                                               \
     LOG_LOC(FATAL, WRAP(CONTEXT_IDENTIFIER << ": Unsupported data type: "                                              \
diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h
@@ -43,7 +43,7 @@ class Module : public std::enable_shared_from_this<Module> {
     explicit Module(const std::string &type);
     Module(const Module &) = default;
 
-    virtual ~Module(){};
+    virtual ~Module() {};
 
     const std::string &type() const;
 
diff --git a/infini_train/include/nn/parallel/tensor_parallel.h b/infini_train/include/nn/parallel/tensor_parallel.h
@@ -119,7 +119,7 @@ class VocabParallelCrossEntropyLoss : public nn::CloneableModule<VocabParallelCr
 public:
     static constexpr char kType[] = "VocabParallelCrossEntropyLoss";
     VocabParallelCrossEntropyLoss(int64_t vocab_size_original = 0, float label_smoothing = 0.f)
-        : CloneableModule(kType), vocab_size_original_(vocab_size_original), label_smoothing_(label_smoothing){};
+        : CloneableModule(kType), vocab_size_original_(vocab_size_original), label_smoothing_(label_smoothing) {};
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
diff --git a/infini_train/src/checkpoint.cc b/infini_train/src/checkpoint.cc
diff --git a/infini_train/src/dataloader.cc b/infini_train/src/dataloader.cc
diff --git a/infini_train/src/kernels/cuda/concat.cu b/infini_train/src/kernels/cuda/concat.cu
diff --git a/infini_train/src/kernels/cuda/slice.cu b/infini_train/src/kernels/cuda/slice.cu
diff --git a/infini_train/src/kernels/cuda/stack.cu b/infini_train/src/kernels/cuda/stack.cu
