运行推理

推理过程概述

运行 Grok-1 推理涉及多个关键组件的协同工作

初始化：加载模型权重、设置设备网格并编译推理函数
分词：将文本提示转换为 token 序列
前向传播：通过 314B 参数的 MoE transformer 处理 token
采样：从模型输出的概率分布中选择 token
文本生成：将生成的 token 转换回文本

来源：runners.py442-577 run.py24-67

核心推理组件

推理系统围绕两个主要类构建

InferenceRunner：用于运行推理的高级接口
ModelRunner：处理模型初始化和检查点加载

来源：runners.py252-270 runners.py136-249

设置推理环境

要运行 Grok-1 推理，您需要

配置模型参数
创建一个 ModelRunner 实例
创建一个 InferenceRunner 实例
初始化运行器
使用您的提示运行推理

以下代码演示了设置过程

来源：run.py50-67 runners.py275-441

基本使用示例

这是一个运行 Grok-1 推理的最小示例

# Set up model configuration
grok_1_model = LanguageModelConfig(...)  # Configuration omitted for brevity

# Create InferenceRunner
inference_runner = InferenceRunner(
    pad_sizes=(1024,),
    runner=ModelRunner(
        model=grok_1_model,
        bs_per_device=0.125,
        checkpoint_path=CKPT_PATH,
    ),
    name="local",
    load=CKPT_PATH,
    tokenizer_path="./tokenizer.model",
    local_mesh_config=(1, 8),
    between_hosts_config=(1, 1),
)

# Initialize
inference_runner.initialize()
gen = inference_runner.run()

# Run inference
prompt = "The answer to life the universe and everything is of course"
output = sample_from_model(gen, prompt, max_len=100, temperature=0.01)
print(output)

来源：run.py24-67 runners.py596-605

关键推理参数

运行推理时，您可以通过多个参数来控制生成过程

参数	描述	默认示例
`prompt`	输入给模型的文本	“生命的答案是...”
`max_len`	生成的序列的最大长度	100
`temperature`	控制随机性（值越低，结果越确定）	0.01
`nucleus_p`	控制 nucleus 采样的截止点	1.0
`rng_seed`	用于可复现生成的随机种子	42