Tiny-LLM项目：采样策略实现与第二周准备工作指南

引言

在Tiny-LLM项目的第一周第七天，我们将重点实现多种采样策略，并为第二周的优化工作做好准备。本文将详细介绍温度采样、Top-k采样和Top-p（核）采样的实现原理，以及如何为第二周的Metal内核优化工作准备开发环境。

采样策略实现

1. 温度采样（Temperature Sampling）

温度采样是最基础的采样策略之一，它通过温度参数控制输出的随机性。

实现原理：

• 当temp=0时，退化为贪心采样（Greedy Sampling），直接选择概率最高的token
• 当temp>0时，将logits除以温度值后，使用随机分类分布选择下一个token
• 温度值越大，分布越均匀，模型输出越"有创意"；温度值越小，分布越尖锐，输出越确定

实现方法：

# 伪代码示例
scaled_logits = logits / temperature
next_token = mx.random.categorical(scaled_logits)

使用示例：

pdm run main --solution tiny_llm --loader week1 --model Qwen/Qwen2-0.5B-Instruct-MLX --sampler-temp 0.5

2. Top-k采样

Top-k采样通过限制候选token数量来提高生成质量。

实现原理：

1. 保留概率最高的k个token
2. 将其余token的概率设为负无穷
3. 在筛选后的分布上进行温度采样

关键步骤：

• 使用mx.argpartition快速找到top-k元素
• 掩码非top-k元素的logits
• 应用温度采样

使用示例：

pdm run main --solution tiny_llm --loader week1 --model Qwen/Qwen2-0.5B-Instruct-MLX --sampler-temp 0.5 --sampler-top-k 10

3. Top-p采样（核采样）

Top-p采样根据累积概率动态选择候选token。

实现原理：

1. 将token按概率从高到低排序
2. 计算累积概率
3. 保留使累积概率达到p的最小token集合
4. 在筛选后的分布上进行温度采样

实现技巧：

• 使用mx.argsort进行排序
• 计算cumsum获取累积概率
• 动态确定截断点

使用示例：

pdm run main --solution tiny_llm --loader week1 --model Qwen/Qwen2-0.5B-Instruct-MLX --sampler-temp 0.5 --sampler-top-p 0.9

第二周准备工作

第二周将重点优化Qwen2模型的服务基础设施，需要准备Metal开发环境。

1. Xcode安装与配置

完整步骤：

1. 从Mac App Store或Apple开发者网站安装Xcode
2. 首次启动Xcode并安装额外组件
3. 安装命令行工具：

   xcode-select --install
   

4. 确保命令行工具指向正确的Xcode路径：

   sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
   

5. 接受Xcode许可协议：

   sudo xcodebuild -license accept
   

2. 构建工具安装

安装CMake构建工具：

brew install cmake

3. 环境验证

测试扩展模块编译：

pdm run build-ext
pdm run build-ext-test

预期输出应包含correct: True。

4. 预备练习建议

为熟悉Metal编程，建议尝试：

• 实现基础元素级运算（exp、sin、cos等）
• 替换模型中的MLX运算
• 理解Metal着色语言基础

总结

通过本日的学习，我们完成了：

1. 三种重要采样策略的实现
2. 为第二周的Metal优化工作准备了开发环境
3. 了解了基本的性能优化方向

第二周将深入底层优化，使Qwen2模型在Apple Silicon设备上获得最佳性能表现。建议不熟悉Metal/C++的开发者提前进行相关基础知识的学习。