mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
470 lines
14 KiB
Python
470 lines
14 KiB
Python
"""Optional MLX backend for Apple Silicon native inference and weight editing.
|
|
|
|
MLX is Apple's array framework that runs natively on the Apple Neural Engine
|
|
and Metal GPU. When available, it provides significantly faster inference and
|
|
weight manipulation than PyTorch's MPS backend on Apple hardware.
|
|
|
|
This module is entirely optional — if ``mlx`` / ``mlx-lm`` are not installed,
|
|
``MLX_AVAILABLE`` is ``False`` and all public functions raise ``RuntimeError``.
|
|
|
|
Install with::
|
|
|
|
pip install mlx>=0.22 mlx-lm>=0.20
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Availability check
|
|
# ---------------------------------------------------------------------------
|
|
|
|
MLX_AVAILABLE = False
|
|
_mx = None # mlx module
|
|
_mlx_lm = None # mlx-lm module
|
|
_mlx_nn = None # mlx.nn module
|
|
|
|
try:
|
|
import mlx.core as _mx_core # type: ignore[import-untyped]
|
|
import mlx.nn as _mlx_nn_mod # type: ignore[import-untyped]
|
|
import mlx_lm # type: ignore[import-untyped]
|
|
|
|
_mx = _mx_core
|
|
_mlx_nn = _mlx_nn_mod
|
|
_mlx_lm = mlx_lm
|
|
MLX_AVAILABLE = True
|
|
logger.info("MLX backend available (mlx %s)", _mx.__version__ if hasattr(_mx, "__version__") else "?")
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def _require_mlx() -> None:
|
|
if not MLX_AVAILABLE:
|
|
raise RuntimeError(
|
|
"MLX backend is not available. Install with: pip install mlx>=0.22 mlx-lm>=0.20"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MLXModelHandle:
|
|
"""Lightweight wrapper around an MLX-loaded model + tokenizer."""
|
|
|
|
def __init__(self, model: Any, tokenizer: Any, model_name: str):
|
|
self.model = model
|
|
self.tokenizer = tokenizer
|
|
self.model_name = model_name
|
|
|
|
@property
|
|
def config(self) -> Any:
|
|
return getattr(self.model, "config", None)
|
|
|
|
|
|
def load_model(
|
|
model_name: str,
|
|
dtype: str = "float16",
|
|
) -> MLXModelHandle:
|
|
"""Load a HuggingFace model via ``mlx-lm`` for Apple-native execution.
|
|
|
|
Parameters
|
|
----------
|
|
model_name : str
|
|
HuggingFace model identifier (e.g. ``"meta-llama/Llama-3.2-3B-Instruct"``).
|
|
dtype : str
|
|
One of ``"float16"``, ``"bfloat16"``, ``"float32"``.
|
|
|
|
Returns
|
|
-------
|
|
MLXModelHandle
|
|
Wrapper with ``.model`` and ``.tokenizer`` attributes.
|
|
"""
|
|
_require_mlx()
|
|
|
|
from mlx_lm import load # type: ignore[import-untyped]
|
|
|
|
logger.info("Loading %s via MLX (dtype=%s)", model_name, dtype)
|
|
model, tokenizer = load(model_name)
|
|
|
|
return MLXModelHandle(model=model, tokenizer=tokenizer, model_name=model_name)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Inference
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def generate(
|
|
handle: MLXModelHandle,
|
|
prompt: str,
|
|
max_tokens: int = 256,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
repetition_penalty: float | None = None,
|
|
) -> str:
|
|
"""Generate text using the MLX model.
|
|
|
|
Parameters
|
|
----------
|
|
handle : MLXModelHandle
|
|
A loaded MLX model handle.
|
|
prompt : str
|
|
The input prompt string.
|
|
max_tokens : int
|
|
Maximum number of tokens to generate.
|
|
temperature : float
|
|
Sampling temperature.
|
|
top_p : float
|
|
Nucleus sampling threshold.
|
|
repetition_penalty : float or None
|
|
Repetition penalty (1.0 = no penalty).
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Generated text completion.
|
|
"""
|
|
_require_mlx()
|
|
|
|
from mlx_lm import generate as mlx_generate # type: ignore[import-untyped]
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"max_tokens": max_tokens,
|
|
"temp": temperature,
|
|
"top_p": top_p,
|
|
}
|
|
if repetition_penalty is not None:
|
|
kwargs["repetition_penalty"] = repetition_penalty
|
|
|
|
return mlx_generate(
|
|
handle.model,
|
|
handle.tokenizer,
|
|
prompt=prompt,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Activation extraction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def get_activations(
|
|
handle: MLXModelHandle,
|
|
prompts: list[str],
|
|
layer_indices: list[int],
|
|
max_length: int = 256,
|
|
) -> dict[int, list[Any]]:
|
|
"""Extract hidden-state activations from specified layers.
|
|
|
|
Uses MLX's computation graph to capture intermediate outputs.
|
|
|
|
Parameters
|
|
----------
|
|
handle : MLXModelHandle
|
|
Loaded model.
|
|
prompts : list[str]
|
|
Input prompts to probe.
|
|
layer_indices : list[int]
|
|
Which transformer layers to capture.
|
|
max_length : int
|
|
Maximum sequence length for tokenization.
|
|
|
|
Returns
|
|
-------
|
|
dict[int, list[mlx.core.array]]
|
|
Mapping from layer index to list of activation arrays (one per prompt).
|
|
Each array has shape ``(hidden_dim,)`` — the last-token hidden state.
|
|
"""
|
|
_require_mlx()
|
|
import mlx.core as mx # type: ignore[import-untyped]
|
|
|
|
model = handle.model
|
|
tokenizer = handle.tokenizer
|
|
|
|
# Identify the transformer block list
|
|
layers = None
|
|
for attr in ("model.layers", "transformer.h", "gpt_neox.layers"):
|
|
obj = model
|
|
try:
|
|
for part in attr.split("."):
|
|
obj = getattr(obj, part)
|
|
layers = obj
|
|
break
|
|
except AttributeError:
|
|
continue
|
|
|
|
if layers is None:
|
|
raise RuntimeError(
|
|
"Cannot locate transformer layers in the MLX model. "
|
|
"Supported architectures: LLaMA, GPT-2, GPT-NeoX."
|
|
)
|
|
|
|
activations: dict[int, list[Any]] = {idx: [] for idx in layer_indices}
|
|
target_set = set(layer_indices)
|
|
|
|
for prompt in prompts:
|
|
tokens = tokenizer.encode(prompt)
|
|
if len(tokens) > max_length:
|
|
tokens = tokens[:max_length]
|
|
|
|
input_ids = mx.array([tokens])
|
|
|
|
# Forward through embedding
|
|
if hasattr(model, "model"):
|
|
# LLaMA-style: model.model.embed_tokens
|
|
embed_module = model.model
|
|
elif hasattr(model, "transformer"):
|
|
embed_module = model.transformer
|
|
else:
|
|
embed_module = model
|
|
|
|
if hasattr(embed_module, "embed_tokens"):
|
|
h = embed_module.embed_tokens(input_ids)
|
|
elif hasattr(embed_module, "wte"):
|
|
h = embed_module.wte(input_ids)
|
|
else:
|
|
raise RuntimeError("Cannot find embedding layer in MLX model")
|
|
|
|
# Walk through layers, capturing activations at target indices
|
|
for i, layer in enumerate(layers):
|
|
h = layer(h)
|
|
# Some layers return tuples (hidden, attention) — take first
|
|
if isinstance(h, tuple):
|
|
h = h[0]
|
|
|
|
if i in target_set:
|
|
# Last token hidden state
|
|
last_hidden = h[0, -1, :]
|
|
mx.eval(last_hidden) # force evaluation
|
|
activations[i].append(last_hidden)
|
|
|
|
return activations
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Weight manipulation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def get_weight(handle: MLXModelHandle, layer_idx: int, param_path: str) -> Any:
|
|
"""Retrieve a weight matrix from the model.
|
|
|
|
Parameters
|
|
----------
|
|
handle : MLXModelHandle
|
|
Loaded model.
|
|
layer_idx : int
|
|
Transformer layer index.
|
|
param_path : str
|
|
Dot-separated path within the layer, e.g. ``"self_attn.o_proj.weight"``.
|
|
|
|
Returns
|
|
-------
|
|
mlx.core.array
|
|
The weight tensor.
|
|
"""
|
|
_require_mlx()
|
|
model = handle.model
|
|
|
|
# Navigate to the layer
|
|
layers = _get_layers(model)
|
|
layer = layers[layer_idx]
|
|
|
|
# Navigate the param path
|
|
obj = layer
|
|
for part in param_path.split("."):
|
|
obj = getattr(obj, part)
|
|
|
|
return obj
|
|
|
|
|
|
def modify_weights(
|
|
handle: MLXModelHandle,
|
|
layer_idx: int,
|
|
param_path: str,
|
|
modifier_fn: Callable[[Any], Any],
|
|
) -> None:
|
|
"""Modify a weight matrix in-place using a function.
|
|
|
|
Parameters
|
|
----------
|
|
handle : MLXModelHandle
|
|
Loaded model.
|
|
layer_idx : int
|
|
Transformer layer index.
|
|
param_path : str
|
|
Dot-separated path within the layer to the weight, e.g.
|
|
``"self_attn.o_proj.weight"``.
|
|
modifier_fn : callable
|
|
Function that takes the current weight (mlx array) and returns the
|
|
modified weight (mlx array). For abliteration, this would project
|
|
out the refusal direction.
|
|
"""
|
|
_require_mlx()
|
|
import mlx.core as mx # type: ignore[import-untyped]
|
|
|
|
model = handle.model
|
|
layers = _get_layers(model)
|
|
layer = layers[layer_idx]
|
|
|
|
# Navigate to the parent module and leaf attribute
|
|
parts = param_path.split(".")
|
|
parent = layer
|
|
for part in parts[:-1]:
|
|
parent = getattr(parent, part)
|
|
leaf_name = parts[-1]
|
|
|
|
old_weight = getattr(parent, leaf_name)
|
|
new_weight = modifier_fn(old_weight)
|
|
|
|
# MLX uses a functional update pattern
|
|
if hasattr(parent, "update"):
|
|
parent.update({leaf_name: new_weight})
|
|
else:
|
|
setattr(parent, leaf_name, new_weight)
|
|
|
|
mx.eval(new_weight) # materialize
|
|
|
|
|
|
def project_out_direction(weight: Any, direction: Any) -> Any:
|
|
"""Project a direction out of a weight matrix (abliteration).
|
|
|
|
Given weight matrix W and unit direction d, computes::
|
|
|
|
W' = W - (W @ d) outer d
|
|
|
|
Parameters
|
|
----------
|
|
weight : mlx.core.array
|
|
Weight matrix, shape ``(out_features, in_features)``.
|
|
direction : mlx.core.array
|
|
Unit direction vector, shape ``(in_features,)``.
|
|
|
|
Returns
|
|
-------
|
|
mlx.core.array
|
|
Modified weight with direction projected out.
|
|
"""
|
|
_require_mlx()
|
|
import mlx.core as mx # type: ignore[import-untyped]
|
|
|
|
d = direction.astype(weight.dtype)
|
|
# W @ d gives the component along d for each row
|
|
proj = mx.matmul(weight, d[:, None]) # (out, 1)
|
|
return weight - mx.matmul(proj, d[None, :]) # (out, in)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Save model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def save_model(
|
|
handle: MLXModelHandle,
|
|
output_dir: str | Path,
|
|
upload_repo: str | None = None,
|
|
) -> Path:
|
|
"""Save the (modified) MLX model to disk.
|
|
|
|
Saves in safetensors format compatible with both MLX and HuggingFace.
|
|
|
|
Parameters
|
|
----------
|
|
handle : MLXModelHandle
|
|
Model handle (possibly with modified weights).
|
|
output_dir : str or Path
|
|
Directory to save into.
|
|
upload_repo : str or None
|
|
If set, also uploads to HuggingFace Hub.
|
|
|
|
Returns
|
|
-------
|
|
Path
|
|
The output directory.
|
|
"""
|
|
_require_mlx()
|
|
|
|
from mlx_lm import convert # type: ignore[import-untyped]
|
|
|
|
out = Path(output_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
|
|
# mlx-lm's save uses safetensors
|
|
if hasattr(_mlx_lm, "save_model"):
|
|
_mlx_lm.save_model(str(out), handle.model, handle.tokenizer)
|
|
else:
|
|
# Fallback: manual save via mlx.core.save_safetensors
|
|
import mlx.core as mx # type: ignore[import-untyped]
|
|
weights = dict(handle.model.parameters())
|
|
flat = {}
|
|
_flatten_dict(weights, "", flat)
|
|
mx.save_safetensors(str(out / "model.safetensors"), flat)
|
|
# Save tokenizer via transformers
|
|
handle.tokenizer.save_pretrained(str(out))
|
|
|
|
logger.info("MLX model saved to %s", out)
|
|
|
|
if upload_repo:
|
|
try:
|
|
from mlx_lm import upload_to_hub # type: ignore[import-untyped]
|
|
upload_to_hub(str(out), upload_repo)
|
|
logger.info("Uploaded to %s", upload_repo)
|
|
except (ImportError, AttributeError):
|
|
logger.warning("mlx-lm upload not available — push manually with huggingface-cli")
|
|
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conversion: PyTorch ↔ MLX
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def torch_tensor_to_mlx(tensor: "torch.Tensor") -> Any: # noqa: F821
|
|
"""Convert a PyTorch tensor to an MLX array."""
|
|
_require_mlx()
|
|
import mlx.core as mx # type: ignore[import-untyped]
|
|
import numpy as np
|
|
|
|
# Move to CPU and convert via numpy
|
|
np_array = tensor.detach().cpu().float().numpy()
|
|
return mx.array(np_array)
|
|
|
|
|
|
def mlx_to_torch_tensor(array: Any, device: str = "cpu") -> "torch.Tensor": # noqa: F821
|
|
"""Convert an MLX array to a PyTorch tensor."""
|
|
import numpy as np
|
|
import torch
|
|
|
|
np_array = np.array(array)
|
|
return torch.from_numpy(np_array).to(device)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internal helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _get_layers(model: Any) -> Any:
|
|
"""Locate the transformer block list in an MLX model."""
|
|
for attr_path in ("model.layers", "transformer.h", "gpt_neox.layers"):
|
|
obj = model
|
|
try:
|
|
for part in attr_path.split("."):
|
|
obj = getattr(obj, part)
|
|
return obj
|
|
except AttributeError:
|
|
continue
|
|
raise RuntimeError("Cannot locate transformer layers in MLX model")
|
|
|
|
|
|
def _flatten_dict(d: dict, prefix: str, out: dict) -> None:
|
|
"""Flatten a nested dict with dot-separated keys."""
|
|
for k, v in d.items():
|
|
key = f"{prefix}{k}" if prefix else k
|
|
if isinstance(v, dict):
|
|
_flatten_dict(v, f"{key}.", out)
|
|
else:
|
|
out[key] = v
|