diff --git a/docs/measurements/checking-overhead-20261027.md b/docs/measurements/checking-overhead-20261027.md new file mode 100644 index 00000000..f02dfd05 --- /dev/null +++ b/docs/measurements/checking-overhead-20261027.md @@ -0,0 +1,132 @@ +# Tiny-Llama selective API instrumentation, 2026-10-27 + +## Result and scope + +On this 6,163,712-parameter TorchTitan Llama workload, selective API instrumentation costs **22.64x the plain steady-step time** (median paired slowdown across three fresh pairs, range 22.62–22.67x). That is **+2164% overhead relative to plain training**. The median of per-run medians is 43.861 ms plain and 994.425 ms selective. This is a deliberately tiny model, batch 1, sequence 256, with every step instrumented; it is not a claim about larger training jobs or interval-10 sampling. + +All three pairs and a different-seed, 20-step holdout have matching SHA256 digests of initial/final parameter, gradient and buffer bytes, matching DTensor metadata, losses, and Python/NumPy/Torch CPU/CUDA RNG-state evidence. The holdout has 22.77x steady-step slowdown. + +This measures **API-only collection for checking**, followed by separate offline validation. The policy selects no model tracker. Native TorchTitan FSDP remains enabled. FSDP model tracking and a concurrently running live checker are not validated by these results and remain separate work. + +## Timing + +Native `train_step` timings include `start_step` and the final CUDA synchronization. The steady statistic uses steps 3 onward. Configuration, model construction, evidence capture, outer train-loop housekeeping, final queue draining, and process shutdown are outside the step timer. The experiment sets sampling_interval=1 and warm_up_steps=0. + +| Scope | Plain | Selective API | +|---|---:|---:| +| Pair 1 median steps 3–10 |43.969 ms|995.449 ms| +| Pair 2 median steps 3–10 |43.858 ms|994.425 ms| +| Pair 3 median steps 3–10 |43.861 ms|992.226 ms| +| Holdout median steps 3–20 |43.801 ms|997.229 ms| +| Pair 1 total 10 training steps, including cold first step |1.824 s|12.531 s| +| Pair 2 total 10 training steps |0.845 s|12.520 s| +| Pair 3 total 10 training steps |0.846 s|12.508 s| +| Holdout total 20 training steps |1.288 s|22.524 s| + +The first plain pair has cold initialization cost, so total-time ratios vary much more than the steady-step ratios. These totals must not be substituted for the steady-step statistic. + +A separate earlier full-API reference pair measured 44.570 ms plain and 5,079.953 ms full tracing per steady step. Selective's 994.425 ms is 80.42% lower than that full-trace time, but this is a cross-job policy comparison with a single full-trace reference, not a repeated optimization comparison. + +## Invariant validation + +The frozen 722-rule set was inferred from the full API reference: 419 APIContain, 156 ConsistentOutput, 142 ConsistentInputOutput, 5 DistinctArgument. One-worker inference took 58.654 s operational process time and 1.89 GB peak process-group RSS; this includes imports and parsing. + +Both the fresh same-seed trace and the different-seed holdout report: + +- 722 completed rules; 719 active (718 pass and 1 violates), 3 dormant. +- 1 violated rule, 0 checker/observer errors, 0 unknown outcomes, 0 skipped variable checks (none requested by this API-only policy). +- Exact complete CheckerResult parity with the direct checker for 722/722 rules. IDs, array contents and order are preserved; identical nonstandard NaN literals are compared only to the same literal at the same path, remaining distinct from null/missing. +- 114,373 actual decisions in the 10-step trace; 218,713 in the 20-step holdout. + +The remaining violation is rule 614, which learned a literal process-specific CUDA stream handle. It is retained unchanged and is a portability false positive, not evidence of divergent training. Dormant rules 216 and 276 require a reference worker thread ID absent from the fresh run; 717 was already dormant in the reference. Context-dependent rules 661, 719 and 720 now activate. + +All four selective traces passed independent raw/parsed-row and complete-call-pair audits, with no missing or duplicate scoped PRE/POST pairs. Matching report input hashes are checked against the actual audited files and the frozen invariant set. The full-reference parser reported skipped _CachedTorchDispatchMode contexts; those warnings were absent from these final selective audits. Row-count parity alone does not certify every possible context relationship. + +The direct offline checker took 29.980 s for 10 steps (1.26 GB peak process-group RSS) and 54.520 s for 20 steps (1.84 GB). The coverage-observed runs took 53.364 s and 105.459 s respectively; these include observer overhead and are not ordinary checking timings. None of these offline durations is added to or presented as concurrent training overhead. + +## What was fixed before using these timings + +The first selective harness bypassed the production JSON-loading boundary, so integer argument-selector keys were not recognized by the dumper. Its approximately 14.1x timing is invalid as a complete checking result. + +After fixing the harness, production policy generation still produced 61 violated rules: 54 missing-return cases, 1 incomplete distinct-argument case, 5 positional-only/variadic selector cases, and the real stream-handle portability issue. Context constructor inputs also caused avoidable dormancy. + +PR #161 repairs the required evidence by unioning return, complete-argument and context requirements, and requesting full argument metadata when the selector cannot represent a required argument. Supported selectors stay selective. It preserves the 722 rules and their predicates. The final trace reduces 61 violations to 1 and raises active rules from 716 to 719; no rule or identity field was removed to obtain that result. This correctness repair costs more than the incomplete policy, so no speedup is claimed for it. + +## Trace volume and actionable follow-up + +| Representative 10-step trace | Full API | Selective API | +|---|---:|---:| +| Observed API names |367|244| +| API records |442,918|416,364| +| Complete calls |221,459|208,182| +| JSONL bytes, including newlines |255,365,856|143,453,745| + +Selective reduces records by 5.995% and bytes by 43.824%. Each steady step still emits 29,864 records for 14,932 calls and about 10.20 MB. Metadata-only records form 74.36% of events and 66.41% of bytes. These are counts and sizes, not runtime percentages. + +The six largest selective API byte contributors are registry lookup, export-state lookup, dispatch-rule lookup, checkpoint assertions, dispatch-rule search, and checkpoint storage-key lookup. Many lightweight framework calls remain selected by the frozen invariants. Dropping those records would change the evidence required by those rules; record reduction needs semantic/coverage review, not a silent filter. + +File writing already runs on a background thread. Metadata conversion and JSON encoding happen on the producer before enqueue. Final queue drain alone cannot establish the cost of overlapping writes. + +## Component profile: main-thread diagnostic + +H100 job 55325428 uses the same frozen training source, policy, invariants and mode-specific ordinary controls. A targeted `sys.setprofile` collector measures Python function boundaries on the main thread during steps 3–10. It includes untracked descendants within a boundary, subtracts tracked child elapsed exactly once, and leaves application execution and unclassified wrapper work in an explicit residual. Final serialization includes its stdlib fallback; the earlier JSON compatibility probe has its own bucket. Function totals reconcile to categories, and categories reconcile to the full captured window. + +**These are diagnostic shares of a separately profiled window, not shares of the ordinary 994 ms step or predicted optimization savings.** The profiler expands the selective eight-step window from 7.961 s to 42.005 s (5.28x), and the plain window from 0.354 s to 0.712 s (2.01x). Callback overhead is included in surrounding categories or the residual; it is not independently estimated. This strong, unequal perturbation makes the profile useful for locating work but unsuitable for calculating an ordinary-run speedup from its percentages. + +| Component, main thread only | Profiled elapsed over 8 steps | Share of profiled window | +|---|---:|---:| +| Metadata conversion, excluding separately timed warning/probe/helpers |22.327 s|53.15%| +| Warning management |7.852 s|18.69%| +| Known wrapper/bookkeeping helpers |1.537 s|3.66%| +| Producer buffering/enqueue |1.292 s|3.08%| +| Final record JSON encoding, including fallback |0.657 s|1.56%| +| JSON compatibility probes |0.300 s|0.71%| +| Application plus unclassified wrapper work and measurement residual |8.040 s|19.14%| +| Total |42.006 s|100%| + +Metadata conversion means turning arguments/returns into recorded fields: walking attributes, reading their values, classifying types, and constructing dictionaries. In this window, `convert_var_to_dict` runs 25,312 times and `safe_getattr` runs 1,243,312 times: **155,414 attribute reads per step**. Each read constructs a warning context, enters it, sets an ignore filter and exits it. Each of those four warning operations also runs 1,243,312 times. The warning category measures these boundaries with their internal work; it excludes the separately reported getter/metadata work. + +The main thread emits 133,440 records in this window, exactly matching its serializer/buffer call counts. The other trace-producing worker emits another 105,472 records, 44.1% of all 238,912 records. **Worker instrumentation and writer-thread CPU time are not profiled.** Main-thread elapsed time may include waiting for them; this profile cannot assign that waiting to their internal components. Both threads' complete step/API/event counts and byte totals match the ordinary run. State-byte digests, DTensor metadata, losses and RNG evidence also match the plain/selective controls. + +### Concrete optimization targets + +1. **Reduce repeated metadata attribute work and warning-context setup.** This is the first implementation target suggested by the measured work. Keep requested fields, dynamic attributes/properties, exceptions and application warning behavior intact; do not simply remove warning suppression or cache mutable snapshots. Any fast path needs matched ordinary runs and output/behavior tests before claiming a speedup. +2. **Avoid predictable JSON fallback exceptions.** `to_dict_args_kwargs` constructs integer positional keys, while the first `orjson.dumps` attempt rejects those keys. In the first diagnostic window, the count of 55,264 nonempty-positional-argument records exactly matches the 55,264 final-serialization stdlib fallback calls; source inspection and a minimal reproduction support the integer-key mechanism. No per-record fallback correlation was recorded. The corrected main-thread window contains 23,792 such records; 31,472 belong to the worker. A safe initial approach can route known integer-key records directly to the existing fallback. Blindly enabling non-string keys or stringifying them to select orjson can change NaN/Infinity output from the existing stdlib representation to null, so those cases, large integers, unsupported keys and defaults require explicit parity/error tests. +3. **Avoid unnecessary serialize-to-test work for known supported primitives.** Main-thread `var_to_serializable` makes 34,104 `json.dumps` compatibility probes during the eight steps. Exact built-in-type fast paths could avoid the eligible primitive subset of those probes while retaining the existing fallback for dynamic/custom objects; that subset was not counted separately. Final record encoding and these probes together occupy 2.28% of this diagnostic window; moving encoding asynchronously would still leave metadata capture and warning handling on the producer. An async design must snapshot immutable evidence before training can mutate it. + +No production optimization was applied in this profiling follow-up, and no improvement percentage is claimed. The ordinary 22.64x slowdown remains the benchmark result. + +### Profiler validation and retained failed attempt + +The allocated Python 3.12.3 runtime self-test executed 20 background markers and profiled 0. Nested accounting, application-exception preservation, lost-hook rejection, strict typing and pinned formatting/lint checks also pass; local tests cover Python 3.10 and 3.12. Raw trace/state validation for the final job passes with the same 416,364 records, 208,182 complete calls and 143,453,745 bytes as the ordinary 10-step selective run. This is record-integrity and path-volume comparability, not equality after removing changing IDs or timestamps. + +The earlier cProfile job 55325378 is retained as rejected attribution evidence. CPython 3.12 captured background-thread callbacks despite its main-thread intent and produced inconsistent caller accounting; a separate marker reproduction captured 20/20 background calls. Its training parity and trace integrity remain valid, but its component percentages are not used. The corrected collector and denominator above replace them. + +Reproduce the final evidence checks from the downloaded study directory: + +```sh +python validate_profile_evidence.py --profile-job job-55325428 +python analyze_final_component_profile.py --job-dir job-55325428 \ + --controls-root profile-controls --output final-component-profile-analysis.json +``` + +The saved `component-profile.sbatch`, `run_profile_phases.py`, `profile_components.py` and its adjacent `profile_selective.py` loader/parity helper reproduce the allocated experiment. They preserve the frozen training archive and exact ordinary-control files. The corrected results archive SHA256 is `a5352de5b6bf7954b5e2400241361d18ace83f6022c8d09d9c24ff14c3fbb4ae`. Detailed ledger, function counts, owner-thread audit and provenance are in `final-component-profile-analysis.json`, `job-55325428/profile-validation.json`, and the complete downloaded job directory. + +## Reproduction and provenance + +- Baseline main: 6dccb47d0a8bc192fbe7cbbea700b66106bc2381. +- Full-reference composed source: b34e309; job 55324690. +- Corrected selective composed source: c2233f119e0eb0f01827a7ebb593b259f82596f2; job 55325270. Later 8cc73ab changes only a test-discovery fixture. +- TorchTitan: 2210414902d84d1219fd63ab8efeb24375a2a1a1. +- One H100 80 GB on Expanse, Python 3.12.3, Torch 2.16.0.dev20261012+cu130, CUDA 13.0, CUDA hashing enabled through PR #158. CPU validation uses labmach Python 3.11.4/Torch 2.16.0.dev20261012+cpu, one worker and single native/intraop/interop thread, bounded to 16 GiB and phase deadlines. +- Seed 42 paired order: plain/selective, selective/plain, plain/selective; 10 steps. Holdout seed 43: plain/selective; 20 steps. No parameter updates or setup math were replaced. +- Invariant SHA256: bad55c222f4eb3406bdfaec527ab1a75eaf66cdbe8f06c42d7f0775f4ee9224b. +- Selective source archive SHA256: 33363cd5dac62e800545b55924ad49d887f1634c5a588cba04a05c16255270bd. +- Final results archive SHA256: f8f7f9077ae322bcb9c766b57496c94fc2eca3d3e551f00d0039345a8976bdff. +- Actual policy SHA256, identical in all four selective runs: df4c91296774b2dfd3af09d0a25e21e25f26ff1de93cb97ab959fa77379c3211. + +Durable downloaded artifacts are in the experiment owner's `outputs/checking-overhead-20261027/`, including source/runtime manifests, exact scripts, results, raw profiles when available, `job-55325270/paired-summary.json`, `final-policy-validation-summary.json`, `trace-volume-final-policy.json`, and `offline/results/final-policy-*`. Large traces stay outside Git. Expanse copies are under `/expanse/lustre/scratch/yjiang21/temp_project/traincheck-checking-20261027/job-JOBID`; this scratch is retention-limited. Lab artifacts are in `/home/yuxuan/experiments/checking-overhead-20261027`. + +Reproduction uses saved `selective.sbatch`/`run_phases.py` on an allocated GPU, then `offline-final-policy.sh` and `audit-final-policy-repeats.sh` on labmach. Locally run `python analyze_final_trials.py` and `python trace-volume-final-policy.py`. The hash-checked coverage summary uses `analyze_reference_validation.py` against the remote study root with `--fresh-prefix final-policy --policy-status final-production-candidate`, because audit records contain remote absolute paths. `python analyze_final_policy_attribution.py` uses the downloaded mirror and local final traces to append all four audits and exact stream-handle differences. The saved source archives pin the actual measured code rather than a future branch tip. + +PR #159 contains the workload harness; PR #158 supplies the H100 hashing correctness fix; PR #160 supplies actual coverage observation; PR #161 supplies the production policy repairs. Each targets main independently. The GPU experiment uses their composed code. No PR was merged by this study. diff --git a/tests/test_checking_workload.py b/tests/test_checking_workload.py new file mode 100644 index 00000000..97c3bea6 --- /dev/null +++ b/tests/test_checking_workload.py @@ -0,0 +1,504 @@ +"""Behavioral contracts for the TorchTitan checking-overhead adapter boundary.""" + +from __future__ import annotations + +import json +import random +import tempfile +import unittest +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import torch + +from tools.checking_overhead import workload +from traincheck.config import config +from traincheck.config.config import InstrOpt +from traincheck.instrumentor.caches import META_VARS +from traincheck.instrumentor.proxy_wrapper import subclass +from traincheck.instrumentor.proxy_wrapper.proxy_registry import VarRegistry +from traincheck.instrumentor.tracer import to_dict_args_kwargs +from traincheck.invariant.base_cls import ( + APIParam, + Arguments, + GroupedPreconditions, + Invariant, + Preconditions, + UnconditionalPrecondition, + VarTypeParam, +) +from traincheck.invariant.consistency_relation import ConsistencyRelation +from traincheck.invariant.contain_relation import APIContainRelation + + +def positional_api(flag: bool, *, label: str) -> bool: + return flag + + +@dataclass +class Optimizers: + optimizers: list[torch.optim.Optimizer] + + +class Metrics: + def __init__(self) -> None: + self.calls: list[int] = [] + + def log( + self, + step: int, + average: float, + maximum: float, + gradient_norm: float, + *, + extra_metrics: dict[str, int], + ) -> None: + self.calls.append(step) + + +class Trainer: + def __init__(self, steps: int = 3) -> None: + self.model_parts = [torch.nn.Linear(2, 1)] + self.optimizers = Optimizers( + [torch.optim.SGD(self.model_parts[0].parameters(), lr=0.01)] + ) + self.metrics_processor = Metrics() + self.step = 0 + self.steps = steps + self.closed = False + self.failure: Exception | None = None + self.disabled_at_step: list[bool] = [] + + def train_step(self, data: torch.Tensor) -> None: + self.disabled_at_step.append(config.DISABLE_WRAPPER) + if self.failure is not None: + raise self.failure + optimizer = self.optimizers.optimizers[0] + optimizer.zero_grad() + loss = self.model_parts[0](data).square().mean() + loss.backward() + norm = torch.nn.utils.clip_grad_norm_(self.model_parts[0].parameters(), 1.0) + optimizer.step() + value = float(loss.item()) + self.metrics_processor.log( + self.step, + value, + value, + float(norm.item()), + extra_metrics={"n_tokens_seen": self.step}, + ) + + def train(self) -> None: + while self.step < self.steps: + self.step += 1 + self.train_step(torch.randn(3, 2)) + + def close(self) -> None: + self.closed = True + + +class CustomParameter(torch.nn.Parameter): + """The production subclass tracker cannot preserve this type's identity.""" + + +class WorkloadTests(unittest.TestCase): + def setUp(self) -> None: + self.registry = VarRegistry() + for target, name, value in ( + (config, "DISABLE_WRAPPER", True), + (config, "INSTRUMENTATION_POLICY", {"interval": 1, "warm_up": 0}), + (subclass, "get_global_registry", lambda: self.registry), + ): + patcher = patch.object(target, name, value) + patcher.start() + self.addCleanup(patcher.stop) + meta_patch = patch.dict(META_VARS, {"step": 0}, clear=True) + meta_patch.start() + self.addCleanup(meta_patch.stop) + + def seeded_trainer(self) -> Trainer: + torch.manual_seed(42) + random.seed(42) + np.random.seed(42) + return Trainer() + + def test_subclass_preserves_state_rng_and_existing_optimizer_references( + self, + ) -> None: + trainer = self.seeded_trainer() + handle = workload.TrainerHandle.from_trainer(trainer) + parameters = tuple(handle.model.parameters()) + initial = workload.model_evidence(handle.model) + rng = workload.rng_evidence() + workload.attach_model_tracker(handle, "subclass") + self.assertEqual(initial, workload.model_evidence(handle.model)) + self.assertEqual(rng, workload.rng_evidence()) + self.assertTrue( + all(a is b for a, b in zip(parameters, handle.model.parameters())) + ) + self.assertTrue( + all( + a is b + for a, b in zip( + parameters, workload.optimizer_parameters(handle.optimizers) + ) + ) + ) + + def test_tracked_training_matches_plain_bytes_losses_and_rng(self) -> None: + plain = self.seeded_trainer() + plain_handle = workload.TrainerHandle.from_trainer(plain) + _, plain_losses = workload.run_training(plain_handle, 3, False, lambda: None) + expected = workload.model_evidence(plain_handle.model) + expected_rng = workload.rng_evidence() + tracked = self.seeded_trainer() + tracked_handle = workload.TrainerHandle.from_trainer(tracked) + workload.attach_model_tracker(tracked_handle, "subclass") + _, tracked_losses = workload.run_training( + tracked_handle, 3, False, lambda: None + ) + self.assertEqual(expected, workload.model_evidence(tracked_handle.model)) + self.assertEqual(expected_rng, workload.rng_evidence()) + self.assertEqual(plain_losses, tracked_losses) + + def test_step_boundary_once_and_sync_outside_tracing(self) -> None: + trainer = self.seeded_trainer() + handle = workload.TrainerHandle.from_trainer(trainer) + sync_states: list[bool] = [] + + def synchronize() -> None: + sync_states.append(config.DISABLE_WRAPPER) + + timings, losses = workload.run_training(handle, 3, True, synchronize) + self.assertEqual(META_VARS["step"], 3) + self.assertEqual([row["step"] for row in timings], [1, 2, 3]) + self.assertEqual([row["step"] for row in losses], [1, 2, 3]) + self.assertEqual(trainer.disabled_at_step, [False, False, False]) + self.assertEqual(sync_states, [True] * 6) + self.assertTrue(config.DISABLE_WRAPPER) + self.assertEqual(trainer.train_step, handle.train_step) + self.assertEqual(trainer.metrics_processor.log, handle.metrics_log) + + def test_failure_propagates_and_restores_boundaries(self) -> None: + trainer = self.seeded_trainer() + failure = RuntimeError("native step failure") + trainer.failure = failure + handle = workload.TrainerHandle.from_trainer(trainer) + with self.assertRaises(RuntimeError) as raised: + workload.run_training(handle, 3, True, lambda: None) + self.assertIs(raised.exception, failure) + self.assertTrue(config.DISABLE_WRAPPER) + self.assertEqual(trainer.train_step, handle.train_step) + self.assertEqual(trainer.metrics_processor.log, handle.metrics_log) + + def test_custom_parameter_rejected_before_mutation(self) -> None: + trainer = self.seeded_trainer() + model = trainer.model_parts[0] + model.weight = CustomParameter(model.weight.detach()) + trainer.optimizers = Optimizers([torch.optim.SGD(model.parameters(), lr=0.1)]) + before = tuple(model.parameters()) + with self.assertRaisesRegex(TypeError, "exact nn.Parameter"): + workload.attach_model_tracker( + workload.TrainerHandle.from_trainer(trainer), "subclass" + ) + self.assertTrue(all(a is b for a, b in zip(before, model.parameters()))) + + def test_malformed_optimizer_boundary_is_rejected(self) -> None: + trainer = self.seeded_trainer() + trainer.optimizers = Optimizers([]) + with self.assertRaisesRegex(TypeError, "nonempty list"): + workload.TrainerHandle.from_trainer(trainer) + + def test_optimizer_outside_model_is_rejected(self) -> None: + trainer = self.seeded_trainer() + trainer.optimizers = Optimizers( + [torch.optim.SGD([torch.nn.Parameter(torch.ones(1))], lr=0.1)] + ) + with self.assertRaisesRegex(ValueError, "outside"): + workload.attach_model_tracker( + workload.TrainerHandle.from_trainer(trainer), "subclass" + ) + + def test_exact_scalar_empty_and_dtype_evidence(self) -> None: + scalar = workload.tensor_evidence(torch.tensor(2.0)) + vector = workload.tensor_evidence(torch.tensor([2.0])) + self.assertEqual(scalar["sha256"], vector["sha256"]) + self.assertNotEqual(scalar["shape"], vector["shape"]) + self.assertEqual(workload.tensor_evidence(torch.empty(0))["shape"], [0]) + self.assertNotEqual( + scalar, workload.tensor_evidence(torch.tensor(2.0, dtype=torch.float64)) + ) + + def test_options_reject_missing_or_extraneous_invariants(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + with self.assertRaisesRegex(ValueError, "required exactly"): + workload.Options("selective", root, root) + with self.assertRaisesRegex(ValueError, "required exactly"): + workload.Options( + "plain", root, root, invariants=root / "invariants.json" + ) + with self.assertRaisesRegex(ValueError, "positive integer"): + workload.Options("plain", root, root, steps=0) + with self.assertRaisesRegex(ValueError, "positive integer"): + workload.Options("plain", root, root, steps=True) + + def test_selective_consistency_rules_preserve_sampler_policy(self) -> None: + invariant = Invariant( + ConsistencyRelation, + [ + VarTypeParam("torch.nn.Parameter", "data"), + VarTypeParam("torch.nn.Parameter", "grad"), + ], + None, + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "traces").mkdir() + options = workload.Options( + "selective", root, root, invariants=root / "invariants.json" + ) + with ( + patch( + "traincheck.invariant.base_cls.read_inv_file", + return_value=[invariant], + ), + patch("traincheck.instrumentor.tracer.Instrumentor") as instrumentor, + patch.object(workload, "install_standard_observers") as observers, + patch.object(config, "INSTR_OPTS", None), + patch.object(config, "INSTR_OPTS_FILE", config.INSTR_OPTS_FILE), + patch.object(config, "MODEL_TRACKER_STYLE", None), + ): + self.assertEqual(workload.install_instrumentation(options), "sampler") + self.assertEqual(config.MODEL_TRACKER_STYLE, "sampler") + self.assertIsNotNone(config.INSTR_OPTS) + instrumentor.return_value.instrument.assert_called_once() + observers.assert_not_called() + self.assertIn( + '"model_tracker_style": "sampler"', + (root / "traces" / "instr_opts.json").read_text(), + ) + + def test_selective_loads_saved_positional_selectors_before_instrumentation( + self, + ) -> None: + function = f"{__name__}.positional_api" + arguments = Arguments( + {"0": {"bool": True}}, {"label": {"str": "example"}}, function + ) + invariant = Invariant( + APIContainRelation, + [APIParam(function, arguments=arguments), APIParam(function)], + GroupedPreconditions( + {"parent_func_call_pre": Preconditions([UnconditionalPrecondition()])} + ), + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "traces").mkdir() + options = workload.Options( + "selective", root, root, invariants=root / "invariants.json" + ) + with ( + patch( + "traincheck.invariant.base_cls.read_inv_file", + return_value=[invariant], + ), + patch("traincheck.instrumentor.tracer.Instrumentor"), + patch.object(config, "INSTR_OPTS", None), + patch.object(config, "INSTR_OPTS_FILE", config.INSTR_OPTS_FILE), + patch.object(config, "MODEL_TRACKER_STYLE", None), + ): + workload.install_instrumentation(options) + loaded: object = config.load_instr_opts() + self.assertIsInstance(loaded, InstrOpt) + assert isinstance(loaded, InstrOpt) + saved: object = json.loads( + (root / "traces" / "instr_opts.json").read_text() + ) + self.assertIsInstance(saved, dict) + assert isinstance(saved, dict) + self.assertEqual(loaded.funcs_instr_opts, saved["funcs_instr_opts"]) + self.assertEqual(set(loaded.funcs_instr_opts), {function}) + selector = loaded.funcs_instr_opts[function]["dump_args_config"] + assert isinstance(selector, dict) + self.assertEqual(set(selector), {"flag", "0", "label"}) + self.assertEqual( + to_dict_args_kwargs((True,), {"label": "example"}, selector), + { + "args": {0: {"bool": True}}, + "kwargs": {"label": {"str": "example"}}, + }, + ) + self.assertFalse(loaded.funcs_instr_opts[function]["dump_ret"]) + + def test_selective_rejects_empty_invariants_before_instrumentation(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + options = workload.Options( + "selective", root, root, invariants=root / "invariants.json" + ) + with ( + patch("traincheck.invariant.base_cls.read_inv_file", return_value=[]), + patch("traincheck.instrumentor.tracer.Instrumentor") as instrumentor, + patch.object(config, "INSTR_OPTS", None), + patch.object(config, "INSTR_OPTS_FILE", config.INSTR_OPTS_FILE), + ): + with self.assertRaisesRegex(ValueError, "nonempty invariants"): + workload.install_instrumentation(options) + instrumentor.assert_not_called() + + def test_api_reference_has_no_model_tracker(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + options = workload.Options("reference_api", root, root) + with ( + patch("traincheck.instrumentor.tracer.Instrumentor") as instrumentor, + patch.object(workload, "install_standard_observers") as observers, + patch.object(config, "INSTR_OPTS", None), + patch.object(config, "INSTR_OPTS_FILE", config.INSTR_OPTS_FILE), + patch.object(config, "MODEL_TRACKER_STYLE", "subclass"), + ): + self.assertIsNone(workload.install_instrumentation(options)) + self.assertIsNone(config.MODEL_TRACKER_STYLE) + self.assertIsNone(config.INSTR_OPTS) + self.assertFalse(instrumentor.call_args.kwargs["scan_proxy_in_args"]) + instrumentor.return_value.instrument.assert_called_once() + observers.assert_not_called() + + def test_dtensor_api_absence_preserves_plain_evidence(self) -> None: + native = torch.tensor([1.0, 2.0]) + expected = workload.tensor_evidence(native) + + def absent(namespace: str) -> object: + raise ModuleNotFoundError(name=namespace) + + with patch.object(workload.importlib, "import_module", side_effect=absent): + self.assertIsNone(workload.dtensor_api()) + self.assertEqual(expected, workload.tensor_evidence(native)) + + def test_dtensor_api_does_not_hide_broken_import_or_exports(self) -> None: + failure = ModuleNotFoundError(name="missing_dependency") + with patch.object(workload.importlib, "import_module", side_effect=failure): + with self.assertRaises(ModuleNotFoundError) as raised: + workload.dtensor_api() + self.assertIs(raised.exception, failure) + malformed = SimpleNamespace(DTensor=object, Shard=object, Replicate=object) + with patch.object(workload.importlib, "import_module", return_value=malformed): + with self.assertRaisesRegex(TypeError, "Tensor subclass"): + workload.dtensor_api() + + def test_real_one_rank_dtensor_evidence_preserves_identity_and_storage( + self, + ) -> None: + api = workload.dtensor_api() + if api is None or not torch.distributed.is_gloo_available(): + self.skipTest( + "CPU DTensor evidence requires the optional DTensor and Gloo APIs" + ) + assert api is not None + from torch.distributed.device_mesh import DeviceMesh + + with tempfile.TemporaryDirectory() as directory: + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{directory}/rendezvous", + rank=0, + world_size=1, + ) + try: + mesh = DeviceMesh("cpu", [0], mesh_dim_names=("shard",)) + for placement in ( + workload.required_callable(api, "shard_type")(0), + workload.required_callable(api, "replicate_type")(), + ): + with self.subTest(placement=type(placement)): + local = torch.arange(6, dtype=torch.float32).reshape(2, 3) + distributed = workload.required_callable( + api.tensor_type, "from_local" + )( + local, + device_mesh=mesh, + placements=[placement], + run_check=False, + ) + self.assertIsInstance(distributed, torch.Tensor) + assert isinstance(distributed, torch.Tensor) + identity = id(distributed) + original_local = workload.required_callable( + distributed, "to_local" + )() + assert isinstance(original_local, torch.Tensor) + storage_pointer = original_local.data_ptr() + expected = workload.tensor_evidence(local) + rng = workload.rng_evidence() + evidence = workload.tensor_evidence(distributed) + self.assertEqual(evidence["sha256"], expected["sha256"]) + self.assertEqual(evidence["shape"], [2, 3]) + self.assertEqual(evidence["local_shape"], [2, 3]) + self.assertEqual( + evidence["storage_scope"], "dtensor_local_shard" + ) + metadata = evidence["distributed"] + assert metadata is not None + self.assertEqual(metadata["mesh_ranks"], [0]) + self.assertEqual(metadata["mesh_dim_names"], ["shard"]) + self.assertEqual(metadata["world_size"], 1) + self.assertEqual(id(distributed), identity) + after_local = workload.required_callable( + distributed, "to_local" + )() + assert isinstance(after_local, torch.Tensor) + self.assertEqual(after_local.data_ptr(), storage_pointer) + self.assertEqual(rng, workload.rng_evidence()) + parameter = torch.nn.Parameter(distributed) + parameter.grad = torch.ones_like(distributed) + model = torch.nn.Module() + model.register_parameter("weight", parameter) + model_state = workload.model_evidence(model) + self.assertEqual( + model_state["parameters"]["weight"]["sha256"], + expected["sha256"], + ) + gradient = model_state["gradients"]["weight"] + assert gradient is not None + self.assertEqual( + gradient["storage_scope"], "dtensor_local_shard" + ) + self.assertEqual( + gradient["sha256"], + workload.tensor_evidence(torch.ones_like(local))["sha256"], + ) + self.assertIs(model.weight, parameter) + with patch.object( + torch.distributed, "get_world_size", return_value=2 + ): + with self.assertRaisesRegex(ValueError, "world_size=1"): + workload.tensor_evidence(distributed) + finally: + torch.distributed.destroy_process_group() + + def test_cli_resolves_explicit_inputs(self) -> None: + options = workload.parse_options( + [ + "--mode", + "selective", + "--output-dir", + "result", + "--torchtitan-root", + "titan", + "--invariants", + "invariants.json", + ] + ) + self.assertEqual(options.mode, "selective") + self.assertTrue(options.output_dir.is_absolute()) + self.assertEqual(options.steps, 10) + self.assertIsNotNone(options.invariants) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/checking_overhead/workload.py b/tools/checking_overhead/workload.py new file mode 100644 index 00000000..1a25764a --- /dev/null +++ b/tools/checking_overhead/workload.py @@ -0,0 +1,743 @@ +"""One-process, one-GPU TorchTitan checking-overhead pilot. + +Invoke the absolute script path with TrainCheck's source root on PYTHONPATH: + + PYTHONPATH=/path/to/TrainCheck torchrun --standalone --nproc-per-node=1 \ + /path/to/TrainCheck/tools/checking_overhead/workload.py [options] + +TorchTitan also has a tools package, so module lookup can select its package. +Each invocation owns a fresh interpreter and output directory; installing +instrumentation is intentionally process-global. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import json +import os +import random +import sys +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypedDict + +import numpy as np +import torch + +import traincheck +from traincheck.config import config as tc_config + +Mode = Literal["plain", "reference", "reference_api", "selective"] +TrackerStyle = Literal["subclass", "sampler"] | None +EXPECTED_PARAMETERS = 6_163_712 + + +@dataclass(frozen=True) +class Options: + mode: Mode + output_dir: Path + torchtitan_root: Path + steps: int = 10 + seq_len: int = 256 + batch_size: int = 1 + seed: int = 42 + invariants: Path | None = None + + def __post_init__(self) -> None: + if self.mode not in ("plain", "reference", "reference_api", "selective"): + raise ValueError( + "mode must be plain, reference, reference_api, or selective" + ) + for name in ("steps", "seq_len", "batch_size"): + value = getattr(self, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if ( + not isinstance(self.seed, int) + or isinstance(self.seed, bool) + or not 0 <= self.seed < 2**32 + ): + raise ValueError("seed must be an integer in [0, 2**32)") + if (self.mode == "selective") != (self.invariants is not None): + raise ValueError("--invariants is required exactly for selective mode") + + +class PlacementEvidence(TypedDict): + kind: Literal["Shard", "Replicate"] + dim: int | None + + +class DistributedEvidence(TypedDict): + api_namespace: str + world_size: int + mesh_device_type: str + mesh_shape: list[int] + mesh_ranks: list[int] + mesh_dim_names: list[str] | None + placements: list[PlacementEvidence] + + +class TensorEvidence(TypedDict): + dtype: str + shape: list[int] + sha256: str + storage_scope: Literal["full_tensor", "dtensor_local_shard"] + local_shape: list[int] + distributed: DistributedEvidence | None + + +class ModelEvidence(TypedDict): + parameters: dict[str, TensorEvidence] + gradients: dict[str, TensorEvidence | None] + buffers: dict[str, TensorEvidence] + + +class LossRecord(TypedDict): + step: int + average: float + maximum: float + gradient_norm: float + + +class StepRecord(TypedDict): + step: int + seconds: float + + +@dataclass(frozen=True) +class TrainerHandle: + """Validated boundary to the pinned TorchTitan trainer, imported lazily.""" + + trainer: object + model: torch.nn.Module + optimizers: tuple[torch.optim.Optimizer, ...] + metrics: object + train: Callable[..., object] + train_step: Callable[..., object] + metrics_log: Callable[..., object] + close: Callable[..., object] + + @classmethod + def from_trainer(cls, trainer: object) -> TrainerHandle: + parts: object = getattr(trainer, "model_parts") + if not isinstance(parts, list) or len(parts) != 1: + raise ValueError("This single-GPU pilot requires exactly one model part") + model: object = parts[0] + if not isinstance(model, torch.nn.Module): + raise TypeError("model_parts[0] must be a torch.nn.Module") + container: object = getattr(trainer, "optimizers") + optimizers: object = getattr(container, "optimizers") + if not isinstance(optimizers, list) or not optimizers: + raise TypeError("TorchTitan optimizers.optimizers must be a nonempty list") + checked: list[torch.optim.Optimizer] = [] + for optimizer in optimizers: + if not isinstance(optimizer, torch.optim.Optimizer): + raise TypeError("Every optimizer must be a torch.optim.Optimizer") + checked.append(optimizer) + metrics: object = getattr(trainer, "metrics_processor") + return cls( + trainer, + model, + tuple(checked), + metrics, + required_callable(trainer, "train"), + required_callable(trainer, "train_step"), + required_callable(metrics, "log"), + required_callable(trainer, "close"), + ) + + +def required_callable(owner: object, name: str) -> Callable[..., object]: + value: object = getattr(owner, name) + if not callable(value): + raise TypeError(f"{name} must be callable") + return value + + +@dataclass(frozen=True) +class DTensorAPI: + namespace: str + tensor_type: type[torch.Tensor] + shard_type: type[object] + replicate_type: type[object] + + +def dtensor_api() -> DTensorAPI | None: + """Support PyTorch's public DTensor API and its explicit 2.2 legacy namespace. + + An absent optional API leaves ordinary tensors supported. Broken imports or + malformed exports from an installed API propagate; they are never absence. + """ + for namespace in ("torch.distributed.tensor", "torch.distributed._tensor"): + try: + module = importlib.import_module(namespace) + except ModuleNotFoundError as error: + if error.name != namespace: + raise + continue + # PyTorch 2.2 has a tensor namespace for parallel APIs without a + # DTensor export; that explicit capability absence selects _tensor. + if "DTensor" not in module.__dict__: + continue + tensor_type: object = getattr(module, "DTensor") + shard_type: object = getattr(module, "Shard") + replicate_type: object = getattr(module, "Replicate") + if not isinstance(tensor_type, type) or not issubclass( + tensor_type, torch.Tensor + ): + raise TypeError(f"{namespace}.DTensor must be a Tensor subclass") + if not isinstance(shard_type, type) or not isinstance(replicate_type, type): + raise TypeError(f"{namespace} placement exports must be classes") + return DTensorAPI(namespace, tensor_type, shard_type, replicate_type) + return None + + +def distributed_evidence(tensor: torch.Tensor, api: DTensorAPI) -> DistributedEvidence: + from torch.distributed.device_mesh import DeviceMesh + + if not torch.distributed.is_initialized(): + raise RuntimeError("DTensor evidence requires an initialized process group") + world_size = torch.distributed.get_world_size() + if world_size != 1: + raise ValueError("This pilot supports DTensor evidence only for world_size=1") + mesh: object = getattr(tensor, "device_mesh") + if not isinstance(mesh, DeviceMesh): + raise TypeError("DTensor device_mesh must be a DeviceMesh") + ranks: object = mesh.mesh.flatten().tolist() + if not isinstance(ranks, list) or len(ranks) != 1 or ranks[0] != 0: + raise ValueError( + "This pilot requires a one-rank DTensor mesh containing rank zero" + ) + # The checked mesh is a one-rank integer tensor; do not coerce unknown values. + if not isinstance(ranks[0], int) or isinstance(ranks[0], bool): + raise TypeError("DeviceMesh ranks must be integers") + device_type: object = mesh.device_type + if not isinstance(device_type, str): + raise TypeError("DeviceMesh device_type must be a string") + names: object = mesh.mesh_dim_names + if names is not None and ( + not isinstance(names, tuple) or not all(isinstance(name, str) for name in names) + ): + raise TypeError("DeviceMesh dimension names must be a tuple of strings or None") + checked_names: list[str] | None = None + if isinstance(names, tuple): + checked_names = [] + for name in names: + if not isinstance(name, str): + raise TypeError("DeviceMesh dimension names must be strings") + checked_names.append(name) + placements: object = getattr(tensor, "placements") + if not isinstance(placements, tuple) or len(placements) != mesh.ndim: + raise TypeError("DTensor placements must match its mesh dimensions") + checked_placements: list[PlacementEvidence] = [] + for placement in placements: + if type(placement) is api.shard_type: + dim: object = getattr(placement, "dim") + if not isinstance(dim, int) or isinstance(dim, bool): + raise TypeError("Shard dimension must be an integer") + checked_placements.append({"kind": "Shard", "dim": dim}) + elif type(placement) is api.replicate_type: + checked_placements.append({"kind": "Replicate", "dim": None}) + else: + raise TypeError( + f"Unsupported DTensor placement for this pilot: {type(placement)}" + ) + return { + "api_namespace": api.namespace, + "world_size": world_size, + "mesh_device_type": device_type, + "mesh_shape": list(mesh.mesh.shape), + "mesh_ranks": [ranks[0]], + "mesh_dim_names": checked_names, + "placements": checked_placements, + } + + +def tensor_evidence(tensor: torch.Tensor) -> TensorEvidence: + """Exact local bytes plus complete one-rank DTensor placement evidence. + + A DTensor hash describes its local shard, never an implicit all-gather. The + one-rank restriction makes this shard complete for supported placements. + """ + api = dtensor_api() + distributed: DistributedEvidence | None = None + scope: Literal["full_tensor", "dtensor_local_shard"] = "full_tensor" + local = tensor.detach() + if api is not None and isinstance(tensor, api.tensor_type): + distributed = distributed_evidence(tensor, api) + local_value = required_callable(local, "to_local")() + if ( + not isinstance(local_value, torch.Tensor) + or type(local_value) is not torch.Tensor + ): + raise TypeError( + "DTensor.to_local() must return an ordinary Tensor for this pilot" + ) + local = local_value + if local.shape != tensor.shape or local.dtype != tensor.dtype: + raise ValueError( + "Single-rank DTensor local shape and dtype must match the global tensor" + ) + scope = "dtensor_local_shard" + raw = local.contiguous().cpu().reshape(-1).view(torch.uint8) + return { + "dtype": str(tensor.dtype), + "shape": list(tensor.shape), + "sha256": hashlib.sha256(raw.numpy().tobytes()).hexdigest(), + "storage_scope": scope, + "local_shape": list(local.shape), + "distributed": distributed, + } + + +def model_evidence(model: torch.nn.Module) -> ModelEvidence: + return { + "parameters": { + name: tensor_evidence(parameter) + for name, parameter in model.named_parameters() + }, + "gradients": { + name: None if parameter.grad is None else tensor_evidence(parameter.grad) + for name, parameter in model.named_parameters() + }, + "buffers": { + name: tensor_evidence(buffer) for name, buffer in model.named_buffers() + }, + } + + +def rng_evidence() -> dict[str, object]: + numpy_state = np.random.get_state(legacy=True) + if not isinstance(numpy_state, tuple) or len(numpy_state) != 5: + raise TypeError("NumPy legacy RNG state must be a five-field tuple") + return { + "python_sha256": hashlib.sha256(repr(random.getstate()).encode()).hexdigest(), + "numpy": { + "algorithm": numpy_state[0], + "state_sha256": hashlib.sha256(numpy_state[1].tobytes()).hexdigest(), + "position": numpy_state[2], + "has_gaussian": numpy_state[3], + "cached_gaussian": numpy_state[4], + }, + "torch_cpu": tensor_evidence(torch.get_rng_state()), + "torch_cuda": [ + tensor_evidence(state) for state in torch.cuda.get_rng_state_all() + ], + } + + +def optimizer_parameters( + optimizers: tuple[torch.optim.Optimizer, ...], +) -> tuple[torch.Tensor, ...]: + parameters: list[torch.Tensor] = [] + for optimizer in optimizers: + for group in optimizer.param_groups: + values: object = group["params"] + if not isinstance(values, list): + raise TypeError("Optimizer param group params must be a list") + for value in values: + if not isinstance(value, torch.Tensor): + raise TypeError("Optimizer parameters must be tensors") + parameters.append(value) + return tuple(parameters) + + +def attach_model_tracker(handle: TrainerHandle, style: TrackerStyle) -> object: + """Use the ordinary tracker implementation without adding sampling points.""" + before = tuple(handle.model.parameters()) + optimizer_before = optimizer_parameters(handle.optimizers) + if not all( + any(value is parameter for value in before) for parameter in optimizer_before + ): + raise ValueError("Optimizer has a parameter outside the single model part") + if style == "subclass": + # Custom subclasses (including DTensor parameters) take an out-of-place + # path in proxy_parameter. Reject that unsupported boundary before mutation. + if any(type(parameter) is not torch.nn.Parameter for parameter in before): + types = sorted({str(type(parameter)) for parameter in before}) + raise TypeError( + f"Subclass pilot requires exact nn.Parameter objects: {types}" + ) + from traincheck.instrumentor.proxy_wrapper.subclass import proxy_parameter + + proxy_parameter(handle.model, parent_name="model") + tracker: object = None + elif style == "sampler": + from traincheck.instrumentor.tracer import VarSampler + + sampler = VarSampler(handle.model, "model") + for optimizer in handle.optimizers: + sampler.register_hook(optimizer) + tracker = sampler + elif style is None: + tracker = None + else: + raise ValueError(f"Unsupported model tracker style: {style}") + after = tuple(handle.model.parameters()) + optimizer_after = optimizer_parameters(handle.optimizers) + if len(before) != len(after) or any(a is not b for a, b in zip(before, after)): + raise RuntimeError("Model tracking changed parameter identity") + if len(optimizer_before) != len(optimizer_after) or any( + a is not b for a, b in zip(optimizer_before, optimizer_after) + ): + raise RuntimeError("Model tracking changed optimizer parameter references") + return tracker + + +def install_instrumentation(options: Options) -> TrackerStyle: + from traincheck.config.config import InstrOpt + from traincheck.instrumentor.caches import META_VARS + + tc_config.DISABLE_WRAPPER = True + tc_config.INSTRUMENTATION_POLICY = {"interval": 1, "warm_up": 0} + if META_VARS["step"] != 0: + raise RuntimeError( + "Workload requires a fresh interpreter with TrainCheck step zero" + ) + tc_config.INSTR_OPTS = None + # Never load an unrelated instr_opts.json from the working directory. + tc_config.INSTR_OPTS_FILE = str(options.output_dir / "traces" / "instr_opts.json") + style: TrackerStyle = None + if options.mode == "plain": + tc_config.MODEL_TRACKER_STYLE = None + return style + if options.mode == "reference": + style = "subclass" + elif options.mode == "selective": + from traincheck.collect_trace import ( + get_model_tracker_instr_opts, + get_per_func_instr_opts, + ) + from traincheck.invariant.base_cls import read_inv_file + + assert options.invariants is not None + invariants = read_inv_file(str(options.invariants)) + if not invariants: + raise ValueError("Selective instrumentation requires nonempty invariants") + selected = get_model_tracker_instr_opts(invariants, "subclass") + if selected not in (None, "subclass", "sampler"): + raise ValueError(f"Unsupported selected tracker style: {selected}") + # Equality checks above validate content; branches narrow the literal type. + style = ( + "subclass" + if selected == "subclass" + else "sampler" if selected == "sampler" else None + ) + instr_opts = InstrOpt( + get_per_func_instr_opts(invariants), style, disable_proxy_dumping=True + ) + Path(tc_config.INSTR_OPTS_FILE).write_text(instr_opts.to_json() + "\n") + # Match collect_trace's generated-file boundary: positional selector keys + # become strings in JSON, as required by the runtime argument converter. + loaded_options: object = tc_config.load_instr_opts() + if not isinstance(loaded_options, InstrOpt): + raise TypeError("Saved instrumentation policy did not load as InstrOpt") + if loaded_options.model_tracker_style != style: + raise ValueError("Saved instrumentation policy changed tracker selection") + tc_config.MODEL_TRACKER_STYLE = style + from traincheck.instrumentor.tracer import Instrumentor + + Instrumentor( + torch, + scan_proxy_in_args=style is not None, + use_full_instr=False, + API_dump_stack_trace=False, + sampling_interval=1, + warm_up_steps=0, + ).instrument() + if style == "subclass": + install_standard_observers() + return style + + +def install_standard_observers() -> None: + from traincheck.instrumentor.proxy_wrapper.proxy_config import auto_observer_config + from traincheck.static_analyzer.graph_generator.call_graph_parser import ( + add_observer_given_call_graph, + ) + + if auto_observer_config["enable_auto_observer"] is not True: + raise ValueError("Reference pilot requires standard automatic model observers") + graph_dir = Path(traincheck.__file__).parent / "static_analyzer" / "func_level" + graphs = sorted(graph_dir.glob("*.log")) + if not graphs: + raise FileNotFoundError(f"No standard observer call graphs in {graph_dir}") + for graph in graphs: + add_observer_given_call_graph( + str(graph), + depth=auto_observer_config["enable_auto_observer_depth"], + observe_up_to_depth=auto_observer_config["observe_up_to_depth"], + neglect_hidden_func=auto_observer_config["neglect_hidden_func"], + neglect_hidden_module=auto_observer_config["neglect_hidden_module"], + observe_then_unproxy=auto_observer_config["observe_then_unproxy"], + ) + + +def run_training( + handle: TrainerHandle, + steps: int, + instrumented: bool, + synchronize: Callable[[], None], +) -> tuple[list[StepRecord], list[LossRecord]]: + from traincheck.instrumentor.control import start_step + + timings: list[StepRecord] = [] + losses: list[LossRecord] = [] + + def capture_metrics(*args: object, **kwargs: object) -> object: + if len(args) != 4: + raise TypeError("Pinned metrics.log must receive four positional arguments") + step, average, maximum, gradient_norm = args + if not isinstance(step, int) or isinstance(step, bool): + raise TypeError("Metrics step must be an integer") + if ( + not isinstance(average, float) + or not isinstance(maximum, float) + or not isinstance(gradient_norm, float) + ): + raise TypeError( + "Pinned metrics loss and gradient norm must be Python floats" + ) + result = handle.metrics_log(*args, **kwargs) + losses.append( + { + "step": step, + "average": average, + "maximum": maximum, + "gradient_norm": gradient_norm, + } + ) + return result + + def measured_step(*args: object, **kwargs: object) -> object: + step: object = getattr(handle.trainer, "step") + if ( + not isinstance(step, int) + or isinstance(step, bool) + or step != len(timings) + 1 + ): + raise ValueError( + "Expected consecutive native train_step calls starting at one" + ) + synchronize() + started = time.perf_counter() + try: + if instrumented: + start_step() + result = handle.train_step(*args, **kwargs) + finally: + tc_config.DISABLE_WRAPPER = True + synchronize() + timings.append({"step": step, "seconds": time.perf_counter() - started}) + return result + + tc_config.DISABLE_WRAPPER = True + setattr(handle.trainer, "train_step", measured_step) + setattr(handle.metrics, "log", capture_metrics) + try: + handle.train() + finally: + tc_config.DISABLE_WRAPPER = True + setattr(handle.trainer, "train_step", handle.train_step) + setattr(handle.metrics, "log", handle.metrics_log) + if len(timings) != steps or [row["step"] for row in losses] != list( + range(1, steps + 1) + ): + raise RuntimeError( + "Native training did not produce exactly one timing and loss per requested step" + ) + return timings, losses + + +def drain_trace_queues() -> float: + """Flush producer buffers and wait for writes; file closure occurs at exit. + + This workload traces only the main Python training thread. TrainCheck's + production writer monitor owns shutdown; joining it here would deadlock. + """ + from traincheck.instrumentor import dumper + + started = time.perf_counter() + for name in ("api_trace_buffer", "var_trace_buffer"): + # Absence means that stream produced no events; it is not an error. + buffer: object = getattr(dumper.THREAD_DATA, name, None) + if buffer is not None: + if not isinstance(buffer, dumper.TraceBuffer): + raise TypeError(f"Unexpected trace buffer: {name}") + buffer._flush() + buffer.queue.join() + return time.perf_counter() - started + + +def configure_trainer(options: Options) -> tuple[TrainerHandle, dict[str, object], str]: + registry = importlib.import_module("torchtitan.models.llama3.config_registry") + logging_module = importlib.import_module("torchtitan.tools.logging") + required_callable(logging_module, "init_logger")() + config = required_callable(registry, "llama3_debugmodel")() + overrides: dict[str, object] = { + "training.steps": options.steps, + "training.seq_len": options.seq_len, + "training.local_batch_size": options.batch_size, + "debug.seed": options.seed, + "checkpoint.enable": False, + "validator.enable": False, + "compile.enable": False, + "profiler.enable_profiling": False, + "metrics.log_freq": 1, + "metrics.enable_tensorboard": False, + "dump_folder": str(options.output_dir / "titan"), + } + for key, value in overrides.items(): + owner = config + components = key.split(".") + for component in components[:-1]: + owner = getattr(owner, component) + # Required pinned fields must exist; missing fields cannot become overrides. + getattr(owner, components[-1]) + setattr(owner, components[-1], value) + if getattr(owner, components[-1]) != value: + raise ValueError(f"TorchTitan did not retain requested override {key}") + representation = repr(config) + return ( + TrainerHandle.from_trainer(required_callable(config, "build")()), + overrides, + representation, + ) + + +def write_json(path: Path, payload: object) -> None: + path.write_text(json.dumps(payload, indent=2, allow_nan=False) + "\n") + + +def main(argv: Sequence[str] | None = None) -> None: + options = parse_options(argv) + if os.environ["WORLD_SIZE"] != "1" or os.environ["RANK"] != "0": + raise ValueError("Run this pilot under torchrun with exactly one rank") + if not torch.cuda.is_available() or torch.cuda.device_count() != 1: + raise RuntimeError("Pilot requires exactly one visible CUDA GPU") + if not options.torchtitan_root.is_dir(): + raise FileNotFoundError(options.torchtitan_root) + options.output_dir.mkdir(parents=True, exist_ok=False) + (options.output_dir / "traces").mkdir() + os.environ["TRAINCHECK_OUTPUT_DIR"] = str(options.output_dir / "traces") + os.chdir(options.torchtitan_root) + sys.path.insert(0, str(options.torchtitan_root)) + setup_started = time.perf_counter() + tc_config.DISABLE_WRAPPER = True + random.seed(options.seed) + np.random.seed(options.seed) + torch.manual_seed(options.seed) + torch.cuda.manual_seed_all(options.seed) + style = install_instrumentation(options) + handle, overrides, representation = configure_trainer(options) + try: + parameter_count = sum( + parameter.numel() for parameter in handle.model.parameters() + ) + if parameter_count != EXPECTED_PARAMETERS: + raise ValueError( + f"Expected {EXPECTED_PARAMETERS} parameters, found {parameter_count}" + ) + torch.cuda.synchronize() + initial = model_evidence(handle.model) + initial_rng = rng_evidence() + parameter_types = sorted( + {str(type(parameter)) for parameter in handle.model.parameters()} + ) + tracker = attach_model_tracker(handle, style) + if initial != model_evidence(handle.model) or initial_rng != rng_evidence(): + raise RuntimeError( + "Model tracker installation changed application state or RNG" + ) + setup_seconds = time.perf_counter() - setup_started + write_json( + options.output_dir / "initial.json", {"model": initial, "rng": initial_rng} + ) + timings, losses = run_training( + handle, options.steps, options.mode != "plain", torch.cuda.synchronize + ) + torch.cuda.synchronize() + final = model_evidence(handle.model) + final_rng = rng_evidence() + drain_seconds = drain_trace_queues() + write_json( + options.output_dir / "result.json", + { + "schema_version": 2, + "status": "completed", + "mode": options.mode, + "model_config": "llama3_debugmodel", + "parameter_count": parameter_count, + "parameter_types_before_tracking": parameter_types, + "selected_model_tracker": style, + "model_tracking_enabled": style is not None, + "sampler_initial_setup_snapshot": style == "sampler", + "ordinary_full_collection": { + "use_full_instr": False, + "sampling_interval": 1, + "warm_up_steps": 0, + }, + "seed_policy": "Seed Python, NumPy, Torch CPU, and every visible CUDA generator before instrumentation/configure; retain TorchTitan debug.seed override during trainer initialization.", + "overrides": overrides, + "effective_config_repr": representation, + "invariants_sha256": ( + None + if options.invariants is None + else hashlib.sha256(options.invariants.read_bytes()).hexdigest() + ), + "python_version": sys.version, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "initial": {"model": initial, "rng": initial_rng}, + "final": {"model": final, "rng": final_rng}, + "losses": losses, + "steps": timings, + "training_seconds": sum(row["seconds"] for row in timings), + "setup_seconds": setup_seconds, + "trace_queue_drain_seconds": drain_seconds, + "timing_scope": "Synchronized native train_step including start_step; excludes initialization, state evidence, outer train-loop housekeeping, queue drain, and process shutdown/file closure.", + }, + ) + # Keep native sampler and its hook alive until all training has completed. + del tracker + finally: + tc_config.DISABLE_WRAPPER = True + handle.close() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def parse_options(argv: Sequence[str] | None = None) -> Options: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--mode", + choices=("plain", "reference", "reference_api", "selective"), + required=True, + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--torchtitan-root", type=Path, required=True) + parser.add_argument("--invariants", type=Path) + parser.add_argument("--steps", type=int, default=10) + parser.add_argument("--seq-len", type=int, default=256) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args(argv) + return Options( + mode=args.mode, + output_dir=args.output_dir.resolve(), + torchtitan_root=args.torchtitan_root.resolve(), + steps=args.steps, + seq_len=args.seq_len, + batch_size=args.batch_size, + seed=args.seed, + invariants=None if args.invariants is None else args.invariants.resolve(), + ) + + +if __name__ == "__main__": + main()