Implement set prediction

configs/training/README.md

@@ -0,0 +1,72 @@
+# Reusable training scenarios
+
+Scientific comparisons live under `scenarios/`; machine-dependent paths and runtime
+tuning live under `platforms/`. Run a scenario locally with:
+
+```bash
+uv run python3 scripts/training/run_scenario.py \
+  --scenario configs/training/scenarios/cld_hits_output_comparison.yaml \
+  --platform configs/training/platforms/local.yaml \
+  --global-batch-size 8 \
+  --dry-run
+```
+
+The scenario declares a global batch size. The runner derives
+`gpu_batch_multiplier` from the number of GPUs and the dataset batch size, and rejects
+non-integral combinations. It also resolves every variant through `MLPFConfig` and
+checks that variants differ only in the fields listed by
+`allowed_variant_differences`.
+
+The local picker discovers the same scenario files and applies the local platform
+profile and short-run defaults:
+
+```bash
+scripts/local/train_scenario.sh --list
+scripts/local/train_scenario.sh cld_hits_output_comparison --dry-run
+scripts/local/train_scenario.sh cld_hits_output_comparison --seed 2468
+```
+
+Additional arguments are forwarded to the generic scenario runner, such as
+`--variant set`, `--global-batch-size 4`, or `--set num_steps=100`.
+
+With multiple variants or seeds, jobs are ordered by seed and then by variant. A
+Slurm array can select one job using `--task-index $SLURM_ARRAY_TASK_ID`.
+`--seed N` replaces the scenario seed list, including when a task index is used.
+The site shell launchers expose this as the `SEED` environment variable. Without an
+override, seeds come from the scenario file and are recorded in both the resolved
+configuration and run manifest.
+
+List the available scenarios and accelerators, then submit using the Flatiron
+picker:
+
+```bash
+scripts/flatiron/train_scenario.sh --list
+scripts/flatiron/train_scenario.sh cld_hits_output_comparison h100 --dry-run
+scripts/flatiron/train_scenario.sh cld_hits_output_comparison h100
+```
+
+The picker reads Slurm resources from the selected platform profile and derives
+the array size from the scenario's variants and seeds. Use `--seed N` to submit
+one comparison pair with an explicit seed.
+
+Tallinn and LUMI use the same interface with site-specific profiles and workers:
+
+```bash
+scripts/tallinn/train_scenario.sh --list
+scripts/tallinn/train_scenario.sh cld_hits_output_comparison l40 --dry-run
+scripts/tallinn/train_scenario.sh cld_hits_output_comparison l40
+
+scripts/lumi/train_scenario.sh --list
+scripts/lumi/train_scenario.sh cld_hits_output_comparison mi250x --dry-run
+scripts/lumi/train_scenario.sh cld_hits_output_comparison mi250x
+```
+
+The Tallinn worker runs the repository's `uv` environment directly. The LUMI
+submitter uses `particleflow-env` (override its interpreter with
+`PYTHON_EXECUTABLE`) and the worker executes that environment in the standard
+PyTorch ROCm container (override the image with `IMG`).
+
+Use repeated `--set KEY=VALUE` options only for explicit one-off overrides. Every
+resolved run writes `scenario-manifest.json` containing the scenario, platform,
+seed, final configuration, command, and git revision. Runs are grouped as
+`<experiments_dir>/<scenario_name>/<variant>_seed<seed>_<timestamp>/`.

configs/training/platforms/flatiron_a100.yaml

@@ -0,0 +1,18 @@
+name: flatiron_a100
+gpus: 4
+data_dir:
+  cld: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds
+  clic: /mnt/ceph/users/${USER}/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds
+experiments_dir: /mnt/home/${USER}/particleflow/experiments
+runtime_overrides:
+  dtype: bfloat16
+  num_workers: 8
+  prefetch_factor: 2
+  model.attention.use_flash_attn_varlen: false
+slurm:
+  partition: gpu
+  constraint: a100
+  time: "12:00:00"
+  nodes: 1
+  tasks_per_node: 1
+  cpus_per_task: 64

configs/training/platforms/flatiron_h100.yaml

@@ -0,0 +1,18 @@
+name: flatiron_h100
+gpus: 8
+data_dir:
+  cld: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds
+  clic: /mnt/ceph/users/${USER}/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds
+experiments_dir: /mnt/home/${USER}/particleflow/experiments
+runtime_overrides:
+  dtype: bfloat16
+  num_workers: 8
+  prefetch_factor: 2
+  model.attention.use_flash_attn_varlen: false
+slurm:
+  partition: gpu
+  constraint: h100
+  time: "48:00:00"
+  nodes: 1
+  tasks_per_node: 1
+  cpus_per_task: 64

configs/training/platforms/flatiron_h200.yaml

@@ -0,0 +1,18 @@
+name: flatiron_h200
+gpus: 8
+data_dir:
+  cld: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds
+  clic: /mnt/ceph/users/${USER}/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds
+experiments_dir: /mnt/home/${USER}/particleflow/experiments
+runtime_overrides:
+  dtype: bfloat16
+  num_workers: 8
+  prefetch_factor: 2
+  model.attention.use_flash_attn_varlen: false
+slurm:
+  partition: gpuxl
+  constraint: h200
+  time: "48:00:00"
+  nodes: 1
+  tasks_per_node: 1
+  cpus_per_task: 64

configs/training/platforms/local.yaml

@@ -0,0 +1,13 @@
+name: local
+gpus: 1
+data_dir:
+  cld: /mnt/work/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds
+  clic: /mnt/work/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds
+experiments_dir: experiments
+environment:
+  PF_SITE: local
+runtime_overrides:
+  dtype: bfloat16
+  num_workers: 8
+  prefetch_factor: 4
+  model.attention.use_flash_attn_varlen: false

configs/training/platforms/lumi_mi250x.yaml

@@ -0,0 +1,27 @@
+name: lumi_mi250x
+gpus: 8
+data_dir: /scratch/project_465001293/${USER}/tensorflow_datasets
+experiments_dir: /scratch/project_465001293/${USER}/particleflow/experiments
+environment:
+  MIOPEN_USER_DB_PATH: /tmp/${USER}-${SLURM_JOB_ID}-miopen-cache
+  MIOPEN_CUSTOM_CACHE_DIR: /tmp/${USER}-${SLURM_JOB_ID}-miopen-cache
+  ROCM_PATH: /opt/rocm
+  KERAS_BACKEND: torch
+  NCCL_SOCKET_IFNAME: hsn
+  NCCL_NET_GDR_LEVEL: "3"
+runtime_overrides:
+  compile: true
+  dtype: bfloat16
+  num_workers: 2
+  prefetch_factor: 2
+  model.attention.use_flash_attn_varlen: true
+slurm:
+  partition: standard-g
+  account: project_465001293
+  time: "1-00:00:00"
+  nodes: 1
+  tasks_per_node: 1
+  cpus_per_task: 32
+  gpu_request: gpus-per-task
+  memory: 450G
+  no_requeue: true

configs/training/platforms/tallinn_l40.yaml

@@ -0,0 +1,22 @@
+name: tallinn_l40
+gpus: 2
+data_dir:
+  cld: /local/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds
+  clic: /local/${USER}/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds
+experiments_dir: /home/${USER}/particleflow/experiments
+environment:
+  PF_SITE: tallinn
+runtime_overrides:
+  dtype: bfloat16
+  num_workers: 4
+  prefetch_factor: 2
+  model.attention.use_flash_attn_varlen: false
+slurm:
+  partition: gpu
+  time: "24:00:00"
+  nodes: 1
+  tasks_per_node: 1
+  cpus_per_task: 8
+  gpu_request: gres
+  gpu_type: l40
+  memory_per_gpu: 80G

configs/training/scenarios/cld_hits_backbone_comparison.yaml

@@ -0,0 +1,46 @@
+name: cld_hits_backbone_comparison
+spec_file: particleflow_spec.yaml
+production_name: cld
+
+variants:
+  attention:
+    model_name: pyg-cld-hits-v1
+  heptv2:
+    model_name: pyg-cld-hits-v1
+    overrides:
+      model.type: heptv2
+      model.attention: null
+      # Match the input padding used by the attention run so HEPTv2 can form
+      # complete hash-attention buckets without changing the training samples.
+      model.heptv2.block_size: 128
+
+# Add seeds here and expand the Slurm array to 2 * len(seeds) tasks.
+seeds: [12345]
+
+training:
+  # Kept fixed across hardware profiles. The runner derives the per-GPU batch.
+  global_batch_size: 512
+  parameters:
+    lr: 0.001
+    num_steps: 20000
+    val_freq: 2000
+    checkpoint_freq: 2000
+    nvalid: 512
+    ntest: 512
+    sampler_mode: interleaved-shards
+    validation_diagnostics_batches: 4
+    pad_to_multiple_elements: 128
+    make_plots: true
+
+common_overrides:
+  # Keep the output head and depth fixed so only the backbone family changes.
+  model.output_mode: elementwise
+  model.task_queries: false
+  model.backbone.mode: shared
+  model.backbone.num_convs: 6
+
+allowed_variant_differences:
+  - conv_type
+  - model.type
+  - model.attention
+  - model.heptv2

configs/training/scenarios/cld_hits_output_comparison.yaml

@@ -0,0 +1,66 @@
+name: cld_hits_output_comparison
+spec_file: particleflow_spec.yaml
+production_name: cld
+
+variants:
+  elementwise:
+    model_name: pyg-cld-hits-v1
+  set:
+    model_name: pyg-cld-hits-set-v1
+    overrides:
+      # Start slots from energetic detector inputs so their initial direction is
+      # physical, then refine within a local eta-phi neighborhood.
+      model.set_decoder.query_init: input-conditioned
+      model.set_decoder.local_attention_radius: 0.4
+      model.set_decoder.tracker_query_fraction: 0.6
+      model.set_decoder.num_layers: 4
+
+      # Treat no-object suppression as a first-class objective and supervise
+      # every refinement layer, not only the final decoder output.
+      model.set_decoder.presence_threshold: 0.5
+      model.set_decoder.no_object_weight: 1.0
+      model.set_decoder.cardinality_loss_weight: 0.05
+      model.set_decoder.auxiliary_loss_weight: 0.25
+
+      # Match in the same variables emphasized by particle/jet evaluation.
+      # A delta-R of 0.1 and a factor-two pT error each cost one unit.
+      model.set_decoder.matcher.presence: 1.0
+      model.set_decoder.matcher.pid: 1.0
+      model.set_decoder.matcher.geometry: 2.0
+      model.set_decoder.matcher.pt: 1.0
+      model.set_decoder.matcher.energy: 0.0
+      model.set_decoder.matcher.dr_scale: 0.1
+      model.set_decoder.matcher.log_pt_scale: 0.6931471805599453
+      model.set_decoder.matcher.log_energy_scale: 0.6931471805599453
+
+# Add seeds here and expand the Slurm array to 2 * len(seeds) tasks.
+seeds: [12345]
+
+training:
+  # Kept fixed across hardware profiles. The runner derives the per-GPU batch.
+  global_batch_size: 512
+  parameters:
+    lr: 0.001
+    num_steps: 20000
+    val_freq: 2000
+    checkpoint_freq: 2000
+    nvalid: 512
+    ntest: 512
+    sampler_mode: interleaved-shards
+    validation_diagnostics_batches: 4
+    pad_to_multiple_elements: 128
+    make_plots: true
+
+common_overrides:
+  model.task_queries: false
+  # Preserve the six-layer detector-aware attention backbone used by the
+  # existing Flatiron H100 training, independent of the accelerator profile.
+  model.backbone.mode: shared
+  model.backbone.num_convs: 6
+  model.backbone.num_tracker_layers: 2
+  model.backbone.num_calo_layers: 2
+  model.backbone.num_common_layers: 2
+
+allowed_variant_differences:
+  - model.output_mode
+  - model.set_decoder

configs/training/scenarios/cld_pf_hits_comparison.yaml

@@ -0,0 +1,85 @@
+name: cld_pf_hits_comparison
+spec_file: particleflow_spec.yaml
+production_name: cld
+
+variants:
+  pf:
+    model_name: pyg-cld-v1
+  elementwise_hits:
+    model_name: pyg-cld-hits-v1
+    overrides:
+      # Only hit inputs have tracker/calo detector partitions. Keep the same
+      # six-layer budget while specializing two layers per detector branch.
+      model.backbone.num_tracker_layers: 2
+      model.backbone.num_calo_layers: 2
+      model.backbone.num_common_layers: 2
+  set_hits:
+    model_name: pyg-cld-hits-set-v1
+    overrides:
+      model.backbone.num_tracker_layers: 2
+      model.backbone.num_calo_layers: 2
+      model.backbone.num_common_layers: 2
+      # Start slots from energetic detector inputs so their initial direction is
+      # physical, then refine within a local eta-phi neighborhood.
+      model.set_decoder.query_init: input-conditioned
+      model.set_decoder.local_attention_radius: 0.4
+      model.set_decoder.tracker_query_fraction: 0.6
+      model.set_decoder.num_layers: 4
+      model.set_decoder.presence_threshold: 0.5
+      model.set_decoder.no_object_weight: 1.0
+      model.set_decoder.cardinality_loss_weight: 0.05
+      model.set_decoder.auxiliary_loss_weight: 0.25
+      model.set_decoder.matcher.presence: 1.0
+      model.set_decoder.matcher.pid: 1.0
+      model.set_decoder.matcher.geometry: 2.0
+      model.set_decoder.matcher.pt: 1.0
+      model.set_decoder.matcher.energy: 0.0
+      model.set_decoder.matcher.dr_scale: 0.1
+      model.set_decoder.matcher.log_pt_scale: 0.6931471805599453
+      model.set_decoder.matcher.log_energy_scale: 0.6931471805599453
+
+# Add seeds here; launchers derive the three jobs per seed automatically.
+seeds: [12345]
+
+training:
+  # Kept fixed across hardware profiles. The runner derives the per-GPU batch.
+  global_batch_size: 512
+  parameters:
+    lr: 0.001
+    num_steps: 40000
+    val_freq: 5000
+    checkpoint_freq: 5000
+    nvalid: 512
+    ntest: 512
+    sampler_mode: interleaved-shards
+    validation_diagnostics_batches: 4
+    pad_to_multiple_elements: 128
+    make_plots: true
+
+common_overrides:
+  # Hold the trainable architecture fixed apart from the output formulation.
+  model.task_queries: false
+  model.backbone.mode: shared
+  model.backbone.num_convs: 6
+  model.attention.use_jagged_attention: true
+
+allowed_variant_differences:
+  # The PF run reads tracks/clusters while the hit runs read detector hits.
+  - dataset
+  - enabled_test_datasets
+  - input_dim
+  - train_dataset
+  - valid_dataset
+  - test_dataset
+  # This is the model-spec default; training.parameters.lr overrides it equally.
+  - hyperparameters.lr
+  # Hit elementwise training uses focal presence classification; PF keeps CE.
+  - model.binary_classification_focal_gamma
+  # Detector-specific branches are meaningful only for raw hit inputs. The PF
+  # model uses all six layers as common layers.
+  - model.backbone.num_tracker_layers
+  - model.backbone.num_calo_layers
+  - model.backbone.num_common_layers
+  # Only the set-based hit variant has a set decoder.
+  - model.output_mode
+  - model.set_decoder