| |
| |
|
|
| 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)) |
|
|
| |
| |
| |
| 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 |
|
|
