LLM 사전 학습 & 스케일링 법칙: Chinchilla, Flash Attention, MoE까지

import math

def chinchilla_optimal(compute_budget_flops: float):
    """
    Chinchilla 법칙에 따른 최적 N, D 계산
    compute_budget_flops: 총 FLOPs (예: 6 * N * D 근사)
    반환: (optimal_params, optimal_tokens)
    """
    # Chinchilla 계수 (Hoffmann et al. Table A3 기준)
    # C = 6 * N * D 근사에서
    # N_opt = (C / (6 * 20))^0.5, D_opt = 20 * N_opt
    N_opt = math.sqrt(compute_budget_flops / (6 * 20))
    D_opt = 20 * N_opt
    return N_opt, D_opt

# 예시: A100 1000개, 30일 학습
# A100: ~312 TFLOPS (bf16), 가동률 40%
flops_per_sec = 1000 * 312e12 * 0.4
duration_sec = 30 * 24 * 3600
total_flops = flops_per_sec * duration_sec

N_opt, D_opt = chinchilla_optimal(total_flops)
print(f"총 FLOPs: {total_flops:.2e}")
print(f"최적 파라미터 수: {N_opt/1e9:.1f}B")
print(f"최적 토큰 수: {D_opt/1e12:.1f}T")

from datasketch import MinHash, MinHashLSH
import re

def get_shingles(text: str, k: int = 5):
    """문자 k-shingle 집합 생성"""
    text = re.sub(r'\s+', ' ', text.lower())
    return {text[i:i+k] for i in range(len(text) - k + 1)}

def build_minhash(text: str, num_perm: int = 128) -> MinHash:
    """텍스트에서 MinHash 서명 생성"""
    m = MinHash(num_perm=num_perm)
    for shingle in get_shingles(text):
        m.update(shingle.encode('utf-8'))
    return m

def deduplicate_corpus(documents: list, threshold: float = 0.8):
    """
    MinHash LSH로 유사 문서 제거
    threshold: Jaccard 유사도 임계값
    """
    lsh = MinHashLSH(threshold=threshold, num_perm=128)
    unique_docs = []
    seen = set()

    for idx, doc in enumerate(documents):
        mh = build_minhash(doc)
        result = lsh.query(mh)

        if len(result) == 0:
            lsh.insert(f"doc_{idx}", mh)
            unique_docs.append(doc)
        # 유사 문서가 이미 존재하면 스킵

    print(f"원본: {len(documents)}개 → 중복 제거 후: {len(unique_docs)}개")
    return unique_docs

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import ByteLevel as ByteLevelProcessor

def train_bpe_tokenizer(
    corpus_files: list,
    vocab_size: int = 32000,
    save_path: str = "tokenizer.json"
):
    """BPE 토크나이저 학습 (GPT-2 스타일 Byte-level BPE)"""
    tokenizer = Tokenizer(BPE(unk_token=None))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
    tokenizer.post_processor = ByteLevelProcessor(trim_offsets=False)

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=2,
        special_tokens=["<pad>", "<s>", "</s>", "<unk>"],
        show_progress=True,
    )

    tokenizer.train(files=corpus_files, trainer=trainer)
    tokenizer.save(save_path)
    print(f"토크나이저 저장 완료: {save_path} (vocab_size={vocab_size})")
    return tokenizer

# 사용 예시
# train_bpe_tokenizer(["data/train.txt"], vocab_size=32000)

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoELayer(nn.Module):
    """간단한 Mixture of Experts 레이어"""
    def __init__(self, d_model: int, d_ff: int, num_experts: int = 8, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.SiLU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor):
        # x: (batch, seq_len, d_model)
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (B*T, D)

        # 라우터 계산
        router_logits = self.router(x_flat)  # (B*T, num_experts)
        router_probs = F.softmax(router_logits, dim=-1)

        # Top-K 전문가 선택
        topk_probs, topk_idx = router_probs.topk(self.top_k, dim=-1)
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)  # 재정규화

        # Load Balancing Loss (전문가 균형 유지)
        # 각 전문가가 균등하게 선택되도록
        expert_load = router_probs.mean(0)  # (num_experts,)
        load_balance_loss = self.num_experts * (expert_load * expert_load.mean()).sum()

        # 전문가 출력 계산
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_indices = topk_idx[:, k]
            expert_weights = topk_probs[:, k].unsqueeze(-1)
            for e_idx in range(self.num_experts):
                mask = (expert_indices == e_idx)
                if mask.any():
                    expert_out = self.experts[e_idx](x_flat[mask])
                    output[mask] += expert_weights[mask] * expert_out

        return output.view(B, T, D), load_balance_loss

import math

def cosine_lr_with_warmup(
    optimizer,
    step: int,
    warmup_steps: int,
    total_steps: int,
    max_lr: float,
    min_lr: float = 0.0,
):
    """Cosine Annealing with Linear Warmup"""
    if step < warmup_steps:
        lr = max_lr * step / warmup_steps
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        lr = min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

# 전형적인 사전 학습 설정
# warmup: 전체 스텝의 1~2%
# max_lr: 3e-4 (7B 모델 기준)
# min_lr: max_lr * 0.1

import torch

def train_step_with_stability(model, optimizer, batch, grad_clip: float = 1.0):
    """손실 스파이크 방지를 위한 안전한 학습 스텝"""
    optimizer.zero_grad()
    loss = model(batch)
    loss.backward()

    # Gradient Norm 계산
    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(), max_norm=grad_clip
    )

    # 비정상적으로 큰 그래디언트 감지
    if grad_norm > 100 * grad_clip:
        print(f"경고: 비정상 grad_norm={grad_norm:.2f}, 스텝 스킵")
        optimizer.zero_grad()
        return None, grad_norm

    optimizer.step()
    return loss.item(), grad_norm

# Flash Attention 2 사용 예시
# pip install flash-attn --no-build-isolation

import torch
from flash_attn import flash_attn_func
from flash_attn.bert_padding import unpad_input, pad_input

def flash_attention_forward(
    q: torch.Tensor,  # (batch, seqlen, nheads, headdim)
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = True,
    softmax_scale: float = None,
):
    """
    Flash Attention 2 순전파
    causal=True: 자동 회귀 언어 모델용 인과적 마스크 적용
    """
    if softmax_scale is None:
        softmax_scale = q.shape[-1] ** (-0.5)

    out = flash_attn_func(
        q, k, v,
        dropout_p=0.0,
        softmax_scale=softmax_scale,
        causal=causal,
    )
    return out  # (batch, seqlen, nheads, headdim)

import torch
import math

def compute_perplexity(model, dataloader, device: str = "cuda"):
    """검증 세트 퍼플렉시티 계산"""
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids=input_ids, labels=labels)
            # CrossEntropy loss는 이미 토큰당 평균
            batch_tokens = (labels != -100).sum().item()
            total_loss += outputs.loss.item() * batch_tokens
            total_tokens += batch_tokens

    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    return perplexity

# 전형적인 사전 학습 진행 기준
# GPT-2 수준: PPL ~20 (WebText)
# 좋은 7B 모델: PPL ~6~8 (검증 세트 기준)

# lm-evaluation-harness 설치 및 실행
pip install lm-eval

# Llama 3 8B 평가 예시
lm_eval --model hf \
    --model_args pretrained=meta-llama/Meta-Llama-3-8B \
    --tasks mmlu,hellaswag,arc_challenge,winogrande \
    --device cuda:0 \
    --batch_size 8 \
    --output_path results/llama3-8b/

# 특정 체크포인트 평가
lm_eval --model hf \
    --model_args pretrained=./checkpoints/step-50000 \
    --tasks hellaswag \
    --num_fewshot 10 \
    --device cuda:0