docs/configuration.md

Configuration guide

The CLI composes the installed gnn4colliders.configs/config.yaml with semantic Hydra groups. A config describes experiment intent; it does not contain arbitrary Python module or class import paths.

Groups

Group	Purpose
data	ROOT source or graph-cache path, batch settings, and fold splits
model	Model family and supported ROOT-GNN constructor settings
task	Binary or multiclass loss, score, and metric semantics
trainer	Device, seed, epochs, optimizer, scheduler, and early stopping
checkpoint	Output directory, resume checkpoint, or pretrained checkpoint
inference	Checkpoint, split, output path, and output format
environment	Device and experiment output root
distributed	Single-process/DDP selection and process-group backend

The active model groups are root_gnn/edge_network and root_gnn/fine_tuned_edge_network. The active task groups are pretraining_multiclass, binary_classification, and tth_cp_finetune.

Preparation

prepare requires data.files, data.cache.path, data.feature_branches, data.object_types, and data.scales. feature_branches follows the shared seven-column feature contract: one branch/constant specification per output column and one entry per configured object type. CALC_E and NODE_TYPE are reserved derived specifications. object_types entries are vector or single. Preparation reads the configured tree in file order and writes a versioned GraphSampleCache.

Example overrides are easiest to maintain in a YAML file for real datasets:

# project-local example: data/my_events.yaml
files: [data/events.root]
tree_name: Events
cache:
  path: cache/events.pt
feature_branches:
  - [jet_pt]
  - [jet_eta]
  - [jet_phi]
  - CALC_E
  - [1.0]
  - [0.0]
  - NODE_TYPE
object_types: [vector]
scales: [1, 1, 1, 1, 1, 1, 1]
fold_var: eventNumber
weight_var: weight

Then compose it with data=my_events. The cache stores processed graph samples, labels, globals, named metadata, and feature/graph/cache schema versions. Changing the feature or graph schema requires a new compatible cache; loading a mismatched schema raises an error.

Common overrides

uv run gnn4colliders train \
  data.cache.path=cache/events.pt \
  data.batch_size=64 \
  trainer.max_epochs=50 \
  trainer.seed=123 \
  environment.output_root=outputs/my_run

data.batch_size is per process. data.splits.train_folds, validation_folds, and test_folds define conventional train/validation/test selection and must be disjoint. Model/task mismatches are rejected during config validation; binary tasks require model.out_size=1, while multiclass tasks require model.out_size=task.num_classes.

Hugging Face ROOT data

Training preparation can resolve ROOT files directly from a pinned Hugging Face dataset. Use data.source instead of local data.files:

source:
  type: huggingface
  repo_id: HWresearch/Delphes
  repo_type: dataset
  revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e
  files:
    - path: testing/ttH_NLO_64.root
      label: 0
      sha256: 89d69eae9cd4d28a5d414d77bbc07185ea5b7de03481fdafe28899e2f3a09dec
      size_bytes: 15050

The resolver uses the Hugging Face cache, downloads only configured files, and verifies optional size and SHA-256 metadata before ROOT ingestion. Ready-to-use configurations are config_hf_smoke for the 64-event fixture and config_hf_delphes for the pinned 12-process pretraining sources:

uv run gnn4colliders prepare --config-name config_hf_smoke
uv run gnn4colliders prepare --config-name config_hf_delphes

The full Delphes configuration may download a large dataset. It is never used by the default configuration and is cached by the Hub client for later runs.

data.num_workers controls local process parallelism during prepare. The default is 0 (single-process). Positive values split the ordered event range into independent temporary graph-cache shards and merge them deterministically; workers do not write the final cache concurrently. This setting does not launch distributed preparation and should be benchmarked on the target machine.

Checkpoints and resolved configuration

Training writes epoch_####.pt and the fully resolved configuration at <environment.output_root>/resolved_config.yaml. A checkpoint includes schema version, model weights/config, task config, trainer/optimizer/scheduler state, early stopping state, metadata, and optional RNG state. Set checkpoint.resume=/path/to/epoch_####.pt to continue a run. Set checkpoint.pretrained=/path/to/epoch_####.pt with the fine-tuned model group to load weights into a new task head; these options are mutually exclusive.

Environment profiles

environment=local selects CPU by default. environment=macos selects the Apple MPS device for the portable PyTorch ROOT-GNN backend. environment=perlmutter selects the CUDA device and a conventional output-root pattern. Profiles should hold device/output policy only; site-specific module loads and filesystem paths belong in a launcher or shell environment.