Skip to content

Commit c0b4361

Browse files
JacobSzwejbkafacebook-github-bot
authored andcommitted
update export_llama (#1565)
Summary: Pull Request resolved: #1565 title Reviewed By: larryliu0820 Differential Revision: D52634973 fbshipit-source-id: 5cdbc873c6b78f15ae1d805dac9b674df2b09f85
1 parent 52a4cec commit c0b4361

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

examples/models/llama2/export_llama.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
import torch
1414
from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
15+
from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
1516

1617
from ...portable.utils import export_to_edge, save_pte_program
1718

@@ -54,7 +55,14 @@ def main() -> None:
5455
)
5556

5657
export_program = edge_manager.to_executorch(
57-
ExecutorchBackendConfig(extract_constant_segment=True)
58+
ExecutorchBackendConfig(
59+
extract_constant_segment=True,
60+
sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),
61+
)
62+
)
63+
print(
64+
"Required memory for activation in bytes: ",
65+
export_program._emitter_output.program.execution_plan[0].non_const_buffer_sizes,
5866
)
5967
save_pte_program(export_program.buffer, "llama2", args.output_dir)
6068
# model.forward(input)

0 commit comments

Comments
 (0)