Transformer 学习记录

typedef struct {
    int dim;           // 词嵌入向量的长度
    int hidden_dim;    // FFN 中间层的宽度
    int n_layers;      // Transformer Block 重复堆叠的次数
    int n_heads;       // Multi-Head Attention 中 Query 的头数
    int n_kv_heads;    // Key 和 Value 的头数（可以小于 n_heads）
    int vocab_size;    // 词表大小
    int seq_len;       // 最大序列长度（上下文窗口）
} Config;

typedef struct {
    float* token_embedding_table;  // (vocab_size, dim) Token 嵌入表
    float* rms_att_weight;         // (layer, dim) Attention 前的归一化权重
    float* rms_ffn_weight;         // (layer, dim) FFN 前的归一化权重
    float* wq;                     // (layer, dim, n_heads * head_size) Query 矩阵
    float* wk;                     // (layer, dim, n_kv_heads * head_size) Key 矩阵
    float* wv;                     // (layer, dim, n_kv_heads * head_size) Value 矩阵
    float* wo;                     // (layer, n_heads * head_size, dim) 输出矩阵
    float* w1;                     // (layer, hidden_dim, dim) FFN 升维矩阵（Gate）
    float* w2;                     // (layer, dim, hidden_dim) FFN 降维矩阵
    float* w3;                     // (layer, hidden_dim, dim) FFN 升维矩阵（Up）
    float* rms_final_weight;       // (dim,) 最终归一化权重
    float* wcls;                   // 输出分类器（可能与 embedding 共享）
} TransformerWeights;

typedef struct {
    float *x;            // 当前激活值 (dim,)
    float *xb;           // 残差分支缓冲区 (dim,)
    float *xb2;          // 额外缓冲区 (dim,)
    float *hb;           // FFN 隐藏层缓冲区 (hidden_dim,)
    float *hb2;          // FFN 隐藏层缓冲区 (hidden_dim,)
    float *q;            // Query (dim,)
    float *k;            // Key (dim,)
    float *v;            // Value (dim,)
    float *att;          // 注意力分数 (n_heads, seq_len)
    float *logits;       // 输出 logits
    float* key_cache;    // KV 缓存 (layer, seq_len, dim)
    float* value_cache;  // KV 缓存 (layer, seq_len, dim)
} RunState;

// 升维
matmul(s->hb, s->xb, w->w1 + l*dim*hidden_dim, dim, hidden_dim);
matmul(s->hb2, s->xb, w->w3 + l*dim*hidden_dim, dim, hidden_dim);

// SwiGLU 激活
for (int i = 0; i < hidden_dim; i++) {
    float val = s->hb[i];
    // SiLU(x) = x * σ(x)
    val *= (1.0f / (1.0f + expf(-val)));
    // 逐元素乘以 w3(x)
    val *= s->hb2[i];
    s->hb[i] = val;
}

// 降维
matmul(s->xb, s->hb, w->w2 + l*dim*hidden_dim, hidden_dim, dim);

void rmsnorm(float* o, float* x, float* weight, int size) {
    // 计算平方和
    float ss = 0.0f;
    for (int j = 0; j < size; j++) {
        ss += x[j] * x[j];
    }
    ss /= size;
    ss += 1e-5f; // epsilon
    ss = 1.0f / sqrtf(ss);
    // 归一化并缩放
    for (int j = 0; j < size; j++) {
        o[j] = weight[j] * (ss * x[j]);
    }
}

// RoPE 相对位置编码：在每个头内旋转 q 和 k
for (int i = 0; i < dim; i+=2) {
    int head_dim = i % head_size;
    float freq = 1.0f / powf(10000.0f, head_dim / (float)head_size);
    float val = pos * freq;
    float fcr = cosf(val);
    float fci = sinf(val);
    int rotn = i < kv_dim ? 2 : 1; // 2 = q & k, 1 = q only
    for (int v = 0; v < rotn; v++) {
        float* vec = v == 0 ? s->q : s->k;
        float v0 = vec[i];
        float v1 = vec[i+1];
        vec[i]   = v0 * fcr - v1 * fci;
        vec[i+1] = v0 * fci + v1 * fcr;
    }
}

float* forward(Transformer* transformer, int token, int pos) {
    // 1. 获取 token 的嵌入向量
    float* content_row = w->token_embedding_table + token * dim;
    memcpy(x, content_row, dim*sizeof(*x));
    
    // 2. 遍历所有层
    for(unsigned long long l = 0; l < p->n_layers; l++) {
        // 2.1 Attention 前的 RMSNorm
        rmsnorm(s->xb, x, w->rms_att_weight + l*dim, dim);
        
        // 2.2 计算 QKV
        matmul(s->q, s->xb, w->wq + l*dim*dim, dim, dim);
        matmul(s->k, s->xb, w->wk + l*dim*kv_dim, dim, kv_dim);
        matmul(s->v, s->xb, w->wv + l*dim*kv_dim, dim, kv_dim);
        
        // 2.3 应用 RoPE
        // ... (旋转 Q 和 K)
        
        // 2.4 多头注意力
        for (h = 0; h < p->n_heads; h++) {
            // 计算注意力分数
            for (int t = 0; t <= pos; t++) {
                score = dot_product(q, k[t]) / sqrt(head_size);
                att[t] = score;
            }
            // Softmax
            softmax(att, pos + 1);
            // 加权求和
            for (int t = 0; t <= pos; t++) {
                xb += att[t] * v[t];
            }
        }
        
        // 2.5 输出投影
        matmul(s->xb2, s->xb, w->wo + l*dim*dim, dim, dim);
        
        // 2.6 残差连接
        for (int i = 0; i < dim; i++) {
            x[i] += s->xb2[i];
        }
        
        // 2.7 FFN 前的 RMSNorm
        rmsnorm(s->xb, x, w->rms_ffn_weight + l*dim, dim);
        
        // 2.8 FFN（SwiGLU）
        matmul(s->hb, s->xb, w->w1 + l*dim*hidden_dim, dim, hidden_dim);
        matmul(s->hb2, s->xb, w->w3 + l*dim*hidden_dim, dim, hidden_dim);
        // SwiGLU 激活
        for (int i = 0; i < hidden_dim; i++) {
            s->hb[i] = silu(s->hb[i]) * s->hb2[i];
        }
        matmul(s->xb, s->hb, w->w2 + l*dim*hidden_dim, hidden_dim, dim);
        
        // 2.9 残差连接
        for (int i = 0; i < dim; i++) {
            x[i] += s->xb[i];
        }
    }
    
    // 3. 最终归一化
    rmsnorm(x, x, w->rms_final_weight, dim);
    
    // 4. 输出分类器
    matmul(s->logits, x, w->wcls, p->dim, p->vocab_size);
    
    return s->logits;
}

void encode(Tokenizer* t, char *text, int8_t bos, int8_t eos, 
            int *tokens, int *n_tokens) {
    // 1. 添加 BOS token（如果需要）
    if (bos) tokens[(*n_tokens)++] = 1;
    
    // 2. 添加虚拟前置空格
    if (text[0] != '\0') {
        int dummy_prefix = str_lookup(" ", t->sorted_vocab, t->vocab_size);
        tokens[(*n_tokens)++] = dummy_prefix;
    }
    
    // 3. UTF-8 初步拆解（原子化）
    for (char *c = text; *c != '\0'; c++) {
        // 判断是否为起始字节
        if ((*c & 0xC0) != 0x80) {
            str_len = 0;
        }
        str_buffer[str_len++] = *c;
        str_buffer[str_len] = '\0';
        
        // 继续读取后续字节
        if ((*(c+1) & 0xC0) == 0x80 && str_len < 4) {
            continue;
        }
        
        // 读完一个完整字符，查表
        int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size);
        if (id != -1) {
            tokens[(*n_tokens)++] = id;
        } else {
            // 字节回退：如果找不到，就按字节编码
            for (int i=0; i < str_len; i++) {
                tokens[(*n_tokens)++] = (unsigned char)str_buffer[i] + 3;
            }
        }
        str_len = 0;
    }
    
    // 4. BPE 合并循环
    while (1) {
        float best_score = -1e10;
        int best_id = -1;
        int best_idx = -1;
        
        // 扫描所有相邻 token 对
        for (int i=0; i < (*n_tokens-1); i++) {
            sprintf(str_buffer, "%s%s", 
                    t->vocab[tokens[i]], t->vocab[tokens[i+1]]);
            int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size);
            if (id != -1 && t->vocab_scores[id] > best_score) {
                best_score = t->vocab_scores[id];
                best_id = id;
                best_idx = i;
            }
        }
        
        if (best_idx == -1) break; // 没有可合并的了
        
        // 合并最佳对
        tokens[best_idx] = best_id;
        for (int i = best_idx+1; i < (*n_tokens-1); i++) {
            tokens[i] = tokens[i+1];
        }
        (*n_tokens)--;
    }
    
    // 5. 添加 EOS token（如果需要）
    if (eos) tokens[(*n_tokens)++] = 2;
}

char* decode(Tokenizer* t, int prev_token, int token) {
    char *piece = t->vocab[token];
    // BOS 后的空格处理
    if (prev_token == 1 && piece[0] == ' ') { piece++; }
    // 处理字节 token（如 <0x01>）
    unsigned char byte_val;
    if (sscanf(piece, "<0x%02hhX>", &byte_val) == 1) {
        piece = (char*)t->byte_pieces + byte_val * 2;
    }
    return piece;
}

int sample_argmax(float* probabilities, int n) {
    int max_i = 0;
    float max_p = probabilities[0];
    for (int i = 1; i < n; i++) {
        if (probabilities[i] > max_p) {
            max_i = i;
            max_p = probabilities[i];
        }
    }
    return max_i;
}

int sample_mult(float* probabilities, int n, float coin) {
    float cdf = 0.0f;
    for (int i = 0; i < n; i++) {
        cdf += probabilities[i];
        if (coin < cdf) {
            return i;
        }
    }
    return n - 1;
}

int sample_topp(float* probabilities, int n, float topp, 
                ProbIndex* probindex, float coin) {
    // 1. 预筛选：剔除极低概率的词
    const float cutoff = (1.0f - topp) / (n - 1);
    int n0 = 0;
    for (int i = 0; i < n; i++) {
        if (probabilities[i] >= cutoff) {
            probindex[n0].index = i;
            probindex[n0].prob = probabilities[i];
            n0++;
        }
    }
    
    // 2. 排序：按概率从高到低
    qsort(probindex, n0, sizeof(ProbIndex), compare);
    
    // 3. 截断：累加到 topp 为止
    float cumulative_prob = 0.0f;
    int last_idx = n0 - 1;
    for (int i = 0; i < n0; i++) {
        cumulative_prob += probindex[i].prob;
        if (cumulative_prob > topp) {
            last_idx = i;
            break;
        }
    }
    
    // 4. 在截断后的列表中采样
    float r = coin * cumulative_prob;
    float cdf = 0.0f;
    for (int i = 0; i <= last_idx; i++) {
        cdf += probindex[i].prob;
        if (r < cdf) {
            return probindex[i].index;
        }
    }
    return probindex[last_idx].index;
}

// 应用温度
for (int q=0; q<sampler->vocab_size; q++) { 
    logits[q] /= sampler->temperature; 
}
// 再应用 softmax
softmax(logits, sampler->vocab_size);

int sample(Sampler* sampler, float* logits) {
    int next;
    if (sampler->temperature == 0.0f) {
        // 贪心采样
        next = sample_argmax(logits, sampler->vocab_size);
    } else {
        // 应用温度
        for (int q=0; q<sampler->vocab_size; q++) { 
            logits[q] /= sampler->temperature; 
        }
        // Softmax
        softmax(logits, sampler->vocab_size);
        // 生成随机数
        float coin = random_f32(&sampler->rng_state);
        // 选择采样方式
        if (sampler->topp <= 0 || sampler->topp >= 1) {
            next = sample_mult(logits, sampler->vocab_size, coin);
        } else {
            next = sample_topp(logits, sampler->vocab_size, 
                              sampler->topp, sampler->probindex, coin);
        }
    }
    return next;
}

void generate(Transformer *transformer, Tokenizer *tokenizer, 
              Sampler *sampler, char *prompt, int steps) {
    // 1. 编码 prompt
    int num_prompt_tokens = 0;
    int* prompt_tokens = (int*)malloc((strlen(prompt)+3) * sizeof(int));
    encode(tokenizer, prompt, 1, 0, prompt_tokens, &num_prompt_tokens);
    
    // 2. 初始化
    int token = prompt_tokens[0];
    int pos = 0;
    long start = 0;
    
    // 3. 主循环
    while (pos < steps) {
        // 3.1 前向传播
        float* logits = forward(transformer, token, pos);
        
        // 3.2 决定下一个 token
        int next;
        if (pos < num_prompt_tokens - 1) {
            // Prompt 填充期：强制使用下一个输入词
            next = prompt_tokens[pos + 1];
        } else {
            // 自主生成期：采样选词
            next = sample(sampler, logits);
        }
        pos++;
        
        // 3.3 退出条件
        if (next == 1) { break; }
        
        // 3.4 解码并显示
        char* piece = decode(tokenizer, token, next);
        safe_printf(piece);
        fflush(stdout); // 流式输出
        
        // 3.5 更新状态
        token = next;
        
        // 3.6 启动计时器
        if (start == 0) { start = time_in_ms(); }
    }
    
    // 4. 报告性能
    if (pos > 1) {
        long end = time_in_ms();
        fprintf(stderr, "achieved tok/s: %f\n", 
                (pos-1) / (double)(end-start)*1000);
    }
    
    free(prompt_tokens);
}

void chat(Transformer *transformer, Tokenizer *tokenizer, Sampler *sampler,
          char *cli_user_prompt, char *cli_system_prompt, int steps) {
    char system_prompt[512];
    char user_prompt[512];
    char rendered_prompt[1152];
    int* prompt_tokens = (int*)malloc(1152 * sizeof(int));
    
    int8_t user_turn = 1; // 1=用户输入，0=模型回复
    int pos = 0;
    
    while (pos < steps) {
        if (user_turn) {
            // 1. 获取 system prompt（仅第一次）
            if (pos == 0) {
                if (cli_system_prompt == NULL) {
                    read_stdin("Enter system prompt (optional): ", 
                              system_prompt, sizeof(system_prompt));
                } else {
                    strcpy(system_prompt, cli_system_prompt);
                }
            }
            
            // 2. 获取 user prompt
            if (pos == 0 && cli_user_prompt != NULL) {
                strcpy(user_prompt, cli_user_prompt);
            } else {
                read_stdin("User: ", user_prompt, sizeof(user_prompt));
            }
            
            // 3. 拼接成 Llama 2 Chat 格式
            if (pos == 0 && system_prompt[0] != '\0') {
                sprintf(rendered_prompt, 
                       "[INST] <<SYS>>\n%s\n<</SYS>>\n\n%s [/INST]",
                       system_prompt, user_prompt);
            } else {
                sprintf(rendered_prompt, "[INST] %s [/INST]", user_prompt);
            }
            
            // 4. 编码
            encode(tokenizer, rendered_prompt, 1, 0, 
                  prompt_tokens, &num_prompt_tokens);
            user_idx = 0;
            user_turn = 0;
            printf("Assistant: ");
        }
        
        // 5. 确定输入 token
        if (user_idx < num_prompt_tokens) {
            token = prompt_tokens[user_idx++]; // 填充 KV Cache
        } else {
            token = next;
        }
        
        // 6. EOS (=2) 结束 Assistant 回合
        if (token == 2) { user_turn = 1; }
        
        // 7. 前向传播 + 采样
        float* logits = forward(transformer, token, pos);
        next = sample(sampler, logits);
        pos++;
        
        // 8. 显示 Assistant 的回复
        if (user_idx >= num_prompt_tokens && next != 2) {
            char* piece = decode(tokenizer, token, next);
            safe_printf(piece);
            fflush(stdout);
        }
        if (next == 2) { printf("\n"); }
    }
    
    free(prompt_tokens);
}

*data = mmap(NULL, *file_size, PROT_READ, MAP_PRIVATE, *fd, 0);

s->x = calloc(p->dim, sizeof(float));