| |
| """Pack pre-quantized Ling-3 assets into the mmap-friendly .l3r format.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import struct |
| import zlib |
| from pathlib import Path |
|
|
|
|
| ALIGNMENT = 4096 |
| HEADER = struct.Struct("<8s6I6Q20sI28I4f24s") |
| ENTRY = struct.Struct("<12I4Q2f4I32s24s") |
|
|
| DTYPES = { |
| "unknown": 0, |
| "bf16": 1, |
| "fp16": 2, |
| "fp32": 3, |
| "f32": 3, |
| "int8": 4, |
| "int4": 5, |
| "int16": 6, |
| "int32": 7, |
| "rknn": 8, |
| } |
| ROLES = { |
| "unknown": 0, |
| "embedding": 1, |
| "norm": 2, |
| "linear_weight": 3, |
| "linear_scale": 4, |
| "bias": 5, |
| "rknn_island": 6, |
| "tokenizer": 7, |
| } |
| QUANTS = {"none": 0, "per_tensor": 1, "per_output_channel": 2, "fixed_activation": 3} |
| LAYOUTS = {"row_major": 0, "rknn_native_b": 1, "packed_int4": 2, "opaque": 3} |
|
|
| MODEL_FIELDS = ( |
| "vocab_size", |
| "hidden_size", |
| "layer_count", |
| "attention_heads", |
| "head_dim", |
| "kv_lora_rank", |
| "q_lora_rank", |
| "qk_nope_dim", |
| "qk_rope_dim", |
| "value_head_dim", |
| "dense_ffn_dim", |
| "expert_ffn_dim", |
| "shared_ffn_dim", |
| "expert_count", |
| "experts_per_token", |
| "expert_group_count", |
| "selected_group_count", |
| "layer_group_size", |
| "leading_dense_layers", |
| "convolution_kernel", |
| "max_context", |
| "eos_token", |
| "pad_token", |
| "bos_token", |
| "mla_layer_count", |
| "kda_layer_count", |
| "default_weight_bits", |
| "default_activation_bits", |
| ) |
|
|
|
|
| def align(value: int) -> int: |
| return (value + ALIGNMENT - 1) & ~(ALIGNMENT - 1) |
|
|
|
|
| def enum_value(mapping: dict[str, int], value: str, label: str) -> int: |
| try: |
| return mapping[value] |
| except KeyError as error: |
| raise ValueError(f"unknown {label} {value!r}") from error |
|
|
|
|
| def resolve(base: Path, value: str) -> Path: |
| path = Path(value) |
| return path if path.is_absolute() else base / path |
|
|
|
|
| def copy_file(source: Path, output, hasher: hashlib._Hash) -> int: |
| total = 0 |
| with source.open("rb") as stream: |
| while chunk := stream.read(8 * 1024 * 1024): |
| output.write(chunk) |
| hasher.update(chunk) |
| total += len(chunk) |
| return total |
|
|
|
|
| def pad_to(output, offset: int) -> None: |
| current = output.tell() |
| if current > offset: |
| raise ValueError("package layout overlapped while writing") |
| output.write(b"\0" * (offset - current)) |
|
|
|
|
| def make_header(manifest: dict, tensor_count: int, string_offset: int, string_bytes: int, |
| payload_offset: int, payload_bytes: int, file_bytes: int, crc: int) -> bytes: |
| model = manifest["model"] |
| source_revision = bytes.fromhex(manifest["source_revision"]) |
| if len(source_revision) != 20: |
| raise ValueError("source_revision must be a 40-character Git SHA-1") |
| integers = [int(model[name]) for name in MODEL_FIELDS] |
| floats = [ |
| float(model["rms_epsilon"]), |
| float(model["rope_theta"]), |
| float(model["routed_scale"]), |
| float(model["kda_lower_bound"]), |
| ] |
| return HEADER.pack( |
| b"L3RKNN1\0", |
| 1, |
| HEADER.size, |
| 0x01020304, |
| tensor_count, |
| ENTRY.size, |
| int(manifest.get("flags", 0)), |
| HEADER.size, |
| string_offset, |
| string_bytes, |
| payload_offset, |
| payload_bytes, |
| file_bytes, |
| source_revision, |
| crc, |
| *integers, |
| *floats, |
| bytes(24), |
| ) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("manifest", type=Path) |
| parser.add_argument("output", type=Path) |
| parser.add_argument( |
| "--include-prefix", |
| action="append", |
| default=[], |
| help="pack only tensors whose names start with this prefix (repeatable)", |
| ) |
| args = parser.parse_args() |
|
|
| manifest_path = args.manifest.resolve() |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) |
| base = manifest_path.parent |
| tensors = manifest.get("tensors", []) |
| if args.include_prefix: |
| tensors = [ |
| item for item in tensors |
| if any(item["name"].startswith(prefix) for prefix in args.include_prefix) |
| ] |
| if not tensors: |
| raise ValueError("--include-prefix did not match any tensors") |
|
|
| names = bytearray() |
| prepared = [] |
| for item in tensors: |
| name = item["name"].encode("utf-8") |
| if not name or b"\0" in name: |
| raise ValueError("tensor names must be non-empty UTF-8 without NUL") |
| dims = [int(value) for value in item.get("dims", [])] |
| if len(dims) > 4: |
| raise ValueError(f"tensor {item['name']!r} has more than four dimensions") |
| source = resolve(base, item["path"]) |
| if not source.is_file(): |
| raise FileNotFoundError(source) |
| aux = resolve(base, item["aux_path"]) if item.get("aux_path") else None |
| if aux is not None and not aux.is_file(): |
| raise FileNotFoundError(aux) |
| prepared.append((item, len(names), name, dims, source, aux)) |
| names.extend(name) |
|
|
| table_bytes = len(prepared) * ENTRY.size |
| string_offset = HEADER.size + table_bytes |
| payload_offset = align(string_offset + len(names)) |
| cursor = payload_offset |
| layouts = [] |
| for item, name_offset, name, dims, source, aux in prepared: |
| data_offset = align(cursor) |
| data_bytes = source.stat().st_size |
| cursor = data_offset + data_bytes |
| aux_offset = align(cursor) if aux is not None else 0 |
| aux_bytes = aux.stat().st_size if aux is not None else 0 |
| if aux is not None: |
| cursor = aux_offset + aux_bytes |
| layouts.append((item, name_offset, name, dims, source, aux, data_offset, data_bytes, aux_offset, aux_bytes)) |
| file_bytes = align(cursor) |
| payload_bytes = file_bytes - payload_offset |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| entries = [] |
| with args.output.open("w+b") as output: |
| output.write(bytes(HEADER.size + table_bytes)) |
| output.write(names) |
| pad_to(output, payload_offset) |
|
|
| for item, name_offset, name, dims, source, aux, data_offset, data_bytes, aux_offset, aux_bytes in layouts: |
| pad_to(output, data_offset) |
| digest = hashlib.sha256() |
| if copy_file(source, output, digest) != data_bytes: |
| raise IOError(f"short read while packing {source}") |
| if aux is not None: |
| pad_to(output, aux_offset) |
| aux_digest = hashlib.sha256() |
| if copy_file(aux, output, aux_digest) != aux_bytes: |
| raise IOError(f"short read while packing {aux}") |
| padded_dims = dims + [0] * (4 - len(dims)) |
| entries.append(ENTRY.pack( |
| name_offset, |
| len(name), |
| enum_value(DTYPES, item["dtype"], "dtype"), |
| len(dims), |
| enum_value(ROLES, item.get("role", "unknown"), "role"), |
| enum_value(QUANTS, item.get("quant", "none"), "quantization"), |
| enum_value(LAYOUTS, item.get("layout", "row_major"), "layout"), |
| int(item.get("flags", 0)), |
| *padded_dims, |
| data_offset, |
| data_bytes, |
| aux_offset, |
| aux_bytes, |
| float(item.get("scale", 1.0)), |
| float(item.get("zero_point", 0.0)), |
| int(item.get("layer", 0xFFFFFFFF)), |
| int(item.get("op", 0)), |
| int(item.get("expert", 0xFFFFFFFF)), |
| int(item.get("core", 0xFFFFFFFF)), |
| digest.digest(), |
| bytes(24), |
| )) |
|
|
| pad_to(output, file_bytes) |
| header = make_header(manifest, len(entries), string_offset, len(names), payload_offset, |
| payload_bytes, file_bytes, 0) |
| crc = zlib.crc32(header) & 0xFFFFFFFF |
| header = make_header(manifest, len(entries), string_offset, len(names), payload_offset, |
| payload_bytes, file_bytes, crc) |
| output.seek(0) |
| output.write(header) |
| for entry in entries: |
| output.write(entry) |
| output.flush() |
| os.fsync(output.fileno()) |
|
|
| print(f"packed {len(entries)} tensors, {file_bytes} bytes -> {args.output}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
