DOCSEARCHbyAlgolia
Submit Docs

MLX LM

repository·main·Indexed 27 days ago

https://github.com/ml-explore/mlx-lm

A Python package for generating text and fine-tuning large language models on Apple silicon using the MLX framework. It provides tools for model discovery, quantization (including DWQ, AWQ, GPTQ, and dynamic quantization), and LoRA/QLoRA fine-tuning with seamless Hugging Face Hub integration. Features include a CLI for chatting and generation, prompt caching for long contexts, and utilities for evaluating model perplexity and MMLU Pro benchmarks.

Tokens
7.5K
Snippets
31
Records
49
Agent score
92%

What's inside mlx-lm

  1. Start the MLX LM HTTP Model Server

    main

    You can use mlx-lm to create an HTTP API for text generation that is compatible with the OpenAI chat API. Note that this server is not recommended for production due to basic security implementations.

    To start the server, use the mlx_lm.server command and provide a path to a local model or a Hugging Face repository ID using the --model flag. The server defaults to running on localhost:8080.

    mlx_lm.server --model mlx-community/Mistral-7B-Instruct-v0.3-4bit

  6. Use prompt caching for long contexts

    main

    To speed up repeated queries with the same long context, use mlx_lm.cache_prompt to save a prompt to a .safetensors file. You can then use this file with mlx_lm.generate to treat the cached content as a prefix.

    # Cache a prompt
cat prompt.txt | mlx_lm.cache_prompt \
  --model mistralai/Mistral-7B-Instruct-v0.3 \
  --prompt - \
  --prompt-cache-file mistral_prompt.safetensors

# Use the cached prompt
mlx_lm.generate \
    --prompt-cache-file mistral_prompt.safetensors \
    --prompt "\nSummarize the above text."

  7. Configure Hugging Face datasets via YAML

    main

    To train on Hugging Face datasets, install datasets (pip install datasets) and use a YAML configuration to map dataset keys to MLX LM features.

    Key configuration options under hf_dataset:

    • path: The Hugging Face dataset ID.
    • prompt_feature: Key for the prompt (for completions).
    • completion_feature: Key for the completion (for completions).
    • text_feature: Key for raw text (for text).
    • chat_feature: Key for chat messages (for chat).
    • {train,valid,test}_split: Specify the split (e.g., train[:90%]).

    You can provide a list of multiple datasets in the config.

    hf_dataset:
  - path: "Open-Orca/OpenOrca"
    train_split: "train[:90%]"
    valid_split: "train[-10%]"
    prompt_feature: "question"
    completion_feature: "response"
  - path: "trl-lib/ultrafeedback_binarized"
    train_split: "train[:90%]"
    valid_split: "train[-10%]"
    chat_feature: "chosen"

  8. Start the MLX LM HTTP Server

    main

    You can run the MLX LM server as a CLI tool to provide an OpenAI-compatible API for your models. The server supports chat completions, text completions, and model listing.

    Note: It is recommended to use python -m mlx_lm.server or the mlx_lm.server entry point rather than calling the module directly.

  9. Handle large models on macOS

    main

    When running models larger than available RAM, performance may be slow. On macOS 15.0 or higher, you can attempt to speed up generation by increasing the system wired memory limit using sysctl. The value N should be larger than the model size in MB but smaller than the total machine memory.

    sudo sysctl iogpu.wired_limit_mb=N

  10. Reduce memory usage during LoRA fine-tuning

    main

    If you encounter memory issues while fine-tuning a large model with LoRA, use the following strategies to reduce consumption:

    1. Use QLoRA: Generate a quantized model using convert.py with the -q flag before fine-tuning.
    2. Adjust Batch Size: Decrease --batch-size (default is 4) to 2 or 1. To maintain effective batch size, use --grad-accumulation-steps <N> to accumulate gradients over <N> batches before updating parameters.
    3. Limit Fine-tuned Layers: Reduce the number of layers targeted for fine-tuning using --num-layers. The default is 16; try 8 or 4 to reduce memory needed for backpropagation (note: this may impact model quality).
    4. Shorten Sequences: Break examples into smaller sequences when preparing your {train, valid, test}.jsonl files.
    5. Enable Gradient Checkpointing: Use the --grad-checkpoint flag to trade computation time for reduced memory usage by recomputing intermediate values instead of storing them.

  11. Run a memory-efficient LoRA fine-tuning command

    main

    For machines with limited memory (e.g., 32 GB), use a reduced batch size and fewer layers to ensure the training process runs reasonably fast.

    mlx_lm.lora \
    --model mistralai/Mistral-7B-v0.1 \
    --train \
    --batch-size 1 \
    --num-layers 4 \
    --data  mlx-community/wikisql

  12. Make a chat completion request to the MLX LM server

    main

    Once the server is running, you can send requests to the /v1/chat/completions endpoint. The request body should be a JSON object containing a messages array, where each object has a role (e.g., user, assistant) and content.

    curl localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{ "messages": [{"role": "user", "content": "Say this is a test!"}], "temperature": 0.7 }'