TACE Scripts Tutorial

TACE installs command-line scripts for training, inference, data preparation, model conversion, and deployment. Most scripts use argparse and support -h directly:

tace-eval -h
tace-export-eval -h

Important

Some scripts, including tace-train and tace-graph, are launched through Hydra. Hydra resolves the configuration before running the command, so the corresponding YAML file must be available even when only requesting -h. Run these commands from the configuration directory and select the YAML file with -cn when needed:

cd example/train
tace-train -cn tace.yaml -h
tace-graph -cn tace.yaml -h

Script Overview

Command	Purpose
tace-train	Train, validate, or test from a Hydra configuration
tace-eval	Run inference on an ASE-readable structure file
tace-export-train	Export an editable model for training or transfer
tace-export-eval	Export state-dict, full-model, or AOTI inference artifacts
tace-export-lammps	Export eager or AOTI-backed LAMMPS ML-IAP artifacts
tace-compile	Alias for tace-export-eval
tace-graph	Pre-build graph data using a training configuration
tace-split	Split an ASE-readable dataset into train/validation/test files
tace-finetune	Generate a fine-tuning and LoRA configuration
tace-convert	Convert model parameters, currently including LoRA merging
tace-update	Update model statistics such as atomic energies, scale, and shift
tace-average	Average parameters from models with identical architectures
tace-copy	Copy shape-compatible parameters between two models
tace-modify	Inspect or override model properties in memory
tace-download	List or download pretrained TACE models
tace-clean	Remove standard training outputs from the current directory

Training

tace-train uses Hydra. The -cn argument selects a YAML configuration without its extension:

tace-train -cn tace

If the configuration is named tace.yaml in the configured search path, the default command is sufficient:

tace-train

Hydra overrides can be appended on the command line. This is useful for short experiments without editing the source YAML:

tace-train -cn tace trainer.max_epochs=10 dataset.batch_size=8

Training checkpoints contain optimizer, scheduler, callback, and model state, so keep .ckpt files when a run may need to be resumed. For resume settings, see resume_from_model.

Inference

tace-eval reads structures through ase.io.read and writes predictions through ase.io.write. At minimum, provide the input and model paths:

tace-eval \
  -i structures.xyz \
  -m model.ckpt \
  -o predict.xyz \
  --device cuda \
  -b 16

Supported model inputs include checkpoints, state-dict packages, serialized full models, and compatible .pt2 AOTI packages. Use -t 1 to report test metrics when references are present and -e 0 to disable EMA checkpoint parameters.

Property names in a dataset can be remapped with arguments such as --energy_key, --forces_key, and the other key options shown by tace-eval -h.

Model Export

Export commands, formats, AOTI requirements, output names, and deployment examples are documented in the dedicated TACE Export Tutorial.

The three primary entry points are:

tace-export-train -m model.ckpt
tace-export-eval -m model.ckpt --backend state_dict
tace-export-lammps -m model.pt --backend mliap

Use tace-export-eval --backend aoti for native ASE/TorchSim deployment and tace-export-lammps --backend aoti for compiled LAMMPS inference. AOTI requires PyTorch 2.11 or newer. Use -f or --fidelity_idx to select a fidelity during export.

Dataset Preparation

tace-split

Split a structure file into train, validation, and test sets by specifying the number of configurations in each output:

tace-split -i dataset.xyz -n 800 100 100 -s 42

For dataset.xyz, the command writes dataset_train.xyz, dataset_valid.xyz, and dataset_test.xyz. Each selected structure keeps its original position in atoms.info["tace_index"]. The requested counts must not exceed the number of input structures.

tace-graph

Pre-building graphs is useful for very large datasets or datasets reused by multiple training runs:

tace-graph -cn tace

The command uses the dataset, element mapping, cutoff, neighbor-list backend, and other graph settings from the training configuration, then exits after graph construction. The input and graph-related settings must exactly match the later training run. When LMDB storage and the configured shard directories already contain graphs, tace-train reads them and skips reconstruction.

Fine-Tuning and Model Conversion

tace-finetune

Generate a fine-tuning template from an existing model:

tace-finetune -m foundation.pt

This writes finetune_config.yaml in the current directory. By default the base parameters are frozen and LoRA adapters are enabled. Review the generated file before launching tace-train.

tace-convert

Merge trained LoRA parameters into the base model for inference or export:

tace-convert -m lora-model.pt -t merge_lora

The output is lora-model.pt-merged_lora.pt. The merged model no longer requires separate LoRA adapter parameters.

tace-copy

Copy parameters shared by a source and destination architecture:

tace-copy -m source.pt destination.pt

Only parameters with matching names and shapes are copied. Destination-only or shape-incompatible parameters are preserved, and the result is written to destination.pt-merged.pt.

Model Statistics and Averaging

tace-update

Update stored statistics using statistics_<fidelity_idx>.yaml files in the current directory:

tace-update \
  -m model.pt \
  -u atomic_energy scale shift

Available update fields are atomic_energy, scale, and shift. The result is written to new_statistics.pt. Verify that every statistics file uses the same element order and fidelity indexing as the model.

tace-average

Average parameters from checkpoints or exported models with identical architectures:

tace-average -m epoch-90.ckpt epoch-95.ckpt epoch-100.ckpt -e 0

The command writes average_model-state.pt. This is a manual stochastic weight averaging workflow. EMA is disabled by default and should generally remain disabled when averaging multiple time-adjacent checkpoints.

Inspection and Maintenance

tace-modify

Load a model and override its requested inference properties in memory:

tace-modify -m model.pt -t energy forces stress

This command currently validates and prints the resulting property selection; it does not save a new model. Atomic-number reduction is reserved by the CLI but is not implemented yet.

tace-download

List registered pretrained models or download one into the TACE cache:

tace-download --list
tace-download -m TACE-OAM-L

Omitting -m requests all registered models. Manual download may be more reliable on machines with restricted network access.

tace-clean

Remove standard run outputs from the current working directory:

tace-clean

Targets include Hydra outputs, Lightning/W&B logs, checkpoint directories, index files, and generated configuration files. The command intentionally refuses cleanup when a target exceeds its size limit. Run it only from the training directory whose generated outputs should be removed.