Skip to content

Commit 4b248b6

Browse files
ArcaLunarArcaLunar
committed
format: format files in examples and infini_train
format: use clang-format-16 instead
1 parent 3b13af4 commit 4b248b6

File tree

6 files changed

+53
-60
lines changed

6 files changed

+53
-60
lines changed

‎example/gpt2/main.cc‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
1+
#include <algorithm>
12
#include <chrono>
23
#include <cstdlib>
34
#include <filesystem>
45
#include <format>
56
#include <limits>
67
#include <memory>
78
#include <optional>
8-
#include <algorithm>
99
#include <unordered_map>
1010
#include <unordered_set>
1111

1212
#include "gflags/gflags.h"
1313
#include "glog/logging.h"
1414

1515
#include "infini_train/include/autocast.h"
16-
#include "infini_train/include/core/runtime/device_guard.h"
1716
#include "infini_train/include/checkpoint.h"
17+
#include "infini_train/include/core/runtime/device_guard.h"
1818
#include "infini_train/include/dataloader.h"
1919
#include "infini_train/include/device.h"
2020
#include "infini_train/include/nn/lora/lora_utils.h"
@@ -376,7 +376,8 @@ void Train(const nn::parallel::Rank &rank) {
376376

377377
LOG(INFO) << "start training";
378378

379-
auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
379+
auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step,
380+
bool prune_step_checkpoints) {
380381
const auto ckpt_start = std::chrono::high_resolution_clock::now();
381382

382383
TrainerState state;
@@ -413,8 +414,7 @@ void Train(const nn::parallel::Rank &rank) {
413414
const auto root = std::filesystem::path(FLAGS_checkpoint_dir);
414415
if (std::filesystem::exists(root)) {
415416
for (const auto &entry : std::filesystem::directory_iterator(root)) {
416-
if (entry.is_directory()
417-
&& entry.path().filename().string().starts_with("checkpoint_step_")) {
417+
if (entry.is_directory() && entry.path().filename().string().starts_with("checkpoint_step_")) {
418418
ckpts.push_back(entry.path());
419419
}
420420
}

‎example/gpt2/net.cc‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -720,10 +720,8 @@ std::shared_ptr<GPT2> GPT2::FromLLMC(const std::string &filepath) {
720720
}
721721

722722
void GPT2::SaveAsLLMC(const std::string &filepath) const {
723-
CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1)
724-
<< "SaveAsLLMC currently supports TP=1 only.";
725-
CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1)
726-
<< "SaveAsLLMC currently supports PP=1 only.";
723+
CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1) << "SaveAsLLMC currently supports TP=1 only.";
724+
CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1) << "SaveAsLLMC currently supports PP=1 only.";
727725

728726
std::ofstream ofs(filepath, std::ios::binary);
729727
CHECK(ofs.is_open()) << "Failed to open model file for write: " << filepath;
@@ -737,7 +735,8 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
737735
header[5] = static_cast<int32_t>(config_.n_head);
738736
header[6] = static_cast<int32_t>(config_.n_embd);
739737
header[7] = static_cast<int32_t>(config_.vocab_size);
740-
ofs.write(reinterpret_cast<const char *>(header.data()), static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
738+
ofs.write(reinterpret_cast<const char *>(header.data()),
739+
static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
741740

742741
const auto state_dict = StateDict();
743742
auto get_tensor = [&](const std::string &name) -> std::shared_ptr<Tensor> {
@@ -763,12 +762,12 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
763762
nn::Embedding::kParamWeightName)));
764763

765764
for (int idx = 0; idx < config_.n_layer; ++idx) {
766-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
767-
idx, Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)));
765+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
766+
Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)));
768767
}
769768
for (int idx = 0; idx < config_.n_layer; ++idx) {
770-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
771-
idx, Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)));
769+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
770+
Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)));
772771
}
773772
for (int idx = 0; idx < config_.n_layer; ++idx) {
774773
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
@@ -791,12 +790,12 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
791790
nn::parallel::RowParallelLinear::kParamBiasName)));
792791
}
793792
for (int idx = 0; idx < config_.n_layer; ++idx) {
794-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
795-
idx, Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)));
793+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
794+
Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)));
796795
}
797796
for (int idx = 0; idx < config_.n_layer; ++idx) {
798-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
799-
idx, Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)));
797+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName, idx,
798+
Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)));
800799
}
801800
for (int idx = 0; idx < config_.n_layer; ++idx) {
802801
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, GPT2Chunk::kHLayerName,
@@ -819,10 +818,10 @@ void GPT2::SaveAsLLMC(const std::string &filepath) const {
819818
nn::parallel::RowParallelLinear::kParamBiasName)));
820819
}
821820

822-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName,
823-
nn::LayerNorm::kParamWeightName)));
824-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName,
825-
nn::LayerNorm::kParamBiasName)));
821+
write_tensor_fp32(get_tensor(
822+
std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName, nn::LayerNorm::kParamWeightName)));
823+
write_tensor_fp32(get_tensor(
824+
std::format("{}.{}.{}", kTransformerLayerName, GPT2LastStage::kLnFLayerName, nn::LayerNorm::kParamBiasName)));
826825

827826
ofs.flush();
828827
CHECK(ofs.good()) << "Failed to flush model file: " << filepath;

‎example/llama3/main.cc‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,18 @@
1+
#include <algorithm>
12
#include <cstdlib>
23
#include <filesystem>
34
#include <format>
45
#include <limits>
56
#include <memory>
67
#include <optional>
7-
#include <algorithm>
88
#include <unordered_set>
99

1010
#include "gflags/gflags.h"
1111
#include "glog/logging.h"
1212

1313
#include "infini_train/include/autocast.h"
14-
#include "infini_train/include/core/runtime/device_guard.h"
1514
#include "infini_train/include/checkpoint.h"
15+
#include "infini_train/include/core/runtime/device_guard.h"
1616
#include "infini_train/include/dataloader.h"
1717
#include "infini_train/include/device.h"
1818
#include "infini_train/include/nn/lora/lora_utils.h"
@@ -339,7 +339,8 @@ void Train(const nn::parallel::Rank &rank) {
339339
}
340340
}
341341

342-
auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
342+
auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step,
343+
bool prune_step_checkpoints) {
343344
const auto ckpt_start = std::chrono::high_resolution_clock::now();
344345

345346
TrainerState state;
@@ -374,8 +375,7 @@ void Train(const nn::parallel::Rank &rank) {
374375
const auto root = std::filesystem::path(FLAGS_checkpoint_dir);
375376
if (std::filesystem::exists(root)) {
376377
for (const auto &entry : std::filesystem::directory_iterator(root)) {
377-
if (entry.is_directory()
378-
&& entry.path().filename().string().starts_with("checkpoint_step_")) {
378+
if (entry.is_directory() && entry.path().filename().string().starts_with("checkpoint_step_")) {
379379
ckpts.push_back(entry.path());
380380
}
381381
}

‎example/llama3/net.cc‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
#include "example/llama3/net.h"
22

33
#include <cmath>
4-
#include <cstring>
54
#include <cstdlib>
5+
#include <cstring>
66
#include <filesystem>
77
#include <fstream>
88
#include <map>
@@ -758,10 +758,8 @@ std::shared_ptr<LLaMA3> LLaMA3::FromLLMC(const std::string &filepath) {
758758
}
759759

760760
void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
761-
CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1)
762-
<< "SaveAsLLMC currently supports TP=1 only.";
763-
CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1)
764-
<< "SaveAsLLMC currently supports PP=1 only.";
761+
CHECK_EQ(nn::parallel::global::GetTensorParallelSize(), 1) << "SaveAsLLMC currently supports TP=1 only.";
762+
CHECK_EQ(nn::parallel::global::GetPipelineParallelSize(), 1) << "SaveAsLLMC currently supports PP=1 only.";
765763

766764
std::ofstream ofs(filepath, std::ios::binary);
767765
CHECK(ofs.is_open()) << "Failed to open model file for write: " << filepath;
@@ -789,7 +787,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
789787
header[13] = static_cast<int32_t>(config_.max_gen_batch_size);
790788
header[14] = 1; // version_major
791789
header[15] = 0; // version_minor
792-
ofs.write(reinterpret_cast<const char *>(header.data()), static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
790+
ofs.write(reinterpret_cast<const char *>(header.data()),
791+
static_cast<std::streamsize>(header.size() * sizeof(int32_t)));
793792

794793
const auto state_dict = StateDict();
795794
auto get_tensor = [&](const std::string &name) -> std::shared_ptr<Tensor> {
@@ -810,8 +809,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
810809
nn::parallel::VocabParallelEmbedding::kParamWeightName)));
811810

812811
for (int idx = 0; idx < config_.n_layer; ++idx) {
813-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
814-
idx, Block::kLn1LayerName, RMSNorm::kParamWeightName)));
812+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName, idx,
813+
Block::kLn1LayerName, RMSNorm::kParamWeightName)));
815814
}
816815
for (int idx = 0; idx < config_.n_layer; ++idx) {
817816
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
@@ -824,8 +823,8 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
824823
nn::parallel::RowParallelLinear::kParamWeightName)));
825824
}
826825
for (int idx = 0; idx < config_.n_layer; ++idx) {
827-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
828-
idx, Block::kLn2LayerName, RMSNorm::kParamWeightName)));
826+
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName, idx,
827+
Block::kLn2LayerName, RMSNorm::kParamWeightName)));
829828
}
830829
for (int idx = 0; idx < config_.n_layer; ++idx) {
831830
write_tensor_fp32(get_tensor(std::format("{}.{}.{}.{}.{}.{}", kTransformerLayerName, LLaMA3Chunk::kHLayerName,
@@ -843,10 +842,10 @@ void LLaMA3::SaveAsLLMC(const std::string &filepath) const {
843842
nn::parallel::RowParallelLinear::kParamWeightName)));
844843
}
845844

846-
write_tensor_fp32(get_tensor(std::format("{}.{}.{}", kTransformerLayerName, LLaMA3LastStage::kLnFLayerName,
847-
RMSNorm::kParamWeightName)));
848-
write_tensor_fp32(get_tensor(std::format("{}.{}", LLaMA3LastStage::kLMHeadLayerName,
849-
nn::parallel::ColumnParallelLinear::kParamWeightName)));
845+
write_tensor_fp32(get_tensor(
846+
std::format("{}.{}.{}", kTransformerLayerName, LLaMA3LastStage::kLnFLayerName, RMSNorm::kParamWeightName)));
847+
write_tensor_fp32(get_tensor(
848+
std::format("{}.{}", LLaMA3LastStage::kLMHeadLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)));
850849

851850
ofs.flush();
852851
CHECK(ofs.good()) << "Failed to flush model file: " << filepath;

‎infini_train/include/checkpoint.h‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,11 @@ class Checkpoint {
4949
TrainerState *state, const CheckpointLoadOptions &options = {});
5050

5151
private:
52-
static void SaveStateDictBinary(
53-
const std::filesystem::path &path,
54-
const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
52+
static void SaveStateDictBinary(const std::filesystem::path &path,
53+
const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
5554

56-
static std::unordered_map<std::string, std::shared_ptr<Tensor>> LoadStateDictBinary(const std::filesystem::path &path);
55+
static std::unordered_map<std::string, std::shared_ptr<Tensor>>
56+
LoadStateDictBinary(const std::filesystem::path &path);
5757

5858
static void SaveTrainerState(const std::filesystem::path &path, const TrainerState &state);
5959
static TrainerState LoadTrainerState(const std::filesystem::path &path);

‎infini_train/src/checkpoint.cc‎

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,7 @@ template <typename T> T ExtractNumberField(const std::string &content, const std
8686

8787
void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, const nn::Module &model, const Optimizer &optimizer,
8888
const TrainerState &state, const CheckpointOptions &options) {
89-
CHECK(options.format == "bin" || options.format == "pth")
90-
<< "Unsupported checkpoint format: " << options.format;
89+
CHECK(options.format == "bin" || options.format == "pth") << "Unsupported checkpoint format: " << options.format;
9190
std::filesystem::create_directories(checkpoint_dir);
9291
LOG(ERROR) << "[CKPT] Save begin: dir=" << checkpoint_dir << ", format=" << options.format
9392
<< ", global_step=" << state.global_step;
@@ -149,12 +148,11 @@ void Checkpoint::Load(const std::filesystem::path &checkpoint_dir, nn::Module *m
149148
}
150149

151150
*state = LoadTrainerState(checkpoint_dir / "trainer_state.json");
152-
LOG(ERROR) << "[CKPT] Load done: global_step=" << state->global_step << ", data_batch_idx="
153-
<< state->data_batch_idx << ", data_batch_stride=" << state->data_batch_stride
154-
<< ", best_loss=" << state->best_loss << ", last_lr=" << state->last_lr
155-
<< ", optimizer_type=" << state->optimizer_type
156-
<< ", topology(ddp,tp,sp,pp)=(" << state->ddp_size << "," << state->tp_size << ","
157-
<< state->sp_size << "," << state->pp_size << ")";
151+
LOG(ERROR) << "[CKPT] Load done: global_step=" << state->global_step << ", data_batch_idx=" << state->data_batch_idx
152+
<< ", data_batch_stride=" << state->data_batch_stride << ", best_loss=" << state->best_loss
153+
<< ", last_lr=" << state->last_lr << ", optimizer_type=" << state->optimizer_type
154+
<< ", topology(ddp,tp,sp,pp)=(" << state->ddp_size << "," << state->tp_size << "," << state->sp_size
155+
<< "," << state->pp_size << ")";
158156
}
159157

160158
void Checkpoint::SaveStateDictBinary(const std::filesystem::path &path,
@@ -178,9 +176,7 @@ void Checkpoint::SaveStateDictBinary(const std::filesystem::path &path,
178176
const auto &dims = tensor->Dims();
179177
uint32_t ndim = static_cast<uint32_t>(dims.size());
180178
ofs.write(reinterpret_cast<const char *>(&ndim), sizeof(ndim));
181-
for (const auto dim : dims) {
182-
ofs.write(reinterpret_cast<const char *>(&dim), sizeof(dim));
183-
}
179+
for (const auto dim : dims) { ofs.write(reinterpret_cast<const char *>(&dim), sizeof(dim)); }
184180

185181
Tensor cpu_tensor = tensor->To(Device());
186182
uint64_t bytes = static_cast<uint64_t>(cpu_tensor.SizeInBytes());
@@ -189,7 +185,8 @@ void Checkpoint::SaveStateDictBinary(const std::filesystem::path &path,
189185
}
190186
}
191187

192-
std::unordered_map<std::string, std::shared_ptr<Tensor>> Checkpoint::LoadStateDictBinary(const std::filesystem::path &path) {
188+
std::unordered_map<std::string, std::shared_ptr<Tensor>>
189+
Checkpoint::LoadStateDictBinary(const std::filesystem::path &path) {
193190
std::ifstream ifs(path, std::ios::binary);
194191
CHECK(ifs.is_open()) << "Failed to open checkpoint file: " << path;
195192

@@ -214,9 +211,7 @@ std::unordered_map<std::string, std::shared_ptr<Tensor>> Checkpoint::LoadStateDi
214211
uint32_t ndim = 0;
215212
ifs.read(reinterpret_cast<char *>(&ndim), sizeof(ndim));
216213
std::vector<int64_t> dims(ndim);
217-
for (uint32_t d = 0; d < ndim; ++d) {
218-
ifs.read(reinterpret_cast<char *>(&dims[d]), sizeof(dims[d]));
219-
}
214+
for (uint32_t d = 0; d < ndim; ++d) { ifs.read(reinterpret_cast<char *>(&dims[d]), sizeof(dims[d])); }
220215

221216
uint64_t bytes = 0;
222217
ifs.read(reinterpret_cast<char *>(&bytes), sizeof(bytes));

0 commit comments

Comments
 (0)