尧图网页设计 
在 Apple silicon 上用 MLX 实现 Llama 推理200 行内完成模型搭建、权重转换与文本生成【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx本文基于 MLX 官方示例文档 llama-inference.rst完整讲解如何用mlx.nn提供的神经网络构件在不到 200 行 Python 内实现一个可推理、可训练的 Llama 系列 Transformer 模型包括基于 RoPE 与 KV Cache 的注意力层、RMSNorm SwiGLU 编码器层、以 Python 生成器实现的自回归采样以及 PyTorch 权重到 MLX 的转换与 NPZ 加载流程。读完本文你将掌握在 Apple silicon 上从零构建并驱动 LLM 推理的完整实战路径并理解 MLX 懒求值与统一内存如何让这段代码同时具备简洁性与高效性。为什么用 MLX 做 LLM 推理MLX 是一个面向 Apple silicon 的数组框架其设计目标是在不牺牲易用性的前提下让中等规模的 Transformer 模型也能在 Mac 上高效推理。支撑这一目标的核心机制有三点统一的设备内存模型CPU 与 GPU 共享内存模型权重与中间张量无需在设备间搬运天然适合大权重、逐 token 生成的推理场景。懒求值Lazy Evaluation算子调用只构建计算图不会立即触发计算。这使我们可以把整段生成逻辑一次性提交给调度器由 MLX 在后台合并与优化 kernel用户几乎不需要手动干预设备同步。完整的神经网络构件库mlx.nnLinear、Embedding、RMSNorm、RoPE等高层模块开箱即用让模型定义保持教科书级的可读性。下面的实现严格复用mlx.nn中的构件模型主体代码控制在 200 行以内。实现模型三个层次我们将按注意力层 → 编码器层 → 完整模型的顺序搭建 Llama 架构。所有实现都可以在examples/python/目录之外独立复刻只依赖 mlx.nn 层实现 中公开的模块。注意力层RoPE 位置编码 可选 KV CacheLlama 注意力层的一个显著特征是不使用绝对位置编码而是采用旋转位置编码RoPERotary Position Embedding[1]。此外为了支持高效推理注意力层需要可选的键/值缓存key/value cache当传入cache时新计算的 keys/values 会与历史缓存拼接避免每生成一个 token 就重复计算整段历史。import mlx.core as mx import mlx.nn as nn class LlamaAttention(nn.Module): def __init__(self, dims: int, num_heads: int): super().__init__() self.num_heads num_heads self.rope nn.RoPE(dims // num_heads, traditionalTrue) self.query_proj nn.Linear(dims, dims, biasFalse) self.key_proj nn.Linear(dims, dims, biasFalse) self.value_proj nn.Linear(dims, dims, biasFalse) self.out_proj nn.Linear(dims, dims, biasFalse) def __call__(self, queries, keys, values, maskNone, cacheNone): queries self.query_proj(queries) keys self.key_proj(keys) values self.value_proj(values) # Extract some shapes num_heads self.num_heads B, L, D queries.shape # Prepare the queries, keys and values for the attention computation queries queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) keys keys.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) values values.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) # Add RoPE to the queries and keys and combine them with the cache if cache is not None: key_cache, value_cache cache queries self.rope(queries, offsetkey_cache.shape[2]) keys self.rope(keys, offsetkey_cache.shape[2]) keys mx.concatenate([key_cache, keys], axis2) values mx.concatenate([value_cache, values], axis2) else: queries self.rope(queries) keys self.rope(keys) # Finally perform the attention computation scale math.sqrt(1 / queries.shape[-1]) scores (queries * scale) keys.transpose(0, 1, 3, 2) if mask is not None: scores scores mask scores mx.softmax(scores, axis-1) values_hat (scores values).transpose(0, 2, 1, 3).reshape(B, L, -1) # Note that we return the keys and values to possibly be used as a cache return self.out_proj(values_hat), (keys, values)几个值得展开的细节RoPE 实例化nn.RoPE(dims // num_heads, traditionalTrue)对每个注意力头的特征维度施加旋转编码。查看 positional_encoding.py 可以看到RoPE模块实际封装了mx.fast.rope并支持traditional、base默认10000即 RoFormer 论文中的频率基数、scale与offset参数。其中traditionalTrue表示按传统方式旋转特征维中相邻的元素对而默认实现则以步长为半个特征维的方式配对以提升效率offset用于告知编码器当前 token 的绝对位置这正是逐 token 生成阶段复用缓存时正确注入位置信息的关键。缓存拼接mx.concatenate([key_cache, keys], axis2)在序列长度维度上拼接历史与新增的 keys/values缓存数组被就地更新后返回供下一轮继续传入。缩放与掩码注意力分数使用scale sqrt(1 / head_dim)缩放mask以加法方式叠加additive mask。这种加法因果掩码的生成方式见下文完整模型部分。编码器层RMSNorm SwiGLULlama 编码器层的另外两个组件是 RMS 归一化 [2] 与 SwiGLU 激活 [3]。RMSNorm 无需计算均值、只对均方根做归一化mlx.nn已内置该模块见 normalization.py其eps默认为1e-5底层调用mx.fast.rms_norm且累加在 32 位精度下完成。SwiGLU 则是把 FFN 的两条支路用门控相乘y a * sigmoid(a) * b。class LlamaEncoderLayer(nn.Module): def __init__(self, dims: int, mlp_dims: int, num_heads: int): super().__init__() self.attention LlamaAttention(dims, num_heads) self.norm1 nn.RMSNorm(dims) self.norm2 nn.RMSNorm(dims) self.linear1 nn.Linear(dims, mlp_dims, biasFalse) self.linear2 nn.Linear(dims, mlp_dims, biasFalse) self.linear3 nn.Linear(mlp_dims, dims, biasFalse) def __call__(self, x, maskNone, cacheNone): y self.norm1(x) y, cache self.attention(y, y, y, mask, cache) x x y y self.norm2(x) a self.linear1(y) b self.linear2(y) y a * mx.sigmoid(a) * b y self.linear3(y) x x y return x, cache注意力使用**预归一化pre-norm**结构先norm1再进入注意力随后残差相加FFN 同理。三路投影linear1/linear2/linear3分别对应 SwiGLU 的 gating、up 与 down 投影与原始 PyTorch Llama 权重中的feed_forward.w1/w3/w2一一对应见下文权重转换。__call__同时返回更新后的cache便于上层逐层收集。完整模型Embedding 层堆叠 输出头把多个LlamaEncoderLayer与 token 嵌入、最终归一化、输出投影组合起来就得到完整的Llama模型。其中mask由nn.MultiHeadAttention.create_additive_causal_mask生成——查看 transformer.py 可以看到它的实现对序列长度N构造indices[:, None] indices[None]的上三角布尔矩阵再乘以mx.finfo(dtype).min使未来位置在 softmax 前被压到-inf附近从而被完全屏蔽。class Llama(nn.Module): def __init__( self, num_layers: int, vocab_size: int, dims: int, mlp_dims: int, num_heads: int ): super().__init__() self.embedding nn.Embedding(vocab_size, dims) self.layers [ LlamaEncoderLayer(dims, mlp_dims, num_heads) for _ in range(num_layers) ] self.norm nn.RMSNorm(dims) self.out_proj nn.Linear(dims, vocab_size, biasFalse) def __call__(self, x): mask nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1]) mask mask.astype(self.embedding.weight.dtype) x self.embedding(x) for l in self.layers: x, _ l(x, mask) x self.norm(x) return self.out_proj(x)注意编码器层存放在一个普通 Python 列表里但这不影响参数管理——model.parameters()依然会递归收集这些层的权重mlx.nn.Module的参数收集是树结构的普通 list 也会被遍历。Embedding的权重形状为(vocab_size, dims)初始化时按scale sqrt(1 / dims)的高斯分布填充见 embedding.py。实现生成Python 生成器 温度采样上面的__call__只能处理单次前向既忽略缓存也不做采样适合训练但不足以支撑推理。因此我们在Llama类上追加generate方法它先像训练前向一样处理完整 prompt同时把每层的缓存保存下来随后进入自回归循环一次只喂入上一个 token逐 token 产出结果。class Llama(nn.Module): ... def generate(self, x, temp1.0): cache [] # Make an additive causal mask. We will need that to process the prompt. mask nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1]) mask mask.astype(self.embedding.weight.dtype) # First we process the prompt x the same way as in __call__ but # save the caches in cache x self.embedding(x) for l in self.layers: x, c l(x, maskmask) cache.append(c) # --- we store the per layer cache in a # simple python list x self.norm(x) y self.out_proj(x[:, -1]) # --- we only care about the last logits # that generate the next token y mx.random.categorical(y * (1/temp)) # y now has size [1] # Since MLX is lazily evaluated nothing is computed yet. # Calling y.item() would force the computation to happen at # this point but we can also choose not to do that and let the # user choose when to start the computation. yield y # Now we parsed the prompt and generated the first token we # need to feed it back into the model and loop to generate the # rest. while True: # Unsqueezing the last dimension to add a sequence length # dimension of 1 x y[:, None] x self.embedding(x) for i in range(len(cache)): # We are overwriting the arrays in the cache list. When # the computation will happen, MLX will be discarding the # old cache the moment it is not needed anymore. x, cache[i] self.layersi x self.norm(x) y self.out_proj(x[:, -1]) y mx.random.categorical(y * (1/temp)) yield y这段生成器代码有三处关键设计先并行、后自回归prompt 阶段一次前向处理整段输入充分利用矩阵乘并行度此后每轮只处理长度为 1 的序列配合缓存把每 token 的计算量降到最低。缓存就地覆盖cache[i]被不断覆盖为新数组。得益于懒求值当计算真正发生时MLX 会在旧缓存不再被引用后立刻回收内存因此长时间生成也不会让缓存无限膨胀。采样与温度mx.random.categorical(y * (1/temp))对 logits 做温度缩放后按类别分布采样temp1.0为原始分布越小越贪婪、越大越发散。把一切组合起来懒求值下的完整流程现在可以实例化一个小 Llama 模型喂入 prompt 并生成 tokenmodel Llama(num_layers12, vocab_size8192, dims512, mlp_dims1024, num_heads8) # Since MLX is lazily evaluated nothing has actually been materialized yet. # We could have set the dims to 20_000 on a machine with 8GB of RAM and the # code above would still run. Lets actually materialize the model. mx.eval(model.parameters()) prompt mx.array([[1, 10, 8, 32, 44, 7]]) # -- Note the double brackets because we # have a batch dimension even # though it is 1 in this case generated [t for i, t in zip(range(10), model.generate(prompt, 0.8))] # Since we havent evaluated anything, nothing is computed yet. The list # generated contains the arrays that hold the computation graph for the # full processing of the prompt and the generation of 10 tokens. # # We can evaluate them one at a time, or all together. Concatenate them or # print them. They would all result in very similar runtimes and give exactly # the same results. mx.eval(generated)这里浓缩了 MLX 懒求值模型的精髓实例化模型时没有任何权重真正落盘只有mx.eval(model.parameters())才会强制物化这也是为什么在 8GB 内存的机器上可以先用很大dims构建模型而不会立即 OOM。model.generate(...)迭代 10 次只是把 10 个 token 的计算图累积起来mx.eval(generated)才一次性触发整段计算。无论逐个 eval、整体 eval、拼接还是打印运行时间几乎相同结果完全一致——计算图合并后的调度由 MLX 统一完成。prompt 需要双括号[[...]]保留 batch 维度即使 batch size 是 1。转换 PyTorch 权重到 MLX 格式上一节的随机初始化模型只是验证流程。要真正使用 Llama 权重需要把 PyTorch 的 checkpoint 转换成 MLX 可加载的 NPZ 文件。假设你已经拥有原始 Llama 权重及其配套的 SentencePiece 词表下面的脚本完成转换import argparse from itertools import starmap import numpy as np import torch def map_torch_to_mlx(key, value): if tok_embedding in key: key embedding.weight elif norm in key: key key.replace(attention_norm, norm1).replace(ffn_norm, norm2) elif wq in key or wk in key or wv in key or wo in key: key key.replace(wq, query_proj) key key.replace(wk, key_proj) key key.replace(wv, value_proj) key key.replace(wo, out_proj) elif w1 in key or w2 in key or w3 in key: # The FFN is a separate submodule in PyTorch key key.replace(feed_forward.w1, linear1) key key.replace(feed_forward.w3, linear2) key key.replace(feed_forward.w2, linear3) elif output in key: key key.replace(output, out_proj) elif rope in key: return None, None return key, value.numpy() if __name__ __main__: parser argparse.ArgumentParser(descriptionConvert Llama weights to MLX) parser.add_argument(torch_weights) parser.add_argument(output_file) args parser.parse_args() state torch.load(args.torch_weights) np.savez( args.output_file, **{k: v for k, v in starmap(map_torch_to_mlx, state.items()) if k is not None} )命名映射规则清晰体现了我们模型与官方权重的一一对应PyTorch 权重名MLX 参数名说明tok_embeddings.weightembedding.weighttoken 嵌入layers.{i}.attention_norm.weightlayers.{i}.norm1.weight注意力前归一化pre-normlayers.{i}.ffn_norm.weightlayers.{i}.norm2.weightFFN 前归一化layers.{i}.attention.wq/wk/wv/wo.weightlayers.{i}.attention.query_proj/key_proj/value_proj/out_proj.weight四路注意力投影layers.{i}.feed_forward.w1/w3/w2.weightlayers.{i}.linear1/linear2/linear3.weightSwiGLU 三路投影output.weightout_proj.weight输出投影rope.freqs跳过旋转频率由mx.fast.rope运行时计算转换时直接return None, None丢弃 PyTorch 预计算的 RoPE 频率表因为 MLX 的RoPE模块会根据base参数自行计算角度频率。权重加载与推理基准转换完成后用mlx.core.load读取 NPZ 文件再借助mlx.utils.tree_unflatten把扁平的 key/value 字典还原成嵌套字典即可直接model.updatefrom mlx.utils import tree_unflatten model.update(tree_unflatten(list(mx.load(weight_file).items())))tree_unflatten会把形如layers.2.attention.query_proj.weight的点分键转换为{layers: [..., ..., {attention: {query_proj: {weight: ...}}}]}进而以与model.parameters()完全同构的嵌套结构更新模型参数实现见 utils.py。需要注意上述加载路径包含磁盘 → numpy → MLX的多次不必要拷贝文档中已说明未来会替换为直接加载到 MLX 的方式。完整的示例脚本可在mlx-examples仓库获得。假设 PyTorch Llama 权重位于llama-7B/目录执行$ python convert.py --torch-path llama-7B/ $ python llama.py --prompt Call me Ishmael. Some years ago never mind how long precisely [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ , having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, ------ [INFO] Prompt processing: 0.437 s [INFO] Full generation: 4.330 s以下是该文档记录的代表性基准数据环境为 M1 Ultra、7B 参数 Llama 模型具体数值会随设备与模型规模变化仅供参考从磁盘加载模型权重约5.247 s处理上述短 prompt 约0.437 s生成 100 个 token 共4.330 s即约39 ms/token将 prompt 显著加长后prompt 处理约0.579 s、完整生成4.690 s——prompt 处理时间与逐 token 生成时间几乎保持不变用--max-tokens 500生成长文本时prompt 处理约0.633 s完整生成约21.475 s。这组数据说明两件事第一由于 prompt 阶段是并行矩阵运算prompt 再长也只占很小一部分时间第二生成阶段每 token 的耗时基本恒定这正是 KV Cache 与 MLX 统一内存调度共同作用的结果——生成过程中无需在 CPU/GPU 间搬运权重从而把单 token 延迟压到毫秒级。延伸阅读本文完整示例代码在mlx-examples仓库的llms/llama目录外部托管未随本仓库分发本仓库内的mlx.nn层实现layers是理解每个构件底层行为的第一手资料。相关论文[1] RoFormer: Enhanced Transformer with Rotary Position EmbeddingarXiv:2104.09864[2] Root Mean Square Layer NormalizationNeurIPS 2019[3] GLU Variants Improve TransformerarXiv:2002.05202。若想进一步了解 KV Cache 的工程细节与推理场景下的使用建议可参考本仓库文档 kv_cache.rst关于 MLX 懒求值与统一内存模型可阅读 lazy_evaluation.rst 与 unified_memory.rst。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 在 Apple silicon 上用 MLX 实现 Llama 推理：200 行内完成模型搭建、权重转换与文本生成
在 Apple silicon 上用 MLX 实现 Llama 推理200 行内完成模型搭建、权重转换与文本生成【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx本文基于 MLX 官方示例文档 llama-inference.rst完整讲解如何用mlx.nn提供的神经网络构件在不到 200 行 Python 内实现一个可推理、可训练的 Llama 系列 Transformer 模型包括基于 RoPE 与 KV Cache 的注意力层、RMSNorm SwiGLU 编码器层、以 Python 生成器实现的自回归采样以及 PyTorch 权重到 MLX 的转换与 NPZ 加载流程。读完本文你将掌握在 Apple silicon 上从零构建并驱动 LLM 推理的完整实战路径并理解 MLX 懒求值与统一内存如何让这段代码同时具备简洁性与高效性。为什么用 MLX 做 LLM 推理MLX 是一个面向 Apple silicon 的数组框架其设计目标是在不牺牲易用性的前提下让中等规模的 Transformer 模型也能在 Mac 上高效推理。支撑这一目标的核心机制有三点统一的设备内存模型CPU 与 GPU 共享内存模型权重与中间张量无需在设备间搬运天然适合大权重、逐 token 生成的推理场景。懒求值Lazy Evaluation算子调用只构建计算图不会立即触发计算。这使我们可以把整段生成逻辑一次性提交给调度器由 MLX 在后台合并与优化 kernel用户几乎不需要手动干预设备同步。完整的神经网络构件库mlx.nnLinear、Embedding、RMSNorm、RoPE等高层模块开箱即用让模型定义保持教科书级的可读性。下面的实现严格复用mlx.nn中的构件模型主体代码控制在 200 行以内。实现模型三个层次我们将按注意力层 → 编码器层 → 完整模型的顺序搭建 Llama 架构。所有实现都可以在examples/python/目录之外独立复刻只依赖 mlx.nn 层实现 中公开的模块。注意力层RoPE 位置编码 可选 KV CacheLlama 注意力层的一个显著特征是不使用绝对位置编码而是采用旋转位置编码RoPERotary Position Embedding[1]。此外为了支持高效推理注意力层需要可选的键/值缓存key/value cache当传入cache时新计算的 keys/values 会与历史缓存拼接避免每生成一个 token 就重复计算整段历史。import mlx.core as mx import mlx.nn as nn class LlamaAttention(nn.Module): def __init__(self, dims: int, num_heads: int): super().__init__() self.num_heads num_heads self.rope nn.RoPE(dims // num_heads, traditionalTrue) self.query_proj nn.Linear(dims, dims, biasFalse) self.key_proj nn.Linear(dims, dims, biasFalse) self.value_proj nn.Linear(dims, dims, biasFalse) self.out_proj nn.Linear(dims, dims, biasFalse) def __call__(self, queries, keys, values, maskNone, cacheNone): queries self.query_proj(queries) keys self.key_proj(keys) values self.value_proj(values) # Extract some shapes num_heads self.num_heads B, L, D queries.shape # Prepare the queries, keys and values for the attention computation queries queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) keys keys.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) values values.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) # Add RoPE to the queries and keys and combine them with the cache if cache is not None: key_cache, value_cache cache queries self.rope(queries, offsetkey_cache.shape[2]) keys self.rope(keys, offsetkey_cache.shape[2]) keys mx.concatenate([key_cache, keys], axis2) values mx.concatenate([value_cache, values], axis2) else: queries self.rope(queries) keys self.rope(keys) # Finally perform the attention computation scale math.sqrt(1 / queries.shape[-1]) scores (queries * scale) keys.transpose(0, 1, 3, 2) if mask is not None: scores scores mask scores mx.softmax(scores, axis-1) values_hat (scores values).transpose(0, 2, 1, 3).reshape(B, L, -1) # Note that we return the keys and values to possibly be used as a cache return self.out_proj(values_hat), (keys, values)几个值得展开的细节RoPE 实例化nn.RoPE(dims // num_heads, traditionalTrue)对每个注意力头的特征维度施加旋转编码。查看 positional_encoding.py 可以看到RoPE模块实际封装了mx.fast.rope并支持traditional、base默认10000即 RoFormer 论文中的频率基数、scale与offset参数。其中traditionalTrue表示按传统方式旋转特征维中相邻的元素对而默认实现则以步长为半个特征维的方式配对以提升效率offset用于告知编码器当前 token 的绝对位置这正是逐 token 生成阶段复用缓存时正确注入位置信息的关键。缓存拼接mx.concatenate([key_cache, keys], axis2)在序列长度维度上拼接历史与新增的 keys/values缓存数组被就地更新后返回供下一轮继续传入。缩放与掩码注意力分数使用scale sqrt(1 / head_dim)缩放mask以加法方式叠加additive mask。这种加法因果掩码的生成方式见下文完整模型部分。编码器层RMSNorm SwiGLULlama 编码器层的另外两个组件是 RMS 归一化 [2] 与 SwiGLU 激活 [3]。RMSNorm 无需计算均值、只对均方根做归一化mlx.nn已内置该模块见 normalization.py其eps默认为1e-5底层调用mx.fast.rms_norm且累加在 32 位精度下完成。SwiGLU 则是把 FFN 的两条支路用门控相乘y a * sigmoid(a) * b。class LlamaEncoderLayer(nn.Module): def __init__(self, dims: int, mlp_dims: int, num_heads: int): super().__init__() self.attention LlamaAttention(dims, num_heads) self.norm1 nn.RMSNorm(dims) self.norm2 nn.RMSNorm(dims) self.linear1 nn.Linear(dims, mlp_dims, biasFalse) self.linear2 nn.Linear(dims, mlp_dims, biasFalse) self.linear3 nn.Linear(mlp_dims, dims, biasFalse) def __call__(self, x, maskNone, cacheNone): y self.norm1(x) y, cache self.attention(y, y, y, mask, cache) x x y y self.norm2(x) a self.linear1(y) b self.linear2(y) y a * mx.sigmoid(a) * b y self.linear3(y) x x y return x, cache注意力使用**预归一化pre-norm**结构先norm1再进入注意力随后残差相加FFN 同理。三路投影linear1/linear2/linear3分别对应 SwiGLU 的 gating、up 与 down 投影与原始 PyTorch Llama 权重中的feed_forward.w1/w3/w2一一对应见下文权重转换。__call__同时返回更新后的cache便于上层逐层收集。完整模型Embedding 层堆叠 输出头把多个LlamaEncoderLayer与 token 嵌入、最终归一化、输出投影组合起来就得到完整的Llama模型。其中mask由nn.MultiHeadAttention.create_additive_causal_mask生成——查看 transformer.py 可以看到它的实现对序列长度N构造indices[:, None] indices[None]的上三角布尔矩阵再乘以mx.finfo(dtype).min使未来位置在 softmax 前被压到-inf附近从而被完全屏蔽。class Llama(nn.Module): def __init__( self, num_layers: int, vocab_size: int, dims: int, mlp_dims: int, num_heads: int ): super().__init__() self.embedding nn.Embedding(vocab_size, dims) self.layers [ LlamaEncoderLayer(dims, mlp_dims, num_heads) for _ in range(num_layers) ] self.norm nn.RMSNorm(dims) self.out_proj nn.Linear(dims, vocab_size, biasFalse) def __call__(self, x): mask nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1]) mask mask.astype(self.embedding.weight.dtype) x self.embedding(x) for l in self.layers: x, _ l(x, mask) x self.norm(x) return self.out_proj(x)注意编码器层存放在一个普通 Python 列表里但这不影响参数管理——model.parameters()依然会递归收集这些层的权重mlx.nn.Module的参数收集是树结构的普通 list 也会被遍历。Embedding的权重形状为(vocab_size, dims)初始化时按scale sqrt(1 / dims)的高斯分布填充见 embedding.py。实现生成Python 生成器 温度采样上面的__call__只能处理单次前向既忽略缓存也不做采样适合训练但不足以支撑推理。因此我们在Llama类上追加generate方法它先像训练前向一样处理完整 prompt同时把每层的缓存保存下来随后进入自回归循环一次只喂入上一个 token逐 token 产出结果。class Llama(nn.Module): ... def generate(self, x, temp1.0): cache [] # Make an additive causal mask. We will need that to process the prompt. mask nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1]) mask mask.astype(self.embedding.weight.dtype) # First we process the prompt x the same way as in __call__ but # save the caches in cache x self.embedding(x) for l in self.layers: x, c l(x, maskmask) cache.append(c) # --- we store the per layer cache in a # simple python list x self.norm(x) y self.out_proj(x[:, -1]) # --- we only care about the last logits # that generate the next token y mx.random.categorical(y * (1/temp)) # y now has size [1] # Since MLX is lazily evaluated nothing is computed yet. # Calling y.item() would force the computation to happen at # this point but we can also choose not to do that and let the # user choose when to start the computation. yield y # Now we parsed the prompt and generated the first token we # need to feed it back into the model and loop to generate the # rest. while True: # Unsqueezing the last dimension to add a sequence length # dimension of 1 x y[:, None] x self.embedding(x) for i in range(len(cache)): # We are overwriting the arrays in the cache list. When # the computation will happen, MLX will be discarding the # old cache the moment it is not needed anymore. x, cache[i] self.layersi x self.norm(x) y self.out_proj(x[:, -1]) y mx.random.categorical(y * (1/temp)) yield y这段生成器代码有三处关键设计先并行、后自回归prompt 阶段一次前向处理整段输入充分利用矩阵乘并行度此后每轮只处理长度为 1 的序列配合缓存把每 token 的计算量降到最低。缓存就地覆盖cache[i]被不断覆盖为新数组。得益于懒求值当计算真正发生时MLX 会在旧缓存不再被引用后立刻回收内存因此长时间生成也不会让缓存无限膨胀。采样与温度mx.random.categorical(y * (1/temp))对 logits 做温度缩放后按类别分布采样temp1.0为原始分布越小越贪婪、越大越发散。把一切组合起来懒求值下的完整流程现在可以实例化一个小 Llama 模型喂入 prompt 并生成 tokenmodel Llama(num_layers12, vocab_size8192, dims512, mlp_dims1024, num_heads8) # Since MLX is lazily evaluated nothing has actually been materialized yet. # We could have set the dims to 20_000 on a machine with 8GB of RAM and the # code above would still run. Lets actually materialize the model. mx.eval(model.parameters()) prompt mx.array([[1, 10, 8, 32, 44, 7]]) # -- Note the double brackets because we # have a batch dimension even # though it is 1 in this case generated [t for i, t in zip(range(10), model.generate(prompt, 0.8))] # Since we havent evaluated anything, nothing is computed yet. The list # generated contains the arrays that hold the computation graph for the # full processing of the prompt and the generation of 10 tokens. # # We can evaluate them one at a time, or all together. Concatenate them or # print them. They would all result in very similar runtimes and give exactly # the same results. mx.eval(generated)这里浓缩了 MLX 懒求值模型的精髓实例化模型时没有任何权重真正落盘只有mx.eval(model.parameters())才会强制物化这也是为什么在 8GB 内存的机器上可以先用很大dims构建模型而不会立即 OOM。model.generate(...)迭代 10 次只是把 10 个 token 的计算图累积起来mx.eval(generated)才一次性触发整段计算。无论逐个 eval、整体 eval、拼接还是打印运行时间几乎相同结果完全一致——计算图合并后的调度由 MLX 统一完成。prompt 需要双括号[[...]]保留 batch 维度即使 batch size 是 1。转换 PyTorch 权重到 MLX 格式上一节的随机初始化模型只是验证流程。要真正使用 Llama 权重需要把 PyTorch 的 checkpoint 转换成 MLX 可加载的 NPZ 文件。假设你已经拥有原始 Llama 权重及其配套的 SentencePiece 词表下面的脚本完成转换import argparse from itertools import starmap import numpy as np import torch def map_torch_to_mlx(key, value): if tok_embedding in key: key embedding.weight elif norm in key: key key.replace(attention_norm, norm1).replace(ffn_norm, norm2) elif wq in key or wk in key or wv in key or wo in key: key key.replace(wq, query_proj) key key.replace(wk, key_proj) key key.replace(wv, value_proj) key key.replace(wo, out_proj) elif w1 in key or w2 in key or w3 in key: # The FFN is a separate submodule in PyTorch key key.replace(feed_forward.w1, linear1) key key.replace(feed_forward.w3, linear2) key key.replace(feed_forward.w2, linear3) elif output in key: key key.replace(output, out_proj) elif rope in key: return None, None return key, value.numpy() if __name__ __main__: parser argparse.ArgumentParser(descriptionConvert Llama weights to MLX) parser.add_argument(torch_weights) parser.add_argument(output_file) args parser.parse_args() state torch.load(args.torch_weights) np.savez( args.output_file, **{k: v for k, v in starmap(map_torch_to_mlx, state.items()) if k is not None} )命名映射规则清晰体现了我们模型与官方权重的一一对应PyTorch 权重名MLX 参数名说明tok_embeddings.weightembedding.weighttoken 嵌入layers.{i}.attention_norm.weightlayers.{i}.norm1.weight注意力前归一化pre-normlayers.{i}.ffn_norm.weightlayers.{i}.norm2.weightFFN 前归一化layers.{i}.attention.wq/wk/wv/wo.weightlayers.{i}.attention.query_proj/key_proj/value_proj/out_proj.weight四路注意力投影layers.{i}.feed_forward.w1/w3/w2.weightlayers.{i}.linear1/linear2/linear3.weightSwiGLU 三路投影output.weightout_proj.weight输出投影rope.freqs跳过旋转频率由mx.fast.rope运行时计算转换时直接return None, None丢弃 PyTorch 预计算的 RoPE 频率表因为 MLX 的RoPE模块会根据base参数自行计算角度频率。权重加载与推理基准转换完成后用mlx.core.load读取 NPZ 文件再借助mlx.utils.tree_unflatten把扁平的 key/value 字典还原成嵌套字典即可直接model.updatefrom mlx.utils import tree_unflatten model.update(tree_unflatten(list(mx.load(weight_file).items())))tree_unflatten会把形如layers.2.attention.query_proj.weight的点分键转换为{layers: [..., ..., {attention: {query_proj: {weight: ...}}}]}进而以与model.parameters()完全同构的嵌套结构更新模型参数实现见 utils.py。需要注意上述加载路径包含磁盘 → numpy → MLX的多次不必要拷贝文档中已说明未来会替换为直接加载到 MLX 的方式。完整的示例脚本可在mlx-examples仓库获得。假设 PyTorch Llama 权重位于llama-7B/目录执行$ python convert.py --torch-path llama-7B/ $ python llama.py --prompt Call me Ishmael. Some years ago never mind how long precisely [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ , having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, ------ [INFO] Prompt processing: 0.437 s [INFO] Full generation: 4.330 s以下是该文档记录的代表性基准数据环境为 M1 Ultra、7B 参数 Llama 模型具体数值会随设备与模型规模变化仅供参考从磁盘加载模型权重约5.247 s处理上述短 prompt 约0.437 s生成 100 个 token 共4.330 s即约39 ms/token将 prompt 显著加长后prompt 处理约0.579 s、完整生成4.690 s——prompt 处理时间与逐 token 生成时间几乎保持不变用--max-tokens 500生成长文本时prompt 处理约0.633 s完整生成约21.475 s。这组数据说明两件事第一由于 prompt 阶段是并行矩阵运算prompt 再长也只占很小一部分时间第二生成阶段每 token 的耗时基本恒定这正是 KV Cache 与 MLX 统一内存调度共同作用的结果——生成过程中无需在 CPU/GPU 间搬运权重从而把单 token 延迟压到毫秒级。延伸阅读本文完整示例代码在mlx-examples仓库的llms/llama目录外部托管未随本仓库分发本仓库内的mlx.nn层实现layers是理解每个构件底层行为的第一手资料。相关论文[1] RoFormer: Enhanced Transformer with Rotary Position EmbeddingarXiv:2104.09864[2] Root Mean Square Layer NormalizationNeurIPS 2019[3] GLU Variants Improve TransformerarXiv:2002.05202。若想进一步了解 KV Cache 的工程细节与推理场景下的使用建议可参考本仓库文档 kv_cache.rst关于 MLX 懒求值与统一内存模型可阅读 lazy_evaluation.rst 与 unified_memory.rst。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考