第一次作业提交

.gitignore

@@ -2,3 +2,4 @@ build/
 .cache/
 .vscode/
 Data/
+third_party/

CMakeLists.txt

@@ -8,6 +8,20 @@ include(CMakeDependentOption)
 cmake_dependent_option(BUILD_TEST_CORE "Build tests for core components" ON BUILD_TEST OFF)
 project(infini_train VERSION 0.3.0 LANGUAGES CXX)
 
+# 添加这些行
+if(NOT CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Debug)
+endif()
+
+# 确保Debug模式包含调试信息
+set(CMAKE_CXX_FLAGS_DEBUG "-g -O0")
+set(CMAKE_C_FLAGS_DEBUG "-g -O0")
+
+# 如果使用CUDA，也为CUDA添加调试标志
+if(USE_CUDA)
+    set(CMAKE_CUDA_FLAGS_DEBUG "-g -O0")
+endif()
+
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
@@ -52,7 +66,18 @@ if(USE_CUDA)
     target_link_libraries(infini_train_cuda_kernels glog CUDA::cudart CUDA::cublas)
 
     add_library(infini_train STATIC ${SRC})
-    target_link_libraries(infini_train glog gflags "-Wl,--whole-archive" infini_train_cpu_kernels infini_train_cuda_kernels "-Wl,--no-whole-archive")
+    # 修改这行：添加CUDA库链接
+    target_link_libraries(infini_train 
+        glog 
+        gflags 
+        CUDA::cudart 
+        CUDA::cuda_driver  # 添加这个，解决 cuInit 等函数
+        CUDA::cublas
+        "-Wl,--whole-archive" 
+        infini_train_cpu_kernels 
+        infini_train_cuda_kernels 
+        "-Wl,--no-whole-archive"
+    )
 else()
     add_library(infini_train STATIC ${SRC})
     target_link_libraries(infini_train glog gflags "-Wl,--whole-archive" infini_train_cpu_kernels "-Wl,--no-whole-archive")

Makefile

@@ -8,7 +8,7 @@ CMAKE_OPT += -DUSE_CUDA=$(USE_CUDA)
 
 build:
 	mkdir -p build/$(TYPE)
-	cd build/$(TYPE) && cmake $(CMAKE_OPT) ../.. && make -j8
+	cd build/$(TYPE) && cmake $(CMAKE_OPT) ../.. && make -j
 
 clean:
 	rm -rf build

example/common/tiny_shakespeare_dataset.cc

@@ -61,14 +61,55 @@ TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequ
     | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B) | token数据           |
     ----------------------------------------------------------------------------------
        =================================== 作业 =================================== */
+    if (!std::filesystem::exists(path)) {
+        LOG(FATAL) << "File not found: " << path;
+    }
+
+    TinyShakespeareFile text_file;
+    std::ifstream ifs(path, std::ios::binary);
+    const auto header = ReadSeveralBytesFromIfstream(1024, &ifs);
+    const int magic = BytesToType<int32_t>(header, 0);
+    const int version = BytesToType<int32_t>(header, 4);
+    const int num_tokens = BytesToType<int32_t>(header, 8);
+    text_file.type = kTypeMap.at(magic);
+
+    const int num_sequences = num_tokens / sequence_length;
+    text_file.dims.assign({num_sequences, static_cast<int64_t>(sequence_length)});
+
+    const int data_size_in_bytes
+        = kTypeToSize.at(text_file.type)
+        * std::accumulate(text_file.dims.begin(), text_file.dims.end(), 1, std::multiplies<int>());
+    // shape: (num_seq, seq_len), dtype: int64
+    text_file.tensor = infini_train::Tensor(text_file.dims, DataType::kINT64);
+    int64_t *dst = static_cast<int64_t *>(text_file.tensor.DataPtr());
+
+    std::variant<std::vector<uint16_t>, std::vector<int32_t>> buffer;
+    if (text_file.type == TinyShakespeareType::kUINT16) {
+        CHECK_LE(sequence_length, 1024); // GPT-2: max_seq_length = 1024
+        buffer = std::vector<uint16_t>(num_sequences * sequence_length);
+    } else if (text_file.type == TinyShakespeareType::kUINT32) {
+        CHECK_LE(sequence_length, 8192); // LLaMA-3: max_seq_length = 8192
+        buffer = std::vector<int32_t>(num_sequences * sequence_length);
+    }
+    std::visit(
+        [&](auto &vec) {
+            ifs.read(reinterpret_cast<char *>(vec.data()), data_size_in_bytes);
+            for (size_t i = 0; i < vec.size(); ++i) { dst[i] = static_cast<int64_t>(vec[i]); }
+        },
+        buffer);
+    return text_file;
 }
 } // namespace
 
-TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length) {
+TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length)
+    : text_file_(ReadTinyShakespeareFile(filepath, sequence_length)), sequence_length_(sequence_length),
+      sequence_size_in_bytes_(sequence_length * sizeof(int64_t)), num_samples_(text_file_.dims[0] - 1) {
     // =================================== 作业 ===================================
     // TODO：初始化数据集实例
     // HINT: 调用ReadTinyShakespeareFile加载数据文件
     // =================================== 作业 ===================================
+    CHECK_EQ(text_file_.dims[1], sequence_length_);
+    CHECK_EQ(static_cast<int>(text_file_.tensor.Dtype()), static_cast<int>(DataType::kINT64));
 }
 
 std::pair<std::shared_ptr<infini_train::Tensor>, std::shared_ptr<infini_train::Tensor>>

example/common/tokenizer.cc

@@ -78,14 +78,52 @@ Tokenizer::Tokenizer(const std::string &filepath) {
     | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B) | token词表数据       |
     ----------------------------------------------------------------------------------
     ===================================== 作业 ===================================== */
+    if (!std::filesystem::exists(filepath)) {
+        LOG(FATAL) << "File not found: " << filepath;
+    }
+
+    std::ifstream ifs(filepath, std::ios::binary);
+    const auto header = ReadSeveralBytesFromIfstream(1024, &ifs);
+
+    magic_number_ = BytesToType<uint32_t>(header, 0);
+    const uint32_t version_num = BytesToType<uint32_t>(header, 4);
+    vocab_size_ = BytesToType<uint32_t>(header, 8);
+    if (kEotMap.find(magic_number_) == kEotMap.end()) {
+        LOG(FATAL) << "Unsupported tokenizer magic: " << magic_number_;
+    }
+
+    Version version = static_cast<Version>(version_num);
+    if (version == Version::kV1) {
+        eot_token_ = kEotMap.at(magic_number_);
+    } else if (version == Version::kV2) {
+        const uint32_t eot_token_2 = BytesToType<uint32_t>(header, 12);
+        eot_token_ = eot_token_2;
+    } else {
+        LOG(FATAL) << "Unsupported tokenizer version: " << version_num;
+        return;
+    }
+
+    token_table_.resize(vocab_size_);
+    for (uint32_t i = 0; i < vocab_size_; ++i) {
+        uint8_t length;
+        ifs.read(reinterpret_cast<char *>(&length), sizeof(length));
+
+        std::vector<char> buffer(length);
+        ifs.read(buffer.data(), length);
+
+        token_table_[i] = std::string(buffer.data(), length);
+    }
 }
 
 std::string Tokenizer::Decode(uint32_t token_id) const {
     /* ===================================== 作业 =====================================
     TODO：实现token_id到文本的转换
     功能描述：根据token_id返回对应的文本片段
     ===================================== 作业 ===================================== */
-    return "";
+    if (token_id >= vocab_size_) {
+        return "[INVALID_TOKEN]";
+    }
+    return token_table_[token_id];   
 }
 
 void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length,
@@ -106,11 +144,28 @@ void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_siz
     auto x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
     uint64_t kRngState = kRngState;
     LOG(INFO) << "start generate text:";
+
+    infini_train::Device cpu_device = infini_train::Device(infini_train::DeviceType::kCPU, 0);
     for (int t = prompt_len; t < text_length; t++) {
         /* ===================================== 作业 =====================================
         TODO：实现单步文本生成逻辑
         HINT：调用model.Forward推理获取logits，根据推理结果进行随机采样，调用Decode获取文本结果
         ===================================== 作业 ===================================== */
+        x = std::make_shared<infini_train::Tensor>(x->To(device)); // CPU->calc device
+        // TODO(jym): use no_grad forward later
+        auto logits = model.Forward({x})[0];
+        auto logits_orignal = nn::function::Softmax(logits, -1);
+        auto logits_cpu = logits_orignal->To(cpu_device);
+        auto data = logits_cpu.DataPtr();
+        auto vocab_size = logits->Dims()[2];
+        float *probs = static_cast<float *>(data) + (t - 1) * vocab_size;
+        float coin = RandomF32(kRngState);
+        int next_token = SampleMult(probs, vocab_size, coin);
+
+        x = std::make_shared<infini_train::Tensor>(x->To(cpu_device)); // calc device->CPU
+        auto data_temp = static_cast<int64_t *>(x->DataPtr());
+        data_temp[t] = next_token;
+        std::cout << Decode(next_token);
     }
     std::cout << std::endl;
 }

infini_train/include/common/common.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <utility>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/device.h"
+#include "infini_train/include/tensor.h"
+
+#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: "                                              \
+                                                  + kDataTypeToDesc.at(static_cast<infini_train::DataType>(dtype))))