Skip to content

Fix FP16 weight detection in NeuralNetworkBuilder.make_updatable (only first layer was checked) - #2858

Open
Anai-Guo wants to merge 1 commit into
apple:mainfrom
Anai-Guo:fix-fp16-updatable-check
Open

Anai-Guo wants to merge 1 commit into
apple:mainfrom
Anai-Guo:fix-fp16-updatable-check

Conversation

@Anai-Guo

Copy link
Copy Markdown

Summary

NeuralNetworkBuilder.make_updatable() is supposed to refuse models that contain FP16 weight params, but _check_fp16_weight_param_exists() only ever looked at the first layer: its return False was indented inside the for layer in layers: loop.

So a model like [ip1 (FP32), ip2 (FP16)] was silently marked updatable.

This PR also fixes two more bugs in the same function:

  • uniDirectionalLSTM: it did return self._check_fp16_weight_params_lstms(...), so the scan stopped at the first LSTM layer even when that LSTM was FP32.
  • branch / loop: it read .float16Value from ifBranch / elseBranch / conditionNetwork / bodyNetwork. Those are nested NeuralNetwork messages with no such field, so this raised AttributeError. It now recurses into their .layers.

Why the test suite didn't catch it

MLModelUpdatableTest.test_nn_partial_fp16_make_updatable_fail was defined twice. The second definition replaced the first, so the first one never ran. That first one covers exactly this case: quantize every layer except ip1 to FP16, then expect make_updatable to raise.

The second definition tests something else (quantizing an already-updatable model), so I renamed it to test_nn_updatable_fp16_quantize_fail. Both tests run now.

Verification

Linux, coremltools==9.0 wheel (its builder.py is byte-identical to main), with test_model_updatable.py from this branch:

pytest test_model_updatable.py
before (old builder.py) 1 failed, 16 passed: test_nn_partial_fp16_make_updatable_fail, AssertionError: ValueError not raised
after 17 passed

I also called _check_fp16_weight_param_exists directly on hand-built layer lists:

layers before after
[fp32 ip, fp16 ip] False True
[fp32 LSTM, fp16 ip] False True
[branch containing fp16 ip] AttributeError True
[loop containing fp32 ip] AttributeError False
[fp32 ip, loop containing fp32 ip] False False

🤖 Generated with Claude Code

`_check_fp16_weight_param_exists` had its `return False` indented inside
the `for layer in layers` loop, so only the first layer was ever
inspected. A model whose first layer is FP32 but a later layer holds
FP16 weights was wrongly allowed to be marked updatable.

The same function also returned early for any `uniDirectionalLSTM`
layer (even when it had no FP16 params), and for `branch`/`loop`
layers it read `.float16Value` off nested `NeuralNetwork` messages,
which raises AttributeError. Recurse into the nested networks instead.

`test_nn_partial_fp16_make_updatable_fail` was defined twice in
MLModelUpdatableTest, so the first definition -- which covers exactly
this case -- never ran (and fails against the old check). Rename the
second definition to `test_nn_updatable_fp16_quantize_fail`, which
matches what it tests.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant