Merge, prune, and export a decoder ONNX Runtime Web can load.
The browser wants the graph shape transformers.js expects: a merged decoder with
past-key inputs, exported at an opset ORT Web actually implements. We produce
that with Optimum, having first merged any LoRA adapter and -- the part that
matters for speed -- sliced the unembedding down to the answer letters.
Everything heavy (torch, transformers, peft, optimum, onnx) is
imported inside the function that needs it, so import typedecide stays fast.
Read docs/SPEED.md before tuning anything here. In particular: a stock export
returns logits for every position, which on a 512-token prefill is a
311 MB fp32 tensor copied out of the GPU on one call.
trim_logits_to_last_position is the fix, and pruning is what shrinks what is
left.
module-attribute
ONNX_TASK = 'text-generation-with-past'
module-attribute
MANIFEST_FILENAME = 'export_manifest.json'
module-attribute
STAGING_DIRNAME = '_merged'
module-attribute
dataclass
ExportConfig(
base_model: str,
adapter: Path | None,
output_dir: Path,
opset: int = 17,
prune_head_to: Sequence[int] | None = None,
keep_merged: bool = False,
)
One ONNX export.
prune_head_to is the list of original vocabulary ids to keep, in answer
order -- normally answer_letter_token_ids(base_model, n_options). Leave it
None to export the full unembedding. When it is set, the export also writes
slot_map.json next to the model, and without that file the resulting
weights are uninterpretable: the readout must index rows 0..K-1, not
vocabulary ids.
keep_merged keeps the intermediate merged checkpoint on disk. It is large
and only useful for debugging a suspicious export.
instance-attribute
instance-attribute
instance-attribute
class-attribute
instance-attribute
class-attribute
instance-attribute
prune_head_to: Sequence[int] | None = None
class-attribute
instance-attribute
keep_merged: bool = False
to_manifest() -> dict[str, Any]
Source code in src/typedecide/export/onnx.py
| def to_manifest(self) -> dict[str, Any]:
return {
"base_model": self.base_model,
"adapter": None if self.adapter is None else str(self.adapter),
"output_dir": str(self.output_dir),
"opset": self.opset,
"prune_head_to": None if self.prune_head_to is None else list(self.prune_head_to),
"keep_merged": self.keep_merged,
}
|
Export the model and return the directory holding it.
Order matters. The adapter is merged first, because pruning a model that
still has an adapter attached would prune the base weights and leave the
adapter pointing at rows that no longer exist. Pruning happens before the
ONNX trace, because the trace freezes the output width into the graph.
Source code in src/typedecide/export/onnx.py
| def export_onnx(config: ExportConfig) -> Path:
"""Export the model and return the directory holding it.
Order matters. The adapter is merged first, because pruning a model that
still has an adapter attached would prune the base weights and leave the
adapter pointing at rows that no longer exist. Pruning happens before the
ONNX trace, because the trace freezes the output width into the graph.
"""
torch = _require("torch")
transformers = _require("transformers")
started = _now()
output_dir = config.output_dir
output_dir.mkdir(parents=True, exist_ok=True)
log.info("loading %s", config.base_model)
tokenizer = _load_tokenizer(transformers, config.base_model)
model = _load_model(transformers, torch, config.base_model)
if config.adapter is not None:
if not config.adapter.exists():
raise ExportError(
f"No adapter at {config.adapter}. Point `adapter` at the directory "
"`finetune` wrote (it holds adapter_config.json), or pass None to export "
"the base model."
)
peft = _require("peft")
log.info("merging adapter %s", config.adapter)
try:
model = peft.PeftModel.from_pretrained(model, str(config.adapter)).merge_and_unload()
except Exception as error: # noqa: BLE001 - peft raises a wide variety here
raise ExportError(
f"Merging {config.adapter} into {config.base_model} failed: {error}. The "
"adapter and the base model must be the same architecture; check "
"manifest.json beside the adapter for the base model it was trained on."
) from error
slot_map: SlotMap | None = None
if config.prune_head_to is not None:
slot_map = _prune(model, tokenizer, config)
write_slot_map(slot_map, output_dir)
staging = output_dir / STAGING_DIRNAME
log.info("staging the merged checkpoint at %s", staging)
model.save_pretrained(staging)
tokenizer.save_pretrained(staging)
try:
_run_optimum_export(staging, output_dir, config.opset)
manifest = {
"config": config.to_manifest(),
"task": ONNX_TASK,
"slot_map": None if slot_map is None else slot_map.to_dict(),
"pruned": slot_map is not None,
"logits_are_row_indices": slot_map is not None,
"can_generate_text": slot_map is None,
"versions": _versions(),
"started": started,
"finished": _now(),
}
(output_dir / MANIFEST_FILENAME).write_text(
json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
)
finally:
if not config.keep_merged:
shutil.rmtree(staging, ignore_errors=True)
log.info("exported to %s", output_dir)
return output_dir
|
trim_logits_to_last_position(
model_path: Path | str,
*,
output_path: Path | str | None = None
) -> Path
Make the graph return only the final position's logits.
A stock Optimum export returns [batch, seq, vocab] for the whole sequence.
A typed readout reads exactly one position, so on a prefill every other
position is computed, serialised, and discarded. At 512 tokens and fp32 that
is 512 x 151936 x 4 = 311,168,512 bytes copied out of the GPU on a single
call. Appending a Slice on the sequence axis turns that into one position.
Decode steps already have seq == 1, so this changes nothing there; it is a
prefill fix. Downstream code that slices the last position itself keeps
working, because slicing the last of one position is a no-op.
Returns the path written. Needs onnx.
Source code in src/typedecide/export/onnx.py
| def trim_logits_to_last_position(
model_path: Path | str, *, output_path: Path | str | None = None
) -> Path:
"""Make the graph return only the final position's logits.
A stock Optimum export returns `[batch, seq, vocab]` for the whole sequence.
A typed readout reads exactly one position, so on a prefill every other
position is computed, serialised, and discarded. At 512 tokens and fp32 that
is 512 x 151936 x 4 = 311,168,512 bytes copied out of the GPU on a single
call. Appending a `Slice` on the sequence axis turns that into one position.
Decode steps already have `seq == 1`, so this changes nothing there; it is a
prefill fix. Downstream code that slices the last position itself keeps
working, because slicing the last of one position is a no-op.
Returns the path written. Needs `onnx`.
"""
onnx = _require("onnx")
numpy = _require("numpy")
from onnx import helper, numpy_helper
source = _resolve_model_file(Path(model_path))
destination = Path(output_path) if output_path is not None else source
graph = (proto := onnx.load(str(source))).graph
target = next((o for o in graph.output if o.name == "logits"), None)
if target is None:
if not graph.output:
raise ExportError(f"{source} has no graph outputs to trim.")
target = graph.output[0]
log.warning("no output named 'logits'; trimming %r instead", target.name)
name = target.name
inner = f"{name}_all_positions"
dims = target.type.tensor_type.shape.dim
if len(dims) != 3:
raise ExportError(
f"Expected {name!r} to be a 3-D [batch, sequence, vocab] tensor; it has "
f"{len(dims)} dimensions. This graph is not shaped the way a decoder export is, "
"so trimming the sequence axis would cut the wrong thing."
)
full = helper.make_tensor_value_info(
inner, target.type.tensor_type.elem_type, [None, None, None]
)
for node in graph.node:
for i, produced in enumerate(node.output):
if produced == name:
node.output[i] = inner
for i, consumed in enumerate(node.input):
if consumed == name:
node.input[i] = inner
constants = [
numpy_helper.from_array(numpy.array([-1], dtype=numpy.int64), "typedecide_trim_starts"),
numpy_helper.from_array(
numpy.array([numpy.iinfo(numpy.int64).max], dtype=numpy.int64), "typedecide_trim_ends"
),
numpy_helper.from_array(numpy.array([1], dtype=numpy.int64), "typedecide_trim_axes"),
]
graph.initializer.extend(constants)
graph.value_info.append(full)
graph.node.append(
helper.make_node(
"Slice",
[inner, *(c.name for c in constants)],
[name],
name="typedecide_trim_last_position",
)
)
dims[1].ClearField("dim_param")
dims[1].dim_value = 1
try:
onnx.checker.check_model(proto, full_check=False)
except Exception as error: # noqa: BLE001 - onnx raises its own checker errors
raise ExportError(
f"The trimmed graph does not pass onnx.checker: {error}. Nothing was written; "
f"{source} is untouched."
) from error
destination.parent.mkdir(parents=True, exist_ok=True)
_save(onnx, proto, destination)
log.info("trimmed %s to the last position -> %s", source, destination)
return destination
|