LLM大厂手撕真题(字节/阿里/百度/美团)

import numpy as npclass LinearReg:    def __init__(self, dim) -> None:        pass    def fit(self, x, y):        """        x: (B, dim)        y: (B, 1)        """        pass    def predict(self, x):        passif __name__ == '__main__':    batch_size, dim = 16, 2    x = np.random.uniform(-5, 5, size=(batch_size, dim))    w = np.random.uniform(-5, 5, size=(dim, 1))    y = x @ w + 10    lr = LinearReg(dim=1)    lr.fit(x, y)    y_hat = lr.predict(x)    assert np.allclose(y, y_hat)

（6）

手撕 MHA、MQA 和 GQA：

import torchimport torch.nn as nnclass MultiHeadAttention(nn.Module):    def __init__(self, d_model, num_heads):        r"""        d_model: Total dimension of the model.        num_heads:        - Number of parallel attention heads.         - Note that embed_dim will be split across num_heads (i.e. each head will have dimension embed_dim // num_heads).        """        super().__init__()    def forward(self, query, key, value):        """        query: (B, L, D)        key: (B, L, D)        value: (B, L, D)        out: (B, L, D)        """        return outif __name__ == '__main__':    B, L, D = 16, 32, 64    H = 4    mha = MultiHeadAttention(D, H)    q = torch.rand(B, L, D)    k = torch.rand(B, L, D)    v = torch.rand(B, L, D)    y = mha(q, k, v)

class RMSNorm(nn.Module):    def __init__(self, dim, eps=1e-6) -> None:        super().__init__()    def _norm(self, x):        return pass    def forward(self, x):        return passclass LayerNorm(nn.Module):    def __init__(self, dim, eps=1e-6) -> None:        super().__init__()    def _norm(self, x):    def forward(self, x):

from email import policyimport torchfrom torch.nn.modules import lossdef gspo_loss(    advantage,  # (B, 1)    policy_logprob, # (B, L)    old_policy_logprob, # (B, L)    clip_low,     clip_high    ):    return def grpo_loss(    advantage,  # (B, 1)    policy_logprob, # (B, L)    old_policy_logprob, # (B, L)    clip_low,     clip_high    ):    returnif __name__ == '__main__':    B, L, D = 16, 32, 64    advantage = torch.rand(B, 1)    policy_logprob = torch.rand(B, L)    old_policy_logprob = torch.rand(B, L)    loss = grpo_loss(advantage, policy_logprob, old_policy_logprob, 0.2, 0.28)    print(loss)    loss = gspo_loss(advantage, policy_logprob, old_policy_logprob, 0.2, 0.28)    print(loss)

import torchdef compute_advantage(    rewards,    values,    gamma,    lamda,):    advantages = []    batch_size, response_length = rewards.shape    for t in reversed(range(response_length)):          # TODO:    advantages = advantages[::-1]    return torch.stack(advantages, dim=1)  # 建议返回 tensor 而不是 list

class Solution:    def lengthOfLIS(self, nums: List[int]) -> int:        tail = []        tail.append(-1e9)        for i in range(0, len(nums)):            left, right = 0, len(tail) - 1            if nums[i] > tail[right]:                tail.append(nums[i])                continue            while left <= right:                if nums[i] > tail[left]:                    left += 1                else:                     right -= 1            tail[left] = nums[i]        return len(tail) - 1

import numpy as npclass LinearReg:    def __init__(self, dim) -> None:        self.W = None    def fit(self, x, y):        """        x: (B, dim)        y: (B, 1)        """        x = np.concat((x,[[1]] * len(x)), axis=-1)        self.W = np.linalg.inv(x.T @ x) @ x.T @ y    def predict(self, x):        return np.concat((x,[[1]] * len(x)), axis=-1) @ self.Wif __name__ == '__main__':    batch_size, dim = 16, 2    x = np.random.uniform(-5, 5, size=(batch_size, dim))    w = np.random.uniform(-5, 5, size=(dim, 1))    y = x @ w + 10    lr = LinearReg(dim=1)    lr.fit(x, y)    y_hat = lr.predict(x)    assert np.allclose(y, y_hat)

import numpy as npclass LinearRegGD:    def __init__(self, dim, lr=0.01, epochs=1000, tol=1e-6) -> None:        self.dim = dim        self.lr = lr        self.epochs = epochs        self.tol = tol        self.W = None    def fit(self, x, y):        """        x: (B, dim)        y: (B, 1)        """        B = x.shape[0]        # 添加偏置列        x_bias = np.hstack([x, np.ones((B, 1))])        # 初始化权重（包含偏置）        self.W = np.random.randn(self.dim + 1, 1) * 0.01        for i in range(self.epochs):            # 预测            y_pred = x_bias @ self.W            # 梯度：X^T (XW - y) / B  (MSE的梯度)            grad = x_bias.T @ (y_pred - y) / B            # 更新            new_W = self.W - self.lr * grad            # 收敛检查            if np.linalg.norm(new_W - self.W) < self.tol:                print(f"Converged at epoch {i}")                break            self.W = new_W    def predict(self, x):        B = x.shape[0]        x_bias = np.hstack([x, np.ones((B, 1))])        return x_bias @ self.Wif __name__ == '__main__':    np.random.seed(42)    batch_size, dim = 16, 2    x = np.random.uniform(-5, 5, size=(batch_size, dim))    w_true = np.random.uniform(-5, 5, size=(dim, 1))    y = x @ w_true + 10    lr = LinearRegGD(dim=dim, lr=0.1, epochs=5000)    lr.fit(x, y)    y_hat = lr.predict(x)    print(f"True weights:\n{w_true}\nTrue bias: 10")    print(f"Learned weights:\n{lr.W[:-1]}\nLearned bias: {lr.W[-1, 0]:.4f}")    # 梯度下降不会完全等于闭式解，但应该很接近    assert np.allclose(y, y_hat, atol=1e-3), f"Max error: {np.max(np.abs(y - y_hat))}"    print("✅ 梯度下降测试通过")

grad = x_bias.T @ (y_pred - y) / B

import torchimport torch.nn as nnclass MultiHeadAttention(nn.Module):    def __init__(self, d_model, num_heads):        r"""        d_model: Total dimension of the model.        num_heads:        - Number of parallel attention heads.         - Note that embed_dim will be split across num_heads (i.e. each head will have dimension embed_dim // num_heads).        """        super().__init__()        self.W_q = nn.Linear(d_model, d_model)        self.W_k = nn.Linear(d_model, d_model)        self.W_v = nn.Linear(d_model, d_model)        self.W_o = nn.Linear(d_model, d_model)        self.num_heads = num_heads    def forward(self, query, key, value):        """        query: (B, L, D)        key: (B, L, D)        value: (B, L, D)        """        B, L, D = query.shape        query = self.W_q(query).view(B, L, self.num_heads, -1).transpose(1, 2)        key = self.W_k(key).view(B, L, self.num_heads, -1).transpose(1, 2).transpose(2, 3)        value = self.W_v(value).view(B, L, self.num_heads, -1).transpose(1, 2)        d = query.size(-1)        dot_product = torch.softmax(query @ key / torch.sqrt(torch.tensor(d)), dim=-1)        out = (dot_product @ value).transpose(1, 2).reshape(B, L, -1)        out = self.W_o(out)        return outclass MultiQueryAttention(nn.Module):    """    不同的head共享一份key和value，但是query还是独立的    """    def __init__(self, d_model, num_heads):        super().__init__()        assert d_model % num_heads == 0        self.num_heads = num_heads        self.head_dim = d_model // num_heads        self.W_q = nn.Linear(d_model, d_model)        self.W_k = nn.Linear(d_model, self.head_dim)        self.W_v = nn.Linear(d_model, self.head_dim)        self.W_o = nn.Linear(d_model, d_model)    def forward(self, query, key, value):        B, L, D = query.shape        query = self.W_q(query).view(B, L, self.num_heads, -1).transpose(1, 2) # (B, num_heads, L, head_dim)        key = self.W_k(key).view(B, L, 1, -1).transpose(1, 2).transpose(2, 3)        value = self.W_v(value).view(B, L, 1, -1).transpose(1, 2)        d = query.size(-1)        dot_product = torch.softmax(query @ key / torch.sqrt(torch.tensor(d)), dim=-1)        out = (dot_product @ value).transpose(1, 2).reshape(B, L, -1)        out = self.W_o(out)        return outclass GroupedQueryAttention(nn.Module):    """    GQA: 同一组内的多个 head 共享同一份 K 和 V，每个 head 仍有独立的 Q。    num_groups 个 group，每组 group_size = num_heads // num_groups 个 head。    """    def __init__(self, d_model, num_heads, num_groups):        super().__init__()        assert d_model % num_heads == 0        self.num_heads = num_heads        assert num_heads % num_groups == 0        self.num_groups = num_groups        self.group_size = num_heads // num_groups        self.head_dim = d_model // num_heads        self.W_q = nn.Linear(d_model, d_model)        self.W_k = nn.Linear(d_model, self.num_groups * self.head_dim)        self.W_v = nn.Linear(d_model, self.num_groups * self.head_dim)        self.W_o = nn.Linear(d_model, d_model)    def forward(self, query, key, value):        B, L, D = query.shape        H = self.num_heads        G = self.num_groups        d = self.head_dim        group_size = self.group_size        # ---- Q ----        # (B, L, D) → (B, H, L, d)        q = self.W_q(query).view(B, L, H, d).transpose(1, 2)        # ---- K ----        # (B, L, G*d) → (B, G, L, d)        k = self.W_k(key).view(B, L, G, d).transpose(1, 2)        # ---- V ----        # (B, L, G*d) → (B, G, L, d)        v = self.W_v(value).view(B, L, G, d).transpose(1, 2)        # ---- broadcast KV to heads ----        # (B,G,L,d) → (B,G,group_size,L,d)        k = k.unsqueeze(2).expand(B, G, group_size, L, d)        v = v.unsqueeze(2).expand(B, G, group_size, L, d)        # (B,G,group_size,L,d) → (B,H,L,d)        k = k.reshape(B, H, L, d)        v = v.reshape(B, H, L, d)        # ---- attention ----        scale = 1.0 / math.sqrt(d)        attn = torch.matmul(q, k.transpose(-2, -1)) * scale        attn = F.softmax(attn, dim=-1)        out = torch.matmul(attn, v)        # ---- merge heads ----        out = out.transpose(1, 2).reshape(B, L, D)        return self.W_o(out)        return outif __name__ == '__main__':    B, L, D = 16, 32, 64    H = 4    mha = MultiHeadAttention(D, H)    q = torch.rand(B, L, D)    k = torch.rand(B, L, D)    v = torch.rand(B, L, D)    y = mha(q, k, v)

class LayerNorm(nn.Module):    def __init__(self, dim, eps=1e-6) -> None:        super().__init__()        self.w = nn.Parameter(torch.ones(dim))        self.b = nn.Parameter(torch.zeros(dim))        self.eps = eps    def _norm(self, x):        return (x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(torch.var(x, dim=-1, keepdim=True, unbiased=False) + self.eps)    def forward(self, x):        return self.w * self._norm(x) + self.bclass RMSNorm(nn.Module):    def __init__(self, dim, eps=1e-6) -> None:        super().__init__()        self.w = nn.Parameter(torch.ones(dim))        self.eps = eps    def _norm(self, x):        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)    def forward(self, x):        return self.w * self._norm(x)

import torchimport torch.nn as nnimport mathimport torch.nn.functional as Fclass Expert(nn.Module):    def __init__(self, n_embd, dropout=0.1):        super().__init__()        self.net = nn.Sequential(            nn.Linear(n_embd, 4 * n_embd),            nn.ReLU(),            nn.Linear(4 * n_embd, n_embd),            nn.Dropout(dropout),        )    def forward(self, x):        return self.net(x)class TopkRouter(nn.Module):    def __init__(self, n_embed, num_experts, top_k):        super(TopkRouter, self).__init__()        self.top_k = top_k        self.linear =nn.Linear(n_embed, num_experts)    def forward(self, mh_output):        """        mh_output: (B, L, D)        return: (B, L, num_experts)        """        logits = self.linear(mh_output)        # 获取前K大的值和索引，沿列。        top_k_logits, indices = logits.topk(self.top_k, dim=-1)        # 创建一个形状和logits相同全'-inf'矩阵        zeros = torch.full_like(logits, float('-inf'))        # 按照索引和值填充上述zeros矩阵        sparse_logits = zeros.scatter(-1, indices, top_k_logits)        # 对其进行softmax，未被填充的位置会为0        router_output = F.softmax(sparse_logits, dim=-1)        return router_output, indicesclass SparseMoE(nn.Module):    def __init__(self, n_embd, n_experts, top_k, dropout=0.1):        super().__init__()        self.top_k = top_k        self.n_experts = n_experts        self.experts = nn.ModuleList([Expert(n_embd, dropout) for _ in range(n_experts)])        self.router = TopkRouter(n_embd, n_experts, top_k)    def forward(self, x):        """        x: (B, L, D)        return: (B, L, D)        """        B, L, D = x.shape        router_output, indices = self.router(x) # (B, L, num_experts), (B, L, top_k)        x = torch.flatten(x, end_dim=1)        router_output = torch.flatten(router_output, end_dim=1) # (F, num_experts)        router_mask = router_output > 0 # (B, L, num_experts) -> (F, num_experts)        output = torch.zeros_like(x)        for i in range(self.n_experts):            current_expert_router = router_mask[:, i] # 当前expert需要处理的那些token的索引            current_expert_input = x[current_expert_router]            current_expert_output = self.experts[i](current_expert_input)            output[current_expert_router] += router_output[current_expert_router, i].unsqueeze(-1) * current_expert_output        return outputif __name__ == '__main__':    B, L, D = 8, 32, 64    num_experts, top_k = 4, 2    x = torch.rand(B, L, D)    smoe = SparseMoE(n_embd=D, n_experts=num_experts, top_k=top_k)    y = smoe.forward(x)    print(y.shape)

from email import policyimport torchfrom torch.nn.modules import lossdef gspo_loss(    advantage,  # (B, 1)    policy_logprob, # (B, L)    old_policy_logprob, # (B, L)    clip_low,     clip_high    ):    # ratio = torch.exp(1/ L * \sum log ratio)    negative_approx_kl = policy_logprob - old_policy_logprob    logratio = policy_logprob - policy_logprob.detach() + negative_approx_kl.detach()    ratio = torch.exp(logratio.mean(dim=-1))    loss2 = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantage    loss1 = ratio * advantage     loss = torch.min(loss1, loss2)    return -loss.mean()def grpo_loss(    advantage,  # (B, 1)    policy_logprob, # (B, L)    old_policy_logprob, # (B, L)    clip_low,     clip_high    ):    ratio = torch.exp(policy_logprob - old_policy_logprob)    # loss = - min(ratio * adv, clamp(ratio, 1-clip_low, 1+clip_high) * adv)    loss2 = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) * advantage    loss1 = ratio * advantage     loss = torch.min(loss1, loss2)    return -loss.mean()if __name__ == '__main__':    B, L, D = 16, 32, 64    advantage = torch.rand(B, 1)    policy_logprob = torch.rand(B, L)    old_policy_logprob = torch.rand(B, L)    loss = grpo_loss(advantage, policy_logprob, old_policy_logprob, 0.2, 0.28)    print(loss)    loss = gspo_loss(advantage, policy_logprob, old_policy_logprob, 0.2, 0.28)    print(loss)

import torchdef compute_advantage(    rewards,    values,    gamma,    lamda,):    last_gae = 0    advantages = []    batch_size, response_length = rewards.shape    for t in reversed(range(response_length)):        next_value = torch.zeros(batch_size, 1) if t == response_length - 1 else values[:, t + 1]        # 修正：r + gamma * V(s_{t+1}) - V(s_t)        delta = rewards[:, t] + gamma * next_value - values[:, t]        gae = delta + gamma * lamda * last_gae        advantages.append(gae)        last_gae = gae  # 更新 last_gae 供下一次迭代使用    advantages = advantages[::-1]    return torch.stack(advantages, dim=1)  # 建议返回 tensor 而不是 list