完成 TinyInfiniTrain 作业：autograd、算子（CPU/CUDA）、Tensor、Dispatcher 与 GPT-2 端到端训练

.gitignore

@@ -2,3 +2,5 @@ build/
 .cache/
 .vscode/
 Data/
+review
+.claude

example/common/tiny_shakespeare_dataset.cc

@@ -61,10 +61,45 @@ TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequ
     | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B) | token数据           |
     ----------------------------------------------------------------------------------
        =================================== 作业 =================================== */
+    std::ifstream ifs(path, std::ios::binary);
+    CHECK(ifs.is_open()) << "Failed to open file: " << path;
+
+    auto header = ReadSeveralBytesFromIfstream(1024, &ifs);
+    CHECK_EQ(ifs.gcount(), 1024) << "Truncated header in file: " << path;
+    const int magic = BytesToType<int>(header, 0);
+    const int num_toks = BytesToType<int>(header, 8);
+    CHECK(kTypeMap.contains(magic)) << "Unsupported magic number: " << magic;
+    CHECK_GE(num_toks, 0) << "Invalid num_toks in file: " << path;
+
+    TinyShakespeareFile file;
+    file.type = kTypeMap.at(magic);
+    const size_t token_size = kTypeToSize.at(file.type);
+
+    // 读 token 流并转为 int64 张量：CrossEntropy 的 target 要求 int64，
+    // 且 operator[] 中 y 相对 x 偏移 sizeof(int64_t) = 8 字节，恰为一个 int64 token（预测下一 token）
+    const size_t num_tok_bytes = static_cast<size_t>(num_toks) * token_size;
+    auto token_bytes = ReadSeveralBytesFromIfstream(num_tok_bytes, &ifs);
+    // 读取完整性校验（损坏/截断文件下避免短向量导致的越界读）
+    CHECK_EQ(static_cast<size_t>(ifs.gcount()), num_tok_bytes) << "Truncated token data in file: " << path;
+    const size_t num_samples = static_cast<size_t>(num_toks) / sequence_length;
+    file.dims = {static_cast<int64_t>(num_samples), static_cast<int64_t>(sequence_length)};
+    file.tensor = infini_train::Tensor(file.dims, DataType::kINT64);
+    int64_t *tensor_data = static_cast<int64_t *>(file.tensor.DataPtr());
+    for (size_t i = 0; i < num_samples * sequence_length; ++i) {
+        tensor_data[i] = (token_size == 2) ? BytesToType<uint16_t>(token_bytes, i * 2)
+                                           : BytesToType<int32_t>(token_bytes, i * 4);
+    }
+    return 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)),
+      // 最后一个序列缺"下一 token"作标签，operator[] 以 CHECK_LT(idx, dims[0]-1) 限定可访问边界，
+      // 故 Size 报告可访问样本数 = dims[0]-1（空数据集防护为 0），与 operator[] 的契约一致
+      num_samples_(text_file_.dims[0] > 0 ? text_file_.dims[0] - 1 : 0) {
     // =================================== 作业 ===================================
     // TODO：初始化数据集实例
     // HINT: 调用ReadTinyShakespeareFile加载数据文件
@@ -73,9 +108,10 @@ TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size
 
 std::pair<std::shared_ptr<infini_train::Tensor>, std::shared_ptr<infini_train::Tensor>>
 TinyShakespeareDataset::operator[](size_t idx) const {
-    CHECK_LT(idx, text_file_.dims[0] - 1);
+    // 以 num_samples_（可访问样本数，构造时已防护空数据集）限定边界，与 Size 契约一致
+    CHECK_LT(idx, num_samples_);
     std::vector<int64_t> dims = std::vector<int64_t>(text_file_.dims.begin() + 1, text_file_.dims.end());
-    // x: (seq_len), y: (seq_len) -> stack -> (bs, seq_len) (bs, seq_len)
+    // x/y 与 text_file_.tensor 共享 buffer：x 取第 idx 个序列，y 偏移一个 token（预测下一 token）
     return {std::make_shared<infini_train::Tensor>(text_file_.tensor, idx * sequence_size_in_bytes_, dims),
             std::make_shared<infini_train::Tensor>(text_file_.tensor, idx * sequence_size_in_bytes_ + sizeof(int64_t),
                                                    dims)};

example/common/tokenizer.cc

@@ -78,21 +78,41 @@ Tokenizer::Tokenizer(const std::string &filepath) {
     | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B) | token词表数据       |
     ----------------------------------------------------------------------------------
     ===================================== 作业 ===================================== */
+    std::ifstream ifs(filepath, std::ios::binary);
+    CHECK(ifs.is_open()) << "Failed to open tokenizer file: " << filepath;
+
+    auto header = ReadSeveralBytesFromIfstream(1024, &ifs);
+    CHECK_EQ(ifs.gcount(), 1024) << "Truncated header in file: " << filepath;
+    magic_number_ = BytesToType<uint32_t>(header, 0);
+    vocab_size_ = BytesToType<uint32_t>(header, 8);
+    eot_token_ = BytesToType<uint32_t>(header, 12);
+    CHECK(kEotMap.contains(magic_number_)) << "Unsupported tokenizer magic number: " << magic_number_;
+
+    // 词表格式（与 llm.c gpt2_tokenizer.bin 一致）：每个 token 为 1 字节长度前缀 + 原始字节
+    token_table_.reserve(vocab_size_);
+    for (uint32_t i = 0; i < vocab_size_; ++i) {
+        const uint8_t len = BytesToType<uint8_t>(ReadSeveralBytesFromIfstream(1, &ifs), 0);
+        auto bytes = ReadSeveralBytesFromIfstream(len, &ifs);
+        token_table_.emplace_back(reinterpret_cast<const char *>(bytes.data()), len);
+    }
+    // 词表读取完整性校验：任一项读取失败（文件截断）都会置位 failbit，快速失败而非静默加载垃圾词表
+    CHECK(ifs) << "Truncated tokenizer vocab table in file: " << filepath;
 }
 
 std::string Tokenizer::Decode(uint32_t token_id) const {
     /* ===================================== 作业 =====================================
     TODO：实现token_id到文本的转换
     功能描述：根据token_id返回对应的文本片段
     ===================================== 作业 ===================================== */
-    return "";
+    CHECK_LT(token_id, vocab_size_) << "token_id out of range: " << token_id;
+    return token_table_[token_id];
 }
 
 void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length,
                              uint32_t text_length, Device device) const {
     std::vector<int64_t> dims;
     dims.assign({batch_size, sequence_length});
-    // x_tensor (FLAGS_batch_size, FLAGS_sequence_length) eq:(4, 64)
+    // x_tensor 形状为 (batch_size, sequence_length)
     infini_train::Tensor x_tensor = infini_train::Tensor(dims, DataType::kINT64);
     int64_t *x_buff = static_cast<int64_t *>(x_tensor.DataPtr());
     for (int i = 0; i < batch_size * sequence_length; ++i) { x_buff[i] = eot_token_; }
@@ -111,6 +131,53 @@ void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_siz
         TODO：实现单步文本生成逻辑
         HINT：调用model.Forward推理获取logits，根据推理结果进行随机采样，调用Decode获取文本结果
         ===================================== 作业 ===================================== */
+        // 生成场景无 Backward：临时禁用参数梯度使前向不建 autograd 图（Function 即时释放），
+        // 避免算子 saved_tensors_ 的循环引用在无 Backward 场景下导致显存逐步泄漏。
+        // 副作用说明：本方法（const 签名）临时修改模型参数的 requires_grad 状态——生成区间禁用、
+        // 末步恢复为 true；若生成中途以异常中断，requires_grad 将残留为 false，需重新调用
+        // set_requires_grad(true) 恢复训练（本框架错误处理为 CHECK 快速失败，正常路径不触发）
+        if (t == prompt_len) {
+            for (auto &param : model.Parameters()) { param->set_requires_grad(false); }
+        }
+        // 同步 host 输入到目标设备并前向推理
+        x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
+        auto outputs = model.Forward({x});
+        auto logits = outputs[0];
+        auto logits_cpu = logits->To(Device(DeviceType::kCPU, 0));
+        const float *logits_data = static_cast<const float *>(logits_cpu.DataPtr());
+        const int64_t vocab_size = logits->Dims()[2];
+        // 恢复参数梯度（后续训练仍需要建图反向）
+        if (t == text_length - 1) {
+            for (auto &param : model.Parameters()) { param->set_requires_grad(true); }
+        }
+
+        // 生成语义（对齐 llm.c）：取位置 t-1 的分布预测位置 t 的 token，每 batch 独立采样写回
+        std::vector<float> probs(static_cast<size_t>(vocab_size));
+        for (uint32_t b = 0; b < batch_size; ++b) {
+            const float *logits_at_t = logits_data + (b * sequence_length + t - 1) * vocab_size;
+            // softmax（数值稳定：先减最大值）
+            float max_logit = logits_at_t[0];
+            for (int64_t i = 1; i < vocab_size; ++i) {
+                if (logits_at_t[i] > max_logit) {
+                    max_logit = logits_at_t[i];
+                }
+            }
+            float sum_exp = 0.0f;
+            for (int64_t i = 0; i < vocab_size; ++i) {
+                probs[i] = std::exp(logits_at_t[i] - max_logit);
+                sum_exp += probs[i];
+            }
+            for (int64_t i = 0; i < vocab_size; ++i) {
+                probs[i] /= sum_exp;
+            }
+
+            // 函数开头声明了与命名空间常量同名的局部变量 `kRngState = kRngState;`（自初始化，未定义行为），
+            // 此处显式引用命名空间常量 kRngState（=1337，与 llm.c 固定种子一致）保证生成可复现
+            static uint64_t rng_state = infini_train::kRngState;
+            const int next_token = SampleMult(probs.data(), static_cast<int>(vocab_size), RandomF32(rng_state));
+            x_buff[b * sequence_length + t] = next_token;
+            std::cout << Decode(next_token);
+        }
     }
     std::cout << std::endl;
 }

infini_train/include/dispatcher.h

@@ -22,6 +22,7 @@ class KernelFunction {
 
         using FuncT = RetT (*)(ArgsT...);
         // TODO: 实现函数调用逻辑
+        return reinterpret_cast<FuncT>(func_ptr_)(args...);
     }
 
 private:
@@ -48,15 +49,24 @@ class Dispatcher {
         // TODO：实现kernel注册机制
         // 功能描述：将kernel函数与设备类型、名称绑定
         // =================================== 作业 ===================================
+        CHECK(!key_to_kernel_map_.contains(key))
+            << "Kernel already registered: " << key.second << " on device: " << static_cast<int>(key.first);
+        key_to_kernel_map_.emplace(key, KernelFunction(std::forward<FuncT>(kernel)));
     }
 
 private:
     std::map<KeyT, KernelFunction> key_to_kernel_map_;
 };
 } // namespace infini_train
 
+#define INFINI_KERNEL_CONCAT_(a, b) a##b
+#define INFINI_KERNEL_CONCAT(a, b) INFINI_KERNEL_CONCAT_(a, b)
 #define REGISTER_KERNEL(device, kernel_name, kernel_func)                                                              \
-    // =================================== 作业 ===================================
-    // TODO：实现自动注册宏
-    // 功能描述：在全局静态区注册kernel，避免显式初始化代码
-    // =================================== 作业 ===================================
+    /* =================================== 作业 =================================== */                               \
+    /* TODO：实现自动注册宏 */                                                                                 \
+    /* 功能描述：在全局静态区注册kernel，避免显式初始化代码 */                                  \
+    /* =================================== 作业 =================================== */                               \
+    static const bool INFINI_KERNEL_CONCAT(kernel_reg_, __COUNTER__) =                                                 \
+        (infini_train::Dispatcher::Instance().Register(                                                                \
+             std::make_pair(device, #kernel_name), kernel_func),                                                       \
+         true);

infini_train/src/autograd/elementwise.cc

@@ -10,17 +10,25 @@ std::vector<std::shared_ptr<Tensor>> Neg::Forward(const std::vector<std::shared_
     // TODO：通过Dispatcher获取设备专属kernel，对输入张量进行取反操作
     // NOTES: 依赖test_dispatcher，Neg kernel实现已给出
     // =================================== 作业 ===================================
+    CHECK_EQ(input_tensors.size(), 1);
+    const auto &input = input_tensors[0];
 
-    return std::vector<std::shared_ptr<Tensor>>();
+    auto device = input->GetDevice().Type();
+    auto kernel = Dispatcher::Instance().GetKernel({device, "NegForward"});
+    return {kernel.Call<std::shared_ptr<Tensor>>(input)};
 }
 
 std::vector<std::shared_ptr<Tensor>> Neg::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
     // =================================== 作业 ===================================
     // TODO：通过Dispatcher获取设备专属的反向传播kernel，计算梯度
     // NOTES: 依赖test_dispatcher，Neg的kernel实现已给出
     // =================================== 作业 ===================================
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &grad_output = grad_outputs[0];
 
-    return std::vector<std::shared_ptr<Tensor>>();
+    auto device = grad_output->GetDevice().Type();
+    auto kernel = Dispatcher::Instance().GetKernel({device, "NegBackward"});
+    return {kernel.Call<std::shared_ptr<Tensor>>(grad_output)};
 }
 
 std::vector<std::shared_ptr<Tensor>> Reciprocal::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {

infini_train/src/kernels/cpu/accumulate_grad.cc

@@ -1,3 +1,4 @@
+#include <cmath>
 #include <cstddef>
 #include <memory>
 
@@ -18,6 +19,27 @@ void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_p
     // TODO：实现Adam优化器的梯度累积和参数更新
     // REF:
     // =================================== 作业 ===================================
+    // 标准 Adam（含偏差校正）：m = β1*m + (1-β1)*g，v = β2*v + (1-β2)*g²，
+    // param -= lr * m̂ / (√v̂ + eps)，其中 m̂ = m/(1-β1^t)，v̂ = v/(1-β2^t)，t 从 1 开始
+    CHECK_EQ(grad->NumElements(), param->NumElements());  // 校验各张量元素数一致，防止直接调用时越界读写
+    CHECK_EQ(m->NumElements(), param->NumElements());
+    CHECK_EQ(v->NumElements(), param->NumElements());
+    CHECK_GE(t, 1);  // Adam 偏差校正要求 t 从 1 开始（t=0 时 1-β^t=0，除零产生 NaN）
+    // 偏差校正因子仅依赖步数，与元素无关，提升到循环外计算一次（double 对齐 PyTorch 标量路径的双精度惯例）
+    const double beta1_t = std::pow(static_cast<double>(beta1), static_cast<double>(t));
+    const double beta2_t = std::pow(static_cast<double>(beta2), static_cast<double>(t));
+    for (int64_t idx = 0; idx < param->NumElements(); ++idx) {
+        const float g = static_cast<const float *>(grad->DataPtr())[idx];
+        float &param_elem = static_cast<float *>(param->DataPtr())[idx];
+        float &m_elem = static_cast<float *>(m->DataPtr())[idx];
+        float &v_elem = static_cast<float *>(v->DataPtr())[idx];
+
+        m_elem = beta1 * m_elem + (1.0f - beta1) * g;
+        v_elem = beta2 * v_elem + (1.0f - beta2) * g * g;
+        const float m_hat = static_cast<float>(m_elem / (1.0 - beta1_t));
+        const float v_hat = static_cast<float>(v_elem / (1.0 - beta2_t));
+        param_elem -= learning_rate * m_hat / (std::sqrt(v_hat) + eps);
+    }
 }
 
 } // namespace infini_train::kernels::cpu