第 5 章 预训练：从随机权重到连贯文本

// inputs 两句话，targets 是各自右移一位的结果
auto inputs  = torch::tensor({{16833L, 3626L, 6100L},   // "every effort moves"
                              {40L,    1107L,  588L}},  // "I really like"
                             torch::kLong);
auto targets = torch::tensor({{3626L, 6100L,  345L},    // " effort moves you"
                              {1107L,  588L, 11311L}},  // " really like chocolate"
                             torch::kLong);

auto probas = torch::softmax(logits, -1);               // (2, 3, 50257)
// probas[0, :, [3626, 6100, 345]] 后取 diagonal —— 每个位置上目标词元的概率
auto target_probas_1 = probas
    .index({0, torch::indexing::Slice(),
            torch::tensor({3626L, 6100L, 345L})})
    .diagonal();
auto log_probas = torch::log(torch::cat({target_probas_1, target_probas_2}));
auto neg_avg = torch::mean(log_probas) * -1;            // 手工交叉熵

inline torch::Tensor calc_loss_batch(const torch::Tensor& input_batch,
                                     const torch::Tensor& target_batch,
                                     ch4::GPTModel& model) {
    auto logits = model->forward(input_batch);
    auto loss = torch::nn::functional::cross_entropy(
        logits.flatten(0, 1),          // (batch*seq, vocab)
        target_batch.to(logits.device()).flatten());
    return loss;
}

inline double calc_loss_loader(ch4::GPTModel& model,
                               const std::vector<ch2::GPTBatch>& loader,
                               c10::optional<int64_t> num_batches = c10::nullopt) {
    torch::NoGradGuard no_grad;
    model->eval();
    for (int64_t i = 0; i < n; ++i) {
        total_loss += calc_loss_batch(loader[i].inputs, loader[i].targets,
                                      model).item<double>();
    }
    model->train();                       // 评估完切回训练模式
    return total_loss / n;
}

auto next_token_logits = torch::tensor(
    {4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79});
//        closer every effort forward inches moves pizza toward you

multinomial 采样 1000 次:  582 x forward, 343 x toward, 73 x closer, 2 x inches
T=1.0   P(forward)=0.5721   P(pizza)=0.0001   // 原始分布
T=0.1   P(forward)=0.9910   P(pizza)=0.0000   // 趋近 one-hot，确定性强
T=5.0   P(forward)=0.2421   P(pizza)=0.0454   // 趋近均匀，"pizza"也有机会

if (top_k.has_value()) {
    auto top_values = std::get<0>(torch::topk(logits, k));
    auto threshold  = top_values.narrow(1, k - 1, 1);      // 第 k 大的值
    logits = torch::where(logits < threshold,
                          torch::full_like(logits, -std::numeric_limits<double>::infinity()),
                          logits);
}
torch::Tensor idx_next;
if (temperature > 0.0) {                                   // 温度采样
    auto probs = torch::softmax(logits / temperature, -1);
    idx_next = torch::multinomial(probs, 1);
} else {                                                   // 贪心
    idx_next = torch::argmax(logits, -1, /*keepdim=*/true);
}
if (eos_id.has_value() && idx_next.item<int64_t>() == eos_id.value()) break;

[8 字节 uint64 LE]  header 长度 N
[N 字节 JSON]       {"wte.weight": {"dtype":"F32","shape":[50257,768],
                                      "data_offsets":[0,154389504]}, ...}
[数据区]            各张量按 data_offsets 连续存放（float32, row-major）

auto c_attn_w = params("blocks." + std::to_string(b) + ".attn.c_attn.w");
auto split = c_attn_w.split(cfg.emb_dim, /*dim=*/1);        // [in,emb] × 3
block.att->W_query->weight =
    assign(block.att->W_query->weight, split[0].t().contiguous());
block.att->W_key->weight =
    assign(block.att->W_key->weight, split[1].t().contiguous());
block.att->W_value->weight =
    assign(block.att->W_value->weight, split[2].t().contiguous());

auto assign = [](torch::Tensor left, const torch::Tensor& right) {
    if (left.sizes() != right.sizes())
        throw std::runtime_error("Shape mismatch");
    { torch::NoGradGuard g; left.copy_(right); }
    return left;
};