From a87dcf166c96307e5f970b291cfb1664519f9124 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Tue, 27 Oct 2026 02:58:05 +0800 Subject: [PATCH] Make subclass the default model parameter tracker --- .../gpt2-pretrain-config/config.yml | 2 +- .../mnist-config/config.yml | 2 +- docs/instr.md | 30 +- tests/test_model_tracker_selection.py | 315 ++++++++++++++++++ traincheck/collect_trace.py | 3 +- 5 files changed, 347 insertions(+), 5 deletions(-) create mode 100644 tests/test_model_tracker_selection.py diff --git a/docs/assets/examples/traincheck-collect/gpt2-pretrain-config/config.yml b/docs/assets/examples/traincheck-collect/gpt2-pretrain-config/config.yml index a014282d..fa277178 100644 --- a/docs/assets/examples/traincheck-collect/gpt2-pretrain-config/config.yml +++ b/docs/assets/examples/traincheck-collect/gpt2-pretrain-config/config.yml @@ -7,4 +7,4 @@ shscript: run.sh copy_all_files: true models_to_track: - model -model_tracker_style: proxy # [Optional] "proxy" (default), "subclass", or "sampler". +model_tracker_style: subclass # [Optional] "subclass" (default), "proxy", or "sampler". Keep "proxy" for existing proxy-derived invariants. diff --git a/docs/assets/examples/traincheck-collect/mnist-config/config.yml b/docs/assets/examples/traincheck-collect/mnist-config/config.yml index 4aa57483..3f579d6d 100644 --- a/docs/assets/examples/traincheck-collect/mnist-config/config.yml +++ b/docs/assets/examples/traincheck-collect/mnist-config/config.yml @@ -4,7 +4,7 @@ pyscript: mnist.py # The Python entry point of your training program. shscript: run.sh # [Optional] Shell script to launch the program with custom arguments or environment setup. models_to_track: # [Optional] List of variable names for models you want to track. If omitted, model tracking is disabled. - model -model_tracker_style: proxy # [Optional] Method for model tracking. Choose between "proxy" (default), "subclass", or "sampler". +model_tracker_style: subclass # [Optional] Choose "subclass" (default), "proxy", or "sampler". Keep "proxy" for existing proxy-derived invariants. copy_all_files: false # [Optional] Set to true if your code uses relative paths (e.g., loading local datasets or configs). # This ensures TrainCheck copies the entire working directory before execution. # Note: TrainCheck automatically handles PYTHONPATH. Default is false. diff --git a/docs/instr.md b/docs/instr.md index e558fafe..8fc03f3d 100644 --- a/docs/instr.md +++ b/docs/instr.md @@ -69,7 +69,7 @@ modules_to_instr: - torch models_to_track: - model -model_tracker_style: proxy +model_tracker_style: subclass copy_all_files: false output_dir: traincheck_trace ``` @@ -87,7 +87,33 @@ Config keys use underscores, not hyphens. For example, the CLI flag `--output-di - `--sampling-interval`: collect every Nth step after warm-up. - `--warm-up-steps`: collect the first N steps. - `--copy-all-files`: copy files beside the training script into the output directory. -- `--model-tracker-style`: choose `proxy`, `subclass`, or `sampler`. +- `--model-tracker-style`: choose `subclass` (default), `proxy`, or `sampler`. + +### Choosing a model tracker + +The default `subclass` tracker keeps the model as a PyTorch `nn.Module` and +tracks its parameters using a `torch.nn.Parameter` subclass. It requires +`--models-to-track` (or `models_to_track` in YAML); choosing a tracker alone does +not enable model tracking. Model tracking must be installed before creating the +optimizer, so the optimizer uses the tracked parameters. + +The `proxy` tracker remains available through `--model-tracker-style proxy` or +`model_tracker_style: proxy` in YAML. It wraps model objects and dynamically +accessed attributes, so its coverage differs from the parameter-focused subclass +tracker. Explicit CLI/YAML selections remain supported, and saved +instrumentation options are not rewritten. + +For invariants inferred from proxy traces, keep `proxy` explicit when collecting +the checking trace. To migrate, collect a new reference trace with `subclass`, +infer new invariants, and verify that the intended rules actually activate during +checking. Zero violations alone does not establish coverage; module or buffer +rules from a proxy trace may not activate with parameter-only tracking. + +Selective collection keeps its existing policy: invariants requiring immediate +variable tracking use the selected `subclass` or `proxy` tracker. When only +cross-variable consistency rules require model tracking, collection uses the +`sampler` tracker to observe parameters after optimizer steps. If no invariant +requires model tracking, selective collection disables it. Run the command help for the complete option list: diff --git a/tests/test_model_tracker_selection.py b/tests/test_model_tracker_selection.py new file mode 100644 index 00000000..cc6177b7 --- /dev/null +++ b/tests/test_model_tracker_selection.py @@ -0,0 +1,315 @@ +"""Collection defaults and compatibility contracts for model tracker selection.""" + +import copy +import json +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pandas as pd +import torch + +import traincheck.collect_trace as collect_trace +import traincheck.config.config as config +import traincheck.instrumentor.proxy_wrapper.proxy_config as proxy_config +from tools.fresh_validation.check_coverage import observe +from traincheck.config.config import InstrOpt +from traincheck.instrumentor import dumper, tracer +from traincheck.instrumentor.proxy_wrapper import subclass +from traincheck.instrumentor.proxy_wrapper.proxy_registry import VarRegistry +from traincheck.invariant.base_cls import ( + APIParam, + GroupedPreconditions, + Invariant, + Preconditions, + UnconditionalPrecondition, + VarTypeParam, +) +from traincheck.invariant.consistency_relation import ConsistencyRelation +from traincheck.invariant.contain_relation import ( + APIContainRelation, + cache_events_scanner, +) +from traincheck.trace.trace_pandas import TracePandas +from traincheck.utils import typename + + +def parameter_rule() -> Invariant: + return Invariant( + APIContainRelation, + [APIParam("torch.optim.SGD.step"), VarTypeParam("torch.nn.Parameter", "data")], + GroupedPreconditions( + {"parent_func_call_pre": Preconditions([UnconditionalPrecondition()])} + ), + ) + + +def freeze_parameter(parameter: torch.nn.Parameter) -> None: + parameter.requires_grad = False + + +class ModelTrackerSelectionTests(unittest.TestCase): + def collect( + self, + arguments: list[str], + *, + yaml_options: str | None = None, + invariants: list[Invariant] | None = None, + ) -> tuple[dict[str, object], str, InstrOpt | None]: + """Generate real instrumented source without running a training process.""" + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + script = directory / "training.py" + script.write_text( + "import torch\nmodel = torch.nn.Linear(2, 1)\n" + "optimizer = torch.optim.SGD(model.parameters(), lr=0.1)\n", + encoding="utf-8", + ) + output = directory / "trace" + argv = [ + "traincheck-collect", + "--pyscript", + str(script), + "--output-dir", + str(output), + "--only-instr", + *arguments, + ] + if yaml_options is not None: + configuration = directory / "traincheck.yml" + configuration.write_text(yaml_options, encoding="utf-8") + argv.extend(["--use-config", "--config", str(configuration)]) + if invariants is not None: + argv.extend(["--invariants", "reference.json"]) + + with ( + patch.object(sys, "argv", argv), + patch.object(collect_trace, "dump_env"), + patch.object(collect_trace, "read_inv_file", return_value=invariants), + patch.object( + proxy_config, + "auto_observer_config", + copy.deepcopy(proxy_config.auto_observer_config), + ), + patch.object( + collect_trace.instrumentor, + "instrument_file", + wraps=collect_trace.instrumentor.instrument_file, + ) as instrument, + patch.object(collect_trace.runner, "ProgramRunner") as runner, + ): + runner.return_value.run.return_value = ("", 0) + collect_trace.main() + assert instrument.call_args is not None + assert runner.call_args is not None + options: dict[str, object] = dict(instrument.call_args.kwargs) + source: object = runner.call_args.args[0] + assert isinstance(source, str) + compile(source, str(script), "exec") + + saved_path = output / "instr_opts.json" + saved_options = ( + InstrOpt.from_json(saved_path.read_text(encoding="utf-8")) + if saved_path.exists() + else None + ) + return options, source, saved_options + + def test_default_tracks_parameters_with_subclass(self) -> None: + options, source, saved = self.collect(["--models-to-track", "model"]) + self.assertEqual(options["model_tracker_style"], "subclass") + self.assertEqual(options["models_to_track"], ["model"]) + self.assertIn("proxy_parameter(model,", source) + self.assertNotIn("model = Proxy(model,", source) + self.assertIsNone(saved) + + def test_explicit_cli_modes_remain_available(self) -> None: + statements = { + "proxy": "model = Proxy(model,", + "subclass": "proxy_parameter(model,", + "sampler": "model_sampler = VarSampler(model,", + } + for style, statement in statements.items(): + with self.subTest(style=style): + options, source, _ = self.collect( + ["--models-to-track", "model", "--model-tracker-style", style] + ) + self.assertEqual(options["model_tracker_style"], style) + self.assertIn(statement, source) + + def test_explicit_proxy_yaml_preserves_legacy_selection(self) -> None: + options, source, _ = self.collect( + [], yaml_options="models_to_track: [model]\nmodel_tracker_style: proxy\n" + ) + self.assertEqual(options["model_tracker_style"], "proxy") + self.assertIn("model = Proxy(model,", source) + + def test_yaml_without_style_uses_subclass(self) -> None: + options, source, _ = self.collect([], yaml_options="models_to_track: [model]\n") + self.assertEqual(options["model_tracker_style"], "subclass") + self.assertIn("proxy_parameter(model,", source) + + def test_default_does_not_activate_unselected_models(self) -> None: + options, source, _ = self.collect([]) + self.assertIsNone(options["models_to_track"]) + self.assertNotIn("proxy_parameter(model,", source) + self.assertNotIn("model = Proxy(model,", source) + self.assertNotIn("model_sampler = VarSampler(model,", source) + + def test_selective_parameter_rules_preserve_selected_tracker(self) -> None: + for explicit_style in (None, "proxy", "subclass"): + with self.subTest(explicit_style=explicit_style): + arguments = ["--models-to-track", "model"] + if explicit_style is not None: + arguments.extend(["--model-tracker-style", explicit_style]) + options, _, saved = self.collect( + arguments, invariants=[parameter_rule()] + ) + expected = "subclass" if explicit_style is None else explicit_style + self.assertEqual(options["model_tracker_style"], expected) + assert saved is not None + self.assertEqual(saved.model_tracker_style, expected) + self.assertTrue(saved.disable_proxy_dumping) + self.assertEqual( + saved.funcs_instr_opts["torch.optim.SGD.step"][ + "var_types_to_track" + ], + {"torch.nn.Parameter": {"dump_unchanged": False}}, + ) + + def test_sampler_cannot_replace_immediate_variable_tracking(self) -> None: + self.assertEqual( + collect_trace.get_model_tracker_instr_opts([parameter_rule()], "sampler"), + "subclass", + ) + + def test_consistency_only_rules_keep_sampler_policy(self) -> None: + invariant = Invariant( + ConsistencyRelation, + [ + VarTypeParam("torch.nn.Parameter", "data"), + VarTypeParam("torch.nn.Parameter", "grad"), + ], + None, + ) + for style in ("subclass", "proxy", "sampler"): + with self.subTest(style=style): + self.assertEqual( + collect_trace.get_model_tracker_instr_opts([invariant], style), + "sampler", + ) + + def test_selective_api_only_rules_disable_model_tracking(self) -> None: + invariant = Invariant( + APIContainRelation, + [APIParam("torch.optim.SGD.step"), APIParam("torch.add")], + GroupedPreconditions( + {"parent_func_call_pre": Preconditions([UnconditionalPrecondition()])} + ), + ) + options, source, saved = self.collect( + ["--models-to-track", "model"], invariants=[invariant] + ) + self.assertIsNone(options["model_tracker_style"]) + self.assertIsNone(options["models_to_track"]) + self.assertNotIn("proxy_parameter(model,", source) + assert saved is not None + self.assertIsNone(saved.model_tracker_style) + + def test_saved_instrumentation_styles_round_trip_without_migration(self) -> None: + for style in ("proxy", "subclass", "sampler", None): + with self.subTest(style=style): + original = InstrOpt({}, style, True) + restored = InstrOpt.from_json(original.to_json()) + self.assertEqual(restored.model_tracker_style, style) + self.assertEqual(restored.to_json(), original.to_json()) + + def test_saved_instrumentation_style_remains_required(self) -> None: + incomplete = {"funcs_instr_opts": {}, "disable_proxy_dumping": True} + with self.assertRaises(KeyError): + InstrOpt.from_json(json.dumps(incomplete)) + + def test_subclass_parameter_rule_activates_and_checks_real_transition(self) -> None: + records: list[dict[str, object]] = [] + + def capture(record: dict[str, object]) -> None: + # Retain production timestamping, metadata conversion, and JSON + # serialization; substitute only the final file/buffer destination. + payload: object = json.loads(dumper.serialize(record)) + assert isinstance(payload, dict) + validated: dict[str, object] = {} + for key, value in payload.items(): + assert isinstance(key, str) + validated[key] = value + records.append(validated) + + destination = SimpleNamespace(add_trace=capture) + with ( + patch.object(subclass, "get_global_registry", return_value=VarRegistry()), + patch.object( + dumper, + "THREAD_DATA", + SimpleNamespace( + api_trace_buffer=destination, var_trace_buffer=destination + ), + ), + patch.object(config, "DISABLE_WRAPPER", False), + patch.object(config, "INSTR_OPTS", None), + ): + model = torch.nn.Module() + model.register_parameter("weight", torch.nn.Parameter(torch.ones(2))) + subclass.proxy_parameter(model, parent_name="model") + wrapped = tracer.wrapper( + freeze_parameter, + is_bound_method=False, + scan_proxy_in_args=True, + dump_stack_trace=False, + dump_args=False, + dump_ret=False, + handle_proxy=False, + ) + wrapped(model.weight) + + for expected_post, should_pass in ((False, True), (True, False)): + with self.subTest(expected_post=expected_post): + invariant = Invariant( + APIContainRelation, + [ + APIParam(typename(freeze_parameter)), + VarTypeParam( + "torch.nn.Parameter", + "requires_grad", + pre_value=True, + post_value=expected_post, + ), + ], + GroupedPreconditions( + { + "parent_func_call_pre": Preconditions( + [UnconditionalPrecondition()] + ) + } + ), + "freezing a parameter changes requires_grad", + ) + cache_events_scanner.clear() + self.addCleanup(cache_events_scanner.clear) + report = observe( + TracePandas( + pd.json_normalize(records), truncate_incomplete_func_calls=False + ), + invariant, + ) + self.assertEqual(report["status"], "completed", report) + self.assertTrue(report["active"]) + self.assertFalse(report["dormant"]) + self.assertTrue(report["fully_validated_active"]) + self.assertEqual(report["decisions_completed"], 1) + self.assertEqual(report["check_passed"], should_pass) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/collect_trace.py b/traincheck/collect_trace.py index 2d307e49..cac7645b 100644 --- a/traincheck/collect_trace.py +++ b/traincheck/collect_trace.py @@ -360,7 +360,8 @@ def main(): "--model-tracker-style", type=str, choices=["sampler", "proxy", "subclass"], - default="proxy", + default="subclass", + help="Model parameter tracker: subclass (default), proxy, or sampler. Requires --models-to-track. Keep proxy explicit when checking invariants collected with proxy tracking.", ) parser.add_argument( "--tensor-dump-format",