Add convert model script

Author: neoremind

convert_model.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+
+import json
+import struct
+import torch
+import argparse
+from safetensors import safe_open
+
+# MXFP4 conversion constants
+FP4_VALUES = [
+    +0.0, +0.5, +1.0, +1.5, +2.0, +3.0, +4.0, +6.0,
+    -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
+]
+
+def convert_mxfp4_to_bf16(blocks_tensor, scales_tensor):
+    """Convert MXFP4 format to BF16"""
+    blocks = blocks_tensor
+    scales = scales_tensor.to(torch.int32) - 127
+
+    lut = torch.tensor(FP4_VALUES, dtype=torch.float32, device=blocks.device)
+
+    # Split into low and high nibbles
+    idx_lo = (blocks_tensor & 0x0F).to(torch.long)   # [..., G, B]
+    idx_hi = (blocks_tensor >> 4).to(torch.long)     # [..., G, B]
+
+    # Convert to FP32 values
+    vals_lo = lut[idx_lo]   # [..., G, B]
+    vals_hi = lut[idx_hi]   # [..., G, B]
+
+    # Interleave along last dimension → [..., G, 2B]
+    interleaved = torch.stack((vals_lo, vals_hi), dim=-1).reshape(*vals_lo.shape[:-1], -1)
+
+    # Apply exponent scaling: x * 2**exp
+    result = torch.ldexp(interleaved, scales.unsqueeze(-1))
+
+    return result.to(torch.bfloat16)
+
+def get_target_dtype(weight_name):
+    """Determine target dtype based on weight name"""
+    if weight_name in ["embedding.weight", "unembedding.weight"]:
+        return "BF16"
+    elif "norm.scale" in weight_name and not weight_name.startswith("block."):
+        return "BF16"  # final layer norm
+    elif any(x in weight_name for x in ["mlp1_weight.blocks", "mlp2_weight.blocks", "mlp1_weight.scales", "mlp2_weight.scales"]):
+        return "BF16"
+    elif any(x in weight_name for x in ["attn.norm.scale", "attn.qkv", "attn.sinks", "attn.out", "mlp.norm.scale", "mlp.gate", "mlp1_bias", "mlp2_bias"]):
+        return "BF16"
+    else:
+        print(f"Not recognized weight name {weight_name}")
+        return "BF16"  # default
+
+def read_model_header(model_path):
+    """Read header from footer of model.bin file"""
+    with open(model_path, 'rb') as f:
+        # Read header size from last 8 bytes
+        f.seek(-8, 2)
+        header_size = struct.unpack('<Q', f.read(8))[0]
+
+        # Read header JSON
+        f.seek(-(8 + header_size), 2)
+        header_json = f.read(header_size).decode('utf-8')
+        return json.loads(header_json)
+
+def read_tensor_by_name(model_path, tensor_name, header=None):
+    """Read specific tensor from model.bin file"""
+    if header is None:
+        header = read_model_header(model_path)
+
+    tensor_info = header[tensor_name]
+    start_offset, end_offset = tensor_info["data_offsets"]
+    shape = tensor_info["shape"]
+    dtype = tensor_info["dtype"]
+
+    with open(model_path, 'rb') as f:
+        f.seek(start_offset)
+        tensor_bytes = f.read(end_offset - start_offset)
+
+    # Convert bytes to tensor
+    if dtype == "FP32":
+        tensor = torch.frombuffer(tensor_bytes, dtype=torch.float32).view(shape)
+    elif dtype == "BF16":
+        tensor = torch.frombuffer(tensor_bytes, dtype=torch.bfloat16).view(shape)
+    else:
+        tensor = torch.frombuffer(tensor_bytes, dtype=torch.float32).view(shape)
+
+    return tensor
+
+def get_dtype_size_multiplier(original_dtype, target_dtype):
+    """Get expected size multiplier for dtype conversion"""
+    dtype_sizes = {"FP32": 4, "BF16": 2, "FP16": 2, "FP4": 0.5}
+
+    orig_size = dtype_sizes.get(original_dtype, 4)
+    target_size = dtype_sizes.get(target_dtype, 4)
+
+    return target_size / orig_size
+
+def validate_data_offset_size(original_size, output_size, original_dtype, target_dtype, tensor_name):
+    """Validate output data offset size matches expected conversion ratio"""
+    expected_multiplier = get_dtype_size_multiplier(original_dtype, target_dtype)
+    expected_size = int(original_size * expected_multiplier)
+
+    if output_size != expected_size:
+        raise ValueError(f"Data offset size mismatch for {tensor_name}: "
+                        f"expected {expected_size} bytes (original {original_size} * {expected_multiplier}), "
+                        f"got {output_size} bytes")
+
+def convert_safetensors_to_modelbin(input_path, output_path):
+    """Convert safetensors to custom model.bin format with streaming writes"""
+
+    header = {}
+    current_offset = 0
+
+    with open(output_path, 'wb') as out_f:
+        with safe_open(input_path, framework="pt", device="cpu") as f:
+            tensor_names = list(f.keys())
+
+            # Sort tensor names: embedding, unembedding, norm.scale, then blocks in ascending order
+            def sort_key(name):
+                if name == "embedding.weight":
+                    return (0, 0, name)
+                elif name == "unembedding.weight":
+                    return (1, 0, name)
+                elif name == "norm.scale":
+                    return (2, 0, name)
+                elif name.startswith("block."):
+                    # Extract block number for proper numeric sorting
+                    block_num = int(name.split('.')[1])
+                    return (3, block_num, name)
+                else:
+                    return (4, 0, name)
+
+            tensor_names.sort(key=sort_key)
+
+            # Process each tensor and write immediately
+            for name in tensor_names:
+                target_dtype = get_target_dtype(name)
+                original_tensor = f.get_tensor(name)
+                original_dtype = str(original_tensor.dtype).upper().replace("TORCH.", "")
+                if original_dtype == "BFLOAT16":
+                    original_dtype = "BF16"
+                elif original_dtype == "FLOAT32":
+                    original_dtype = "FP32"
+
+                # Calculate original size
+                original_size = original_tensor.numel() * original_tensor.element_size()
+
+                # Handle MXFP4 weights
+                if name.endswith(".mlp1_weight.blocks") or name.endswith(".mlp2_weight.blocks"):
+                    base_name = name.replace(".blocks", "")
+                    scales_name = name.replace(".blocks", ".scales")
+
+                    if scales_name in tensor_names:
+                        blocks = f.get_tensor(name)
+                        scales = f.get_tensor(scales_name)
+                        tensor = convert_mxfp4_to_bf16(blocks, scales)
+                        dtype_str = "BF16"
+                        original_dtype = "FP4"  # Override for FP4 blocks
+                    else:
+                        continue
+                elif name.endswith(".scales"):
+                    continue  # Skip scales, handled with blocks
+                else:
+                    tensor = original_tensor
+
+                    # Convert to target dtype
+                    if target_dtype == "FP32":
+                        tensor = tensor.float()
+                        dtype_str = "FP32"
+                    elif target_dtype == "BF16":
+                        tensor = tensor.to(torch.bfloat16)
+                        dtype_str = "BF16"
+                    else:
+                        dtype_str = original_dtype  # Keep original
+
+                # Write tensor data immediately
+                if tensor.dtype == torch.bfloat16:
+                    tensor_bytes = tensor.view(torch.uint16).numpy().tobytes()
+                else:
+                    tensor_bytes = tensor.numpy().tobytes()
+
+                out_f.write(tensor_bytes)
+
+                # Validate data offset size
+                output_size = len(tensor_bytes)
+                validate_data_offset_size(original_size, output_size, original_dtype, dtype_str, name)
+
+                # Update header with offset info
+                header[name] = {
+                    "dtype": dtype_str,
+                    "shape": list(tensor.shape),
+                    "data_offsets": [current_offset, current_offset + len(tensor_bytes)]
+                }
+                current_offset += len(tensor_bytes)
+                print(f"Wrote tensor: {name} {header[name]}")
+
+                # Free memory
+                del tensor
+                if 'blocks' in locals():
+                    del blocks
+                if 'scales' in locals():
+                    del scales
+
+        # Write footer with header info
+        header_json = json.dumps(header, separators=(',', ':')).encode('utf-8')
+        header_size = len(header_json)
+        out_f.write(header_json)
+        out_f.write(struct.pack('<Q', header_size))  # Footer: header size at end
+        print(f"Wrote header as footer")
+
+    print(f"Converted {len(header)} tensors from {input_path} to {output_path}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert GPT-OSS safetensors to model.bin format")
+    parser.add_argument("input_path", help="Path to input safetensors file")
+    parser.add_argument("output_path", help="Path to output model.bin file")
+
+    args = parser.parse_args()
+
+    convert_safetensors_to_modelbin(args.input_path, args.output_path)