Skip to content

Revert export_for_training migration in llm/export/builder.py #6029

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
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 extension/llm/export/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@

from executorch.extension.export_util.utils import export_to_edge, save_pte_program
from executorch.extension.llm.tokenizer.utils import get_tokenizer
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
from torch.ao.quantization.quantizer import Quantizer
from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
from torch.export import export_for_training
from torch.nn.attention import SDPBackend

FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
Expand Down Expand Up @@ -190,9 +190,9 @@ def capture_pre_autograd_graph(self) -> "LLMEdgeManager":
strict=True,
).module()
else:
self.pre_autograd_graph_module = export_for_training(
self.pre_autograd_graph_module = capture_pre_autograd_graph(
self.model, self.example_inputs, dynamic_shapes=dynamic_shape
).module()
)

return self

Expand Down
Loading