[update] Several Python Models

Author: jamesnulliu

.vscode/extensions.json

@@ -5,7 +5,6 @@
         "josetr.cmake-language-support-vscode",
         "ms-vscode.cpptools",  // C/C++
         "ms-python.python",  // Python
-        "ms-python.black-formatter",  // Python formatter
         "njpwerner.autodocstring",  // Python docstring generator
     ]
 }

csrc/CMakeLists.txt

@@ -32,7 +32,7 @@ include(${PROJECT_SOURCE_DIR}/cmake/libraries/libtorch.cmake)
 find_package(cxxopts CONFIG REQUIRED)
 find_package(fmt CONFIG REQUIRED)
 find_package(spdlog CONFIG REQUIRED)
-find_package(proxy CONFIG REQUIRED)
+find_package(msft_proxy4 CONFIG REQUIRED)
 find_package(yaml-cpp CONFIG REQUIRED)
 enable_testing()
 find_package(GTest CONFIG REQUIRED)

csrc/cmake/compilers/cxx-compiler-configs.cmake

@@ -33,7 +33,7 @@ if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
     string(APPEND CMAKE_EXE_LINKER_FLAGS " /STACK:${STACK_SIZE}")
 # Clang ---------------------------------------------------------------------------------------------------------------
 elseif(CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
-    string(APPEND CMAKE_CXX_FLAGS " -fopenmp -Wall -Wextra -Werror")
+    string(APPEND CMAKE_CXX_FLAGS " -stdlib=libc++ -fopenmp -Wall -Wextra -Werror")
     if (WIN32)
         string(APPEND CMAKE_EXE_LINKER_FLAGS " -Wl,-stack,${STACK_SIZE}")
     else()

csrc/lib/CMakeLists.txt

@@ -25,7 +25,7 @@ target_link_libraries(
     Python::Python
     ${TORCH_LIBRARIES}
     fmt::fmt
-    msft_proxy
+    msft_proxy4::proxy
 )
 
 target_compile_definitions(${LIB_NAME} 

pmpp/models/attention.py

@@ -1,86 +1,10 @@
-import math
 from typing import Optional
 import torch
 from torch import nn
-from torch.nn import functional as F
-import numpy as np
+from .mha_kernels import MHAKernel
 
 
-def set_random_seed(
-    seed: int, rank: int = 0, force_deterministic: bool = False
-) -> None:
-    """
-    Set the random seed for numpy and torch.
-    """
-    np.random.seed(seed + rank)
-    torch.manual_seed(seed + rank)
-    if force_deterministic:
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-
-
-class MultiHeadSelfAttentionKernel(nn.Module):
-    def __init__(self, hidden_dim: int, num_heads: int):
-        super().__init__()
-
-        self.hidden_dim: int = hidden_dim
-        self.num_heads: int = num_heads
-        self.head_size: int = hidden_dim // num_heads
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ):
-        """
-        Calculates softmax(Q @ KT / sqrt(dk)) @ V .
-
-        Parameters
-        ----------
-        q : torch.Tensor; Shape: (q_len, hidden_dim)
-
-        k : torch.Tensor; Shape: (kv_len, hidden_dim)
-
-        v : torch.Tensor; Shape: (kv_len, hidden_dim)
-
-        mask: torch.Tensor; Shape: (q_len, kv_len), optional
-
-        Note
-        ----
-        When prefilling, q_len equals to seq_len (number of tokens in the input
-        seq);
-        When decoding, q_len equals to 1, refering to the newly generated
-        token. (Based on different sampling strategies, q_len could be larger
-        than 1.)
-        """
-
-        q_len, kv_len = q.size(0), k.size(0)
-        # q -> (num_heads, q_len, head_size)
-        q = q.view(q_len, self.num_heads, self.head_size).transpose(0, 1)
-        # k -> (num_heads, kv_len, head_size)
-        k = k.view(kv_len, self.num_heads, self.head_size).transpose(0, 1)
-        # v -> (num_heads, kv_len, head_size)
-        v = v.view(kv_len, self.num_heads, self.head_size).transpose(0, 1)
-        # scores -> (num_heads, q_len, kv_len)
-        scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_size**0.5)
-        scores = (
-            scores.masked_fill(mask == 0, float("-inf"))
-            if mask is not None
-            else scores
-        )
-        # scores -> (num_heads, q_len, kv_len)
-        attn_probs = F.softmax(scores, dim=-1)
-        # out -> (num_heads, q_len, head_size)
-        out = torch.matmul(attn_probs, v)
-        # out -> (q_len, num_heads, head_size)
-        out = out.transpose(0, 1).reshape(q_len, self.hidden_dim)
-
-        return out
-
-
-class MultiHeadSelfAttention(nn.Module):
+class MHA(nn.Module):
     def __init__(
         self,
         embed_dim: int,
@@ -97,7 +21,7 @@ def __init__(
         self.Wv = nn.Linear(embed_dim, hidden_dim)
         self.Wo = nn.Linear(hidden_dim, embed_dim)
 
-        self.attn_kernel = MultiHeadSelfAttentionKernel(hidden_dim, num_heads)
+        self.attn_kernel = MHAKernel(hidden_dim, num_heads)
 
     def forward(
         self,
@@ -135,7 +59,7 @@ def forward(
         v = self.Wv(seq)
 
         # k_cache -> (kv_len + seq_len, hidden_dim)
-        k = k if k_cache is None else torch.cat([k_cache, k.detach()], dim=0) 
+        k = k if k_cache is None else torch.cat([k_cache, k.detach()], dim=0)
         # v_cache -> (kv_len + seq_len, hidden_dim)
         v = v if v_cache is None else torch.cat([v_cache, v.detach()], dim=0)
 
@@ -148,9 +72,7 @@ def forward(
 class TransformerBlock(nn.Module):
     def __init__(self, embed_dim, num_heads, hidden_dim, mlp_dim, dropout=0.1):
         super().__init__()
-        self.attention = MultiHeadSelfAttention(
-            embed_dim, num_heads, hidden_dim
-        )
+        self.attention = MHA(embed_dim, num_heads, hidden_dim)
         self.norm1 = nn.RMSNorm(embed_dim)
         self.norm2 = nn.RMSNorm(embed_dim)
         self.mlp = nn.Sequential(
@@ -293,8 +215,6 @@ def forward(
 
 
 if __name__ == "__main__":
-    set_random_seed(114514)
-
     seq_len = 4
     vocab_size = 1024
     embed_dim = 128
@@ -347,5 +267,5 @@ def forward(
     for i in range(1, n_generate):
         probs = lm(token, is_prefilling=False)
         token = torch.argmax(probs[-1, :], dim=-1, keepdim=True)
-        print(f"The {i+1}th predicted token: {token}")
+        print(f"The {i + 1}th predicted token: {token}")
         print(f"|- Token Shape: {token.shape}")

pmpp/models/grpo.py

@@ -0,0 +1,86 @@
+"""
+PSEUDO-CODE FOR GRPO TRAINING (Group Relative Policy Optimization):
+    p_model   <- policy model (the one we're training)
+    ref_model <- reference policy (frozen; for KL regularization)
+    reward_model <- frozen model / callable that scores (prompt, resp)
+    # [NOTE] No value model / critic. GRPO uses group-relative baselines.
+
+    for i in range(num_iterations):
+        prompts <- [B, prompt_len]; Sampled from prompts dataset
+
+        # 1) Rollout: sample a GROUP of responses per prompt
+        # Let G = num_generations_per_prompt (a.k.a. group size)
+        resps <- [B, G, T]; Rollout from p_model on prompts (G samples each)
+        old_logp <- [B, G, T]; Logprobs under rollout policy (snapshot of p_model)
+        ref_logp <- [B, G, T]; Logprobs under ref_model for same tokens
+        action_mask <- [B, G, T]; 1 on valid action tokens, 0 on padding
+
+        # 2) Compute rewards (often sequence-level)
+        # reward_seq is per (prompt, resp) scalar score from reward_model
+        reward_seq <- [B, G]; reward_model(prompts, resps)
+
+        # Optional: add formatting penalties, stop-token penalties, etc.
+        # reward_seq <- reward_seq + extra_terms
+
+        # If you want token-shaped rewards, place seq reward on last valid token
+        reward_tok <- zeros([B, G, T])
+        last_idx <- last_valid_index(action_mask)        # [B, G]
+        reward_tok[b,g,last_idx[b,g]] += reward_seq[b,g] # scatter add
+
+        # 3) KL term (tokenwise, rollout policy vs reference)
+        kl_tok <- [B, G, T]; kl_tok = old_logp - ref_logp
+
+        # 4) Optional KL shaping of reward (same idea as PPO RLHF)
+        shaped_reward_tok <- [B, G, T];
+                             shaped_reward_tok = reward_tok - beta * kl_tok
+
+        # 5) Construct group-relative advantage (baseline from the GROUP)
+        # Most common GRPO: baseline is mean reward within the group (per
+        # prompt).
+        # Use shaped (sequence) reward or unshaped, depending on your design.
+        # Here: use shaped sequence reward = sum over tokens of
+        # shaped_reward_tok on valid tokens.
+        shaped_reward_seq <- [B, G]
+        shaped_reward_seq[b,g] = sum_t(shaped_reward_tok[b,g,t] * action_mask[b,g,t])
+
+        group_mean <- [B, 1]; group_mean[b,1] = mean_g(shaped_reward_seq[b,g])
+        group_std  <- [B, 1]; group_std[b,1]  = std_g(shaped_reward_seq[b,g]) + eps
+
+        adv <- [B, G]; adv = (shaped_reward_seq - group_mean) / group_std
+            # Alternatively: adv = shaped_reward_seq - group_mean (no normalization)
+
+        # Broadcast to tokens if doing token-level PPO-style objective
+        adv_tok <- [B, G, T]; adv_tok = adv[..., None] * action_mask
+
+        batch <- {prompts, resps, old_logp, ref_logp, adv_tok, action_mask}
+
+        # 6) Policy optimization (PPO-style clipped objective, but no value loss)
+        for epoch in range(num_epochs_per_rollout):
+            for minibatch in iterate_minibatches(batch, mb_size):
+                new_logp <- [mb, T]; logprobs from (updated) p_model on minibatch prompts+resps
+                old_logp_mb <- [mb, T]; from minibatch old_logp
+                adv_tok_mb <- [mb, T]; from minibatch adv_tok
+                mask_mb <- [mb, T]; from minibatch action_mask
+
+                # PPO ratio per token
+                log_ratio <- new_logp - old_logp_mb
+                ratio <- exp(log_ratio)
+
+                # GRPO policy gradient loss (clipped), averaged over valid tokens
+                unclipped <- -adv_tok_mb * ratio
+                clipped   <- -adv_tok_mb * clip(ratio, 1-eps_clip, 1+eps_clip)
+                pg_loss_tok <- max(unclipped, clipped)
+                pg_loss <- sum(pg_loss_tok * mask_mb) / sum(mask_mb)
+
+                # Optional entropy bonus
+                ent_bonus <- entropy_from_logits(...)   # or from new_logp if available
+                loss <- pg_loss - ent_coef * ent_bonus
+
+                optimizer.zero_grad()
+                loss.backward()
+                clip_grad_norm_(p_model.parameters(), max_grad_norm)  # common
+                optimizer.step()
+
+        # Optional: monitor approximate KL to ref, early stop if KL too large
+        # approx_kl = mean( (new_logp - ref_logp_mb) * mask_mb )
+"""

pmpp/models/kldiv.py

@@ -0,0 +1,46 @@
+import torch
+
+
+def kl_divergence_from_log_probs(
+    log_p: torch.Tensor,
+    log_q: torch.Tensor,
+    reduction: str = "batchmean",
+    eps: float = 0.0,
+) -> torch.Tensor:
+    """
+    KL(P || Q) = sum(P * (log P - log Q))
+    where log_p = log P, log_q = log Q along the last dim.
+
+    Args:
+        log_p: (..., K) log-probabilities for P (target)
+        log_q: (..., K) log-probabilities for Q (prediction)
+        reduction: "none" | "sum" | "mean" | "batchmean"
+        eps: optional additive smoothing in prob space; usually keep 0.0
+
+    Returns:
+        KL divergence with the chosen reduction.
+    """
+    if eps != 0.0:
+        # Smooth in prob space then re-normalize
+        p = log_p.exp()
+        q = log_q.exp()
+        p = p + eps
+        q = q + eps
+        p = p / p.sum(dim=-1, keepdim=True)
+        q = q / q.sum(dim=-1, keepdim=True)
+        log_p = (p.clamp_min(1e-30)).log()
+        log_q = (q.clamp_min(1e-30)).log()
+
+    p = log_p.exp()
+    kl_per_elem = p * (log_p - log_q)  # (..., K)
+    kl = kl_per_elem.sum(dim=-1)  # (...,)
+
+    if reduction == "none":
+        return kl  # Shape: (...)
+    if reduction == "sum":
+        return kl.sum()  # Scalar
+    if reduction == "mean":
+        return kl.mean()  # Scalar
+    if reduction == "batchmean":
+        return kl.sum() / kl.shape[0]  # Scalar
+    raise ValueError(f"Unknown reduction: {reduction}")

pmpp/models/mha_kernels.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional
+
+
+class MHAKernel(nn.Module):
+    def __init__(self, hidden_dim: int, num_heads: int):
+        super().__init__()
+
+        self.hidden_dim: int = hidden_dim
+        self.num_heads: int = num_heads
+        self.head_size: int = hidden_dim // num_heads
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Calculates softmax(Q @ KT / sqrt(dk)) @ V .
+
+        Parameters
+        ----------
+        q : torch.Tensor; Shape: (q_len, hidden_dim)
+
+        k : torch.Tensor; Shape: (kv_len, hidden_dim)
+
+        v : torch.Tensor; Shape: (kv_len, hidden_dim)
+
+        mask: torch.Tensor; Shape: (q_len, kv_len), optional
+
+        Note
+        ----
+        When prefilling, q_len equals to seq_len (number of tokens in the input
+        seq);
+        When decoding, q_len equals to 1, refering to the newly generated
+        token. (Based on different sampling strategies, q_len could be larger
+        than 1.)
+        """
+
+        q_len, kv_len = q.size(0), k.size(0)
+        # q -> (num_heads, q_len, head_size)
+        q = q.reshape(q_len, self.num_heads, self.head_size).transpose(0, 1)
+        # k -> (num_heads, kv_len, head_size)
+        k = k.reshape(kv_len, self.num_heads, self.head_size).transpose(0, 1)
+        # v -> (num_heads, kv_len, head_size)
+        v = v.reshape(kv_len, self.num_heads, self.head_size).transpose(0, 1)
+        # scores -> (num_heads, q_len, kv_len)
+        scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_size**0.5)
+        scores = (
+            scores.masked_fill(mask == 0, float("-inf"))
+            if mask is not None
+            else scores
+        )
+        # scores -> (num_heads, q_len, kv_len)
+        attn_probs = F.softmax(scores.to(torch.float32), dim=-1).type_as(
+            scores
+        )
+        # out -> (num_heads, q_len, head_size)
+        out = torch.matmul(attn_probs, v)
+        # out -> (q_len, num_heads, head_size)
+        out = out.transpose(0, 1).reshape(q_len, self.hidden_dim)
+
+        return out