diff --git a/docs/measurements/ring18-distributed-20261031.md b/docs/measurements/ring18-distributed-20261031.md new file mode 100644 index 00000000..6dab1e9d --- /dev/null +++ b/docs/measurements/ring18-distributed-20261031.md @@ -0,0 +1,196 @@ +# Ring18 two-GPU training and tracing qualification + +Thirteen bounded two-GPU runs completed on 2026-10-31. DDP training and +checkpoint interchange work for this small workload, but the learned API rules +did not produce a fault-specific alarm for incorrect distributed loss +normalization. Active parameter sampling also exposes stale post-update +timestamps. This qualifies infrastructure and identifies gaps for #165; it is +not an agent-benefit experiment or an official RSI result. + +## Workload and provenance + +- Ring18, two NVIDIA A16 GPUs (indices 0 and 1), driver 535.309.01. All eight + GPUs were idle at discovery and after completion; other devices were unused. +- Python 3.10.14, PyTorch 2.2.2+cu121, NumPy 1.26.0, pandas 2.2.2, + polars 1.7.1, numba 0.60.0, orjson 3.10.18. Existing environment reused; + source and outputs live in an isolated experiment directory. +- TrainCheck base `754277ae87e2189cebf16b405940bef1fee5904c` (PR #164); + workload/tracing commit `15ed99cc2d72853ebeb44a665b3a91d8923dd99e`; + complete-state comparator `ebd66a5f7062952ee0f059f2df8a61b2dd2a2575`. + Both commits are signed. Runtime workload, tracing and core source hashes + match across all runs; comparator/test additions between phases are recorded + in each invocation's source manifest. +- A 5,504-parameter, two-layer Transformer, width 16, two attention heads, + vocabulary 32, sequence length 8, dropout zero. Synthetic token targets are + `(token + 1) % 32`; this is a numerical contract probe, not a model-quality task. +- DDP/NCCL, two microbatches of two sequences per rank per optimizer update, + four updates, AdamW, clipping and StepLR. Each update contains 10 valid target + tokens on rank 0 and 22 on rank 1. Full accumulation windows only. +- Seeds, deterministic algorithms, disabled TF32 and + `CUBLAS_WORKSPACE_CONFIG=:4096:8` are fixed. FP32 and BF16 autocast are tested. + BF16 uses no GradScaler; FP16 overflow behavior is outside scope. + +## Independent correctness and checkpoint results + +The independent first-update oracle concatenates the logical global batch and +computes a token-mean gradient in a separate model. Correct DDP accumulation +scales each local loss sum by `world_size / global_valid_tokens`. The injected +fault instead divides by each rank's local token count before DDP averaging. + +| Configuration | Maximum absolute gradient error, both ranks | Oracle outcome | +| --- | ---: | --- | +| Healthy FP32, plain/API/subclass | 4.47035e-8 | Pass (`atol=2e-6`, `rtol=2e-4`) | +| Faulty FP32, plain/API | 0.0695832 | Fail with the same declared tolerance | +| Healthy BF16, plain/subclass | 0.000564575 | Pass (`atol=0.003`, `rtol=0.03`) | + +The faulty model's two ranks still agree exactly. Rank agreement alone therefore +does not establish correct global loss normalization. The fault also changes the +final model relative to the healthy run. + +All nine complete checkpoint comparisons pass: + +| Left | Right | +| --- | --- | +| Plain FP32 continuous | API FP32 continuous | +| Plain FP32 continuous | Subclass FP32 continuous | +| Plain FP32 continuous | Plain save at update 2, plain resume to update 4 | +| Plain FP32 continuous | Plain save at update 2, subclass resume to update 4 | +| Plain FP32 continuous | Subclass save at update 2, plain resume to update 4 | +| Plain FP32 save at update 2 | Subclass FP32 save at update 2 | +| Plain BF16 continuous | Subclass BF16 continuous | +| API FP32 healthy reference | Fresh API FP32 healthy control | +| Plain FP32 fault | API FP32 fault | + +Each comparison covers both ranks: 274 tensor pairs, 112 scalar pairs and +78 container pairs. Model, stored gradients, optimizer, scheduler, CPU/CUDA RNG +and configuration fields are compared without excluded fields or normalization. +All value bits and whole checkpoint-file hashes match. Storage aliasing and +original device placement are not reconstructed by the artifact comparison. + +Resume uses fresh processes and the same world size. Checkpoints are taken at +complete update boundaries; stored gradients are evidence, and the next update +clears gradients before use. Data is regenerated from absolute step/rank indices, +so DataLoader cursor recovery is not tested. RNG is saved/restored, but dropout +is zero and these runs do not exercise stochastic model operations. + +## Trace coverage and timestamp defect + +API and subclass modes use the normal API instrumenter. Subclass tracking is +attached before DDP and retains parameter identity; model parameters are sampled +explicitly before and after each update. Proxy argument scanning is enabled in +subclass mode; API mode disables it. Standard tensor hashing is enabled in both +instrumented modes. Bookkeeping, reference calculations, checkpoint I/O +and rank-consistency measurements are suppressed from user tracing. Optimizer +state is not registered, and no cross-rank TrainCheck checker is installed. + +Per rank, a four-update FP32 run has 1,876 API pre/post pairs in both instrumented +modes. Subclass mode additionally has 297 variable records across 27 parameter +names, including 108 before/after sample pairs. No duplicate API event keys were +observed. These counts are record audits, not a proof that every desired +application event was captured. + +Every post-update parameter sample in both full FP32 and BF16 subclass runs +carries a timestamp **earlier than that update's `AdamW.step` pre-event**, despite +a changed recorded data hash. At the first executed update, all 27 parameters +also have identical before/after timestamps on each rank. Subsequent pairs need +not have equal timestamps, but their post-update timestamps are still stale. + +Source inspection explains the observation: `ProxyParameter.dump_trace` uses +`last_update_timestamp`; explicit registry sampling does not refresh it, and +in-place optimizer mutation does not necessarily pass through the Python +assignment hook. This prevents treating these records as reliable evidence of +the temporal relationship between an optimizer call and its updates. No timestamp +rewrite or core fix was applied to make the experiment pass. Existing hashes +cover the instrumenter's first-row semantics, not a full tensor-content proof; +full-state equality above comes from the independent checkpoint comparison. + +## Learned diagnostics on a fresh healthy run and the fault + +Learn rules separately from each rank's healthy API reference. Use pandas and +one CPU worker; disable `FunctionLeadRelation` and `FunctionCoverRelation`, as in +the established pilot workflow. Freeze these rule files, then run the real +checker with coverage observation on the fresh healthy control and the fault. +Do not concatenate ranks and call the result a cross-rank check. + +| Rank | Learned / active rules in each check | Healthy violations | Fault violations | Fault-only violated rules | +| --- | ---: | ---: | ---: | ---: | +| 0 | 592 / 592 | 2 | 1 | 0 | +| 1 | 593 / 593 | 2 | 2 | 0 | + +All four checks complete with no dormant rules, unobserved outcomes or checker +errors. Activation does not imply all possible comparisons ran: individual +checkers can stop at their first violation. No model-variable rules are checked +in this API-only inference comparison. + +The shared alarm is `DistributedDataParallel.forward()` containing +`Embedding.forward()`. Rank 0's healthy-only alarm is a distinct-arguments rule +for `LayerNorm.forward()`; rank 1's healthy and faulty runs both violate a +distinct-arguments rule for `TransformerEncoderLayer.forward()`. The full rule +identities and failing evidence are retained; no alarms were removed. Counts of +identity preconditions alone do not establish that relation parameters are +portable across processes. + +Thus this configuration gives no fault-specific diagnostic for the independently +confirmed normalization error. Fewer alarms in rank 0's faulty run do not mean +better training. This is one fault and one fresh healthy control, not an estimate +of general detection accuracy. + +## Cost, validation and remaining work + +The thirteen supervised workload phases total 211.29 seconds of wall time. +For context, maximum per-rank loop durations were 0.348 s plain FP32, 3.549 s +API FP32 reference, 4.513 s API FP32 healthy repeat, and 5.753 s subclass FP32. +Plain/subclass BF16 took 0.345/4.739 s. These short runs include oracle snapshots, +first-use costs and, for subclass mode, hashing. They are operational observations, +not repeated steady-state overhead estimates. CPU inference overlapped part of +the later matrix. Full FP32 API/subclass traces total 4,946,510/5,595,602 bytes +across both ranks, excluding instrumentation logs and checkpoints. + +- 28 focused tests pass locally on Python 3.10/Torch 2.2.2 and on Ring18 with + CUDA visible. Configured hooks and explicit strict mypy over seven new + production/test modules pass; the normal mypy hook does not cover `tools/`. +- An initial local sandbox attempt blocked OpenMP shared memory; the same + composed suite passed with the required access. This was not bypassed by + skipping the tracing tests. +- With `CUDA_VISIBLE_DEVICES` empty on Ring18's CUDA-enabled Torch build, the + suite has 26 passes and two tracer-startup failures. Metadata inspection of + `torch.backends.cuda` calls `current_device()` and attempts CUDA initialization. + This additional compatibility failure is preserved; exposing a GPU permits + all 28 tests to pass but does not fix the hidden-GPU case. +- The reused supervisor measures the launcher's process group; torchrun workers + use separate sessions. Its RSS figure/8 GiB setting is not a complete worker + memory bound, and cleanup relies on torchrun signal forwarding. Its trace-size + limit covers API files only. Final process/GPU audit found no remaining study + processes and all GPUs idle. Future unattended trials need stronger tree-wide + supervision before claiming complete memory/cleanup enforcement. + +Next gates for #165 are to repair and validate parameter observation timestamps, +establish relevant state/normalization diagnostic coverage, and qualify FSDP and +the selected feature tasks. Partial windows, FP16 overflow, stochastic resume, +different world sizes and natural agent implementations remain untested. Do not +start the paired agent study on the assumption that these results cover them. + +## Reproduction and artifacts + +Use `python -m torch.distributed.run --standalone --nproc-per-node=2 +tools/distributed_probe/workload.py --output RUN --steps 4 --device cuda +--instrumentation plain` with the fixed environment above. Choose `api` or +`subclass`, `--fault local-mean`, or `--precision bf16` for the named variants. +Save with `--stop-after 2`, then use `--resume SAVE_RUN` in a fresh run. Compare +with `python tools/distributed_probe/compare.py --left LEFT --right RIGHT +--output comparison.json`. Invocation manifests retain the exact commands, +runtime source hashes and supervisor settings for every phase. + +Remote root: `/home/yuxuan/experiments/traincheck-distributed-20261031` on Ring18. +Local evidence root: `/Users/matt/gitrepos/traincheck/outputs/ring18-distributed-20261031`. +The downloaded `evidence.tar.gz` contains 188 verified files, including complete +checkpoints, variable traces, rule/checker outputs, phase logs and source hashes. +Large API traces remain under the remote `runs/*/payload/rank*/trace` directories; +their hashes and sizes are in `results.json`. + +Archive SHA256: `ba267cddd70abb54b7aae8a7ead085e6bb3f4fbe91771475d703659841425922`. +`evidence-bundle/evidence/diagnostics.json` records alarm identities; +`results.json` records checkpoint comparisons and raw timestamp evidence; +`retrieval-verification.json` records archive verification and the final 160 +Python source-file hash matches. No traces, checkpoints or machine-specific +operator scripts are committed to Git. diff --git a/tests/test_distributed_probe_compare.py b/tests/test_distributed_probe_compare.py new file mode 100644 index 00000000..a609d05c --- /dev/null +++ b/tests/test_distributed_probe_compare.py @@ -0,0 +1,101 @@ +"""Checkpoint parity must include optimizer/RNG state and exact scalar contracts.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from tools.distributed_probe.compare import ( + Comparison, + checkpoint_files, + compare_files, + compare_paths, + compare_values, + validate_state, +) + + +def state() -> dict[str, object]: + return { + "model": {"weight": torch.tensor([1.0, 2.0])}, + "optimizer": {"state": {0: {"moment": torch.tensor([0.1, 0.2])}}}, + "rng": torch.tensor([1, 2, 3], dtype=torch.uint8), + "step": 4, + "options": (False, ["fp32", None]), + } + + +def test_complete_state_parity_and_optimizer_difference(tmp_path: Path) -> None: + left = tmp_path / "left.pt" + right = tmp_path / "right.pt" + torch.save(state(), left) + torch.save(state(), right) + report = compare_files(left, right) + assert report["equal"] is True + assert report["tensor_pairs_compared"] == 3 + changed = state() + changed["optimizer"] = {"state": {0: {"moment": torch.tensor([0.1, 0.3])}}} + torch.save(changed, right) + report = compare_files(left, right) + assert report["equal"] is False + assert report["differences"] == [ + { + "path": "$['optimizer']['state'][0]['moment']", + "reason": "tensor value bits differ", + } + ] + + +@pytest.mark.parametrize( + ("left", "right"), + [ + (1, True), + ([1], (1,)), + ({"a": 1}, {}), + ({"a": 1, "b": 2}, {"b": 2, "a": 1}), + (0.0, -0.0), + (torch.tensor([0.0]), torch.tensor([-0.0])), + (torch.tensor([1]), torch.tensor([1.0])), + (torch.tensor([1]), torch.tensor([[1]])), + ([1, 2], [1]), + ], +) +def test_meaningful_differences_are_preserved(left: object, right: object) -> None: + validate_state(left) + validate_state(right) + report = Comparison() + compare_values(left, right, report) + assert report.differences + + +def test_identical_nan_payload_is_byte_equal_not_a_finiteness_claim() -> None: + left = torch.tensor([float("nan")]) + report = Comparison() + compare_values(left, left.clone(), report) + assert not report.differences + + +@pytest.mark.parametrize("value", [{True: 1}, {"value": object()}, {1, 2}]) +def test_unsupported_state_is_rejected(value: object) -> None: + with pytest.raises(TypeError): + validate_state(value) + + +def test_phase_payload_discovery_requires_matching_rank_sets(tmp_path: Path) -> None: + left = tmp_path / "left" + right = tmp_path / "right" + for rank in (0, 1): + for root in (left, right / "payload"): + path = root / f"rank{rank}" / "checkpoint.pt" + path.parent.mkdir(parents=True) + torch.save(state(), path) + assert compare_paths(left, right)["equal"] is True + (right / "payload" / "rank1" / "checkpoint.pt").unlink() + with pytest.raises(ValueError, match="sets differ"): + compare_paths(left, right) + with pytest.raises(ValueError, match="No rank checkpoints"): + checkpoint_files(tmp_path) + with pytest.raises(FileNotFoundError): + checkpoint_files(tmp_path / "absent") diff --git a/tests/test_distributed_probe_tracing.py b/tests/test_distributed_probe_tracing.py new file mode 100644 index 00000000..4d14e4a7 --- /dev/null +++ b/tests/test_distributed_probe_tracing.py @@ -0,0 +1,155 @@ +"""The probe must delimit real events and preserve its application computation.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from tools.distributed_probe.tracing import TraceSession + +ROOT = Path(__file__).resolve().parents[1] + +WORKER = """ +import copy +import json +import sys +from pathlib import Path + +import torch +from tools.distributed_probe.tracing import start_tracing + +torch.set_num_threads(1) +torch.manual_seed(731) +model = torch.nn.Linear(2, 1) +reference = copy.deepcopy(model) +values = torch.tensor([[1., 2.], [3., 4.]]) +ref_opt = torch.optim.SGD(reference.parameters(), lr=0.03) +ref_loss = reference(values).square().mean() +ref_loss.backward() +ref_opt.step() +identities = [id(p) for p in model.parameters()] +session = start_tracing(sys.argv[1], Path(sys.argv[2]), model) +assert identities == [id(p) for p in model.parameters()] +optimizer = torch.optim.SGD(model.parameters(), lr=0.03) +with session.suppress(): + torch.sin(values) +with session.step(4): + torch.cos(values) + with session.suppress(): + torch.sin(values) + loss = model(values).square().mean() + loss.backward() + optimizer.step() +with session.suppress(): + assert torch.equal(loss, ref_loss) + for key, value in model.state_dict().items(): + assert torch.equal(value, reference.state_dict()[key]), key +try: + with session.step(5): + with session.suppress(): + raise RuntimeError('expected application failure') +except RuntimeError as exc: + assert str(exc) == 'expected application failure' +else: + raise AssertionError('application failure was swallowed') +report = session.close() +assert report['completed_trace_steps'] == [4] +Path(sys.argv[3]).write_text(json.dumps(report)) +""" + + +def json_object(path: Path) -> dict[str, object]: + value: object = json.loads(path.read_text()) + if not isinstance(value, dict) or not all(isinstance(key, str) for key in value): + raise TypeError("Expected JSON object with string keys") + return {key: item for key, item in value.items()} + + +def trace_records(root: Path, kind: str) -> list[dict[str, object]]: + records: list[dict[str, object]] = [] + for path in sorted(root.glob(f"trace_{kind}_*.log")): + for line in path.read_text().splitlines(): + value: object = json.loads(line) + if not isinstance(value, dict) or not all( + isinstance(key, str) for key in value + ): + raise TypeError("Expected trace record with string keys") + records.append({key: item for key, item in value.items()}) + return records + + +@pytest.mark.parametrize("mode", ["api", "subclass"]) +def test_actual_events_and_bookkeeping_suppression(tmp_path: Path, mode: str) -> None: + output = tmp_path / "traces" + report_path = tmp_path / "report.json" + environment = dict(os.environ) + environment["PYTHONPATH"] = str(ROOT) + environment["PYTHONDONTWRITEBYTECODE"] = "1" + result = subprocess.run( + [sys.executable, "-c", WORKER, mode, str(output), str(report_path)], + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + report = json_object(report_path) + assert report["mode"] == mode + assert report["scan_proxy_in_args"] is (mode == "subclass") + assert report["completed_trace_steps"] == [4] + assert report["optimizer_state_registered"] is False + assert report["cross_rank_checks_executed"] is False + assert report["trace_record_counts"] is None + records = trace_records(output, "API") + assert records + functions = [record["function"] for record in records] + assert any(isinstance(name, str) and name.endswith(".cos") for name in functions) + assert not any( + isinstance(name, str) and name.endswith(".sin") for name in functions + ) + variables = trace_records(output, "VAR") + if mode == "subclass": + assert any("proxy_obj_names" in record for record in records) + assert report["registered_parameter_names"] == ["model.weight", "model.bias"] + assert {record["var_name"] for record in variables} == { + "model.weight", + "model.bias", + } + assert any( + record["dump_loc"] == "distributed_probe.update.after" + for record in variables + ) + else: + assert report["registered_parameter_names"] == [] + assert variables == [] + + +def test_scope_contracts_and_failure_not_completed(tmp_path: Path) -> None: + session = TraceSession("plain", tmp_path, 0, ()) + for bad_step in (0, -1, True): + with pytest.raises(ValueError, match="positive integer"): + with session.step(bad_step): + pass + with session.step(3): + with pytest.raises(RuntimeError, match="overlap"): + with session.step(4): + pass + with pytest.raises(RuntimeError, match="inside an update"): + session.close() + with pytest.raises(ValueError, match="consecutive"): + with session.step(5): + pass + with pytest.raises(LookupError): + with session.step(4): + raise LookupError("application") + assert session.close()["completed_trace_steps"] == [3] + with pytest.raises(RuntimeError, match="already closed"): + with session.suppress(): + pass diff --git a/tests/test_distributed_probe_workload.py b/tests/test_distributed_probe_workload.py new file mode 100644 index 00000000..ac0579b5 --- /dev/null +++ b/tests/test_distributed_probe_workload.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +from tools.distributed_probe.workload import ( + Fault, + Options, + TinyTransformer, + checkpoint_payload, + compare_gradients, + copy_state, + gradient_reference, + load_checkpoint, + loss_scale, + make_batch, + model_gradients, + parse_options, + token_loss, +) + + +def test_global_token_oracle_distinguishes_local_mean_fault() -> None: + torch.set_num_threads(1) + torch.manual_seed(41) + model = TinyTransformer() + device = torch.device("cpu") + expected = gradient_reference(model, 2, 2, device, "fp32") + rank_batches = [ + [make_batch(1, rank, microbatch, device) for microbatch in range(2)] + for rank in range(2) + ] + local_counts = [sum(int(b.mask.sum()) for b in batches) for batches in rank_batches] + assert local_counts[0] != local_counts[1] + faults: tuple[Fault, ...] = ("none", "local-mean") + for fault in faults: + rank_gradients: list[dict[str, torch.Tensor]] = [] + for rank, batches in enumerate(rank_batches): + replica = copy.deepcopy(model) + scale = loss_scale(fault, 2, sum(local_counts), local_counts[rank]) + for batch in batches: + (token_loss(replica(batch.tokens), batch) * scale).backward() + rank_gradients.append(model_gradients(replica)) + averaged = { + name: (rank_gradients[0][name] + rank_gradients[1][name]) / 2 + for name in expected + } + comparison = compare_gradients(averaged, expected, "fp32") + assert comparison["passed"] is (fault == "none") + + +def test_state_snapshot_is_independent_and_rejects_unsupported_values() -> None: + source = torch.tensor([1.0, 2.0]) + snapshot = copy_state({"tensor": source, "empty": [], "optional": None}) + source.add_(3) + assert isinstance(snapshot, dict) + assert isinstance(snapshot["tensor"], torch.Tensor) + assert torch.equal(snapshot["tensor"], torch.tensor([1.0, 2.0])) + with pytest.raises(TypeError, match="unsupported"): + copy_state({"bad": object()}) + with pytest.raises(TypeError, match="dictionary key"): + copy_state({True: 1}) + + +def test_oracle_rejects_nonfinite_gradients() -> None: + with pytest.raises(ValueError, match="Nonfinite"): + compare_gradients( + {"p": torch.tensor([float("nan")])}, {"p": torch.zeros(1)}, "fp32" + ) + + +def checkpoint_fixture( + path: Path, +) -> tuple[nn.Module, torch.optim.AdamW, torch.optim.lr_scheduler.StepLR, Options]: + model = nn.Linear(3, 2) + optimizer = torch.optim.AdamW(model.parameters(), foreach=False) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2) + model(torch.ones(2, 3)).sum().backward() + optimizer.step() + scheduler.step() + options = Options(path, None, 4, 4, "fp32", "none", "plain", "cpu", 2) + return model, optimizer, scheduler, options + + +def test_checkpoint_restores_model_optimizer_scheduler_and_rng(tmp_path: Path) -> None: + torch.manual_seed(7) + model, optimizer, scheduler, options = checkpoint_fixture(tmp_path) + payload = checkpoint_payload( + model, optimizer, scheduler, 2, 0, 2, options, torch.device("cpu") + ) + path = tmp_path / "checkpoint.pt" + torch.save(payload, path) + expected_rng = torch.rand(5) + expected_weights = model.weight.detach().clone() + with torch.no_grad(): + model.weight.add_(10) + optimizer.state.clear() + scheduler.last_epoch = 90 + loaded = load_checkpoint( + path, model, optimizer, scheduler, 0, 2, options, torch.device("cpu") + ) + assert loaded == 2 + assert torch.equal(model.weight, expected_weights) + assert torch.equal(torch.rand(5), expected_rng) + assert len(optimizer.state) == 2 + assert scheduler.last_epoch == 1 + + +@pytest.mark.parametrize( + "field,value", + [("rank", 1), ("rank", False), ("precision", "bf16"), ("completed_steps", 4)], +) +def test_checkpoint_rejects_incompatible_provenance( + tmp_path: Path, field: str, value: object +) -> None: + model, optimizer, scheduler, options = checkpoint_fixture(tmp_path) + raw: dict[str, object] = dict( + checkpoint_payload( + model, optimizer, scheduler, 2, 0, 2, options, torch.device("cpu") + ) + ) + raw[field] = value + path = tmp_path / "checkpoint.pt" + torch.save(raw, path) + with pytest.raises(ValueError, match="Checkpoint"): + load_checkpoint( + path, model, optimizer, scheduler, 0, 2, options, torch.device("cpu") + ) + + +def test_checkpoint_missing_required_state_is_an_error(tmp_path: Path) -> None: + model, optimizer, scheduler, options = checkpoint_fixture(tmp_path) + payload = checkpoint_payload( + model, optimizer, scheduler, 2, 0, 2, options, torch.device("cpu") + ) + del payload["optimizer"] + path = tmp_path / "checkpoint.pt" + torch.save(payload, path) + with pytest.raises(KeyError, match="optimizer"): + load_checkpoint( + path, model, optimizer, scheduler, 0, 2, options, torch.device("cpu") + ) + + +def test_cli_rejects_invalid_update_boundary() -> None: + with pytest.raises(SystemExit): + parse_options(["--output", "/unused", "--steps", "2", "--stop-after", "3"]) diff --git a/tools/distributed_probe/__init__.py b/tools/distributed_probe/__init__.py new file mode 100644 index 00000000..beac88e1 --- /dev/null +++ b/tools/distributed_probe/__init__.py @@ -0,0 +1 @@ +"""Bounded distributed-training contract and tracing qualification tools.""" diff --git a/tools/distributed_probe/compare.py b/tools/distributed_probe/compare.py new file mode 100644 index 00000000..a4d7f5a9 --- /dev/null +++ b/tools/distributed_probe/compare.py @@ -0,0 +1,199 @@ +"""Exact, complete comparison of trusted probe checkpoint artifacts. + +The weights-only loader maps storage to CPU. Tensor comparison preserves dtype, +shape and every stored value bit (including signed zero/NaN payloads); device +placement and storage aliasing are not reconstructed by this artifact comparison. +Dictionary order and key types, sequence types, and primitive types also matter. +No checkpoint field is excluded. Missing artifacts are errors, not equality. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import struct +from dataclasses import dataclass, field +from pathlib import Path +from typing import TypedDict + +import torch + + +@dataclass +class Difference: + path: str + reason: str + + +@dataclass +class Comparison: + tensor_pairs_compared: int = 0 + scalar_pairs_compared: int = 0 + container_pairs_compared: int = 0 + differences: list[Difference] = field(default_factory=list) + + +class FileComparison(TypedDict): + left: str + right: str + left_sha256: str + right_sha256: str + equal: bool + tensor_pairs_compared: int + scalar_pairs_compared: int + container_pairs_compared: int + differences: list[dict[str, str]] + + +class PathComparison(TypedDict): + equal: bool + pairs: list[FileComparison] + scope: str + + +class Options(argparse.Namespace): + left: Path + right: Path + output: Path + + +def validate_state(value: object, path: str = "$") -> None: + """Reject unsupported state rather than coercing it for a comparison.""" + if isinstance(value, torch.Tensor): + if value.layout != torch.strided or value.is_quantized: + raise TypeError(f"{path}: only dense nonquantized tensors are supported") + return + if value is None or type(value) in (bool, int, float, str): + return + if isinstance(value, (tuple, list)): + for index, item in enumerate(value): + validate_state(item, f"{path}[{index}]") + return + if isinstance(value, dict): + for key, item in value.items(): + if type(key) not in (int, str): + raise TypeError(f"{path}: unsupported mapping key type") + validate_state(item, f"{path}[{key!r}]") + return + raise TypeError(f"{path}: unsupported state type {type(value).__name__}") + + +def tensor_bytes(value: torch.Tensor) -> bytes: + return ( + value.detach() + .cpu() + .contiguous() + .reshape(-1) + .view(torch.uint8) + .numpy() + .tobytes() + ) + + +def compare_values( + left: object, right: object, result: Comparison, path: str = "$" +) -> None: + """Compare values after boundary validation; collect differences without truncation.""" + if type(left) is not type(right): + result.differences.append(Difference(path, "types differ")) + elif isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor): + result.tensor_pairs_compared += 1 + if left.dtype != right.dtype or left.shape != right.shape: + result.differences.append(Difference(path, "tensor shape/dtype differ")) + elif tensor_bytes(left) != tensor_bytes(right): + result.differences.append(Difference(path, "tensor value bits differ")) + elif isinstance(left, dict) and isinstance(right, dict): + result.container_pairs_compared += 1 + left_keys = [(type(key), key) for key in left] + right_keys = [(type(key), key) for key in right] + if left_keys != right_keys: + result.differences.append( + Difference(path, "mapping keys/order/types differ") + ) + for key in left: + if key in right and (type(key), key) in right_keys: + compare_values(left[key], right[key], result, f"{path}[{key!r}]") + elif isinstance(left, (tuple, list)) and isinstance(right, (tuple, list)): + result.container_pairs_compared += 1 + if len(left) != len(right): + result.differences.append(Difference(path, "sequence lengths differ")) + for index, (left_item, right_item) in enumerate(zip(left, right)): + compare_values(left_item, right_item, result, f"{path}[{index}]") + else: + result.scalar_pairs_compared += 1 + if isinstance(left, float) and isinstance(right, float): + equal = struct.pack("!d", left) == struct.pack("!d", right) + else: + equal = left == right + if not equal: + result.differences.append(Difference(path, "scalar value differs")) + + +def compare_files(left: Path, right: Path) -> FileComparison: + left_value: object = torch.load(left, map_location="cpu", weights_only=True) + right_value: object = torch.load(right, map_location="cpu", weights_only=True) + validate_state(left_value) + validate_state(right_value) + result = Comparison() + compare_values(left_value, right_value, result) + return { + "left": str(left.resolve()), + "right": str(right.resolve()), + "left_sha256": hashlib.sha256(left.read_bytes()).hexdigest(), + "right_sha256": hashlib.sha256(right.read_bytes()).hexdigest(), + "equal": not result.differences, + "tensor_pairs_compared": result.tensor_pairs_compared, + "scalar_pairs_compared": result.scalar_pairs_compared, + "container_pairs_compared": result.container_pairs_compared, + "differences": [ + {"path": item.path, "reason": item.reason} for item in result.differences + ], + } + + +def checkpoint_files(path: Path) -> dict[str, Path]: + """Accept a checkpoint, worker directory, or phase/payload with rank directories.""" + if path.is_file(): + return {"checkpoint.pt": path} + if not path.is_dir(): + raise FileNotFoundError(path) + if (path / "checkpoint.pt").is_file(): + return {"checkpoint.pt": path / "checkpoint.pt"} + direct = list(path.glob("rank*/checkpoint.pt")) + nested = list((path / "payload").glob("rank*/checkpoint.pt")) + if direct and nested: + raise ValueError("Both direct and payload checkpoint sets exist") + paths = direct or nested + if not paths: + raise ValueError(f"No rank checkpoints found under {path}") + return {item.parent.name: item for item in sorted(paths)} + + +def compare_paths(left: Path, right: Path) -> PathComparison: + left_files, right_files = checkpoint_files(left), checkpoint_files(right) + if left_files.keys() != right_files.keys(): + raise ValueError("Checkpoint rank/artifact sets differ") + comparisons = [ + compare_files(left_files[key], right_files[key]) for key in left_files + ] + return { + "equal": all(item["equal"] is True for item in comparisons), + "pairs": comparisons, + "scope": "complete stored checkpoint; CPU mapped; storage aliasing not checked", + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--left", type=Path, required=True) + parser.add_argument("--right", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(namespace=Options()) + report = compare_paths(args.left, args.right) + args.output.write_text(json.dumps(report, indent=2) + "\n") + return 0 if report["equal"] is True else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/distributed_probe/tracing.py b/tools/distributed_probe/tracing.py new file mode 100644 index 00000000..c74c4b5f --- /dev/null +++ b/tools/distributed_probe/tracing.py @@ -0,0 +1,224 @@ +"""Rank-local TrainCheck setup for the distributed qualification workload. + +Install once in a fresh worker process, before DDP/optimizer construction. The +API wrappers persist until process exit; only explicitly delimited update scopes +emit events. Subclass mode samples the registered model parameters at update +boundaries in addition to their ordinary assignment observations. It does not +register optimizer state or implement cross-rank invariant checking. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Literal, TypedDict + +if TYPE_CHECKING: + from torch import nn + +Mode = Literal["plain", "api", "subclass"] + + +class TraceReport(TypedDict): + mode: Mode + instrumented_apis: int + scan_proxy_in_args: bool + registered_parameter_names: list[str] + completed_trace_steps: list[int] + tensor_hash_enabled: bool + model_sampling: str + optimizer_state_registered: bool + cross_rank_checks_executed: bool + trace_directory: str + buffers_drained: bool + trace_record_counts: None + + +@dataclass +class TraceSession: + """Own one worker's tracing gates; suppression restores its prior state.""" + + mode: Mode + output: Path + instrumented_apis: int + parameter_names: tuple[str, ...] + _closed: bool = False + _in_step: bool = False + _steps: list[int] = field(default_factory=list) + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("Trace session is already closed") + + @contextmanager + def suppress(self) -> Iterator[None]: + """Exclude bookkeeping from API and subclass records, including errors.""" + self._require_open() + if self.mode == "plain": + yield + return + from traincheck.config import config + from traincheck.instrumentor.tracer import suppress_tracing + + previous = config.DISABLE_WRAPPER + config.DISABLE_WRAPPER = True + try: + with suppress_tracing(): + yield + finally: + config.DISABLE_WRAPPER = previous + + def _sample_parameters(self, boundary: Literal["before", "after"]) -> None: + if self.mode != "subclass": + return + from traincheck.instrumentor.proxy_wrapper.proxy_registry import ( + get_global_registry, + ) + from traincheck.instrumentor.tracer import suppress_tracing + + # Keep model dumping enabled while excluding sampling's own Torch APIs. + # An empty per-type value means full metadata in the existing registry. + with suppress_tracing(): + get_global_registry().dump_sample( + dump_loc=f"distributed_probe.update.{boundary}", + dump_config={"torch.nn.Parameter": {}}, + ) + + @contextmanager + def step(self, step: int) -> Iterator[None]: + """Trace an absolute update number; a fresh resume may begin after one.""" + self._require_open() + if not isinstance(step, int) or isinstance(step, bool) or step < 1: + raise ValueError("Trace step must be a positive integer") + if self._in_step: + raise RuntimeError("Trace update scopes cannot overlap") + if self._steps and step != self._steps[-1] + 1: + raise ValueError("Trace update numbers must be consecutive") + self._in_step = True + try: + if self.mode == "plain": + yield + else: + from traincheck.config import config + from traincheck.instrumentor.caches import META_VARS + from traincheck.instrumentor.control import start_step + + META_VARS["step"] = step - 1 + start_step() + try: + self._sample_parameters("before") + yield + self._sample_parameters("after") + finally: + config.DISABLE_WRAPPER = True + # A failed update is not reported as a completed trace scope. + self._steps.append(step) + finally: + self._in_step = False + + def close(self) -> TraceReport: + """Drain main-thread buffers; actual record counts require exited workers. + + The dumper owns its writer threads and closes their files at process exit. + Queue completion does not itself prove filesystem flushing, nor coverage + of buffers created by other application threads. The launcher audits all + trace files only after the worker process exits successfully. + """ + self._require_open() + if self._in_step: + raise RuntimeError("Cannot close tracing inside an update") + if self.mode != "plain": + from traincheck.config import config + from traincheck.instrumentor import dumper + + config.DISABLE_WRAPPER = True + with self.suppress(): + for name in ("api_trace_buffer", "var_trace_buffer"): + # No buffer means this thread emitted no events of that kind. + 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() + self._closed = True + return { + "mode": self.mode, + "instrumented_apis": self.instrumented_apis, + "scan_proxy_in_args": self.mode == "subclass", + "registered_parameter_names": list(self.parameter_names), + "completed_trace_steps": list(self._steps), + "tensor_hash_enabled": self.mode != "plain", + "model_sampling": ( + "explicit before/after update plus subclass assignment observations" + if self.mode == "subclass" + else "disabled" + ), + "optimizer_state_registered": False, + "cross_rank_checks_executed": False, + "trace_directory": str(self.output), + "buffers_drained": True, + "trace_record_counts": None, + } + + +def start_tracing(mode: Mode, output: Path, model: nn.Module) -> TraceSession: + """Install normal full-collection wrappers without source rewriting/env dump.""" + if mode not in ("plain", "api", "subclass"): + raise ValueError("Unknown tracing mode") + output.mkdir(parents=True, exist_ok=False) + if mode == "plain": + return TraceSession(mode, output, 0, ()) + + import torch + + if not isinstance(model, torch.nn.Module): + raise TypeError("Tracked model must be a torch.nn.Module") + os.environ["TRAINCHECK_OUTPUT_DIR"] = str(output.resolve()) + + from traincheck.config import config + from traincheck.instrumentor import tracer + from traincheck.instrumentor.caches import META_VARS + from traincheck.instrumentor.proxy_wrapper.proxy_config import tensor_dump_format + + if META_VARS["step"] != 0 or tracer.modules_or_cls_id_instrumented: + raise RuntimeError("Trace adapter requires a fresh worker process") + config.DISABLE_WRAPPER = True + config.MODEL_TRACKER_STYLE = "subclass" if mode == "subclass" else None + config.INSTRUMENTATION_POLICY = {"interval": 1, "warm_up": 0} + config.INSTR_OPTS = None + config.INSTR_OPTS_FILE = str(output / "instr_opts.json") + tensor_dump_format.update( + dump_tensor_hash=True, dump_tensor_stats=False, dump_tensor_full=False + ) + count = tracer.Instrumentor( + torch, + scan_proxy_in_args=mode == "subclass", + use_full_instr=False, + API_dump_stack_trace=False, + sampling_interval=1, + warm_up_steps=0, + ).instrument() + if count <= 0: + raise RuntimeError("Instrumentor did not wrap any API") + parameter_names: tuple[str, ...] = () + session = TraceSession(mode, output, count, parameter_names) + if mode == "subclass": + from traincheck.instrumentor.proxy_wrapper.subclass import ( + ProxyParameter, + proxy_parameter, + ) + + with session.suppress(): + proxy_parameter(model, parent_name="model", should_dump_trace=False) + named = list(model.named_parameters()) + if not named: + raise ValueError("Model must contain parameters for subclass tracking") + if any(not isinstance(parameter, ProxyParameter) for _, parameter in named): + raise RuntimeError("Some model parameters were not subclass-tracked") + parameter_names = tuple("model." + name for name, _ in named) + session.parameter_names = parameter_names + return session diff --git a/tools/distributed_probe/workload.py b/tools/distributed_probe/workload.py new file mode 100644 index 00000000..5baff6f5 --- /dev/null +++ b/tools/distributed_probe/workload.py @@ -0,0 +1,592 @@ +"""Small DDP contract probe, not an RSI score or performance benchmark. + +Synthetic token batches have unequal valid-token counts on each rank. A separate +global-batch calculation supplies the first-update gradient oracle. Checkpoints +are made only at complete update boundaries. Data is regenerated from absolute +step/rank/microbatch indices, so this probe does not exercise a DataLoader cursor. +""" + +from __future__ import annotations + +import argparse +import copy +import datetime +import hashlib +import json +import os +import sys +import time +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypeAlias + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel + +if __package__ in (None, ""): + sys.path.append(str(Path(__file__).resolve().parents[2])) + +Precision: TypeAlias = Literal["fp32", "bf16"] +Fault: TypeAlias = Literal["none", "local-mean"] +Instrumentation: TypeAlias = Literal["plain", "api", "subclass"] +StateValue: TypeAlias = ( + torch.Tensor + | None + | bool + | int + | float + | str + | list["StateValue"] + | tuple["StateValue", ...] + | dict[str | int, "StateValue"] +) + +SEED = 20261031 +VOCAB_SIZE = 32 +SEQUENCE_LENGTH = 8 +BATCH_SIZE = 2 + + +@dataclass(frozen=True) +class Options: + output: Path + resume: Path | None + steps: int + stop_after: int + precision: Precision + fault: Fault + instrumentation: Instrumentation + device: Literal["cpu", "cuda"] + accumulation: int + + +@dataclass(frozen=True) +class Batch: + tokens: torch.Tensor + labels: torch.Tensor + mask: torch.Tensor + + +class TinyTransformer(nn.Module): + """Dropout-free model so independent batching has an exact mathematical oracle.""" + + def __init__(self) -> None: + super().__init__() + self.embedding = nn.Embedding(VOCAB_SIZE, 16) + layer = nn.TransformerEncoderLayer( + d_model=16, + nhead=2, + dim_feedforward=32, + dropout=0.0, + batch_first=True, + ) + self.encoder = nn.TransformerEncoder(layer, num_layers=2) + self.readout = nn.Linear(16, VOCAB_SIZE) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + result: object = self.readout(self.encoder(self.embedding(tokens))) + if not isinstance(result, torch.Tensor): + raise TypeError("Transformer must return a Tensor") + return result + + +def make_batch(step: int, rank: int, microbatch: int, device: torch.device) -> Batch: + if step < 1 or rank < 0 or microbatch < 0: + raise ValueError( + "step must be positive; rank and microbatch must be nonnegative" + ) + generator = torch.Generator(device="cpu") + generator.manual_seed(SEED + 10000 * step + 100 * rank + microbatch) + tokens = torch.randint( + VOCAB_SIZE, (BATCH_SIZE, SEQUENCE_LENGTH), generator=generator + ) + labels = (tokens + 1) % VOCAB_SIZE + # Ranks 0 and 1 have disjoint low/high valid-token-count ranges. + lengths = torch.tensor( + [2 + 3 * rank + (step + microbatch + row) % 2 for row in range(BATCH_SIZE)] + ).clamp(max=SEQUENCE_LENGTH) + mask = torch.arange(SEQUENCE_LENGTH).unsqueeze(0) < lengths.unsqueeze(1) + return Batch(tokens.to(device), labels.to(device), mask.to(device)) + + +def token_loss(logits: torch.Tensor, batch: Batch) -> torch.Tensor: + losses = F.cross_entropy( + logits.reshape(-1, VOCAB_SIZE), batch.labels.reshape(-1), reduction="none" + ).reshape_as(batch.labels) + return (losses * batch.mask).sum() + + +def loss_scale( + fault: Fault, world_size: int, global_tokens: int, local_tokens: int +) -> float: + if world_size < 1 or local_tokens < 1 or global_tokens < local_tokens: + raise ValueError("Invalid world size or valid-token counts") + if fault == "none": + return world_size / global_tokens + if fault == "local-mean": + return 1.0 / local_tokens + raise ValueError(f"Unknown fault: {fault}") + + +def copy_state(value: object, context: str = "state") -> StateValue: + """Validate and snapshot framework state without dropping values or coercing types.""" + if isinstance(value, torch.Tensor): + snapshot = value.detach().cpu().clone() + if type(snapshot) is not torch.Tensor: + raise TypeError(f"{context} did not detach into a plain Tensor") + return snapshot + if value is None or type(value) in (bool, int, float, str): + # Explicit branches below narrow primitive values for the type checker. + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, list): + return [copy_state(child, context) for child in value] + if isinstance(value, tuple): + return tuple(copy_state(child, context) for child in value) + if isinstance(value, dict): + result: dict[str | int, StateValue] = {} + for key, child in value.items(): + if type(key) not in (str, int) or not isinstance(key, (str, int)): + raise TypeError(f"{context} contains an unsupported dictionary key") + result[key] = copy_state(child, f"{context}.{key}") + return result + raise TypeError(f"{context} contains unsupported {type(value).__name__}") + + +def tensor_mapping(value: object, context: str) -> dict[str, torch.Tensor]: + if not isinstance(value, dict): + raise TypeError(f"{context} must be a dictionary") + result: dict[str, torch.Tensor] = {} + for key, child in value.items(): + if not isinstance(key, str) or not isinstance(child, torch.Tensor): + raise TypeError(f"{context} must map names to tensors") + result[key] = child + return result + + +def string_mapping(value: object, context: str) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError(f"{context} must be a dictionary") + result: dict[str, object] = {} + for key, child in value.items(): + if not isinstance(key, str): + raise TypeError(f"{context} must have string keys") + result[key] = child + return result + + +def model_gradients(model: nn.Module) -> dict[str, torch.Tensor]: + gradients: dict[str, torch.Tensor] = {} + for name, parameter in model.named_parameters(): + if parameter.grad is None: + raise ValueError(f"Missing gradient for {name}") + value = copy_state(parameter.grad, name) + assert isinstance(value, torch.Tensor) + gradients[name] = value + return gradients + + +def tensor_hash(values: dict[str, torch.Tensor]) -> str: + digest = hashlib.sha256() + for name, value in values.items(): + digest.update(name.encode()) + digest.update(str(value.dtype).encode()) + digest.update(str(tuple(value.shape)).encode()) + digest.update( + value.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + ) + return digest.hexdigest() + + +def gradient_reference( + model: nn.Module, + world_size: int, + accumulation: int, + device: torch.device, + precision: Precision, +) -> dict[str, torch.Tensor]: + reference = copy.deepcopy(model) + batches = [ + make_batch(1, rank, microbatch, device) + for rank in range(world_size) + for microbatch in range(accumulation) + ] + combined = Batch( + torch.cat([batch.tokens for batch in batches]), + torch.cat([batch.labels for batch in batches]), + torch.cat([batch.mask for batch in batches]), + ) + with torch.autocast( + device_type=device.type, dtype=torch.bfloat16, enabled=precision == "bf16" + ): + logits: object = reference(combined.tokens) + if not isinstance(logits, torch.Tensor): + raise TypeError("Reference output must be a tensor") + loss = token_loss(logits, combined) / combined.mask.sum() + loss.backward() + return model_gradients(reference) + + +def compare_gradients( + actual: dict[str, torch.Tensor], + expected: dict[str, torch.Tensor], + precision: Precision, +) -> dict[str, bool | float | int]: + if actual.keys() != expected.keys(): + raise ValueError("Oracle parameter names differ") + atol, rtol = (2e-6, 2e-4) if precision == "fp32" else (3e-3, 3e-2) + maximum = 0.0 + passed = True + for name, value in actual.items(): + reference = expected[name] + if value.shape != reference.shape or value.dtype != reference.dtype: + raise ValueError(f"Oracle metadata differs for {name}") + if not bool(torch.isfinite(value).all()) or not bool( + torch.isfinite(reference).all() + ): + raise ValueError(f"Nonfinite oracle gradient for {name}") + maximum = max(maximum, float((value - reference).abs().max().item())) + passed = passed and bool(torch.allclose(value, reference, atol=atol, rtol=rtol)) + return { + "passed": passed, + "max_abs_difference": maximum, + "atol": atol, + "rtol": rtol, + "parameter_tensors": len(actual), + } + + +def checkpoint_payload( + model: nn.Module, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.StepLR, + completed: int, + rank: int, + world_size: int, + options: Options, + device: torch.device, +) -> dict[str, StateValue]: + return { + "format_version": 1, + "model": copy_state(model.state_dict()), + "gradients": copy_state(model_gradients(model)), + "optimizer": copy_state(optimizer.state_dict()), + "scheduler": copy_state(scheduler.state_dict()), + "cpu_rng": torch.get_rng_state().clone(), + "cuda_rng": ( + torch.cuda.get_rng_state(device).cpu() if device.type == "cuda" else None + ), + "completed_steps": completed, + "rank": rank, + "world_size": world_size, + "total_steps": options.steps, + "precision": options.precision, + "fault": options.fault, + "accumulation": options.accumulation, + "data_seed": SEED, + } + + +def load_checkpoint( + path: Path, + model: nn.Module, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.StepLR, + rank: int, + world_size: int, + options: Options, + device: torch.device, +) -> int: + raw: object = torch.load(path, map_location="cpu", weights_only=True) + payload = string_mapping(raw, "checkpoint") + expected: dict[str, int | str] = { + "format_version": 1, + "rank": rank, + "world_size": world_size, + "total_steps": options.steps, + "precision": options.precision, + "fault": options.fault, + "accumulation": options.accumulation, + "data_seed": SEED, + } + for key, expected_value in expected.items(): + if ( + key not in payload + or not isinstance(payload[key], type(expected_value)) + or isinstance(payload[key], bool) + or payload[key] != expected_value + ): + raise ValueError(f"Checkpoint field {key} does not match this run") + completed = payload["completed_steps"] + if ( + not isinstance(completed, int) + or isinstance(completed, bool) + or not 1 <= completed < options.stop_after + ): + raise ValueError( + "Checkpoint completed_steps must precede this run's stop boundary" + ) + # Validate the complete tree before handing the library its owned structures. + copy_state(payload) + model.load_state_dict(tensor_mapping(payload["model"], "model"), strict=True) + optimizer.load_state_dict(string_mapping(payload["optimizer"], "optimizer")) + scheduler.load_state_dict(string_mapping(payload["scheduler"], "scheduler")) + cpu_rng = payload["cpu_rng"] + if ( + not isinstance(cpu_rng, torch.Tensor) + or cpu_rng.dtype != torch.uint8 + or cpu_rng.ndim != 1 + ): + raise TypeError("cpu_rng must be a one-dimensional uint8 tensor") + torch.set_rng_state(cpu_rng) + cuda_rng = payload["cuda_rng"] + if device.type == "cuda": + if ( + not isinstance(cuda_rng, torch.Tensor) + or cuda_rng.dtype != torch.uint8 + or cuda_rng.ndim != 1 + ): + raise TypeError("cuda_rng must be a one-dimensional uint8 tensor") + torch.cuda.set_rng_state(cuda_rng, device) + elif cuda_rng is not None: + raise ValueError("CPU checkpoint must not have CUDA RNG state") + return completed + + +def parse_options(argv: list[str] | None = None) -> Options: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--resume", type=Path) + parser.add_argument("--steps", type=int, default=4) + parser.add_argument("--stop-after", type=int) + parser.add_argument("--accumulation", type=int, default=2) + parser.add_argument("--precision", choices=("fp32", "bf16"), default="fp32") + parser.add_argument("--fault", choices=("none", "local-mean"), default="none") + parser.add_argument( + "--instrumentation", choices=("plain", "api", "subclass"), default="plain" + ) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + args = parser.parse_args(argv) + stop_after = args.steps if args.stop_after is None else args.stop_after + if not 1 <= stop_after <= args.steps or not 1 <= args.accumulation <= 4: + parser.error("Require 1 <= stop-after <= steps and 1 <= accumulation <= 4") + return Options( + args.output, + args.resume, + args.steps, + stop_after, + args.precision, + args.fault, + args.instrumentation, + args.device, + args.accumulation, + ) + + +def run(options: Options) -> None: + from tools.distributed_probe.tracing import start_tracing + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if rank < 0 or world_size < 1 or rank >= world_size or local_rank < 0: + raise ValueError("Invalid torchrun rank metadata") + device = ( + torch.device("cuda", local_rank) + if options.device == "cuda" + else torch.device("cpu") + ) + torch.set_num_threads(1) + if device.type == "cuda": + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.use_deterministic_algorithms(True) + torch.manual_seed(SEED) + if world_size > 1: + dist.init_process_group( + "nccl" if device.type == "cuda" else "gloo", + timeout=datetime.timedelta(seconds=60), + ) + rank_output = options.output / f"rank{rank}" + rank_output.mkdir(parents=True, exist_ok=False) + model = TinyTransformer().to(device).train() + reference = gradient_reference( + model, world_size, options.accumulation, device, options.precision + ) + session = start_tracing(options.instrumentation, rank_output / "trace", model) + try: + with session.suppress(): + wrapped: nn.Module = model + if world_size > 1: + wrapped = DistributedDataParallel( + model, device_ids=[local_rank] if device.type == "cuda" else None + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, foreach=False) + scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=2, gamma=0.8 + ) + completed = 0 + if options.resume is not None: + completed = load_checkpoint( + options.resume / f"rank{rank}" / "checkpoint.pt", + model, + optimizer, + scheduler, + rank, + world_size, + options, + device, + ) + initial = tensor_mapping(copy_state(model.state_dict()), "initial model") + start_step = completed + 1 + batches = { + step: [ + make_batch(step, rank, microbatch, device) + for microbatch in range(options.accumulation) + ] + for step in range(start_step, options.stop_after + 1) + } + rows: list[dict[str, object]] = [] + oracle: dict[str, bool | float | int] | None = None + started = time.perf_counter() + for step in range(start_step, options.stop_after + 1): + local_batches = batches[step] + with session.suppress(): + local_tokens = sum( + int(batch.mask.sum().item()) for batch in local_batches + ) + count_tensor = torch.tensor( + local_tokens, device=device, dtype=torch.int64 + ) + if world_size > 1: + dist.all_reduce(count_tensor) + global_tokens = int(count_tensor.item()) + scale = loss_scale( + options.fault, world_size, global_tokens, local_tokens + ) + losses: list[torch.Tensor] = [] + with session.step(step): + optimizer.zero_grad(set_to_none=True) + for microbatch, batch in enumerate(local_batches): + context = ( + wrapped.no_sync() + if isinstance(wrapped, DistributedDataParallel) + and microbatch + 1 < len(local_batches) + else nullcontext() + ) + with context: + with torch.autocast( + device_type=device.type, + dtype=torch.bfloat16, + enabled=options.precision == "bf16", + ): + logits: object = wrapped(batch.tokens) + if not isinstance(logits, torch.Tensor): + raise TypeError("Training output must be a Tensor") + loss = token_loss(logits, batch) * scale + loss.backward() + losses.append(loss.detach()) + if step == 1: + with session.suppress(): + actual = model_gradients(model) + oracle = compare_gradients(actual, reference, options.precision) + torch.save( + {"actual": actual, "expected": reference}, + rank_output / "oracle.pt", + ) + gradient_norm = nn.utils.clip_grad_norm_( + model.parameters(), max_norm=1.0 + ) + optimizer.step() + scheduler.step() + with session.suppress(): + rows.append( + { + "step": step, + "local_valid_tokens": local_tokens, + "global_valid_tokens": global_tokens, + "scaled_local_loss": float(torch.stack(losses).sum().item()), + "gradient_norm_before_clip": float(gradient_norm.item()), + "learning_rate": scheduler.get_last_lr()[0], + } + ) + with session.suppress(): + if device.type == "cuda": + torch.cuda.synchronize(device) + training_seconds = time.perf_counter() - started + final_state = tensor_mapping(copy_state(model.state_dict()), "final model") + rank_maximum = 0.0 + if world_size > 1: + for parameter in model.parameters(): + if not bool(torch.isfinite(parameter).all()): + raise ValueError("Nonfinite final model parameter") + peer = parameter.detach().clone() + dist.broadcast(peer, src=0) + rank_maximum = max( + rank_maximum, float((peer - parameter).abs().max().item()) + ) + payload = checkpoint_payload( + model, + optimizer, + scheduler, + options.stop_after, + rank, + world_size, + options, + device, + ) + torch.save(payload, rank_output / "checkpoint.pt") + metrics: dict[str, object] = { + "rank": rank, + "world_size": world_size, + "device": str(device), + "torch_version": torch.__version__, + "precision": options.precision, + "fault": options.fault, + "instrumentation": options.instrumentation, + "start_step": start_step, + "completed_steps": options.stop_after, + "total_steps": options.steps, + "accumulation": options.accumulation, + "parameter_count": sum(p.numel() for p in model.parameters()), + "initial_model_sha256": tensor_hash(initial), + "final_model_sha256": tensor_hash(final_state), + "rank_parameter_max_abs_difference_from_rank0": rank_maximum, + "gradient_oracle": oracle, + "training_seconds_including_oracle_snapshot": training_seconds, + "steps": rows, + "limitations": [ + "dropout=0; saved RNG is not exercised by stochastic model operations", + "synthetic deterministic absolute-step batches; no DataLoader cursor", + "BF16 autocast without GradScaler; no FP16 overflow coverage", + "update-boundary checkpoints only", + "saved gradients are comparison evidence; resumed updates zero gradients before use", + "all accumulation windows are full; partial-window behavior is not exercised", + "no inferred invariant or cross-rank TrainCheck checker in this workload", + ], + } + (rank_output / "metrics.json").write_text( + json.dumps(metrics, indent=2, allow_nan=False) + "\n" + ) + finally: + try: + report = session.close() + (rank_output / "trace-session.json").write_text( + json.dumps(report, indent=2, allow_nan=False) + "\n" + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def main() -> None: + run(parse_options()) + + +if __name__ == "__main__": + main()