Metal backend: loading a model after another was destroyed fails with "Memory address ... is already being tracked by another tensor"

Start in backends/apple/metal/runtime/shims/memory.cpp, then trace aoti_torch_delete_tensor_object and cleanup_memory to understand how tensors and memory_to_n_tensor are cleared. Reproduce with the provided extension::Module loop over multiple Metal .pte files; done means repeated model destruction and loading no longer reports a tracked memory address, with the regression test passing.

🐛 Describe the bug

In one process, load a Metal-delegated model, destroy it, then load another one (or the same one again). The later load fails about half the time:

E executorch:memory.cpp:126] Memory address 0x10a214000 is already being tracked by another tensor
Error: aoti_torch_create_tensor_from_blob_v2(...) API call failed at .../aoti_runtime/model_container...
E executorch:method.cpp:132] Init failed for backend MetalBackend: 0x1

Any app that lets the user switch between Metal models hits this. executor_runner loads a single model per process, so the module tests cannot see it.

Repro. Six small Metal .pte files, each loaded and destroyed in turn, twice, with extension::Module:

for (int round = 0; round < 2; round++) {
  for (int i = 1; i < argc; i++) {
    Module m(argv[i]);
    auto err = m.load_method("forward");   // fails intermittently from the 2nd load on
  }
}

Result: 12 loads, 7 failures (the first load always succeeds). Each model loads fine in a fresh process.

Cause. backends/apple/metal/runtime/shims/memory.cpp. Tensors created from the constants blob are registered with memory_to_n_tensor[adjusted_data] = NOT_OWN. In aoti_torch_delete_tensor_object the NOT_OWN branch erases the tensor from tensors and returns without touching memory_to_n_tensor:

if (ref_count == NOT_OWN) {
  tensors.erase(it);
  return Error::Ok;          // the address stays in memory_to_n_tensor
}

cleanup_memory(), which MetalBackend::destroy calls, clears tensors but not memory_to_n_tensor either. The stale addresses outlive the model, and when the next model's constants are mapped at an address used before, aoti_torch_create_tensor_from_blob_v2 trips its "already being tracked" check. Whether that happens depends on address reuse, hence the intermittency.

I have a fix with a test ready and will open a PR.

Versions

ExecuTorch: main @ 11120c8dff (also reproduced with the 1.5.0 release: PyPI wheel for export, runtime built from the v1.5.0 tag)
PyTorch version: 2.14.0
torchao: built from source with TORCHAO_BUILD_EXPERIMENTAL_MPS=1
OS: macOS 27.0 (arm64), Apple M2 Pro
Clang version: 21.0.0 (clang-2100.1.1.101)
CMake version: 4.4.3
Python version: 3.10.11