DOCSEARCHbyAlgolia
Submit Docs

MLX

repository·main·Indexed 12 days ago

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

An array framework for machine learning optimized for Apple silicon, featuring a unified memory model and lazy computation. It provides Python and C++ interfaces, including the JACCL library for distributed computing over RDMA and Thunderbolt.

Tokens
50K
Snippets
160
Records
195
Agent score
96%

What's inside MLX

  1. Use learning rate schedulers in mlx.optimizers

    main

    MLX provides several scheduler functions within the mlx.optimizers module to dynamically adjust the learning rate during training. These schedulers can be used to implement common decay strategies such as linear, exponential, or cosine decay, or to combine multiple schedules using join_schedules.

    import mlx.optimizers as optim

# Example of available scheduler types:
# - cosine_decay
# - exponential_decay
# - join_schedules
# - linear_schedule
# - step_decay

  2. Use neural network layers from mlx.nn

    main

    The mlx.nn module provides a comprehensive collection of neural network layers and modules for building machine learning models. These layers include standard linear transformations, convolutional layers, normalization techniques, activation functions, and complex architectures like Transformers. Most layers are designed to be used within a Sequential container or as components of a custom class inheriting from mlx.nn.Module.

    import mlx.nn as nn

# Example of using standard layers
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

  3. Overview of MLX Distributed Communication Backends

    main

    MLX supports distributed communication to share the computational cost of training or inference across multiple physical machines. It provides several backends depending on your environment and hardware requirements:

    • MPI: A mature, full-featured distributed communications library.
    • RING: Uses TCP sockets for ring all-reduce and all-gather operations. It is always available and typically faster than MPI.
    • JACCL: Provides low-latency communication with RDMA over Thunderbolt, specifically useful for tensor parallelism.
    • NCCL: The preferred backend for CUDA environments.

  5. Use `mx.compile` to optimize computation graphs

    main

    The mx.compile function transforms a function by compiling its computation graph. This process merges common work and fuses operations, which can significantly improve runtime performance and reduce memory usage.

    Key behaviors:

    • First call overhead: The first time a compiled function is called, MLX builds the graph, optimizes it, and generates code. This is relatively slow.
    • Caching: MLX caches compiled functions. Subsequent calls with the same input characteristics will use the cached version without re-compiling.
    • Re-compilation triggers: A function will be re-compiled if you change:
      • The shape or number of dimensions of inputs.
      • The data type (dtype) of any input.
      • The number of inputs to the function.

    Best Practice: Avoid compiling functions inside loops (e.g., compiling a lambda inside a for loop), as this creates a new function object and triggers compilation at every iteration.

    def fun(x, y):
    return mx.exp(-x) + y

x = mx.array(1.0)
y = mx.array(2.0)

# Regular call
print(fun(x, y))

# Compile the function for reuse
compiled_fun = mx.compile(fun)
print(compiled_fun(x, y))

  6. Configure the NCCL backend for CUDA

    main

    When launching from a Mac to a Linux machine with CUDA, use --backend nccl.

    To launch multi-node and multi-GPU jobs, use the --repeat-hosts (or -n) argument. This specifies how many processes to launch per host.

    # Launch 16 processes total (8 processes on each of the 2 nodes)
mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 -- ./my-job.sh

  7. Use the mlx.core.array object

    main

    The mlx.core.array is the fundamental data structure in MLX. It represents a multi-dimensional array of elements. You can create arrays, inspect their properties (like shape, dtype, and ndim), and perform a wide range of mathematical, statistical, and transformation operations directly on them. Most operations in MLX are designed to work efficiently on these array objects.

    import mlx.core as mx

# Example of creating and inspecting an array
a = mx.array([1, 2, 3])
print(a.shape)  # (3,)
print(a.dtype)  # float32 (default depending on input)
print(a.ndim)   # 1

  8. Core concepts of MLX

    main

    MLX is an array framework designed for machine learning on Apple silicon. Key architectural concepts include:

    • Unified Memory Model: Arrays live in shared memory. Operations can be performed on different device types (CPU/GPU) without explicit data transfers.
    • Lazy Computation: Computations are deferred; arrays are only materialized when their values are actually needed.
    • Dynamic Graph Construction: Computation graphs are built dynamically at runtime. Changing argument shapes does not trigger slow recompilations.
    • Composable Transformations: Supports automatic differentiation, automatic vectorization, and computation graph optimization.
    • Familiar APIs: The Python API follows NumPy, while higher-level packages like mlx.nn and mlx.optimizers follow PyTorch patterns.

  9. Avoid shape-dependent errors in shapeless compiled functions

    main

    When using mx.compile(..., shapeless=True), avoid using reshape with logic that relies on the specific dimensions of the first input encountered. This is because the compiled graph may capture those specific dimensions as constants.

    Instead of using reshape to manipulate dimensions based on x.shape, use flatten to maintain shape-agnostic behavior. This allows the compiled function to work correctly across different input shapes.

    import mlx.core as mx

# BAD: reshape uses static shape logic that fails on new shapes
def bad_fun(x):
    return x.reshape(x.shape[0] * x.shape[1], -1)

# GOOD: flatten is shape-agnostic and works with shapeless=True
def good_fun(x):
    return x.flatten(0, 1)

compiled_fun = mx.compile(good_fun, shapeless=True)

# Works for multiple shapes
x1 = mx.random.uniform(shape=(2, 3, 4))
print(compiled_fun(x1))

x2 = mx.random.uniform(shape=(5, 5, 3))
print(compiled_fun(x2))

  10. How lazy evaluation works in MLX

    main

    In MLX, performing operations does not trigger immediate computation. Instead, MLX records a compute graph. The actual computation is deferred until an explicit mx.eval is called or until the data is implicitly required by an external operation.

    Benefits of Lazy Evaluation

    • Transformations: Enables function transformations like grad and vmap and graph optimizations by recording the graph without executing it.
    • Efficiency: Avoids computing outputs that are never used. For example, if a function returns (a, b) but you only use a, the graph for b is built but its computation is never triggered.
    • Memory Management: Model initialization (e.g., model = Model()) does not consume significant memory for weights until mx.eval is called. This allows you to load weights in a lower precision (like float16) before the memory is actually allocated for the computation.
    def fun(x):
    a = fun1(x)
    b = expensive_fun(a)
    return a, b

y, _ = fun(x)  # expensive_fun is never actually computed

  11. Compose transformations with `mx.compile`

    main

    MLX function transformations (like mx.grad) are composable. You can apply mx.compile to the result of a transformation.

    Important Note: To maximize optimization, apply mx.compile to the outermost function. While you can compile a transformed function, it is often more efficient to compile the function that calls the transformed logic to allow the compiler to see the largest possible computation graph.

    # Compiling a transformed function
grad_fn = mx.grad(mx.exp)
compiled_grad_fn = mx.compile(grad_fn)

# Best practice: Compile the outermost function
@mx.compile
def inner(x):
    return mx.exp(-mx.abs(x))

def outer(x):
    return inner(inner(x))

fun = mx.compile(outer)

  12. Define a custom neural network module with `nn.Module`

    main

    To create a custom neural network in MLX, inherit from mlx.nn.Module. You must follow two main steps:

    1. __init__: Initialize parameters and submodules (like nn.Linear). MLX automatically registers parameters assigned to self during initialization.
    2. __call__: Implement the forward pass logic. This method defines how input data flows through the layers.

    Example implementation of a Multi-Layer Perceptron (MLP) with ReLU activations:

    import mlx.core as mx
import mlx.nn as nn

class MLP(nn.Module):
    def __init__(self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
        self.layers = [
            nn.Linear(idim, odim)
            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
        ]

    def __call__(self, x):
        for l in self.layers[:-1]:
            x = mx.maximum(l(x), 0.0)  # ReLU activation
        return self.layers[-1](x)