Repository navigation
Core ML: run ops that read uint8 and int8 inputs on the CPU, refuse ops coremltools gets wrong, and fix fills, tanh GELU, RMSNorm and repeated outputs - #24317
Open
shoumikhin wants to merge 38 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24317
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 1 PendingAs of commit c315323 with merge base 681e493 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Contributor
Author
|
@claude /pr-review |
meta-cla
Bot
added
the
CLA Signed
|
Claude finished @shoumikhin's task in 28m 3s —— View job Summary: This fixes the reported uint8 camera-frame case, but some uint8/int8 inputs still fail and the new tests leave most of the claimed behaviour unchecked. A Detailed reviewPR Review: #24317SummaryThe PR keeps ops that read a uint8/int8 tensor from outside their Core ML partition on the CPU.
The problems:
Nothing was run for this review, because this job is not allowed to run Code Quality
Backward Compatibility
Testing
RecommendationRequest Changes
|
Contributor
Author
|
@metascroy @cymbalrush @YifanShenSZ could you take a look? This changes the Core ML backend. |
shoumikhin
force-pushed
the
coreml-uint8-inputs
branch
from
October 30, 2026 04:55
245afe3 to
8d3480d
Compare
shoumikhin
force-pushed
the
coreml-uint8-inputs
branch
from
October 30, 2026 05:20
8d3480d to
61600ea
Compare
shoumikhin
force-pushed
the
coreml-uint8-inputs
branch
from
October 30, 2026 05:21
61600ea to
0863813
Compare
Contributor
Author
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 42m 31s —— View job This keeps ops that read uint8/int8 tensors on the CPU as intended, but at the default deployment target it can also move a state update to the CPU while a Core ML delegate keeps reading a frozen copy of that state, so the recommendation is Request Changes. Detailed reviewPR Review: #24317SummaryThe PR makes the Core ML partitioner run on the CPU two kinds of ops: those that read a uint8/int8 tensor from outside their partition, and those in any partition this leaves with no tensor input. It also makes Code Quality
Testing
Performance
RecommendationRequest Changes The stale-state case gives wrong results with no error. It happens with a uint8/int8 input, state and the default deployment target, a combination that crashed with a
|
shoumikhin
changed the title
shoumikhin
force-pushed
the
coreml-uint8-inputs
branch
from
October 30, 2026 14:51
0863813 to
fb99778
Compare
Contributor
Author
|
@claude /pr-review |
Contributor
Author
|
@claude /pr-review |
coremltools 9.0 often ignores the dtype of an op that fills a tensor
with one value. torch.full((1, 16), 0) has no dtype, so torch makes it
int64, but coremltools makes it float. zeros and new_zeros come out
float even with an int dtype, ones_like does from iOS 16, and
full_like, fill and scalar_tensor take the type of their fill value.
An int or bool tensor made as float breaks in two ways. An op that does
not promote types, such as index_put through a bool mask, fails to
lower ("updates has dtype int32 whereas data has dtype fp32"). The
SmolVLM vision embeddings hit this. Other ops run at float16 precision
and silently lose integers above 2048.
The partitioner now refuses those fills when they make an int or bool
tensor. full, full_like, fill and scalar_tensor run on the portable
kernels. zeros, new_zeros, zeros_like, ones_like and new_full are no
longer kept whole, so lowering decomposes them into a full or full_like
with the dtype given, which Core ML handles. Float fills are unchanged.
With lower_full_graph=True, a model with such a fill now raises instead
of lowering with the wrong dtype.
shoumikhin
changed the title
Contributor
Author
|
@claude /pr-review |
PR Review: #24317
|
Contributor
Author
|
@claude /pr-review |
PR Review: #24317
|
torch.nn.RMSNorm with its default eps=None exports rms_norm without the eps argument. coremltools 9.0 requires that argument, so lowering fails with "node aten_rms_norm_default (rms_norm) got 3 input(s), expected [4]", and it turns a None eps into 1e-5. Torch computes RMSNorm in float32 and uses the float32 machine epsilon, about 1.19e-7, also for float16 and bfloat16 inputs. The Core ML backend now gives coremltools that eps when the model leaves it out, so float32 models with such an RMSNorm lower and match torch. A float16 model still fails in coremltools' rms_norm, with or without eps. The branch for an older coremltools, which has no rms_norm conversion, is unchanged.
torch.export passes the approximate argument of gelu as a keyword. The coremltools 9.0 converter reads it only from the positional inputs, so gelu(approximate="tanh") silently becomes the exact gelu. The two differ by up to 4.7e-4, at float32 and float16 precision alike, which looks like poor float32 accuracy in Core ML. The Core ML backend now registers a gelu conversion that also reads the keyword. A tanh gelu written out with pow was not affected, since coremltools fuses that formula into the right mode.
shoumikhin
changed the title
coremltools converts ops such as clone into nothing, so a Core ML partition can return one of its own inputs, as in "return x.clone(), x * 2", or in "y = x.clone()" right before an op that runs on the CPU. It can also return one value twice, as in "return y, y.clone()". When the partition has other outputs, Core ML hands back such an output wrong: all zeros, or the bytes of another output. The Core ML backend now runs one more coremltools pass, which gives each output that repeats an input or an earlier output an identity op of its own, so every output comes from an op and has its own name. The clone test now returns such outputs. It checks them at run time, and checks the pass on the converted program, which needs no runtime.
When a constant has users in different partitions, or in a partition and on the CPU, to_backend gives each user its own copy of the placeholder. duplicate_constant_node only replaced the constant in the top-level args and kwargs of a user, so an op that reads it from a list, such as cat or the indices of index_put, kept the original. Lowering then failed with "constant data node ... is tagged with ... but has user ... which has tag (None)". For example, a Core ML model whose index buffer is read by a gather in Core ML and by an index_put on the CPU failed this way. The pass now replaces the constant wherever the user reads it.
coremltools 9.0 converts index_put with a whole slice, such as y[:, idx] = v, by writing v into the part of y the indices select, with each indexed dim kept as size 1. That only works with accumulate=False, int indices of shape () or (1,), a static size in every other dim, a v that broadcasts to that part, and a non-bool y below iOS 18. With a 0-d index or several indices, torch lays out that part differently, so v must not vary before the last index, or the values land in the wrong place. With two entries in idx, as in the X-VLA policy, lowering fails with "Input squeeze_0 has rank 1 != other inputs rank 0". Without a whole slice, it fails when several index tensors have different shapes, as in y[rows[:, None], cols] = v, or include a bool mask. The partitioner now refuses those forms, so they run on the portable kernels. The forms coremltools handles stay in Core ML, such as a single-entry index next to a slice for a cache update, a scalar written through several single-entry indices, or index tensors of one shape without a slice. A negative index next to a slice still stays in Core ML, since the check cannot see index values: it fails, or gives wrong values when it comes from an input. The fill test now checks how many calls of an op stay on the CPU, since its index_put model has several forms that must all be refused.
The tanh gelu and RMSNorm tests and the int zeros case of the fill test only failed without their fixes when the Core ML runtime ran the program, so they passed on runners without it. The gelu test now also checks that the converted program uses the tanh mode, the RMSNorm test checks the eps in the converted program, and the fill test checks that the int zeros is not kept whole, so lowering decomposes it. The RMSNorm override now also says why it has to stay once coremltools accepts a missing eps: coremltools would use 1e-5, not the torch eps. The RMSNorm test is skipped in fbcode, whose coremltools has no rms_norm converter to check.
to_edge_transform_and_lower asks the partitioner twice in a row for the ops to keep whole, with the same program. When the program has an op Core ML rejects that has a decomposition, such as an int zeros or a bfloat16 linear, each answer decomposes the whole program in a dry run, which is slow for a large model. The Core ML partitioner now keeps its last answer and returns it again for the same program, as the MLX partitioner does. It holds the program through a weak reference, so it never keeps a program alive.
The partitioner raised for every state buffer whose readers and update
were in different places. But tag_mutated_buffer only moves a buffer
into Core ML when an op in a partition reads it. When only CPU ops read
it, for example an op Core ML rejects or skips, the buffer stays a
program buffer and the delegate returns its new value, which worked
before the check. Such a model, like
y = torch.acosh(self.state)
self.state.copy_(x * 2)
failed to lower with "used in more than one place".
The check now skips a buffer that no partition reads.
coremltools makes full without a dtype, zeros, new_zeros and, from iOS 16, ones_like as float, zeros_like with the dtype of its input, and full_like, new_full, fill and scalar_tensor with the type of their fill value, instead of with the dtype torch gives the result. So an int or bool tensor came out as float: index math on it failed to lower, or lost integers above 2048 at float16 precision. A float tensor could come out as int or bool too: the zeros_like(mask, dtype=float) that nn.MultiheadAttention makes from a bool mask stayed bool, and attention came out wrong. The partitioner refused the int and bool cases, which moved them to the CPU, and with lower_full_graph=True it failed models that worked before, such as x + torch.full((1, 16), 0). The backend now converts these ops with the dtype in the node's metadata, and the partitioner no longer refuses them. A TorchScript node, which has no such metadata, and a fill value known only at run time keep the coremltools converter.
The partitioner refuses some ops that coremltools converts, such as an index_put form coremltools gets wrong, or an op on a bfloat16 tensor. With lower_full_graph=True, lowering then failed with "CoreML does not support the op index_put" and advice to ask coremltools for the op, although coremltools has it. The real reason only went to the info log. Lowering now fails with the reason, for example "aten.index_put.default with these indices is not supported by CoreML. Set lower_full_graph=False in the CoreMLPartitioner to run it outside of CoreML." The info log keeps the same reasons, now with op names instead of their schemas.
The dry run in ops_to_not_decompose drops kept ops that end up on the CPU, and runs again, since a dropped op gets decomposed and can change the program. A dropped op with no decomposition, such as the cast of a uint8 input, changes nothing, so the next round repeated the same decomposition and search only to find nothing more to drop. The dry run now stops there, and builds the decomposition table once. Its docstring also says that it does not see transform passes, or the ops an earlier partitioner takes.
When ops in more than one place use a buffer that Core ML takes over, lowering fails, since each place would use its own copy. The error said only that the buffer is used in more than one place. For a model whose cache update runs on the CPU, such as an index_put Core ML cannot convert, it did not say which op that is. The error now lists each op that uses the buffer and where it runs, and the docs for take_over_mutable_buffer mention the error.
shoumikhin
requested review from
JacobSzwejbka and
larryliu0820
as code owners
November 1, 2026 12:01
shoumikhin
changed the title
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Core ML lowering crashes, or gives wrong values, for several common patterns. This change fixes them in the Core ML partitioner and backend, plus an
exir(ExecuTorch export library) bug that the new partitioner refusals run into.Problems and fixes
Lowering crashes, or the program fails to load:
KeyError: torch.uint8(issue CoreML partitioner errors out with 8-bit integer types #11686). Ops that read such an input now run on the CPU up to the conversion to float, and the ops after it still go to Core ML. So does an op that makes an int8 tensor used outside Core ML, which the Core ML runtime cannot return, and a partition left with only size inputs, or no inputs below iOS 18. Withlower_full_graph=True, lowering fails and names the tensor.take_over_mutable_buffer=False, instead of leaving one side with a stale copy. State that only CPU ops read stays on the CPU, also when Core ML computes its new value.min,max,amin,amax,minimumormaximumover a bool tensor, such as the checktorch.distributions.MultivariateNormalexports, crashes lowering. These ops now run on the portable CPU kernels.matmulkeepsaten.matmulwhole, a bfloat16lineardecomposes into anothermatmulthat also stays whole, on the CPU, and loading fails with "kernel 'aten::matmul.out' not found". The partitioner now decomposes a copy of the program the way lowering will and runs the same search there, so an op that ends up on the CPU is not kept whole (an op skipped by name still is). This dry run only happens with a uint8 or int8 tensor, or when an op Core ML rejects gets decomposed, and once per program.fullwithout a dtype,zerosandnew_zerosas float,ones_likeas float from iOS 16,zeros_likewith the dtype of its input, andfull_like,new_full,fillandscalar_tensorwith the type of their fill value. Sotorch.full((1, 16), 0)(int64 in torch) is made as float, and index math on it fails to lower, as in the SmolVLM vision embeddings, or loses integers above 2048 at float16 precision. The backend now converts these ops, andempty,new_emptyandempty_like, with the dtype torch gives them, so they stay in Core ML. A fill value known only at run time keeps the coremltools conversion.torch.nn.RMSNorm(root mean square normalization) leaves out eps (the small number added before the square root) by default, and coremltools requires it: "got 3 input(s), expected [4]". The backend now passes the float32 machine epsilon, which is what PyTorch uses. A model whose tensors are float16 still fails inside coremltools, with or without eps.y[:, idx] = vonly whenidxhas a single entry. Withidx = torch.tensor([0, 2]), as in the X-VLA policy, lowering fails with "Input squeeze_0 has rank 1 != other inputs rank 0".accumulate=True, a bool mask, or, below iOS 18, a bool tensor next to a slice fail too, and so do index tensors that broadcast. These forms now run on the portable kernels, while forms coremltools handles, such as a single-entry index for a cache update, stay in Core ML. So a cache update at several positions, as in a prefill, now lowers, with the update on the CPU. A negative index next to a slice still stays in Core ML, since the partitioner cannot see index values: it fails, or, for an index computed at run time, can give wrong values.to_backendgives each side a copy, but it missed ops that read it from a list, such ascatorindex_put. Lowering then failed with "constant data node ... is tagged with ... but has user ... which has tag (None)".duplicate_constant_nodeinexirnow replaces the constant there too.Wrong values without any error:
gelu(approximate="tanh")ran as the exact GELU (Gaussian Error Linear Unit), because torch.export passesapproximateas a keyword and coremltools reads it only from the positional arguments. The two differ by up to 4.7e-4, which in pi0 and pi0.5 was the whole float32 error of the policy. The backend now reads the keyword.cloneinto nothing, so a partition can return its own input, as inreturn x.clone(), x * 2, or one value twice. Next to other outputs, such an output came back as zeros or as the bytes of another output. The backend now gives each one a copy of its own.torch.zeros_like(mask, dtype=...), which coremltools kept bool, so attention came out wrong. The fill fix above covers this too.With
lower_full_graph=True, an op the partitioner refuses now fails with the reason, such as "aten.index_put.default with these indices is not supported by CoreML. Set lower_full_graph=False in the CoreMLPartitioner to run it outside of CoreML.", instead of saying that Core ML does not support the op.A uint8 buffer that the model casts still fails with the default settings, and lowers with
take_over_constant_data=False, as the docs now say.How to review
Start with
coreml_partitioner.py: the uint8 and int8 search in_propose_partitions, the state check, and the dry run in_drop_ops_left_outside. The refusals are small predicates near the top of the file, one per problem. The fixes intorch_ops.py,coreml_preprocess.pyandexirdo not depend on the partitioner changes.Tests
With coremltools 9.0, the tests below cover these cases. On macOS 15 or later they also run the program with the Core ML runtime and compare it with eager. With the Core ML code of main, the fill,
index_put, state and compiler tests fail (the nn.MultiheadAttention case only with the runtime). The cases:zeros_like,masked_fill, a channel swap,max(dim=1)orsplit, an int8 frame read by a permute, and a roll with a dynamic height.min/amax/maximum, and a bfloat16 linear next to a float32 matmul.index_putwith a two-entry index, broadcast index tensors,accumulate=Trueor a bool tensor written next to a slice, and 0-d indices whose values do not fit or would land in the wrong place. Withlower_full_graph=True, the error gives the reason.exir, a buffer read by twocatops gets a copy for the second one.test_coreml_partitioner.py,test_torch_ops.pyandtest_coreml_recipes.pypass with the Core ML runtime on macOS, and the newexirtest passes, and fails without its fix. For a transformer block and conv models, with and without skipped ops, the partitions match main.LeRobot policies, measured on an earlier version of this change with wheels built from the 1.6 release branch: without it, ACT, Diffusion, VQ-BeT, Multi-Task DiT and pi0.5 (bfloat16) fail to lower with
KeyError: torch.uint8. With it, ACT, Diffusion, VQ-BeT and Multi-Task DiT export, run without loading torch, and match eager within 4.2e-06 over 16 frames. pi0.5 exports and passes its one-chunk check, but runs very slowly, since its bfloat16 linears run on the portable CPU kernels.Not tested: Core ML state at run time, and a model that builds
torch.distributions.MultivariateNormalitself (the test copies its bool check).