From 88b6df2401e187ca6bd81bf0a83bc7c5d9a5cd54 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 02:55:52 +0800 Subject: [PATCH 1/4] Fix subclass tracking of module attribute assignments --- tests/test_subclass_assignment.py | 375 ++++++++++++++++++ .../instrumentor/proxy_wrapper/subclass.py | 171 +++++--- 2 files changed, 493 insertions(+), 53 deletions(-) create mode 100644 tests/test_subclass_assignment.py diff --git a/tests/test_subclass_assignment.py b/tests/test_subclass_assignment.py new file mode 100644 index 00000000..34e77520 --- /dev/null +++ b/tests/test_subclass_assignment.py @@ -0,0 +1,375 @@ +"""Subclass tracking must preserve PyTorch module assignment contracts.""" + +import gc +import unittest +import weakref +from unittest.mock import patch + +import torch +from torch import nn +from torch.nn.modules.module import ( + register_module_module_registration_hook, + register_module_parameter_registration_hook, +) + +import traincheck.config.config as config +from traincheck.instrumentor.proxy_wrapper import subclass +from traincheck.instrumentor.proxy_wrapper.proxy_registry import VarRegistry +from traincheck.instrumentor.proxy_wrapper.subclass import ( + ProxyParameter, + proxy_parameter, +) + + +class MaskedTransformer(nn.Module): + """A small training workload that caches a causal mask during forward.""" + + def __init__(self) -> None: + super().__init__() + self.encoder = nn.TransformerEncoderLayer( + d_model=4, nhead=2, dim_feedforward=8, dropout=0.0, batch_first=True + ) + self.src_mask: torch.Tensor | None = None + + def forward(self, source: torch.Tensor) -> torch.Tensor: + self.src_mask = nn.Transformer.generate_square_subsequent_mask(source.size(1)) + return self.encoder(source, src_mask=self.src_mask) + + +class SubclassAssignmentTests(unittest.TestCase): + def setUp(self) -> None: + # Keep tests independent of the process-wide registry and trace files. + self.registry = VarRegistry() + registry_patch = patch.object( + subclass, "get_global_registry", return_value=self.registry + ) + registry_patch.start() + self.addCleanup(registry_patch.stop) + dumping_patch = patch.object(config, "DISABLE_WRAPPER", True) + dumping_patch.start() + self.addCleanup(dumping_patch.stop) + + def track(self, module: nn.Module) -> None: + proxy_parameter(module, parent_name="model", should_dump_trace=False) + + def test_transformer_mask_forward_backward_and_update_match(self) -> None: + torch.manual_seed(71) + plain = MaskedTransformer() + tracked = MaskedTransformer() + tracked.load_state_dict(plain.state_dict()) + rng_before = torch.get_rng_state().clone() + self.track(tracked) + torch.testing.assert_close(torch.get_rng_state(), rng_before, rtol=0, atol=0) + source = torch.randn(2, 3, 4) + optimizers = [ + torch.optim.SGD(model.parameters(), lr=0.1) for model in (plain, tracked) + ] + for model, optimizer in zip((plain, tracked), optimizers): + model(source).square().mean().backward() + optimizer.step() + self.assertIsInstance(tracked.src_mask, torch.Tensor) + for name, expected in plain.state_dict().items(): + torch.testing.assert_close( + tracked.state_dict()[name], expected, rtol=0, atol=0 + ) + for actual, expected in zip(tracked.parameters(), plain.parameters()): + torch.testing.assert_close(actual.grad, expected.grad, rtol=0, atol=0) + + def test_plain_tensor_and_unrelated_attributes_are_unchanged(self) -> None: + module = nn.Module() + self.track(module) + tensor = torch.ones(2) + module.mask = tensor + module.counter = 3 + module.label = None + self.assertIs(module.mask, tensor) + self.assertEqual(module.counter, 3) + self.assertIsNone(module.label) + self.assertEqual(list(module.named_parameters()), []) + self.assertEqual(list(module.named_buffers()), []) + + def test_buffer_replacement_preserves_identity_and_persistence(self) -> None: + module = nn.Module() + module.register_buffer("persistent", torch.zeros(2)) + module.register_buffer("temporary", torch.zeros(2), persistent=False) + self.track(module) + replacement = torch.ones(2) + module.persistent = replacement + module.temporary = replacement + self.assertIs(module.persistent, replacement) + self.assertIs(module.temporary, replacement) + self.assertEqual(list(module.state_dict()), ["persistent"]) + module.persistent = None + self.assertIsNone(module.persistent) + self.assertEqual(list(module.state_dict()), []) + + def test_new_and_replaced_parameters_are_tracked(self) -> None: + module = nn.Module() + self.track(module) + module.weight = nn.Parameter(torch.ones(2)) + first = module.weight + self.assertIsInstance(first, ProxyParameter) + replacement = nn.Parameter(torch.zeros(2), requires_grad=False) + gradient = torch.full((2,), 7.0) + replacement.grad = gradient + marker = object() + replacement.user_marker = marker + module.weight = replacement + self.assertIsInstance(module.weight, ProxyParameter) + self.assertIsNot(module.weight, first) + self.assertFalse(module.weight.requires_grad) + self.assertIs(module.weight.grad, gradient) + self.assertIs(module.weight.user_marker, marker) + self.assertIs(self.registry.registry["model.weight"].var, module.weight) + module.weight = None + self.assertIsNone(module.weight) + self.assertEqual(list(module.parameters()), []) + + def test_parameter_backward_hook_survives_assignment(self) -> None: + calls: list[torch.Tensor] = [] + + def double_gradient(gradient: torch.Tensor) -> torch.Tensor: + calls.append(gradient.clone()) + return gradient * 2 + + parameter = nn.Parameter(torch.ones(2)) + handle = parameter.register_hook(double_gradient) + self.addCleanup(handle.remove) + module = nn.Module() + self.track(module) + module.weight = parameter + module.weight.sum().backward() + self.assertEqual(len(calls), 1) + torch.testing.assert_close(module.weight.grad, torch.full((2,), 2.0)) + handle.remove() + module.weight.grad = None + module.weight.sum().backward() + self.assertEqual(len(calls), 1) + torch.testing.assert_close(module.weight.grad, torch.ones(2)) + + def test_post_accumulation_hook_survives_assignment(self) -> None: + if not hasattr(torch.Tensor, "register_post_accumulate_grad_hook"): + self.skipTest("This PyTorch version has no post-accumulation hook API") + calls: list[torch.Tensor] = [] + + def observe(parameter: torch.Tensor) -> None: + calls.append(parameter) + + parameter = nn.Parameter(torch.ones(2)) + handle = parameter.register_post_accumulate_grad_hook(observe) + self.addCleanup(handle.remove) + module = nn.Module() + self.track(module) + module.weight = parameter + module.weight.sum().backward() + self.assertEqual(len(calls), 1) + self.assertIs(calls[0], module.weight) + handle.remove() + module.weight.grad = None + module.weight.sum().backward() + self.assertEqual(len(calls), 1) + + def test_assigned_child_and_replacement_parameters_are_tracked(self) -> None: + module = nn.Module() + self.track(module) + for child in (nn.Sequential(nn.Linear(2, 2)), nn.Sequential(nn.Linear(2, 1))): + module.child = child + self.assertIs(module.child, child) + self.assertTrue( + all(isinstance(p, ProxyParameter) for p in child.parameters()) + ) + child.mask = torch.ones(2) + child[0].weight = nn.Parameter(torch.zeros_like(child[0].weight)) + self.assertIsInstance(child[0].weight, ProxyParameter) + module.child = None + self.assertIsNone(module.child) + self.assertEqual(list(module.children()), []) + + def test_parameter_aliases_survive_initial_and_later_tracking(self) -> None: + parameter = nn.Parameter(torch.ones(2)) + module = nn.Module() + module.weight = parameter + module.alias = parameter + child = nn.Module() + child.weight = parameter + module.child = child + self.track(module) + self.assertIs(module.weight, module.alias) + self.assertIs(module.weight, child.weight) + module.later_alias = parameter + self.assertIs(module.weight, module.later_alias) + metadata_name = module.weight.var_name + module.proxy_alias = module.weight + self.track(module) + self.assertIs(module.weight, module.proxy_alias) + self.assertEqual(module.weight.var_name, metadata_name) + self.assertEqual(len(list(module.parameters())), 1) + + def test_alias_cache_does_not_keep_replaced_source_alive(self) -> None: + module = nn.Module() + self.track(module) + parameter = nn.Parameter(torch.ones(2)) + reference = weakref.ref(parameter) + module.weight = parameter + del parameter + gc.collect() + self.assertIsNone(reference()) + module.weight = nn.Parameter(torch.zeros(2)) + torch.testing.assert_close(module.weight, torch.zeros(2)) + + def test_replaced_proxy_is_not_retained_by_alias_cache(self) -> None: + module = nn.Module() + self.track(module) + source = nn.Parameter(torch.ones(2)) + module.weight = source + proxy_reference = weakref.ref(module.weight) + module.weight = nn.Parameter(torch.zeros(2)) + gc.collect() + self.assertIsNone(proxy_reference()) + module.weight = source + self.assertIsInstance(module.weight, ProxyParameter) + torch.testing.assert_close(module.weight, torch.ones(2)) + + def test_assignment_trace_contains_each_new_parameter_once(self) -> None: + module = nn.Module() + source = nn.Parameter(torch.ones(2)) + module.weight = source + with ( + patch.object(config, "DISABLE_WRAPPER", False), + patch.object(subclass, "dump_trace_VAR") as dump, + ): + proxy_parameter(module, parent_name="model") + self.assertEqual(dump.call_count, 1) + self.assertEqual(dump.call_args.args[0]["var_name"], "model.weight") + module.mask = torch.zeros(2) + module.alias = source + self.assertEqual(dump.call_count, 1) + module.weight = nn.Parameter(torch.zeros(2)) + self.assertEqual(dump.call_count, 2) + self.assertEqual(dump.call_args.args[0]["var_name"], "model.weight") + with self.assertRaises(TypeError): + module.weight = torch.ones(2) + self.assertEqual(dump.call_count, 2) + + def test_rejected_assignment_preserves_native_error_and_state(self) -> None: + for name, original, invalid in ( + ("weight", nn.Parameter(torch.ones(2)), torch.zeros(2)), + ("child", nn.Linear(2, 2), torch.zeros(2)), + ("weight", nn.Parameter(torch.ones(2)), nn.Linear(2, 2)), + ): + with self.subTest(name=name, invalid_type=type(invalid)): + plain = nn.Module() + tracked = nn.Module() + setattr(plain, name, original) + setattr(tracked, name, original) + self.track(tracked) + before = getattr(tracked, name) + with self.assertRaises(TypeError) as expected: + setattr(plain, name, invalid) + with self.assertRaises(TypeError) as actual: + setattr(tracked, name, invalid) + self.assertEqual(str(actual.exception), str(expected.exception)) + self.assertIs(getattr(tracked, name), before) + if isinstance(invalid, nn.Module): + self.assertNotIn(subclass.SUBCLASS_HOOK_KEY, invalid.__dict__) + self.assertFalse( + any(isinstance(p, ProxyParameter) for p in invalid.parameters()) + ) + + def test_rejected_buffer_assignment_keeps_native_error_and_state(self) -> None: + module = nn.Module() + buffer = torch.ones(2) + module.register_buffer("mask", buffer) + self.track(module) + with self.assertRaisesRegex(TypeError, "as buffer 'mask'"): + module.mask = 42 + self.assertIs(module.mask, buffer) + + def test_parameter_registration_hook_replacement_is_tracked_once(self) -> None: + module = nn.Module() + self.track(module) + incoming = nn.Parameter(torch.zeros(2)) + replacement = nn.Parameter(torch.ones(2), requires_grad=False) + calls: list[nn.Parameter | None] = [] + + def replace( + owner: nn.Module, name: str, value: nn.Parameter | None + ) -> nn.Parameter: + self.assertIs(owner, module) + self.assertEqual(name, "weight") + calls.append(value) + return replacement + + handle = register_module_parameter_registration_hook(replace) + try: + module.weight = incoming + finally: + handle.remove() + self.assertEqual(len(calls), 1) + self.assertIs(calls[0], incoming) + self.assertIsInstance(module.weight, ProxyParameter) + self.assertFalse(module.weight.requires_grad) + torch.testing.assert_close(module.weight, replacement) + + def test_module_registration_hook_replacement_is_tracked_once(self) -> None: + module = nn.Module() + self.track(module) + incoming = nn.Linear(2, 2) + replacement = nn.Linear(2, 1) + calls: list[nn.Module | None] = [] + + def replace(owner: nn.Module, name: str, value: nn.Module | None) -> nn.Module: + self.assertIs(owner, module) + self.assertEqual(name, "child") + calls.append(value) + return replacement + + handle = register_module_module_registration_hook(replace) + try: + module.child = incoming + finally: + handle.remove() + self.assertEqual(len(calls), 1) + self.assertIs(calls[0], incoming) + self.assertIs(module.child, replacement) + self.assertTrue( + all(isinstance(p, ProxyParameter) for p in replacement.parameters()) + ) + self.assertFalse( + any(isinstance(p, ProxyParameter) for p in incoming.parameters()) + ) + + def test_registration_hook_failure_does_not_instrument_rejected_child(self) -> None: + module = nn.Module() + self.track(module) + child = nn.Linear(2, 2) + error = RuntimeError("registration failed") + + def reject(owner: nn.Module, name: str, value: nn.Module | None) -> None: + raise error + + handle = register_module_module_registration_hook(reject) + try: + with self.assertRaises(RuntimeError) as caught: + module.child = child + finally: + handle.remove() + self.assertIs(caught.exception, error) + self.assertNotIn("child", module._modules) + self.assertNotIn(subclass.SUBCLASS_HOOK_KEY, child.__dict__) + self.assertFalse(any(isinstance(p, ProxyParameter) for p in child.parameters())) + + def test_malformed_present_tracking_hook_is_rejected_before_assignment( + self, + ) -> None: + for malformed in (None, 3, "hook"): + with self.subTest(hook=malformed): + module = nn.Module() + module.__dict__[subclass.SUBCLASS_HOOK_KEY] = malformed + with self.assertRaisesRegex(TypeError, "tracking hook"): + module.label = "new" + self.assertNotIn("label", module.__dict__) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index 1e059fb7..ea0c741a 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -2,6 +2,8 @@ import logging import os import threading +import weakref +from typing import Callable, cast import torch from torch import nn @@ -215,76 +217,139 @@ def dump_trace(self, phase, dump_loc): def proxy_parameter( module: nn.Module, - logdir="proxy_log.log", - log_level=logging.INFO, - # TODO - # recurse=False, - parent_name="", - should_dump_trace=True, - from_call=False, - from_iter=False, - # TODO - # from_copy=False, -): + logdir: str = "proxy_log.log", + log_level: int = logging.INFO, + parent_name: str = "", + should_dump_trace: bool = True, + from_call: bool = False, + from_iter: bool = False, +) -> None: + """Track a module's parameters and subsequent assignments to its module tree. + + Aliases within this tracking session share one proxy and its first registered + name. Weak identity entries avoid retaining replaced source parameters or + comparing Tensor values. Separate calls do not share this cache. Callers must + serialize model tracking and assignment; concurrent model mutation is unsupported. + """ + if not isinstance(module, nn.Module): + raise TypeError("proxy_parameter expects a torch.nn.Module") if in_dynamo(): return - for name, t in list(module.named_parameters(recurse=False)): - var_type = typename(t, is_runtime=True) - logger.debug( - f"[ProxyParameter] Proxying parameter: {parent_name}.{name} of type {var_type}" - ) - module._parameters[name] = ProxyParameter( - t, - logdir, - log_level, - parent_name + "." + name, - var_type, - should_dump_trace, - from_call, - from_iter, - ) - for name, child in module.named_children(): - proxy_parameter( - child, + + aliases: dict[ + int, + tuple[weakref.ReferenceType[nn.Parameter], weakref.ReferenceType[nn.Parameter]], + ] = {} + + def track_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: + if not isinstance(parameter, nn.Parameter): + raise TypeError( + f"Registered parameter {name!r} is not a torch.nn.Parameter" + ) + # Calling ProxyParameter again reruns __init__, overwriting its canonical + # name and emitting another initialization record even for the same object. + if isinstance(parameter, ProxyParameter): + return parameter + identity = id(parameter) + cached = aliases.get(identity) # Absence means this source has no live proxy. + if cached is not None and cached[0]() is parameter: + existing = cached[1]() + if existing is not None: + return existing + proxied = ProxyParameter( + parameter, logdir, log_level, - parent_name + "." + name, + name, + typename(parameter, is_runtime=True), should_dump_trace, from_call, from_iter, ) - # we need to instrument the __setattr__ of the module to capture parameter updates - def subclass_setattr_hook(self, name, value): - logger.debug( - f"[ProxyParameter] Module __setattr__ called: {parent_name}.{name} = {type(value)}" - ) - if isinstance(value, torch.Tensor) or isinstance(value, torch.nn.Module): - proxy_parameter( - value, - logdir, - log_level, - parent_name + "." + name, - should_dump_trace, - from_call, - from_iter, + # Tensor autograd hook dictionaries are slots, not entries in __dict__. + # Share them rather than copying so existing removable handles still work. + torch.Tensor.__setattr__(proxied, "_backward_hooks", parameter._backward_hooks) + # Post-accumulation hooks are absent on older supported PyTorch versions. + if hasattr(parameter, "_post_accumulate_grad_hooks"): + torch.Tensor.__setattr__( + proxied, + "_post_accumulate_grad_hooks", + parameter._post_accumulate_grad_hooks, ) - module.__dict__[SUBCLASS_HOOK_KEY] = subclass_setattr_hook + def forget(reference: weakref.ReferenceType[nn.Parameter]) -> None: + current = aliases.get(identity) + if current is not None and ( + current[0] is reference or current[1] is reference + ): + del aliases[identity] + + aliases[identity] = ( + weakref.ref(parameter, forget), + weakref.ref(proxied, forget), + ) + return proxied + + def track_module(current: nn.Module, path: str, visited: set[int]) -> None: + if not isinstance(current, nn.Module): + raise TypeError(f"Registered child {path!r} is not a torch.nn.Module") + if id(current) in visited: + return + visited.add(id(current)) + # Include aliases rather than named_parameters' default deduplication. + for name, parameter in list(current._parameters.items()): + if parameter is not None: + current._parameters[name] = track_parameter( + parameter, path + "." + name + ) + for name, child in current._modules.items(): + if child is not None: + track_module(child, path + "." + name, visited) + + def subclass_setattr_hook(owner: nn.Module, name: str) -> None: + if in_dynamo(): + return + # The native setter has already validated and registered the value, + # including any replacement made by PyTorch registration hooks. Buffer + # and ordinary Tensor assignments must never become parameters. + if name in owner._parameters: + parameter = owner._parameters[name] + if parameter is not None: + owner._parameters[name] = track_parameter( + parameter, path + "." + name + ) + elif name in owner._modules: + child = owner._modules[name] + if child is not None: + track_module(child, path + "." + name, set()) + + current.__dict__[SUBCLASS_HOOK_KEY] = subclass_setattr_hook + + track_module(module, parent_name, set()) # instrument torch.nn.Module's setattr -orig_setattr = torch.nn.Module.__setattr__ +# PyTorch 2.2 annotates only Tensor | Module, but its native setter also accepts +# arbitrary Python attributes (for example training=False and cached masks=None). +# This cast states that existing runtime contract; it does not coerce the value. +orig_setattr = cast( + Callable[[nn.Module, str, object], None], torch.nn.Module.__setattr__ +) @functools.wraps(orig_setattr) -def wrapped_setattr(self, name, value): - #hook = getattr(self, SUBCLASS_HOOK_KEY, None) - hook = self.__dict__.get(SUBCLASS_HOOK_KEY) - if hook is not None: - # If hook returns True, skip the original setattr; otherwise continue. - hook(self, name, value) - return orig_setattr(self, name, value) +def wrapped_setattr(self: nn.Module, name: str, value: object) -> None: + # Absence identifies an untracked module. A present malformed hook is a + # broken tracking contract, not an instruction to silently disable tracking. + if SUBCLASS_HOOK_KEY not in self.__dict__: + orig_setattr(self, name, value) + return + hook: object = self.__dict__[SUBCLASS_HOOK_KEY] + if not callable(hook): + raise TypeError("TrainCheck module tracking hook must be callable") + orig_setattr(self, name, value) + hook(self, name) torch.nn.Module.__setattr__ = wrapped_setattr From 3557f81699ac7e224e7ec87e008ad7425312e710 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 03:01:38 +0800 Subject: [PATCH 2/4] Check subclass assignment regressions against PyTorch types --- tests/test_subclass_assignment.py | 53 +++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/tests/test_subclass_assignment.py b/tests/test_subclass_assignment.py index 34e77520..7e99a7bf 100644 --- a/tests/test_subclass_assignment.py +++ b/tests/test_subclass_assignment.py @@ -3,6 +3,7 @@ import gc import unittest import weakref +from typing import Callable, TypeAlias, cast from unittest.mock import patch import torch @@ -11,6 +12,7 @@ register_module_module_registration_hook, register_module_parameter_registration_hook, ) +from torch.utils.hooks import RemovableHandle import traincheck.config.config as config from traincheck.instrumentor.proxy_wrapper import subclass @@ -20,6 +22,25 @@ proxy_parameter, ) +ParameterRegistrationHook: TypeAlias = Callable[ + [nn.Module, str, nn.Parameter | None], nn.Parameter | None +] +ModuleRegistrationHook: TypeAlias = Callable[ + [nn.Module, str, nn.Module | None], nn.Module | None +] +# PyTorch 2.2 declares these callbacks as returning None, but both API docstrings +# explicitly allow replacement values and the native setters use those values. +# These narrow boundary casts preserve that documented callback/handle contract; +# the tests below verify the replacement is registered and tracked exactly once. +register_parameter_hook = cast( + Callable[[ParameterRegistrationHook], RemovableHandle], + register_module_parameter_registration_hook, +) +register_module_hook = cast( + Callable[[ModuleRegistrationHook], RemovableHandle], + register_module_module_registration_hook, +) + class MaskedTransformer(nn.Module): """A small training workload that caches a causal mask during forward.""" @@ -80,8 +101,11 @@ def test_plain_tensor_and_unrelated_attributes_are_unchanged(self) -> None: self.track(module) tensor = torch.ones(2) module.mask = tensor - module.counter = 3 - module.label = None + # Module.__setattr__ in PyTorch 2.2 omits ordinary attributes and None + # from its annotation. setattr exercises the same runtime setter without + # misrepresenting those values as Tensor | Module. + setattr(module, "counter", 3) + setattr(module, "label", None) self.assertIs(module.mask, tensor) self.assertEqual(module.counter, 3) self.assertIsNone(module.label) @@ -99,7 +123,7 @@ def test_buffer_replacement_preserves_identity_and_persistence(self) -> None: self.assertIs(module.persistent, replacement) self.assertIs(module.temporary, replacement) self.assertEqual(list(module.state_dict()), ["persistent"]) - module.persistent = None + setattr(module, "persistent", None) self.assertIsNone(module.persistent) self.assertEqual(list(module.state_dict()), []) @@ -113,15 +137,15 @@ def test_new_and_replaced_parameters_are_tracked(self) -> None: gradient = torch.full((2,), 7.0) replacement.grad = gradient marker = object() - replacement.user_marker = marker + setattr(replacement, "user_marker", marker) module.weight = replacement self.assertIsInstance(module.weight, ProxyParameter) self.assertIsNot(module.weight, first) self.assertFalse(module.weight.requires_grad) self.assertIs(module.weight.grad, gradient) - self.assertIs(module.weight.user_marker, marker) + self.assertIs(module.weight.__dict__["user_marker"], marker) self.assertIs(self.registry.registry["model.weight"].var, module.weight) - module.weight = None + setattr(module, "weight", None) self.assertIsNone(module.weight) self.assertEqual(list(module.parameters()), []) @@ -181,7 +205,7 @@ def test_assigned_child_and_replacement_parameters_are_tracked(self) -> None: child.mask = torch.ones(2) child[0].weight = nn.Parameter(torch.zeros_like(child[0].weight)) self.assertIsInstance(child[0].weight, ProxyParameter) - module.child = None + setattr(module, "child", None) self.assertIsNone(module.child) self.assertEqual(list(module.children()), []) @@ -198,11 +222,12 @@ def test_parameter_aliases_survive_initial_and_later_tracking(self) -> None: self.assertIs(module.weight, child.weight) module.later_alias = parameter self.assertIs(module.weight, module.later_alias) - metadata_name = module.weight.var_name + metadata_name: object = module.weight.__dict__["var_name"] + self.assertIsInstance(metadata_name, str) module.proxy_alias = module.weight self.track(module) self.assertIs(module.weight, module.proxy_alias) - self.assertEqual(module.weight.var_name, metadata_name) + self.assertEqual(module.weight.__dict__["var_name"], metadata_name) self.assertEqual(len(list(module.parameters())), 1) def test_alias_cache_does_not_keep_replaced_source_alive(self) -> None: @@ -282,7 +307,7 @@ def test_rejected_buffer_assignment_keeps_native_error_and_state(self) -> None: module.register_buffer("mask", buffer) self.track(module) with self.assertRaisesRegex(TypeError, "as buffer 'mask'"): - module.mask = 42 + setattr(module, "mask", 42) self.assertIs(module.mask, buffer) def test_parameter_registration_hook_replacement_is_tracked_once(self) -> None: @@ -300,7 +325,7 @@ def replace( calls.append(value) return replacement - handle = register_module_parameter_registration_hook(replace) + handle = register_parameter_hook(replace) try: module.weight = incoming finally: @@ -324,7 +349,7 @@ def replace(owner: nn.Module, name: str, value: nn.Module | None) -> nn.Module: calls.append(value) return replacement - handle = register_module_module_registration_hook(replace) + handle = register_module_hook(replace) try: module.child = incoming finally: @@ -348,7 +373,7 @@ def test_registration_hook_failure_does_not_instrument_rejected_child(self) -> N def reject(owner: nn.Module, name: str, value: nn.Module | None) -> None: raise error - handle = register_module_module_registration_hook(reject) + handle = register_module_hook(reject) try: with self.assertRaises(RuntimeError) as caught: module.child = child @@ -367,7 +392,7 @@ def test_malformed_present_tracking_hook_is_rejected_before_assignment( module = nn.Module() module.__dict__[subclass.SUBCLASS_HOOK_KEY] = malformed with self.assertRaisesRegex(TypeError, "tracking hook"): - module.label = "new" + setattr(module, "label", "new") self.assertNotIn("label", module.__dict__) From 3cfee7ec5f4c35eccf2b164598a128205ccbddcb Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 03:23:32 +0800 Subject: [PATCH 3/4] Preserve parameter identity and live alias observation ownership --- tests/test_subclass_assignment.py | 289 ++++++++++++++++- .../instrumentor/proxy_wrapper/subclass.py | 304 +++++++++++++----- 2 files changed, 514 insertions(+), 79 deletions(-) diff --git a/tests/test_subclass_assignment.py b/tests/test_subclass_assignment.py index 7e99a7bf..6fcafb8d 100644 --- a/tests/test_subclass_assignment.py +++ b/tests/test_subclass_assignment.py @@ -1,6 +1,7 @@ """Subclass tracking must preserve PyTorch module assignment contracts.""" import gc +import io import unittest import weakref from typing import Callable, TypeAlias, cast @@ -57,6 +58,10 @@ def forward(self, source: torch.Tensor) -> torch.Tensor: return self.encoder(source, src_mask=self.src_mask) +class CustomParameter(nn.Parameter): + """A distinct framework/user subclass that uses the legacy wrapping path.""" + + class SubclassAssignmentTests(unittest.TestCase): def setUp(self) -> None: # Keep tests independent of the process-wide registry and trace files. @@ -230,30 +235,300 @@ def test_parameter_aliases_survive_initial_and_later_tracking(self) -> None: self.assertEqual(module.weight.__dict__["var_name"], metadata_name) self.assertEqual(len(list(module.parameters())), 1) - def test_alias_cache_does_not_keep_replaced_source_alive(self) -> None: + def test_untying_parameter_keeps_both_live_parameters_observable(self) -> None: + source = nn.Parameter(torch.ones(2)) module = nn.Module() + module.weight = source + module.alias = source self.track(module) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], module.weight) + dump.reset_mock() + module.weight = nn.Parameter(torch.zeros(2)) + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 2) + observed = [call.args[0] for call in dump.call_args_list] + self.assertEqual(sum(value is module.weight for value in observed), 1) + self.assertEqual(sum(value is module.alias for value in observed), 1) + self.assertNotEqual( + module.weight.__dict__["var_name"], module.alias.__dict__["var_name"] + ) + dump.reset_mock() + weight = module.weight + alias = module.alias + assert isinstance(weight, ProxyParameter) + assert isinstance(alias, ProxyParameter) + weight.register_object() + alias.register_object() + self.registry.dump_modified(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 2) + observed = [call.args[0] for call in dump.call_args_list] + self.assertEqual(sum(value is module.weight for value in observed), 1) + self.assertEqual(sum(value is module.alias for value in observed), 1) + + def test_retained_source_updates_are_preserved_when_aliased(self) -> None: + source = nn.Parameter(torch.ones(2)) + module = nn.Module() + module.weight = source + self.track(module) + source.requires_grad_(False) + gradient = torch.full((2,), 7.0) + source.grad = gradient + marker = object() + setattr(source, "user_marker", marker) + module.alias = source + self.assertIs(module.alias, module.weight) + self.assertFalse(module.alias.requires_grad) + self.assertIs(module.alias.grad, gradient) + self.assertIs(module.alias.__dict__["user_marker"], marker) + + def test_unchanged_source_alias_does_not_erase_accumulated_gradient(self) -> None: + source = nn.Parameter(torch.ones(2)) + module = nn.Module() + module.weight = source + self.track(module) + module.weight.sum().backward() + gradient = module.weight.grad + self.assertIsNotNone(gradient) + module.alias = source + self.assertIs(module.alias, module.weight) + self.assertIs(module.weight.grad, gradient) + torch.testing.assert_close(module.weight.grad, torch.ones(2)) + + def test_hooks_added_to_retained_source_apply_after_alias_assignment(self) -> None: + source = nn.Parameter(torch.ones(2)) + module = nn.Module() + module.weight = source + self.track(module) + calls: list[torch.Tensor] = [] + + def double_gradient(gradient: torch.Tensor) -> torch.Tensor: + calls.append(gradient.clone()) + return gradient * 2 + + handle = source.register_hook(double_gradient) + self.addCleanup(handle.remove) + module.alias = source + module.alias.sum().backward() + self.assertEqual(len(calls), 1) + torch.testing.assert_close(module.weight.grad, torch.full((2,), 2.0)) + handle.remove() + module.weight.grad = None + module.alias.sum().backward() + self.assertEqual(len(calls), 1) + torch.testing.assert_close(module.weight.grad, torch.ones(2)) + + def test_parameter_identity_lives_until_last_model_binding_is_removed(self) -> None: + module = nn.Module() parameter = nn.Parameter(torch.ones(2)) reference = weakref.ref(parameter) module.weight = parameter + self.track(module) + self.assertIs(module.weight, parameter) del parameter gc.collect() + self.assertIs(reference(), module.weight) + setattr(module, "weight", None) + gc.collect() self.assertIsNone(reference()) + self.assertEqual(self.registry.registry, {}) + + def test_removed_source_cannot_overwrite_replacement_registry_entry(self) -> None: + module = nn.Module() + source = nn.Parameter(torch.ones(2)) + module.weight = source + self.track(module) + reference = weakref.ref(source) module.weight = nn.Parameter(torch.zeros(2)) - torch.testing.assert_close(module.weight, torch.zeros(2)) + replacement = module.weight + source.grad = torch.full((2,), 7.0) + self.assertIs(self.registry.registry["model.weight"].var, replacement) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], replacement) + module.weight = source + self.assertIs(module.weight, source) + self.assertIs(self.registry.registry["model.weight"].var, source) + setattr(module, "weight", None) + del source + gc.collect() + self.assertIsNone(reference()) - def test_replaced_proxy_is_not_retained_by_alias_cache(self) -> None: + def test_replaced_child_parameters_leave_model_observation(self) -> None: + old_child = nn.Linear(2, 2) module = nn.Module() + module.child = old_child self.track(module) + module.child = nn.Linear(2, 2) + old_child.weight.grad = torch.ones_like(old_child.weight) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + observed = [call.args[0] for call in dump.call_args_list] + self.assertEqual(len(observed), 2) + self.assertEqual(sum(value is module.child.weight for value in observed), 1) + self.assertEqual(sum(value is module.child.bias for value in observed), 1) + self.assertIs( + self.registry.registry["model.child.weight"].var, module.child.weight + ) + + def test_tracking_does_not_retain_root_module(self) -> None: + module = nn.Module() source = nn.Parameter(torch.ones(2)) module.weight = source - proxy_reference = weakref.ref(module.weight) - module.weight = nn.Parameter(torch.zeros(2)) + self.track(module) + root_reference = weakref.ref(module) + del module gc.collect() - self.assertIsNone(proxy_reference()) + self.assertIsNone(root_reference()) + source.grad = torch.ones(2) + self.assertEqual(self.registry.registry, {}) + + def test_old_session_cleanup_cannot_disable_retracked_parameter(self) -> None: + source = nn.Parameter(torch.ones(2)) + old_root = nn.Module() + old_root.weight = source + self.track(old_root) + new_root = nn.Module() + new_root.weight = source + self.track(new_root) + del old_root + gc.collect() + source.grad = torch.ones(2) + self.assertIs(self.registry.registry["model.weight"].var, new_root.weight) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], source) + + def test_shared_parameter_survives_collection_of_either_tracked_root(self) -> None: + for removed_index in (0, 1): + with self.subTest(removed_index=removed_index): + source = nn.Parameter(torch.ones(2)) + roots = [nn.Module(), nn.Module()] + roots[0].weight = source + roots[1].weight = source + proxy_parameter(roots[0], parent_name="left", should_dump_trace=False) + proxy_parameter(roots[1], parent_name="right", should_dump_trace=False) + reference = weakref.ref(roots[removed_index]) + del roots[removed_index] + gc.collect() + self.assertIsNone(reference()) + expected_name = "right.weight" if removed_index == 0 else "left.weight" + self.assertEqual(source.__dict__["var_name"], expected_name) + source.grad = torch.ones(2) + self.assertEqual(list(self.registry.registry), [expected_name]) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], roots[0].weight) + roots.clear() + gc.collect() + self.assertEqual(self.registry.registry, {}) + + def test_overlapping_roots_both_observe_parameter_replacement(self) -> None: + child = nn.Module() + source = nn.Parameter(torch.ones(2)) + child.weight = source + parent = nn.Module() + parent.child = child + proxy_parameter(parent, parent_name="parent", should_dump_trace=False) + proxy_parameter(child, parent_name="child", should_dump_trace=False) + child.weight = nn.Parameter(torch.zeros(2)) + source.grad = torch.ones(2) + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], child.weight) + del parent + gc.collect() + self.assertEqual(child.weight.__dict__["var_name"], "child.weight") + self.assertEqual(list(self.registry.registry), ["child.weight"]) + setattr(child, "weight", None) + self.assertEqual(self.registry.registry, {}) + + def test_retracking_same_root_moves_namespace_without_stale_entries(self) -> None: + module = nn.Module() + source = nn.Parameter(torch.ones(2)) + module.weight = source + proxy_parameter(module, parent_name="old", should_dump_trace=False) + proxy_parameter(module, parent_name="new", should_dump_trace=False) + gc.collect() + self.assertIs(module.weight, source) + self.assertEqual(list(self.registry.registry), ["new.weight"]) + self.assertEqual(source.__dict__["var_name"], "new.weight") + module.alias = source + with patch.object(ProxyParameter, "dump_trace", autospec=True) as dump: + self.registry.dump_sample(dump_config={"torch.nn.Parameter": True}) + self.assertEqual(dump.call_count, 1) + self.assertIs(dump.call_args.args[0], source) + del module + gc.collect() + self.assertEqual(self.registry.registry, {}) + + def test_custom_subclass_keeps_legacy_conversion_with_current_source_state( + self, + ) -> None: + source = CustomParameter(torch.ones(2)) + marker = object() + setattr(source, "user_marker", marker) + module = nn.Module() module.weight = source + module.alias = source + self.track(module) self.assertIsInstance(module.weight, ProxyParameter) - torch.testing.assert_close(module.weight, torch.ones(2)) + self.assertIsNot(module.weight, source) + self.assertIs(module.weight, module.alias) + self.assertIs(module.weight.__dict__["user_marker"], marker) + # The legacy custom-subclass path has distinct source/proxy identities. + # A retained source is converted anew, never reused from a stale cache. + source.requires_grad_(False) + gradient = torch.full((2,), 7.0) + source.grad = gradient + module.later = source + self.assertIsNot(module.later, module.weight) + self.assertFalse(module.later.requires_grad) + self.assertIs(module.later.grad, gradient) + self.assertTrue(module.weight.requires_grad) + + def test_optimizer_created_before_tracking_and_checkpoint_match_plain(self) -> None: + torch.manual_seed(9) + plain = nn.Linear(2, 1) + tracked = nn.Linear(2, 1) + tracked.load_state_dict(plain.state_dict()) + plain_optimizer = torch.optim.SGD(plain.parameters(), lr=0.1, momentum=0.9) + tracked_optimizer = torch.optim.SGD(tracked.parameters(), lr=0.1, momentum=0.9) + original = tracked.weight + pointer = original.data_ptr() + self.track(tracked) + self.assertIs(tracked.weight, original) + self.assertEqual(tracked.weight.data_ptr(), pointer) + self.assertIs(tracked_optimizer.param_groups[0]["params"][0], original) + source = torch.randn(3, 2) + for _ in range(3): + for model, optimizer in ( + (plain, plain_optimizer), + (tracked, tracked_optimizer), + ): + optimizer.zero_grad() + model(source).square().mean().backward() + optimizer.step() + torch.testing.assert_close( + tracked.state_dict(), plain.state_dict(), rtol=0, atol=0 + ) + torch.testing.assert_close( + tracked_optimizer.state_dict(), plain_optimizer.state_dict(), rtol=0, atol=0 + ) + checkpoint = io.BytesIO() + torch.save(tracked.state_dict(), checkpoint) + checkpoint.seek(0) + restored = nn.Linear(2, 1) + restored.load_state_dict(torch.load(checkpoint, weights_only=True)) + torch.testing.assert_close( + restored.state_dict(), plain.state_dict(), rtol=0, atol=0 + ) def test_assignment_trace_contains_each_new_parameter_once(self) -> None: module = nn.Module() diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index ea0c741a..074ef4d1 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -3,6 +3,7 @@ import os import threading import weakref +from dataclasses import dataclass, field from typing import Callable, cast import torch @@ -16,9 +17,11 @@ from traincheck.utils import get_timestamp_ns, typename from .proxy_basics import is_fake_tensor -from .proxy_registry import get_global_registry +from .proxy_registry import RegistryEntry, get_global_registry SUBCLASS_HOOK_KEY = "_tc_setattr_hook" +MODEL_BINDINGS_KEY = "_tc_model_bindings" +MODEL_CLEANUP_KEY = "_tc_model_tracking_cleanup" logger = logging.getLogger(__name__) @@ -32,7 +35,27 @@ def in_dynamo() -> bool: return False +@dataclass +class _ParameterBindings: + paths: dict[object, str] = field(default_factory=dict) + + +@dataclass +class _ModuleTrackingHooks: + callbacks: dict[object, Callable[[nn.Module, str], None]] = field( + default_factory=dict + ) + + def __call__(self, owner: nn.Module, name: str) -> None: + for callback in tuple(self.callbacks.values()): + callback(owner, name) + + class ProxyParameter(torch.nn.Parameter): + process_id: int + thread_id: int | None + var_name: str + loglevel = logging.INFO def __new__( @@ -175,13 +198,25 @@ def update_timestamp(self): # TODO: # Proxy.var_dict[self.__dict__["var_name"]].last_update_timestamp = current_time - def register_object(self): + def _model_tracking_enabled(self) -> bool: + # Direct construction has no model ownership metadata and remains active. + # Model-owned proxies remain observable while any tracked root owns them. + if MODEL_BINDINGS_KEY not in self.__dict__: + return True + bindings: object = self.__dict__[MODEL_BINDINGS_KEY] + if not isinstance(bindings, _ParameterBindings): + raise TypeError("ProxyParameter model bindings must be _ParameterBindings") + return bool(bindings.paths) + + def register_object(self) -> None: + if not self._model_tracking_enabled(): + return get_global_registry().add_var( self, self.__dict__["var_name"], self.__dict__["var_type"] ) - def dump_trace(self, phase, dump_loc): - if config.DISABLE_WRAPPER: + def dump_trace(self, phase: str, dump_loc: str | None) -> None: + if config.DISABLE_WRAPPER or not self._model_tracking_enabled(): return # TODO @@ -224,38 +259,138 @@ def proxy_parameter( from_call: bool = False, from_iter: bool = False, ) -> None: - """Track a module's parameters and subsequent assignments to its module tree. - - Aliases within this tracking session share one proxy and its first registered - name. Weak identity entries avoid retaining replaced source parameters or - comparing Tensor values. Separate calls do not share this cache. Callers must - serialize model tracking and assignment; concurrent model mutation is unsupported. + """Track a module tree while preserving ordinary Parameter identity. + + Exact nn.Parameter objects are promoted in place: source references, aliases, + autograd state, and existing optimizer references remain one object. Custom + Parameter subclasses keep the legacy out-of-place conversion. Their aliases + within one traversal share a proxy, but retained custom-source references stay + separate; assigning a source again converts its current state anew. + + Observation uses the first registered live root's canonical parameter path. + Attribute assignments (including None) reconcile paths; ordinary attributes + do not rescan the tree. Direct register_parameter/register_module calls and + delattr remain outside this setter hook. Collecting one root preserves other + roots' ownership. Roots and modules are held weakly; retracking a root replaces + its prior session. + Distinct roots should use distinct parent_name namespaces. Callers must + serialize tracking/assignment; concurrent model mutation is unsupported. """ if not isinstance(module, nn.Module): raise TypeError("proxy_parameter expects a torch.nn.Module") if in_dynamo(): return + if MODEL_CLEANUP_KEY in module.__dict__: + previous_cleanup: object = module.__dict__[MODEL_CLEANUP_KEY] + if not callable(previous_cleanup): + raise TypeError("TrainCheck model tracking cleanup must be callable") + previous_cleanup() + + ownership = object() + tracked: dict[int, weakref.ReferenceType[ProxyParameter]] = {} + hooked_modules: dict[int, weakref.ReferenceType[nn.Module]] = {} + + def bindings_for(parameter: ProxyParameter) -> _ParameterBindings: + if MODEL_BINDINGS_KEY not in parameter.__dict__: + parameter.__dict__[MODEL_BINDINGS_KEY] = _ParameterBindings() + bindings: object = parameter.__dict__[MODEL_BINDINGS_KEY] + if not isinstance(bindings, _ParameterBindings): + raise TypeError("ProxyParameter model bindings must be _ParameterBindings") + return bindings + + def reconcile(parameters: dict[int, ProxyParameter]) -> None: + registry = get_global_registry() + with registry.registry_lock: + names: list[tuple[ProxyParameter, str | None]] = [] + for parameter in parameters.values(): + old_name: object = parameter.__dict__["var_name"] + if not isinstance(old_name, str): + raise TypeError("ProxyParameter observation name must be str") + paths = bindings_for(parameter).paths + new_name = next(iter(paths.values())) if paths else None + if new_name is not None and not isinstance(new_name, str): + raise TypeError("ProxyParameter model binding path must be str") + if old_name != new_name: + entry = registry.registry.get(old_name) + if entry is not None and entry.var is parameter: + del registry.registry[old_name] + names.append((parameter, new_name)) + # Remove all departed paths before publishing replacements, including + # an old tied parameter moving to its surviving alias's name. + for parameter, name in names: + if name is None: + continue + parameter.__dict__["var_name"] = name + entry = registry.registry.get(name) + if entry is None or entry.var is not parameter: + registry.registry[name] = RegistryEntry( + parameter, name, "torch.nn.Parameter", stale=False + ) - aliases: dict[ - int, - tuple[weakref.ReferenceType[nn.Parameter], weakref.ReferenceType[nn.Parameter]], - ] = {} - - def track_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: + def remove_hook(owner: nn.Module) -> None: + if SUBCLASS_HOOK_KEY not in owner.__dict__: + return + hook: object = owner.__dict__[SUBCLASS_HOOK_KEY] + if not isinstance(hook, _ModuleTrackingHooks): + raise TypeError( + "TrainCheck model tracking hook must be _ModuleTrackingHooks" + ) + hook.callbacks.pop(ownership, None) + if not hook.callbacks: + del owner.__dict__[SUBCLASS_HOOK_KEY] + + def cleanup() -> None: + affected: dict[int, ProxyParameter] = {} + for identity, reference in tracked.items(): + parameter = reference() + if parameter is not None: + bindings_for(parameter).paths.pop(ownership, None) + affected[identity] = parameter + reconcile(affected) + tracked.clear() + for reference_to_module in hooked_modules.values(): + owner = reference_to_module() + if owner is None: + continue + remove_hook(owner) + hooked_modules.clear() + root = root_reference() + if root is not None and root.__dict__.get(MODEL_CLEANUP_KEY) is cleanup: + del root.__dict__[MODEL_CLEANUP_KEY] + + def release_root(reference: weakref.ReferenceType[nn.Module]) -> None: + cleanup() + + root_reference = weakref.ref(module, release_root) + module.__dict__[MODEL_CLEANUP_KEY] = cleanup + + def convert_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: if not isinstance(parameter, nn.Parameter): raise TypeError( f"Registered parameter {name!r} is not a torch.nn.Parameter" ) - # Calling ProxyParameter again reruns __init__, overwriting its canonical - # name and emitting another initialization record even for the same object. - if isinstance(parameter, ProxyParameter): + if isinstance(parameter, ProxyParameter) or is_fake_tensor(parameter): return parameter - identity = id(parameter) - cached = aliases.get(identity) # Absence means this source has no live proxy. - if cached is not None and cached[0]() is parameter: - existing = cached[1]() - if existing is not None: - return existing + if type(parameter) is nn.Parameter: + # A fresh Tensor wrapper splits source and proxy autograd metadata. + # Promotion preserves the Python object, TensorImpl, and every slot. + torch.Tensor.__setattr__(parameter, "__class__", ProxyParameter) + assert isinstance(parameter, ProxyParameter) + bindings_for(parameter).paths[ownership] = name + ProxyParameter.__init__( + parameter, + parameter, + logdir, + log_level, + name, + "torch.nn.Parameter", + should_dump_trace, + from_call, + from_iter, + ) + return parameter + # Preserve the established conversion for custom Parameter subclasses; + # their class/layout/dispatch semantics cannot safely be promoted in place. proxied = ProxyParameter( parameter, logdir, @@ -266,67 +401,92 @@ def track_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: from_call, from_iter, ) - - # Tensor autograd hook dictionaries are slots, not entries in __dict__. - # Share them rather than copying so existing removable handles still work. + # Autograd hook dictionaries are slots, not __dict__ entries. Sharing them + # preserves removal by handles registered before the legacy conversion. torch.Tensor.__setattr__(proxied, "_backward_hooks", parameter._backward_hooks) - # Post-accumulation hooks are absent on older supported PyTorch versions. if hasattr(parameter, "_post_accumulate_grad_hooks"): torch.Tensor.__setattr__( proxied, "_post_accumulate_grad_hooks", parameter._post_accumulate_grad_hooks, ) - - def forget(reference: weakref.ReferenceType[nn.Parameter]) -> None: - current = aliases.get(identity) - if current is not None and ( - current[0] is reference or current[1] is reference - ): - del aliases[identity] - - aliases[identity] = ( - weakref.ref(parameter, forget), - weakref.ref(proxied, forget), - ) return proxied - def track_module(current: nn.Module, path: str, visited: set[int]) -> None: - if not isinstance(current, nn.Module): - raise TypeError(f"Registered child {path!r} is not a torch.nn.Module") - if id(current) in visited: + def on_assignment(owner: nn.Module, name: str) -> None: + if in_dynamo() or ( + name not in owner._parameters and name not in owner._modules + ): return - visited.add(id(current)) - # Include aliases rather than named_parameters' default deduplication. - for name, parameter in list(current._parameters.items()): - if parameter is not None: - current._parameters[name] = track_parameter( - parameter, path + "." + name - ) - for name, child in current._modules.items(): - if child is not None: - track_module(child, path + "." + name, visited) - - def subclass_setattr_hook(owner: nn.Module, name: str) -> None: - if in_dynamo(): + root = root_reference() + if root is not None: + synchronize(root) + + def synchronize(root: nn.Module) -> None: + observed: dict[int, tuple[ProxyParameter, str]] = {} + # This cache is traversal-local, so custom-source state cannot become + # stale between assignments. Strong references prevent identity reuse. + converted: dict[int, tuple[nn.Parameter, nn.Parameter]] = {} + visited: dict[int, weakref.ReferenceType[nn.Module]] = {} + + def visit(current: nn.Module, path: str) -> None: + if not isinstance(current, nn.Module): + raise TypeError(f"Registered child {path!r} is not a torch.nn.Module") + if id(current) in visited: return - # The native setter has already validated and registered the value, - # including any replacement made by PyTorch registration hooks. Buffer - # and ordinary Tensor assignments must never become parameters. - if name in owner._parameters: - parameter = owner._parameters[name] - if parameter is not None: - owner._parameters[name] = track_parameter( - parameter, path + "." + name - ) - elif name in owner._modules: - child = owner._modules[name] + visited[id(current)] = weakref.ref(current) + for name, parameter in list(current._parameters.items()): + if parameter is None: + continue + identity = id(parameter) + cached = converted.get(identity) + if cached is None: + proxied = convert_parameter(parameter, path + "." + name) + converted[identity] = (parameter, proxied) + else: + assert cached[0] is parameter + proxied = cached[1] + current._parameters[name] = proxied + if isinstance(proxied, ProxyParameter) and id(proxied) not in observed: + observed[id(proxied)] = (proxied, path + "." + name) + for name, child in current._modules.items(): if child is not None: - track_module(child, path + "." + name, set()) + visit(child, path + "." + name) + if SUBCLASS_HOOK_KEY not in current.__dict__: + current.__dict__[SUBCLASS_HOOK_KEY] = _ModuleTrackingHooks() + hook: object = current.__dict__[SUBCLASS_HOOK_KEY] + if not isinstance(hook, _ModuleTrackingHooks): + raise TypeError( + "TrainCheck model tracking hook must be _ModuleTrackingHooks" + ) + hook.callbacks[ownership] = on_assignment - current.__dict__[SUBCLASS_HOOK_KEY] = subclass_setattr_hook + visit(root, parent_name) + affected: dict[int, ProxyParameter] = {} + for identity, reference in tracked.items(): + parameter = reference() + if parameter is not None: + if identity not in observed: + bindings_for(parameter).paths.pop(ownership, None) + affected[identity] = parameter + for identity, (parameter, path) in observed.items(): + bindings_for(parameter).paths[ownership] = path + affected[identity] = parameter + reconcile(affected) + for identity, reference_to_module in hooked_modules.items(): + if identity in visited: + continue + owner = reference_to_module() + if owner is not None: + remove_hook(owner) + hooked_modules.clear() + hooked_modules.update(visited) + tracked.clear() + tracked.update( + (identity, weakref.ref(parameter)) + for identity, (parameter, _) in observed.items() + ) - track_module(module, parent_name, set()) + synchronize(module) # instrument torch.nn.Module's setattr From ba266c920d66d40e7ee277b08274de298ae74a31 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 03:37:38 +0800 Subject: [PATCH 4/4] Retire cyclic model roots safely across registry sampling --- tests/fixtures/registry_gc_probe.py | 140 ++++++++++++++++ tests/test_registry_lifecycle.py | 149 +++++++++++++++++ .../proxy_wrapper/proxy_registry.py | 158 +++++++++++------- .../instrumentor/proxy_wrapper/subclass.py | 27 ++- 4 files changed, 406 insertions(+), 68 deletions(-) create mode 100644 tests/fixtures/registry_gc_probe.py create mode 100644 tests/test_registry_lifecycle.py diff --git a/tests/fixtures/registry_gc_probe.py b/tests/fixtures/registry_gc_probe.py new file mode 100644 index 00000000..22e6752f --- /dev/null +++ b/tests/fixtures/registry_gc_probe.py @@ -0,0 +1,140 @@ +"""Bounded subprocess probe for finalization while registry dumping holds its lock.""" + +import gc +import sys +import weakref +from collections.abc import Callable +from unittest.mock import patch + +from torch import nn + +import traincheck.config.config as config +from traincheck.instrumentor.proxy_wrapper import subclass +from traincheck.instrumentor.proxy_wrapper.proxy_registry import VarRegistry +from traincheck.instrumentor.proxy_wrapper.subclass import proxy_parameter + + +class CallbackVariable: + def __init__(self, name: str, calls: list[str]) -> None: + self.name = name + self.calls = calls + self.callback: Callable[[], None] | None = None + + def dump_trace(self, phase: str, dump_loc: str | None) -> None: + self.calls.append(self.name) + callback = self.callback + self.callback = None + if callback is not None: + callback() + + +def run_reentrant() -> None: + registry = VarRegistry() + calls: list[str] = [] + first = CallbackVariable("first", calls) + second = CallbackVariable("second", calls) + registry.add_var(first, "first", "kind") + registry.add_var(second, "second", "kind") + + def nested_sample() -> None: + registry.dump_modified(dump_config={"kind": {}}) + + first.callback = nested_sample + registry.dump_modified(dump_config={"kind": {}}) + assert calls == ["first", "second"] + calls.clear() + registry.dump_modified(dump_config={"kind": {}}) + assert calls == [] + + def mark_second_modified() -> None: + registry.add_var(second, "second", "kind") + + first.callback = mark_second_modified + registry.add_var(first, "first", "kind") + registry.dump_modified(dump_config={"kind": {}}) + assert calls == ["first", "second"] + calls.clear() + + replacement = CallbackVariable("replacement", calls) + added = CallbackVariable("added", calls) + + def change_membership() -> None: + registry.add_var(replacement, "second", "kind") + registry.add_var(added, "added", "kind") + + first.callback = change_membership + registry.dump_sample(dump_config={"kind": {}}) + assert calls == ["first"] + calls.clear() + registry.dump_modified(dump_config={"kind": {}}) + assert calls == ["replacement", "added"] + calls.clear() + registry.dump_modified(dump_config={"kind": {}}) + assert calls == [] + + +def run(mode: str) -> None: + if mode == "reentrant": + run_reentrant() + print("OK reentrant") + return + if mode not in ("sample", "modified"): + raise ValueError(f"Unknown sampling mode: {mode}") + registry = VarRegistry() + records: list[dict[str, object]] = [] + + def collect_during_conversion(obj: object, value: object) -> dict[str, object]: + gc.collect() + return {} + + with ( + patch.object(subclass, "get_global_registry", return_value=registry), + patch.object(config, "DISABLE_WRAPPER", False), + ): + shared = nn.Linear(2, 2) + first = nn.Sequential(shared) + second = nn.Sequential(shared) + setattr(first, "saved_forward", first.forward) + proxy_parameter(first, parent_name="first", should_dump_trace=False) + proxy_parameter(second, parent_name="second", should_dump_trace=False) + first_reference = weakref.ref(first) + gc_was_enabled = gc.isenabled() + gc.disable() + try: + del first + assert first_reference() is not None + with ( + patch.object( + subclass, "dump_attributes", side_effect=collect_during_conversion + ), + patch.object(subclass, "dump_trace_VAR", side_effect=records.append), + ): + if mode == "sample": + registry.dump_sample(dump_config={"torch.nn.Parameter": {}}) + else: + registry.dump_modified(dump_config={"torch.nn.Parameter": {}}) + assert first_reference() is None + # The current sample completes under its original canonical paths; + # deferred retirement publishes surviving-root paths afterwards. + assert [record["var_name"] for record in records] == [ + "first.0.weight", + "first.0.bias", + ] + assert list(registry.registry) == ["second.0.weight", "second.0.bias"] + records.clear() + registry.dump_modified(dump_config={"torch.nn.Parameter": {}}) + assert [record["var_name"] for record in records] == [ + "second.0.weight", + "second.0.bias", + ] + records.clear() + registry.dump_modified(dump_config={"torch.nn.Parameter": {}}) + assert records == [] + finally: + if gc_was_enabled: + gc.enable() + print(f"OK {mode}") + + +if __name__ == "__main__": + run(sys.argv[1]) diff --git a/tests/test_registry_lifecycle.py b/tests/test_registry_lifecycle.py new file mode 100644 index 00000000..49231b60 --- /dev/null +++ b/tests/test_registry_lifecycle.py @@ -0,0 +1,149 @@ +"""Registry sampling and model-finalization contracts.""" + +import gc +import os +import subprocess +import sys +import unittest +import weakref +from pathlib import Path +from unittest.mock import patch + +from torch import nn + +import traincheck.config.config as config +from traincheck.instrumentor.proxy_wrapper import subclass +from traincheck.instrumentor.proxy_wrapper.proxy_registry import VarRegistry +from traincheck.instrumentor.proxy_wrapper.subclass import ( + ProxyParameter, + proxy_parameter, +) + + +class RecordingVariable: + def __init__( + self, name: str, calls: list[tuple[str, str]], error: Exception | None = None + ) -> None: + self.name = name + self.calls = calls + self.error = error + + def dump_trace(self, phase: str, dump_loc: str | None) -> None: + self.calls.append((self.name, phase)) + if self.error is not None: + raise self.error + + +class RegistryLifecycleTests(unittest.TestCase): + def test_cyclic_model_is_retired_after_collection(self) -> None: + registry = VarRegistry() + with ( + patch.object(subclass, "get_global_registry", return_value=registry), + patch.object(config, "DISABLE_WRAPPER", True), + ): + model = nn.Linear(2, 2) + setattr(model, "saved_forward", model.forward) + proxy_parameter(model, parent_name="model", should_dump_trace=False) + reference = weakref.ref(model) + parameter = model.weight + del model + gc.collect() + self.assertIsNone(reference()) + self.assertEqual(registry.registry, {}) + assert isinstance(parameter, ProxyParameter) + parameter.register_object() + self.assertEqual(registry.registry, {}) + + def test_gc_during_both_sampling_modes_finishes_and_preserves_shared_owner( + self, + ) -> None: + probe = Path(__file__).parent / "fixtures" / "registry_gc_probe.py" + root = Path(__file__).resolve().parents[1] + for mode in ("sample", "modified"): + with self.subTest(mode=mode): + result = subprocess.run( + [sys.executable, str(probe), mode], + cwd=root, + text=True, + capture_output=True, + timeout=20, + check=False, + env={**os.environ, "PYTHONPATH": str(root)}, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn(f"OK {mode}", result.stdout) + + def test_nested_sampling_and_reentrant_updates_preserve_live_stale_status( + self, + ) -> None: + probe = Path(__file__).parent / "fixtures" / "registry_gc_probe.py" + root = Path(__file__).resolve().parents[1] + result = subprocess.run( + [sys.executable, str(probe), "reentrant"], + cwd=root, + text=True, + capture_output=True, + timeout=20, + check=False, + env={**os.environ, "PYTHONPATH": str(root)}, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertIn("OK reentrant", result.stdout) + + def test_missing_configuration_remains_an_error_without_stale_changes(self) -> None: + registry = VarRegistry() + calls: list[tuple[str, str]] = [] + registry.add_var(RecordingVariable("variable", calls), "variable", "kind") + with self.assertRaises(AttributeError): + registry.dump_sample() + with self.assertRaises(AttributeError): + registry.dump_modified() + self.assertEqual(calls, []) + self.assertFalse(registry.registry["variable"].stale) + + def test_sampling_preserves_order_filter_and_stale_flags(self) -> None: + registry = VarRegistry() + calls: list[tuple[str, str]] = [] + first = RecordingVariable("first", calls) + second = RecordingVariable("second", calls) + filtered = RecordingVariable("filtered", calls) + registry.add_var(first, "first", "selected") + registry.add_var(second, "second", "selected") + registry.add_var(filtered, "filtered", "other") + registry.dump_sample(dump_config={"selected": True}) + self.assertEqual(calls, [("first", "sample"), ("second", "sample")]) + self.assertTrue(registry.registry["first"].stale) + self.assertTrue(registry.registry["second"].stale) + self.assertFalse(registry.registry["filtered"].stale) + registry.dump_modified(dump_config={"selected": True}) + self.assertEqual(len(calls), 2) + registry.add_var(first, "first", "selected") + registry.dump_modified(dump_config={"selected": True}) + self.assertEqual(calls[-1], ("first", "selective-sample")) + self.assertEqual(len(calls), 3) + + def test_sampling_propagates_error_without_marking_later_entries_stale( + self, + ) -> None: + for mode in ("sample", "modified"): + with self.subTest(mode=mode): + registry = VarRegistry() + calls: list[tuple[str, str]] = [] + error = RuntimeError("dump failed") + registry.add_var( + RecordingVariable("first", calls, error), "first", "kind" + ) + registry.add_var(RecordingVariable("later", calls), "later", "kind") + with self.assertRaises(RuntimeError) as caught: + if mode == "sample": + registry.dump_sample(dump_config={"kind": True}) + else: + registry.dump_modified(dump_config={"kind": True}) + self.assertIs(caught.exception, error) + self.assertTrue(registry.registry["first"].stale) + self.assertFalse(registry.registry["later"].stale) + self.assertEqual(len(calls), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/instrumentor/proxy_wrapper/proxy_registry.py b/traincheck/instrumentor/proxy_wrapper/proxy_registry.py index 75c464d8..66c640f3 100644 --- a/traincheck/instrumentor/proxy_wrapper/proxy_registry.py +++ b/traincheck/instrumentor/proxy_wrapper/proxy_registry.py @@ -1,10 +1,19 @@ import threading +from collections import deque +from collections.abc import Callable, Mapping +from typing import Literal, Protocol + + +class TraceableVariable(Protocol): + def dump_trace(self, phase: str, dump_loc: str | None) -> None: ... class RegistryEntry: - """A class to store the tracked object and its associated metadata""" + """A tracked object and its associated observation metadata.""" - def __init__(self, obj, var_name, var_type, stale): + def __init__( + self, obj: TraceableVariable, var_name: str, var_type: str, stale: bool + ) -> None: self.var = obj self.var_name = var_name self.var_type = var_type @@ -12,22 +21,25 @@ def __init__(self, obj, var_name, var_type, stale): class VarRegistry: - """A helper class managing all variables being tracked and allow for controlled dumps of - the variable states. - - A variable is uniquely identified by its "name" - When a variable is added to the registry, it is marked as "not stale". - When a variable is dumped through `dump_sample` or `dump_modified`, it is marked as "stale". - A variable is only dumped through `dump_modified` if it is not stale. + """Name-indexed variables with ordered full and modified-only sampling. + Sampling snapshots membership and order, checking type and stale status at + each entry's turn. Additions wait for the next call; replaced entries are skipped. + Root finalization is deferred until the outermost sampling call ends, keeping + a record's canonical name stable throughout conversion and serialization. """ - def __init__(self): + def __init__(self) -> None: self.registry: dict[str, RegistryEntry] = {} - self.registry_lock = threading.Lock() - - def add_var(self, var, var_name: str, var_type: str): - """Add a new proxy variable to the registry""" + # GC may request retirement synchronously while the sampling thread owns + # the lock. Reentry only queues retirement until sampling is finished. + self.registry_lock = threading.RLock() + self._sampling_depth = 0 + self._pending_cleanups: deque[Callable[[], None]] = deque() + self._draining_cleanups = False + + def add_var(self, var: TraceableVariable, var_name: str, var_type: str) -> None: + """Register a variable or mark its current observation as modified.""" with self.registry_lock: if var_name in self.registry: self.registry[var_name].var = var @@ -39,59 +51,83 @@ def add_var(self, var, var_name: str, var_type: str): var, var_name, var_type, stale=False ) - def dump_sample(self, dump_loc=None, dump_config=None): - """A complete dump of all present proxy objects - - Calling this API mark all proxy objects as stale which - will affect the `dump_modified` API. - """ - to_dump_types = set(dump_config.keys()) + def defer_cleanup(self, cleanup: Callable[[], None]) -> None: + """Retire immediately when idle, or after the current sample completes.""" with self.registry_lock: - for _, entry in self.registry.items(): - var_type = entry.var_type - if var_type not in to_dump_types: - continue - entry.stale = True - entry.var.dump_trace(phase="sample", dump_loc=dump_loc) - - def dump_modified(self, dump_loc=None, dump_config=None): - """Dump only the proxy variables that might be modified since last dump - - args: - dump_loc: the location to dump the trace, an optional string to add to trace records - dump_config: the config for dumping, each key would be the type of the variable and the value - would be whether to dump all changed vars or just one - - ** This is a middle ground between blindly dump everything everytime v.s. fully-accurate delta dumping ** - fully-accuracy dumping is hard as for each "modifications" to the variable, you will need to compare - the new state v.s. the old state to ensure the state has actually changed, which introduces great overhead. - - This function implements delta dumping but does not guarantee two consecutive dumps will be different, - we only guarantee that between two dumps there has been attempts (e.g. through __setattr__ or observer) - to modify the variable. - - - Side effects: - when calling the function, all dumped proxy vars will be marked as stale and will not be dumped next time - unless there are new modification attempts to t - """ + self._pending_cleanups.append(cleanup) + self._drain_cleanups() + + def _drain_cleanups(self) -> None: + # Caller holds registry_lock. Preserve queue order during nested requests. + if self._sampling_depth or self._draining_cleanups: + return + self._draining_cleanups = True + try: + while self._pending_cleanups: + self._pending_cleanups.popleft()() + finally: + self._draining_cleanups = False + + def _dump( + self, + phase: Literal["sample", "selective-sample"], + dump_loc: str | None, + dump_config: Mapping[str, object] | None, + ) -> None: + if dump_config is None: + # Absence remains an error, as in the original .keys() call; it must + # not silently become an empty filter and a successful empty sample. + raise AttributeError("dump_config is required for registry sampling") to_dump_types = set(dump_config.keys()) with self.registry_lock: - for _, entry in self.registry.items(): - var_type = entry.var_type - if var_type not in to_dump_types: - continue - - if entry.stale: - continue - - entry.stale = True - entry.var.dump_trace(phase="selective-sample", dump_loc=dump_loc) + self._sampling_depth += 1 + try: + entries = [ + (name, entry, entry.var) for name, entry in self.registry.items() + ] + for name, entry, variable in entries: + if ( + self.registry.get(name) is not entry + or entry.var is not variable + ): + continue + if entry.var_type not in to_dump_types: + continue + if phase == "selective-sample" and entry.stale: + continue + # Preserve the original failure contract: the attempted entry + # is stale even if dumping raises; later entries stay unchanged. + entry.stale = True + variable.dump_trace(phase=phase, dump_loc=dump_loc) + finally: + self._sampling_depth -= 1 + self._drain_cleanups() + + def dump_sample( + self, + dump_loc: str | None = None, + dump_config: Mapping[str, object] | None = None, + ) -> None: + """Dump all selected entries in registry order and mark attempts stale.""" + self._dump("sample", dump_loc, dump_config) + + def dump_modified( + self, + dump_loc: str | None = None, + dump_config: Mapping[str, object] | None = None, + ) -> None: + """Dump selected entries still marked modified when their turn arrives. + + Modified means an update was attempted, not that tensor contents differ. + Successful and failed dump attempts are marked stale until add_var marks + them modified again. Configuration values are retained for compatibility; + only its type-name keys currently select entries. + """ + self._dump("selective-sample", dump_loc, dump_config) -# Global dictionary to store registered objects global_registry = VarRegistry() -def get_global_registry(): +def get_global_registry() -> VarRegistry: return global_registry diff --git a/traincheck/instrumentor/proxy_wrapper/subclass.py b/traincheck/instrumentor/proxy_wrapper/subclass.py index 074ef4d1..baf2838a 100644 --- a/traincheck/instrumentor/proxy_wrapper/subclass.py +++ b/traincheck/instrumentor/proxy_wrapper/subclass.py @@ -287,6 +287,7 @@ def proxy_parameter( previous_cleanup() ownership = object() + cleanup_requested = False tracked: dict[int, weakref.ReferenceType[ProxyParameter]] = {} hooked_modules: dict[int, weakref.ReferenceType[nn.Module]] = {} @@ -339,7 +340,7 @@ def remove_hook(owner: nn.Module) -> None: if not hook.callbacks: del owner.__dict__[SUBCLASS_HOOK_KEY] - def cleanup() -> None: + def retire() -> None: affected: dict[int, ProxyParameter] = {} for identity, reference in tracked.items(): parameter = reference() @@ -358,10 +359,20 @@ def cleanup() -> None: if root is not None and root.__dict__.get(MODEL_CLEANUP_KEY) is cleanup: del root.__dict__[MODEL_CLEANUP_KEY] - def release_root(reference: weakref.ReferenceType[nn.Module]) -> None: - cleanup() - - root_reference = weakref.ref(module, release_root) + def cleanup() -> None: + nonlocal cleanup_requested + if cleanup_requested: + return + cleanup_requested = True + finalizer.detach() + get_global_registry().defer_cleanup(retire) + + # finalize keeps the death observer outside the module/session closure cycle, + # so ordinary model cycles (for example a saved bound forward method) retire + # correctly. The callback captures only a weak root, never the root itself. + root_reference = weakref.ref(module) + finalizer = weakref.finalize(module, cleanup) + finalizer.atexit = False module.__dict__[MODEL_CLEANUP_KEY] = cleanup def convert_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: @@ -413,8 +424,10 @@ def convert_parameter(parameter: nn.Parameter, name: str) -> nn.Parameter: return proxied def on_assignment(owner: nn.Module, name: str) -> None: - if in_dynamo() or ( - name not in owner._parameters and name not in owner._modules + if ( + cleanup_requested + or in_dynamo() + or (name not in owner._parameters and name not in owner._modules) ): return root = root_reference()