Skip to content

fix(quantization): align eager-mode op call counts between capture and optimization passes - #78

Merged
pkmandke merged 5 commits into
apple:mainfrom
pkmandke:dev/fix_eager_act_quant_pass_counts
Aug 26, 2026
Merged

fix(quantization): align eager-mode op call counts between capture and optimization passes#78
pkmandke merged 5 commits into
apple:mainfrom
pkmandke:dev/fix_eager_act_quant_pass_counts

Conversation

@pkmandke

@pkmandke pkmandke commented Aug 26, 2026

Copy link
Copy Markdown
Member

An op call whose arguments hold no quantizable tensor is not counted during capture (modes.py:798) but is counted during the optimization pass (modes.py:1123).
Both passes run inside Quantizer.prepare(): the capture forward at handler.py:63, the optimization forward at quantizer.py:207. The counts then disagree and RuntimeError is raised (registered_optimizers_tracker.py:214-221).

"No quantizable tensor" here means no argument is a float tensor with more than one element, per is_optimizable_tensor (the torch.numel(tensor) <= 1 clause). This is what Qwen2 hits in the toy config in coreai-models (please see this PR), it has no head_dim, so RoPE's freqs * scale multiplies a single-element tensor by a Python float.

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
@pkmandke pkmandke self-assigned this Aug 26, 2026
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
@pkmandke pkmandke added the bug Something isn't working label Aug 26, 2026
@pkmandke
pkmandke marked this pull request as ready for review August 26, 2026 20:28
@pkmandke
pkmandke requested review from crowbat and u-simha August 26, 2026 20:28
@pkmandke pkmandke removed their assignment Aug 26, 2026
@pkmandke pkmandke changed the title [WIP] fix(quantization): align eager-mode op call counts between capture and optimization passes Aug 26, 2026
Comment thread src/coreai_opt/_utils/insertion/torch_function/modes.py Outdated
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
@pkmandke
pkmandke merged commit ee85c40 into apple:main Aug 26, 2026
13 checks passed
@pkmandke
pkmandke deleted the dev/fix_eager_act_quant_pass_counts branch August 27, 2026 20:31
pkmandke added a commit that referenced this pull request Aug 28, 2026
---------

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

2 participants