Train a model

This guide trains the standard CLD track-and-cluster model on configuration 1, version 3.3.0 of the ttbar, WW, and qq datasets. The catalog lists compatible recipes and published versions. Start with the two-step CPU run to check the complete path. Use the GPU example only after that check succeeds.

A training step updates the model once using one dataset-group batch. A validation cycle pauses training to measure loss on a separate set of held-out events. A checkpoint is a saved model and training state from a particular step.

Prerequisites¶

• Complete the main installation.

• Download and verify configuration 1, version 3.3.0 of the three CLD training datasets. You can also produce matching data using Key4HEP production.

• The commands below use the download guide’s data root, data/tfds/tensorflow_datasets/cld. For locally produced data, replace it with the directory containing the three dataset-name directories (typically <workspace>/tfds). For a small integration check without downloading the full partition, use the quickstart.

• Run commands from the repository root.

The standard pyg-cld-v1 recipe selects these three event types from configurations 1 through 10. --data_config 1 restricts all three to the downloaded configuration.

The configuration reference explains recipe selection, command-line overrides, and how to inspect the resolved settings before training.

Step 1: run a CPU software check¶

uv run python3 mlpf/pipeline.py \
  --spec-file particleflow_spec.yaml \
  --model-name pyg-cld-v1 \
  --production-name cld \
  --data-dir data/tfds/tensorflow_datasets/cld \
  --prefix docs-cld-cpu- \
  train \
  --data_config 1 \
  --gpus 0 \
  --dtype float32 \
  --model.attention.attention_type math \
  --gpu_batch_multiplier 1 \
  --num_workers 1 \
  --prefetch_factor 1 \
  --ntrain 10 \
  --nvalid 10 \
  --ntest 10 \
  --num_steps 2 \
  --val_freq 2 \
  --checkpoint_freq 2

This run checks the software integration. A useful trained model requires a full training and validation campaign. CPU time depends strongly on the machine. The run should end with a VALIDATION line and create an experiment directory beginning with experiments/docs-cld-cpu-.

The essential options are:

All other values come from three layers, in order: typed defaults in mlpf/conf.py, the selected recipes in particleflow_spec.yaml, and command-line overrides. The program logs the final flattened configuration before it starts.

Step 2: inspect the result¶

Find the newest matching directory:

ls -dt experiments/docs-cld-cpu-* | head -n 1

A successful two-step run contains at least:

experiments/docs-cld-cpu-.../
├── checkpoints/checkpoint-02.pth
├── history/step_2.json
├── model_kwargs.pkl
├── hyperparameters.json
├── particleflow_spec.yaml
├── train-config.yaml
└── train.log

It may also contain TensorBoard data, predictions, and plots. See Artifacts and directories before moving or publishing a checkpoint.

Check that the history contains finite training and validation losses:

uv run python3 - <<'PY'
import json
from pathlib import Path

experiment = sorted(Path("experiments").glob("docs-cld-cpu-*"), key=lambda p: p.stat().st_mtime)[-1]
history = json.loads((experiment / "history" / "step_2.json").read_text())
print("train total:", history["train"]["Total"])
print("valid total:", history["valid"]["Total"])
PY

Finite numbers show that data loading, forward and backward passes, and validation completed. Model convergence requires a longer loss history and held-out performance checks.

Step 3: start a GPU training¶

The following bounded engineering example reads configuration 1 and stops after 2,000 steps. Publication reproduction uses the archived recipe and artifacts associated with that publication.

uv run python3 mlpf/pipeline.py \
  --spec-file particleflow_spec.yaml \
  --model-name pyg-cld-v1 \
  --production-name cld \
  --data-dir data/tfds/tensorflow_datasets/cld \
  --prefix cld-ttbar-config1- \
  train \
  --data_config 1 \
  --gpus 1 \
  --dtype bfloat16 \
  --gpu_batch_multiplier 4 \
  --sampler_mode interleaved-shards \
  --num_workers 2 \
  --prefetch_factor 2 \
  --num_steps 2000 \
  --val_freq 1000 \
  --checkpoint_freq 1000

Memory use depends on event size, attention implementation, precision, and per-rank batch size. Lower --gpu_batch_multiplier first if the process runs out of GPU memory. Assess physics quality with the held-out checks in the validation overview.

Attention and padding¶

Native Flash attention automatically packs valid elements into a jagged representation so padding does not enter attention. It requires a compatible GPU and a head dimension divisible by 8, between 8 and 256. Use --model.attention.attention_type math for CPU runs; dense math attention applies a key-padding mask. model.attention.use_flash_attn_varlen optionally selects the external FlashAttention variable-length kernel and requires the matching package.

Keep the configured padding multiple for HEPTv2 and other block-based models. HEPTv2 uses deterministic integer hashing by default (model.heptv2.deterministic_hashing: true) and stable tie ordering. Its coordinates and bucket assignments can differ from older floating-point hashing, so validate older checkpoints before resuming or comparing predictions.

For sample ordering, weighted groups, and per-group losses, see Dataset mixing. The ONNX guide describes the separate export contract and runtime profiles.

Experimental CLD hit-set model¶

The recommended global-attention preset is pyg-cld-hits-heptv2-set-v1, with supervised hard hit grouping and tracker-neighborhood features. Use the cld_pf_heptv2_default_50k scenario described in the scenario guide. The sectorized variant below is experimental and is not the recommended default.

The pyg-cld-hits-heptv2-sector-set-v1 recipe uses a HEPTv2 hit backbone and an input-conditioned particle-set decoder. Decoder slots are assigned to periodic azimuthal sectors from their current reference direction. Each slot cross-attends only to hits in its sector and sector_neighbors sectors on either side; slot self-attention is also sector-local. This avoids dense slot-by-hit masks and global slot self-attention. A sector with no nearby hits uses its nearest hit so attention is always defined.

The recipe starts with 32 sectors, one neighboring sector on each side, and 256 slots. These settings can be overridden with model.set_decoder.num_sectors, model.set_decoder.sector_neighbors, and model.set_decoder.num_slots. In this mode, local_attention_radius limits the reference update per decoder layer; it is not an exact geometric cross-attention cutoff. The model requires hit counts padded to a multiple of the HEPTv2 block size (128 in this recipe).

HEPTv2 GPU bucket attention requires compiled torch.nn.attention.flex_attention and never falls back to materializing attention scores. Bucket counts use a dynamic batch dimension so varying padded event lengths do not create a separate compiled kernel for every number of buckets. Failures retain the original exception and report the input shape, dtype, and PyTorch/ROCm versions; the explicit score calculation is used only on CPU.

On an allocated GPU with the training environment activated, check fused attention and checkpoint replay before a campaign:

python -m pytest -q tests/test_heptv2.py tests/test_heptv2_checkpointing.py

These tests check fused outputs and gradients, including log-normalizers used to weight hashes, and checkpoint replay. GPU cases skip when no GPU is visible. They do not measure full training memory or distributed throughput. For a LUMI environment and RCCL check, submit scripts/lumi/smoke_test.sh; for training memory measurements, see HEPTv2 checkpointing.

Train CMS or CLIC¶

The command shape is the same, but the production, model recipe, data root, dataset version, and locally available samples must agree:

Use the dataset catalog to pair names and versions. Hit inputs require the dedicated *-hits-v1 recipes and substantially more memory.

Next steps¶

• Continue or fine-tune from a checkpoint.

• Run the test command and detector-level plots described in Key4HEP validation.

• Compare PyTorch and ONNX execution with ONNX validation.