tzhouamandClaude Opus 5
67bee59fda
Reading the device profile of the first export showed a fifth of every
sliding layer went nowhere useful: Output 9.9%, its quantise conversion
3.1%, and the window Slice 7.4%. The layer took a 512-entry window in,
appended the new key, sliced the oldest off and handed the whole recomputed
window back out -- 21 of 28 layers paying to move a cache-sized tensor
across the graph boundary once per token. Full layers paid a similar 9.5%
concatenating the new key onto the cache.
Both are the same mistake, and both go away by joining in score space
instead of key space: score the cache and the new token separately, then
concatenate two vectors of length n and 1 rather than an n x 256 key matrix.
Nothing cache-sized is copied, and each layer returns just the new K/V entry
for the runtime to drop into its ring slot. Softmax over keys does not care
what order the slots are in and each cached key already carries its own
rotary phase, so a ring buffer needs no reordering -- only a mask for slots
not yet written. The sliding cache correspondingly holds window-1 entries,
since a 512-token window counts the current token (transformers' own sliding
cache holds 511 for the same reason).
Measured on a real Galaxy S25, fp16, ms per layer:
sliding 1.689 -> 1.483 (-12.2%)
full 1.559 -> 1.567 (unchanged: the concat saved is spent on the
extra score matmul)
Two shapes had to be written as matmuls rather than the obvious elementwise
forms -- the new token's score (broadcast multiply + ReduceSum) and its
value contribution (a two-axis broadcast multiply). Both compile at fp16 and
both make QNN's context-binary converter exit 14 under w4a16; as matmuls
they compile and are the cheaper op on HTP anyway.
Also adds a selectable GELU approximation, and the measurement that closes
the question: erf GELU is 11-13% of every layer, but tanh is *slower* on
this NPU (1.535 vs 1.483 ms) despite holding 65.2 dB. The exact form stays,
which is also the better-fidelity choice.
Parity against the reference implementation is 136.8 dB with identical
top-1/top-5, and a test pins ring and roll to the same output on a shared
history so the device numbers cannot silently measure a different model.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
490 lines
19 KiB
Python
490 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Export a Spark-X2.5 decode step for phone accelerators (QNN / TFLite / ORT).
|
|
|
|
The graph is one autoregressive step with the KV cache passed in and out as
|
|
explicit tensors, which is the shape every mobile NPU runtime wants: static
|
|
shapes, no control flow, no in-place cache mutation.
|
|
|
|
What makes Spark worth exporting this way is its hybrid attention. Three of
|
|
every four layers attend to a 512-token sliding window, so their cache is a
|
|
fixed 512 entries *no matter how long the conversation is*; only the seven
|
|
full-attention layers grow with context. At a 4k context that is 81 MiB of
|
|
KV instead of the 235 MiB a uniform-attention model of the same shape would
|
|
carry, and the gap widens linearly from there. The two layer types are
|
|
therefore exported with separately sized caches, masks and rotary angles --
|
|
sliding layers rotate all 256 head dims with theta 1e4, full layers only the
|
|
first 64 with theta 5e6.
|
|
|
|
``python -m vllm_omni.edge.spark_export --help`` exports one layer of each
|
|
type, or the whole 28-layer step, to ONNX.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm_omni.edge.qnn_export import sanitize_onnx
|
|
|
|
__all__ = [
|
|
"SparkStepConfig",
|
|
"SparkDecodeStep",
|
|
"config_from_spark",
|
|
"example_inputs",
|
|
"input_names",
|
|
"output_names",
|
|
"export_onnx",
|
|
"load_spark_weights",
|
|
"kv_cache_bytes",
|
|
]
|
|
|
|
SLIDING = "sliding_attention"
|
|
FULL = "full_attention"
|
|
|
|
|
|
@dataclass
|
|
class SparkStepConfig:
|
|
"""Shape of one Spark-X2.5 decode step."""
|
|
|
|
hidden_size: int
|
|
num_attention_heads: int
|
|
num_key_value_heads: int
|
|
head_dim: int
|
|
intermediate_size: int
|
|
vocab_size: int
|
|
sliding_window: int
|
|
layer_types: tuple[str, ...]
|
|
rope_theta: dict[str, float]
|
|
partial_rotary_factor: dict[str, float]
|
|
rms_norm_eps: float = 1e-6
|
|
include_lm_head: bool = True
|
|
gelu_mode: str = "exact"
|
|
"""``exact`` is the erf GELU the model was trained with (and the only one
|
|
its own config accepts). ``tanh`` and ``sigmoid`` are the usual cheaper
|
|
approximations, exposed because the erf form is 11-13% of every layer on
|
|
the Hexagon NPU -- adopt one only if its fidelity is measured."""
|
|
cache_layout: str = "ring"
|
|
"""``ring``: the cache is a fixed ring buffer the runtime writes into, and
|
|
the graph returns only the new K/V entry. ``roll``: the graph concatenates
|
|
the new entry, re-slices the window and hands the whole window back. Roll
|
|
was the first layout measured on device and is kept for A/B only -- it
|
|
spends 20% of a sliding layer moving the window in and out."""
|
|
first_layer: int = 0
|
|
"""Index of the first exported layer, so a subset can be profiled on its own."""
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.num_attention_heads % self.num_key_value_heads:
|
|
raise ValueError(
|
|
f"num_attention_heads {self.num_attention_heads} is not a multiple "
|
|
f"of num_key_value_heads {self.num_key_value_heads}"
|
|
)
|
|
bad = set(self.layer_types) - {SLIDING, FULL}
|
|
if bad:
|
|
raise ValueError(f"unknown layer types: {sorted(bad)}")
|
|
if self.cache_layout not in ("ring", "roll"):
|
|
raise ValueError(f"unknown cache_layout: {self.cache_layout}")
|
|
if self.gelu_mode not in ("exact", "tanh", "sigmoid"):
|
|
raise ValueError(f"unknown gelu_mode: {self.gelu_mode}")
|
|
|
|
@property
|
|
def group_size(self) -> int:
|
|
return self.num_attention_heads // self.num_key_value_heads
|
|
|
|
@property
|
|
def num_layers(self) -> int:
|
|
return len(self.layer_types)
|
|
|
|
def rotary_dim(self, layer_type: str) -> int:
|
|
return int(self.head_dim * self.partial_rotary_factor[layer_type])
|
|
|
|
def cache_len(self, layer_type: str, context: int) -> int:
|
|
"""KV entries a layer of this type holds at ``context`` tokens.
|
|
|
|
A 512-token window counts the current token, so a ring buffer holds
|
|
511 previous entries and the graph scores those plus the new one.
|
|
The roll layout instead keeps 512 and drops the oldest after
|
|
appending. (transformers' own sliding cache holds 511, same reason.)
|
|
"""
|
|
if layer_type != SLIDING:
|
|
return context
|
|
window = self.sliding_window - 1 if self.cache_layout == "ring" else self.sliding_window
|
|
return min(context, window)
|
|
|
|
|
|
def _gelu(x: torch.Tensor, mode: str) -> torch.Tensor:
|
|
if mode == "exact":
|
|
return torch.nn.functional.gelu(x)
|
|
if mode == "tanh":
|
|
return torch.nn.functional.gelu(x, approximate="tanh")
|
|
return x * torch.sigmoid(1.702 * x)
|
|
|
|
|
|
def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
|
|
return weight * (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps))
|
|
|
|
|
|
def _rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
|
"""Neox-style rotation of the leading ``cos.shape[-1]`` dims; the rest pass through."""
|
|
rot = cos.shape[-1]
|
|
if rot < x.shape[-1]:
|
|
x_rot, x_pass = x[..., :rot], x[..., rot:]
|
|
else:
|
|
x_rot, x_pass = x, None
|
|
half = rot // 2
|
|
x1, x2 = x_rot[..., :half], x_rot[..., half:]
|
|
out = x_rot * cos + torch.cat([-x2, x1], dim=-1) * sin
|
|
return out if x_pass is None else torch.cat([out, x_pass], dim=-1)
|
|
|
|
|
|
class _Layer(nn.Module):
|
|
def __init__(self, cfg: SparkStepConfig, layer_type: str):
|
|
super().__init__()
|
|
self.cfg = cfg
|
|
self.layer_type = layer_type
|
|
self.cache_layout = cfg.cache_layout
|
|
h, d = cfg.hidden_size, cfg.head_dim
|
|
q_dim = cfg.num_attention_heads * d
|
|
kv_dim = cfg.num_key_value_heads * d
|
|
self.input_layernorm = nn.Parameter(torch.ones(h))
|
|
self.post_attention_layernorm = nn.Parameter(torch.ones(h))
|
|
self.q_k_v_proj = nn.Linear(h, q_dim + 2 * kv_dim, bias=False)
|
|
self.g_proj = nn.Linear(h, cfg.num_attention_heads, bias=False)
|
|
self.out_proj = nn.Linear(q_dim, h, bias=False)
|
|
self.gate_proj = nn.Linear(h, cfg.intermediate_size, bias=False)
|
|
self.up_proj = nn.Linear(h, cfg.intermediate_size, bias=False)
|
|
self.down_proj = nn.Linear(cfg.intermediate_size, h, bias=False)
|
|
|
|
def forward(self, x, cos, sin, k_cache, v_cache, mask):
|
|
cfg = self.cfg
|
|
heads, kv, d, g = (
|
|
cfg.num_attention_heads,
|
|
cfg.num_key_value_heads,
|
|
cfg.head_dim,
|
|
cfg.group_size,
|
|
)
|
|
h = _rms(x, self.input_layernorm, cfg.rms_norm_eps)
|
|
qkv = self.q_k_v_proj(h)
|
|
q = qkv[..., : heads * d].view(1, heads, 1, d)
|
|
k = qkv[..., heads * d : (heads + kv) * d].view(1, kv, 1, d)
|
|
v = qkv[..., (heads + kv) * d :].view(1, kv, 1, d)
|
|
gate = torch.sigmoid(self.g_proj(h)).view(1, heads, 1, 1)
|
|
q, k = _rope(q, cos, sin), _rope(k, cos, sin)
|
|
|
|
# GQA without materialising expanded K/V: fold the query heads that
|
|
# share a kv head into their own axis.
|
|
scale = 1.0 / math.sqrt(d)
|
|
qg = q.view(1, kv, g, d)
|
|
|
|
if self.cache_layout == "ring":
|
|
# Score the cache and the new token separately, then join in score
|
|
# space. The join is over a length-(N+1) vector per head instead of
|
|
# an N x head_dim key matrix, so nothing the size of the cache is
|
|
# ever copied, and the layer returns one entry for the runtime to
|
|
# drop into the ring slot. Softmax over keys does not care what
|
|
# order the slots are in, and each cached key already carries its
|
|
# own rotary phase, so a ring buffer needs no reordering -- only a
|
|
# mask marking slots that have not been written yet.
|
|
# Both score terms are matmuls against a key matrix -- the cache
|
|
# (n entries) and the new token (1). Writing the second as a
|
|
# broadcast multiply plus ReduceSum is equivalent but QNN's
|
|
# context-binary converter rejects that shape at w4a16.
|
|
s_cache = torch.matmul(qg, k_cache.transpose(2, 3)) * scale
|
|
s_new = torch.matmul(qg, k.transpose(2, 3)) * scale
|
|
scores = torch.cat([s_cache, s_new], dim=-1) + mask
|
|
w = torch.softmax(scores, dim=-1)
|
|
n = k_cache.shape[2]
|
|
w_cache, w_new = w[..., :n], w[..., n:]
|
|
# The new token's contribution is an outer product, written as a
|
|
# matmul rather than a broadcast multiply: [g,1] @ [1,d] -> [g,d].
|
|
# The broadcast form (which broadcasts along two axes at once)
|
|
# compiles at fp16 but fails QNN's context-binary converter under
|
|
# w4a16, and a matmul is the cheaper op on HTP anyway.
|
|
attn = torch.matmul(w_cache, v_cache) + torch.matmul(w_new, v)
|
|
k_out, v_out = k, v
|
|
else:
|
|
# Roll the window: the oldest entry falls off the front.
|
|
keys = torch.cat([k_cache, k], dim=2)
|
|
vals = torch.cat([v_cache, v], dim=2)
|
|
if self.layer_type == SLIDING:
|
|
keys, vals = keys[:, :, 1:], vals[:, :, 1:]
|
|
scores = torch.matmul(qg, keys.transpose(2, 3)) * scale + mask
|
|
attn = torch.matmul(torch.softmax(scores, dim=-1), vals)
|
|
k_out = keys if self.layer_type == SLIDING else k
|
|
v_out = vals if self.layer_type == SLIDING else v
|
|
|
|
attn = (attn.view(1, heads, 1, d) * gate).view(1, 1, heads * d)
|
|
|
|
x = x + self.out_proj(attn)
|
|
h = _rms(x, self.post_attention_layernorm, cfg.rms_norm_eps)
|
|
x = x + self.down_proj(_gelu(self.gate_proj(h), cfg.gelu_mode) * self.up_proj(h))
|
|
return x, k_out, v_out
|
|
|
|
|
|
class SparkDecodeStep(nn.Module):
|
|
"""One decode step.
|
|
|
|
Under the default ``ring`` cache layout every layer returns just the new
|
|
K/V entry and the runtime drops it into the ring slot for this position.
|
|
Under ``roll`` a sliding layer hands back its whole recomputed window,
|
|
which costs a fifth of the layer on device.
|
|
"""
|
|
|
|
def __init__(self, cfg: SparkStepConfig):
|
|
super().__init__()
|
|
self.cfg = cfg
|
|
self.layers = nn.ModuleList(_Layer(cfg, t) for t in cfg.layer_types)
|
|
if cfg.include_lm_head:
|
|
self.norm = nn.Parameter(torch.ones(cfg.hidden_size))
|
|
self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
|
|
|
|
def forward(self, x, cos_sw, sin_sw, cos_full, sin_full, mask_sw, mask_full, *caches):
|
|
cfg = self.cfg
|
|
if len(caches) != 2 * cfg.num_layers:
|
|
raise ValueError(
|
|
f"expected {2 * cfg.num_layers} cache tensors, got {len(caches)}"
|
|
)
|
|
new: list[torch.Tensor] = []
|
|
for i, layer in enumerate(self.layers):
|
|
sliding = cfg.layer_types[i] == SLIDING
|
|
cos, sin = (cos_sw, sin_sw) if sliding else (cos_full, sin_full)
|
|
mask = mask_sw if sliding else mask_full
|
|
x, k, v = layer(x, cos, sin, caches[2 * i], caches[2 * i + 1], mask)
|
|
new += [k, v]
|
|
if not cfg.include_lm_head:
|
|
return (x, *new)
|
|
logits = self.lm_head(_rms(x, self.norm, cfg.rms_norm_eps))[:, 0]
|
|
return (logits, *new)
|
|
|
|
|
|
def config_from_spark(
|
|
hf_config: dict[str, Any],
|
|
layer_slice: slice | None = None,
|
|
include_lm_head: bool = True,
|
|
cache_layout: str = "ring",
|
|
gelu_mode: str = "exact",
|
|
) -> SparkStepConfig:
|
|
"""Build a step config from a Spark ``config.json``.
|
|
|
|
``layer_slice`` selects a contiguous subset of layers, which is how a
|
|
single sliding or full layer gets profiled on its own.
|
|
"""
|
|
types = list(hf_config["layer_types"])
|
|
first = 0
|
|
if layer_slice is not None:
|
|
first = layer_slice.start or 0
|
|
types = types[layer_slice]
|
|
rope = hf_config["rope_parameters"]
|
|
return SparkStepConfig(
|
|
hidden_size=hf_config["hidden_size"],
|
|
num_attention_heads=hf_config["num_attention_heads"],
|
|
num_key_value_heads=hf_config["num_key_value_heads"],
|
|
head_dim=hf_config["head_dim"],
|
|
intermediate_size=hf_config["intermediate_size"],
|
|
vocab_size=hf_config["vocab_size"],
|
|
sliding_window=hf_config["sliding_window"],
|
|
layer_types=tuple(types),
|
|
rope_theta={k: float(v.get("rope_theta", 10000)) for k, v in rope.items()},
|
|
partial_rotary_factor={
|
|
k: float(v.get("partial_rotary_factor", 1.0)) for k, v in rope.items()
|
|
},
|
|
rms_norm_eps=hf_config.get("rms_norm_eps", 1e-6),
|
|
include_lm_head=include_lm_head,
|
|
cache_layout=cache_layout,
|
|
gelu_mode=gelu_mode,
|
|
first_layer=first,
|
|
)
|
|
|
|
|
|
def input_names(cfg: SparkStepConfig) -> list[str]:
|
|
return [
|
|
"x",
|
|
"cos_sw",
|
|
"sin_sw",
|
|
"cos_full",
|
|
"sin_full",
|
|
"mask_sw",
|
|
"mask_full",
|
|
] + [f"{kv}_cache_{i}" for i in range(cfg.num_layers) for kv in ("k", "v")]
|
|
|
|
|
|
def output_names(cfg: SparkStepConfig) -> list[str]:
|
|
head = "logits" if cfg.include_lm_head else "hidden"
|
|
return [head] + [
|
|
f"{kv}_new_{i}" for i in range(cfg.num_layers) for kv in ("k", "v")
|
|
]
|
|
|
|
|
|
def _angles(position: int, rotary_dim: int, theta: float, dtype: torch.dtype):
|
|
inv = 1.0 / (
|
|
theta ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
|
|
)
|
|
ang = position * inv
|
|
cos = torch.cat([ang.cos(), ang.cos()]).view(1, 1, 1, rotary_dim).to(dtype)
|
|
sin = torch.cat([ang.sin(), ang.sin()]).view(1, 1, 1, rotary_dim).to(dtype)
|
|
return cos, sin
|
|
|
|
|
|
def example_inputs(
|
|
cfg: SparkStepConfig,
|
|
context: int,
|
|
position: int | None = None,
|
|
dtype: torch.dtype = torch.float32,
|
|
) -> tuple[torch.Tensor, ...]:
|
|
"""Concrete inputs for tracing at a given context length."""
|
|
pos = context if position is None else position
|
|
sw_len = cfg.cache_len(SLIDING, context)
|
|
full_len = cfg.cache_len(FULL, context)
|
|
cos_sw, sin_sw = _angles(pos, cfg.rotary_dim(SLIDING), cfg.rope_theta[SLIDING], dtype)
|
|
cos_full, sin_full = _angles(pos, cfg.rotary_dim(FULL), cfg.rope_theta[FULL], dtype)
|
|
# Ring layout scores every cache slot plus the new token, for both layer
|
|
# types. Roll layout drops the oldest sliding slot, so its row is sw_len.
|
|
if cfg.cache_layout == "ring":
|
|
mask_sw = torch.zeros(1, 1, 1, sw_len + 1, dtype=dtype)
|
|
else:
|
|
mask_sw = torch.zeros(1, 1, 1, sw_len, dtype=dtype)
|
|
mask_full = torch.zeros(1, 1, 1, full_len + 1, dtype=dtype)
|
|
x = (torch.randn(1, 1, cfg.hidden_size) * 0.1).to(dtype)
|
|
caches: list[torch.Tensor] = []
|
|
for t in cfg.layer_types:
|
|
n = cfg.cache_len(t, context)
|
|
for _ in range(2):
|
|
caches.append(
|
|
torch.randn(1, cfg.num_key_value_heads, n, cfg.head_dim, dtype=dtype)
|
|
* 0.5
|
|
)
|
|
return (x, cos_sw, sin_sw, cos_full, sin_full, mask_sw, mask_full, *caches)
|
|
|
|
|
|
def kv_cache_bytes(cfg: SparkStepConfig, context: int, bytes_per_elem: int = 2) -> dict[str, int]:
|
|
"""KV footprint of the hybrid layout versus an all-full-attention model."""
|
|
per_entry = 2 * cfg.num_key_value_heads * cfg.head_dim * bytes_per_elem
|
|
hybrid = sum(cfg.cache_len(t, context) for t in cfg.layer_types) * per_entry
|
|
uniform = cfg.num_layers * context * per_entry
|
|
return {"hybrid": hybrid, "uniform": uniform, "saved": uniform - hybrid}
|
|
|
|
|
|
def load_spark_weights(
|
|
module: SparkDecodeStep, weights_dir: str | Path, dtype: torch.dtype = torch.float32
|
|
) -> int:
|
|
"""Load real Spark weights into the exported step. Returns tensors copied."""
|
|
from safetensors.torch import load_file
|
|
|
|
weights_dir = Path(weights_dir)
|
|
shards = sorted(weights_dir.glob("*.safetensors"))
|
|
if not shards:
|
|
raise FileNotFoundError(f"no safetensors under {weights_dir}")
|
|
state: dict[str, torch.Tensor] = {}
|
|
for shard in shards:
|
|
state.update(load_file(str(shard)))
|
|
|
|
cfg = module.cfg
|
|
copied = 0
|
|
with torch.no_grad():
|
|
for i, layer in enumerate(module.layers):
|
|
src = f"model.layers.{cfg.first_layer + i}."
|
|
for dst_name, key in (
|
|
("input_layernorm", "input_layernorm.weight"),
|
|
("post_attention_layernorm", "post_attention_layernorm.weight"),
|
|
):
|
|
getattr(layer, dst_name).copy_(state[src + key].to(dtype))
|
|
copied += 1
|
|
for mod, key in (
|
|
(layer.q_k_v_proj, "self_attn.q_k_v_proj.weight"),
|
|
(layer.g_proj, "self_attn.g_proj.weight"),
|
|
(layer.out_proj, "self_attn.out_proj.weight"),
|
|
(layer.gate_proj, "mlp.gate_proj.weight"),
|
|
(layer.up_proj, "mlp.up_proj.weight"),
|
|
(layer.down_proj, "mlp.down_proj.weight"),
|
|
):
|
|
mod.weight.copy_(state[src + key].to(dtype))
|
|
copied += 1
|
|
if cfg.include_lm_head:
|
|
module.norm.copy_(state["model.norm.weight"].to(dtype))
|
|
# Spark ties the output projection to the input embedding.
|
|
emb = state.get("lm_head.weight", state["model.embedding.weight"])
|
|
module.lm_head.weight.copy_(emb.to(dtype))
|
|
copied += 2
|
|
return copied
|
|
|
|
|
|
def export_onnx(
|
|
module: SparkDecodeStep, context: int, path: str | Path, opset: int = 18
|
|
) -> dict[str, Any]:
|
|
"""Trace to ONNX at a fixed context length and sanitize for Qualcomm tools."""
|
|
cfg = module.cfg
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
args = example_inputs(cfg, context)
|
|
module.eval()
|
|
with torch.no_grad():
|
|
torch.onnx.export(
|
|
module,
|
|
args,
|
|
str(path),
|
|
input_names=input_names(cfg),
|
|
output_names=output_names(cfg),
|
|
opset_version=opset,
|
|
dynamo=False,
|
|
)
|
|
stats = sanitize_onnx(path)
|
|
return {
|
|
"path": str(path),
|
|
"context": context,
|
|
"layers": cfg.num_layers,
|
|
"layer_types": list(cfg.layer_types),
|
|
"include_lm_head": cfg.include_lm_head,
|
|
"sanitized": stats,
|
|
"kv_bytes": kv_cache_bytes(cfg, context),
|
|
}
|
|
|
|
|
|
def _main(argv: list[str] | None = None) -> int:
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument("--config", required=True, help="path to Spark config.json")
|
|
ap.add_argument("--weights", default=None, help="directory of safetensors shards")
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--context", type=int, default=1024)
|
|
ap.add_argument(
|
|
"--layers",
|
|
default="all",
|
|
help="'all', 'sliding', 'full', or START:STOP",
|
|
)
|
|
ap.add_argument("--no-lm-head", action="store_true")
|
|
ap.add_argument("--cache-layout", default="ring", choices=["ring", "roll"])
|
|
ap.add_argument("--gelu", default="exact", choices=["exact", "tanh", "sigmoid"])
|
|
a = ap.parse_args(argv)
|
|
|
|
hf = json.loads(Path(a.config).read_text())
|
|
types = hf["layer_types"]
|
|
if a.layers == "all":
|
|
sl = None
|
|
elif a.layers == "sliding":
|
|
i = types.index(SLIDING)
|
|
sl = slice(i, i + 1)
|
|
elif a.layers == "full":
|
|
i = types.index(FULL)
|
|
sl = slice(i, i + 1)
|
|
else:
|
|
start, stop = (int(v) for v in a.layers.split(":"))
|
|
sl = slice(start, stop)
|
|
|
|
cfg = config_from_spark(hf, sl, include_lm_head=not a.no_lm_head,
|
|
cache_layout=a.cache_layout, gelu_mode=a.gelu)
|
|
step = SparkDecodeStep(cfg)
|
|
if a.weights:
|
|
print(f"loaded {load_spark_weights(step, a.weights)} tensors")
|
|
print(json.dumps(export_onnx(step, a.context, a.out), indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover
|
|
raise SystemExit(_main())
|