[Mutable Buffer] [Core ML Delegate] Let Core ML Handle Mutable Buffer

[YifanShenSZ]

🚀 The feature, motivation and pitch

Starting from iOS 18, Core ML has state, which is the counterpart of mutable buffer. As a result, ExecuTorch can now let Core ML handle buffer mutation

Additional context

The change will probably base on https://github.com/pytorch/executorch/pull/2876/

[lucylq]

@cccclai on Core ML

[cccclai]

@YifanShenSZ any chance we know eta for this feature?

[YifanShenSZ]

I'm actively working on coremltools side, aiming toward our next release (end of July)

1. (Done) conversion of executorch exported mutable-buffer model
2. (Ongoing) make sure the in-place ops have correct pattern

[YifanShenSZ]

Coremltools side change is almost done, I can now export stateful executorch model to CoreML. We will need some change in executorch to_backend implementation to let the delegation path work

[cccclai]

From this error log:

  File "/Volumes/Models/LLM/Framework/CoreMLTools-Dev_ExecuTorch-0.2/envs/llama-py310/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 113, in _
    copied_edge_program = copy.deepcopy(edge_program)

Looks like the program is not copiable. Can you try running copy.deepcopy(exported_program) before running to_backend api?

[YifanShenSZ]

Dug a bit and found the issue to be more involved. Please reproduce by

git clone --branch coreml-state https://github.com/YifanShenSZ/executorch.git

cd executorch
git submodule sync
git submodule update --init
./install_requirements.sh

./backends/apple/coreml/scripts/install_requirements.sh
pip install numpy==1.26.4

cd backends/apple/coreml
python test/test_coreml_partitioner.py

[kimishpatel]

Are you asking for state management completely handed over to delegate? If so, would delegate allow access to this state? We have had requests from users who wanted to manipulate kv cache state and not sure how this will line up with that

[YifanShenSZ]

Are you asking for state management completely handed over to delegate?

Yes

If so, would delegate allow access to this state? We have had requests from users who wanted to manipulate kv cache state and not sure how this will line up with that

@cymbalrush does Core ML runtime allow user to access state?

[cymbalrush]

Core ML allows access to state - https://developer.apple.com/documentation/coreml/mlstate/getmultiarrayforstatenamed:handler:?language=objc

[cccclai]

I was able to repro, looking

Traceback (most recent call last):
  File "/Users/chenlai/coreml_debug/executorch/backends/apple/coreml/test/test_coreml_partitioner.py", line 118, in <module>
    test_runner.test_buffer()
  File "/Users/chenlai/coreml_debug/executorch/backends/apple/coreml/test/test_coreml_partitioner.py", line 102, in test_buffer
    delegated_program_manager = edge_program_manager.to_backend(CoreMLPartitioner())
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/executorch/exir/program/_program.py", line 1204, in to_backend
    new_edge_programs[name] = to_backend(program, partitioner)
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/functools.py", line 878, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 394, in _
    return ExportedProgram(
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/torch/export/exported_program.py", line 682, in __init__
    self._validate()
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/torch/export/exported_program.py", line 1101, in _validate
    v().check(self)
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/torch/_export/verifier.py", line 157, in check
    _verify_exported_program_signature(ep)
  File "/opt/homebrew/anaconda3/envs/coreml_debug/lib/python3.10/site-packages/torch/_export/verifier.py", line 408, in _verify_exported_program_signature
    raise SpecViolationError(
torch._export.verifier.SpecViolationError: Buffer output getitem does not point to a buffer that exists.
Dict of buffers that are mutated, in order: {'getitem': 'state_1'}
Buffer nodes available: []

[cccclai]

After checking, some changes need to be added to the delegate infra to support consuming the mutable buffer. It is our first case to have the backend to consume the mutable buffer. @angelayi will help add the feature.

[kimishpatel]

@cccclai can you describe what is the nature of the change?

[cccclai]

@cccclai can you describe what is the nature of the change?

We'd need to add some changes in lowered_backend_module.py and backend_api.py to support backend consuming the in-place ops. Since we didn't have a backend that can consume in place ops before, it's the first time we test this branch.

[cccclai]

There are some progress from @angelayi on this. https://github.com/pytorch/executorch/pull/4566 is first pr and there will be one more pr to resolve it.

[YifanShenSZ]

Verified. Thanks team!