File size: 3,068 Bytes
656477e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
import gc
import ttnn
from models.common.modules.lazy_weight import LazyWeight
def cleanup_ttnn_value(value):
if value is None:
return
if isinstance(value, ttnn.Tensor):
ttnn.deallocate(value)
return
if isinstance(value, dict):
for nested_value in value.values():
cleanup_ttnn_value(nested_value)
return
if isinstance(value, (list, tuple, set)):
for nested_value in value:
cleanup_ttnn_value(nested_value)
def cleanup_object_graph(obj, seen=None):
if obj is None:
return
if seen is None:
seen = set()
obj_id = id(obj)
if obj_id in seen:
return
seen.add(obj_id)
if isinstance(obj, ttnn.Tensor):
cleanup_ttnn_value(obj)
return
if isinstance(obj, LazyWeight):
if obj._value is not None:
cleanup_ttnn_value(obj._value)
obj._value = None
return
if isinstance(obj, dict):
for value in obj.values():
cleanup_object_graph(value, seen)
return
if isinstance(obj, (list, tuple, set)):
for value in obj:
cleanup_object_graph(value, seen)
return
state = getattr(obj, "__dict__", None)
if state is None:
return
for name, value in list(state.items()):
cleanup_object_graph(value, seen)
if isinstance(value, ttnn.Tensor):
setattr(obj, name, None)
if hasattr(obj, "_device_weights_loaded"):
obj._device_weights_loaded = False
def cleanup_model_case(model, mesh_device):
ttnn.synchronize_device(mesh_device)
if model is not None:
cleanup_object_graph(model)
ttnn.synchronize_device(mesh_device)
gc.collect()
def cleanup_dp_model_case(group, lanes, models, parent_mesh, submeshes):
"""Release one carved DP case and restore command-queue ownership to its parent."""
failures = []
def run(action):
try:
action()
except Exception as error:
failures.append(error)
if group is not None:
run(group.cleanup)
else:
for lane in lanes:
run(lane.cleanup)
for model, submesh in models:
run(lambda model=model, submesh=submesh: cleanup_model_case(model, submesh))
# A carved child owns a distinct MeshCommandQueue over its parent's physical
# devices. Drain every live child through the parent before closing the child
# handles and returning command-queue ownership to the fixture-owned parent.
run(parent_mesh.quiesce_devices)
for submesh in submeshes:
run(lambda submesh=submesh: ttnn.close_mesh_device(submesh))
if failures:
primary = failures[0]
add_note = getattr(primary, "add_note", None)
if add_note is not None:
for failure in failures[1:]:
add_note(f"additional DP cleanup failure: {type(failure).__name__}: {failure}")
raise primary
|