ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

MLX 数组序列化完全指南:掌握 save / load / savez / save_safetensors / save_gguf 四种格式

MLX 数组序列化完全指南:掌握 save / load / savez / save_safetensors / save_gguf 四种格式 MLX 数组序列化完全指南掌握 save / load / savez / save_safetensors / save_gguf 四种格式【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx导读MLX 为 Apple silicon 上的数组框架在模型训练、权重保存与推理部署中数组序列化是绕不开的基础环节。本文以官方文档 docs/src/usage/saving_and_loading.rst 为主线系统讲解 MLX 支持的四种数组序列化格式.npy、.npz、.safetensors、.gguf的保存与加载方法并结合仓库源码剖析load按扩展名自动识别格式、惰性加载、端序处理、GGUF 量化张量转换等底层细节。读完本文你将能够根据业务场景单数组备份、多权重打包、Hugging Face 模型互操作、LLM GGUF 文件读写准确选用 API并理解其背后的实现机制。一、四种序列化格式总览MLX 官方文档给出了一个非常清晰的格式对照表这里完整保留并补充实现细节格式扩展名保存函数说明NumPy.npymx.save仅支持单个数组NumPy 压缩包.npzmx.savez/mx.savez_compressed支持多个数组Safetensors.safetensorsmx.save_safetensors支持多个数组带元数据GGUF.ggufmx.save_gguf支持多个数组带元数据其中mx.load是一个统一入口它能加载上述任意一种格式根据文件扩展名自动推断格式不同格式的返回值类型不同详见下文各节。所有保存函数都接受字符串路径、pathlib.Path其中save、savez、save_safetensors还接受已打开的二进制文件对象file-like object。二、单数组持久化mx.save 与 .npy 格式2.1 基本用法保存一个数组到.npy文件import mlx.core as mx a mx.array([1.0]) mx.save(array, a) # 生成文件 array.npy注意扩展名是自动补全的。官方文档明确说明如果缺少扩展名会自动加上因此mx.save(array, a)与mx.save(array.npy, a)等价。加载时既可以用补全后的文件名也可以直接写不带扩展名的名字加载同样会自动补全mx.load(array.npy) # array([1], dtypefloat32)2.2 源码视角.npy 文件是怎么写出来的save的 C 实现在 mlx/io/load.cpp 中关键步骤为先求值再打开文件a contiguous(a, true); a.eval();之后才调用out_stream-open()。之所以刻意保持这个顺序是因为FileWriter打开文件时会O_TRUNC截断文件见 mlx/io/load.h如果输入数组本身是从同一文件惰性加载来的lazy input先截断会导致读不到数据。这一点在测试 python/tests/test_load.py 中有专门验证先mx.save再mx.load然后对惰性数组a 2后覆盖保存结果必须为3。写魔数与版本写入 6 字节魔数\x93NUMPYMAGIC常量见 mlx/io/load.cpp再根据头部长度选择 NPY 1.x2 字节长度字段或 2.x4 字节长度字段版本头部按 16 字节对齐并补空格和换行——这与 NumPy 官方格式规范完全一致。写数据头部包含descrdtype 描述串、fortran_order、shape。非空数组才写数据指针空数组只写头部。load对应的读取实现mlx/io/load.cpp会校验魔数、解析版本、从头部还原 dtype 与 shape并通过swap_endianness read_is_big_endian ! is_big_endian()处理大端文件与当前平台的差异若fortran_order为 True还会在加载后自动做transpose保证返回数组与原始数组等价。这意味着MLX 加载由 NumPy 保存的 .npy 文件、以及 NumPy 加载 MLX 保存的 .npy 文件都可以互通——这是官方测试 python/tests/test_load.py 中交叉验证过的行为。2.3 为什么读取是惰性的load返回的数组并不立即把数据全部读进内存而是构造一个Load原始算子primitive数据真正被使用时例如求值、参与计算才从文件按需读取。为此 MLX 在 mlx/io/load.h 中实现了ParallelFileReader超过batch_size_ 1 25约 32 MiB的大读取会被拆成多个批次提交到共享线程池并行pread线程数按n_cores / 2收敛在 416 之间见 mlx/io/load.cpp。对一个分片sharded大模型所有 Reader 共享同一个磁盘读线程池避免多文件并发读互相放大。这种设计让加载大权重文件时打开即返回、按需取数配合 MLX 的统一内存模型可显著降低峰值内存占用。三、多数组打包mx.savez 与 .npz 格式3.1 基本用法把多个数组保存到一个.npz压缩包import mlx.core as mx a mx.array([1.0]) b mx.array([2.0]) mx.savez(arrays, a, bb) # 生成文件 arrays.npz为兼容numpy.savezMLX 的savez接收位置参数与关键字参数两种形式。未指定关键字时自动生成默认名arr_0、arr_1……加载结果为名字 → 数组的字典mx.load(arrays.npz) # {b: array([2], dtypefloat32), arr_0: array([1], dtypefloat32)}3.2 位置参数与关键字的命名规则Python 绑定实现在 python/src/load.cpp位置参数按顺序命名为arr_0、arr_1……若用户同时用关键字arr_0等与自动命名冲突会抛出Cannot use un-named variables and keyword arr_0异常。.npz本质是一个 ZIP 容器每个数组以名字.npy条目存放加载时会剥掉.npy后缀还原名字这正是 NumPy 兼容的基础。3.3 savez 与 savez_compressed 的区别mx.savez使用zipfile.ZIP_STORED不压缩写入速度快mx.savez_compressed使用zipfile.ZIP_DEFLATEDDEFLATE 压缩文件更小但保存更慢。两者在 python/src/ops.cpp 中注册为两个独立函数底层都走同一个mlx_savez_helper仅compressed标志不同。3.4 实际场景保存整个模型权重.npz非常适合保存模型全部参数。官方 API 文档给出了一个完整的实战范例见 python/src/ops.cppimport mlx.core as mx import mlx.nn as nn from mlx.utils import tree_flatten model nn.TransformerEncoder(6, 128, 4) flat_params tree_flatten(model.parameters()) mx.savez(model.npz, **dict(flat_params))借助mlx.utils.tree_flatten把嵌套的参数字典展平为(name, array)列表再用**解包成关键字参数传给savez每个参数名如layers.0.self_attention.wq就自动成为.npz中的条目名。加载后得到{名字: 数组}字典配合tree_unflatten即可恢复模型状态。四、Hugging Face 生态互操作mx.save_safetensors4.1 基本用法与savez不同save_safetensors接收的是一个dict[str, array]而非位置参数import mlx.core as mx a mx.array([1.0]) b mx.array([2.0]) mx.save_safetensors(arrays, {a: a, b: b})加载mx.load(arrays.safetensors) # {a: array([1], dtypefloat32), b: array([2], dtypefloat32)}4.2 可选元数据参数save_safetensors的第三个参数metadata接收dict[str, str]会写入 safetensors 头部的__metadata__字段mx.save_safetensors( model.safetensors, {weight: mx.ones((4, 4))}, {format: mlx, testing: test}, ) arrays, metadata mx.load(model.safetensors, return_metadataTrue) print(metadata) # {format: mlx, testing: test}加载端通过return_metadataTrue把(arrays_dict, metadata_dict)一并取回对.npy/.npz使用return_metadataTrue会直接报错见 python/src/load.cpp因为这两个格式没有元数据概念。注意元数据值必须是字符串——若传入非字符串如{testing: 0}绑定层会抛出Metadata must be a dictionary with string keys and values见 python/src/load.cpp测试 python/tests/test_load.py 覆盖了该行为。4.3 格式细节与 dtype 映射safetensors 的文件布局是头部uint64长度字段 JSON 头部含__metadata__与每个张量的dtype、shape、data_offsets 连续张量数据区。C 实现在 mlx/io/safetensors.cpp保存前会先把所有数组contiguous化并统一eval同样是为了避免覆盖正在被惰性读取的文件测试见 python/tests/test_load.py然后按序排布数据并记录data_offsets头部 JSON 长度上限与官方 Rust 实现保持一致kMaxJsonHeaderLength 100000000加载时会做多层防御性校验头部长度合法性、data_offsets必须恰好覆盖dtype × shape的字节数、数据区不得越过文件末尾任一不满足都会提示Perhaps an incomplete download or corrupt file?见 mlx/io/safetensors.cppdtype 字符串采用 safetensors 规范F32/F16/BF16/I8/I16/I32/I64/U8/U16/U32/U64/BOOL/C64并额外识别F8_E4M3、F8_E8M0映射为uint8存储见 mlx/io/safetensors.cpp。由于 safetensors 是 Hugging Face 生态的默认权重格式mx.save_safetensors/mx.load让你可以直接在 MLX 与transformers、safetensors等工具之间搬运权重无需中间格式转换。五、LLM 权重格式mx.save_gguf 与 .gguf 文件5.1 基本用法import mlx.core as mx a mx.array([1.0]) b mx.array([2.0]) mx.save_gguf(arrays, {a: a, b: b})5.2 元数据不仅限于字符串与 safetensors 不同save_gguf的元数据值支持三种类型标量/一维array、str、list[str]签名见 python/src/ops.cpp。GGUF 规范中的标量类型uint8/int8/uint16/int16/uint32/int32/uint64/int64/float32/bool/string/float64/array在 mlx/io/gguf.cpp 中逐一映射为 MLX 类型标量数组值要求一维且非空多维数组或空数组会抛出异常见 mlx/io/gguf.cppfloat64元数据在读取时转为float32数组值只支持一层嵌套多层嵌套数组会报Multiple levels of nested arrays are not supported写入前会把非行连续row-contiguous的数组通过reshape(flatten(v), v.shape())规整仍不连续的则报错见 mlx/io/gguf.cpp。加载时用return_metadataTrue同样可以取回元数据arrays, metadata mx.load(model.gguf, return_metadataTrue)5.3 张量 dtype 与维度顺序save_gguf对张量 dtype 的支持是有选择的仅float32、float16、int8、int16、int32五种可直接保存见 mlx/io/gguf.cpp 与测试 python/tests/test_load.py其他 dtype 会报[save_gguf] dtype ... is not supported。写入时维度顺序会做一次反转GGML 规范以ne[0]为最内层维度而 MLX 使用行主序因此保存时dim[i] arr.shape()[num_dim - 1 - i]加载时再反转回来get_shape见 mlx/io/gguf.cpp。这是 MLX 与 GGML 生态互操作的关键细节自行拼接 GGUF 时最容易在这里出错。5.4 量化张量的自动反量化GGUF 中常见的 Q4_0/Q4_1/Q8_0 等量化张量加载时会通过gguf_load_quantized解包为(weights, scales, biases)三元组见 mlx/io/gguf.cpp而其他暂不支持的量化格式如 Q5、Q6、Q2_K、Q3_K 等在读取时会自动转为mx.float16这是官方文档与 API 文档都明确标注的警告行为见 python/src/ops.cpp。也就是说从llama.cpp生态下载的 GGUF 模型即使量化格式不被 MLX 原生支持也能先以 fp16 加载出来继续使用。测试 python/tests/test_load.py 展示了手工构造最小 GGUF v3 文件并校验 Q4 反量化数值的方法可作为理解该流程的参考。5.5 其他限制save_gguf目前只接受文件路径不接受 file-like 对象python/src/load.cpp且 Windows 平台不启用 GGUF 支持测试中以unittest.skipIf(platform.system() Windows, ...)跳过写入文件时使用GGUF_OVERWRITE模式创建若打开失败会抛出[save_gguf] gguf_create failed。六、统一加载入口 mx.load 的完整签名mx.load的完整 Python 签名见 python/src/ops.cpp为mx.load( file, # file | str | pathlib.Path formatNone, # 显式指定格式默认按扩展名推断 return_metadataFalse, # 是否返回元数据 *, streamNone, # 指定加载到哪个流/设备如 mx.cpu )6.1 格式推断与显式指定格式推断逻辑在 python/src/load.cpp取文件名最后一个.之后的子串作为扩展名npy/npz/safetensors/gguf找不到扩展名时抛出[load] Could not infer file format from extension。当文件对象没有name属性或文件名不带扩展名时可以显式传formatnpy或npz/safetensors/gguf绕过推断传入未知格式则报[load] Unknown file format ...。6.2 返回值取决于格式格式load返回值return_metadataTrue时.npy单个array不支持报错.npzdict[str, array]不支持报错.safetensorsdict[str, array](dict, metadata).ggufdict[str, array](dict, metadata)6.3 加载到指定设备stream 参数load支持可选的stream关键字把数组直接加载到指定流或设备。例如测试中常用的mx.load(save_file, streammx.cpu)见 python/tests/test_load.py。在 C 层当 CUDA 后端可用时会用传入的 stream否则回退到 CPU stream见 mlx/io/safetensors.cpp 与 mlx/io/load.cpp。这意味着在多设备环境中你可以控制大权重直接落在目标设备上避免先加载到 CPU 再搬运。6.4 文件对象支持load、save、savez、save_safetensors都支持二进制模式打开的文件对象。绑定层通过is_istream_object检测readinto/seek/tell/closed与is_ostream_object检测write/seek/tell/closed来识别 file-like 对象见 python/src/load.cpp并包装为PyFileReader/PyFileWriter供 C 读写。典型用法with open(model.safetensors, wb) as f: mx.save_safetensors(f, {w: mx.ones((4, 4))}) with open(model.safetensors, rb) as f: weights mx.load(f)对文件对象执行 load 时绑定层会立即eval数组因为对象生命周期由 Python 侧管理见 python/src/load.cpp而对路径加载则是惰性的可在后续按需读取。此外.npz的加载还会借助 Pythonzipfile模块判断文件是否为合法 ZIP 容器见 python/src/load.cpp。七、格式选型建议单个数组快速落盘 / 与 NumPy 互通选mx.save.npy。体积最小、零依赖、双向兼容 NumPy多个参数打包、需要人类可读的名字选mx.savez不压缩快或mx.savez_compressedDEFLATE 压缩省空间配合tree_flatten一键保存整个模型与 Hugging Face 生态对接、需要元数据、追求零拷贝加载选mx.save_safetensors。其头部含长度与偏移信息非常适合按需读取指定张量的分片加载场景LLM 权重 / llama.cpp 生态互操作选mx.save_gguf.gguf并注意它只接受路径、dtype 支持有限、量化张量加载时可能自动转 fp16 等前提限制。八、可继续深入阅读的仓库入口官方文档docs/src/usage/saving_and_loading.rstNPY 读写实现mlx/io/load.cpp文件 IO 抽象Reader/Writer/并行读取mlx/io/load.hSafetensors 实现mlx/io/safetensors.cppGGUF 实现mlx/io/gguf.cppPython 绑定python/src/load.cpp、python/src/ops.cpp测试用例python/tests/test_load.py相关主题保存/加载模型状态与 checkpoints 之外还可参考 docs/src/usage/export.rst模型导出、docs/src/usage/saving_and_loading.rst 同目录下的 docs/src/usage/quick_start.rst 与 docs/src/usage/function_transforms.rst 了解 MLX 的求值模型与变换机制。结语MLX 的序列化 API 设计高度对齐 NumPy 与深度学习生态惯例.npy/.npz保证与 NumPy 双向兼容.safetensors与 Hugging Face 权重无缝衔接.gguf打通 LLM 推理工具链而统一的mx.load按扩展名自动路由、支持惰性加载与指定设备加载。理解这些格式的底层布局与约束能帮助你在实际项目中选对 API、规避坑点写出更稳健的模型存取代码。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表