Runtime Model Surgery¶

Runtime model surgery applies YAML-defined weight operations while a checkpoint loads, changing model behavior without retraining the model or writing a converted checkpoint back to disk. The main user-facing entry point is --surgery <FILE>, which loads a YAML pipeline and applies ordered weight operations to the in-memory WeightMap before the model graph is constructed.

Use this when you want to test a controlled weight-space change quickly: scaling selected tensors, applying a task vector, pruning structural components, replacing tensors from a donor checkpoint, or interpolating between donor checkpoints.

Table of Contents¶

• When to Use It
• How It Runs
• YAML Schema
• Operations
• Examples
• Interactions with Other Features
• Limitations
• Troubleshooting

When to Use It¶

Use load-time surgery when the change is naturally expressed as a tensor operation:

• Model arithmetic: add a task vector to selected weights.
• Model merging: interpolate between two donor checkpoints with LERP or SLERP.
• Ablation: zero selected layers, attention heads, or MLP channels.
• Targeted patching: replace one or more tensors with tensors from another checkpoint.
• Gain control: scale selected projections or normalization parameters.

This is different from prompt engineering, language bias, or structured output. Those features operate at prompt or decode time. Model surgery changes the loaded weights themselves, so every later request uses the modified model until the process exits or the server is restarted.

How It Runs¶

--surgery is accepted by mlxcel generate and mlxcel serve. The same value can also be provided through MLXCEL_SURGERY.

mlxcel generate \
  --model ./models/qwen3-4b-4bit \
  --surgery ./surgery.yaml \
  --prompt "Explain this model change in one paragraph."

MLXCEL_SURGERY=./surgery.yaml \
mlxcel serve \
  --model ./models/qwen3-4b-4bit \
  --host 0.0.0.0 \
  --port 8080

The surgery pipeline is parsed before model loading starts. Missing files, malformed YAML, invalid glob patterns, and missing donor files fail fast, before large model weights are loaded.

For text models, the hook runs after basic weight sanitization such as tied-embedding restoration and quantized checkpoint normalization, and before Apple Silicon bf16-to-f16 conversion. For VLMs, the hook runs in the common VLM weight loader before the family-specific VLM remapping code. In practice, write patterns against the tensor keys as they appear in the checkpoint or in the post-sanitized text loader.

Operations are applied in YAML order. If any operation fails, model loading fails and the process does not continue with a partially modified model.

YAML Schema¶

A surgery file has version: 1 and an ordered operations list:

version: 1
operations:
  - op: scale
    pattern: "model.layers.*.self_attn.o_proj.weight"
    factor: 1.2

  - op: add
    pattern: "model.layers.*.mlp.down_proj.weight"
    source: "./task_vectors/friendly.safetensors"
    alpha: 0.5

Relative source, source_a, and source_b paths are resolved relative to the YAML file. Absolute paths are used as-is.

The file is parsed with YAML 1.2 semantics by serde-saphyr. A field repeated within one mapping is rejected, and malformed-input error messages come from that parser.

Operations¶

scale¶

version: 1
operations:
  - op: scale
    pattern: "model.layers.*.self_attn.o_proj.weight"
    factor: 0.9

scale is useful for gain experiments. On quantized affine layouts, matching a .weight key scales the effective dequantized tensor by changing the adjacent .scales and .biases metadata rather than corrupting packed integer weights.

add¶

version: 1
operations:
  - op: add
    pattern: "model.layers.*.mlp.down_proj.weight"
    source: "./task_vectors/domain_delta.safetensors"
    alpha: 0.5

add is the task-vector operation. The donor file must contain tensors with the same keys and shapes as the matched base tensors. Floating donor tensors are cast to the base dtype when needed. Packed integer base tensors are not modified directly; for quantized checkpoints, target float metadata such as .scales and .biases, or use a floating checkpoint.

prune¶

version: 1
operations:
  - op: prune
    granularity: attention_head
    pattern: "model.layers.12.self_attn.*"
    head_ids: [3, 7]

prune zero-masks selected structure while preserving tensor shapes, so the model architecture does not need to be rebuilt. The supported granularities are:

replace¶

version: 1
operations:
  - op: replace
    pattern: "model.layers.*.self_attn.q_proj.weight"
    source: "./donors/patch.safetensors"
    source_key: "donor.layers.*.attention.q.weight"

replace is for targeted tensor substitution. The number of * wildcards in pattern and source_key must match. For quantized tensors, replacement is atomic across .weight, .scales, and .biases siblings when those siblings exist, and donor shape and dtype must match.

interpolate¶

version: 1
operations:
  - op: interpolate
    pattern: "model.layers.*.*"
    source_a: "./donors/model_a.safetensors"
    source_b: "./donors/model_b.safetensors"
    ratio: 0.3
    method: slerp

interpolate blends matching tensors from two donor checkpoints with ratio as the weight of source_b (ratio: 0.3 keeps 70% of source_a); ratio must be a finite number from 0.0 to 1.0, and method is required. method: lerp uses linear interpolation. method: slerp uses spherical interpolation and falls back to LERP for near-parallel or zero-magnitude tensors. Interpolation is defined for floating tensors; packed quantized tensors should be dequantized or handled through another operation.

Examples¶

Run a No-op Config¶

An empty operation list is valid and should produce the same output as no surgery:

version: 1
operations: []

This is useful for verifying deployment wiring.

Combine Operations¶

Operations are ordered, so later operations see the results of earlier ones:

version: 1
operations:
  - op: scale
    pattern: "model.layers.*.self_attn.o_proj.weight"
    factor: 0.95
  - op: add
    pattern: "model.layers.*.mlp.down_proj.weight"
    source: "./task_vectors/domain_delta.safetensors"
    alpha: 0.25

Start a Server with Surgery¶

mlxcel serve \
  --model ./models/qwen3-4b-4bit \
  --surgery ./surgery.yaml \
  --port 8080

All requests served by that process use the modified in-memory weights. To change the surgery config, restart the server with a different YAML file.

Interactions with Other Features¶

Language Bias is decode-time logit control. It does not change weights and can suppress or promote language/script token sets at sampling time. Use Language Bias when the desired change is output-language policy rather than tensor manipulation.

LoRA adapters are also weight-space composition, but they use the dedicated --adapter path and PEFT-style adapter directories. The current adapter path fuses LoRA weights through a separate preloaded-weight route and does not run the --surgery pipeline. Use one load-time composition path at a time unless you have verified the exact route you are using.

Structured output is constrained decoding. response_format: {"type": "json_schema", ...} masks invalid tokens before sampling but leaves weights unchanged.

Pipeline and tensor parallel runtimes use the common text weight loader for supported families, so the surgery hook is available on those paths when they call the consolidated loader. Unsupported or special load paths may bypass it.

Limitations¶

• Surgery is applied once at model load time. Runtime hot-swap is not supported.
• The modified weights are in memory only. No checkpoint is written.
• Patterns must match actual tensor keys in the weight map seen by the loader.
• add and interpolate require compatible tensor shapes. Interpolation is for floating tensors.
• replace requires exact shape and dtype compatibility.
• Pruning zeroes slices but does not shrink model dimensions, so it is useful for ablation and behavior control, not for reducing compute shape.

Troubleshooting¶

The config fails before loading starts. Check YAML syntax, version: 1, operation names, glob patterns, and donor file paths. Relative paths are resolved from the YAML file's directory.

The pattern matched zero tensors. Inspect the checkpoint tensor names. VLM family loaders may remap keys after the common loader, so patterns for VLM surgery should usually target original checkpoint keys.

A quantized tensor fails with a dtype error. Avoid arithmetic on packed integer .weight tensors unless the operation explicitly supports the quantized layout. For task-vector addition, target float metadata such as .scales and .biases, or use a floating checkpoint.

The server output did not change. Confirm the server was restarted with --surgery, the log reports a loaded surgery pipeline, and the matched tensors are on the active model path.