Multifunction consumes linearly increasing memory in the number of functions, even with all the weights shared
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🌱 Describe your Feature Request
Hi, when creating a multifunction with many functions that reuses the same weights for all cases, the memory increases on the number of functions. It would be useful to reduce the memory usage for speed and accessibility to use.
## Example code to reproduce the problem
This code creates a simple model with 1 billion parameters, and creates many functions with different input sizes for the same model
```python
import os
import shutil
import numpy as np
import coremltools as ct
from coremltools.converters.mil import Builder as mb
import coremltools.converters.mil as mil
w1 = np.random.normal(loc=0.01, size=(16_384, 16_384, 1)).astype(np.float16)
w2 = np.random.normal(loc=0.01, size=(16_384, 16_384, 1)).astype(np.float16)
w3 = np.random.normal(loc=0.01, size=(16_384, 16_384, 1)).astype(np.float16)
w4 = np.random.normal(loc=0.01, size=(16_384, 16_384, 1)).astype(np.float16)
def make_model(length):
@mb.program(
input_specs=[
mb.TensorSpec(
(1, 16_384, length),
dtype=mil.input_types.types.fp16,
),
],
opset_version=mil.builder.AvailableTarget.iOS18,
)
def program(x):
x = mb.conv(x=x, weight=w1)
x = mb.conv(x=x, weight=w2)
x = mb.conv(x=x, weight=w3)
return mb.conv(x=x, weight=w4)
cml_converted = ct.convert(
program,
compute_units=ct.ComputeUnit.CPU_AND_NE,
compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.iOS18,
skip_model_load=True,
)
cml_converted.save(f"./model_{length}")
def merge_mfs(mf_filename, new_model, length):
if os.path.isdir(mf_filename):
desc = ct.utils.MultiFunctionDescriptor(mf_filename)
else:
desc = ct.utils.MultiFunctionDescriptor(None)
print(f"Adding length {length}, already created lengths: {desc._functions()}")
desc.add_function(
new_model,
src_function_name="main",
target_function_name=f"length_{length}",
)
desc.default_function_name = "length_1"
ct.utils.save_multifunction(desc, mf_filename)
shutil.rmtree(new_model)
if __name__ == "__main__":
mf_name = "mf.mlpackage"
for i in [1, 2, 4, 6, 8]:
make_model(i)
merge_mfs(mf_name, f"model_{i}.mlpackage", i)
```
Included a [video](https://drive.google.com/file/d/1QaqLFye4hCgBoJnru1Q9UOZD_hbw0tLl/view?usp=sharing) showing the RAM usage.
- Minute 1:15 peak RAM usage of 4.74GB for 1 function
- Min 2:06 peak RAM usage of 7.74GB for 2 functions
- min 4:07 peak RAM usage 9.74GB for 3 functions
- min 6:31 peak RAM usage of 11.74GB for 4 functions
Contributor guide
Research direction
Start with the provided reproduction and inspect ct.utils.MultiFunctionDescriptor.add_function and ct.utils.save_multifunction, which are the entry points named in the report. Compare memory usage as functions are added with shared weights; done means the resulting multifunction no longer consumes memory linearly with the number of functions while preserving the reported behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100