llm/algorithms.py

配套源码，运行方法见同目录 README。返回实验总览。

"""Small CPU reference algorithms for A2-A5 concepts, not GPU benchmarks."""
import hashlib
import re
import numpy as np


def online_attention(q, k, v, block_size=8):
    """One head, shapes (T,D), exact causal attention without a T*T output buffer.

    Iterating per query is for clarity. A real FlashAttention kernel tiles query
    rows too, runs on a GPU, and fuses the work within a Triton/CUDA program.
    """
    if block_size < 1 or q.shape != k.shape or v.shape[0] != q.shape[0]:
        raise ValueError("invalid attention shapes or block_size")
    out = np.zeros((len(q), v.shape[-1]), dtype=np.result_type(q, k, v))
    for i, query in enumerate(q):
        maximum, denominator = -np.inf, 0.0
        numerator = np.zeros(v.shape[-1])
        for j in range(0, i + 1, block_size):
            end = min(j + block_size, i + 1)
            scores = query @ k[j:end].T / np.sqrt(q.shape[-1])
            new_max = max(maximum, float(scores.max()))
            correction = np.exp(maximum - new_max)
            weights = np.exp(scores - new_max)
            denominator = correction * denominator + weights.sum()
            numerator = correction * numerator + weights @ v[j:end]
            maximum = new_max
        out[i] = numerator / denominator
    return out


def normalize_document(text):
    return re.sub(r"\s+", " ", text).strip()


def exact_deduplicate(documents):
    seen, kept = set(), []
    for text in documents:
        clean = normalize_document(text)
        digest = hashlib.sha256(clean.encode("utf-8")).digest()
        if clean and digest not in seen:
            seen.add(digest)
            kept.append(clean)
    return kept


def shingles(text, n=3):
    words = normalize_document(text).split()
    return {tuple(words[i:i + n]) for i in range(max(0, len(words) - n + 1))}


def jaccard(a, b):
    union = a | b
    return len(a & b) / len(union) if union else 1.0


def fit_power_law(x, loss, floor):
    """Fit loss = floor + a*x**(-alpha) with a KNOWN floor.

    Unknown floors require nonlinear fitting; this deliberately refuses
    nonpositive residuals rather than hiding them with clipping.
    """
    x, residual = np.asarray(x), np.asarray(loss) - floor
    if np.any(x <= 0) or np.any(residual <= 0):
        raise ValueError("x and loss-floor must be positive")
    slope, intercept = np.polyfit(np.log(x), np.log(residual), 1)
    return float(np.exp(intercept)), float(-slope)


def group_advantages(rewards, eps=1e-8):
    rewards = np.asarray(rewards, dtype=np.float64)
    return (rewards - rewards.mean(axis=-1, keepdims=True)) / (
        rewards.std(axis=-1, keepdims=True) + eps)


def masked_sft_loss(logits, targets, response_mask):
    """Token-normalized response-only NLL plus its logits gradient."""
    from model import softmax
    mask = np.asarray(response_mask, dtype=np.float64)
    if targets.shape != logits.shape[:-1] or mask.shape != targets.shape:
        raise ValueError("targets and response_mask must match logits positions")
    if mask.sum() <= 0 or np.any((mask != 0) & (mask != 1)):
        raise ValueError("response_mask must be binary with at least one token")
    shifted = logits - logits.max(axis=-1, keepdims=True)
    nll = np.log(np.exp(shifted).sum(axis=-1)) - np.take_along_axis(
        shifted, targets[..., None], axis=-1)[..., 0]
    grad = softmax(logits)
    flat = grad.reshape(-1, grad.shape[-1])
    flat[np.arange(targets.size), targets.reshape(-1)] -= 1
    grad *= mask[..., None] / mask.sum()
    return float((nll * mask).sum() / mask.sum()), grad


def clipped_policy_objective(new_logprob, old_logprob, advantage, mask, clip=.2):
    """Positive PPO-style token objective; maximize it or minimize its negative.

    Treat old_logprob and advantage as constants in differentiation.
    Uses a token-normalized reduction, not per-sequence normalization.
    """
    mask = np.asarray(mask, dtype=float)
    if mask.sum() <= 0:
        raise ValueError("need at least one response token")
    ratio = np.exp(np.asarray(new_logprob) - np.asarray(old_logprob))
    term = np.minimum(ratio * advantage, np.clip(ratio, 1 - clip, 1 + clip) * advantage)
    return float((term * mask).sum() / mask.sum())


def dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=.1):
    margin = beta * ((np.asarray(policy_chosen) - policy_rejected) -
                     (np.asarray(ref_chosen) - ref_rejected))
    return float(np.logaddexp(0, -margin).mean())


if __name__ == "__main__":
    rng = np.random.default_rng(4)
    q, k, v = (rng.normal(size=(11, 8)) for _ in range(3))
    from model import causal_attention
    dense = causal_attention(q[None, None], k[None, None], v[None, None])[0][0, 0]
    print("online attention max absolute error:", np.max(np.abs(dense - online_attention(q, k, v, 3))))
    x = np.array([1., 2., 4., 8., 16.])
    print("synthetic scaling recovery (a, alpha):", fit_power_law(x, 1.2 + 3 * x**(-.4), 1.2))
    print("group advantages:", group_advantages([[0, 0, 1, 1]]))
    print("DPO neutral loss:", dpo_loss(-5., -5., -5., -5.))