Skip to content

Forge writes three drafter norm tensors shifted by exactly +1.0 when extracting an embedded head (draft acceptance collapses to 0-2%) #301

Description

@Blakeolson21

Summary

Running Forge against a checkpoint that embeds its MTP (multi-token
prediction) tensors in the main shards produces an mtp.safetensors sidecar
in which exactly three of the seven F32 RMSNorm gain tensors are wrong:

  • mtp.layers.0.self_attn.q_norm.weight (256,)
  • mtp.layers.0.self_attn.k_norm.weight (256,)
  • mtp.norm.weight (5120,)

Each of the three comes out as the source value plus exactly 1.0, elementwise
and bit-exact in F32. The other 26 tensors in the sidecar, including the
remaining four F32 norm gains, are byte-identical to the values stored in the
source shards. The trunk loads correctly, so the drafter runs against a trunk
it no longer matches. Base decode is healthy (12.67 to 12.89 tok/s on this
machine) while draft acceptance collapses to 0 to 2 percent. That collapse
makes every draft depth slower than plain decoding.

The extraction is deterministic: two independent Forge runs on 2026-08-19
produced tensor-identical corrupted sidecars.

Environment

  • mtplx 2.8.1, installed binary at ~/.mtplx/bin/mtplx. Per the changelog,
    2.8.2 and 2.8.3 (current main) change nothing in the forge extraction or
    sanitize path, so this applies to current main.
  • macOS 26.3, Apple M3 Max, 64 GB
  • Source model: chimingw/Qwen3.8-27B-Uncensored-OrcaRouter-MLX-6bit
    (Qwen3.8-27B family, 6-bit MLX quantization, group size 64). The
    checkpoint's index lists 29 language_model.mtp.* tensors, and its norm
    gains are stored in the absolute convention, the same one the official
    sidecar-layout builds use.

Repro

  1. Download chimingw/Qwen3.8-27B-Uncensored-OrcaRouter-MLX-6bit (the
    embedded-MTP layout: 29 language_model.mtp.* keys inside
    model-0000x-of-00006.safetensors).
  2. Run mtplx forge build on it to produce an MTPLX artifact with an
    mtp.safetensors sidecar.
  3. Compare the sidecar tensors against the same tensors read directly from the
    source shards (method below). Three norm gains differ by exactly +1.0; the
    other 26 tensors match byte for byte.
  4. Serve the forged artifact with speculative decoding enabled. Acceptance sits
    at 0 to 2 percent and depths 1 through 3 all decode slower than the base
    model, while base decoding itself is healthy.

Tensor-diff method

Read both files with safetensors.safe_open(path, framework="np"), compare
per key. For the source side, resolve each mtp.* sidecar key to its
language_model.mtp.* shard key through model.safetensors.index.json.

import numpy as np
from safetensors import safe_open

with safe_open(forged_sidecar, framework="np") as fb, \
     safe_open(source_shard, framework="np") as fs:
    b = fb.get_tensor("mtp.norm.weight")
    r = fs.get_tensor("language_model.mtp.norm.weight")
    assert np.array_equal(b, r + np.float32(1.0))   # holds, bit-exact

Measured on the forged output (good here is the raw shard value):

mtp.layers.0.self_attn.q_norm.weight  f32 (256,)   good mean 1.7906  forged mean 2.7906  diff exactly 1.0 on all 256 elements
mtp.layers.0.self_attn.k_norm.weight  f32 (256,)   good mean 1.7795  forged mean 2.7795  diff exactly 1.0 on all 256 elements
mtp.norm.weight                       f32 (5120,)  good mean 2.2520  forged mean 3.2520  diff exactly 1.0 on all 5120 elements

The four F32 norm gains that survive intact are
mtp.layers.0.input_layernorm.weight,
mtp.layers.0.post_attention_layernorm.weight,
mtp.pre_fc_norm_embedding.weight, and mtp.pre_fc_norm_hidden.weight.

Mechanism

The offset is exactly the zero-centered-to-absolute RMSNorm convention shift,
and it is readable directly in the extraction path. mtplx forge build
extracts the embedded MTP block through
mtplx/commands/forge.py::_extract_embedded_mtp, which copies the tensors via
_copy_safetensors_subset_sanitized, which applies
mtplx/compressed_tensors.py::sanitize_plain_weight to every copied tensor.
That function shifts norm gains by +1.0 in two tiers. Keys ending in
MTP_RMSNORM_ALWAYS_SHIFT_SUFFIXES (self_attn.q_norm.weight,
self_attn.k_norm.weight, mtp.norm.weight; compressed_tensors.py lines
53-56) are shifted unconditionally. Keys ending in
MTP_RMSNORM_SHIFT_IF_LOW_SUFFIXES (the two layernorms and the two pre-fc
norms; lines 47-52) are shifted only when the tensor mean is below 0.5.

This checkpoint stores its gains absolute, so every norm gain has a mean at or
above 0.5. The conditional tier therefore correctly leaves its four tensors
alone, and the unconditional tier double-shifts exactly the three tensors
observed. The two suffix lists fully explain the measured split.

Workaround that confirms the diagnosis

Building the sidecar by moving the same 29 tensors out of the shards with no
value changes (a pure repack) produces a drafter with 97.6 percent acceptance
at depth 2 on the same trunk, same machine, same prompts. The forged and
repacked sidecars are tensor-identical except for the three norm gains above.

I kept both corrupted sidecars and can attach header dumps, full per-tensor
diffs, or the acceptance logs if useful.

mtplx doctor --json

All 14 diagnostic checks pass on this machine. Full output attached below.

doctor output
{
 "compiled_verify": {
  "above_fence_behavior": "eager verify per call",
  "default_model": "~/.mtplx/models/Youssofal--Qwen3.8-27B-MTPLX-Optimized-Speed",
  "fenced": true,
  "max_context_source": "turbo profile",
  "max_context_tokens": 32768,
  "mode": "on",
  "mode_source": "turbo profile",
  "resolved_default_profile": "turbo"
 },
 "diagnostics": {
  "checks": [
   {
    "command": null,
    "docs_url": "https://ml-explore.github.io/mlx/build/html/install.html",
    "expected": "macOS >= 14.0 on Apple Silicon",
    "fix": "Upgrade to macOS 14+; MLX does not support older macOS.",
    "id": "os.macos_version",
    "observed": "26.3.1",
    "severity": "error",
    "status": "pass"
   },
   {
    "command": "python3 -c \"import platform; print(platform.machine(), platform.processor())\"",
    "docs_url": "https://ml-explore.github.io/mlx/build/html/install.html",
    "expected": "native arm64 Python, not Rosetta",
    "fix": "Install/use a native arm64 Python. If needed, reinstall via Homebrew arm64 or uv.",
    "id": "python.native_arm64",
    "observed": {
     "machine": "arm64",
     "processor": "arm"
    },
    "severity": "error",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": "https://ml-explore.github.io/mlx/build/html/install.html",
    "expected": "Python >= 3.11",
    "fix": "Install Python 3.11 or newer.",
    "id": "python.version",
    "observed": "3.14.5",
    "severity": "error",
    "status": "pass"
   },
   {
    "command": "python3 -m pip install mlx",
    "docs_url": "https://ml-explore.github.io/mlx/build/html/install.html",
    "expected": "mlx importable",
    "fix": "Install MLX into this same native Python environment.",
    "id": "mlx.import",
    "observed": {
     "default_device": "Device(gpu, 0)",
     "get_active_memory": 0,
     "get_peak_memory": 0,
     "mlx": "0.32.1",
     "mlx_lm": "0.31.3"
    },
    "severity": "error",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "the model this Mac's default routes to fits: measured peak <= unified memory (comfortable at 1.5x)",
    "fix": null,
    "id": "resource.memory",
    "observed": {
     "default_model": "Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
     "estimated_peak_gib": 25.0,
     "unified_memory_gib": 64.0
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "free space for model + temp download + safety headroom",
    "fix": "Free disk space or set MTPLX_MODEL_DIR to a larger volume.",
    "id": "resource.model_cache_disk",
    "observed": {
     "cache_dir": "~/.mtplx/models",
     "free_gib": 552.02,
     "required_gib": 49.63
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": "mtplx pull Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "docs_url": "https://huggingface.co/Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "expected": "default model available in the HF cache or as the verified local startup model",
    "fix": "No action needed.",
    "id": "model.cache",
    "observed": {
     "hf_cache_exists": true,
     "hf_cache_path": "~/.mtplx/models/Youssofal--Qwen3.8-27B-MTPLX-Optimized-Speed",
     "hf_cache_validation": {
      "contract_arch_id": "qwen3-next-mtp",
      "contract_error": null,
      "contract_present": true,
      "missing_files": [],
      "mtp_sidecar_candidates": [
       "mtp.safetensors",
       "mtp/weights.safetensors",
       "model-mtp.safetensors"
      ],
      "ok": true,
      "required_files": [
       "config.json",
       "tokenizer.json",
       "model.safetensors.index.json",
       "mtplx_runtime.json",
       "mtp.safetensors"
      ]
     },
     "startup_default_model": null
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": "mtplx pull Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "docs_url": "https://huggingface.co/Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "expected": "a published Youssofal/... repo (not a local mtplx/ or models/ path)",
    "fix": "Pull the default model, or pass --model to serve a different one.",
    "id": "model.default_repo",
    "observed": "Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "severity": "error",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": "https://docs.docker.com/desktop/setup/install/mac-install/",
    "expected": "Docker Desktop installed for Open WebUI Docker path",
    "fix": "Install Docker Desktop if you want the Open WebUI Docker integration.",
    "id": "docker.binary",
    "observed": "~/.docker/bin/docker",
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "port free before starting mtplx serve, or already a healthy MTPLX server",
    "fix": "A healthy MTPLX server already on this port is fine to keep using; if something else holds it, stop that process or use --port 8001.",
    "id": "port.mtplx_server",
    "observed": {
     "host": "127.0.0.1",
     "open": false,
     "port": 8000
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "port free before starting Open WebUI, or already an Open WebUI container",
    "fix": "Use a different Open WebUI host port or stop the process on 3000.",
    "id": "port.openwebui",
    "observed": {
     "host": "127.0.0.1",
     "open": false,
     "port": 3000
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "ThermalForge or TG Pro available for explicit --max only",
    "fix": "Install ThermalForge only if you want opt-in fan boost.",
    "id": "thermal.control",
    "observed": {
     "kind": "thermalforge",
     "path": "~/.mtplx/bin/thermalforge",
     "version": {
      "command": [
       "~/.mtplx/bin/thermalforge",
       "--version"
      ],
      "ok": true,
      "returncode": 0,
      "stderr": "",
      "stdout": "0.1.0"
     }
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "Low Power Mode off for best sustained decode",
    "fix": "Turn off Low Power Mode before benchmarking or serving long responses.",
    "id": "power.low_power_mode",
    "observed": {
     "available": true,
     "lowpowermode": null,
     "powermode": "0",
     "thermal": "Note: No thermal warning level has been recorded\nNote: No performance warning level has been recorded\nNote: No CPU power status has been recorded",
     "thermal_ok": true
    },
    "severity": "warning",
    "status": "pass"
   },
   {
    "command": null,
    "docs_url": null,
    "expected": "no recorded thermal or performance warning",
    "fix": "Let the Mac cool down or improve airflow before sustained benchmarks.",
    "id": "power.thermal_pressure",
    "observed": "Note: No thermal warning level has been recorded\nNote: No performance warning level has been recorded\nNote: No CPU power status has been recorded",
    "severity": "warning",
    "status": "pass"
   }
  ],
  "created_at": "2026-08-19T17:25:46-0500",
  "host": {
   "cache_dir": "~/.mtplx/models",
   "chip": "Apple M3 Max",
   "disk_free_bytes": 592727658496,
   "disk_free_gib": 552.02,
   "mac_model": "Mac15,9",
   "machine": "arm64",
   "macos_version": "26.3.1",
   "memory_bytes": 68719476736,
   "memory_gib": 64.0,
   "platform": "macOS-26.3.1-arm64-arm-64bit-Mach-O",
   "processor": "arm",
   "python_executable": "~/Library/Application Support/MTPLX/runtime-venv/bin/python",
   "python_version": "3.14.5",
   "system": "Darwin"
  },
  "overall": "pass",
  "resources": {
   "default_model_size_bytes": 21313949792,
   "estimated_runtime_memory_bytes": 42788786272,
   "required_download_free_bytes": 53284874480
  },
  "schema_version": 1,
  "support_matrix": {
   "preview_test_targets": [
    "M3 Max",
    "M4 Max",
    "M3 Ultra / Mac Studio",
    "M5 Max"
   ],
   "supported": {
    "default_model": "Youssofal/Qwen3.8-27B-MTPLX-Optimized-Speed",
    "default_profile": "turbo",
    "docker": "Docker Desktop current plus previous two macOS major releases",
    "macos": ">= 14.0",
    "platform": "Apple Silicon arm64 Mac",
    "python": "native arm64 Python >= 3.11"
   }
  }
 },
 "environment": {
  "git_branch": "not a git worktree",
  "git_status": "not a git worktree",
  "hf_path": null,
  "mlx": {
   "default_device": "Device(gpu, 0)",
   "get_active_memory": 0,
   "get_peak_memory": 0,
   "mlx": "0.32.1",
   "mlx_lm": "0.31.3"
  },
  "platform": "macOS-26.3.1-arm64-arm-64bit-Mach-O",
  "project_root": "~",
  "python_executable": "~/Library/Application Support/MTPLX/runtime-venv/bin/python",
  "python_version": "3.14.5 (main, Jun  2 2026, 22:28:56) [Clang 22.1.3 ]",
  "uv_path": "~/.local/bin/uv"
 },
 "huggingface": {
  "cache_dir": "~/.mtplx/models",
  "cache_exists": true,
  "cache_writable": true,
  "cached_models": 12,
  "disk_free_bytes": 592727658496,
  "disk_free_gb": 592.728,
  "token_present": true,
  "token_source": "huggingface_hub"
 },
 "policy": {
  "benchmark_exactness_smoke_context": 2048,
  "fanmax_counts_for_product_gate": false
 },
 "thermal_control": {
  "available": true,
  "clock_anchor_enabled": false,
  "clock_anchor_policy": "explicit experimental only; never used for product claims",
  "instructions": "Install ThermalForge and ensure the thermalforge CLI is on PATH.",
  "selected": {
   "kind": "thermalforge",
   "path": "~/.mtplx/bin/thermalforge",
   "version": {
    "command": [
     "~/.mtplx/bin/thermalforge",
     "--version"
    ],
    "ok": true,
    "returncode": 0,
    "stderr": "",
    "stdout": "0.1.0"
   }
  },
  "tools": [
   {
    "kind": "thermalforge",
    "path": "~/.mtplx/bin/thermalforge",
    "version": {
     "command": [
      "~/.mtplx/bin/thermalforge",
      "--version"
     ],
     "ok": true,
     "returncode": 0,
     "stderr": "",
     "stdout": "0.1.0"
    }
   }
  ]
 },
 "tools": {
  "powermetrics": "/usr/bin/powermetrics",
  "python": "~/Library/Application Support/MTPLX/runtime-venv/bin/python",
  "smc_atlas": null,
  "smc_atlas_exists": false,
  "sovereign": null,
  "sovereign_exists": false,
  "sudo": "/usr/bin/sudo"
 }
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions