推理API参考

import json
from model import ModelArgs
 
# 从文件加载配置
with open("configs/config_671B.json") as f:
    args = ModelArgs(**json.load(f))

import torch
from model import Transformer, ModelArgs
 
# 根据配置将默认数据类型设置为bfloat16或float8_e4m3fn
torch.set_default_dtype(torch.bfloat16)  # 或torch.float8_e4m3fn用于FP8
 
# 初始化模型
with torch.device("cuda"):
    model = Transformer(args)

from safetensors.torch import load_model
import os
 
# 加载模型权重
load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))

@torch.inference_mode()
def forward(self, tokens: torch.Tensor, start_pos: int = 0):
    """
    Transformer模型的正向传递。
 
    Args:
        tokens (torch.Tensor): 形状为(batch_size, seq_len)的令牌ID输入张量。
        start_pos (int, optional): 序列中的起始位置。默认为0。
 
    Returns:
        torch.Tensor: 下一个令牌预测的logits张量。
    """
    # 实现细节...
    return logits

@torch.inference_mode()
def generate(
    model: Transformer,
    prompt_tokens: List[List[int]],
    max_new_tokens: int,
    eos_id: int,
    temperature: float = 1.0
) -> List[List[int]]:
    """
    基于给定的提示令牌生成新令牌。
 
    Args:
        model (Transformer): 用于令牌生成的变换器模型。
        prompt_tokens (List[List[int]]): 包含提示令牌的列表。
        max_new_tokens (int): 要生成的最大新令牌数。
        eos_id (int): 结束序列令牌ID。
        temperature (float): 采样温度（1.0=标准，<1.0=更窄，>1.0=更宽）。
 
    Returns:
        List[List[int]]: 包含生成令牌的列表。
    """
    # 实现细节...
    return completion_tokens

from transformers import AutoTokenizer
 
tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
 
# 初始化消息历史
messages = []
 
# 添加用户消息
messages.append({"role": "user", "content": "Tell me about DeepSeek-V3"})
 
# 应用聊天模板并生成响应
prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
completion_tokens = generate(model, [prompt_tokens], max_new_tokens, 
                            tokenizer.eos_token_id, temperature)
completion = tokenizer.decode(completion_tokens[0], skip_special_tokens=True)
 
# 将助手响应添加到消息历史
messages.append({"role": "assistant", "content": completion})

# 从文件读取多个提示
with open(input_file) as f:
    prompts = [line.strip() for line in f.readlines()]
 
# 编码所有提示
prompt_tokens = [tokenizer.apply_chat_template([{"role": "user", "content": prompt}], 
                add_generation_prompt=True) for prompt in prompts]
 
# 为所有提示生成补全
completion_tokens = generate(model, prompt_tokens, max_new_tokens, 
                            tokenizer.eos_token_id, temperature)
completions = tokenizer.batch_decode(completion_tokens, skip_special_tokens=True)

import torch.distributed as dist
 
# 初始化分布式环境
world_size = int(os.getenv("WORLD_SIZE", "1"))
rank = int(os.getenv("RANK", "0"))
local_rank = int(os.getenv("LOCAL_RANK", "0"))
 
if world_size > 1:
    dist.init_process_group("nccl")
    torch.cuda.set_device(local_rank)
 
# 使用torchrun启动分布式推理
# torchrun --nnodes 2 --nproc-per-node 8 --node-rank $RANK --master-addr $ADDR generate.py 
#   --ckpt-path /path/to/model --config configs/config_671B.json --interactive

# 在配置中将FP8设置为模型的数据类型
args = ModelArgs(dtype="fp8", ...)
 
# 模型将内部使用torch.float8_e4m3fn
model = Transformer(args)

# 从命令行：
python fp8_cast_bf16.py --input-fp8-hf-path /path/to/fp8_weights --output-bf16-hf-path /path/to/bf16_weights

python convert.py --hf-ckpt-path /path/to/DeepSeek-V3 --save-path /path/to/DeepSeek-V3-Demo --n-experts 256 --model-parallel 16