260 lines
12 KiB
Python
260 lines
12 KiB
Python
"""
|
||
# 如果使用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}")
|