Skip to content

fix(quantization): clear unused dense weights in graph mode finalized model - #101

Merged
pkmandke merged 3 commits into
apple:mainfrom
pkmandke:dev/clear_dense_weights_graph_mode
Sep 18, 2026
Merged

pkmandke merged 3 commits into
apple:mainfrom
pkmandke:dev/clear_dense_weights_graph_mode

Conversation

@pkmandke

Copy link
Copy Markdown
Member

Similar to eager mode, clear the unused dense weights from the finalized model in graph mode.

Details:

  • _process_mlir_weight_quantization registers the quantized weight, scale and offset as new buffers, but the dense float parameter stays registered on its owning module. It then survives into state_dict().
  • To release the dense weights, this PR keeps track of the dense get_attr nodes and only replaces them if they are not used anywhere in the graph.

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
raise FileExistsError(f"mmap_dir {mmap_dir!r} is non-empty. Pass an empty directory.")


def _zero_size_placeholder(tensor: torch.Tensor) -> torch.Tensor:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

refactored here so both eager and graph can use it

@pkmandke
pkmandke marked this pull request as ready for review September 16, 2026 20:04
@pkmandke
pkmandke requested review from crowbat and u-simha September 16, 2026 20:04
@pkmandke pkmandke added the bug Something isn't working label Sep 16, 2026
@pkmandke
pkmandke requested a review from vineet-g September 16, 2026 20:06
Comment thread src/coreai_opt/_utils/export_utils.py
Comment thread src/coreai_opt/quantization/_graph/_prepare_for_export.py Outdated
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
@pkmandke
pkmandke enabled auto-merge (squash) September 18, 2026 21:07
@pkmandke
pkmandke merged commit 293bf10 into apple:main Sep 18, 2026
14 checks passed
@pkmandke
pkmandke deleted the dev/clear_dense_weights_graph_mode branch September 18, 2026 22:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants