Fix: prevent model graph and kwargs mutation in AutoModel.export() /修复 AutoModel 连续导出导致的计算图破坏与字典污染Fixes a critical destructive bug in AutoModel.export().#2835
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves a significant issue in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request addresses a critical bug in AutoModel.export() that caused crashes on consecutive calls. The fix correctly isolates the model and its configuration using deep copies, which prevents in-place mutations of self.model and self.kwargs. The change also includes a safer way to handle dictionary key deletion, improving robustness. My review includes a minor suggestion to improve code style by removing a redundant import.
Fixes a critical destructive bug in
AutoModel.export(). Previously, callingexport()modified bothself.model(by in-place replacing layers with ONNX-specific modules likeSANMEncoderExport) andself.kwargs. This caused consecutiveexport()calls (e.g., exporting FP32 then Int8) to crash with eitherKeyError: 'model'orAttributeError: 'SANMEncoderExport' object has no attribute 'encoders'.Implemented a dual deepcopy strategy:
model = copy.deepcopy(self.model): Isolates the graph mutation, ensuring the original PyTorch model remains untouched.kwargs = copy.deepcopy(self.kwargs): Prevents parameter reference pollution.修复了
AutoModel.export()中破坏性的 Bug。原逻辑在导出时会就地(in-place)修改self.model的计算图(替换为SANMEncoderExport等 ONNX 算子)并污染self.kwargs。这导致同一脚本连续调用export()时必定崩溃,抛出KeyError或AttributeError。采用了双重深拷贝策略:
model = copy.deepcopy(self.model):对模型进行深拷贝,隔离 ONNX 算子替换对原计算图的物理破坏。kwargs = copy.deepcopy(self.kwargs):深拷贝配置字典,防止引用污染。Reproduction bug/ 复现bug