Skip to content

Latest commit

 

History

History
924 lines (735 loc) · 51.6 KB

File metadata and controls

924 lines (735 loc) · 51.6 KB

TinyInfiniTrain 作业报告

一、test 通过截图

此处粘贴 make test-cpp 完整输出截图(远程服务器:RTX 5090 / CUDA 12.8 / Linux)。

最终结论:8 个测例中 7 个稳定通过(65/100 分);test_gpt2 经环境升级后处于容差边界(约 75-80% 通过率),诊断确认是测例容差问题而非实现缺陷(详见作业六)。

测例 分值 结果
test_elementwise 5 ✅ 通过
test_matmul 5 ✅ 通过
test_matmul_cuda 10 ✅ 通过
test_adam 5 ✅ 通过
test_adam_cuda 10 ✅ 通过
test_tensor 10 ✅ 通过
test_dispatcher 20 ✅ 通过
test_gpt2 35 ⚠️ 边界抖动(环境升级后 75-80% 通过,见作业六诊断)

test_gpt2 关键诊断证据(逐步收敛曲线)

step 0:  max_abs=51.34   step 4:  18.06   step 8:  5.34
step 2:  30.60           step 6:  11.23   step 10: 0.0065

误差从 51.34 单调收敛到 0.0065,与参考快照精确吻合到 float32 极限(相对误差 5e-5),说明训练轨迹正确。将环境升级至与往届通过者完全一致(CUDA 13.3.1 + native sm_120 + RTX 5090)后,误差从稳定超差 0.0023 降至边界 ~0.00107(仅超容差 7%),实测约 75-80% 通过率——进一步证明未通过仅因测例使用绝对容差 1e-3 衡量 |value|~200 的 logits。

二、作业步骤

将代码填入下面代码块中指定位置,并详细描述完成该作业的解决思路和遇到的问题。

作业一:autograd机制调用Neg kernel的实现

难度:⭐

对应测例:TEST(ElementwiseTest, NegForward)TEST(ElementwiseTest, NegBackward)

需要实现的代码块位置:infini_train/src/autograd/elementwise.cc

std::vector<std::shared_ptr<Tensor>> Neg::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
    CHECK_EQ(input_tensors.size(), 1);
    const auto &input = input_tensors[0];

    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) {
    // d(-x)/dx == -1, so the input is not needed here and Neg has no SetupContext.
    CHECK_EQ(grad_outputs.size(), 1);
    const auto &grad_output = grad_outputs[0];

    auto device = grad_output->GetDevice().Type();
    auto kernel = Dispatcher::Instance().GetKernel({device, "NegBackward"});
    return {kernel.Call<std::shared_ptr<Tensor>>(grad_output)};
}

解决思路

  1. Forward 只做一件事:从输入张量的 GetDevice().Type() 拿到设备类型,以 {device, "NegForward"} 为 key 从 Dispatcher::Instance().GetKernel() 取 kernel,再通过 kernel.Call<std::shared_ptr<Tensor>>(input) 调用。返回类型由调用点显式指定,参数类型由实参推导(机制见作业五)。
  2. Backward 的关键是意识到 d(-x)/dx = -1 与输入无关:梯度只依赖 grad_output,所以 Neg 不需要 SetupContext 保存输入,直接对 grad_output 调用 NegBackward kernel 即可。
  3. 对照了同文件里 Reciprocal / Sin / Cos / Tanh 等算子:凡是导数形式包含输入的算子(如 Reciprocal 的 -1/x²),都通过 SetupContext 保存输入(或输出),Backward 时从 saved_tensors_ 取出。这帮助我理解了"是否需要保存上下文由导数形式决定"这一 autograd 设计。

遇到问题

  • 最初容易把 Backward 想复杂(以为必须保存输入)。读 kernel 注册表和其它算子的实现后,发现 NegBackward 只接收 grad_output 一个参数,从而理解了 Neg 是"无需上下文"的最简算子,这个问题本身帮我建立了对 SetupContext 机制的正确理解。

作业二:实现矩阵乘法

难度:⭐⭐

CPU实现

对应测例:TEST(MatmulTest, BasicMatrixMultiply)TEST(MatmulTest, BatchedMatrixMultiply), TEST(MatmulTest, BackwardPass)

需要实现的代码块位置:infini_train/src/kernels/cpu/linear.cc

// Broadcasts two batch shapes following numpy/PyTorch rules: shapes are right-aligned and a
// dimension of 1 stretches to match the other operand.
std::vector<int64_t> BroadcastBatchDims(const std::vector<int64_t> &a_dims, const std::vector<int64_t> &b_dims) {
    const size_t ndim = std::max(a_dims.size(), b_dims.size());
    const size_t a_pad = ndim - a_dims.size();
    const size_t b_pad = ndim - b_dims.size();

    std::vector<int64_t> out_dims(ndim, 1);
    for (size_t i = 0; i < ndim; ++i) {
        const int64_t a_dim = i < a_pad ? 1 : a_dims[i - a_pad];
        const int64_t b_dim = i < b_pad ? 1 : b_dims[i - b_pad];
        CHECK(a_dim == b_dim || a_dim == 1 || b_dim == 1)
            << "Matmul batch dimensions are not broadcastable: " << a_dim << " vs " << b_dim;
        out_dims[i] = std::max(a_dim, b_dim);
    }
    return out_dims;
}

// Maps a flat index over the broadcast batch shape back to the flat batch index of one operand.
// Broadcast dimensions (size 1) contribute nothing, so several output batches read the same input.
int64_t BroadcastBatchOffset(int64_t flat_idx, const std::vector<int64_t> &out_dims,
                             const std::vector<int64_t> &in_dims) {
    const size_t pad = out_dims.size() - in_dims.size();
    int64_t offset = 0;
    int64_t in_stride = 1;
    for (int i = static_cast<int>(out_dims.size()) - 1; i >= 0; --i) {
        const int64_t coord = flat_idx % out_dims[i];
        flat_idx /= out_dims[i];
        if (i < static_cast<int>(pad)) {
            continue;
        }
        const int64_t in_dim = in_dims[i - pad];
        if (in_dim != 1) {
            offset += coord * in_stride;
        }
        in_stride *= in_dim;
    }
    return offset;
}

std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other) {
    /*
        input: [*, m, k], other: [*, k, n] -> output: [broadcast(*), m, n]
        The leading batch dimensions are broadcast against each other; the trailing two dimensions
        are the matrix being multiplied.
    */
    const auto &input_dims = input->Dims();
    const auto &other_dims = other->Dims();
    CHECK_GE(input_dims.size(), 2);
    CHECK_GE(other_dims.size(), 2);

    const int64_t m = input_dims[input_dims.size() - 2];
    const int64_t k = *input_dims.rbegin();
    const int64_t n = *other_dims.rbegin();
    CHECK_EQ(other_dims[other_dims.size() - 2], k)
        << "Matmul inner dimensions mismatch: " << k << " vs " << other_dims[other_dims.size() - 2];

    const std::vector<int64_t> input_batch_dims(input_dims.begin(), input_dims.end() - 2);
    const std::vector<int64_t> other_batch_dims(other_dims.begin(), other_dims.end() - 2);
    const auto batch_dims = BroadcastBatchDims(input_batch_dims, other_batch_dims);
    const int64_t num_batches = NumBatches(batch_dims);

    auto output_dims = batch_dims;
    output_dims.push_back(m);
    output_dims.push_back(n);
    auto output = std::make_shared<Tensor>(output_dims, DataType::kFLOAT32, input->GetDevice());

    const float *input_ptr = static_cast<const float *>(input->DataPtr());
    const float *other_ptr = static_cast<const float *>(other->DataPtr());
    float *output_ptr = static_cast<float *>(output->DataPtr());

    for (int64_t batch = 0; batch < num_batches; ++batch) {
        const int64_t input_offset = BroadcastBatchOffset(batch, batch_dims, input_batch_dims) * m * k;
        const int64_t other_offset = BroadcastBatchOffset(batch, batch_dims, other_batch_dims) * k * n;

        Eigen::Map<const RowMajorMatrix> a(input_ptr + input_offset, m, k);
        Eigen::Map<const RowMajorMatrix> b(other_ptr + other_offset, k, n);
        Eigen::Map<RowMajorMatrix> c(output_ptr + batch * m * n, m, n);
        c.noalias() = a * b;
    }

    return output;
}

std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
MatmulBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other,
               const std::shared_ptr<Tensor> &grad_output) {
    /*
        grad_input[*, m, k] = grad_output[*, m, n] * other[*, k, n]^T
        grad_other[*, k, n] = input[*, m, k]^T * grad_output[*, m, n]

        When a batch dimension was broadcast in Forward, the same operand block was reused by
        several output batches, so its gradient is the sum over those batches -- hence the
        zero-initialization and `+=` below.
    */
    const auto &input_dims = input->Dims();
    const auto &other_dims = other->Dims();
    CHECK_GE(input_dims.size(), 2);
    CHECK_GE(other_dims.size(), 2);

    const int64_t m = input_dims[input_dims.size() - 2];
    const int64_t k = *input_dims.rbegin();
    const int64_t n = *other_dims.rbegin();
    CHECK_EQ(other_dims[other_dims.size() - 2], k);

    const std::vector<int64_t> input_batch_dims(input_dims.begin(), input_dims.end() - 2);
    const std::vector<int64_t> other_batch_dims(other_dims.begin(), other_dims.end() - 2);
    const auto batch_dims = BroadcastBatchDims(input_batch_dims, other_batch_dims);
    const int64_t num_batches = NumBatches(batch_dims);

    const auto &grad_output_dims = grad_output->Dims();
    CHECK_EQ(grad_output_dims.size(), batch_dims.size() + 2);
    CHECK_EQ(grad_output_dims[grad_output_dims.size() - 2], m);
    CHECK_EQ(*grad_output_dims.rbegin(), n);

    auto grad_input = std::make_shared<Tensor>(input_dims, DataType::kFLOAT32, input->GetDevice());
    auto grad_other = std::make_shared<Tensor>(other_dims, DataType::kFLOAT32, other->GetDevice());
    grad_input->Fill<float>(0.0f);
    grad_other->Fill<float>(0.0f);

    const float *input_ptr = static_cast<const float *>(input->DataPtr());
    const float *other_ptr = static_cast<const float *>(other->DataPtr());
    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
    float *grad_input_ptr = static_cast<float *>(grad_input->DataPtr());
    float *grad_other_ptr = static_cast<float *>(grad_other->DataPtr());

    for (int64_t batch = 0; batch < num_batches; ++batch) {
        const int64_t input_offset = BroadcastBatchOffset(batch, batch_dims, input_batch_dims) * m * k;
        const int64_t other_offset = BroadcastBatchOffset(batch, batch_dims, other_batch_dims) * k * n;

        Eigen::Map<const RowMajorMatrix> a(input_ptr + input_offset, m, k);
        Eigen::Map<const RowMajorMatrix> b(other_ptr + other_offset, k, n);
        Eigen::Map<const RowMajorMatrix> grad_c(grad_output_ptr + batch * m * n, m, n);

        Eigen::Map<RowMajorMatrix> grad_a(grad_input_ptr + input_offset, m, k);
        Eigen::Map<RowMajorMatrix> grad_b(grad_other_ptr + other_offset, k, n);

        grad_a.noalias() += grad_c * b.transpose();
        grad_b.noalias() += a.transpose() * grad_c;
    }

    return {grad_input, grad_other};
}

解决思路

  1. 完整的 numpy 广播语义:matmul 的 batch 维按 numpy/PyTorch 规则右对齐,维数为 1 的维拉伸到对方大小。BroadcastBatchDims 计算输出 batch 形状;BroadcastBatchOffset 把输出 batch 的平坦索引逐维反解出各操作数的平坦 batch 索引(广播维不贡献 offset)。
  2. 零拷贝的 Eigen 写法:用 Eigen::Map 以 RowMajor 视图直接映射底层 buffer,避免复制数据;c.noalias() = a * b 告诉 Eigen 目标矩阵与操作数无别名关系,结果直接写进输出矩阵,避免 a*b 产生临时矩阵再赋值。
  3. 反向传播公式grad_input = grad_output × other^Tgrad_other = input^T × grad_output,直接由链式法则得出,用转置免去任何数据搬移。
  4. 广播场景的反向累加(本作业最核心的细节):前向中同一个操作数块被多个输出 batch 复用,因此它的梯度等于所有复用 batch 的梯度之和。实现上 grad_input/grad_otherFill(0),循环里用 += 累加。

遇到问题

  • 广播反向漏加:最初在循环里用 = 赋值而非 +=,广播维的梯度只来自最后一个 batch,数值错误。意识到"被复用的块必须累加梯度"后才修正——这个问题让我真正理解了 broadcast 与反向传播的相互作用。
  • noalias 语义:直接写 c = a * b 时 Eigen 会保守地先算临时矩阵再赋值,白白多一次 O(mnk) 的拷贝;需要 noalias() 显式声明别名安全。

CUDA实现

对应测例:TEST(MatmulTest, BasicMatrixMultiplyCuda),TEST(MatmulTest, BatchedMatrixMultiplyCuda),TEST(MatmulTest, BackwardPassCuda)

需要实现的代码块位置:infini_train/src/kernels/cuda/linear.cu

// An operand whose batch shape either matches the output exactly or is entirely broadcast can be
// walked with a single constant stride, which lets us use one StridedBatched call instead of a
// loop. `stride == 0` means every output batch reads the same block.
bool TryUniformBatchStride(const std::vector<int64_t> &batch_dims, const std::vector<int64_t> &in_batch_dims,
                           int64_t matrix_size, int64_t *stride) {
    if (NumBatches(in_batch_dims) == 1) {
        *stride = 0;
        return true;
    }
    if (in_batch_dims == batch_dims) {
        *stride = matrix_size;
        return true;
    }
    return false;
}

std::shared_ptr<Tensor> MatmulForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other) {
    /*
        input: [*, m, k], other: [*, k, n] -> output: [broadcast(*), m, n]

        cuBLAS is column-major, so a row-major [r, c] buffer is read as a column-major [c, r] matrix
        -- i.e. already transposed. Writing the product in that view:
            output^T = other^T * input^T
        which is a plain OP_N/OP_N gemm on the buffers we already have, with the operands swapped.
    */
    const auto &input_dims = input->Dims();
    const auto &other_dims = other->Dims();
    CHECK_GE(input_dims.size(), 2);
    CHECK_GE(other_dims.size(), 2);

    const int64_t m = input_dims[input_dims.size() - 2];
    const int64_t k = *input_dims.rbegin();
    const int64_t n = *other_dims.rbegin();
    CHECK_EQ(other_dims[other_dims.size() - 2], k)
        << "Matmul inner dimensions mismatch: " << k << " vs " << other_dims[other_dims.size() - 2];

    const std::vector<int64_t> input_batch_dims(input_dims.begin(), input_dims.end() - 2);
    const std::vector<int64_t> other_batch_dims(other_dims.begin(), other_dims.end() - 2);
    const auto batch_dims = BroadcastBatchDims(input_batch_dims, other_batch_dims);
    const int64_t num_batches = NumBatches(batch_dims);

    auto output_dims = batch_dims;
    output_dims.push_back(m);
    output_dims.push_back(n);
    auto output = std::make_shared<Tensor>(output_dims, DataType::kFLOAT32, input->GetDevice());

    const float *input_ptr = static_cast<const float *>(input->DataPtr());
    const float *other_ptr = static_cast<const float *>(other->DataPtr());
    float *output_ptr = static_cast<float *>(output->DataPtr());

    const float alpha = 1.0f;
    const float beta = 0.0f;
    cublasHandle_t handle;
    CUBLAS_CHECK(cublasCreate(&handle));

    int64_t input_stride = 0;
    int64_t other_stride = 0;
    const bool uniform = TryUniformBatchStride(batch_dims, input_batch_dims, m * k, &input_stride)
                      && TryUniformBatchStride(batch_dims, other_batch_dims, k * n, &other_stride);

    if (uniform) {
        // C[n, m] = other^T[n, k] * input^T[k, m], batched
        CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, other_ptr, n,
                                               other_stride, input_ptr, k, input_stride, &beta, output_ptr, n, m * n,
                                               num_batches));
    } else {
        for (int64_t batch = 0; batch < num_batches; ++batch) {
            const int64_t input_offset = BroadcastBatchOffset(batch, batch_dims, input_batch_dims) * m * k;
            const int64_t other_offset = BroadcastBatchOffset(batch, batch_dims, other_batch_dims) * k * n;
            CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, other_ptr + other_offset, n,
                                     input_ptr + input_offset, k, &beta, output_ptr + batch * m * n, n));
        }
    }

    CUBLAS_CHECK(cublasDestroy(handle));
    return output;
}

std::tuple<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>>
MatmulBackward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &other,
               const std::shared_ptr<Tensor> &grad_output) {
    /*
        grad_input[*, m, k] = grad_output[*, m, n] * other[*, k, n]^T
        grad_other[*, k, n] = input[*, m, k]^T * grad_output[*, m, n]

        In the column-major view (see MatmulForward):
            grad_input^T[k, m] = other[k, n] * grad_output^T[n, m]   -> OP_T, OP_N
            grad_other^T[n, k] = grad_output^T[n, m] * input[m, k]   -> OP_N, OP_T

        A batch dimension that was broadcast in Forward reused the same operand block across
        several output batches, so its gradient is the sum over those batches. That case runs the
        accumulating loop (beta = 1) instead of StridedBatched, which would race on one output.
    */
    const auto &input_dims = input->Dims();
    const auto &other_dims = other->Dims();
    CHECK_GE(input_dims.size(), 2);
    CHECK_GE(other_dims.size(), 2);

    const int64_t m = input_dims[input_dims.size() - 2];
    const int64_t k = *input_dims.rbegin();
    const int64_t n = *other_dims.rbegin();
    CHECK_EQ(other_dims[other_dims.size() - 2], k);

    const std::vector<int64_t> input_batch_dims(input_dims.begin(), input_dims.end() - 2);
    const std::vector<int64_t> other_batch_dims(other_dims.begin(), other_dims.end() - 2);
    const auto batch_dims = BroadcastBatchDims(input_batch_dims, other_batch_dims);
    const int64_t num_batches = NumBatches(batch_dims);

    const auto &grad_output_dims = grad_output->Dims();
    CHECK_EQ(grad_output_dims.size(), batch_dims.size() + 2);
    CHECK_EQ(grad_output_dims[grad_output_dims.size() - 2], m);
    CHECK_EQ(*grad_output_dims.rbegin(), n);

    auto grad_input = std::make_shared<Tensor>(input_dims, DataType::kFLOAT32, input->GetDevice());
    auto grad_other = std::make_shared<Tensor>(other_dims, DataType::kFLOAT32, other->GetDevice());
    grad_input->Fill<float>(0.0f);
    grad_other->Fill<float>(0.0f);

    const float *input_ptr = static_cast<const float *>(input->DataPtr());
    const float *other_ptr = static_cast<const float *>(other->DataPtr());
    const float *grad_output_ptr = static_cast<const float *>(grad_output->DataPtr());
    float *grad_input_ptr = static_cast<float *>(grad_input->DataPtr());
    float *grad_other_ptr = static_cast<float *>(grad_other->DataPtr());

    const float alpha = 1.0f;
    const float beta = 1.0f; // grads are zero-initialized above, so accumulate throughout
    cublasHandle_t handle;
    CUBLAS_CHECK(cublasCreate(&handle));

    int64_t input_stride = 0;
    int64_t other_stride = 0;
    const bool input_uniform = TryUniformBatchStride(batch_dims, input_batch_dims, m * k, &input_stride);
    const bool other_uniform = TryUniformBatchStride(batch_dims, other_batch_dims, k * n, &other_stride);

    // grad_input
    if (input_uniform && input_stride != 0 && other_uniform) {
        CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, k, m, n, &alpha, other_ptr, n,
                                               other_stride, grad_output_ptr, n, m * n, &beta, grad_input_ptr, k,
                                               input_stride, num_batches));
    } else {
        for (int64_t batch = 0; batch < num_batches; ++batch) {
            const int64_t input_offset = BroadcastBatchOffset(batch, batch_dims, input_batch_dims) * m * k;
            const int64_t other_offset = BroadcastBatchOffset(batch, batch_dims, other_batch_dims) * k * n;
            CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, k, m, n, &alpha, other_ptr + other_offset, n,
                                     grad_output_ptr + batch * m * n, n, &beta, grad_input_ptr + input_offset, k));
        }
    }

    // grad_other
    if (other_uniform && other_stride != 0 && input_uniform) {
        CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_T, n, k, m, &alpha, grad_output_ptr, n,
                                               m * n, input_ptr, k, input_stride, &beta, grad_other_ptr, n,
                                               other_stride, num_batches));
    } else {
        for (int64_t batch = 0; batch < num_batches; ++batch) {
            const int64_t input_offset = BroadcastBatchOffset(batch, batch_dims, input_batch_dims) * m * k;
            const int64_t other_offset = BroadcastBatchOffset(batch, batch_dims, other_batch_dims) * k * n;
            CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, n, k, m, &alpha,
                                     grad_output_ptr + batch * m * n, n, input_ptr + input_offset, k, &beta,
                                     grad_other_ptr + other_offset, n));
        }
    }

    CUBLAS_CHECK(cublasDestroy(handle));
    return {grad_input, grad_other};
}

解决思路

  1. 列主序推导(本作业的难点):cuBLAS 是列主序,行主序的 [r, c] buffer 被它读成列主序的 [c, r] 矩阵,即"天然已转置"。在这个视角下重写前向:output^T = other^T × input^T,恰好是一次 OP_N/OP_N 的 gemm、操作数交换位置,无需任何物理转置,也不用手写 tile 乘法。
  2. 反向的列主序推导grad_input^T = other × grad_output^T(OP_T/OP_N);grad_other^T = grad_output^T × input(OP_N/OP_T)。
  3. StridedBatched 优化:操作数的 batch 形状若与输出完全一致(stride = 单个矩阵大小)或完全广播(stride = 0,所有 batch 读同一块),就用 cublasSgemmStridedBatched 一次调用完成全部 batch,避免逐 batch 循环启动 kernel。
  4. 反向广播场景的并发安全:被广播的操作数同一块会被多个 batch 并发写入——StridedBatched 会 race。该场景退化为普通 gemm 循环,梯度先零初始化、beta=1 逐 batch 累加(与 CPU 版本的 += 同理)。
  5. Linear 系列:bias 先用自定义 BiasCopyKernel 铺到输出(beta=1 直接累加到预置的 bias 上);grad_biascub::BlockReduce 按行归约(一行一个 block、256 线程步进累加再 block reduce),避免逐元素原子操作。

遇到问题

  • 列主序方向搞反:最初直接按行主序直觉填 cublas 参数,结果张量形状对但数值全错。通过写出"行主序 buffer 在列主序视角 = 转置"这一等式、从 output^T 视角重新推导后才通过——这比记忆"cublas 要传转置"的咒语可靠得多。
  • StridedBatched 与广播的冲突:广播场景用 StridedBatched 时同一块输出被多次写,结果随机。理解"被复用的块 = 需要累加"后,改为 beta=1 的循环累加。
  • CUDA 部分本地从未编译过:本地 CMakeCache.txtUSE_CUDA=OFF,两个 CUDA 测例的二进制根本不存在。在租用的 RTX 5090 服务器(CUDA 12.8)上开 CUDA 编译后,test_matmul_cuda / test_adam_cuda 零错误通过——实现是好的,只是此前没验证过。

作业三:实现Adam优化器

难度:⭐

CPU实现

对应测例:TEST(AdamOptimizerTest, BasicParameterUpdate),TEST(AdamOptimizerTest, MomentumAccumulation)

代码位置:infini_train/src/kernels/cpu/accumulate_grad.cc

void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_ptr<Tensor> &param,
                        const std::shared_ptr<Tensor> &m, const std::shared_ptr<Tensor> &v, float learning_rate,
                        float beta1, float beta2, float eps, int64_t t) {
    /*
        m_t = beta1 * m_{t-1} + (1 - beta1) * g
        v_t = beta2 * v_{t-1} + (1 - beta2) * g^2
        m_hat = m_t / (1 - beta1^t)      // bias correction: m_0 = 0 biases early steps toward zero
        v_hat = v_t / (1 - beta2^t)
        param -= learning_rate * m_hat / (sqrt(v_hat) + eps)

        t is 1-based (Adam::Step increments before calling), so the correction terms are non-zero.
    */
    const float *grad_ptr = static_cast<const float *>(grad->DataPtr());
    float *param_ptr = static_cast<float *>(param->DataPtr());
    float *m_ptr = static_cast<float *>(m->DataPtr());
    float *v_ptr = static_cast<float *>(v->DataPtr());

    const float bias_correction1 = 1.0f - std::pow(beta1, static_cast<float>(t));
    const float bias_correction2 = 1.0f - std::pow(beta2, static_cast<float>(t));

    for (int64_t idx = 0; idx < grad->NumElements(); ++idx) {
        const float g = grad_ptr[idx];

        m_ptr[idx] = beta1 * m_ptr[idx] + (1.0f - beta1) * g;
        v_ptr[idx] = beta2 * v_ptr[idx] + (1.0f - beta2) * g * g;

        const float m_hat = m_ptr[idx] / bias_correction1;
        const float v_hat = v_ptr[idx] / bias_correction2;

        param_ptr[idx] -= learning_rate * m_hat / (std::sqrt(v_hat) + eps);
    }
}

CUDA实现

对应测例:TEST(AdamOptimizerTest, BasicParameterUpdateCuda),TEST(AdamOptimizerTest, MomentumAccumulationCuda)

代码位置:infini_train/src/kernels/cuda/accumulate_grad.cu

__global__ void AdamAccumulateGradKernel(const float *grad_ptr, float *param_ptr, float *m_ptr, float *v_ptr,
                                         size_t num_elements, float learning_rate, float beta1, float beta2, float eps,
                                         float bias_correction1, float bias_correction2) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_elements) {
        return;
    }

    const float g = grad_ptr[idx];

    const float m = beta1 * m_ptr[idx] + (1.0f - beta1) * g;
    const float v = beta2 * v_ptr[idx] + (1.0f - beta2) * g * g;
    m_ptr[idx] = m;
    v_ptr[idx] = v;

    const float m_hat = m / bias_correction1;
    const float v_hat = v / bias_correction2;

    param_ptr[idx] -= learning_rate * m_hat / (sqrtf(v_hat) + eps);
}

void AdamAccumulateGrad(const std::shared_ptr<Tensor> &grad, const std::shared_ptr<Tensor> &param,
                        const std::shared_ptr<Tensor> &m, const std::shared_ptr<Tensor> &v, float learning_rate,
                        float beta1, float beta2, float eps, int64_t t) {
    /*
        m_t = beta1 * m_{t-1} + (1 - beta1) * g
        v_t = beta2 * v_{t-1} + (1 - beta2) * g^2
        m_hat = m_t / (1 - beta1^t)      // bias correction: m_0 = 0 biases early steps toward zero
        v_hat = v_t / (1 - beta2^t)
        param -= learning_rate * m_hat / (sqrt(v_hat) + eps)

        The bias-correction denominators only depend on t, so they are computed once on the host.
    */
    size_t num_elements = grad->NumElements();

    const float *grad_ptr = static_cast<const float *>(grad->DataPtr());
    float *param_ptr = static_cast<float *>(param->DataPtr());
    float *m_ptr = static_cast<float *>(m->DataPtr());
    float *v_ptr = static_cast<float *>(v->DataPtr());

    const float bias_correction1 = 1.0f - std::pow(beta1, static_cast<float>(t));
    const float bias_correction2 = 1.0f - std::pow(beta2, static_cast<float>(t));

    int threads_per_block = 256;
    int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block;

    AdamAccumulateGradKernel<<<num_blocks, threads_per_block>>>(grad_ptr, param_ptr, m_ptr, v_ptr, num_elements,
                                                                learning_rate, beta1, beta2, eps, bias_correction1,
                                                                bias_correction2);
}

解决思路

  1. 按论文公式逐元素实现
    • 一阶矩:m_t = β1·m_{t-1} + (1-β1)·g
    • 二阶矩:v_t = β2·v_{t-1} + (1-β2)·g²
    • 偏置校正:m̂ = m_t/(1-β1^t)v̂ = v_t/(1-β2^t)
    • 更新:param -= lr·m̂/(√v̂ + eps)
  2. 偏置校正的动机:m、v 初始为 0,前几步会被严重"低估"(尤其 β2 接近 1 时 v 起步极慢,导致初始步长被放大);除以 (1-β^t) 消除这一初始偏差,这是 Adam 相比 RMSProp 的关键改进。
  3. t 的语义:t 是 1-based(Adam::Step 先自增再调用 kernel),保证校正分母非零、公式与论文一致。
  4. CUDA 化:偏置校正分母只依赖 t,host 算一次传入 kernel 即可(避免每元素重复 pow);kernel 采用一维索引、每线程一个元素、256 线程/block、grid 覆盖全部元素,元素间完全独立无需同步。

遇到问题

  • 容易漏掉偏置校正:最初实现直接 param -= lr·m/(√v+eps),前几步有效学习率被严重缩小、收敛明显变慢;对照 Adam 论文公式(Algorithm 1)补上 (1-β^t) 后数值与参考一致。
  • CPU/CUDA 一致性:两套实现保持完全相同的公式与计算顺序,CPU 与 CUDA 测例对拍通过。

作业四:实现Tensor基础操作

实现Tensor的Flatten操作

难度:⭐

对应测例:TEST(TensorTransformTest, Flatten2DTo1D),TEST(TensorTransformTest, FlattenWithRange) ,TEST(TensorTransformTest, FlattenNonContiguous)

代码位置:infini_train/src/tensor.cc

std::shared_ptr<Tensor> Tensor::Flatten(int64_t start, int64_t end) {
    // Collapses dims [start, end] (both inclusive, negatives count from the back) into one.
    // Contiguous() first because View only reinterprets the shape of an already-packed buffer --
    // a transposed tensor's elements are not in row-major order, so viewing it directly would
    // scramble them.
    const int64_t ndim = static_cast<int64_t>(dims_.size());

    // A 0-d tensor flattens to a single-element 1-d tensor, matching torch.flatten.
    if (ndim == 0) {
        return Contiguous()->View({1});
    }

    if (start < 0) {
        start += ndim;
    }
    if (end < 0) {
        end += ndim;
    }
    CHECK_GE(start, 0);
    CHECK_LT(start, ndim);
    CHECK_GE(end, 0);
    CHECK_LT(end, ndim);
    CHECK_LE(start, end) << "Flatten requires start <= end, got start=" << start << ", end=" << end;

    std::vector<int64_t> new_shape;
    new_shape.reserve(ndim - (end - start));
    new_shape.insert(new_shape.end(), dims_.begin(), dims_.begin() + start);
    new_shape.push_back(
        std::accumulate(dims_.begin() + start, dims_.begin() + end + 1, int64_t{1}, std::multiplies<int64_t>()));
    new_shape.insert(new_shape.end(), dims_.begin() + end + 1, dims_.end());

    return Contiguous()->View(new_shape);
}

实现Tensor的反向传播机制

难度:⭐

对应测例:TEST(TensorAutogradTest, BackwardComputesGradient),TEST(TensorAutogradTest, BackwardWithMultipleOutputs)

代码位置:infini_train/src/tensor.cc

void Tensor::Backward(std::shared_ptr<Tensor> gradient, bool retain_graph, bool create_graph) const {
    // Seeds the reverse pass. The traversal itself lives in Function::BackwardPartial, which counts
    // how many consumers each node has (recorded by Function::Apply) and only invokes a node's
    // Backward once every consumer has handed back its gradient -- that is what makes a tensor used
    // by several branches accumulate rather than fire twice with partial gradients.
    if (!gradient) {
        // d(self)/d(self) == 1. Only meaningful for a scalar; a non-scalar root needs an explicit
        // gradient, matching torch's "grad can be implicitly created only for scalar outputs".
        CHECK_EQ(num_elements_, 1) << "Backward() on a non-scalar tensor requires an explicit gradient";
        gradient = std::make_shared<Tensor>(dims_, dtype_, GetDevice());
        gradient->Fill<float>(1.0f);
    }
    CHECK_EQ(gradient->NumElements(), num_elements_)
        << "Gradient has " << gradient->NumElements() << " elements but tensor has " << num_elements_;

    if (!grad_fn_) {
        // A leaf has no graph behind it, so the seed gradient is its whole gradient.
        if (requires_grad_ && grad_) {
            auto kernel = Dispatcher::Instance().GetKernel({GetDevice().Type(), "AccumulateGrad"});
            kernel.Call<void>(gradient, 1.0f, grad_);
        }
        return;
    }

    grad_fn_->BackwardPartial(gradient, output_idx_);
}

解决思路

Flatten

  1. 先 Contiguous 再 View(关键)View 只是"重新解释 shape",不搬数据。transpose 过的张量内存不是行主序紧凑排布,直接 View 会打乱元素顺序;Contiguous() 先保证底层是行主序紧凑存储,View 才是安全的。
  2. 兼容负索引(-1 表示最后一维),并处理 0-d 张量 flat 成 [1](对标 torch.flatten)。
  3. 新 shape 构造:[start 之前的各维] + [start..end 维的乘积] + [end 之后的各维]

Backward

  1. 种子梯度:无显式 gradient 时创建全 1 梯度(d(self)/d(self)=1);对标 torch 只允许对 scalar 隐式创建,非 scalar 根节点必须显式给梯度(CHECK 报错)。
  2. 叶子节点:图后面没有节点,种子梯度即它的完整梯度——通过 AccumulateGrad kernel 累进 grad_
  3. 非叶子节点grad_fn_->BackwardPartial(gradient, output_idx_)多输出累加的机制在 Function::Apply:每次 Apply 通过 IncreaseDependenciesNumber() 记录节点被消费的次数;BackwardPartial 只有在一个节点的所有消费者都交回梯度后才调用它的 Backward——这保证共享节点只执行一次、且拿到的是完整累加后的梯度,而不是多次以部分梯度触发。

遇到问题

  • FlattenNonContiguous 测例:第一次实现直接 View(new_shape),transpose 后的张量数据全错。读 NoOp(View 的算子)实现后才明白 View 不拷贝、只重解释 shape,必须先 Contiguous()。这是"为什么必须连续"的最直观教材。
  • 多输出场景的梯度累加:最初想直接在 Backward 里递归调用各算子的 Backward,但被多个分支共享的节点会以"部分梯度"执行多次。理解 Function 里的 dependencies 计数(Apply 时登记、BackwardPartial 时等齐)之后,用框架自带的 BackwardPartial 正确实现了"消费全部完成才反向"的语义。

作业五 注册算子kernel的实现

难度:⭐⭐⭐

对应测例:TEST(DispatcherTest, RegisterAndGetKernel),TEST(DispatcherTest, DuplicateRegistration),TEST(DispatcherTest, GetNonexistentKernel)

代码位置:infini_train/include/dispatcher.h

template <typename RetT, class... ArgsT> RetT Call(ArgsT... args) const {
    // The kernel is stored type-erased as a void *, so the caller declares the signature it
    // expects and we cast back to it before invoking. ArgsT is deduced from the call site, so
    // a kernel declared as `f(const std::shared_ptr<Tensor> &)` ends up being called through a
    // `f(std::shared_ptr<Tensor>)` pointer: both pass class types by address in the platform
    // ABI, which is what makes the round trip work.
    CHECK(func_ptr_ != nullptr) << "Calling an empty kernel function";

    using FuncT = RetT (*)(ArgsT...);
    return reinterpret_cast<FuncT>(func_ptr_)(std::forward<ArgsT>(args)...);
}

template <typename FuncT> void Register(const KeyT &key, FuncT &&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)));
}

// Binds `kernel_func` to (device, "kernel_name"). The Register() call is the initializer of a
// variable, so at namespace scope it runs before main() with no explicit setup code, and inside a
// function body it runs when control reaches that line. __LINE__ keeps the variable name unique so
// the same kernel name can appear twice in one translation unit.
#define REGISTER_KERNEL(device, kernel_name, kernel_func)                                                              \
    [[maybe_unused]] static const bool INFINI_TRAIN_CONCAT(kernel_registrar_##kernel_name##_, __LINE__) = []() {       \
        infini_train::Dispatcher::Instance().Register({(device), #kernel_name}, kernel_func);                          \
        return true;                                                                                                  \
    }();

解决思路

  1. 类型擦除的 KernelFunction:函数指针以 void* 存进容器(std::map<KeyT, KernelFunction> 要求可拷贝的成员)。Call 的签名由调用点模板参数确定:返回类型 RetT 显式指定、参数类型 ArgsT 由实参推导,再 reinterpret_cast<RetT(*)(ArgsT...)> 还原后调用。
  2. 为什么跨 ABI 可行(Call 里注释的点):kernel 声明为 f(const shared_ptr<Tensor>&),调用点推导出的指针却是 f(shared_ptr<Tensor>)——类类型参数在平台 ABI 下都按地址传递,二者是兼容的,所以这趟类型擦除往返是安全的。
  3. Register 的防覆盖:重复注册直接 CHECK 报错(DuplicateRegistration 测例就靠它),避免同 key 静默覆盖导致难排查的 bug;插入用 emplace 不产生额外拷贝。
  4. REGISTER_KERNEL 宏 = 静态变量 + 立即 lambda:声明一个 static const bool 变量,初始化器是立即执行的 lambda,lambda 内调用 Register。namespace 作用域的静态变量在 main 之前完成初始化 → 零显式初始化代码;写在函数体内则执行到该行时注册。__LINE__ 拼接保证同一翻译单元里同名 kernel 的变量名不冲突。

遇到问题

  • 宏变量名冲突:最初用固定变量名,同一 TU 注册两个 kernel 就"重定义"编译失败;拼接 __LINE__ 后解决。
  • 未注册 key 的报错质量GetKernel 对不存在的 key 直接 CHECK,且错误信息带上 device 与 kernel 名,方便定位(GetNonexistentKernel 测例正是验证这一点)。
  • 顺带理解了整个框架的架构:kernel 实现(kernels/cpu| cuda/)与调用方(autograd/)完全解耦,新增算子只需"实现 + 一行 REGISTER_KERNEL",这正是 test_dispatcher 作为"核心基础设施"值 20 分的原因。

作业六:实现GPT-2整体训练

难度:⭐⭐⭐⭐

对应测例:TEST_F(GPT2TrainingTest, LogitsConsistency)

训练过程logits对比

完成以上所有作业,补齐训练框架的所有实现,理论上TEST_F(GPT2TrainingTest, LogitsConsistency)可以通过,在用例中判断比较预置的值和单步正向传播计算结果是否在误差允许范围内相等。

数据读取实现

代码位置:example/common/tiny_shakespeare_dataset.cc

TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequence_length) {
    /* File layout (produced by llm.c's tinyshakespeare.py):
    ----------------------------------------------------------------------------------
    | HEADER (1024 bytes)                                        | DATA (tokens)     |
    | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B)    | token stream      |
    ----------------------------------------------------------------------------------
       The on-disk tokens are uint16 (GPT-2) or uint32 (LLaMA 3), but they are widened to int64
       here: Embedding requires kINT64 indices, and TinyShakespeareDataset::operator[] shifts the
       label window by exactly sizeof(int64_t) to get y = x advanced by one token.
    */
    if (!std::filesystem::exists(path)) {
        LOG(FATAL) << "Dataset file not found: " << path;
    }
    CHECK_GT(sequence_length, 0);

    std::ifstream ifs(path, std::ios::binary);
    CHECK(ifs.is_open()) << "Failed to open dataset file: " << path;

    const auto header = ReadSeveralBytesFromIfstream(256 * sizeof(int32_t), &ifs);
    const auto magic = BytesToType<int32_t>(header, 0);
    CHECK(kTypeMap.contains(magic)) << "Unsupported dataset magic number: " << magic;
    const auto type = kTypeMap.at(magic);
    const auto version = BytesToType<int32_t>(header, 4);
    CHECK_EQ(version, 1) << "Unsupported dataset version: " << version;
    const auto num_tokens = static_cast<size_t>(BytesToType<int32_t>(header, 8));
    CHECK_GT(num_tokens, sequence_length) << "Dataset holds " << num_tokens << " tokens, which is not enough for a "
                                          << sequence_length << "-token sequence plus its shifted label";

    // One sample is a sequence_length window; the labels of the last window would run past the end
    // of the file, so operator[] exposes dims[0] - 1 samples.
    const size_t num_sequences = num_tokens / sequence_length;
    const size_t num_used_tokens = num_sequences * sequence_length;

    TinyShakespeareFile text_file;
    text_file.type = type;
    text_file.dims = {static_cast<int64_t>(num_sequences), static_cast<int64_t>(sequence_length)};
    text_file.tensor = infini_train::Tensor(text_file.dims, DataType::kINT64);

    const size_t bytes_per_token = kTypeToSize.at(type);
    const auto raw = ReadSeveralBytesFromIfstream(num_used_tokens * bytes_per_token, &ifs);
    CHECK_EQ(ifs.gcount(), static_cast<std::streamsize>(num_used_tokens * bytes_per_token))
        << "Dataset file is shorter than its header claims: " << path;

    int64_t *tokens = static_cast<int64_t *>(text_file.tensor.DataPtr());
    for (size_t idx = 0; idx < num_used_tokens; ++idx) {
        tokens[idx] = type == TinyShakespeareType::kUINT16
                        ? static_cast<int64_t>(BytesToType<uint16_t>(raw, idx * bytes_per_token))
                        : static_cast<int64_t>(BytesToType<uint32_t>(raw, idx * bytes_per_token));
    }

    return text_file;
}

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)),
      // The final window has no room for its shifted labels, so it is not a usable sample.
      num_samples_(text_file_.dims[0] - 1) {
    CHECK_GT(num_samples_, 0) << "Dataset is too small to produce a single sample";
}

Tokenizer功能实现

代码位置:example/common/tokenizer.cc

Tokenizer::Tokenizer(const std::string &filepath) {
    /* File layout (produced by llm.c):
    ----------------------------------------------------------------------------------
    | HEADER (1024 bytes)                                          | VOCAB TABLE     |
    | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B)    | token entries   |
    ----------------------------------------------------------------------------------
       Each vocab entry is a length byte followed by that many raw bytes of text. v1 files predate
       the eot token being recorded in the header, so it comes from kEotMap keyed on the magic.
    */
    if (!std::filesystem::exists(filepath)) {
        LOG(FATAL) << "Tokenizer file not found: " << filepath;
    }

    std::ifstream ifs(filepath, std::ios::binary);
    CHECK(ifs.is_open()) << "Failed to open tokenizer file: " << filepath;

    const auto header = ReadSeveralBytesFromIfstream(256 * sizeof(uint32_t), &ifs);
    magic_number_ = BytesToType<uint32_t>(header, 0);
    CHECK(kEotMap.contains(magic_number_)) << "Unsupported tokenizer magic number: " << magic_number_;

    const auto version = BytesToType<uint32_t>(header, 4);
    vocab_size_ = BytesToType<uint32_t>(header, 8);
    CHECK_GT(vocab_size_, 0);

    if (version == static_cast<uint32_t>(Version::kV1)) {
        eot_token_ = kEotMap.at(magic_number_);
    } else if (version == static_cast<uint32_t>(Version::kV2)) {
        // v2 records the eot token explicitly in the header.
        eot_token_ = BytesToType<uint32_t>(header, 12);
    } else {
        LOG(FATAL) << "Unsupported tokenizer version: " << version;
    }

    token_table_.reserve(vocab_size_);
    for (uint32_t i = 0; i < vocab_size_; ++i) {
        const auto length = ReadSeveralBytesFromIfstream(1, &ifs);
        CHECK(ifs.good()) << "Tokenizer file truncated while reading length of token " << i;
        const uint8_t token_length = length[0];
        const auto token_bytes = ReadSeveralBytesFromIfstream(token_length, &ifs);
        CHECK_EQ(ifs.gcount(), static_cast<std::streamsize>(token_length))
            << "Tokenizer file truncated while reading token " << i;
        token_table_.emplace_back(reinterpret_cast<const char *>(token_bytes.data()), token_length);
    }
}

std::string Tokenizer::Decode(uint32_t token_id) const {
    if (token_id >= vocab_size_) {
        LOG(ERROR) << "token_id " << token_id << " is out of range (vocab_size=" << vocab_size_ << ")";
        return "";
    }
    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});
    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_; }

    // Give some contexts: "The meaning of life is "
    auto prompt = kPromptMap.at(magic_number_);
    auto prompt_len = prompt.size();
    for (int i = 0; i < prompt_len; ++i) { x_buff[i] = prompt[i]; }
    std::cout << "The meaning of life is";

    auto x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
    uint64_t rng_state = kRngState;
    LOG(INFO) << "start generate text:";
    for (int t = prompt_len; t < text_length; t++) {
        // Autoregressive step: run the whole window, then read the distribution the model predicts
        // for the slot after the last token we have written, i.e. position t - 1.
        auto logits = model.Forward({x})[0];
        const auto &logits_dims = logits->Dims();
        CHECK_EQ(logits_dims.size(), 3);
        const int64_t vocab_size = *logits_dims.rbegin();
        CHECK_EQ(logits_dims[1], static_cast<int64_t>(sequence_length));

        // Sampling happens on the host, so pull the logits back before reading them.
        auto logits_cpu = logits->To(Device());
        const float *logits_ptr = static_cast<const float *>(logits_cpu.DataPtr()) + (t - 1) * vocab_size;

        // Softmax over the vocabulary, shifted by the max for numerical stability.
        std::vector<float> probabilities(vocab_size);
        float max_logit = logits_ptr[0];
        for (int64_t v = 1; v < vocab_size; ++v) { max_logit = std::max(max_logit, logits_ptr[v]); }
        float sum_exp = 0.0f;
        for (int64_t v = 0; v < vocab_size; ++v) {
            probabilities[v] = std::exp(logits_ptr[v] - max_logit);
            sum_exp += probabilities[v];
        }
        for (int64_t v = 0; v < vocab_size; ++v) { probabilities[v] /= sum_exp; }

        const float coin = RandomF32(rng_state);
        const uint32_t next_token = SampleMult(probabilities.data(), static_cast<int>(vocab_size), coin);

        std::cout << Decode(next_token) << std::flush;

        // Append the sampled token to every sequence in the batch and re-upload the window. Only
        // batch row 0 is sampled from, so all rows stay identical and the whole batch advances
        // together.
        for (uint32_t b = 0; b < batch_size; ++b) { x_buff[b * sequence_length + t] = next_token; }
        x = std::make_shared<infini_train::Tensor>(x_tensor.To(device));
    }
    std::cout << std::endl;
}

解决思路

数据读取

  1. 校验 1024 字节头部:magic(区分 GPT-2 20240520 的 uint16 与 LLaMA3 20240801 的 uint32 两种格式)、version、num_tokens。
  2. token 拓宽到 int64:磁盘上是 uint16/uint32,但 Embedding 需要 kINT64 索引,读入时直接转 int64 存储;operator[] 的 label 窗口只偏移 sizeof(int64_t) 就得到 y = x 后移一个 token。
  3. 样本切窗:每样本是 sequence_length 的窗口,最后一个窗口没有 label 的空间,所以 num_samples = num_sequences - 1operator[] 通过 Tensor 共享 buffer + offset 构造 x/y,零拷贝

Tokenizer

  1. 头部解析 magic/version/vocab_size;v1 文件头部不记录 eot token,按 magic 从 kEotMap 查(GPT-2=50256 / LLaMA3=128001),v2 从头部偏移 12 处显式读。
  2. vocab 表逐项读:1 字节长度 + 对应字节数文本。
  3. Decode 查表返回,越界记 ERROR 并返回空串(防御 GPT-2 的 50257 边界)。
  4. GenerateText 单步生成:整窗前向 → 取"最后一个已写 token 的下一位置"(t-1)的 logits → 拷回 host → softmax(减 max 保证数值稳定)→ 确定性 rng(RandomU32 位运算 LCG,固定种子 1337,保证可复现)采样 → Decode 输出 → 采样 token 写回 batch 所有行并重新上传设备(仅行 0 被采样,batch 内保持一致)。

端到端训练链路:DataLoader → GPT2 Forward → CrossEntropyLoss → Backward(梯度累积 total_batch_size / tokens_per_fwdbwd 步)→ SGD Step,迭代 11 步后验证 logits。

遇到问题

  1. logits 位置索引差一:生成时易把 t-1 写成 t(读到的会是"下一个位置的分布"而不是"最后已写 token 的分布")。对照"整窗输入、输出每个位置对下一 token 的预测"语义后修正。
  2. test_gpt2 未通过(LogitsConsistency),经完整诊断确认是测例容差问题而非实现 bug
    • 现象:失败点 idx=385973,ref=-204.265、cur=-204.263,diff=0.0023 > 绝对容差 1e-3

    • 误差分布:100 个采样点相对误差全部恒定在 ~1e-5,仅 3 个点超过绝对容差,且全部位于 |ref|>130 的大数值处——绝对误差随数值线性放大,是 float32 精度特征,而非逻辑错误。

    • 排除 TF32:强制 CUBLAS_PEDANTIC_MATH 后误差几乎不变(0.0023→0.0021),Blackwell 的 TF32 不是原因。

    • 决定性证据——逐步收敛曲线(每步取全部采样点 max |diff|):

      step 0:  max_abs=51.34   step 4:  18.06   step 8:  5.34
      step 2:  30.60           step 6:  11.23   step 10: 0.0065
      

      误差从 51.34 单调收敛到 0.0065。reference 就是"11 步训练后"的快照,我的收敛轨迹与它精确吻合到 float32 极限(相对误差 5e-5)。有 bug 的实现不会这样逐步收敛。

    • 环境对照实验(决定性证据):将环境升级至与往届通过者完全一致(CUDA 13.3.1 + native sm_120 编译 + RTX 5090 + 驱动 610.43.02)后,失败点 diff 从稳定超差的 0.0023 降至 0.00107(仅超容差 7%),实测 8 次运行 6 过 2 挂(约 75-80% 通过率),与往届通过者在报告中记录的"约 1e-3 量级的 CUDA 浮点边界差异,单测重跑可通过"完全吻合。在与通过者完全相同的环境(也是当前可获得的最新环境:CUDA 13.3.1 为官方最新版本,无更新版本可升)中依然偶发失败,说明剩余差异是 float32 累加顺序非确定性的固有行为,而非实现问题。

    • 结论:训练本身正确,卡点在于测例用绝对容差 1e-3 衡量 |value|~200 的 logits——float32 只有约 7 位有效数字,多步训练的浮点累加顺序差异(不同 GPU 架构 / cuBLAS 版本 / CPU vs GPU)会被迭代逐步放大到 1e-3 量级。按作业说明"发现问题不改用例、反馈助教",已将此结论与收敛曲线、环境对照实验数据作为证据反馈给助教。

  3. CUDA 路径此前未验证:本地 CMakeCache 为 USE_CUDA=OFF,CUDA 测例从未编译。在租用的 RTX 5090 服务器上全量开 CUDA 编译后,test_matmul_cuda / test_adam_cuda 零错误通过。
  4. CUDA 13 编译兼容问题:升级到 CUDA 13.3 后 softmax.cu / reduction.cu / cross_entropy.cu 编译失败——CCCL 3.x 的 <cuda/std/functional> 头文件在 namespace infini_train::kernels::cuda(命名空间本身名为 cuda)内包含时,头文件内部的相对引用 cuda::std::* 会被解析到 infini_train::kernels::cuda::std,导致 void_t/integral_constant 等模板全部解析失败。修复方式:将 #include <cuda/std/functional> 移到文件顶部全局作用域(对 CUDA 12.8 亦兼容)。