pytorch / pytorch/executorch

[Llama] [Dynamic Shape] [Core ML Delegate] Do Not Delegate Symbol Manipulation

Open
#4,658 19 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

module: exir triaged
Dominant language
Python
Stars
5k
Forks
1.2k
Avg merge
2d 10h
Merged PRs (30d)
581

Description

When export dynamic-shape llama2, in the attached piece of the partitioned model, there are symbol manipulations such as

add: "Sym(s0 + u156)" = _local_scalar_dense + sym_size
...
le_1: "Sym(s0 + u156 <= 128)" = add <= 128

Which came from dynamic shape check assertions in the original model

add: "Sym(s0 + u156)" = _local_scalar_dense + sym_size
le_1: "Sym(s0 + u156 <= 128)" = add <= 128
_assert_scalar_2 = torch.ops.aten._assert_scalar.default(le_1, "Runtime assertion failed for expression s0 + u0 <= 128 on node 'le_27'");  le_1 = None

This leads to 2 problems:

  1. Core ML does not support symbol itself as model io
  2. torch.ops.aten._assert_scalar.default is scattered all over the program, which is not supported in Core ML so may lead to undesirable graph breaks

Imho:

  1. These dynamic shape checks are very small, so doing them in executorch runtime may be more efficient than paying delegate overhead
  2. All assertions may be performed at the beginning of the program in executorch runtime, so the rest heavy computation can be delegated to Core ML as a whole graph

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the attached partitioned-model piece and original model to trace how dynamic-shape checks become Core ML delegate inputs and scattered _assert_scalar operations. Then inspect the Core ML delegate's partitioning and assertion handling; done means symbol manipulation and these assertions stay in ExecuTorch runtime while the remaining heavy computation can be delegated as a whole graph.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, mobile-dev
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.