Initial commit: E4B-MarkBase model integration with passing tests

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully
- All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance)
- All stress tests passed (5/5 in 127.6s)
  - Concurrent inference
  - Memory stress (67.5 tok/s, 0 NaN)
  - Continuous generation
  - Batch processing
  - Long-running stability
- Swift Metal inference engine with multimodal support

convert_e4b.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+"""
+Convert E4B model weights to MarkBase format.
+
+Reads the MLX-community 4-bit quantized Gemma 4 E4B model and re-quantizes
+the vision and audio tower weights into the MarkBase 4-bit packed format.
+"""
+
+import os
+import sys
+import shutil
+import struct
+import json
+import numpy as np
+import safetensors
+
+SRC_DIR = os.path.expanduser(
+    "~/.cache/huggingface/hub/models--mlx-community--gemma-4-e4b-it-4bit/"
+    "snapshots/deb1db712068b1c9f83fb1c97f08c1204b9459a1"
+)
+DST_DIR = "/Users/accusys/MarkBase12B/models/E4B-MarkBase"
+GROUP_SIZE = 64
+
+SRC_MODEL = os.path.join(SRC_DIR, "model.safetensors")
+DST_MODEL = os.path.join(DST_DIR, "model.safetensors")
+DST_INDEX = os.path.join(DST_DIR, "model.safetensors.index.json")
+
+
+def bf16_bytes_to_f32(data: bytes) -> np.ndarray:
+    """Convert BF16 bytes to float32 numpy array."""
+    raw = np.frombuffer(data, dtype=np.uint16)
+    f32 = (raw.astype(np.uint32) << 16).view(np.float32)
+    return f32
+
+
+def quantize_weight_4bit(weight_f32: np.ndarray, group_size: int = GROUP_SIZE):
+    """
+    Group-wise 4-bit affine quantization.
+
+    weight_f32: shape [outDim, inDim], float32
+    Returns (packed_u32, scales_f32, biases_f32)
+      packed_u32: uint32 [outDim, inDim/8]
+      scales_f32: float32 [outDim, inDim/64]
+      biases_f32: float32 [outDim, inDim/64]
+    """
+    out_dim, in_dim = weight_f32.shape
+    assert in_dim % group_size == 0, f"in_dim {in_dim} not divisible by group_size {group_size}"
+    assert in_dim % 8 == 0
+
+    # Reshape to [outDim, inDim/64, 64]
+    num_groups = in_dim // group_size
+    groups = weight_f32.reshape(out_dim, num_groups, group_size)
+
+    # Per-group min/max
+    group_min = groups.min(axis=-1)  # [outDim, num_groups]
+    group_max = groups.max(axis=-1)  # [outDim, num_groups]
+
+    scale = (group_max - group_min) / 15.0
+    bias = group_min
+
+    # Handle zero-range groups (scale == 0)
+    zero_scale_mask = scale == 0.0
+    scale[zero_scale_mask] = 1.0  # avoid division by zero
+
+    # Quantize
+    quantized = np.round((groups - bias[:, :, np.newaxis]) / scale[:, :, np.newaxis])
+    quantized = np.clip(quantized, 0, 15).astype(np.uint32)
+
+    # Pack 8 values into one U32
+    # pack order: q[0] at bits 0-3, q[1] at bits 4-7, ..., q[7] at bits 28-31
+    packed_shape = (out_dim, in_dim // 8)
+    packed = np.zeros(packed_shape, dtype=np.uint32)
+
+    # Reshape quantized to [outDim, inDim/8, 8]
+    quantized_reshaped = quantized.reshape(out_dim, in_dim // 8, 8)
+    for j in range(8):
+        packed |= (quantized_reshaped[:, :, j] << (j * 4))
+
+    return packed, scale.astype(np.float32), bias.astype(np.float32)
+
+
+def is_linear_weight(name: str) -> bool:
+    """Check if a tensor name corresponds to a linear weight that needs quantization."""
+    if name.startswith("language_model."):
+        return False
+    if name.startswith("embed_vision.") or name.startswith("embed_audio."):
+        return False
+    if name.endswith(".linear.weight"):
+        return True
+    if name == "vision_tower.patch_embedder.input_proj.weight":
+        return True
+    if name == "audio_tower.output_proj.weight":
+        return True
+    return False
+
+
+def is_norm_weight(name: str) -> bool:
+    """Check if a tensor name corresponds to a norm/weight that should stay as F32."""
+    if name.startswith("language_model."):
+        return False
+    if name.startswith("embed_vision.") or name.startswith("embed_audio."):
+        return False
+    if "norm" in name.lower() and name.endswith(".weight"):
+        return True
+    if name == "vision_tower.patch_embedder.position_embedding_table":
+        return True
+    if "per_dim_scale" in name:
+        return True
+    if "relative_k_proj.weight" in name:
+        return True
+    if "depthwise_conv1d.weight" in name:
+        return True
+    if "conv_norm.weight" in name:
+        return True
+    if "pre_layer_norm.weight" in name:
+        return True
+    if "post_layer_norm.weight" in name:
+        return True
+    if "norm_out.weight" in name:
+        return True
+    if "norm_pre_attn.weight" in name:
+        return True
+    if "norm_post_attn.weight" in name:
+        return True
+    if name.startswith("audio_tower.subsample_conv"):
+        return True
+    if name.startswith("audio_tower.subsample_conv_projection"):
+        return True
+    if name.endswith("output_proj.bias"):
+        return True
+    return False
+
+
+def process_tensors():
+    """Read source tensors, convert, and write output."""
+    os.makedirs(DST_DIR, exist_ok=True)
+
+    print("Reading source safetensor...")
+    with open(SRC_MODEL, "rb") as f:
+        src_data = f.read()
+    tensors = safetensors.deserialize(src_data)
+    print(f"  Found {len(tensors)} tensors")
+
+    output_tensor_dict = {}
+    text_copied = 0
+    embed_copied = 0
+    vision_quantized = 0
+    vision_f32 = 0
+    audio_quantized = 0
+    audio_f32 = 0
+    skipped = 0
+
+    # Rust safetensors backend only accepts lowercase dtype names for writing.
+    DTYPE_MAP = {
+        "U32": "uint32",
+        "BF16": "bfloat16",
+        "F16": "float16",
+        "F32": "float32",
+        "I32": "int32",
+        "I64": "int64",
+    }
+
+    def add_tensor(name, shape, dtype, data):
+        if isinstance(data, bytearray):
+            data = bytes(data)
+        elif not isinstance(data, bytes):
+            data = data.tobytes()
+        py_dtype = DTYPE_MAP.get(dtype, dtype)
+        output_tensor_dict[name] = {
+            "shape": list(shape) if isinstance(shape, (list, tuple)) else shape,
+            "dtype": py_dtype,
+            "data": data,
+        }
+
+    for name, info in tensors:
+        dtype_tag = info["dtype"]
+        shape = info["shape"]
+        data = info["data"]
+
+        # --- Text model: copy as-is ---
+        if name.startswith("language_model."):
+            add_tensor(name, shape, dtype_tag, data)
+            text_copied += 1
+            continue
+
+        # --- Embed vision/audio projections: already quantized, copy as-is ---
+        if name.startswith("embed_vision.") or name.startswith("embed_audio."):
+            add_tensor(name, shape, dtype_tag, data)
+            embed_copied += 1
+            continue
+
+        # --- Vision/Audio tower input_min/max/output_min/output_max: skip ---
+        if ("input_min" in name or "input_max" in name or
+            "output_min" in name or "output_max" in name):
+            skipped += 1
+            continue
+
+        # --- Quantize linear weights ---
+        if is_linear_weight(name):
+            print(f"  Quantizing: {name}  shape={shape}")
+            weight_f32 = bf16_bytes_to_f32(data).reshape(shape)
+            packed, scales, biases = quantize_weight_4bit(weight_f32)
+
+            out_dim, in_dim = shape
+
+            # Strip `.linear` if present; otherwise use name directly.
+            # e.g. "q_proj.linear.weight" → "q_proj.weight" / ".scales" / ".biases"
+            #       "input_proj.weight"   → "input_proj.weight" / ".scales" / ".biases"
+            if ".linear.weight" in name:
+                weight_name = name.replace(".linear.weight", ".weight")
+                scales_name = name.replace(".linear.weight", ".scales")
+                biases_name = name.replace(".linear.weight", ".biases")
+            else:
+                weight_name = name
+                scales_name = name.replace(".weight", ".scales")
+                biases_name = name.replace(".weight", ".biases")
+
+            add_tensor(weight_name, [out_dim, in_dim // 8], "U32", packed)
+            add_tensor(scales_name, [out_dim, in_dim // GROUP_SIZE], "F32", scales)
+            add_tensor(biases_name, [out_dim, in_dim // GROUP_SIZE], "F32", biases)
+
+            if name.startswith("vision_tower."):
+                vision_quantized += 1
+            else:
+                audio_quantized += 1
+            continue
+
+        # --- Convert BF16 norms and special tensors to float32 ---
+        if is_norm_weight(name):
+            f32_array = bf16_bytes_to_f32(data).reshape(shape)
+            add_tensor(name, shape, "F32", f32_array)
+            if name.startswith("vision_tower."):
+                vision_f32 += 1
+            else:
+                audio_f32 += 1
+            continue
+
+        # --- Catch-all: convert any remaining BF16 to F32 ---
+        if dtype_tag == "BF16":
+            print(f"  Converting BF16->F32: {name}  shape={shape}")
+            f32_array = bf16_bytes_to_f32(data).reshape(shape)
+            add_tensor(name, shape, "F32", f32_array)
+            if name.startswith("vision_tower."):
+                vision_f32 += 1
+            else:
+                audio_f32 += 1
+            continue
+
+        # Should not reach here for any known tensor
+        print(f"  WARNING: unhandled tensor: {name}  dtype={dtype_tag}  shape={shape}")
+        add_tensor(name, shape, dtype_tag, data)
+
+    print()
+    print(f"Text tensors copied: {text_copied}")
+    print(f"Embed tensors copied: {embed_copied}")
+    print(f"Vision quantized: {vision_quantized}")
+    print(f"Vision f32: {vision_f32}")
+    print(f"Audio quantized: {audio_quantized}")
+    print(f"Audio f32: {audio_f32}")
+    print(f"Input/output minmax skipped: {skipped}")
+    print(f"Total output tensors: {len(output_tensor_dict)}")
+
+    print("\nSerializing output safetensor...")
+    serialized = safetensors.serialize(output_tensor_dict)
+    with open(DST_MODEL, "wb") as f:
+        f.write(serialized)
+
+    # Write index file
+    total_bytes = os.path.getsize(DST_MODEL)
+    weight_map = {name: "model.safetensors" for name in output_tensor_dict}
+    index_data = {
+        "metadata": {"total_size": total_bytes},
+        "weight_map": weight_map,
+    }
+    with open(DST_INDEX, "w") as f:
+        json.dump(index_data, f, indent=2)
+
+    print(f"Output model: {DST_MODEL} ({total_bytes} bytes)")
+    print(f"Output index: {DST_INDEX}")
+
+
+def copy_metadata_files():
+    """Copy non-weight files from source to destination."""
+    os.makedirs(DST_DIR, exist_ok=True)
+    metadata_files = [
+        "config.json",
+        "generation_config.json",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "processor_config.json",
+        "chat_template.jinja",
+    ]
+    for fn in metadata_files:
+        src = os.path.join(SRC_DIR, fn)
+        dst = os.path.join(DST_DIR, fn)
+        if os.path.exists(src):
+            shutil.copy2(src, dst)
+            print(f"  Copied {fn}")
+        else:
+            print(f"  WARNING: {fn} not found in source")
+
+
+def verify_output():
+    """Verify the output safetensor can be read back."""
+    print("\nVerifying output...")
+    try:
+        with open(DST_MODEL, "rb") as f:
+            data = f.read()
+        tensors = safetensors.deserialize(data)
+        keys = [t[0] for t in tensors]
+        print(f"  Output contains {len(keys)} tensors")
+
+        # Build a lookup by key
+        tensor_map = {t[0]: t[1] for t in tensors}
+
+        # Verify a text tensor
+        sample = "language_model.model.layers.0.input_layernorm.weight"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+
+        # Verify a vision quantized tensor
+        sample = "vision_tower.encoder.layers.0.self_attn.q_proj.weight"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+        sample = "vision_tower.encoder.layers.0.self_attn.q_proj.scales"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+
+        # Verify a vision norm tensor
+        sample = "vision_tower.encoder.layers.0.input_layernorm.weight"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+
+        # Verify audio quantized
+        sample = "audio_tower.layers.0.self_attn.q_proj.weight"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+
+        # Verify embed_vision
+        sample = "embed_vision.embedding_projection.weight"
+        if sample in tensor_map:
+            info = tensor_map[sample]
+            print(f"  {sample}: dtype={info['dtype']} shape={info['shape']}")
+
+        print("  All checks passed!")
+    except Exception as e:
+        print(f"  ERROR during verification: {e}")
+        sys.exit(1)
+
+
+def main():
+    print("=== E4B to MarkBase Converter ===\n")
+    print(f"Source: {SRC_MODEL}")
+    print(f"Destination: {DST_DIR}\n")
+
+    copy_metadata_files()
+    process_tensors()
+    verify_output()
+
+    print("\nDone! Model converted successfully.")
+
+
+if __name__ == "__main__":
+    main()