""" # 如果使用sglang加速，需通过以下命令首先启动（transformers格式）模型： python -m sglang.launch_server --model-path ./minimind-3 --attention-backend triton --host 0.0.0.0 --port 8998 sglang引擎是一个高性能LLM推理引擎 1.RadixAttention 前缀树缓存，共享前缀时自动复用kv cache 2.连续批处理 3.结构化输出 """ import os import sys __package__ = "trainer" sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import requests import torch import torch.distributed as dist from abc import ABC, abstractmethod from contextlib import nullcontext from dataclasses import dataclass from typing import List, Optional, Tuple from torch import Tensor from torch.nn.parallel import DistributedDataParallel from transformers import AutoTokenizer def unwrap_policy_model(model: torch.nn.Module) -> torch.nn.Module: raw_model = model.module if isinstance(model, DistributedDataParallel) else getattr(model, "module", model) return getattr(raw_model, "_orig_mod", raw_model) # ===== 计算每个 token 的 logprob ===== def compute_per_token_logps(model, input_ids: Tensor, n_keep: int, attention_mask: Optional[Tensor] = None) -> Tensor: if n_keep <= 0: return input_ids.new_empty((input_ids.size(0), 0), dtype=torch.float32) unwrapped = unwrap_policy_model(model) input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids logits = unwrapped(input_ids, attention_mask=attention_mask, logits_to_keep=n_keep + 1).logits[:, :-1, :] per_token_logps = [] for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]): ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row per_token_logps.append( torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1) ) return torch.stack(per_token_logps) # ===== Rollout 结果 ===== @dataclass class RolloutResult: output_ids: Tensor completion_ids: Tensor per_token_logps: Tensor completions: List[str] prompt_lens: Tensor completion_mask: Tensor # ===== Rollout 引擎抽象基类 ===== class RolloutEngine(ABC): tokenizer = None @abstractmethod def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: pass @abstractmethod def update_policy(self, model: torch.nn.Module): pass # ===== PyTorch 原生推理引擎 ===== class TorchRolloutEngine(RolloutEngine): def __init__(self, policy_model: torch.nn.Module, tokenizer, device: str = "cuda", autocast_ctx=None): self.policy_model = policy_model self.tokenizer = tokenizer self.device = device self.autocast_ctx = autocast_ctx def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: model = unwrap_policy_model(self.policy_model) ctx = self.autocast_ctx if self.autocast_ctx else nullcontext() was_training = model.training model.eval() try: with torch.no_grad(), ctx: output_ids = model.generate( input_ids=prompt_ids.repeat_interleave(num_generations, dim=0), attention_mask=attention_mask.repeat_interleave(num_generations, dim=0), max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature, num_return_sequences=1, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, ).clone() # [B*num_gen, P+R] prompt_len = prompt_ids.size(1) completion_ids = output_ids[:, prompt_len:] # [B*num_gen, R] full_mask = (output_ids != self.tokenizer.pad_token_id).long() per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask) finally: if was_training: model.train() completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True) return RolloutResult(output_ids, completion_ids, per_token_logps, completions, prompt_ids.new_full((output_ids.size(0),), prompt_len), attention_mask.new_ones(output_ids.size(0), completion_ids.size(1))) def update_policy(self, model: torch.nn.Module): self.policy_model = model # ===== SGLang HTTP API 推理引擎 ===== class SGLangRolloutEngine(RolloutEngine): def __init__(self, base_url: str, model_path: str, shared_ckpt_path: str = "./sglang_ckpt", timeout: int = 120): #rstrip，用来删除字符串右边的指定字符 self.base_url = base_url.rstrip('/') self.shared_ckpt_path = shared_ckpt_path self.timeout = timeout self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) self.http = requests def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult: # 去除左侧 padding tokens，只保留有效 token input_ids_list = [] for ids, mask in zip(prompt_ids, attention_mask): valid_ids = ids[mask.bool()].tolist() input_ids_list.append(valid_ids) all_input_ids = [ids for ids in input_ids_list for _ in range(num_generations)] payload = { "input_ids": all_input_ids, "sampling_params": { "temperature": temperature, "max_new_tokens": max_new_tokens, "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [], }, "return_logprob": True, } #像SGLang服务器发送HTTP POST请求，让他根据prompt里的payload生成结果 #调用了python的request库发HTTP请求 resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout) #响应处理，如果响应失败会返回404，500 resp.raise_for_status() #把SGLang返回的json对象解析为python对象 results = resp.json() if not isinstance(results, list): results = [results] all_output_ids, all_completion_ids, all_logprobs = [], [], [] completions = [] for i, result in enumerate(results): meta = result.get("meta_info", {}) completion_ids = meta.get("output_ids", result.get("output_ids", [])) raw_logprobs = meta.get("output_token_logprobs", []) logprobs = [] for item in raw_logprobs: if isinstance(item, (list, tuple)) and len(item) >= 1: logprobs.append(item[0]) elif isinstance(item, (int, float)): logprobs.append(item) if len(logprobs) < len(completion_ids): logprobs = [0.0] * (len(completion_ids) - len(logprobs)) + logprobs elif len(logprobs) > len(completion_ids): logprobs = logprobs[-len(completion_ids):] if completion_ids else [] prompt = all_input_ids[i] full_output = prompt + completion_ids all_output_ids.append(full_output) #模型生成的ids all_completion_ids.append(completion_ids) all_logprobs.append(logprobs) completions.append(self.tokenizer.decode(completion_ids, skip_special_tokens=True)) device = prompt_ids.device #统计最长回答长度 max_comp_len = max(1, max(len(ids) for ids in all_completion_ids)) #统计最长输入输出长度 max_out_len = max(len(ids) for ids in all_input_ids) + max_comp_len def pad_to_tensor(seqs, max_len, pad_val=0): return torch.tensor([s + [pad_val] * (max_len - len(s)) for s in seqs], device=device) pad_id = self.tokenizer.pad_token_id return RolloutResult( output_ids=pad_to_tensor(all_output_ids, max_out_len, pad_val=pad_id), completion_ids=pad_to_tensor(all_completion_ids, max_comp_len, pad_val=pad_id), per_token_logps=pad_to_tensor(all_logprobs, max_comp_len, pad_val=0.0), completions=completions, prompt_lens=torch.tensor([len(ids) for ids in all_input_ids], device=device), completion_mask=torch.tensor([[1] * len(ids) + [0] * (max_comp_len - len(ids)) for ids in all_completion_ids], device=device), ) def update_policy(self, model: torch.nn.Module): #把当前训练的policy模型的权重保存到磁盘上，然后通知SGLang服务器从磁盘加载新的权重 ok = True #dist.is_initialized()判断当前进程是否已经初始化了分布式环境，如果单卡就为false，多卡为true #dist.get_rank()获取当前进程的编号，单卡时为0，多卡时为0,1,2... if not dist.is_initialized() or dist.get_rank() == 0: try: unwrapped = unwrap_policy_model(model) abs_path = os.path.abspath(self.shared_ckpt_path) state_dict = {k: v.detach().half().cpu() for k, v in unwrapped.state_dict().items()} #safe_serialization=False表示保存成Pytorch的model.bin格式 #safe_serialization=True表示保存成model.safetensors格式 unwrapped.save_pretrained(abs_path, state_dict=state_dict, safe_serialization=False) self.tokenizer.save_pretrained(abs_path) resp = self.http.post(f"{self.base_url}/update_weights_from_disk", json={"model_path": abs_path}, timeout=self.timeout) if resp.status_code != 200: print(f"[SGLANG WARNING] update_weights 失败: {resp.status_code}, {resp.text}") ok = resp.status_code == 200 except Exception as e: print(f"[SGLANG WARNING] update_weights 异常: {e}"); ok = False if dist.is_initialized(): #将ok转为0/1，然后放到模型所在的设备上 ok_t = torch.tensor(int(ok), device=next(model.parameters()).device) #在分布式训练里同步ok的值 #broadcast把rank0上ok广播到所有其他rank #dist.barrier()让所有rank等一下，直到所有rank都执行到这里，才继续一起往下走 #item（）把tensor重新转成python的bool dist.broadcast(ok_t, src=0); dist.barrier(); ok = bool(ok_t.item()) if not ok: raise RuntimeError("SGLang update_policy failed") return ok def flush_cache(self) -> bool: #向SGLang服务器发请求，清理缓存 resp = self.http.post(f"{self.base_url}/flush_cache", timeout=30) return resp.status_code == 200 def health(self) -> bool: #向SGLang服务器发请求，检查服务器是否健康 try: resp = self.http.get(f"{self.base_url}/health", timeout=5) return resp.status_code == 200 except: return False # ===== 工厂函数 ===== def create_rollout_engine( engine_type: str = "torch", policy_model: torch.nn.Module = None, tokenizer = None, device: str = "cuda", autocast_ctx = None, sglang_base_url: str = None, sglang_model_path: str = None, sglang_shared_path: str = None, ) -> RolloutEngine: if engine_type == "torch": return TorchRolloutEngine(policy_model, tokenizer, device, autocast_ctx) elif engine_type == "sglang": return SGLangRolloutEngine(sglang_base_url, sglang_model_path, sglang_shared_path) else: raise ValueError(f"不支持的引擎类型: {engine_type}")