Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions examples/models/llama/export_llama_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -855,16 +855,16 @@ def _to_edge_and_lower_llama_xnnpack(

# TODO: Enable generating ETRecord with XNNPack and to_edge_transform_and_lower().
if generate_etrecord:
raise NotImplementedError(
"export_llama does not support XNNPack and generating ETRecord at the moment."
)
builder_exported.generate_etrecord = True

builder = builder_exported.pt2e_quantize(quantizers).to_edge_transform_and_lower(
partitioners
)
if verbose:
print_delegation_info(builder.edge_manager.exported_program().graph_module)

# we need builder.export_program

return builder.to_executorch(passes=additional_passes)


Expand Down
3 changes: 3 additions & 0 deletions extension/llm/export/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ def __init__(
metadata: Optional[dict] = None,
dynamic_shapes: Optional[Any] = None,
save_exported_program: bool = False,
generate_etrecord: bool = False,
):
# Store necessary constructor arguments.
self.model = model
Expand All @@ -116,6 +117,7 @@ def __init__(
self.metadata = metadata
self.dynamic_shapes = dynamic_shapes
self.save_exported_program = save_exported_program
self.generate_etrecord = generate_etrecord

# Note: treat this as the source of truth for the result of
# torch.export'ing a model. If the overall ExportedProgram is needed,
Expand Down Expand Up @@ -481,6 +483,7 @@ def to_edge_transform_and_lower(
partitioner=partitioners,
compile_config=edge_config,
constant_methods=self.metadata,
generate_etrecord=self.generate_etrecord,
)
if self.verbose:
logging.info(f"Exported graph:\n{self.edge_manager.exported_program()}")
Expand Down
Loading