ML-SYS 学习宝典：从 RLHF 系统到 SGLang 深入解析

# 伪代码，示意 SGLang 内部流程
# 1. 请求接收与 Tokenization
# 用户请求通过 FastAPI endpoint 进入
async def v1_chat_completions(request: ChatCompletionRequest):
    req_input = convert_to_generate_req_input(request)
    # TokenizerManager 进行分词和多模态数据处理
    tokenized = tokenizer_manager.tokenize(req_input)
    # 将请求放入 Scheduler 的等待队列
    scheduler.add_request(tokenized)

# 2. 调度与批处理 (Scheduler)
# Scheduler 事件循环
while True:
    # 从等待队列中根据优先级（如最长前缀）选取请求组成批次
    batch = scheduler.get_next_batch()
    if batch.prefill_tokens > 0:
        # Prefill 阶段：处理输入提示，填充 KV Cache
        run_prefill_batch(batch)
    # Decode 阶段：自回归生成 token
    next_token_ids = run_decode_batch(batch)
    # 处理生成结果，更新请求状态
    scheduler.process_batch_results(batch, next_token_ids)
    # 如果请求完成，将结果发送给 DetokenizerManager
    for req in batch.finished_requests:
        detokenizer_manager.send_output(req)

# 3. 结果返回
# DetokenizerManager 将 token ID 解码为文本，通过 HTTP 返回给用户。

# 代码片段位于 slime 相关 PR (#850)
# 核心思想：将 GAE 计算转化为可并行的前缀扫描问题
def chunk_scan_gae(full_rewards, full_values, gamma, lambd, chunk_size):
    """
    full_rewards: [B, T]
    full_values: [B, T+1]
    """
    B, T = full_rewards.shape
    num_chunks = (T + chunk_size - 1) // chunk_size

    # 1. 将数据分块
    reward_chunks = full_rewards.split(chunk_size, dim=1)
    value_chunks = full_values.split(chunk_size, dim=1)

    # 2. 并行计算每个 chunk 的局部 delta 和 GAE（伪代码）
    # 每个 chunk 内部是串行的，但 chunk 之间可以并行处理
    chunk_results = []
    for i in range(num_chunks):
        chunk_rewards = reward_chunks[i]
        chunk_values = value_chunks[i]
        next_values = value_chunks[i+1] if i+1 < len(value_chunks) else 0.0
        # 计算该 chunk 的 delta 和局部 GAE
        local_delta = chunk_rewards + gamma * next_values - chunk_values
        local_gae = compute_sequential_gae(local_delta, gamma*lambd) # 内部串行
        chunk_results.append((local_delta, local_gae))

    # 3. 前缀扫描合并 chunk (简化示意)
    full_advantages = torch.zeros_like(full_rewards)
    carry = 0.0  # 跨 chunk 的累积因子
    for i in range(num_chunks):
        local_delta, local_gae = chunk_results[i]
        # 将上一个 chunk 的尾部影响加到当前 chunk 的 GAE 上
        adjusted_gae = local_gae + carry
        # 更新 carry 用于下一个 chunk
        carry = (gamma * lambd) ** chunk_size * local_gae[:, -1:]
        # 存储结果
        start_idx = i * chunk_size
        end_idx = start_idx + chunk_size
        full_advantages[:, start_idx:end_idx] = adjusted_gae
    return full_advantages

# 代码思想基于 sglang/runtime/multimodal_extensions/qwen_vl.py
# 关键步骤：Tokenizer 扩展与 M-RoPE 位置编码
def process_qwen_vl_request(generate_req_input):
    """
    generate_req_input: 包含 text 和 image_data 的请求输入
    """
    text = generate_req_input.text
    image_data_list = generate_req_input.image_data

    # 1. 并发加载和预处理图像
    pixel_values_list = []
    for img_data in image_data_list:
        # 加载图像，并应用模型特定的 resize (如 smart_resize)
        pixel_values = load_and_preprocess_image(img_data)
        pixel_values_list.append(pixel_values)

    # 2. Tokenization 与即时 Token 扩展
    # 原始 prompt 可能包含类似 `<|vision_start|><image><|vision_end|>` 的占位符
    # Tokenizer 会直接将 `<image>` 替换为一连串特定的 image placeholder tokens (如 <|image_pad|>)
    input_ids = tokenizer.encode(text)
    # 此时 input_ids 中已经包含了代表图像区域的特殊 token 序列

    # 3. 计算 M-RoPE (Multimodal Rotary Position Embedding) 位置
    # 这为图像 token 和文本 token 提供了统一的、精确的位置信息
    mrope_positions = compute_mrope_positions(input_ids, pixel_values_list)

    # 4. 构建多模态数据项
    mm_items = []
    for pv in pixel_values_list:
        mm_items.append(MultimodalDataItem(pixel_values=pv))

    # 5. 返回给调度器
    return {
        "input_ids": input_ids,
        "mm_items": mm_items,
        "mrope_positions": mrope_positions
    }

# 在模型前向传播时 (如 qwen2_5_vl.py)
def forward_in_model(input_ids, mrope_positions, mm_items):
    # 获取常规的词嵌入（包含特殊 token）
    embeddings = word_embedding(input_ids)
    # 应用 RoPE，位置信息由 mrope_positions 提供
    embeddings = apply_rotary_pos_emb(embeddings, mrope_positions)

    # 识别出 input_ids 中图像占位符 token 的位置
    image_token_indices = find_image_token_indices(input_ids)
    # 获取视觉特征
    image_features = vision_transformer(mm_items.pixel_values)
    # 将视觉特征投影到语言模型嵌入空间
    projected_image_features = project_to_lm_dim(image_features)
    # 将视觉特征注入到对应位置的 embeddings 中
    embeddings[image_token_indices] = projected_image_features

    # 后续送入 LLM backbone 进行计算
    logits = llm_backbone(embeddings)
    return logits

# 代码思想基于 slime 中训推不一致缓解的实现
# Masked Importance Sampling (MIS) 核心公式应用
def compute_mis_corrected_loss(actor_log_probs,              # 训练引擎计算的 log prob
                               rollout_log_probs,            # 推理引擎 rollout 时计算的 log prob
                               advantages,                   # 优势函数估计
                               clip_epsilon=0.2,
                               mismatch_threshold=0.1):      # 失配阈值
    """
    通过重要性采样权重来修正策略梯度，抑制训推差异过大的样本。
    """
    # 1. 计算重要性权重 (importance weight)
    #   即 rollout 策略与当前训练策略的概率比
    log_ratio = rollout_log_probs - actor_log_probs
    ratio = torch.exp(log_ratio)

    # 2. 计算原始的 PPO 裁剪损失
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * advantages
    original_pg_loss = -torch.min(surr1, surr2).mean()

    # 3. 计算训推差异度量（例如，每个 token 的 KL 散度）
    kl_per_token = actor_log_probs - rollout_log_probs  # 近似 KL
    avg_kl = kl_per_token.mean(dim=-1)  # 序列平均 KL

    # 4. 构造 Mask (基于阈值)
    #    差异过大的样本，其梯度会被抑制
    mask = (avg_kl < mismatch_threshold).float().unsqueeze(-1)  # 扩张维度以匹配 token 级

    # 5. 应用 MIS：将 mask 作为权重乘到原始损失上
    #    或者更精细地，调整重要性权重
    mis_ratio = ratio * mask
    surr1_mis = mis_ratio * advantages
    surr2_mis = torch.clamp(mis_ratio, 1 - clip_epsilon, 1 + clip_epsilon) * advantages
    mis_corrected_loss = -torch.min(surr1_mis, surr2_mis).mean()

    # 也可以选择将 mask 直接作用于损失
    # mis_corrected_loss = (original_pg_loss * mask.mean())

    return mis_corrected_loss, avg_kl, mask