CLI Reference: Collect Traces

Start with Use TrainCheck if you want the full workflow. This page explains the traincheck-collect command.

traincheck-collect instruments a PyTorch training script and writes trace files. Use it for two jobs:

• Full reference collection for invariant inference.
• Selective target collection for checking with an existing invariant file.

Full Reference Collection

Use full collection on a known-good run:

traincheck-collect \
  --pyscript train.py \
  --models-to-track model \
  --output-dir reference_trace

This command runs train.py, tracks the Python variable named model, and writes trace files into reference_trace/.

Selective Target Collection

Use selective collection when you already have an invariant file:

traincheck-collect \
  --pyscript train.py \
  --models-to-track model \
  --invariants invariants.json \
  --output-dir target_trace

--invariants tells TrainCheck which APIs and variables matter for checking. This usually reduces target-run overhead compared with full reference collection.

Do not combine --invariants with --use-full-instr when you want selective collection.

Step Sampling

Sampling is also configured on traincheck-collect:

traincheck-collect \
  --pyscript train.py \
  --models-to-track model \
  --invariants invariants.json \
  --sampling-interval 10 \
  --warm-up-steps 10 \
  --output-dir target_trace

This traces the warm-up steps, then traces every tenth step. Use sampling for long target runs after you have confirmed TrainCheck works on a short run.

Config Files

Use --use-config when the collection command needs repeated options:

traincheck-collect --use-config --config traincheck.yml

Example:

pyscript: ./train.py
shscript: ./run.sh
modules_to_instr:
  - torch
models_to_track:
  - model
model_tracker_style: subclass
copy_all_files: false
output_dir: traincheck_trace

Config keys use underscores, not hyphens. For example, the CLI flag --output-dir becomes output_dir in YAML.

Useful Options

• --pyscript: Python entry point for the training program.
• --shscript: shell script used to launch the Python program.
• --models-to-track: model variable names to track.
• --modules-to-instr: Python modules to instrument, usually torch.
• --invariants: invariant files for selective collection.
• --output-dir: directory for traces and logs.
• --sampling-interval: collect every Nth step after warm-up.
• --warm-up-steps: collect the first N steps.
• --copy-all-files: copy files beside the training script into the output directory.
• --model-tracker-style: choose subclass (default), proxy, or sampler.

Choosing a model tracker

The default subclass tracker keeps the model as a PyTorch nn.Module and tracks its parameters using a torch.nn.Parameter subclass. It requires --models-to-track (or models_to_track in YAML); choosing a tracker alone does not enable model tracking. Model tracking must be installed before creating the optimizer, so the optimizer uses the tracked parameters.

The proxy tracker remains available through --model-tracker-style proxy or model_tracker_style: proxy in YAML. It wraps model objects and dynamically accessed attributes, so its coverage differs from the parameter-focused subclass tracker. Explicit CLI/YAML selections remain supported, and saved instrumentation options are not rewritten.

For invariants inferred from proxy traces, keep proxy explicit when collecting the checking trace. To migrate, collect a new reference trace with subclass, infer new invariants, and verify that the intended rules actually activate during checking. Zero violations alone does not establish coverage; module or buffer rules from a proxy trace may not activate with parameter-only tracking.

Selective collection keeps its existing policy: invariants requiring immediate variable tracking use the selected subclass or proxy tracker. When only cross-variable consistency rules require model tracking, collection uses the sampler tracker to observe parameters after optimizer steps. If no invariant requires model tracking, selective collection disables it.

Run the command help for the complete option list:

traincheck-collect --help

Output Files

The output directory contains trace files, environment metadata, logs, and the instrumented training script. The checker accepts the full output directory through -f or --trace-folders.

traincheck-check -f target_trace -i invariants.json