mlx-weightlifter 0.1.1

Lift any checkpoint into MLX - PyTorch, Safetensors, GGUF

mlx-weightlifter

Lift any checkpoint into MLX

---

Universal weight loading library for MLX. Load checkpoints from any format directly into your nn.Module without framework dependencies.

Supported Formats

Format	Extensions	Features
GGUF	.gguf	Auto architecture detection, MoE expansion, built-in tokenizer
Safetensors	.safetensors	HuggingFace shard support, index files
PyTorch	.pt, .bin, .jit	Custom unpickler (no torch import), auto-cache to safetensors

Installation

pip install mlx-weightlifter

Or from source:

git clone https://github.com/MLXPorts/mlx-weightlifter.git
cd mlx-weightlifter
pip install -e .

Quick Start

Load a GGUF Model

from mlx_weightlifter import load_gguf_with_mlx

# Auto-detects architecture (qwen3moe, qwen3, gemma3, llama)
model, tokenizer = load_gguf_with_mlx("model.gguf")

# Use it
tokens = tokenizer.encode("Hello, world!")

Load Weights into nn.Module

from mlx_weightlifter import load_weights_from_safetensors
import mlx.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(512, 512) for _ in range(12)]

model = MyModel()

# Load with automatic key mapping and dtype conversion
load_weights_from_safetensors(
    model,
    "checkpoint_dir",
    target_dtype=mx.bfloat16,
    strip_prefix="model."
)

Load PyTorch Checkpoints (No Torch Required)

from mlx_weightlifter import load_pytorch_bin

# Loads .pt/.bin files without importing torch
weights = load_pytorch_bin("pytorch_model.bin")

# Auto-caches to safetensors for faster subsequent loads

Features

GGUF Loading

• Architecture Detection: Automatically identifies model architecture from metadata
• Weight Mapping: Translates GGUF tensor names to mlx-lm conventions
• MoE Expansion: Expands consolidated expert weights to per-expert tensors
• Tokenizer Extraction: Builds BPE tokenizer directly from GGUF metadata

Supported architectures:

• qwen3moe - Qwen3 Mixture of Experts
• qwen3 - Qwen3 dense models
• gemma3 - Gemma 3 models
• llama - LLaMA-style models

nn.Module Weight Loading

from mlx_weightlifter import load_weights_from_dict, get_parameter_dict

# Extract parameters as flat dict
params = get_parameter_dict(model)
# {'layers.0.weight': array(...), 'layers.0.bias': array(...), ...}

# Load with custom key mapping
def hf_to_local(key):
    return key.replace("transformer.h.", "layers.")

load_weights_from_dict(
    model,
    weights,
    key_mapper=hf_to_local,
    target_dtype=mx.bfloat16,
    strict=False  # Allow partial loading
)

Config Inference

from mlx_weightlifter import infer_config_from_checkpoint, infer_config_from_safetensors

# Infer model dimensions from weight shapes
config = infer_config_from_safetensors("model_dir")
print(config)
# {'hidden_size': 4096, 'num_layers': 32, 'num_heads': 32, ...}

API Reference

Format Loaders

Function	Description
load_gguf_with_mlx(path)	Load GGUF model + tokenizer
load_pytorch_bin(path)	Load PyTorch .pt/.bin file
load_jit(path)	Load TorchScript .jit file
load_safetensors_weights(path)	Load single safetensors file

nn.Module Loading

Function	Description
load_weights_from_safetensors(module, dir)	Load from HF checkpoint directory
load_weights_from_dict(module, weights)	Load from in-memory dict
get_parameter_dict(module)	Extract flat parameter dict
set_parameter(module, path, value)	Set single parameter by path

Utilities

Function	Description
resolve_dtype(name)	Convert dtype string to mx.Dtype
load_config(path)	Load JSON config file
compute_derived_dims(config)	Compute derived dimensions

Requirements

• Python 3.10+
• MLX 0.10+
• mlx-lm (optional, for GGUF model classes)

License

MIT License - see LICENSE for details.

Contributing

Contributions welcome! Please open an issue or PR on GitHub.