From a0152910d31efa52a989262920fd5fb74b312d16 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 10:11:08 +0300 Subject: [PATCH 01/29] docs: plan set-prediction MLPF implementation --- DOING.md | 319 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 319 insertions(+) create mode 100644 DOING.md diff --git a/DOING.md b/DOING.md new file mode 100644 index 000000000..462edb39b --- /dev/null +++ b/DOING.md @@ -0,0 +1,319 @@ +# Set-prediction MLPF for hit-based training + +## Overall goal + +Add an optional slot/cross-attention MLPF training scheme that predicts a compact, +unordered set of particles directly from the input objects. In this mode, the model +must learn the relationships between detector inputs and output particles instead of +emitting one particle candidate for every input object. + +The first implementation targets the CLD and CLIC hit datasets, where an event can +contain more than 10,000 input hits but only about 100 target particles. It must be +computationally and memory efficient at that scale, while preserving the existing +elementwise MLPF path as the default and keeping existing datasets usable. + +Relevant references: + +- [Better Queries, Cheaper Attention: Adapting Transformers for Efficient Sparse Reconstruction](https://arxiv.org/abs/2606.17631) +- [HEPTv2: End-to-End Efficient Point Transformer for Charged Particle Reconstruction](https://arxiv.org/abs/2606.20437) + +## Initial scope and decisions + +- Do not regenerate the existing TFDS datasets for the first implementation. +- Derive a compact target particle collection in memory from the existing `ytarget`. + Existing hit datasets contain one row with full particle properties for each target + particle; other associated hits contain only `particle_number`. Therefore rows with + a nonzero particle class form the compact target set. +- Extract the compact target before the current input-relative `pt` and energy + transformations. Give set targets their own absolute, slot-compatible transforms. +- Treat input padding and target padding as independent axes, with separate masks. +- Add an architecture-level output mode. The existing elementwise mode remains the + default and must retain its present behavior. +- For the first model, use a fixed bank of learned particle queries and a small + cross-attention decoder. Dynamic queries, sectorization, and explicitly sparse + cross-attention are follow-up optimizations. +- Use a scalable existing encoder (initially HEPTv2) for large hit collections. A + dense all-to-all hit encoder is not acceptable for the 10k-hit use case. +- Use permutation-invariant Hungarian matching between predicted slots and compact + target particles. This output-to-target matching does not associate targets with + individual input hits. +- Do not require hit-to-particle association labels for the initial set-prediction + loss. Existing `particle_number` information may be used later for optional + auxiliary supervision. + +This compatibility path cannot undo decisions already made during dataset +preprocessing. In particular, some particles may already have been merged, dropped, +or assigned a representation-dependent PID. The initial model will therefore learn +the same particle collection as the legacy model, but without using the stored +input-target alignment in its forward pass or loss. Producing truly pre-assignment +targets is a separate, later dataset revision. + +## Proposed data flow + +```text +X [B, N, F] + input_mask + | + v +scalable hit encoder + | + v +H [B, N, D] + | + +-------------------------------+ + | +learned particle queries [B, S, D] | + | | + v | +cross-attention: queries attend to H <--+ +slot self-attention + feed-forward + | + v +presence [B, S, 2], PID [B, S, C], momentum [B, S, 5] + +compact targets [B, K, Y] + target_mask + | + v +Hungarian matching and set loss +``` + +Here `N` is the number of input hits, `K` is the event's target multiplicity, and +`S` is the configured maximum number of particle slots. Start with `S = 256`, then +choose the production value from the observed target multiplicity distribution. An +event with `K > S` must fail explicitly and must never be silently truncated. + +For a local/block-sparse encoder with fixed neighborhood size `w`, the intended +complexity is approximately + +```text +O(L_encoder * N * w + L_decoder * S * N + L_decoder * S^2). +``` + +The cross-attention implementation must use a memory-efficient SDPA/FlashAttention +path and must not materialize or retain the full attention matrix. + +## Implementation plan + +### 1. Build compact targets from existing datasets + +In `TFDSDataSource.__getitem__`, before modifying `ytarget` in place: + +1. Select rows whose particle class is nonzero. +2. Copy those rows to a new in-memory `ytarget_set` field. +3. Validate that every selected row has a unique, nonzero `particle_number` for + dataset versions that provide it. +4. Apply absolute set-target transformations, independent of `X`. Candidate initial + parameterization: `log(pt)`, `eta`, `sin(phi)`, `cos(phi)`, and `log(energy)`. +5. Continue applying the current input-relative transformations only to the legacy + `ytarget` field. + +Sorting or padding the input axis must not reorder or pad `ytarget_set` to the input +length. + +Extend `PFBatch` and `Collater` so that `X` and `ytarget_set` are padded independently. +Expose both: + +- `batch.mask`: valid input objects; +- `batch.target_mask`: valid compact target particles. + +Keep the original `ytarget`, `ytarget_pt_orig`, and `ytarget_e_orig` fields intact for +legacy training and evaluation. + +### 2. Add configuration for set prediction + +Add an output mode orthogonal to the encoder type, conceptually: + +```yaml +architecture: + type: heptv2 + output_mode: set + set_decoder: + num_slots: 256 + num_layers: 2 + num_heads: 8 + query_init: learned + cross_attention: flash +``` + +Requirements: + +- Default `output_mode` to `elementwise`. +- Initially enable `set` only for `cld_hits` and `clic_hits`. +- Validate dimensions and require `num_slots > 0`. +- Do not overload the existing `task_queries` option; those queries are per-element + task readouts, not output-particle slots. +- Give unsupported combinations, including initial ONNX export if necessary, clear + configuration errors. + +### 3. Implement the set decoder + +Create `mlpf/model/set_decoder.py` and reuse/refactor `MLPF.encode_backbone()` as the +common encoder interface. + +The initial decoder should contain: + +- a learned bank of `S` query embeddings; +- two pre-normalized decoder layers; +- slots-to-input cross-attention; +- slot self-attention to coordinate and suppress duplicate predictions; +- a feed-forward block and residual connections; +- presence, PID, and absolute-momentum output heads. + +Use packed variable-length Flash cross-attention where available so padded hits do +not consume decoder compute or memory. Provide a dense masked SDPA fallback for CPU +and unit tests. Do not request attention weights in the production path. + +Keep the external prediction representation close to the existing one, but make its +particle axis the slot axis rather than the input axis. + +### 4. Add Hungarian matching and set losses + +Create `mlpf/model/set_losses.py`. For each event, construct a detached matching cost +between its valid slots and targets using: + +- PID classification cost; +- distance in `log(pt)`; +- distance in `eta`; +- cyclic phi cost, for example `1 - cos(delta_phi)`; +- distance in `log(energy)`. + +Use `scipy.optimize.linear_sum_assignment` initially. The matching problem is small +and independent of the number of input hits. Profile its CPU/GPU synchronization +overhead before considering a GPU matcher. + +After matching, optimize: + +- particle-presence loss over every slot; +- PID loss over matched slots; +- kinematic losses over matched slots only; +- no-particle targets for unmatched slots. + +Normalize losses per event or per target particle so high-multiplicity events do not +dominate. Keep matcher cost weights separate from optimized loss weights. Reuse the +existing task-loss calibration only if its assumptions remain valid for the new loss +normalization. + +### 5. Route training, validation, and inference + +Route the model and loss in `training.py` according to `output_mode`, without changing +the elementwise path. + +For set-mode inference: + +- select active slots using the presence prediction; +- restore physical `pt` and energy without referring to an input object; +- unpad targets using `target_mask`; +- serialize predictions using the number of selected slots, not the input-hit count; +- keep serializing the complete input hit collection separately; +- build jets and MET from the selected particle slots. + +Update diagnostic tables and particle-quality metrics so they no longer assume that +predictions, targets, and inputs share an axis. + +### 6. Verify correctness and scalability + +Add unit tests for: + +- compact-target extraction from representative and `particle_number`-only rows; +- independent input and target padding; +- target permutation invariance; +- known Hungarian assignments; +- unmatched slots and zero-target events; +- `K = S` and explicit `K > S` failure; +- padding masks and finite mixed-precision forward/backward results; +- set output shapes being independent of `N`; +- unchanged legacy configuration, output shapes, and losses. + +Add an integration test with approximately `N = 10,000`, `K = 100`, and `S = 256`, +including backward propagation. + +Extend `scripts/benchmark.py` to sweep at least `N = 1k, 5k, 10k, 20k` and record: + +- encoder, decoder, loss, forward, and backward time; +- peak allocated and reserved GPU memory; +- matcher time; +- valid input and target multiplicities. + +The implementation should scale approximately linearly with `N` for fixed encoder +block size and fixed `S`. Profiling must confirm that no persistent `[B, S, N]` +attention tensor is allocated. + +Compare elementwise and set prediction using: + +- particle multiplicity, efficiency, and fake rate; +- PID confusion and performance versus `pt` and `eta`; +- particle response and resolution; +- jet matching, response, and resolution; +- MET response and resolution; +- summed event energy; +- training throughput and peak memory. + +## TODO + +### Data and batching + +- [ ] Add an `output_mode` or equivalent argument to the data-loading path. +- [ ] Extract `ytarget_set` before legacy relative-target transformations. +- [ ] Add uniqueness and consistency checks using `particle_number`. +- [ ] Define and test the absolute set-target transformation and its inverse. +- [ ] Add `ytarget_set` and `target_mask` to `PFBatch`. +- [ ] Pad input and target collections independently in `Collater`. +- [ ] Verify compact target counts, PIDs, and summed energy against legacy nonzero + target rows on existing CLD/CLIC hit samples. + +### Configuration and model + +- [ ] Add `output_mode` configuration with backward-compatible defaults. +- [ ] Add a validated `SetDecoderConfig`. +- [ ] Add a hit-dataset set-mode example to `particleflow_spec.yaml`. +- [ ] Refactor/reuse the backbone encoder without changing legacy forward behavior. +- [ ] Implement learned fixed queries and two decoder layers. +- [ ] Implement memory-efficient packed cross-attention plus a CPU test fallback. +- [ ] Implement presence, PID, and absolute-momentum heads. +- [ ] Assert and log target-slot overflow instead of truncating. + +### Matching and loss + +- [ ] Implement per-event Hungarian matching. +- [ ] Implement configurable matching costs with cyclic phi handling. +- [ ] Implement matched presence, PID, and regression losses. +- [ ] Decide and test event/particle normalization and no-particle weighting. +- [ ] Integrate or replace calibrated task-loss weighting for set mode. +- [ ] Log target count, active-slot count, matched cost, and unmatched-slot statistics. + +### Training and inference + +- [ ] Route training and validation by output mode. +- [ ] Update validation diagnostics for independent axes. +- [ ] Update `predict_particles` for set outputs and absolute inverse transforms. +- [ ] Update parquet inference serialization to use separate input, target, and + prediction counts. +- [ ] Update particle, jet, and MET metrics for set outputs. +- [ ] Add an explicit error or support path for set-mode ONNX export. + +### Testing and performance + +- [ ] Add compact-target extraction and batching tests. +- [ ] Add matcher and target-permutation tests. +- [ ] Add decoder masking, shape, and numerical-stability tests. +- [ ] Add legacy regression tests. +- [ ] Add a 10k-hit forward/backward integration test. +- [ ] Extend the benchmark script with set-mode timing and memory measurements. +- [ ] Run a small CLD-hits overfit test and confirm that loss and matching converge. +- [ ] Run a short CLD-hits training comparison against the elementwise baseline. +- [ ] Document observed accuracy, throughput, memory, and scaling results here. + +## Follow-up work after the initial baseline + +- [ ] Add input-conditioned dynamic queries. Charged-particle queries may use + innermost tracker-hit candidates; neutral-particle queries will need a separate + calorimeter seeding strategy. +- [ ] Evaluate phi-sector decoding and boundary behavior. +- [ ] Evaluate local strided/block-sparse cross-attention if dense Flash + cross-attention remains a material compute cost. +- [ ] Add optional encoder contrastive/background supervision from + `particle_number` without making it necessary for set prediction. +- [ ] Add optional chunked or local slot-to-hit assignment heads for interpretability + and auxiliary losses. +- [ ] If the baseline is successful, generate a new dataset version containing a + truly independent pre-assignment particle collection and corresponding target + jets. From e01c599473817df427fc8c17102235f51266ddeb Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 11:40:19 +0300 Subject: [PATCH 02/29] feat: add set-prediction MLPF training --- DOING.md | 89 ++++-- mlpf/conf.py | 26 ++ mlpf/model/PFDataset.py | 50 ++- mlpf/model/inference.py | 48 +-- mlpf/model/mlpf.py | 67 ++-- mlpf/model/set_decoder.py | 104 +++++++ mlpf/model/set_losses.py | 216 +++++++++++++ mlpf/model/training.py | 85 +++++- mlpf/model/validation_metrics.py | 338 +++++++++++++++++++++ particleflow_spec.yaml | 28 +- scripts/local/make_local_available_spec.py | 3 +- scripts/local/train.sh | 88 +++--- tests/test_pfdataset_logic.py | 65 ++++ tests/test_set_prediction.py | 234 ++++++++++++++ tests/test_training_diagnostics.py | 33 +- tests/test_validation_metrics.py | 164 ++++++++++ 16 files changed, 1495 insertions(+), 143 deletions(-) create mode 100644 mlpf/model/set_decoder.py create mode 100644 mlpf/model/set_losses.py create mode 100644 mlpf/model/validation_metrics.py create mode 100644 tests/test_set_prediction.py create mode 100644 tests/test_validation_metrics.py diff --git a/DOING.md b/DOING.md index 462edb39b..d2f99b47f 100644 --- a/DOING.md +++ b/DOING.md @@ -32,7 +32,7 @@ Relevant references: - For the first model, use a fixed bank of learned particle queries and a small cross-attention decoder. Dynamic queries, sectorization, and explicitly sparse cross-attention are follow-up optimizations. -- Use a scalable existing encoder (initially HEPTv2) for large hit collections. A +- Use a scalable existing encoder (currently packed attention) for large hit collections. A dense all-to-all hit encoder is not acceptable for the 10k-hit use case. - Use permutation-invariant Hungarian matching between predicted slots and compact target particles. This output-to-target matching does not associate targets with @@ -124,7 +124,7 @@ Add an output mode orthogonal to the encoder type, conceptually: ```yaml architecture: - type: heptv2 + type: attention output_mode: set set_decoder: num_slots: 256 @@ -251,56 +251,83 @@ Compare elementwise and set prediction using: ### Data and batching -- [ ] Add an `output_mode` or equivalent argument to the data-loading path. -- [ ] Extract `ytarget_set` before legacy relative-target transformations. -- [ ] Add uniqueness and consistency checks using `particle_number`. -- [ ] Define and test the absolute set-target transformation and its inverse. -- [ ] Add `ytarget_set` and `target_mask` to `PFBatch`. -- [ ] Pad input and target collections independently in `Collater`. -- [ ] Verify compact target counts, PIDs, and summed energy against legacy nonzero +- [x] Add an `output_mode` or equivalent argument to the data-loading path. +- [x] Extract `ytarget_set` before legacy relative-target transformations. +- [x] Add uniqueness and consistency checks using `particle_number`. +- [x] Define and test the absolute set-target transformation and its inverse. +- [x] Add `ytarget_set` and `target_mask` to `PFBatch`. +- [x] Pad input and target collections independently in `Collater`. +- [x] Verify compact target counts, PIDs, and summed energy against legacy nonzero target rows on existing CLD/CLIC hit samples. ### Configuration and model -- [ ] Add `output_mode` configuration with backward-compatible defaults. -- [ ] Add a validated `SetDecoderConfig`. -- [ ] Add a hit-dataset set-mode example to `particleflow_spec.yaml`. -- [ ] Refactor/reuse the backbone encoder without changing legacy forward behavior. -- [ ] Implement learned fixed queries and two decoder layers. +- [x] Add `output_mode` configuration with backward-compatible defaults. +- [x] Add a validated `SetDecoderConfig`. +- [x] Add a hit-dataset set-mode example to `particleflow_spec.yaml`. +- [x] Refactor/reuse the backbone encoder without changing legacy forward behavior. +- [x] Implement learned fixed queries and two decoder layers. - [ ] Implement memory-efficient packed cross-attention plus a CPU test fallback. -- [ ] Implement presence, PID, and absolute-momentum heads. -- [ ] Assert and log target-slot overflow instead of truncating. +- [x] Implement presence, PID, and absolute-momentum heads. +- [x] Assert target-slot overflow instead of truncating; add aggregate logging later. ### Matching and loss -- [ ] Implement per-event Hungarian matching. -- [ ] Implement configurable matching costs with cyclic phi handling. -- [ ] Implement matched presence, PID, and regression losses. -- [ ] Decide and test event/particle normalization and no-particle weighting. -- [ ] Integrate or replace calibrated task-loss weighting for set mode. +- [x] Implement per-event Hungarian matching. +- [x] Implement matching costs with cyclic phi handling; expose them in model configuration later. +- [x] Implement matched presence, PID, and regression losses. +- [x] Decide and test event/particle normalization and no-particle weighting. +- [x] Integrate calibrated task-loss weighting for set mode. - [ ] Log target count, active-slot count, matched cost, and unmatched-slot statistics. ### Training and inference -- [ ] Route training and validation by output mode. +- [x] Route training and validation loss calculation by output mode. +- [x] Log scheme-independent particle matching, count, PID, kinematic, and event + closure metrics during validation. - [ ] Update validation diagnostics for independent axes. -- [ ] Update `predict_particles` for set outputs and absolute inverse transforms. -- [ ] Update parquet inference serialization to use separate input, target, and +- [x] Update `predict_particles` for set outputs and absolute inverse transforms. +- [x] Update parquet inference serialization to use separate input, target, and prediction counts. - [ ] Update particle, jet, and MET metrics for set outputs. - [ ] Add an explicit error or support path for set-mode ONNX export. ### Testing and performance -- [ ] Add compact-target extraction and batching tests. -- [ ] Add matcher and target-permutation tests. -- [ ] Add decoder masking, shape, and numerical-stability tests. -- [ ] Add legacy regression tests. -- [ ] Add a 10k-hit forward/backward integration test. +- [x] Add compact-target extraction and batching tests. +- [x] Add matcher and target-permutation tests. +- [x] Add decoder masking, shape, and numerical-stability tests. +- [x] Run the existing legacy model and loss regression tests. +- [x] Add a 10k-hit forward/backward integration test. - [ ] Extend the benchmark script with set-mode timing and memory measurements. -- [ ] Run a small CLD-hits overfit test and confirm that loss and matching converge. +- [x] Run a small CLD-hits overfit test and confirm that loss and matching converge. +- [x] Add a local ttbar launcher for paired elementwise and set-output training. - [ ] Run a short CLD-hits training comparison against the elementwise baseline. -- [ ] Document observed accuracy, throughput, memory, and scaling results here. +- [x] Document initial correctness, timing, memory, and scaling measurements here; + add physics accuracy after training. + +## Initial implementation measurements + +Measurements from 2026-09-04 using the existing CLD `cld_edm_ttbar_hits/1:3.2.1` +dataset: + +- The first 50 events contain 4,388--13,219 valid hits and 33--148 compact target + particles. Compact target counts matched the number of nonzero legacy target rows + in every event. +- A real event with 6,423 valid hits, 52 targets, and 256 slots completed the full + four-layer HEPTv2 forward pass, Hungarian loss, and backward pass on an NVIDIA + GeForce RTX 5060 Ti. Peak allocated GPU memory was approximately 0.70 GB. +- A real event with 13,219 valid hits, 135 targets, and 256 slots completed the same + path with approximately 1.39 GB peak allocated GPU memory. +- A 20-step single-event AdamW overfit check with the full HEPTv2 encoder reduced the + uncalibrated total loss from 41.04 to 16.49, confirming end-to-end gradients through + the encoder, decoder, and matched loss. +- The current comparison recipes use a three-layer packed attention backbone. Both + elementwise and set modes completed a GPU forward-pass smoke test on the 6,423-hit + event, using approximately 0.10 GB and 0.08 GB of allocated GPU memory respectively. +- `uv run pytest -q tests` passed 241 tests with 3 skips. Running pytest from the + repository root without restricting collection still encounters two unrelated + pre-existing collection errors under `baselines/HEPTv2` and `scripts/legacy`. ## Follow-up work after the initial baseline diff --git a/mlpf/conf.py b/mlpf/conf.py index 19b889eda..1709fd3c1 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -39,6 +39,11 @@ class BackboneMode(Enum): SPLIT = "split" +class OutputMode(Enum): + ELEMENTWISE = "elementwise" + SET = "set" + + class DatasetSamplerMode(Enum): SHARD_CONSECUTIVE = "shard-consecutive" INTERLEAVED_SHARDS = "interleaved-shards" @@ -586,6 +591,18 @@ class HitFeatureEngineeringConfig(BaseModel): calorimeter_neighborhood: bool = True +class SetDecoderConfig(BaseModel): + """Configuration for permutation-invariant particle-set prediction.""" + + model_config = ConfigDict(extra="forbid") + + num_slots: int = Field(default=256, gt=0) + num_layers: int = Field(default=2, gt=0) + num_heads: int = Field(default=8, gt=0) + ffn_multiplier: float = Field(default=4.0, gt=0.0) + dropout: float = Field(default=0.0, ge=0.0, lt=1.0) + + class ModelArchitectureConfig(BaseModel): model_config = ConfigDict(extra="forbid") @@ -599,8 +616,10 @@ class ModelArchitectureConfig(BaseModel): energy_mode: RegressionMode = RegressionMode.DIRECT_ELEMTYPE_SPLIT trainable: str = "all" task_queries: bool = True + output_mode: OutputMode = OutputMode.ELEMENTWISE backbone: Optional[BackboneConfig] = None hit_feature_engineering: HitFeatureEngineeringConfig = Field(default_factory=HitFeatureEngineeringConfig) + set_decoder: Optional[SetDecoderConfig] = None # Nested configs gnnlsh: Optional[GNNLSHConfig] = None @@ -741,6 +760,13 @@ def populate_defaults(self) -> "MLPFConfig": self.num_classes = len(CLASS_LABELS[self.dataset.value]) if self.elemtypes_nonzero is None: self.elemtypes_nonzero = ELEM_TYPES_NONZERO[self.dataset.value] + if self.model.output_mode == OutputMode.SET: + if self.dataset not in (Dataset.CLD_HITS, Dataset.CLIC_HITS): + raise ValueError("model.output_mode='set' is currently supported only for CLD/CLIC hit datasets") + if self.model.set_decoder is None: + self.model.set_decoder = SetDecoderConfig() + if self.model.backbone.mode != BackboneMode.SHARED: + raise ValueError("model.output_mode='set' currently requires model.backbone.mode='shared'") return self def flatten_config(self, prefix=""): diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index 7cedc6637..875a25dca 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -10,7 +10,7 @@ import torch.utils.data from mlpf.logger import _logger -from mlpf.conf import DatasetSamplerMode, MLPFConfig, dataset_input_type_id, dataset_source_id +from mlpf.conf import DatasetSamplerMode, MLPFConfig, OutputMode, Y_FEATURES, dataset_input_type_id, dataset_source_id # https://github.com/pytorch/pytorch/issues/11201#issuecomment-895047235 @@ -30,7 +30,7 @@ class TFDSDataSource: - def __init__(self, ds, sort, pad_to_multiple=None, feature_dim=None, max_open_readers=None): + def __init__(self, ds, sort, pad_to_multiple=None, feature_dim=None, max_open_readers=None, build_target_set=False): self.ds = ds tmp = self.ds.dataset_info self.ds.dataset_info = SimpleNamespace() @@ -41,6 +41,7 @@ def __init__(self, ds, sort, pad_to_multiple=None, feature_dim=None, max_open_re self.pad_to_multiple = pad_to_multiple self.feature_dim = feature_dim self.max_open_readers = max_open_readers + self.build_target_set = build_target_set def _close_extra_readers(self): # ArrayRecordDataSource caches one reader per shard file forever, so the @@ -143,6 +144,29 @@ def __getitem__(self, item): e = ret["X"][:, 5][msk_ho] ret["X"][:, 1][msk_ho] = np.sqrt(e**2 - (np.tanh(eta) * e) ** 2) + if self.build_target_set: + # Existing hit datasets store one full target row on an exclusive + # representative hit for every particle. Other related hits contain + # only particle_number, so nonzero class rows form the compact set. + target_rows = ret["ytarget"][:, 0] != 0 + ytarget_set = ret["ytarget"][target_rows].copy() + + particle_number_idx = Y_FEATURES.index("particle_number") + if ytarget_set.shape[1] > particle_number_idx and len(ytarget_set) > 0: + particle_numbers = ytarget_set[:, particle_number_idx] + # Some old datasets predate particle_number. If it is present, + # require the representative rows to be one-per-particle. + if np.any(particle_numbers != 0): + if np.any(particle_numbers == 0) or len(np.unique(particle_numbers)) != len(particle_numbers): + raise ValueError("Full ytarget rows must have unique, nonzero particle_number values in set mode") + + # Set outputs have no input-object anchor. Train pt and energy in + # absolute log space instead of using the legacy log(target/input). + if len(ytarget_set) > 0: + ytarget_set[:, 2] = np.log(np.clip(ytarget_set[:, 2], 1e-8, None)) + ytarget_set[:, 6] = np.log(np.clip(ytarget_set[:, 6], 1e-8, None)) + ret["ytarget_set"] = ytarget_set + # transform pt -> log(pt / elem pt), same for energy # where target does not exist, set to 0 with np.errstate(divide="ignore"): @@ -172,7 +196,18 @@ def __repr__(self): class PFDataset: """Builds a DataSource from tensorflow datasets.""" - def __init__(self, data_dir, name, split, num_samples=None, sort=False, pad_to_multiple=512, feature_dim=None, max_open_readers=None): + def __init__( + self, + data_dir, + name, + split, + num_samples=None, + sort=False, + pad_to_multiple=512, + feature_dim=None, + max_open_readers=None, + build_target_set=False, + ): """ Args data_dir: path to tensorflow_datasets (e.g. `../data/tensorflow_datasets/`) @@ -198,6 +233,7 @@ def __init__(self, data_dir, name, split, num_samples=None, sort=False, pad_to_m pad_to_multiple=pad_to_multiple, feature_dim=feature_dim, max_open_readers=max_open_readers, + build_target_set=build_target_set, ) if num_samples and num_samples < len(self.ds): @@ -214,6 +250,7 @@ def __init__(self, **kwargs): # write out the possible attributes here explicitly self.X = kwargs["X"] self.ytarget = kwargs.get("ytarget") + self.ytarget_set = kwargs.get("ytarget_set") self.ytarget_pt_orig = kwargs.get("ytarget_pt_orig", None) self.ytarget_e_orig = kwargs.get("ytarget_e_orig", None) self.pythia = kwargs.get("pythia", None) @@ -224,6 +261,7 @@ def __init__(self, **kwargs): self.source_id = kwargs.get("source_id", None) self.input_type_id = kwargs.get("input_type_id", None) self.mask = self.X[:, :, 0] != 0 + self.target_mask = self.ytarget_set[:, :, 0] != 0 if self.ytarget_set is not None else None def to(self, device, **kwargs): attrs = {} @@ -668,6 +706,7 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, pad_to_multiple=config.pad_to_multiple_elements, feature_dim=config.input_dim, max_open_readers=config.max_open_readers, + build_target_set=config.model.output_mode == OutputMode.SET, ).ds if (rank == 0) or (rank == "cpu"): @@ -695,10 +734,13 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, # build dataloaders batch_size = physical_ds.batch_size * config.gpu_batch_multiplier + per_particle_keys = ["X", "ytarget"] + if config.model.output_mode == OutputMode.SET: + per_particle_keys.append("ytarget_set") loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, - collate_fn=Collater(["X", "ytarget"], ["genmet", "source_id", "input_type_id"]), + collate_fn=Collater(per_particle_keys, ["genmet", "source_id", "input_type_id"]), sampler=sampler, num_workers=config.num_workers, prefetch_factor=config.prefetch_factor, diff --git a/mlpf/model/inference.py b/mlpf/model/inference.py index 3fa22bf5f..aa354bf3f 100644 --- a/mlpf/model/inference.py +++ b/mlpf/model/inference.py @@ -32,6 +32,7 @@ from mlpf.logger import _logger from mlpf.model.utils import unpack_target +from mlpf.conf import OutputMode def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_match_dr, outpath, dir_name, sample): @@ -51,10 +52,18 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m ypred = model_module.predict_particles(batch.X, batch.mask) - batch.ytarget[..., 2] = batch.ytarget_pt_orig - batch.ytarget[..., 6] = batch.ytarget_e_orig - - ytarget = unpack_target(batch.ytarget.to(torch.float32), model_module) + if model_module.output_mode == OutputMode.SET: + ytarget = unpack_target(batch.ytarget_set.to(torch.float32), model_module) + ytarget["pt"] = torch.exp(ytarget["pt"]) + ytarget["energy"] = torch.exp(ytarget["energy"]) + ytarget["momentum"] = torch.stack( + [ytarget["pt"], ytarget["eta"], ytarget["sin_phi"], ytarget["cos_phi"], ytarget["energy"]], dim=-1 + ) + ytarget["p4"] = torch.stack([ytarget["pt"], ytarget["eta"], ytarget["phi"], ytarget["energy"]], dim=-1) + else: + batch.ytarget[..., 2] = batch.ytarget_pt_orig + batch.ytarget[..., 6] = batch.ytarget_e_orig + ytarget = unpack_target(batch.ytarget.to(torch.float32), model_module) ycand = unpack_target(batch.ycand.to(torch.float32), model_module) genjets_msk = batch.genjets[:, :, 0].cpu() > jet_ptcut @@ -74,23 +83,28 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m jets_coll = {} jets_coll["gen"] = genjets - # now cluster jets - # first, flatten events across batch dim with padding mask + # Flatten each independently padded collection with its own mask. X = batch.X[batch.mask].cpu().float().contiguous().numpy() - for k, v in ytarget.items(): - ytarget[k] = v[batch.mask].detach().cpu().float().contiguous().numpy() - for k, v in ycand.items(): - ycand[k] = v[batch.mask].detach().cpu().float().contiguous().numpy() - for k, v in ypred.items(): - ypred[k] = v[batch.mask].detach().cpu().float().contiguous().numpy() - - # second, create awkward arrays according to the counts of not padded elements - counts = torch.sum(batch.mask, axis=1).cpu().numpy() + input_counts = torch.sum(batch.mask, axis=1).cpu().numpy() + if model_module.output_mode == OutputMode.SET: + target_mask = batch.target_mask.bool() + prediction_mask = ypred["cls_id"] != 0 + else: + target_mask = batch.mask.bool() + prediction_mask = batch.mask.bool() + + collection_masks = {"target": target_mask, "cand": batch.mask.bool(), "pred": prediction_mask} awkvals = {} for flat_arr, typ in [(ytarget, "target"), (ycand, "cand"), (ypred, "pred")]: - awk_arr = awkward.Array({k: flat_arr[k] for k in flat_arr.keys()}) + collection_mask = collection_masks[typ] + counts = collection_mask.sum(dim=1).cpu().numpy() + values = { + key: value[collection_mask].detach().cpu().float().contiguous().numpy() + for key, value in flat_arr.items() + } + awk_arr = awkward.Array(values) awkvals[typ] = awkward.unflatten(awk_arr, counts) - Xs = awkward.unflatten(awkward.from_numpy(X), counts) + Xs = awkward.unflatten(awkward.from_numpy(X), input_counts) # now cluster jets for typ, ydata in zip( diff --git a/mlpf/model/mlpf.py b/mlpf/model/mlpf.py index 290e1c198..fd6067dda 100644 --- a/mlpf/model/mlpf.py +++ b/mlpf/model/mlpf.py @@ -16,6 +16,7 @@ from mlpf.model.hept import HEPTLayer, trunc_normal_ from mlpf.model.heptv2 import HEPTv2Layer +from mlpf.model.set_decoder import ParticleSetDecoder try: from flash_attn import flash_attn_varlen_func as _flash_attn_varlen_func @@ -28,6 +29,7 @@ AttentionType, BackboneMode, ModelType, + OutputMode, InputEncoding, LearnedRepresentationMode, RegressionMode, @@ -1136,6 +1138,7 @@ def __init__( super(MLPF, self).__init__() self.config = config.model + self.output_mode = OutputMode(self.config.output_mode) self.is_hit_dataset = config.dataset in (Dataset.CLD_HITS, Dataset.CLIC_HITS) self.raw_input_dim = config.input_dim hit_feature_config = config.model.hit_feature_engineering @@ -1384,7 +1387,7 @@ def __init__( self.classification_norm = torch.nn.LayerNorm(decoding_dim) if self.use_pre_layernorm else None self.regression_norm = torch.nn.LayerNorm(decoding_dim) if self.use_pre_layernorm else None - if self.task_queries and not self.use_split_backbone: + if self.task_queries and not self.use_split_backbone and self.output_mode == OutputMode.ELEMENTWISE: self.classification_query = nn.Parameter(torch.zeros(1, 1, decoding_dim), requires_grad=True) self.regression_query = nn.Parameter(torch.zeros(1, 1, decoding_dim), requires_grad=True) trunc_normal_(self.classification_query, std=0.02) @@ -1422,14 +1425,25 @@ def __init__( self.classification_readout = None self.regression_readout = None - self.nn_binary_particle = ffn(decoding_dim, 2, width, self.act, head_dropout_ff) - self.nn_pid = ffn(decoding_dim, self.num_classes, width, self.act, head_dropout_ff) - - self.nn_pt = RegressionOutput(pt_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) - self.nn_eta = RegressionOutput(eta_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) - self.nn_sin_phi = RegressionOutput(sin_phi_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) - self.nn_cos_phi = RegressionOutput(cos_phi_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) - self.nn_energy = RegressionOutput(energy_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) + self.set_decoder = None + self.nn_binary_particle = None + self.nn_pid = None + self.nn_pt = None + self.nn_eta = None + self.nn_sin_phi = None + self.nn_cos_phi = None + self.nn_energy = None + if self.output_mode == OutputMode.SET: + self.set_decoder = ParticleSetDecoder(decoding_dim, self.num_classes, self.config.set_decoder) + else: + self.nn_binary_particle = ffn(decoding_dim, 2, width, self.act, head_dropout_ff) + self.nn_pid = ffn(decoding_dim, self.num_classes, width, self.act, head_dropout_ff) + + self.nn_pt = RegressionOutput(pt_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) + self.nn_eta = RegressionOutput(eta_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) + self.nn_sin_phi = RegressionOutput(sin_phi_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) + self.nn_cos_phi = RegressionOutput(cos_phi_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) + self.nn_energy = RegressionOutput(energy_mode, decoding_dim, width, self.act, head_dropout_ff, regression_selector_values) _logger.info("Output DNNs initialization took {:.2f}s".format(time.time() - t0)) _logger.info("backbone_mode={}".format(self.backbone_mode)) @@ -1452,13 +1466,15 @@ def __init__( _logger.info( "regression_readout parameters: {}".format(count_parameters(self.regression_readout) if self.regression_readout is not None else 0) ) - _logger.info("nn_binary_particle parameters: {}".format(count_parameters(self.nn_binary_particle))) - _logger.info("nn_pid parameters: {}".format(count_parameters(self.nn_pid))) - _logger.info("nn_pt parameters: {}".format(count_parameters(self.nn_pt))) - _logger.info("nn_eta parameters: {}".format(count_parameters(self.nn_eta))) - _logger.info("nn_sin_phi parameters: {}".format(count_parameters(self.nn_sin_phi))) - _logger.info("nn_cos_phi parameters: {}".format(count_parameters(self.nn_cos_phi))) - _logger.info("nn_energy parameters: {}".format(count_parameters(self.nn_energy))) + _logger.info("output_mode={}".format(self.output_mode.value)) + _logger.info("set_decoder parameters: {}".format(count_parameters(self.set_decoder) if self.set_decoder is not None else 0)) + _logger.info("nn_binary_particle parameters: {}".format(count_parameters(self.nn_binary_particle) if self.nn_binary_particle is not None else 0)) + _logger.info("nn_pid parameters: {}".format(count_parameters(self.nn_pid) if self.nn_pid is not None else 0)) + _logger.info("nn_pt parameters: {}".format(count_parameters(self.nn_pt) if self.nn_pt is not None else 0)) + _logger.info("nn_eta parameters: {}".format(count_parameters(self.nn_eta) if self.nn_eta is not None else 0)) + _logger.info("nn_sin_phi parameters: {}".format(count_parameters(self.nn_sin_phi) if self.nn_sin_phi is not None else 0)) + _logger.info("nn_cos_phi parameters: {}".format(count_parameters(self.nn_cos_phi) if self.nn_cos_phi is not None else 0)) + _logger.info("nn_energy parameters: {}".format(count_parameters(self.nn_energy) if self.nn_energy is not None else 0)) _logger.info("Total MLPF parameters: {}".format(count_parameters(self))) _logger.info("MLPF __init__ done") @@ -1758,6 +1774,9 @@ def final_norm_reg(self): # @torch.compile def forward(self, X_features, mask): + if self.output_mode == OutputMode.SET: + return self.set_decoder(self.encode_backbone(X_features, mask), mask) + X_features = self._engineer_input_features(X_features, mask) if self.use_split_backbone: x_id = self._encode_inputs(X_features, mask=mask, encoder=self._nn0_id) @@ -1819,14 +1838,18 @@ def predict_particles(self, X_features, mask): from mlpf.model.utils import unpack_predictions ypred_raw = self.forward(X_features, mask) - ypred_raw = tuple([y.to(torch.float32) for y in ypred_raw]) + ypred_raw = [y.to(torch.float32) for y in ypred_raw] - # transform log (pt/elempt) -> pt - ypred_raw[2][..., 0] = torch.exp(ypred_raw[2][..., 0]) * X_features[..., 1] - # transform log (E/elemE) -> E - ypred_raw[2][..., 4] = torch.exp(ypred_raw[2][..., 4]) * X_features[..., 5] + if self.output_mode == OutputMode.SET: + ypred_raw[2][..., 0] = torch.exp(ypred_raw[2][..., 0].clamp(-20.0, 20.0)) + ypred_raw[2][..., 4] = torch.exp(ypred_raw[2][..., 4].clamp(-20.0, 20.0)) + else: + # transform log (pt/elempt) -> pt + ypred_raw[2][..., 0] = torch.exp(ypred_raw[2][..., 0]) * X_features[..., 1] + # transform log (E/elemE) -> E + ypred_raw[2][..., 4] = torch.exp(ypred_raw[2][..., 4]) * X_features[..., 5] - ypred = unpack_predictions(ypred_raw) + ypred = unpack_predictions(tuple(ypred_raw)) ypred["ispu"] = torch.softmax(ypred["ispu"], axis=-1)[:, :, -1] # By default, use standard argmax diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py new file mode 100644 index 000000000..33f46f321 --- /dev/null +++ b/mlpf/model/set_decoder.py @@ -0,0 +1,104 @@ +import torch +from torch import nn +from torch.nn import functional as F + + +class ParticleSetDecoderLayer(nn.Module): + """Pre-norm particle-query decoder layer. + + Cross-attention is evaluated event by event on only the valid input embeddings. + With ``need_weights=False``, PyTorch dispatches through scaled-dot-product + attention and can select a fused FlashAttention kernel on CUDA without retaining + the slot-by-input attention matrix. + """ + + def __init__(self, embedding_dim, num_heads, ffn_dim, dropout=0.0): + super().__init__() + self.query_norm = nn.LayerNorm(embedding_dim) + self.memory_norm = nn.LayerNorm(embedding_dim) + self.cross_attention = nn.MultiheadAttention( + embedding_dim, num_heads, dropout=dropout, batch_first=True + ) + self.self_norm = nn.LayerNorm(embedding_dim) + self.self_attention = nn.MultiheadAttention( + embedding_dim, num_heads, dropout=dropout, batch_first=True + ) + self.ffn_norm = nn.LayerNorm(embedding_dim) + self.ffn = nn.Sequential( + nn.Linear(embedding_dim, ffn_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(ffn_dim, embedding_dim), + nn.Dropout(dropout), + ) + + def forward(self, slots, memory, memory_mask): + cross_queries = self.query_norm(slots) + normalized_memory = self.memory_norm(memory) + cross_outputs = [] + for event_idx in range(memory.shape[0]): + event_memory = normalized_memory[ + event_idx : event_idx + 1, memory_mask[event_idx] + ] + if event_memory.shape[1] == 0: + cross_outputs.append( + torch.zeros_like(cross_queries[event_idx : event_idx + 1]) + ) + continue + event_output, _ = self.cross_attention( + cross_queries[event_idx : event_idx + 1], + event_memory, + event_memory, + need_weights=False, + ) + cross_outputs.append(event_output) + slots = slots + torch.cat(cross_outputs, dim=0) + + normalized_slots = self.self_norm(slots) + self_output, _ = self.self_attention( + normalized_slots, normalized_slots, normalized_slots, need_weights=False + ) + slots = slots + self_output + return slots + self.ffn(self.ffn_norm(slots)) + + +class ParticleSetDecoder(nn.Module): + """Decode a fixed bank of learned queries into an unordered particle set.""" + + def __init__(self, embedding_dim, num_classes, config): + super().__init__() + if embedding_dim % config.num_heads != 0: + raise ValueError( + f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}" + ) + + self.num_slots = config.num_slots + self.queries = nn.Parameter(torch.empty(1, config.num_slots, embedding_dim)) + nn.init.trunc_normal_(self.queries, std=0.02) + ffn_dim = int(config.ffn_multiplier * embedding_dim) + self.layers = nn.ModuleList( + ParticleSetDecoderLayer( + embedding_dim, config.num_heads, ffn_dim, config.dropout + ) + for _ in range(config.num_layers) + ) + self.output_norm = nn.LayerNorm(embedding_dim) + self.presence_head = nn.Linear(embedding_dim, 2) + self.pid_head = nn.Linear(embedding_dim, num_classes) + self.momentum_head = nn.Linear(embedding_dim, 5) + + def forward(self, memory, memory_mask): + slots = self.queries.expand(memory.shape[0], -1, -1) + for layer in self.layers: + slots = layer(slots, memory, memory_mask.bool()) + slots = self.output_norm(slots) + + presence = self.presence_head(slots) + pid = self.pid_head(slots) + momentum = self.momentum_head(slots) + phi_direction = F.normalize(momentum[..., 2:4], dim=-1, eps=1e-6) + momentum = torch.cat( + [momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1 + ) + pileup = torch.zeros_like(presence) + return presence, pid, momentum, pileup diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py new file mode 100644 index 000000000..c92f60c22 --- /dev/null +++ b/mlpf/model/set_losses.py @@ -0,0 +1,216 @@ +from dataclasses import dataclass + +import torch +from scipy.optimize import linear_sum_assignment +from torch.nn import functional as F + +from mlpf.logger import _logger +from mlpf.model.losses import LOSS_TASKS, REGRESSION_FEATURES + + +@dataclass(frozen=True) +class SetMatcherWeights: + presence: float = 1.0 + pid: float = 1.0 + pt: float = 1.0 + eta: float = 1.0 + phi: float = 1.0 + energy: float = 1.0 + + +def _pairwise_matching_cost(target, prediction, weights): + """Return the [num_slots, num_targets] detached matching cost.""" + + target_cls = target["cls_id"].long() + presence_cost = -F.log_softmax(prediction["cls_binary"].float(), dim=-1)[:, 1:2] + pid_cost = -F.log_softmax(prediction["cls_id_onehot"].float(), dim=-1)[ + :, target_cls + ] + + def l1_cost(feature): + return torch.abs( + prediction[feature].float()[:, None] - target[feature].float()[None, :] + ) + + pred_direction = F.normalize( + torch.stack([prediction["sin_phi"], prediction["cos_phi"]], dim=-1).float(), + dim=-1, + eps=1e-6, + ) + target_direction = F.normalize( + torch.stack([target["sin_phi"], target["cos_phi"]], dim=-1).float(), + dim=-1, + eps=1e-6, + ) + phi_cost = 1.0 - pred_direction @ target_direction.transpose(0, 1) + + return ( + weights.presence * presence_cost + + weights.pid * pid_cost + + weights.pt * l1_cost("pt") + + weights.eta * l1_cost("eta") + + weights.phi * phi_cost + + weights.energy * l1_cost("energy") + ).detach() + + +def hungarian_match(targets, predictions, target_mask, weights=None): + """Match particle slots to targets independently for each event.""" + + weights = weights or SetMatcherWeights() + matches = [] + num_slots = predictions["cls_binary"].shape[1] + for event_idx in range(predictions["cls_binary"].shape[0]): + valid = target_mask[event_idx].bool() + num_targets = int(valid.sum().item()) + if num_targets > num_slots: + raise ValueError( + f"Event {event_idx} has {num_targets} targets but the decoder has only {num_slots} slots" + ) + if num_targets == 0: + empty = torch.empty( + 0, dtype=torch.long, device=predictions["cls_binary"].device + ) + matches.append((empty, empty)) + continue + + event_targets = {key: value[event_idx][valid] for key, value in targets.items()} + event_predictions = { + key: value[event_idx] for key, value in predictions.items() + } + cost = _pairwise_matching_cost(event_targets, event_predictions, weights) + slot_indices, target_indices = linear_sum_assignment(cost.float().cpu().numpy()) + matches.append( + ( + torch.as_tensor(slot_indices, dtype=torch.long, device=cost.device), + torch.as_tensor(target_indices, dtype=torch.long, device=cost.device), + ) + ) + return matches + + +def set_event_loss( + targets, + predictions, + target_mask, + regression_weights, + matcher_weights=None, + no_object_weight=0.1, +): + """Permutation-invariant particle-set loss for a padded event batch.""" + + matches = hungarian_match(targets, predictions, target_mask, matcher_weights) + device = predictions["cls_binary"].device + presence_targets = torch.zeros( + predictions["cls_binary"].shape[:2], dtype=torch.long, device=device + ) + + matched_predictions = {key: [] for key in ("cls_id_onehot", *REGRESSION_FEATURES)} + matched_targets = {key: [] for key in ("cls_id", *REGRESSION_FEATURES)} + for event_idx, (slot_indices, target_indices) in enumerate(matches): + if len(slot_indices) == 0: + continue + presence_targets[event_idx, slot_indices] = 1 + valid_targets = target_mask[event_idx].bool() + for key in matched_predictions: + matched_predictions[key].append(predictions[key][event_idx, slot_indices]) + for key in matched_targets: + matched_targets[key].append( + targets[key][event_idx, valid_targets][target_indices] + ) + + presence_class_weights = predictions["cls_binary"].new_tensor( + [no_object_weight, 1.0] + ) + losses = { + "Classification_binary": 10.0 + * F.cross_entropy( + predictions["cls_binary"].reshape(-1, 2), + presence_targets.reshape(-1), + weight=presence_class_weights, + ) + } + + num_matched = int(presence_targets.sum().item()) + if num_matched == 0: + zero = predictions["cls_binary"].sum() * 0.0 + losses["Classification"] = zero + for feature in REGRESSION_FEATURES: + losses[f"Regression_{feature}"] = zero + return losses, matches + + matched_predictions = { + key: torch.cat(value, dim=0) for key, value in matched_predictions.items() + } + matched_targets = { + key: torch.cat(value, dim=0) for key, value in matched_targets.items() + } + losses["Classification"] = F.cross_entropy( + matched_predictions["cls_id_onehot"], matched_targets["cls_id"] + ) + + sqrt_target_pt = torch.sqrt( + torch.exp(matched_targets["pt"].float()).clamp_min(1e-6) + ) + for feature in REGRESSION_FEATURES: + prediction = torch.nan_to_num(matched_predictions[feature].float()) + per_particle = regression_weights[feature] * F.mse_loss( + prediction, matched_targets[feature].float(), reduction="none" + ) + losses[f"Regression_{feature}"] = ( + per_particle * sqrt_target_pt + ).sum() / num_matched + return losses, matches + + +def set_mlpf_loss( + targets, + predictions, + batch, + regression_weights, + task_loss_weighter=None, + matcher_weights=None, + no_object_weight=0.1, +): + """Compute the set-prediction objective with the standard task names.""" + + if batch.target_mask is None: + raise ValueError( + "Set prediction requires batch.ytarget_set and batch.target_mask" + ) + + effective_regression_weights = ( + regression_weights + if task_loss_weighter is None + else {feature: 1.0 for feature in REGRESSION_FEATURES} + ) + losses, _ = set_event_loss( + targets, + predictions, + batch.target_mask, + effective_regression_weights, + matcher_weights=matcher_weights, + no_object_weight=no_object_weight, + ) + if task_loss_weighter is None: + loss_opt = sum(losses.values()) + diagnostics = None + else: + # Keep the same task names so the existing one-time calibration can be + # evaluated for set mode rather than introducing a second mechanism. + assert tuple(losses) == LOSS_TASKS + loss_opt, diagnostics = task_loss_weighter(losses) + + losses["Total"] = loss_opt + if not torch.isfinite(loss_opt): + _logger.error(predictions) + _logger.error(losses) + raise RuntimeError("Set-prediction loss became non-finite") + + detached_losses = {key: value.detach() for key, value in losses.items()} + if diagnostics is not None: + diagnostics = { + name: {task: value.detach() for task, value in values.items()} + for name, values in diagnostics.items() + } + return loss_opt, detached_losses, diagnostics diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 5d64f04bd..05bb870c5 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -86,8 +86,10 @@ mlpf_loss, particle_loss, ) +from mlpf.model.set_losses import set_mlpf_loss +from mlpf.model.validation_metrics import compute_validation_particle_metrics, validation_particle_collections from mlpf.utils import create_comet_experiment -from mlpf.conf import INPUT_TYPE_LABELS, MLPFConfig, SOURCE_LABELS +from mlpf.conf import INPUT_TYPE_LABELS, MLPFConfig, OutputMode, SOURCE_LABELS from mlpf.jet_utils import get_jet_config UNIT_REGRESSION_WEIGHTS = {feature: 1.0 for feature in REGRESSION_FEATURES} @@ -115,8 +117,10 @@ def _log_batch_composition(batch, tensorboard_writer, step, prefix): tensorboard_writer.add_scalar(f"{prefix}/valid_elements_mean", valid_counts.float().mean().item(), step) tensorboard_writer.add_scalar(f"{prefix}/valid_elements_max", valid_counts.max().item(), step) - if batch.ytarget is not None: - target_counts = ((batch.ytarget[..., 0] != 0) & batch.mask).sum(dim=1).detach().to("cpu") + target_values = batch.ytarget_set if batch.target_mask is not None else batch.ytarget + target_valid = batch.target_mask if batch.target_mask is not None else batch.mask + if target_values is not None: + target_counts = ((target_values[..., 0] != 0) & target_valid).sum(dim=1).detach().to("cpu") tensorboard_writer.add_scalar(f"{prefix}/target_particles_mean", target_counts.float().mean().item(), step) tensorboard_writer.add_scalar(f"{prefix}/target_particles_max", target_counts.max().item(), step) @@ -131,7 +135,7 @@ def _log_batch_composition(batch, tensorboard_writer, step, prefix): def _add_accumulator(accum, key, value, count=1.0): - if count <= 0: + if count < 0: return if key not in accum: accum[key] = [torch.zeros((), device=value.device, dtype=torch.float32), torch.zeros((), device=value.device, dtype=torch.float32)] @@ -209,6 +213,15 @@ def _format_task_diagnostic(task_diagnostic): return " | ".join(f"{name}: {value:.4f}" for name, value in sorted(task_diagnostic.items())) +def _log_validation_results_to_tensorboard(tensorboard_writer, validation_results, step): + for name, value in validation_results.items(): + if name.startswith("metrics/"): + tag = f"validation/{name.removeprefix('metrics/')}" + else: + tag = f"step/loss_{name}" + tensorboard_writer.add_scalar(tag, value, step) + + def _format_compact_diagnostic(values, keys): if not values: return "" @@ -316,9 +329,17 @@ def model_step(batch, model, loss_fn, regression_weights): _logger.debug(f"model_step X={batch.X.shape}") ypred_raw = model(batch.X, batch.mask) ypred = unpack_predictions(ypred_raw) - ytarget = unpack_target(batch.ytarget, model) - - loss_opt, losses_detached, task_loss_diagnostics = loss_fn(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + model_module = model.module if hasattr(model, "module") else model + if model_module.output_mode == OutputMode.SET: + ytarget = unpack_target(batch.ytarget_set, model_module) + loss_opt, losses_detached, task_loss_diagnostics = set_mlpf_loss( + ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) + ) + else: + ytarget = unpack_target(batch.ytarget, model_module) + loss_opt, losses_detached, task_loss_diagnostics = loss_fn( + ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) + ) return loss_opt, losses_detached, task_loss_diagnostics, ypred_raw, ypred, ytarget @@ -424,8 +445,17 @@ def train_step( phase_start = time.perf_counter() with torch.autocast(device_type=device_type, dtype=dtype, enabled=device_type == "cuda"): ypred = unpack_predictions(ypred_raw) - ytarget = unpack_target(batch.ytarget, model) - loss_opt, loss, task_loss_diagnostics = mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + model_module = model.module if hasattr(model, "module") else model + if model_module.output_mode == OutputMode.SET: + ytarget = unpack_target(batch.ytarget_set, model_module) + loss_opt, loss, task_loss_diagnostics = set_mlpf_loss( + ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) + ) + else: + ytarget = unpack_target(batch.ytarget, model_module) + loss_opt, loss, task_loss_diagnostics = mlpf_loss( + ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) + ) phase_start = _record_phase_time_if_enabled(diagnostics.get("time", {}), "loss", phase_start, device_type, log_this_step) if log_this_step: _collect_step_memory(rank, "after_loss", diagnostics) @@ -723,11 +753,31 @@ def evaluate( model_module = model.module if hasattr(model, "module") else model ypred_particles = model_module.predict_particles(batch.X, batch.mask) + metric_collections = validation_particle_collections( + batch, + ypred_particles, + model_module.output_mode, + ) + particle_metrics = compute_validation_particle_metrics( + *metric_collections, + num_classes=config.num_classes, + ) + for metric_name, (metric_total, metric_count) in particle_metrics.items(): + _add_accumulator( + diagnostic_accum, + f"metrics/particle/{metric_name}", + torch.as_tensor(metric_total, device=batch.X.device), + count=metric_count, + ) - if ival == 0 and (rank == 0 or rank == "cpu"): + if model_module.output_mode == OutputMode.ELEMENTWISE and ival == 0 and (rank == 0 or rank == "cpu"): print_event_table(batch, ytarget, ypred_particles, config) - if config.validation_diagnostics_batches > 0 and ival < config.validation_diagnostics_batches: + if ( + model_module.output_mode == OutputMode.ELEMENTWISE + and config.validation_diagnostics_batches > 0 + and ival < config.validation_diagnostics_batches + ): _accumulate_domain_losses_and_stats( batch, ytarget, @@ -737,7 +787,12 @@ def evaluate( ) # Save validation plots for first batch - if (rank == 0 or rank == "cpu") and ival == 0 and config.make_plots: + if ( + model_module.output_mode == OutputMode.ELEMENTWISE + and (rank == 0 or rank == "cpu") + and ival == 0 + and config.make_plots + ): validation_plots(batch, ypred_raw, ytarget, ypred, tensorboard_writer, step, outdir) # Accumulate losses @@ -911,8 +966,7 @@ def _run_validation_cycle( stale_steps += 1 # Log validation losses to TensorBoard - for loss, value in losses_valid.items(): - tensorboard_writer_valid.add_scalar(f"step/loss_{loss}", value, step) + _log_validation_results_to_tensorboard(tensorboard_writer_valid, losses_valid, step) # Save step statistics to a JSON file history_path = Path(outdir) / "history" @@ -1234,6 +1288,7 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi pad_to_multiple=config.pad_to_multiple_elements, feature_dim=config.input_dim, max_open_readers=config.max_open_readers, + build_target_set=config.model.output_mode == OutputMode.SET, ).ds dataset.append(ds) ds = torch.utils.data.ConcatDataset(dataset) @@ -1247,6 +1302,8 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi sampler = torch.utils.data.SequentialSampler(ds) vals_for_test = ["X", "ytarget", "ytarget_pt_orig", "ytarget_e_orig", "ycand", "genjets", "targetjets"] + if config.model.output_mode == OutputMode.SET: + vals_for_test.append("ytarget_set") # pythia branch was introduced for cms in version 2.8.0 if sample.startswith("cms_") and version and Version(version) >= Version("2.8.0"): diff --git a/mlpf/model/validation_metrics.py b/mlpf/model/validation_metrics.py new file mode 100644 index 000000000..1acfeb5ba --- /dev/null +++ b/mlpf/model/validation_metrics.py @@ -0,0 +1,338 @@ +"""Scheme-independent particle metrics for validation. + +Matching minimizes a fixed combination of delta-R and absolute log-pT ratio. A +match is accepted for efficiency/purity when delta-R < 0.1 and relative pT error +< 0.5. Neither model's native target association or training cost is used. +""" + +import math + +import torch +from scipy.optimize import linear_sum_assignment + +from mlpf.conf import OutputMode +from mlpf.model.utils import unpack_target + + +MATCH_DR = 0.1 +MATCH_LOG_PT = math.log(2.0) +MATCH_REL_PT = 0.5 +STRICT_DR = 0.05 +STRICT_REL_PT = 0.2 + + +def validation_particle_collections(batch, predictions, output_mode): + """Return physical target/prediction collections with independent masks.""" + + if output_mode == OutputMode.SET: + targets = unpack_target(batch.ytarget_set.to(torch.float32), None) + targets["pt"] = torch.exp(targets["pt"].clamp(-20.0, 20.0)) + targets["energy"] = torch.exp(targets["energy"].clamp(-20.0, 20.0)) + target_mask = batch.target_mask.bool() + prediction_mask = predictions["cls_id"] != 0 + else: + targets = unpack_target(batch.ytarget.to(torch.float32), None) + if batch.ytarget_pt_orig is not None and batch.ytarget_e_orig is not None: + targets["pt"] = batch.ytarget_pt_orig.to(torch.float32) + targets["energy"] = batch.ytarget_e_orig.to(torch.float32) + else: + # Legacy/custom collaters may omit the cached absolute values. The + # elementwise targets store log(target/input), so reconstruct them + # without requiring the input dataset to be regenerated. + targets["pt"] = torch.exp(targets["pt"].clamp(-20.0, 20.0)) * batch.X[ + ..., 1 + ].to(torch.float32) + targets["energy"] = torch.exp( + targets["energy"].clamp(-20.0, 20.0) + ) * batch.X[..., 5].to(torch.float32) + target_mask = batch.mask.bool() & (targets["cls_id"] != 0) + prediction_mask = batch.mask.bool() & (predictions["cls_id"] != 0) + + return targets, target_mask, predictions, prediction_mask + + +def _empty_metrics(num_classes): + names = [ + "count/target_mean", + "count/prediction_mean", + "count/bias_mean", + "count/mae", + "matching/target_coverage_dr0p05", + "matching/target_coverage_dr0p10", + "matching/target_coverage_rel_pt0p20", + "matching/target_coverage_rel_pt0p50", + "matching/efficiency", + "matching/purity", + "matching/f1", + "matching/duplicate_fraction", + "matched/pid_accuracy", + "matched/delta_r_mean", + "matched/delta_eta_abs_mean", + "matched/delta_phi_abs_mean", + "matched/pt_relative_abs_mean", + "matched/energy_relative_abs_mean", + "event/energy_response_mean", + "event/energy_relative_abs_error", + "event/scalar_pt_response_mean", + "event/scalar_pt_relative_abs_error", + "event/vector_pt_closure", + "event/met_abs_error", + ] + for class_id in range(1, num_classes): + names.extend( + [ + f"class_{class_id}/efficiency", + f"class_{class_id}/pid_accuracy", + ] + ) + return {name: [0.0, 0.0] for name in names} + + +def _add(metrics, name, total, count): + metrics[name][0] += float(total) + metrics[name][1] += float(count) + + +def _clean_kinematics(collection, mask): + selected = {} + limits = { + "pt": (0.0, 1.0e6), + "eta": (-10.0, 10.0), + "energy": (0.0, 1.0e6), + "sin_phi": (-1.0, 1.0), + "cos_phi": (-1.0, 1.0), + } + for name, (minimum, maximum) in limits.items(): + value = collection[name][mask].detach().to(device="cpu", dtype=torch.float32) + selected[name] = torch.nan_to_num( + value, nan=0.0, posinf=maximum, neginf=minimum + ).clamp(minimum, maximum) + selected["cls_id"] = ( + collection["cls_id"][mask].detach().to(device="cpu", dtype=torch.long) + ) + selected["phi"] = torch.atan2(selected["sin_phi"], selected["cos_phi"]) + return selected + + +def _pairwise_geometry(targets, predictions): + delta_eta = predictions["eta"][:, None] - targets["eta"][None, :] + delta_phi = predictions["phi"][:, None] - targets["phi"][None, :] + delta_phi = torch.remainder(delta_phi + math.pi, 2.0 * math.pi) - math.pi + delta_r = torch.sqrt(delta_eta.square() + delta_phi.square()) + log_pt_ratio = torch.abs( + torch.log(predictions["pt"].clamp_min(1.0e-8))[:, None] + - torch.log(targets["pt"].clamp_min(1.0e-8))[None, :] + ) + relative_pt = ( + torch.abs(predictions["pt"][:, None] - targets["pt"][None, :]) + / targets["pt"].clamp_min(1.0e-8)[None, :] + ) + return delta_eta, delta_phi, delta_r, log_pt_ratio, relative_pt + + +def _accumulate_event_metrics(metrics, targets, predictions, num_classes): + num_targets = len(targets["pt"]) + num_predictions = len(predictions["pt"]) + _add(metrics, "count/target_mean", num_targets, 1) + _add(metrics, "count/prediction_mean", num_predictions, 1) + _add(metrics, "count/bias_mean", num_predictions - num_targets, 1) + _add(metrics, "count/mae", abs(num_predictions - num_targets), 1) + + matched_prediction_indices = torch.empty(0, dtype=torch.long) + matched_target_indices = torch.empty(0, dtype=torch.long) + delta_eta = delta_phi = delta_r = log_pt_ratio = pairwise_relative_pt = None + if num_targets and num_predictions: + delta_eta, delta_phi, delta_r, log_pt_ratio, pairwise_relative_pt = ( + _pairwise_geometry(targets, predictions) + ) + cost = (delta_r / MATCH_DR).square() + (log_pt_ratio / MATCH_LOG_PT).square() + prediction_indices, target_indices = linear_sum_assignment(cost.numpy()) + matched_prediction_indices = torch.as_tensor( + prediction_indices, dtype=torch.long + ) + matched_target_indices = torch.as_tensor(target_indices, dtype=torch.long) + + matched_dr = ( + delta_r[matched_prediction_indices, matched_target_indices] + if delta_r is not None + else torch.empty(0) + ) + relative_pt = ( + pairwise_relative_pt[matched_prediction_indices, matched_target_indices] + if pairwise_relative_pt is not None + else torch.empty(0) + ) + accepted = (matched_dr < MATCH_DR) & (relative_pt < MATCH_REL_PT) + num_accepted = int(accepted.sum()) + + _add( + metrics, + "matching/target_coverage_dr0p05", + int((matched_dr < STRICT_DR).sum()), + num_targets, + ) + _add( + metrics, + "matching/target_coverage_dr0p10", + int((matched_dr < MATCH_DR).sum()), + num_targets, + ) + + _add( + metrics, + "matching/target_coverage_rel_pt0p20", + int((relative_pt < STRICT_REL_PT).sum()), + num_targets, + ) + _add( + metrics, + "matching/target_coverage_rel_pt0p50", + int((relative_pt < MATCH_REL_PT).sum()), + num_targets, + ) + _add(metrics, "matching/efficiency", num_accepted, num_targets) + _add(metrics, "matching/purity", num_accepted, num_predictions) + _add(metrics, "matching/f1", 2 * num_accepted, num_targets + num_predictions) + + if delta_r is not None: + close_to_any_target = ( + (delta_r < MATCH_DR) & (pairwise_relative_pt < MATCH_REL_PT) + ).any(dim=1) + num_duplicates = max(int(close_to_any_target.sum()) - num_accepted, 0) + else: + num_duplicates = 0 + _add(metrics, "matching/duplicate_fraction", num_duplicates, num_predictions) + + accepted_prediction_indices = matched_prediction_indices[accepted] + accepted_target_indices = matched_target_indices[accepted] + if num_accepted: + accepted_target_pt = targets["pt"][accepted_target_indices] + accepted_prediction_pt = predictions["pt"][accepted_prediction_indices] + accepted_target_energy = targets["energy"][accepted_target_indices] + accepted_prediction_energy = predictions["energy"][accepted_prediction_indices] + accepted_relative_pt = torch.abs( + accepted_prediction_pt - accepted_target_pt + ) / accepted_target_pt.clamp_min(1.0e-8) + accepted_relative_energy = torch.abs( + accepted_prediction_energy - accepted_target_energy + ) / accepted_target_energy.clamp_min(1.0e-8) + pid_correct = ( + predictions["cls_id"][accepted_prediction_indices] + == targets["cls_id"][accepted_target_indices] + ) + + _add(metrics, "matched/pid_accuracy", pid_correct.sum(), num_accepted) + _add(metrics, "matched/delta_r_mean", matched_dr[accepted].sum(), num_accepted) + _add( + metrics, + "matched/delta_eta_abs_mean", + delta_eta[matched_prediction_indices, matched_target_indices][accepted] + .abs() + .sum(), + num_accepted, + ) + _add( + metrics, + "matched/delta_phi_abs_mean", + delta_phi[matched_prediction_indices, matched_target_indices][accepted] + .abs() + .sum(), + num_accepted, + ) + _add( + metrics, + "matched/pt_relative_abs_mean", + accepted_relative_pt.sum(), + num_accepted, + ) + _add( + metrics, + "matched/energy_relative_abs_mean", + accepted_relative_energy.sum(), + num_accepted, + ) + else: + pid_correct = torch.empty(0, dtype=torch.bool) + + for class_id in range(1, num_classes): + class_target_count = int((targets["cls_id"] == class_id).sum()) + if num_accepted: + accepted_in_class = targets["cls_id"][accepted_target_indices] == class_id + class_accepted_count = int(accepted_in_class.sum()) + class_pid_correct = int((accepted_in_class & pid_correct).sum()) + else: + class_accepted_count = 0 + class_pid_correct = 0 + _add( + metrics, + f"class_{class_id}/efficiency", + class_accepted_count, + class_target_count, + ) + _add( + metrics, + f"class_{class_id}/pid_accuracy", + class_pid_correct, + class_accepted_count, + ) + + target_energy = targets["energy"].to(torch.float64).sum() + prediction_energy = predictions["energy"].to(torch.float64).sum() + target_scalar_pt = targets["pt"].to(torch.float64).sum() + prediction_scalar_pt = predictions["pt"].to(torch.float64).sum() + energy_denominator = target_energy.clamp_min(1.0e-8) + pt_denominator = target_scalar_pt.clamp_min(1.0e-8) + energy_residual = (prediction_energy - target_energy) / energy_denominator + scalar_pt_residual = (prediction_scalar_pt - target_scalar_pt) / pt_denominator + _add( + metrics, "event/energy_response_mean", prediction_energy / energy_denominator, 1 + ) + _add(metrics, "event/energy_relative_abs_error", energy_residual.abs(), 1) + _add( + metrics, + "event/scalar_pt_response_mean", + prediction_scalar_pt / pt_denominator, + 1, + ) + _add(metrics, "event/scalar_pt_relative_abs_error", scalar_pt_residual.abs(), 1) + + target_px = ( + targets["pt"].to(torch.float64) * torch.cos(targets["phi"].to(torch.float64)) + ).sum() + target_py = ( + targets["pt"].to(torch.float64) * torch.sin(targets["phi"].to(torch.float64)) + ).sum() + prediction_px = ( + predictions["pt"].to(torch.float64) + * torch.cos(predictions["phi"].to(torch.float64)) + ).sum() + prediction_py = ( + predictions["pt"].to(torch.float64) + * torch.sin(predictions["phi"].to(torch.float64)) + ).sum() + vector_pt_error = torch.hypot(prediction_px - target_px, prediction_py - target_py) + target_met = torch.hypot(target_px, target_py) + prediction_met = torch.hypot(prediction_px, prediction_py) + _add(metrics, "event/vector_pt_closure", vector_pt_error / pt_denominator, 1) + _add(metrics, "event/met_abs_error", torch.abs(prediction_met - target_met), 1) + + +def compute_validation_particle_metrics( + targets, target_mask, predictions, prediction_mask, num_classes +): + """Compute additive, scheme-independent particle metrics for one batch.""" + + metrics = _empty_metrics(num_classes) + for event_idx in range(target_mask.shape[0]): + event_targets = _clean_kinematics( + {name: value[event_idx] for name, value in targets.items()}, + target_mask[event_idx], + ) + event_predictions = _clean_kinematics( + {name: value[event_idx] for name, value in predictions.items()}, + prediction_mask[event_idx], + ) + _accumulate_event_metrics( + metrics, event_targets, event_predictions, num_classes + ) + return {name: tuple(values) for name, values in metrics.items()} diff --git a/particleflow_spec.yaml b/particleflow_spec.yaml index 9b18b20e5..de43cd2d9 100644 --- a/particleflow_spec.yaml +++ b/particleflow_spec.yaml @@ -589,22 +589,25 @@ models: version: "3.2.1" splits: ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] - pyg-cld-hits-v1: + pyg-cld-hits-v1: &pyg_cld_hits_v1 <<: *defaults dataset: cld_hits - gpu_batch_multiplier: 12 + gpu_batch_multiplier: 1 + pad_to_multiple_elements: 128 hyperparameters: batch_size: 1 lr: 0.0001 - architecture: + architecture: &pyg_cld_hits_architecture type: "attention" input_encoding: "split" + output_mode: "elementwise" attention: num_convs: 3 head_dim: 16 num_heads: 16 + use_jagged_attention: true # Dataset Selection train_datasets: @@ -646,6 +649,25 @@ models: version: "3.2.1" splits: ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] + pyg-cld-hits-set-v1: + <<: *pyg_cld_hits_v1 + gpu_batch_multiplier: 1 + pad_to_multiple_elements: 128 + + hyperparameters: + batch_size: 1 + lr: 0.0001 + + architecture: + <<: *pyg_cld_hits_architecture + output_mode: "set" + set_decoder: + num_slots: 256 + num_layers: 2 + num_heads: 8 + ffn_multiplier: 4.0 + dropout: 0.0 + # CLIC Model pyg-clic-v1: <<: *defaults diff --git a/scripts/local/make_local_available_spec.py b/scripts/local/make_local_available_spec.py index b7da0cf6d..481b49153 100755 --- a/scripts/local/make_local_available_spec.py +++ b/scripts/local/make_local_available_spec.py @@ -17,7 +17,7 @@ def main(): parser = argparse.ArgumentParser(description="Restrict CLD/CLIC models to locally available ttbar datasets.") parser.add_argument("input_spec", type=Path) parser.add_argument("output_spec", type=Path) - parser.add_argument("--hit-version", default="3.2.0") + parser.add_argument("--hit-version", default="3.2.1") parser.add_argument("--hit-splits", nargs="+", default=["1"]) parser.add_argument("--pf-version", default="3.2.0") parser.add_argument("--pf-splits", nargs="+", default=[str(i) for i in range(1, 11)]) @@ -31,6 +31,7 @@ def main(): local_datasets = { "pyg-cld-hits-v1": ("cld_hits", "cld_edm_ttbar_hits", args.hit_version, args.hit_splits), + "pyg-cld-hits-set-v1": ("cld_hits", "cld_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-clic-hits-v1": ("clic_hits", "clic_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-cld-v1": ("cld", "cld_edm_ttbar_pf", args.pf_version, args.pf_splits), "pyg-clic-v1": ("clic", "clic_edm_ttbar_pf", args.pf_version, args.pf_splits), diff --git a/scripts/local/train.sh b/scripts/local/train.sh index 7b0d39d05..f54654923 100755 --- a/scripts/local/train.sh +++ b/scripts/local/train.sh @@ -5,54 +5,52 @@ export PF_SITE=local SPEC_FILE=${SPEC_FILE:-particleflow_spec.yaml} USE_LOCAL_AVAILABLE_SPEC=${USE_LOCAL_AVAILABLE_SPEC:-true} -LOCAL_SPEC_FILE=${LOCAL_SPEC_FILE:-/tmp/particleflow_local_available_spec.yaml} -TARGETS=${TARGETS:-cld-hits} -DATA_CONFIG=${DATA_CONFIG:-1} +LOCAL_SPEC_FILE=${LOCAL_SPEC_FILE:-/tmp/particleflow_local_ttbar_comparison_spec.yaml} +OUTPUT_MODES=${OUTPUT_MODES:-elementwise,set} +HIT_VERSION=${HIT_VERSION:-3.2.1} +HIT_SPLITS=${HIT_SPLITS:-1} +DATA_CONFIG=${DATA_CONFIG:-${HIT_SPLITS// /,}} NUM_STEPS=${NUM_STEPS:-2000} -VAL_FREQ=${VAL_FREQ:-1000} -CHECKPOINT_FREQ=${CHECKPOINT_FREQ:-1000} -GPU_BATCH_MULTIPLIER=${GPU_BATCH_MULTIPLIER:-4} +VAL_FREQ=${VAL_FREQ:-200} +CHECKPOINT_FREQ=${CHECKPOINT_FREQ:-200} +NVALID=${NVALID:-100} +NTEST=${NTEST:-100} +GPU_BATCH_MULTIPLIER=${GPU_BATCH_MULTIPLIER:-8} NUM_WORKERS=${NUM_WORKERS:-8} PREFETCH_FACTOR=${PREFETCH_FACTOR:-4} VALIDATION_DIAGNOSTICS_BATCHES=${VALIDATION_DIAGNOSTICS_BATCHES:-4} EXPERIMENTS_DIR=${EXPERIMENTS_DIR:-experiments} -NUM_TRACKER_LAYERS=${NUM_TRACKER_LAYERS:-2} -NUM_CALO_LAYERS=${NUM_CALO_LAYERS:-2} -NUM_COMMON_LAYERS=${NUM_COMMON_LAYERS:-2} +PAD_TO_MULTIPLE_ELEMENTS=${PAD_TO_MULTIPLE_ELEMENTS:-128} -IFS=',' read -r -a TARGET_LIST <<< "$TARGETS" +IFS=',' read -r -a OUTPUT_MODE_LIST <<< "$OUTPUT_MODES" +read -r -a HIT_SPLIT_LIST <<< "$HIT_SPLITS" if [[ "$USE_LOCAL_AVAILABLE_SPEC" == "true" ]]; then - uv run python3 scripts/local/make_local_available_spec.py "$SPEC_FILE" "$LOCAL_SPEC_FILE" + uv run python3 scripts/local/make_local_available_spec.py \ + "$SPEC_FILE" "$LOCAL_SPEC_FILE" \ + --hit-version "$HIT_VERSION" \ + --hit-splits "${HIT_SPLIT_LIST[@]}" SPEC_FILE="$LOCAL_SPEC_FILE" fi -set_target() { - local target=$1 - case "$target" in - cld-hits) +PRODUCTION_NAME=cld +DATA_DIR=${DATA_DIR:-$(uv run python3 scripts/get_param.py "$SPEC_FILE" productions."$PRODUCTION_NAME".workspace_dir)/tfds/} + +set_output_mode() { + local output_mode=$1 + case "$output_mode" in + elementwise) MODEL_NAME=pyg-cld-hits-v1 - PRODUCTION_NAME=cld - ;; - clic-hits) - MODEL_NAME=pyg-clic-hits-v1 - PRODUCTION_NAME=clic - ;; - cld-pf) - MODEL_NAME=pyg-cld-v1 - PRODUCTION_NAME=cld ;; - clic-pf) - MODEL_NAME=pyg-clic-v1 - PRODUCTION_NAME=clic + set) + MODEL_NAME=pyg-cld-hits-set-v1 ;; *) - echo "Unknown target '$target'. Valid targets: cld-hits, clic-hits, cld-pf, clic-pf" >&2 + echo "Unknown output mode '$output_mode'. Valid modes: elementwise, set" >&2 exit 1 ;; esac - DATA_DIR=$(python3 scripts/get_param.py "$SPEC_FILE" productions."$PRODUCTION_NAME".workspace_dir)/tfds/ } make_common_args() { @@ -68,37 +66,27 @@ make_common_args() { --val_freq "$VAL_FREQ" --checkpoint_freq "$CHECKPOINT_FREQ" --num_steps "$NUM_STEPS" + --nvalid "$NVALID" + --ntest "$NTEST" --num_workers "$NUM_WORKERS" --prefetch_factor "$PREFETCH_FACTOR" --sampler_mode interleaved-shards --validation_diagnostics_batches "$VALIDATION_DIAGNOSTICS_BATCHES" --make_plots - --model.attention.use_jagged_attention false - --pad_to_multiple_elements 100 + --pad_to_multiple_elements "$PAD_TO_MULTIPLE_ELEMENTS" ) } -run_detector_scenario() { - local target=$1 - if [[ "$target" != "cld-hits" && "$target" != "clic-hits" ]]; then - echo "Detector-specific training is only valid for cld-hits and clic-hits" >&2 - exit 1 - fi - local num_detector_layers=$((NUM_TRACKER_LAYERS + NUM_CALO_LAYERS + NUM_COMMON_LAYERS)) +run_comparison_training() { + local output_mode=$1 + echo "Starting CLD ttbar hit training with output_mode=$output_mode model=$MODEL_NAME" uv run python3 mlpf/pipeline.py \ - --prefix "${target}_detector-backbone_" \ - "${COMMON_ARGS[@]}" \ - --model.backbone.mode shared \ - --model.backbone.num_convs "$num_detector_layers" \ - --model.backbone.num_tracker_layers "$NUM_TRACKER_LAYERS" \ - --model.backbone.num_calo_layers "$NUM_CALO_LAYERS" \ - --model.backbone.num_common_layers "$NUM_COMMON_LAYERS" \ - --model.attention.use_jagged_attention true \ - --model.task_queries false + --prefix "ttbar-${output_mode}_" \ + "${COMMON_ARGS[@]}" } -for target in "${TARGET_LIST[@]}"; do - set_target "$target" +for output_mode in "${OUTPUT_MODE_LIST[@]}"; do + set_output_mode "$output_mode" make_common_args - run_detector_scenario "$target" + run_comparison_training "$output_mode" done diff --git a/tests/test_pfdataset_logic.py b/tests/test_pfdataset_logic.py index be4cd0985..20f283259 100644 --- a/tests/test_pfdataset_logic.py +++ b/tests/test_pfdataset_logic.py @@ -112,6 +112,71 @@ def test_padding(self): self.assertEqual(ret["ytarget"].shape[0], 4) self.assertEqual(ret["X"][1, 0], 0) # Padded with zero + def test_compact_set_target_extraction(self): + ytarget = np.zeros((4, 14), dtype=np.float32) + ytarget[0, [0, 2, 5, 6, 13]] = [1, 20.0, 1.0, 40.0, 1] + ytarget[1, 13] = 1 # related hit carrying only particle_number + ytarget[2, [0, 2, 5, 6, 13]] = [2, 5.0, 1.0, 8.0, 2] + data = { + "X": np.array( + [ + [1, 10.0, 0.0, 0.0, 1.0, 10.0], + [1, 8.0, 0.0, 0.0, 1.0, 8.0], + [2, 4.0, 0.0, 0.0, 1.0, 4.0], + [2, 1.0, 0.0, 0.0, 1.0, 1.0], + ], + dtype=np.float32, + ), + "ytarget": ytarget, + "ycand": np.zeros_like(ytarget), + } + ds = TFDSDataSource(MockTFDS([data], name="cld_hits"), sort=False, build_target_set=True) + + ret = ds[0] + + self.assertEqual(ret["ytarget_set"].shape, (2, 14)) + np.testing.assert_array_equal(ret["ytarget_set"][:, 0], [1, 2]) + np.testing.assert_array_equal(ret["ytarget_set"][:, 13], [1, 2]) + np.testing.assert_allclose(ret["ytarget_set"][:, 2], np.log([20.0, 5.0])) + np.testing.assert_allclose(ret["ytarget_set"][:, 6], np.log([40.0, 8.0])) + self.assertAlmostEqual(ret["ytarget"][0, 2], np.log(2.0)) + + def test_compact_set_target_rejects_duplicate_particle_numbers(self): + ytarget = np.zeros((2, 14), dtype=np.float32) + ytarget[:, 0] = [1, 2] + ytarget[:, 2] = [1.0, 2.0] + ytarget[:, 6] = [1.0, 2.0] + ytarget[:, 13] = 1 + data = { + "X": np.array([[1, 1, 0, 0, 1, 1], [2, 2, 0, 0, 1, 2]], dtype=np.float32), + "ytarget": ytarget, + "ycand": np.zeros_like(ytarget), + } + ds = TFDSDataSource(MockTFDS([data], name="cld_hits"), sort=False, build_target_set=True) + + with self.assertRaisesRegex(ValueError, "unique, nonzero particle_number"): + ds[0] + + +def test_collater_pads_inputs_and_set_targets_independently(): + collater = Collater(["X", "ytarget", "ytarget_set"], []) + item1 = { + "X": np.ones((4, 2), dtype=np.float32), + "ytarget": np.ones((4, 2), dtype=np.float32), + "ytarget_set": np.ones((2, 2), dtype=np.float32), + } + item2 = { + "X": np.ones((2, 2), dtype=np.float32), + "ytarget": np.ones((2, 2), dtype=np.float32), + "ytarget_set": np.ones((1, 2), dtype=np.float32), + } + + batch = collater([item1, item2]) + + assert batch.X.shape == (2, 4, 2) + assert batch.ytarget_set.shape == (2, 2, 2) + torch.testing.assert_close(batch.target_mask, torch.tensor([[True, True], [True, False]])) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py new file mode 100644 index 000000000..9bc1c3380 --- /dev/null +++ b/tests/test_set_prediction.py @@ -0,0 +1,234 @@ +import pytest +import torch + +from mlpf.conf import MLPFConfig +from mlpf.model.PFDataset import PFBatch +from mlpf.model.mlpf import MLPF +from mlpf.model.set_losses import hungarian_match, set_event_loss +from mlpf.model.utils import unpack_predictions, unpack_target + + +REGRESSION_WEIGHTS = { + feature: 1.0 for feature in ("pt", "eta", "sin_phi", "cos_phi", "energy") +} + + +def make_config(num_slots=4): + return MLPFConfig.model_validate( + { + "dataset": "cld_hits", + "data_dir": "/tmp", + "model": { + "type": "heptv2", + "output_mode": "set", + "input_encoding": "joint", + "heptv2": { + "num_convs": 0, + "num_heads": 2, + "embedding_dim": 16, + "width": 16, + "block_size": 8, + }, + "set_decoder": { + "num_slots": num_slots, + "num_layers": 2, + "num_heads": 2, + }, + "hit_feature_engineering": {"enabled": False}, + }, + "conv_type": "heptv2", + } + ) + + +def make_target_tensor(batch_size=1, num_targets=2): + target = torch.zeros(batch_size, num_targets, 14) + phi = torch.linspace(-torch.pi, torch.pi, num_targets + 1)[:-1] + pt = torch.linspace(1.0, 20.0, num_targets) + target[..., 0] = (torch.arange(num_targets) % 5) + 1 + target[..., 2] = torch.log(pt) + target[..., 3] = torch.linspace(-2.0, 2.0, num_targets) + target[..., 4] = torch.sin(phi) + target[..., 5] = torch.cos(phi) + target[..., 6] = torch.log(pt + 2.0) + target[..., 13] = torch.arange(1, num_targets + 1) + return target + + +def test_set_config_populates_decoder_defaults(): + config = MLPFConfig.model_validate( + { + "dataset": "cld_hits", + "data_dir": "/tmp", + "model": {"type": "heptv2", "output_mode": "set", "heptv2": {}}, + "conv_type": "heptv2", + } + ) + + assert config.model.set_decoder is not None + assert config.model.set_decoder.num_slots == 256 + + +def test_set_config_rejects_non_hit_datasets(): + with pytest.raises(ValueError, match="supported only for CLD/CLIC hit datasets"): + MLPFConfig.model_validate( + { + "dataset": "cld", + "data_dir": "/tmp", + "model": {"type": "heptv2", "output_mode": "set", "heptv2": {}}, + "conv_type": "heptv2", + } + ) + + +def test_set_model_output_axis_is_num_slots(): + config = make_config(num_slots=4) + model = MLPF(config) + X = torch.randn(2, 11, config.input_dim) + X[..., 0] = 1 + X[..., 1] = X[..., 1].abs() + 0.1 + X[..., 5] = X[..., 5].abs() + 0.1 + mask = torch.ones(2, 11, dtype=torch.bool) + mask[1, 7:] = False + + presence, pid, momentum, pileup = model(X, mask) + + assert presence.shape == (2, 4, 2) + assert pid.shape == (2, 4, config.num_classes) + assert momentum.shape == (2, 4, 5) + assert pileup.shape == (2, 4, 2) + torch.testing.assert_close( + torch.linalg.vector_norm(momentum[..., 2:4], dim=-1), torch.ones(2, 4) + ) + + +def test_hungarian_match_finds_permuted_particles(): + ytarget_tensor = make_target_tensor() + targets = unpack_target(ytarget_tensor, None) + predictions = { + "cls_binary": torch.tensor([[[-5.0, 5.0], [-5.0, 5.0]]]), + "cls_id_onehot": torch.tensor( + [[[-5.0, -5.0, 5.0, -5.0, -5.0, -5.0], [-5.0, 5.0, -5.0, -5.0, -5.0, -5.0]]] + ), + "pt": targets["pt"].flip(1), + "eta": targets["eta"].flip(1), + "sin_phi": targets["sin_phi"].flip(1), + "cos_phi": targets["cos_phi"].flip(1), + "energy": targets["energy"].flip(1), + } + + matches = hungarian_match(targets, predictions, torch.ones(1, 2, dtype=torch.bool)) + + slot_indices, target_indices = matches[0] + assert slot_indices.tolist() == [0, 1] + assert target_indices.tolist() == [1, 0] + + +def test_set_loss_is_target_permutation_invariant(): + torch.manual_seed(3) + target_tensor = make_target_tensor() + batch = PFBatch(X=torch.ones(1, 5, 15), ytarget_set=target_tensor) + predictions = { + "cls_binary": torch.randn(1, 4, 2, requires_grad=True), + "cls_id_onehot": torch.randn(1, 4, 6, requires_grad=True), + "pt": torch.randn(1, 4, requires_grad=True), + "eta": torch.randn(1, 4, requires_grad=True), + "sin_phi": torch.randn(1, 4, requires_grad=True), + "cos_phi": torch.randn(1, 4, requires_grad=True), + "energy": torch.randn(1, 4, requires_grad=True), + } + targets = unpack_target(target_tensor, None) + losses, _ = set_event_loss( + targets, predictions, batch.target_mask, REGRESSION_WEIGHTS + ) + + permutation = torch.tensor([1, 0]) + permuted_tensor = target_tensor[:, permutation] + permuted_targets = unpack_target(permuted_tensor, None) + permuted_mask = permuted_tensor[..., 0] != 0 + permuted_losses, _ = set_event_loss( + permuted_targets, predictions, permuted_mask, REGRESSION_WEIGHTS + ) + + for key in losses: + torch.testing.assert_close(losses[key], permuted_losses[key]) + sum(losses.values()).backward() + assert predictions["cls_binary"].grad is not None + + +def test_set_loss_rejects_target_overflow(): + targets_tensor = make_target_tensor(num_targets=2) + targets = unpack_target(targets_tensor, None) + predictions = { + "cls_binary": torch.zeros(1, 1, 2), + "cls_id_onehot": torch.zeros(1, 1, 6), + "pt": torch.zeros(1, 1), + "eta": torch.zeros(1, 1), + "sin_phi": torch.zeros(1, 1), + "cos_phi": torch.ones(1, 1), + "energy": torch.zeros(1, 1), + } + + with pytest.raises(ValueError, match="2 targets.*1 slots"): + hungarian_match(targets, predictions, torch.ones(1, 2, dtype=torch.bool)) + + +def test_set_loss_supports_an_event_without_targets(): + target_tensor = torch.zeros(1, 0, 14) + batch = PFBatch(X=torch.ones(1, 3, 15), ytarget_set=target_tensor) + targets = unpack_target(target_tensor, None) + predictions = { + "cls_binary": torch.randn(1, 4, 2, requires_grad=True), + "cls_id_onehot": torch.randn(1, 4, 6, requires_grad=True), + "pt": torch.randn(1, 4, requires_grad=True), + "eta": torch.randn(1, 4, requires_grad=True), + "sin_phi": torch.randn(1, 4, requires_grad=True), + "cos_phi": torch.randn(1, 4, requires_grad=True), + "energy": torch.randn(1, 4, requires_grad=True), + } + + losses, matches = set_event_loss( + targets, predictions, batch.target_mask, REGRESSION_WEIGHTS + ) + loss = sum(losses.values()) + loss.backward() + + assert torch.isfinite(loss) + assert matches[0][0].numel() == 0 + assert losses["Classification"] == 0 + assert losses["Regression_pt"] == 0 + + +def test_predict_particles_restores_absolute_set_kinematics(): + model = MLPF(make_config(num_slots=3)).eval() + X = torch.ones(1, 6, 15) + with torch.no_grad(): + prediction = model.predict_particles(X, torch.ones(1, 6, dtype=torch.bool)) + + assert prediction["pt"].shape == (1, 3) + assert prediction["energy"].shape == (1, 3) + assert torch.all(prediction["pt"] >= 0) + assert torch.all(prediction["energy"] >= 0) + + +def test_set_model_10k_inputs_forward_backward(): + torch.manual_seed(7) + config = make_config(num_slots=256) + model = MLPF(config) + X = torch.randn(1, 10_000, config.input_dim) + X[..., 0] = 1 + X[..., 1] = X[..., 1].abs() + 0.1 + X[..., 5] = X[..., 5].abs() + 0.1 + target_tensor = make_target_tensor(num_targets=100) + batch = PFBatch(X=X, ytarget_set=target_tensor) + + raw_predictions = model(batch.X, batch.mask) + predictions = unpack_predictions(raw_predictions) + targets = unpack_target(batch.ytarget_set, model) + losses, _ = set_event_loss(targets, predictions, batch.target_mask, REGRESSION_WEIGHTS) + loss = sum(losses.values()) + loss.backward() + + assert torch.isfinite(loss) + assert model.set_decoder.queries.grad is not None + assert torch.isfinite(model.set_decoder.queries.grad).all() diff --git a/tests/test_training_diagnostics.py b/tests/test_training_diagnostics.py index fb59dad76..51ad469ef 100644 --- a/tests/test_training_diagnostics.py +++ b/tests/test_training_diagnostics.py @@ -1,7 +1,12 @@ import torch from mlpf.model.PFDataset import PFBatch -from mlpf.model.training import _accumulate_domain_losses_and_stats, _event_domain_labels, _finalize_diagnostics +from mlpf.model.training import ( + _accumulate_domain_losses_and_stats, + _event_domain_labels, + _finalize_diagnostics, + _log_validation_results_to_tensorboard, +) def make_batch(): @@ -78,3 +83,29 @@ def test_domain_loss_and_regression_diagnostics_are_grouped(): assert metrics[f"diagnostic/regression/{label}/energy_residual_mean"] == -0.5 assert metrics[f"diagnostic/regression/{label}/pt_residual_rms"] == 0.25 assert metrics[f"diagnostic/regression/{label}/energy_residual_rms"] == 0.5 + + +def test_common_validation_metrics_use_scheme_independent_tensorboard_tags(): + class RecordingWriter: + def __init__(self): + self.scalars = [] + + def add_scalar(self, tag, value, step): + self.scalars.append((tag, value, step)) + + writer = RecordingWriter() + _log_validation_results_to_tensorboard( + writer, + { + "Total": 3.0, + "metrics/particle/matching/f1": 0.75, + "metrics/particle/count/mae": 2.0, + }, + step=100, + ) + + assert writer.scalars == [ + ("step/loss_Total", 3.0, 100), + ("validation/particle/matching/f1", 0.75, 100), + ("validation/particle/count/mae", 2.0, 100), + ] diff --git a/tests/test_validation_metrics.py b/tests/test_validation_metrics.py new file mode 100644 index 000000000..71ed74f25 --- /dev/null +++ b/tests/test_validation_metrics.py @@ -0,0 +1,164 @@ +import pytest +import torch + +from mlpf.conf import OutputMode +from mlpf.model.PFDataset import PFBatch +from mlpf.model.validation_metrics import ( + compute_validation_particle_metrics, + validation_particle_collections, +) + + +def make_collection(cls_id, pt, eta, phi, energy): + phi = torch.tensor([phi], dtype=torch.float32) + return { + "cls_id": torch.tensor([cls_id], dtype=torch.long), + "pt": torch.tensor([pt], dtype=torch.float32), + "eta": torch.tensor([eta], dtype=torch.float32), + "sin_phi": torch.sin(phi), + "cos_phi": torch.cos(phi), + "energy": torch.tensor([energy], dtype=torch.float32), + } + + +def finalized(metrics): + return {name: total / count for name, (total, count) in metrics.items() if count} + + +def test_common_particle_metrics_match_permuted_particles_and_find_duplicate(): + targets = make_collection( + cls_id=[1, 2], + pt=[10.0, 20.0], + eta=[0.0, 1.0], + phi=[0.0, 0.5], + energy=[12.0, 24.0], + ) + predictions = make_collection( + cls_id=[2, 3, 1], + pt=[20.0, 10.0, 10.0], + eta=[1.0, 0.0, 0.02], + phi=[0.5, 0.0, 0.0], + energy=[24.0, 12.0, 12.0], + ) + target_mask = torch.ones(1, 2, dtype=torch.bool) + prediction_mask = torch.ones(1, 3, dtype=torch.bool) + + values = finalized( + compute_validation_particle_metrics( + targets, + target_mask, + predictions, + prediction_mask, + num_classes=6, + ) + ) + + assert values["count/target_mean"] == 2 + assert values["count/prediction_mean"] == 3 + assert values["count/bias_mean"] == 1 + assert values["count/mae"] == 1 + assert values["matching/efficiency"] == 1 + assert values["matching/purity"] == pytest.approx(2 / 3) + assert values["matching/f1"] == pytest.approx(0.8) + assert values["matching/duplicate_fraction"] == pytest.approx(1 / 3) + assert values["matched/pid_accuracy"] == 0.5 + assert values["matched/delta_r_mean"] == 0 + assert values["event/energy_response_mean"] == pytest.approx(4 / 3) + + +def test_common_particle_metrics_are_prediction_permutation_invariant(): + targets = make_collection( + cls_id=[1, 2], + pt=[10.0, 20.0], + eta=[0.0, 1.0], + phi=[0.0, 0.5], + energy=[12.0, 24.0], + ) + predictions = make_collection( + cls_id=[2, 1], + pt=[20.0, 10.0], + eta=[1.0, 0.0], + phi=[0.5, 0.0], + energy=[24.0, 12.0], + ) + masks = (torch.ones(1, 2, dtype=torch.bool),) * 2 + original = compute_validation_particle_metrics( + targets, masks[0], predictions, masks[1], num_classes=6 + ) + + permutation = torch.tensor([1, 0]) + permuted = {name: value[:, permutation] for name, value in predictions.items()} + reordered = compute_validation_particle_metrics( + targets, masks[0], permuted, masks[1], num_classes=6 + ) + + assert original == reordered + + +def test_validation_collections_restore_same_physical_targets_for_both_modes(): + element_target = torch.zeros(1, 3, 14) + element_target[0, :2, 0] = torch.tensor([1, 2]) + element_target[0, :2, 3] = torch.tensor([0.1, -0.2]) + element_target[0, :2, 4] = torch.sin(torch.tensor([0.3, -0.4])) + element_target[0, :2, 5] = torch.cos(torch.tensor([0.3, -0.4])) + element_batch = PFBatch( + X=torch.ones(1, 3, 15), + ytarget=element_target, + ytarget_pt_orig=torch.tensor([[10.0, 20.0, 0.0]]), + ytarget_e_orig=torch.tensor([[12.0, 24.0, 0.0]]), + ) + + set_target = element_target[:, :2].clone() + set_target[..., 2] = torch.log(torch.tensor([[10.0, 20.0]])) + set_target[..., 6] = torch.log(torch.tensor([[12.0, 24.0]])) + set_batch = PFBatch(X=torch.ones(1, 3, 15), ytarget_set=set_target) + predictions = make_collection( + cls_id=[1, 2, 0], + pt=[10.0, 20.0, 0.0], + eta=[0.1, -0.2, 0.0], + phi=[0.3, -0.4, 0.0], + energy=[12.0, 24.0, 0.0], + ) + + element_collections = validation_particle_collections( + element_batch, predictions, OutputMode.ELEMENTWISE + ) + set_collections = validation_particle_collections( + set_batch, predictions, OutputMode.SET + ) + + torch.testing.assert_close( + element_collections[0]["pt"][:, :2], set_collections[0]["pt"] + ) + torch.testing.assert_close( + element_collections[0]["energy"][:, :2], set_collections[0]["energy"] + ) + assert element_collections[1].sum() == set_collections[1].sum() == 2 + assert element_collections[3].sum() == set_collections[3].sum() == 2 + + +def test_elementwise_validation_reconstructs_targets_without_cached_values(): + X = torch.ones(1, 2, 15) + X[..., 1] = torch.tensor([[2.0, 4.0]]) + X[..., 5] = torch.tensor([[3.0, 6.0]]) + target = torch.zeros(1, 2, 14) + target[..., 0] = torch.tensor([[1, 2]]) + target[..., 2] = torch.log(torch.tensor([[5.0, 2.0]])) + target[..., 4] = 0.0 + target[..., 5] = 1.0 + target[..., 6] = torch.log(torch.tensor([[4.0, 3.0]])) + batch = PFBatch(X=X, ytarget=target) + predictions = make_collection( + cls_id=[1, 2], + pt=[10.0, 8.0], + eta=[0.0, 0.0], + phi=[0.0, 0.0], + energy=[12.0, 18.0], + ) + + targets, _, _, _ = validation_particle_collections( + batch, predictions, OutputMode.ELEMENTWISE + ) + + torch.testing.assert_close(targets["pt"], torch.tensor([[10.0, 8.0]])) + torch.testing.assert_close(targets["energy"], torch.tensor([[12.0, 18.0]])) From 2dca3a14e5e974ac52c5f7b618e1c7be80ab1715 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 12:22:09 +0300 Subject: [PATCH 03/29] Add reusable seeded training scenario launchers --- DOING.md | 2 + configs/training/README.md | 54 ++ configs/training/platforms/flatiron_a100.yaml | 16 + configs/training/platforms/flatiron_b200.yaml | 16 + configs/training/platforms/flatiron_h100.yaml | 16 + configs/training/platforms/local.yaml | 11 + .../scenarios/cld_hits_output_comparison.yaml | 41 ++ docs/map.md | 3 +- mlpf/conf.py | 1 + mlpf/model/PFDataset.py | 17 +- mlpf/model/distributed_ray.py | 4 +- mlpf/model/training.py | 25 +- mlpf/training_scenarios.py | 526 ++++++++++++++++++ mlpf/training_submission.py | 163 ++++++ particleflow_spec.yaml | 1 + scripts/flatiron/run_uv_scenario.sh | 55 ++ scripts/flatiron/submit_scenario.py | 7 + scripts/flatiron/train_a100_uv.sh | 76 --- scripts/flatiron/train_h100_uv.sh | 81 --- scripts/flatiron/train_scenario.sh | 8 + scripts/local/train.sh | 92 --- scripts/local/train_scenario.sh | 103 ++++ scripts/training/run_scenario.py | 7 + tests/test_training_scenarios.py | 141 +++++ tests/test_training_seed.py | 46 ++ tests/test_training_submission.py | 48 ++ 26 files changed, 1304 insertions(+), 256 deletions(-) create mode 100644 configs/training/README.md create mode 100644 configs/training/platforms/flatiron_a100.yaml create mode 100644 configs/training/platforms/flatiron_b200.yaml create mode 100644 configs/training/platforms/flatiron_h100.yaml create mode 100644 configs/training/platforms/local.yaml create mode 100644 configs/training/scenarios/cld_hits_output_comparison.yaml create mode 100644 mlpf/training_scenarios.py create mode 100644 mlpf/training_submission.py create mode 100755 scripts/flatiron/run_uv_scenario.sh create mode 100755 scripts/flatiron/submit_scenario.py delete mode 100644 scripts/flatiron/train_a100_uv.sh delete mode 100644 scripts/flatiron/train_h100_uv.sh create mode 100755 scripts/flatiron/train_scenario.sh delete mode 100755 scripts/local/train.sh create mode 100755 scripts/local/train_scenario.sh create mode 100755 scripts/training/run_scenario.py create mode 100644 tests/test_training_scenarios.py create mode 100644 tests/test_training_seed.py create mode 100644 tests/test_training_submission.py diff --git a/DOING.md b/DOING.md index d2f99b47f..dfb6fe263 100644 --- a/DOING.md +++ b/DOING.md @@ -302,6 +302,8 @@ Compare elementwise and set prediction using: - [ ] Extend the benchmark script with set-mode timing and memory measurements. - [x] Run a small CLD-hits overfit test and confirm that loss and matching converge. - [x] Add a local ttbar launcher for paired elementwise and set-output training. +- [x] Add reusable seeded comparison scenarios with local and Flatiron hardware + profiles. - [ ] Run a short CLD-hits training comparison against the elementwise baseline. - [x] Document initial correctness, timing, memory, and scaling measurements here; add physics accuracy after training. diff --git a/configs/training/README.md b/configs/training/README.md new file mode 100644 index 000000000..ce1f32ba9 --- /dev/null +++ b/configs/training/README.md @@ -0,0 +1,54 @@ +# 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 local and Flatiron 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. + +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. diff --git a/configs/training/platforms/flatiron_a100.yaml b/configs/training/platforms/flatiron_a100.yaml new file mode 100644 index 000000000..e6d60ae19 --- /dev/null +++ b/configs/training/platforms/flatiron_a100.yaml @@ -0,0 +1,16 @@ +name: flatiron_a100 +gpus: 4 +data_dir: /mnt/ceph/users/${USER}/mlpf/cld/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 diff --git a/configs/training/platforms/flatiron_b200.yaml b/configs/training/platforms/flatiron_b200.yaml new file mode 100644 index 000000000..0fc04efd3 --- /dev/null +++ b/configs/training/platforms/flatiron_b200.yaml @@ -0,0 +1,16 @@ +name: flatiron_b200 +gpus: 8 +data_dir: /mnt/ceph/users/${USER}/mlpf/cld/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: b200 + time: "12:00:00" + nodes: 1 + tasks_per_node: 1 + cpus_per_task: 64 diff --git a/configs/training/platforms/flatiron_h100.yaml b/configs/training/platforms/flatiron_h100.yaml new file mode 100644 index 000000000..bfaac20af --- /dev/null +++ b/configs/training/platforms/flatiron_h100.yaml @@ -0,0 +1,16 @@ +name: flatiron_h100 +gpus: 8 +data_dir: /mnt/ceph/users/${USER}/mlpf/cld/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: "12:00:00" + nodes: 1 + tasks_per_node: 1 + cpus_per_task: 64 diff --git a/configs/training/platforms/local.yaml b/configs/training/platforms/local.yaml new file mode 100644 index 000000000..be1f813c3 --- /dev/null +++ b/configs/training/platforms/local.yaml @@ -0,0 +1,11 @@ +name: local +gpus: 1 +data_dir: /mnt/work/mlpf/cld/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 diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml new file mode 100644 index 000000000..917df21bd --- /dev/null +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -0,0 +1,41 @@ +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 + +# 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: 128 + parameters: + lr: 0.001 + num_steps: 20000 + val_freq: 2000 + checkpoint_freq: 2000 + nvalid: 100 + ntest: 100 + 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 diff --git a/docs/map.md b/docs/map.md index 9329e9fae..1a51af67b 100644 --- a/docs/map.md +++ b/docs/map.md @@ -9,7 +9,7 @@ The project uses a hierarchical configuration system. - **`particleflow_spec.yaml`**: The single source of truth for the entire project. It defines machine-specific paths (sites), data production scenarios, and model architectures. - **`mlpf/conf.py`**: Defines the Pydantic models for the configuration, ensuring type safety, path resolution, and model-type schemas including `attention`, `gnnlsh`, `litept`, `hept`, and `heptv2`. - **`mlpf/pipeline.py`**: Implements hierarchical configuration resolution: base defaults in `mlpf/conf.py` are overridden by scenario-specific values in `particleflow_spec.yaml`, which can further be overridden via command-line arguments (e.g., `--model.num_convs 6`). -- **`configs/`**: Site-specific Pixi environment configurations (`local/`, `lxplus/`, `tallinn/`). The root `pixi.toml` is a symlink into this directory. +- **`configs/`**: Site-specific Pixi environment configurations (`local/`, `lxplus/`, `tallinn/`) and reusable training scenarios/platform profiles under `configs/training/`. The root `pixi.toml` is a symlink into this directory. - **`pixi.toml` / `pixi.lock` / `uv.lock` / `uv.singularity`**: Project environment, container definitions, and task management. Defines common tasks like `gen`, `post`, `train`, and `validation`. - **`envs/`**: Isolated virtual environment specifications (e.g., `ort-cpu`, `ort-gpu`) for specific runtimes like ONNX. - **`validation_cms.yaml` / `validation_key4hep.yaml`**: Specification files for validation scenarios. @@ -22,6 +22,7 @@ Complex data production and training pipelines are managed using Snakemake or si - **`produce_cms_validation_snakemake.py`**: Orchestrates validation workflows specifically for CMS. - **`produce_validation_snakemake.py`**: Orchestrates validation workflows for Key4Hep detectors (CLD, CLIC). - **`mlpf/pipeline.py`**: The main CLI for training, testing, and hyperparameter optimization. Supports standard and Ray-based execution. +- **`scripts/training/run_scenario.py`**: Resolves a generic scientific training scenario against a hardware profile, validates comparison invariants and global batch size, and launches reproducible seeded jobs. ## 3. Data Production & Preprocessing - **`mlpf/data/`**: Simulator-specific code for generating and preprocessing data. diff --git a/mlpf/conf.py b/mlpf/conf.py index 1709fd3c1..16e67ea26 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -691,6 +691,7 @@ class MLPFConfig(BaseModel): elemtypes_nonzero: Optional[List[int]] = None # Training parameters + seed: int = Field(default=12345, ge=0) num_steps: int = 100000 patience: int = 10000 checkpoint_freq: int = 10000 diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index 875a25dca..792e6aee0 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -721,13 +721,17 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, sampler_mode = DatasetSamplerMode(config.sampler_mode) _logger.info(f"{split}_dataset sampler_mode={sampler_mode.value} shuffle={shuffle}") if world_size > 1 and sampler_mode == DatasetSamplerMode.INTERLEAVED_SHARDS: - sampler = DistributedInterleavedShardSampler(dataset, world_size=world_size, rank=rank, shuffle=shuffle) + sampler = DistributedInterleavedShardSampler( + dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed + ) elif world_size > 1: - sampler = DistributedShardConsecutiveSampler(dataset, world_size=world_size, rank=rank, shuffle=shuffle) + sampler = DistributedShardConsecutiveSampler( + dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed + ) elif sampler_mode == DatasetSamplerMode.INTERLEAVED_SHARDS: - sampler = InterleavedShardSampler(dataset, shuffle=shuffle) + sampler = InterleavedShardSampler(dataset, shuffle=shuffle, seed=config.seed) else: - sampler = ShardConsecutiveSampler(dataset, shuffle=shuffle) + sampler = ShardConsecutiveSampler(dataset, shuffle=shuffle, seed=config.seed) sampler = ResumableSampler(sampler) sampler.name = f"{type_}:{split}" @@ -737,6 +741,10 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, per_particle_keys = ["X", "ytarget"] if config.model.output_mode == OutputMode.SET: per_particle_keys.append("ytarget_set") + loader_generator = torch.Generator() + rank_index = int(rank) if isinstance(rank, int) else 0 + split_offset = 0 if split == "train" else 1000 + loader_generator.manual_seed(config.seed + 10_000 * rank_index + split_offset + len(loaders[split])) loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, @@ -748,6 +756,7 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, # pin_memory_device="cuda:{}".format(rank) if use_cuda else "", drop_last=True, worker_init_fn=set_worker_sharing_strategy, + generator=loader_generator, persistent_workers=config.num_workers > 0, ) diff --git a/mlpf/model/distributed_ray.py b/mlpf/model/distributed_ray.py index 1ae65c073..3d5dfebca 100644 --- a/mlpf/model/distributed_ray.py +++ b/mlpf/model/distributed_ray.py @@ -13,7 +13,7 @@ from mlpf.logger import _logger, _configLogger from mlpf.model.PFDataset import get_interleaved_dataloaders from mlpf.utils import create_comet_experiment -from mlpf.model.training import train_all_steps, get_optimizer +from mlpf.model.training import get_optimizer, seed_everything, train_all_steps from mlpf.conf import MLPFConfig from mlpf.model.utils import ( @@ -277,12 +277,14 @@ def train_ray_trial(config, args, outdir=None): world_size = ray.train.get_context().get_world_size() mlpf_config = MLPFConfig.model_validate(config) + seed_everything(mlpf_config.seed) model = MLPF(mlpf_config) if world_size > 1: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) # optimizer should be created after distributing the model to devices with ray.train.torch.prepare_model(model) model = ray.train.torch.prepare_model(model) + seed_everything(mlpf_config.seed + world_rank) optimizer = get_optimizer(model, mlpf_config) trainable_params, nontrainable_params, table = count_parameters(model) diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 05bb870c5..aac1b2069 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -32,6 +32,7 @@ import os import os.path as osp +import random import time import logging from pathlib import Path @@ -95,6 +96,16 @@ UNIT_REGRESSION_WEIGHTS = {feature: 1.0 for feature in REGRESSION_FEATURES} +def seed_everything(seed): + """Seed Python, NumPy, and PyTorch RNGs for the current process.""" + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + def _domain_label(source_id, input_type_id): source = SOURCE_LABELS.get(int(source_id), f"source{int(source_id)}") input_type = INPUT_TYPE_LABELS.get(int(input_type_id), f"input{int(input_type_id)}") @@ -1297,7 +1308,7 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi _logger.info(f"test_dataset: {sample}, {len(ds)}", color="blue") if world_size > 1: - sampler = torch.utils.data.distributed.DistributedSampler(ds, shuffle=False) + sampler = torch.utils.data.distributed.DistributedSampler(ds, shuffle=False, seed=config.seed) else: sampler = torch.utils.data.SequentialSampler(ds) @@ -1309,6 +1320,9 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi if sample.startswith("cms_") and version and Version(version) >= Version("2.8.0"): vals_for_test += ["pythia"] + test_loader_generator = torch.Generator() + rank_index = int(rank) if isinstance(rank, int) else 0 + test_loader_generator.manual_seed(config.seed + 10_000 * rank_index + 2000) test_loader = torch.utils.data.DataLoader( ds, batch_size=batch_size, @@ -1316,6 +1330,7 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi sampler=sampler, num_workers=config.num_workers, prefetch_factor=config.prefetch_factor, + generator=test_loader_generator, # pin_memory=use_cuda, # pin_memory_device="cuda:{}".format(rank) if use_cuda else "", ) @@ -1352,6 +1367,12 @@ def run(rank: int | str, world_size: int, config: MLPFConfig, outdir: str, logfi _configLogger("mlpf", rank, filename=f"{logfile}.{rank}", loglevel=loglevel) use_cuda = rank != "cpu" + rank_index = int(rank) if isinstance(rank, int) else 0 + + # All ranks initialize the same model. After DDP synchronizes parameters, + # use rank-specific streams for stochastic layers and data workers. + seed_everything(config.seed) + _logger.info(f"Initializing model with seed={config.seed}; process seed={config.seed + rank_index}") dtype = getattr(torch, config.dtype) _logger.info("configured dtype={} for autocast".format(dtype)) @@ -1445,6 +1466,8 @@ def run(rank: int | str, world_size: int, config: MLPFConfig, outdir: str, logfi model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) _logger.info("Configured model for DistributedDataParallel rank={}".format(rank)) + seed_everything(config.seed + rank_index) + trainable_params, nontrainable_params, table = count_parameters(model) _logger.info(str(table)) _logger.info(model) diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py new file mode 100644 index 000000000..9cd88ba83 --- /dev/null +++ b/mlpf/training_scenarios.py @@ -0,0 +1,526 @@ +"""Resolve reusable scientific training scenarios against hardware profiles.""" + +import argparse +import contextlib +import datetime +import json +import os +import shlex +import subprocess +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import yaml +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from mlpf.conf import MLPFConfig + + +PLATFORM_OVERRIDE_KEYS = { + "compile", + "dtype", + "max_open_readers", + "model.attention.use_flash_attn_varlen", + "num_workers", + "prefetch_factor", +} +DERIVED_KEYS = {"gpu_batch_multiplier", "gpus", "seed"} +KNOWN_BOOLEAN_FLAGS = {"comet", "comet_offline", "compile", "make_plots"} +KNOWN_VALUE_FLAGS = {"dtype", "load", "test_datasets"} + + +class ScenarioVariant(BaseModel): + model_config = ConfigDict(extra="forbid") + + model_name: str + overrides: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def reject_derived_overrides(self): + invalid = DERIVED_KEYS.intersection(self.overrides) + if invalid: + raise ValueError( + f"Variant overrides must not set derived keys: {sorted(invalid)}" + ) + return self + + +class ScenarioTraining(BaseModel): + model_config = ConfigDict(extra="forbid") + + global_batch_size: int = Field(gt=0) + parameters: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def reject_derived_parameters(self): + invalid = DERIVED_KEYS.intersection(self.parameters) + if invalid: + raise ValueError( + f"Scenario parameters must not set derived keys: {sorted(invalid)}" + ) + return self + + +class TrainingScenario(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: str + spec_file: str = "particleflow_spec.yaml" + production_name: str + variants: dict[str, ScenarioVariant] + seeds: list[int] = Field(min_length=1) + training: ScenarioTraining + common_overrides: dict[str, Any] = Field(default_factory=dict) + allowed_variant_differences: list[str] = Field( + default_factory=lambda: ["model.output_mode", "model.set_decoder"] + ) + + @model_validator(mode="after") + def validate_scenario(self): + if len(self.variants) < 2: + raise ValueError("A comparison scenario requires at least two variants") + if len(set(self.seeds)) != len(self.seeds): + raise ValueError("Scenario seeds must be unique") + if any(seed < 0 for seed in self.seeds): + raise ValueError("Scenario seeds must be non-negative") + invalid = DERIVED_KEYS.intersection(self.common_overrides) + if invalid: + raise ValueError( + f"Common overrides must not set derived keys: {sorted(invalid)}" + ) + return self + + +class SlurmProfile(BaseModel): + model_config = ConfigDict(extra="forbid") + + partition: str = "gpu" + constraint: str + time: str = "12:00:00" + nodes: int = Field(default=1, gt=0) + tasks_per_node: int = Field(default=1, gt=0) + cpus_per_task: int = Field(default=64, gt=0) + + +class PlatformProfile(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: str + gpus: int = Field(gt=0) + data_dir: str + experiments_dir: str + environment: dict[str, str] = Field(default_factory=dict) + runtime_overrides: dict[str, Any] = Field(default_factory=dict) + slurm: SlurmProfile | None = None + + @model_validator(mode="after") + def validate_runtime_overrides(self): + invalid = set(self.runtime_overrides).difference(PLATFORM_OVERRIDE_KEYS) + if invalid: + raise ValueError( + "Platform profiles may only set runtime-specific overrides; " + f"invalid keys: {sorted(invalid)}" + ) + return self + + +class ResolvedScenarioJob(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + scenario_name: str + platform_name: str + variant_name: str + model_name: str + seed: int + global_batch_size: int + per_gpu_batch_size: int + gpu_batch_multiplier: int + settings: dict[str, Any] + resolved_config: MLPFConfig + + +def _read_yaml(path): + with Path(path).open() as handle: + return yaml.safe_load(handle) + + +def load_training_scenario(path): + return TrainingScenario.model_validate(_read_yaml(path)) + + +def load_platform_profile(path): + profile = PlatformProfile.model_validate(_read_yaml(path)) + profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) + profile.experiments_dir = os.path.expandvars( + os.path.expanduser(profile.experiments_dir) + ) + profile.environment = { + key: os.path.expandvars(os.path.expanduser(value)) + for key, value in profile.environment.items() + } + return profile + + +@contextlib.contextmanager +def _temporary_environment(values): + previous = {key: os.environ.get(key) for key in values} + os.environ.update(values) + try: + yield + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _serialize_cli_value(value): + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, list): + return ",".join(str(item) for item in value) + return str(value) + + +def _settings_as_extra_args(settings): + args = [] + for key, value in settings.items(): + if key in KNOWN_BOOLEAN_FLAGS or key in KNOWN_VALUE_FLAGS: + continue + args.extend([f"--{key}", _serialize_cli_value(value)]) + return args + + +def _config_args(profile, settings): + return SimpleNamespace( + train=True, + test=True, + pipeline=False, + data_dir=profile.data_dir, + gpus=profile.gpus, + compile=settings.get("compile"), + comet=settings.get("comet"), + comet_offline=settings.get("comet_offline"), + dtype=settings.get("dtype"), + load=settings.get("load"), + make_plots=settings.get("make_plots"), + test_datasets=settings.get("test_datasets", []), + ) + + +def _training_batch_size(config): + physical_datasets = config.train_dataset[config.dataset.value].values() + batch_sizes = {dataset.batch_size for dataset in physical_datasets} + if len(batch_sizes) != 1: + raise ValueError( + "Automatic global-batch resolution requires every physical training dataset " + f"to use the same batch size, got {sorted(batch_sizes)}" + ) + return next(iter(batch_sizes)) + + +def _merge_settings(scenario, platform, variant, extra_overrides): + return { + **scenario.training.parameters, + **scenario.common_overrides, + **platform.runtime_overrides, + **variant.overrides, + **extra_overrides, + } + + +def resolve_scenario_job( + scenario, + platform, + variant_name, + seed, + *, + spec_file=None, + global_batch_size=None, + extra_overrides=None, +): + if variant_name not in scenario.variants: + raise ValueError( + f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}" + ) + variant = scenario.variants[variant_name] + extra_overrides = extra_overrides or {} + invalid = DERIVED_KEYS.intersection(extra_overrides) + if invalid: + raise ValueError( + f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}" + ) + settings = _merge_settings(scenario, platform, variant, extra_overrides) + settings["seed"] = seed + selected_spec = str(spec_file or scenario.spec_file) + + with _temporary_environment(platform.environment): + config = MLPFConfig.from_spec( + selected_spec, + variant.model_name, + scenario.production_name, + args=_config_args(platform, settings), + extra_args=_settings_as_extra_args(settings), + ) + + target_global_batch = ( + global_batch_size + if global_batch_size is not None + else scenario.training.global_batch_size + ) + if target_global_batch <= 0: + raise ValueError("global_batch_size must be positive") + dataset_batch_size = _training_batch_size(config) + divisor = platform.gpus * dataset_batch_size + if target_global_batch % divisor: + raise ValueError( + f"global_batch_size={target_global_batch} is not divisible by " + f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" + ) + multiplier = target_global_batch // divisor + settings["gpu_batch_multiplier"] = multiplier + + with _temporary_environment(platform.environment): + config = MLPFConfig.from_spec( + selected_spec, + variant.model_name, + scenario.production_name, + args=_config_args(platform, settings), + extra_args=_settings_as_extra_args(settings), + ) + + return ResolvedScenarioJob( + scenario_name=scenario.name, + platform_name=platform.name, + variant_name=variant_name, + model_name=variant.model_name, + seed=seed, + global_batch_size=target_global_batch, + per_gpu_batch_size=dataset_batch_size * multiplier, + gpu_batch_multiplier=multiplier, + settings=settings, + resolved_config=config, + ) + + +def _flatten(value, prefix=""): + flattened = {} + if isinstance(value, dict): + for key, child in value.items(): + child_prefix = f"{prefix}.{key}" if prefix else key + flattened.update(_flatten(child, child_prefix)) + else: + flattened[prefix] = value + return flattened + + +def _difference_allowed(path, allowed_paths): + return any( + path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths + ) + + +def validate_variant_invariants(jobs, allowed_paths): + if len(jobs) < 2: + return + reference = _flatten(jobs[0].resolved_config.model_dump(mode="json")) + for job in jobs[1:]: + candidate = _flatten(job.resolved_config.model_dump(mode="json")) + differences = { + path: (reference.get(path), candidate.get(path)) + for path in sorted(set(reference) | set(candidate)) + if reference.get(path) != candidate.get(path) + and not _difference_allowed(path, allowed_paths) + } + if differences: + details = ", ".join( + f"{path}: {values[0]!r} != {values[1]!r}" + for path, values in differences.items() + ) + raise ValueError( + f"Scenario variants differ outside allowed fields: {details}" + ) + + +def resolve_scenario_jobs( + scenario, + platform, + *, + spec_file=None, + global_batch_size=None, + extra_overrides=None, +): + jobs = [ + resolve_scenario_job( + scenario, + platform, + variant_name, + seed, + spec_file=spec_file, + global_batch_size=global_batch_size, + extra_overrides=extra_overrides, + ) + for seed in scenario.seeds + for variant_name in scenario.variants + ] + for seed in scenario.seeds: + seed_jobs = [job for job in jobs if job.seed == seed] + validate_variant_invariants(seed_jobs, scenario.allowed_variant_differences) + return jobs + + +def _pipeline_command(job, scenario, platform, spec_file, experiment_dir): + settings = dict(job.settings) + command = [ + "uv", + "run", + "python3", + "-u", + "mlpf/pipeline.py", + "--spec-file", + str(spec_file), + "--model-name", + job.model_name, + "--production-name", + scenario.production_name, + "--data-dir", + platform.data_dir, + "--experiment-dir", + str(experiment_dir), + "train", + "--gpus", + str(platform.gpus), + ] + for name in sorted(KNOWN_VALUE_FLAGS): + value = settings.pop(name, None) + if value is None: + continue + flag = f"--{name.replace('_', '-')}" + if isinstance(value, list): + command.extend([flag, *(str(item) for item in value)]) + else: + command.extend([flag, str(value)]) + for name in sorted(KNOWN_BOOLEAN_FLAGS): + value = settings.pop(name, None) + if value: + command.append(f"--{name.replace('_', '-')}") + for key, value in settings.items(): + command.extend([f"--{key}", _serialize_cli_value(value)]) + return command + + +def _git_revision(): + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def _experiment_path(platform, job, timestamp=None): + timestamp = timestamp or datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") + name = f"{job.scenario_name}_{job.variant_name}_seed{job.seed}_{timestamp}" + return Path(platform.experiments_dir) / name + + +def run_scenario_job(job, scenario, platform, spec_file, *, dry_run=False): + experiment_dir = _experiment_path( + platform, job, timestamp="TIMESTAMP" if dry_run else None + ) + command = _pipeline_command(job, scenario, platform, spec_file, experiment_dir) + print(shlex.join(command), flush=True) + if dry_run: + return + + experiment_dir.mkdir(parents=True, exist_ok=False) + manifest = { + "scenario": scenario.model_dump(mode="json"), + "platform": platform.model_dump(mode="json"), + "job": job.model_dump(mode="json", exclude={"resolved_config"}), + "resolved_config": job.resolved_config.model_dump(mode="json"), + "command": command, + "git_revision": _git_revision(), + } + with (experiment_dir / "scenario-manifest.json").open("w") as handle: + json.dump(manifest, handle, indent=2) + + environment = os.environ.copy() + environment.update(platform.environment) + subprocess.run(command, check=True, env=environment) + + +def _parse_set_overrides(values): + overrides = {} + for item in values: + if "=" not in item: + raise ValueError(f"--set expects KEY=VALUE, got {item!r}") + key, value = item.split("=", 1) + overrides[key] = yaml.safe_load(value) + return overrides + + +def _validate_slurm_allocation(platform): + allocated = os.environ.get("SLURM_GPUS_PER_NODE") + if allocated and allocated.isdigit() and int(allocated) != platform.gpus: + raise ValueError( + f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}" + ) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--scenario", required=True) + parser.add_argument("--platform", required=True) + parser.add_argument("--spec-file") + selection = parser.add_mutually_exclusive_group() + selection.add_argument("--variant") + selection.add_argument("--task-index", type=int) + parser.add_argument("--seed", type=int) + parser.add_argument("--global-batch-size", type=int) + parser.add_argument("--data-dir") + parser.add_argument("--experiments-dir") + parser.add_argument("--set", action="append", default=[], metavar="KEY=VALUE") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args(argv) + + scenario = load_training_scenario(args.scenario) + if args.seed is not None: + if args.seed < 0: + raise ValueError("seed must be non-negative") + # Apply the same seed override to direct and Slurm-array execution. + scenario.seeds = [args.seed] + platform = load_platform_profile(args.platform) + if args.data_dir: + platform.data_dir = args.data_dir + if args.experiments_dir: + platform.experiments_dir = args.experiments_dir + _validate_slurm_allocation(platform) + + spec_file = args.spec_file or scenario.spec_file + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=spec_file, + global_batch_size=args.global_batch_size, + extra_overrides=_parse_set_overrides(args.set), + ) + if args.task_index is not None: + if args.task_index < 0 or args.task_index >= len(jobs): + raise ValueError(f"task-index must be between 0 and {len(jobs) - 1}") + jobs = [jobs[args.task_index]] + else: + if args.variant: + jobs = [job for job in jobs if job.variant_name == args.variant] + if not jobs: + raise ValueError("No jobs matched the requested variant") + + for job in jobs: + run_scenario_job(job, scenario, platform, spec_file, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py new file mode 100644 index 000000000..c9ccb4216 --- /dev/null +++ b/mlpf/training_submission.py @@ -0,0 +1,163 @@ +"""Build and submit Slurm jobs for reusable training scenarios.""" + +import argparse +import shlex +import subprocess +from pathlib import Path + +from mlpf.training_scenarios import ( + load_platform_profile, + load_training_scenario, + resolve_scenario_jobs, +) + + +def resolve_scenario_path(reference, repo_root): + path = Path(reference).expanduser() + if path.is_file(): + return path.resolve() + if path.suffix != ".yaml": + path = path.with_suffix(".yaml") + candidate = repo_root / "configs/training/scenarios" / path.name + if not candidate.is_file(): + raise ValueError(f"Unknown training scenario {reference!r}") + return candidate.resolve() + + +def resolve_flatiron_profile_path(reference, repo_root): + path = Path(reference).expanduser() + if path.is_file(): + return path.resolve() + name = path.stem + if not name.startswith("flatiron_"): + name = f"flatiron_{name}" + candidate = repo_root / "configs/training/platforms" / f"{name}.yaml" + if not candidate.is_file(): + raise ValueError(f"Unknown Flatiron accelerator/profile {reference!r}") + return candidate.resolve() + + +def available_choices(repo_root): + scenarios = sorted( + path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml") + ) + accelerators = sorted( + path.stem.removeprefix("flatiron_") + for path in (repo_root / "configs/training/platforms").glob("flatiron_*.yaml") + ) + return scenarios, accelerators + + +def build_slurm_submission( + scenario_path, + profile_path, + repo_root, + *, + seed=None, +): + scenario = load_training_scenario(scenario_path) + if seed is not None: + if seed < 0: + raise ValueError("seed must be non-negative") + scenario.seeds = [seed] + profile = load_platform_profile(profile_path) + if profile.slurm is None: + raise ValueError( + f"Platform profile {profile.name!r} has no Slurm configuration" + ) + + spec_file = Path(scenario.spec_file) + if not spec_file.is_absolute(): + spec_file = repo_root / spec_file + jobs = resolve_scenario_jobs(scenario, profile, spec_file=spec_file) + if not jobs: + raise ValueError("Scenario did not resolve to any jobs") + + slurm = profile.slurm + logs_dir = repo_root / "logs_slurm" + worker = repo_root / "scripts/flatiron/run_uv_scenario.sh" + command = [ + "sbatch", + "--time", + slurm.time, + "--nodes", + str(slurm.nodes), + "--ntasks-per-node", + str(slurm.tasks_per_node), + "--partition", + slurm.partition, + "--gpus-per-node", + str(profile.gpus), + "--cpus-per-task", + str(slurm.cpus_per_task), + "--constraint", + slurm.constraint, + "--array", + f"0-{len(jobs) - 1}", + "--job-name", + scenario.name, + "--output", + str(logs_dir / "log_%x_%A_%a.out"), + "--error", + str(logs_dir / "log_%x_%A_%a.err"), + "--chdir", + str(repo_root), + str(worker), + str(Path(scenario_path).resolve()), + str(Path(profile_path).resolve()), + ] + if seed is not None: + command.extend(["--seed", str(seed)]) + return command, jobs + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("scenario", nargs="?", help="Scenario name or YAML path") + parser.add_argument( + "accelerator", + nargs="?", + help="Accelerator name (for example h100) or profile path", + ) + parser.add_argument("--seed", type=int, help="Replace the scenario seed list") + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the sbatch command without submitting", + ) + parser.add_argument( + "--list", action="store_true", help="List available scenarios and accelerators" + ) + args = parser.parse_args(argv) + + repo_root = Path(__file__).resolve().parents[1] + scenarios, accelerators = available_choices(repo_root) + if args.list or args.scenario is None or args.accelerator is None: + print("Scenarios: " + ", ".join(scenarios)) + print("Accelerators: " + ", ".join(accelerators)) + if args.list: + return + parser.error("scenario and accelerator are required") + + scenario_path = resolve_scenario_path(args.scenario, repo_root) + profile_path = resolve_flatiron_profile_path(args.accelerator, repo_root) + command, jobs = build_slurm_submission( + scenario_path, + profile_path, + repo_root, + seed=args.seed, + ) + print( + f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" + + shlex.join(command), + flush=True, + ) + if args.dry_run: + return + + (repo_root / "logs_slurm").mkdir(parents=True, exist_ok=True) + subprocess.run(command, check=True) + + +if __name__ == "__main__": + main() diff --git a/particleflow_spec.yaml b/particleflow_spec.yaml index de43cd2d9..63dfc6cf4 100644 --- a/particleflow_spec.yaml +++ b/particleflow_spec.yaml @@ -407,6 +407,7 @@ models: optimizer: lamb lr_schedule: cosinedecay dtype: bfloat16 + seed: 12345 load: null num_steps: 100000 comet: false diff --git a/scripts/flatiron/run_uv_scenario.sh b/scripts/flatiron/run_uv_scenario.sh new file mode 100755 index 000000000..e9ffc6eac --- /dev/null +++ b/scripts/flatiron/run_uv_scenario.sh @@ -0,0 +1,55 @@ +#!/bin/bash +set -euo pipefail + +SCENARIO_FILE=${1:?scenario file is required} +PLATFORM_FILE=${2:?platform profile is required} +shift 2 + +TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} +SEED_OVERRIDE=${SEED:-} +while [[ $# -gt 0 ]]; do + case "$1" in + --task-index) + TASK_INDEX=${2:?--task-index requires a value} + shift 2 + ;; + --seed) + SEED_OVERRIDE=${2:?--seed requires a value} + shift 2 + ;; + *) + echo "Unknown argument: $1" >&2 + exit 2 + ;; + esac +done + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +cd "$REPO_ROOT" + +module --force purge +module load modules/2.4-20250724 +module load slurm gcc cmake cuda/12.8.0 cudnn/9.2.0.82-12 nccl openmpi apptainer + +nvidia-smi +export PYTHONPATH="$REPO_ROOT" +export SRUN_CPUS_PER_TASK=${SLURM_CPUS_PER_TASK} +export RAY_USAGE_STATS_DISABLE=1 +export RAY_TRAIN_V2_ENABLED=1 + +echo "SLURM_JOB_ID=${SLURM_JOB_ID:-none}" +echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID:-none}" +echo "SLURM_GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-unknown}" +echo "scenario=$SCENARIO_FILE platform=$PLATFORM_FILE task_index=$TASK_INDEX" + +RUN_ARGS=( + --scenario "$SCENARIO_FILE" + --platform "$PLATFORM_FILE" + --task-index "$TASK_INDEX" +) +if [[ -n "$SEED_OVERRIDE" ]]; then + RUN_ARGS+=(--seed "$SEED_OVERRIDE") +fi + +uv run python3 scripts/training/run_scenario.py "${RUN_ARGS[@]}" diff --git a/scripts/flatiron/submit_scenario.py b/scripts/flatiron/submit_scenario.py new file mode 100755 index 000000000..9741bc563 --- /dev/null +++ b/scripts/flatiron/submit_scenario.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from mlpf.training_submission import main + + +if __name__ == "__main__": + main() diff --git a/scripts/flatiron/train_a100_uv.sh b/scripts/flatiron/train_a100_uv.sh deleted file mode 100644 index c0bef2bf6..000000000 --- a/scripts/flatiron/train_a100_uv.sh +++ /dev/null @@ -1,76 +0,0 @@ -#!/bin/sh -#SBATCH -t 12:00:00 -#SBATCH -N 1 -#SBATCH --tasks-per-node=1 -#SBATCH -p gpu -#SBATCH --gpus-per-node=4 -#SBATCH --cpus-per-task=64 -#SBATCH --constraint=a100 - -# Job name -#SBATCH -J train - -# Output and error logs -#SBATCH -o logs_slurm/log_%x_%j.out -#SBATCH -e logs_slurm/log_%x_%j.err - -# Add jobscript to job output -echo "#################### Job submission script. #############################" -cat $0 -echo "################# End of job submission script. #########################" - - -module --force purge; module load modules/2.4-20250724 -module load slurm gcc cmake cuda/12.8.0 cudnn/9.2.0.82-12 nccl openmpi apptainer - -nvidia-smi -export PYTHONPATH=`pwd` - -export CUDA_VISIBLE_DEVICES=0,1,2,3 -num_gpus=$((SLURM_GPUS_PER_NODE)) # gpus per compute node - -export SRUN_CPUS_PER_TASK=${SLURM_CPUS_PER_TASK} # necessary on JURECA for Ray to work - -## Disable Ray Usage Stats -export RAY_USAGE_STATS_DISABLE=1 - -echo "DEBUG: SLURM_JOB_ID: $SLURM_JOB_ID" -echo "DEBUG: SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" -echo "DEBUG: SLURM_NNODES: $SLURM_NNODES" -echo "DEBUG: SLURM_NTASKS: $SLURM_NTASKS" -echo "DEBUG: SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE" -echo "DEBUG: SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST" -echo "DEBUG: SLURMD_NODENAME: $SLURMD_NODENAME" -echo "DEBUG: SLURM_NODEID: $SLURM_NODEID" -echo "DEBUG: SLURM_LOCALID: $SLURM_LOCALID" -echo "DEBUG: SLURM_PROCID: $SLURM_PROCID" -echo "DEBUG: CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" -echo "DEBUG: SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES" -echo "DEBUG: SLURM_CPUS_PER_TASK: $SLURM_CPUS_PER_TASK" -echo "DEBUG: SLURM_GPUS_PER_TASK: $SLURM_GPUS_PER_TASK" -echo "DEBUG: SLURM_GPUS_PER_NODE: $SLURM_GPUS_PER_NODE" -echo "DEBUG: SLURM_GPUS: $SLURM_GPUS" -echo "DEBUG: num_gpus: $num_gpus" - -export RAY_TRAIN_V2_ENABLED=1 - -DATA_DIR="/mnt/ceph/users/jpata/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds" - -echo 'Starting training.' - -uv run python3 -u mlpf/pipeline.py \ - --spec-file particleflow_spec.yaml --model-name pyg-cld-hits-v1 --production cld \ - --data-dir $DATA_DIR \ - --experiments-dir /mnt/home/jpata/particleflow/experiments \ - train \ - --gpus $num_gpus \ - --gpu_batch_multiplier 32 \ - --model.attention.use_jagged_attention True \ - --model.attention.use_flash_attn_varlen False \ - --pad_to_multiple_elements 100 \ - --model.attention.num_convs 6 --model.type attention \ - --model.task_queries false \ - --lr 0.0005 --num_steps 20000 --val_freq 2000 --checkpoint_freq 2000 - -# --compile \ # does not work on multiple H100 currently, needs debugging -echo 'Training done.' diff --git a/scripts/flatiron/train_h100_uv.sh b/scripts/flatiron/train_h100_uv.sh deleted file mode 100644 index bfd593850..000000000 --- a/scripts/flatiron/train_h100_uv.sh +++ /dev/null @@ -1,81 +0,0 @@ -#!/bin/sh -#SBATCH -t 12:00:00 -#SBATCH -N 1 -#SBATCH --tasks-per-node=1 -#SBATCH -p gpu -#SBATCH --gpus-per-node=8 -#SBATCH --cpus-per-task=64 -#SBATCH --constraint=h100 - -# Job name -#SBATCH -J train - -# Output and error logs -#SBATCH -o logs_slurm/log_%x_%j.out -#SBATCH -e logs_slurm/log_%x_%j.err - -# Add jobscript to job output -echo "#################### Job submission script. #############################" -cat $0 -echo "################# End of job submission script. #########################" - - -module --force purge; module load modules/2.4-20250724 -module load slurm gcc cmake cuda/12.8.0 cudnn/9.2.0.82-12 nccl openmpi apptainer - -nvidia-smi -export PYTHONPATH=`pwd` - -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 -num_gpus=$((SLURM_GPUS_PER_NODE)) # gpus per compute node - -export SRUN_CPUS_PER_TASK=${SLURM_CPUS_PER_TASK} # necessary on JURECA for Ray to work - -## Disable Ray Usage Stats -export RAY_USAGE_STATS_DISABLE=1 - -echo "DEBUG: SLURM_JOB_ID: $SLURM_JOB_ID" -echo "DEBUG: SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" -echo "DEBUG: SLURM_NNODES: $SLURM_NNODES" -echo "DEBUG: SLURM_NTASKS: $SLURM_NTASKS" -echo "DEBUG: SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE" -echo "DEBUG: SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST" -echo "DEBUG: SLURMD_NODENAME: $SLURMD_NODENAME" -echo "DEBUG: SLURM_NODEID: $SLURM_NODEID" -echo "DEBUG: SLURM_LOCALID: $SLURM_LOCALID" -echo "DEBUG: SLURM_PROCID: $SLURM_PROCID" -echo "DEBUG: CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" -echo "DEBUG: SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES" -echo "DEBUG: SLURM_CPUS_PER_TASK: $SLURM_CPUS_PER_TASK" -echo "DEBUG: SLURM_GPUS_PER_TASK: $SLURM_GPUS_PER_TASK" -echo "DEBUG: SLURM_GPUS_PER_NODE: $SLURM_GPUS_PER_NODE" -echo "DEBUG: SLURM_GPUS: $SLURM_GPUS" -echo "DEBUG: num_gpus: $num_gpus" - -export RAY_TRAIN_V2_ENABLED=1 - -DATA_DIR="/mnt/ceph/users/jpata/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds" - -echo 'Starting training.' - -uv run python3 -u mlpf/pipeline.py \ - --spec-file particleflow_spec.yaml --model-name pyg-cld-hits-v1 --production cld \ - --data-dir $DATA_DIR \ - --experiments-dir /mnt/home/jpata/particleflow/experiments \ - train \ - --gpus $num_gpus \ - --gpu_batch_multiplier 32 \ - --model.attention.use_jagged_attention True \ - --model.attention.use_flash_attn_varlen False \ - --pad_to_multiple_elements 100 \ - --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 \ - --model.type attention \ - --model.task_queries false \ - --lr 0.001 --num_steps 20000 --val_freq 2000 --checkpoint_freq 2000 - -# --compile \ # does not work on multiple H100 currently, needs debugging -echo 'Training done.' diff --git a/scripts/flatiron/train_scenario.sh b/scripts/flatiron/train_scenario.sh new file mode 100755 index 000000000..93ddf0f92 --- /dev/null +++ b/scripts/flatiron/train_scenario.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +cd "$REPO_ROOT" + +exec uv run python3 scripts/flatiron/submit_scenario.py "$@" diff --git a/scripts/local/train.sh b/scripts/local/train.sh deleted file mode 100755 index f54654923..000000000 --- a/scripts/local/train.sh +++ /dev/null @@ -1,92 +0,0 @@ -#!/bin/bash -set -euo pipefail - -export PF_SITE=local - -SPEC_FILE=${SPEC_FILE:-particleflow_spec.yaml} -USE_LOCAL_AVAILABLE_SPEC=${USE_LOCAL_AVAILABLE_SPEC:-true} -LOCAL_SPEC_FILE=${LOCAL_SPEC_FILE:-/tmp/particleflow_local_ttbar_comparison_spec.yaml} -OUTPUT_MODES=${OUTPUT_MODES:-elementwise,set} -HIT_VERSION=${HIT_VERSION:-3.2.1} -HIT_SPLITS=${HIT_SPLITS:-1} -DATA_CONFIG=${DATA_CONFIG:-${HIT_SPLITS// /,}} - -NUM_STEPS=${NUM_STEPS:-2000} -VAL_FREQ=${VAL_FREQ:-200} -CHECKPOINT_FREQ=${CHECKPOINT_FREQ:-200} -NVALID=${NVALID:-100} -NTEST=${NTEST:-100} -GPU_BATCH_MULTIPLIER=${GPU_BATCH_MULTIPLIER:-8} -NUM_WORKERS=${NUM_WORKERS:-8} -PREFETCH_FACTOR=${PREFETCH_FACTOR:-4} -VALIDATION_DIAGNOSTICS_BATCHES=${VALIDATION_DIAGNOSTICS_BATCHES:-4} -EXPERIMENTS_DIR=${EXPERIMENTS_DIR:-experiments} -PAD_TO_MULTIPLE_ELEMENTS=${PAD_TO_MULTIPLE_ELEMENTS:-128} - -IFS=',' read -r -a OUTPUT_MODE_LIST <<< "$OUTPUT_MODES" -read -r -a HIT_SPLIT_LIST <<< "$HIT_SPLITS" - -if [[ "$USE_LOCAL_AVAILABLE_SPEC" == "true" ]]; then - uv run python3 scripts/local/make_local_available_spec.py \ - "$SPEC_FILE" "$LOCAL_SPEC_FILE" \ - --hit-version "$HIT_VERSION" \ - --hit-splits "${HIT_SPLIT_LIST[@]}" - SPEC_FILE="$LOCAL_SPEC_FILE" -fi - -PRODUCTION_NAME=cld -DATA_DIR=${DATA_DIR:-$(uv run python3 scripts/get_param.py "$SPEC_FILE" productions."$PRODUCTION_NAME".workspace_dir)/tfds/} - -set_output_mode() { - local output_mode=$1 - case "$output_mode" in - elementwise) - MODEL_NAME=pyg-cld-hits-v1 - ;; - set) - MODEL_NAME=pyg-cld-hits-set-v1 - ;; - *) - echo "Unknown output mode '$output_mode'. Valid modes: elementwise, set" >&2 - exit 1 - ;; - esac -} - -make_common_args() { - COMMON_ARGS=( - --spec-file "$SPEC_FILE" - --model-name "$MODEL_NAME" - --production-name "$PRODUCTION_NAME" - --data-dir "$DATA_DIR" - --experiments-dir "$EXPERIMENTS_DIR" - train - --data_config "$DATA_CONFIG" - --gpu_batch_multiplier "$GPU_BATCH_MULTIPLIER" - --val_freq "$VAL_FREQ" - --checkpoint_freq "$CHECKPOINT_FREQ" - --num_steps "$NUM_STEPS" - --nvalid "$NVALID" - --ntest "$NTEST" - --num_workers "$NUM_WORKERS" - --prefetch_factor "$PREFETCH_FACTOR" - --sampler_mode interleaved-shards - --validation_diagnostics_batches "$VALIDATION_DIAGNOSTICS_BATCHES" - --make_plots - --pad_to_multiple_elements "$PAD_TO_MULTIPLE_ELEMENTS" - ) -} - -run_comparison_training() { - local output_mode=$1 - echo "Starting CLD ttbar hit training with output_mode=$output_mode model=$MODEL_NAME" - uv run python3 mlpf/pipeline.py \ - --prefix "ttbar-${output_mode}_" \ - "${COMMON_ARGS[@]}" -} - -for output_mode in "${OUTPUT_MODE_LIST[@]}"; do - set_output_mode "$output_mode" - make_common_args - run_comparison_training "$output_mode" -done diff --git a/scripts/local/train_scenario.sh b/scripts/local/train_scenario.sh new file mode 100755 index 000000000..67d859e8f --- /dev/null +++ b/scripts/local/train_scenario.sh @@ -0,0 +1,103 @@ +#!/bin/bash +# Pick and run a reusable training scenario on the local platform. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +SCENARIO_DIR="$REPO_ROOT/configs/training/scenarios" + +print_choices() { + echo "Scenarios:" + local scenario + for scenario in "$SCENARIO_DIR"/*.yaml; do + [[ -e "$scenario" ]] || continue + basename "$scenario" .yaml + done +} + +if [[ ${1:-} == "--list" ]]; then + print_choices + exit 0 +fi +if [[ $# -eq 0 ]]; then + print_choices + echo "Usage: $0 SCENARIO [RUNNER_OPTIONS...]" >&2 + exit 2 +fi + +SCENARIO_REFERENCE=$1 +shift +if [[ -f "$SCENARIO_REFERENCE" ]]; then + SCENARIO_FILE=$(cd "$(dirname "$SCENARIO_REFERENCE")" && pwd)/$(basename "$SCENARIO_REFERENCE") +else + SCENARIO_NAME=${SCENARIO_REFERENCE%.yaml} + SCENARIO_FILE="$SCENARIO_DIR/$SCENARIO_NAME.yaml" +fi +if [[ ! -f "$SCENARIO_FILE" ]]; then + echo "Unknown training scenario: $SCENARIO_REFERENCE" >&2 + print_choices >&2 + exit 2 +fi + +cd "$REPO_ROOT" +export PF_SITE=local + +PLATFORM_FILE=${PLATFORM_FILE:-configs/training/platforms/local.yaml} +SPEC_FILE=${SPEC_FILE:-$(uv run python3 scripts/get_param.py "$SCENARIO_FILE" spec_file particleflow_spec.yaml)} +PRODUCTION_NAME=$(uv run python3 scripts/get_param.py "$SCENARIO_FILE" production_name) +USE_LOCAL_AVAILABLE_SPEC=${USE_LOCAL_AVAILABLE_SPEC:-true} +LOCAL_SPEC_FILE=${LOCAL_SPEC_FILE:-/tmp/particleflow_local_available_spec.yaml} +SEED=${SEED:-} +HIT_VERSION=${HIT_VERSION:-3.2.1} +HIT_SPLITS=${HIT_SPLITS:-1} +DATA_CONFIG=${DATA_CONFIG:-${HIT_SPLITS// /,}} + +# Local defaults intentionally shorten the generic comparison scenario. Every +# value remains overridable through the existing environment interface or by +# supplying a later runner option on this command line. +NUM_STEPS=${NUM_STEPS:-2000} +VAL_FREQ=${VAL_FREQ:-200} +CHECKPOINT_FREQ=${CHECKPOINT_FREQ:-200} +NVALID=${NVALID:-100} +NTEST=${NTEST:-100} +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-${GPU_BATCH_MULTIPLIER:-8}} +NUM_WORKERS=${NUM_WORKERS:-8} +PREFETCH_FACTOR=${PREFETCH_FACTOR:-4} +VALIDATION_DIAGNOSTICS_BATCHES=${VALIDATION_DIAGNOSTICS_BATCHES:-4} +EXPERIMENTS_DIR=${EXPERIMENTS_DIR:-experiments} +PAD_TO_MULTIPLE_ELEMENTS=${PAD_TO_MULTIPLE_ELEMENTS:-128} + +read -r -a HIT_SPLIT_LIST <<< "$HIT_SPLITS" +if [[ "$USE_LOCAL_AVAILABLE_SPEC" == "true" ]]; then + uv run python3 scripts/local/make_local_available_spec.py \ + "$SPEC_FILE" "$LOCAL_SPEC_FILE" \ + --hit-version "$HIT_VERSION" \ + --hit-splits "${HIT_SPLIT_LIST[@]}" + SPEC_FILE="$LOCAL_SPEC_FILE" +fi + +DATA_DIR=${DATA_DIR:-$(uv run python3 scripts/get_param.py "$SPEC_FILE" productions."$PRODUCTION_NAME".workspace_dir)/tfds/} + +RUN_ARGS=( + --scenario "$SCENARIO_FILE" + --platform "$PLATFORM_FILE" + --spec-file "$SPEC_FILE" + --global-batch-size "$GLOBAL_BATCH_SIZE" + --data-dir "$DATA_DIR" + --experiments-dir "$EXPERIMENTS_DIR" + --set "data_config=$DATA_CONFIG" + --set "num_steps=$NUM_STEPS" + --set "val_freq=$VAL_FREQ" + --set "checkpoint_freq=$CHECKPOINT_FREQ" + --set "nvalid=$NVALID" + --set "ntest=$NTEST" + --set "num_workers=$NUM_WORKERS" + --set "prefetch_factor=$PREFETCH_FACTOR" + --set "validation_diagnostics_batches=$VALIDATION_DIAGNOSTICS_BATCHES" + --set "pad_to_multiple_elements=$PAD_TO_MULTIPLE_ELEMENTS" +) +if [[ -n "$SEED" ]]; then + RUN_ARGS+=(--seed "$SEED") +fi + +exec uv run python3 scripts/training/run_scenario.py "${RUN_ARGS[@]}" "$@" diff --git a/scripts/training/run_scenario.py b/scripts/training/run_scenario.py new file mode 100755 index 000000000..0fe422382 --- /dev/null +++ b/scripts/training/run_scenario.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from mlpf.training_scenarios import main + + +if __name__ == "__main__": + main() diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py new file mode 100644 index 000000000..2854f3ec9 --- /dev/null +++ b/tests/test_training_scenarios.py @@ -0,0 +1,141 @@ +from copy import deepcopy +from pathlib import Path + +import pytest + +from mlpf.training_scenarios import ( + PlatformProfile, + ScenarioVariant, + ScenarioTraining, + load_platform_profile, + load_training_scenario, + resolve_scenario_jobs, + validate_variant_invariants, +) + + +ROOT = Path(__file__).resolve().parents[1] +SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" +PLATFORMS = ROOT / "configs/training/platforms" + + +def test_comparison_scenario_resolves_both_output_modes_with_same_seed(): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + + assert [job.variant_name for job in jobs] == ["elementwise", "set"] + assert {job.seed for job in jobs} == {12345} + assert {job.gpu_batch_multiplier for job in jobs} == {8} + assert {job.resolved_config.model.output_mode.value for job in jobs} == { + "elementwise", + "set", + } + assert all(job.resolved_config.seed == 12345 for job in jobs) + + +@pytest.mark.parametrize( + ("profile_name", "expected_multiplier"), + [ + ("flatiron_h100.yaml", 16), + ("flatiron_a100.yaml", 32), + ("flatiron_b200.yaml", 16), + ], +) +def test_platform_profiles_preserve_global_batch(profile_name, expected_multiplier): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / profile_name) + + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + ) + + assert {job.global_batch_size for job in jobs} == {128} + assert {job.gpu_batch_multiplier for job in jobs} == {expected_multiplier} + assert {job.per_gpu_batch_size for job in jobs} == {128 // platform.gpus} + + +def test_variant_invariant_check_rejects_unapproved_difference(): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + bad_jobs = deepcopy(jobs) + bad_jobs[1].resolved_config.lr *= 2 + + with pytest.raises(ValueError, match="variants differ.*lr"): + validate_variant_invariants(bad_jobs, scenario.allowed_variant_differences) + + +def test_global_batch_must_be_divisible_by_hardware_layout(): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / "flatiron_a100.yaml") + + with pytest.raises(ValueError, match="global_batch_size=130 is not divisible"): + resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=130, + ) + + +def test_scenario_and_platform_reject_misplaced_settings(): + with pytest.raises(ValueError, match="derived keys"): + ScenarioTraining( + global_batch_size=8, + parameters={"gpu_batch_multiplier": 8}, + ) + + with pytest.raises(ValueError, match="runtime-specific"): + PlatformProfile( + name="bad", + gpus=1, + data_dir="/tmp/data", + experiments_dir="/tmp/experiments", + runtime_overrides={"lr": 0.1}, + ) + + with pytest.raises(ValueError, match="derived keys"): + ScenarioVariant( + model_name="pyg-cld-hits-v1", + overrides={"seed": 17}, + ) + + +def test_cli_seed_replaces_scenario_seed_for_task_selection(capsys): + from mlpf.training_scenarios import main + + main( + [ + "--scenario", + str(SCENARIO), + "--platform", + str(PLATFORMS / "local.yaml"), + "--spec-file", + str(ROOT / "particleflow_spec.yaml"), + "--global-batch-size", + "8", + "--seed", + "17", + "--task-index", + "1", + "--dry-run", + ] + ) + + command = capsys.readouterr().out + assert "--seed 17" in command + assert "--model-name pyg-cld-hits-set-v1" in command diff --git a/tests/test_training_seed.py b/tests/test_training_seed.py new file mode 100644 index 000000000..110ff41ad --- /dev/null +++ b/tests/test_training_seed.py @@ -0,0 +1,46 @@ +import random + +import numpy as np +import pytest +import torch + +from mlpf.conf import MLPFConfig +from mlpf.model.training import seed_everything + + +def sample_rngs(): + return random.random(), np.random.random(), torch.rand(3) + + +def test_seed_everything_reproduces_python_numpy_and_torch_streams(): + seed_everything(2468) + first = sample_rngs() + seed_everything(2468) + second = sample_rngs() + + assert first[:2] == second[:2] + torch.testing.assert_close(first[2], second[2]) + + +def test_seed_is_exposed_and_validated_in_mlpf_config(): + config = MLPFConfig.model_validate( + { + "dataset": "cld_hits", + "data_dir": "/tmp", + "model": {"type": "attention", "attention": {}}, + "conv_type": "attention", + "seed": 17, + } + ) + assert config.seed == 17 + + with pytest.raises(ValueError, match="greater than or equal to 0"): + MLPFConfig.model_validate( + { + "dataset": "cld_hits", + "data_dir": "/tmp", + "model": {"type": "attention", "attention": {}}, + "conv_type": "attention", + "seed": -1, + } + ) diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py new file mode 100644 index 000000000..a7a9c63d8 --- /dev/null +++ b/tests/test_training_submission.py @@ -0,0 +1,48 @@ +from pathlib import Path + +from mlpf.training_submission import ( + available_choices, + build_slurm_submission, + resolve_flatiron_profile_path, + resolve_scenario_path, +) + + +ROOT = Path(__file__).resolve().parents[1] + + +def test_picker_discovers_scenarios_and_accelerators(): + scenarios, accelerators = available_choices(ROOT) + + assert "cld_hits_output_comparison" in scenarios + assert {"a100", "b200", "h100"}.issubset(accelerators) + + +def test_h100_submission_is_derived_from_scenario_and_profile(): + scenario = resolve_scenario_path("cld_hits_output_comparison", ROOT) + profile = resolve_flatiron_profile_path("h100", ROOT) + + command, jobs = build_slurm_submission( + scenario, + profile, + ROOT, + seed=2468, + ) + + assert [job.variant_name for job in jobs] == ["elementwise", "set"] + assert {job.seed for job in jobs} == {2468} + assert command[command.index("--gpus-per-node") + 1] == "8" + assert command[command.index("--constraint") + 1] == "h100" + assert command[command.index("--array") + 1] == "0-1" + assert command[-2:] == ["--seed", "2468"] + + +def test_array_size_includes_all_scenario_seeds(): + scenario = resolve_scenario_path("cld_hits_output_comparison", ROOT) + profile = resolve_flatiron_profile_path("a100", ROOT) + + command, jobs = build_slurm_submission(scenario, profile, ROOT) + + assert len(jobs) == 2 + assert command[command.index("--array") + 1] == "0-1" + assert command[command.index("--gpus-per-node") + 1] == "4" From 389da4ee5df884476792010e49c4518cb00ab844 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 12:28:18 +0300 Subject: [PATCH 04/29] Fix Flatiron scenario worker repository path --- mlpf/training_submission.py | 2 ++ scripts/flatiron/run_uv_scenario.sh | 11 +++++++++-- tests/test_training_submission.py | 1 + 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py index c9ccb4216..66d644a33 100644 --- a/mlpf/training_submission.py +++ b/mlpf/training_submission.py @@ -105,6 +105,8 @@ def build_slurm_submission( str(worker), str(Path(scenario_path).resolve()), str(Path(profile_path).resolve()), + "--repo-root", + str(repo_root), ] if seed is not None: command.extend(["--seed", str(seed)]) diff --git a/scripts/flatiron/run_uv_scenario.sh b/scripts/flatiron/run_uv_scenario.sh index e9ffc6eac..5ae40bf6f 100755 --- a/scripts/flatiron/run_uv_scenario.sh +++ b/scripts/flatiron/run_uv_scenario.sh @@ -7,6 +7,7 @@ shift 2 TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} SEED_OVERRIDE=${SEED:-} +REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} while [[ $# -gt 0 ]]; do case "$1" in --task-index) @@ -17,6 +18,10 @@ while [[ $# -gt 0 ]]; do SEED_OVERRIDE=${2:?--seed requires a value} shift 2 ;; + --repo-root) + REPO_ROOT=${2:?--repo-root requires a value} + shift 2 + ;; *) echo "Unknown argument: $1" >&2 exit 2 @@ -24,8 +29,10 @@ while [[ $# -gt 0 ]]; do esac done -SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) -REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +if [[ ! -f "$REPO_ROOT/scripts/training/run_scenario.py" ]]; then + echo "Invalid repository root '$REPO_ROOT': scripts/training/run_scenario.py is missing" >&2 + exit 2 +fi cd "$REPO_ROOT" module --force purge diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index a7a9c63d8..18cdbbe9a 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -34,6 +34,7 @@ def test_h100_submission_is_derived_from_scenario_and_profile(): assert command[command.index("--gpus-per-node") + 1] == "8" assert command[command.index("--constraint") + 1] == "h100" assert command[command.index("--array") + 1] == "0-1" + assert command[command.index("--repo-root") + 1] == str(ROOT) assert command[-2:] == ["--seed", "2468"] From dc6524a3e5b66210d2e80fd04f64869f3f8a6c8d Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 12:38:00 +0300 Subject: [PATCH 05/29] factorize experiment path --- configs/training/README.md | 3 ++- mlpf/training_scenarios.py | 4 ++-- tests/test_training_scenarios.py | 18 ++++++++++++++++++ 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/configs/training/README.md b/configs/training/README.md index ce1f32ba9..3727cecec 100644 --- a/configs/training/README.md +++ b/configs/training/README.md @@ -51,4 +51,5 @@ one comparison pair with an explicit seed. 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. +seed, final configuration, command, and git revision. Runs are grouped as +`//_seed_/`. diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index 9cd88ba83..28f846ec6 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -423,8 +423,8 @@ def _git_revision(): def _experiment_path(platform, job, timestamp=None): timestamp = timestamp or datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") - name = f"{job.scenario_name}_{job.variant_name}_seed{job.seed}_{timestamp}" - return Path(platform.experiments_dir) / name + experiment_name = f"{job.variant_name}_seed{job.seed}_{timestamp}" + return Path(platform.experiments_dir) / job.scenario_name / experiment_name def run_scenario_job(job, scenario, platform, spec_file, *, dry_run=False): diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index 2854f3ec9..9bff63a47 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -7,6 +7,7 @@ PlatformProfile, ScenarioVariant, ScenarioTraining, + _experiment_path, load_platform_profile, load_training_scenario, resolve_scenario_jobs, @@ -139,3 +140,20 @@ def test_cli_seed_replaces_scenario_seed_for_task_selection(capsys): command = capsys.readouterr().out assert "--seed 17" in command assert "--model-name pyg-cld-hits-set-v1" in command + + +def test_experiments_are_grouped_under_the_scenario_directory(): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + job = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + )[0] + + path = _experiment_path(platform, job, timestamp="TIMESTAMP") + + assert path == Path( + "experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP" + ) From 58aa0101beb03d7161efd70edf2f2e47e9a0361a Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 12:43:12 +0300 Subject: [PATCH 06/29] Fix unused set-output parameters in DDP --- .../scenarios/cld_hits_output_comparison.yaml | 2 +- mlpf/model/mlpf.py | 5 +- mlpf/model/set_losses.py | 6 ++- tests/test_set_prediction.py | 49 ++++++++++++++++++- 4 files changed, 57 insertions(+), 5 deletions(-) diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index 917df21bd..4e07d1e5c 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -13,7 +13,7 @@ seeds: [12345] training: # Kept fixed across hardware profiles. The runner derives the per-GPU batch. - global_batch_size: 128 + global_batch_size: 512 parameters: lr: 0.001 num_steps: 20000 diff --git a/mlpf/model/mlpf.py b/mlpf/model/mlpf.py index fd6067dda..2506e44de 100644 --- a/mlpf/model/mlpf.py +++ b/mlpf/model/mlpf.py @@ -1384,8 +1384,9 @@ def __init__( _logger.info("Initializing output DNNs") t0 = time.time() - self.classification_norm = torch.nn.LayerNorm(decoding_dim) if self.use_pre_layernorm else None - self.regression_norm = torch.nn.LayerNorm(decoding_dim) if self.use_pre_layernorm else None + use_elementwise_output_norms = self.use_pre_layernorm and self.output_mode == OutputMode.ELEMENTWISE + self.classification_norm = torch.nn.LayerNorm(decoding_dim) if use_elementwise_output_norms else None + self.regression_norm = torch.nn.LayerNorm(decoding_dim) if use_elementwise_output_norms else None if self.task_queries and not self.use_split_backbone and self.output_mode == OutputMode.ELEMENTWISE: self.classification_query = nn.Parameter(torch.zeros(1, 1, decoding_dim), requires_grad=True) diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py index c92f60c22..4d61aa6cb 100644 --- a/mlpf/model/set_losses.py +++ b/mlpf/model/set_losses.py @@ -133,7 +133,11 @@ def set_event_loss( num_matched = int(presence_targets.sum().item()) if num_matched == 0: - zero = predictions["cls_binary"].sum() * 0.0 + # Keep every set-output head in the autograd graph even for a batch with + # no target particles. This produces zero gradients for PID and momentum + # rather than making their parameters unused under DDP. + output_keys = ("cls_binary", "cls_id_onehot", *REGRESSION_FEATURES) + zero = sum(predictions[key].sum() * 0.0 for key in output_keys) losses["Classification"] = zero for feature in REGRESSION_FEATURES: losses[f"Regression_{feature}"] = zero diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py index 9bc1c3380..96815509c 100644 --- a/tests/test_set_prediction.py +++ b/tests/test_set_prediction.py @@ -41,6 +41,33 @@ def make_config(num_slots=4): ) +def make_attention_config(num_slots=4): + return MLPFConfig.model_validate( + { + "dataset": "cld_hits", + "data_dir": "/tmp", + "model": { + "type": "attention", + "output_mode": "set", + "input_encoding": "joint", + "attention": { + "num_convs": 1, + "num_heads": 2, + "head_dim": 8, + "use_pre_layernorm": True, + }, + "set_decoder": { + "num_slots": num_slots, + "num_layers": 1, + "num_heads": 2, + }, + "hit_feature_engineering": {"enabled": False}, + }, + "conv_type": "attention", + } + ) + + def make_target_tensor(batch_size=1, num_targets=2): target = torch.zeros(batch_size, num_targets, 14) phi = torch.linspace(-torch.pi, torch.pi, num_targets + 1)[:-1] @@ -102,6 +129,23 @@ def test_set_model_output_axis_is_num_slots(): ) +def test_attention_set_model_has_no_unused_elementwise_parameters(): + config = make_attention_config() + model = MLPF(config) + X = torch.randn(2, 8, config.input_dim) + X[..., 0] = 1 + mask = torch.ones(2, 8, dtype=torch.bool) + + predictions = model(X, mask) + sum(prediction.square().mean() for prediction in predictions).backward() + + assert model.classification_norm is None + assert model.regression_norm is None + assert [ + name for name, parameter in model.named_parameters() if parameter.grad is None + ] == [] + + def test_hungarian_match_finds_permuted_particles(): ytarget_tensor = make_target_tensor() targets = unpack_target(ytarget_tensor, None) @@ -197,6 +241,7 @@ def test_set_loss_supports_an_event_without_targets(): assert matches[0][0].numel() == 0 assert losses["Classification"] == 0 assert losses["Regression_pt"] == 0 + assert all(prediction.grad is not None for prediction in predictions.values()) def test_predict_particles_restores_absolute_set_kinematics(): @@ -225,7 +270,9 @@ def test_set_model_10k_inputs_forward_backward(): raw_predictions = model(batch.X, batch.mask) predictions = unpack_predictions(raw_predictions) targets = unpack_target(batch.ytarget_set, model) - losses, _ = set_event_loss(targets, predictions, batch.target_mask, REGRESSION_WEIGHTS) + losses, _ = set_event_loss( + targets, predictions, batch.target_mask, REGRESSION_WEIGHTS + ) loss = sum(losses.values()) loss.backward() From 00b4f8b9dd0fb94c2781843d7f0f815243f43eab Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 13:48:12 +0300 Subject: [PATCH 07/29] Add Tallinn and LUMI scenario launchers --- DOING.md | 4 +- configs/training/README.md | 23 ++++- configs/training/platforms/lumi_mi250x.yaml | 27 +++++ configs/training/platforms/tallinn_l40.yaml | 20 ++++ mlpf/training_scenarios.py | 18 +++- mlpf/training_submission.py | 108 ++++++++++++++------ scripts/lumi/run_scenario.sh | 72 +++++++++++++ scripts/lumi/submit_scenario.py | 7 ++ scripts/lumi/train_scenario.sh | 15 +++ scripts/tallinn/run_scenario.sh | 58 +++++++++++ scripts/tallinn/submit_scenario.py | 7 ++ scripts/tallinn/train_scenario.sh | 8 ++ tests/test_training_scenarios.py | 12 ++- tests/test_training_submission.py | 49 +++++++++ 14 files changed, 384 insertions(+), 44 deletions(-) create mode 100644 configs/training/platforms/lumi_mi250x.yaml create mode 100644 configs/training/platforms/tallinn_l40.yaml create mode 100755 scripts/lumi/run_scenario.sh create mode 100755 scripts/lumi/submit_scenario.py create mode 100755 scripts/lumi/train_scenario.sh create mode 100755 scripts/tallinn/run_scenario.sh create mode 100755 scripts/tallinn/submit_scenario.py create mode 100755 scripts/tallinn/train_scenario.sh diff --git a/DOING.md b/DOING.md index dfb6fe263..6a92a69b9 100644 --- a/DOING.md +++ b/DOING.md @@ -302,8 +302,8 @@ Compare elementwise and set prediction using: - [ ] Extend the benchmark script with set-mode timing and memory measurements. - [x] Run a small CLD-hits overfit test and confirm that loss and matching converge. - [x] Add a local ttbar launcher for paired elementwise and set-output training. -- [x] Add reusable seeded comparison scenarios with local and Flatiron hardware - profiles. +- [x] Add reusable seeded comparison scenarios with local, Tallinn, LUMI, and + Flatiron hardware profiles. - [ ] Run a short CLD-hits training comparison against the elementwise baseline. - [x] Document initial correctness, timing, memory, and scaling measurements here; add physics accuracy after training. diff --git a/configs/training/README.md b/configs/training/README.md index 3727cecec..da3d09616 100644 --- a/configs/training/README.md +++ b/configs/training/README.md @@ -32,9 +32,9 @@ Additional arguments are forwarded to the generic scenario runner, such as 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 local and Flatiron 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. +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: @@ -49,6 +49,23 @@ 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 diff --git a/configs/training/platforms/lumi_mi250x.yaml b/configs/training/platforms/lumi_mi250x.yaml new file mode 100644 index 000000000..bf18bcddb --- /dev/null +++ b/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 diff --git a/configs/training/platforms/tallinn_l40.yaml b/configs/training/platforms/tallinn_l40.yaml new file mode 100644 index 000000000..b841a0ee9 --- /dev/null +++ b/configs/training/platforms/tallinn_l40.yaml @@ -0,0 +1,20 @@ +name: tallinn_l40 +gpus: 2 +data_dir: /local/${USER}/mlpf/cld/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 diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index 28f846ec6..78d19fa77 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -9,7 +9,7 @@ import subprocess from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Any, Literal import yaml from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -96,11 +96,25 @@ class SlurmProfile(BaseModel): model_config = ConfigDict(extra="forbid") partition: str = "gpu" - constraint: str + constraint: str | None = None + account: str | None = None time: str = "12:00:00" nodes: int = Field(default=1, gt=0) tasks_per_node: int = Field(default=1, gt=0) cpus_per_task: int = Field(default=64, gt=0) + gpu_request: Literal["gpus-per-node", "gpus-per-task", "gres"] = "gpus-per-node" + gpu_type: str | None = None + memory: str | None = None + memory_per_gpu: str | None = None + no_requeue: bool = False + + @model_validator(mode="after") + def validate_resources(self): + if self.gpu_type is not None and self.gpu_request != "gres": + raise ValueError("gpu_type is only valid with gpu_request='gres'") + if self.memory is not None and self.memory_per_gpu is not None: + raise ValueError("Set only one of memory and memory_per_gpu") + return self class PlatformProfile(BaseModel): diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py index 66d644a33..4c8cd9ec7 100644 --- a/mlpf/training_submission.py +++ b/mlpf/training_submission.py @@ -1,4 +1,4 @@ -"""Build and submit Slurm jobs for reusable training scenarios.""" +"""Build and submit site-specific Slurm jobs for reusable training scenarios.""" import argparse import shlex @@ -24,36 +24,66 @@ def resolve_scenario_path(reference, repo_root): return candidate.resolve() -def resolve_flatiron_profile_path(reference, repo_root): +def resolve_platform_profile_path(reference, repo_root, site): path = Path(reference).expanduser() if path.is_file(): return path.resolve() name = path.stem - if not name.startswith("flatiron_"): - name = f"flatiron_{name}" + prefix = f"{site}_" + if not name.startswith(prefix): + name = f"{prefix}{name}" candidate = repo_root / "configs/training/platforms" / f"{name}.yaml" if not candidate.is_file(): - raise ValueError(f"Unknown Flatiron accelerator/profile {reference!r}") + raise ValueError(f"Unknown {site} accelerator/profile {reference!r}") return candidate.resolve() -def available_choices(repo_root): +def resolve_flatiron_profile_path(reference, repo_root): + """Backward-compatible Flatiron profile resolver.""" + return resolve_platform_profile_path(reference, repo_root, "flatiron") + + +def available_choices(repo_root, site="flatiron"): scenarios = sorted( path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml") ) accelerators = sorted( - path.stem.removeprefix("flatiron_") - for path in (repo_root / "configs/training/platforms").glob("flatiron_*.yaml") + path.stem.removeprefix(f"{site}_") + for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml") ) return scenarios, accelerators +def _worker_for_site(repo_root, site): + relative_paths = { + "flatiron": "scripts/flatiron/run_uv_scenario.sh", + "tallinn": "scripts/tallinn/run_scenario.sh", + "lumi": "scripts/lumi/run_scenario.sh", + } + try: + return repo_root / relative_paths[site] + except KeyError as exc: + raise ValueError(f"No scenario worker is configured for site {site!r}") from exc + + +def _gpu_request_args(slurm, gpus): + if slurm.gpu_request == "gpus-per-node": + return ["--gpus-per-node", str(gpus)] + if slurm.gpu_request == "gpus-per-task": + return ["--gpus-per-task", str(gpus)] + resource = f"gpu:{gpus}" + if slurm.gpu_type: + resource = f"gpu:{slurm.gpu_type}:{gpus}" + return ["--gres", resource] + + def build_slurm_submission( scenario_path, profile_path, repo_root, *, seed=None, + worker=None, ): scenario = load_training_scenario(scenario_path) if seed is not None: @@ -75,7 +105,9 @@ def build_slurm_submission( slurm = profile.slurm logs_dir = repo_root / "logs_slurm" - worker = repo_root / "scripts/flatiron/run_uv_scenario.sh" + worker = ( + Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") + ) command = [ "sbatch", "--time", @@ -86,40 +118,51 @@ def build_slurm_submission( str(slurm.tasks_per_node), "--partition", slurm.partition, - "--gpus-per-node", - str(profile.gpus), + *_gpu_request_args(slurm, profile.gpus), "--cpus-per-task", str(slurm.cpus_per_task), - "--constraint", - slurm.constraint, - "--array", - f"0-{len(jobs) - 1}", - "--job-name", - scenario.name, - "--output", - str(logs_dir / "log_%x_%A_%a.out"), - "--error", - str(logs_dir / "log_%x_%A_%a.err"), - "--chdir", - str(repo_root), - str(worker), - str(Path(scenario_path).resolve()), - str(Path(profile_path).resolve()), - "--repo-root", - str(repo_root), ] + if slurm.constraint: + command.extend(["--constraint", slurm.constraint]) + if slurm.account: + command.extend(["--account", slurm.account]) + if slurm.memory: + command.extend(["--mem", slurm.memory]) + if slurm.memory_per_gpu: + command.extend(["--mem-per-gpu", slurm.memory_per_gpu]) + if slurm.no_requeue: + command.append("--no-requeue") + command.extend( + [ + "--array", + f"0-{len(jobs) - 1}", + "--job-name", + scenario.name, + "--output", + str(logs_dir / "log_%x_%A_%a.out"), + "--error", + str(logs_dir / "log_%x_%A_%a.err"), + "--chdir", + str(repo_root), + str(worker), + str(Path(scenario_path).resolve()), + str(Path(profile_path).resolve()), + "--repo-root", + str(repo_root), + ] + ) if seed is not None: command.extend(["--seed", str(seed)]) return command, jobs -def main(argv=None): +def main(argv=None, *, site="flatiron"): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("scenario", nargs="?", help="Scenario name or YAML path") parser.add_argument( "accelerator", nargs="?", - help="Accelerator name (for example h100) or profile path", + help="Accelerator name or platform profile path", ) parser.add_argument("--seed", type=int, help="Replace the scenario seed list") parser.add_argument( @@ -133,7 +176,7 @@ def main(argv=None): args = parser.parse_args(argv) repo_root = Path(__file__).resolve().parents[1] - scenarios, accelerators = available_choices(repo_root) + scenarios, accelerators = available_choices(repo_root, site) if args.list or args.scenario is None or args.accelerator is None: print("Scenarios: " + ", ".join(scenarios)) print("Accelerators: " + ", ".join(accelerators)) @@ -142,12 +185,13 @@ def main(argv=None): parser.error("scenario and accelerator are required") scenario_path = resolve_scenario_path(args.scenario, repo_root) - profile_path = resolve_flatiron_profile_path(args.accelerator, repo_root) + profile_path = resolve_platform_profile_path(args.accelerator, repo_root, site) command, jobs = build_slurm_submission( scenario_path, profile_path, repo_root, seed=args.seed, + worker=_worker_for_site(repo_root, site), ) print( f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" diff --git a/scripts/lumi/run_scenario.sh b/scripts/lumi/run_scenario.sh new file mode 100755 index 000000000..efae3292b --- /dev/null +++ b/scripts/lumi/run_scenario.sh @@ -0,0 +1,72 @@ +#!/bin/bash +set -euo pipefail + +SCENARIO_FILE=${1:?scenario file is required} +PLATFORM_FILE=${2:?platform profile is required} +shift 2 + +TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} +SEED_OVERRIDE=${SEED:-} +REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} +while [[ $# -gt 0 ]]; do + case "$1" in + --task-index) + TASK_INDEX=${2:?--task-index requires a value} + shift 2 + ;; + --seed) + SEED_OVERRIDE=${2:?--seed requires a value} + shift 2 + ;; + --repo-root) + REPO_ROOT=${2:?--repo-root requires a value} + shift 2 + ;; + *) + echo "Unknown argument: $1" >&2 + exit 2 + ;; + esac +done + +if [[ ! -f "$REPO_ROOT/scripts/training/run_scenario.py" ]]; then + echo "Invalid repository root '$REPO_ROOT': scripts/training/run_scenario.py is missing" >&2 + exit 2 +fi +cd "$REPO_ROOT" + +module use /appl/local/containers/ai-modules +module load singularity-AI-bindings +module load aws-ofi-rccl + +export IMG=${IMG:-/appl/local/containers/sif-images/lumi-pytorch-rocm-6.2.4-python-3.12-pytorch-v2.7.0.sif} +export MIOPEN_USER_DB_PATH=${MIOPEN_USER_DB_PATH:-/tmp/${USER}-${SLURM_JOB_ID}-miopen-cache} +export MIOPEN_CUSTOM_CACHE_DIR=${MIOPEN_CUSTOM_CACHE_DIR:-$MIOPEN_USER_DB_PATH} +export ROCM_PATH=${ROCM_PATH:-/opt/rocm} +export KERAS_BACKEND=torch +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-hsn} +export NCCL_NET_GDR_LEVEL=${NCCL_NET_GDR_LEVEL:-3} +export NCCL_DEBUG=${NCCL_DEBUG:-INFO} +export PYTHONPATH="$REPO_ROOT" + +rocm-smi --showdriverversion +echo "SLURM_JOB_ID=${SLURM_JOB_ID:-none}" +echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID:-none}" +echo "scenario=$SCENARIO_FILE platform=$PLATFORM_FILE task_index=$TASK_INDEX" + +RUN_ARGS=( + "$REPO_ROOT/scripts/training/run_scenario.py" + --scenario "$SCENARIO_FILE" + --platform "$PLATFORM_FILE" + --task-index "$TASK_INDEX" +) +if [[ -n "$SEED_OVERRIDE" ]]; then + RUN_ARGS+=(--seed "$SEED_OVERRIDE") +fi + +singularity exec \ + -B /scratch/project_465001293 \ + -B /tmp \ + "$IMG" \ + bash -lc 'source "$1/particleflow-env/bin/activate"; shift; exec python3 "$@"' \ + bash "$REPO_ROOT" "${RUN_ARGS[@]}" diff --git a/scripts/lumi/submit_scenario.py b/scripts/lumi/submit_scenario.py new file mode 100755 index 000000000..26efe2471 --- /dev/null +++ b/scripts/lumi/submit_scenario.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from mlpf.training_submission import main + + +if __name__ == "__main__": + main(site="lumi") diff --git a/scripts/lumi/train_scenario.sh b/scripts/lumi/train_scenario.sh new file mode 100755 index 000000000..3a9bcedaa --- /dev/null +++ b/scripts/lumi/train_scenario.sh @@ -0,0 +1,15 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +cd "$REPO_ROOT" +export PYTHONPATH="$REPO_ROOT${PYTHONPATH:+:$PYTHONPATH}" + +PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE:-$REPO_ROOT/particleflow-env/bin/python3} +if [[ ! -x "$PYTHON_EXECUTABLE" ]]; then + echo "Python environment not found at '$PYTHON_EXECUTABLE'; set PYTHON_EXECUTABLE" >&2 + exit 2 +fi + +exec "$PYTHON_EXECUTABLE" scripts/lumi/submit_scenario.py "$@" diff --git a/scripts/tallinn/run_scenario.sh b/scripts/tallinn/run_scenario.sh new file mode 100755 index 000000000..20c5c78d1 --- /dev/null +++ b/scripts/tallinn/run_scenario.sh @@ -0,0 +1,58 @@ +#!/bin/bash +set -euo pipefail + +SCENARIO_FILE=${1:?scenario file is required} +PLATFORM_FILE=${2:?platform profile is required} +shift 2 + +TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} +SEED_OVERRIDE=${SEED:-} +REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} +while [[ $# -gt 0 ]]; do + case "$1" in + --task-index) + TASK_INDEX=${2:?--task-index requires a value} + shift 2 + ;; + --seed) + SEED_OVERRIDE=${2:?--seed requires a value} + shift 2 + ;; + --repo-root) + REPO_ROOT=${2:?--repo-root requires a value} + shift 2 + ;; + *) + echo "Unknown argument: $1" >&2 + exit 2 + ;; + esac +done + +if [[ ! -f "$REPO_ROOT/scripts/training/run_scenario.py" ]]; then + echo "Invalid repository root '$REPO_ROOT': scripts/training/run_scenario.py is missing" >&2 + exit 2 +fi +cd "$REPO_ROOT" + +export PF_SITE=tallinn +export NCCL_P2P_DISABLE=1 +export NCCL_DEBUG=${NCCL_DEBUG:-INFO} +export NCCL_IB_DISABLE=1 +export PYTHONPATH="$REPO_ROOT" + +nvidia-smi topo -m +echo "SLURM_JOB_ID=${SLURM_JOB_ID:-none}" +echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID:-none}" +echo "scenario=$SCENARIO_FILE platform=$PLATFORM_FILE task_index=$TASK_INDEX" + +RUN_ARGS=( + --scenario "$SCENARIO_FILE" + --platform "$PLATFORM_FILE" + --task-index "$TASK_INDEX" +) +if [[ -n "$SEED_OVERRIDE" ]]; then + RUN_ARGS+=(--seed "$SEED_OVERRIDE") +fi + +exec uv run python3 scripts/training/run_scenario.py "${RUN_ARGS[@]}" diff --git a/scripts/tallinn/submit_scenario.py b/scripts/tallinn/submit_scenario.py new file mode 100755 index 000000000..1ececc6ac --- /dev/null +++ b/scripts/tallinn/submit_scenario.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from mlpf.training_submission import main + + +if __name__ == "__main__": + main(site="tallinn") diff --git a/scripts/tallinn/train_scenario.sh b/scripts/tallinn/train_scenario.sh new file mode 100755 index 000000000..fc2a4a872 --- /dev/null +++ b/scripts/tallinn/train_scenario.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd) +cd "$REPO_ROOT" + +exec uv run python3 scripts/tallinn/submit_scenario.py "$@" diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index 9bff63a47..a97c85372 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -44,9 +44,11 @@ def test_comparison_scenario_resolves_both_output_modes_with_same_seed(): @pytest.mark.parametrize( ("profile_name", "expected_multiplier"), [ - ("flatiron_h100.yaml", 16), - ("flatiron_a100.yaml", 32), - ("flatiron_b200.yaml", 16), + ("flatiron_h100.yaml", 64), + ("flatiron_a100.yaml", 128), + ("flatiron_b200.yaml", 64), + ("tallinn_l40.yaml", 256), + ("lumi_mi250x.yaml", 64), ], ) def test_platform_profiles_preserve_global_batch(profile_name, expected_multiplier): @@ -59,9 +61,9 @@ def test_platform_profiles_preserve_global_batch(profile_name, expected_multipli spec_file=ROOT / "particleflow_spec.yaml", ) - assert {job.global_batch_size for job in jobs} == {128} + assert {job.global_batch_size for job in jobs} == {512} assert {job.gpu_batch_multiplier for job in jobs} == {expected_multiplier} - assert {job.per_gpu_batch_size for job in jobs} == {128 // platform.gpus} + assert {job.per_gpu_batch_size for job in jobs} == {512 // platform.gpus} def test_variant_invariant_check_rejects_unapproved_difference(): diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index 18cdbbe9a..2b6a759a4 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -4,6 +4,7 @@ available_choices, build_slurm_submission, resolve_flatiron_profile_path, + resolve_platform_profile_path, resolve_scenario_path, ) @@ -47,3 +48,51 @@ def test_array_size_includes_all_scenario_seeds(): assert len(jobs) == 2 assert command[command.index("--array") + 1] == "0-1" assert command[command.index("--gpus-per-node") + 1] == "4" + + +def test_tallinn_submission_uses_typed_gres_and_site_worker(): + scenario = resolve_scenario_path("cld_hits_output_comparison", ROOT) + profile = resolve_platform_profile_path("l40", ROOT, "tallinn") + worker = ROOT / "scripts/tallinn/run_scenario.sh" + + command, jobs = build_slurm_submission( + scenario, + profile, + ROOT, + seed=2468, + worker=worker, + ) + + assert len(jobs) == 2 + assert command[command.index("--gres") + 1] == "gpu:l40:2" + assert command[command.index("--mem-per-gpu") + 1] == "80G" + assert str(worker) in command + assert "--constraint" not in command + + +def test_lumi_submission_uses_task_gpus_account_and_container_worker(): + scenario = resolve_scenario_path("cld_hits_output_comparison", ROOT) + profile = resolve_platform_profile_path("mi250x", ROOT, "lumi") + worker = ROOT / "scripts/lumi/run_scenario.sh" + + command, jobs = build_slurm_submission( + scenario, + profile, + ROOT, + worker=worker, + ) + + assert len(jobs) == 2 + assert command[command.index("--gpus-per-task") + 1] == "8" + assert command[command.index("--account") + 1] == "project_465001293" + assert command[command.index("--mem") + 1] == "450G" + assert "--no-requeue" in command + assert str(worker) in command + + +def test_picker_discovers_site_specific_accelerators(): + _, tallinn_accelerators = available_choices(ROOT, "tallinn") + _, lumi_accelerators = available_choices(ROOT, "lumi") + + assert tallinn_accelerators == ["l40"] + assert lumi_accelerators == ["mi250x"] From 289522e24ca160eb87c099e65a1b9b7fb2fd2c48 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 13:54:32 +0300 Subject: [PATCH 08/29] Fix distributed validation batching --- DOING.md | 2 + .../scenarios/cld_hits_output_comparison.yaml | 8 +-- mlpf/model/PFDataset.py | 6 +- tests/test_dataloader.py | 61 +++++++++++++++++++ 4 files changed, 72 insertions(+), 5 deletions(-) diff --git a/DOING.md b/DOING.md index 6a92a69b9..08a8fc612 100644 --- a/DOING.md +++ b/DOING.md @@ -304,6 +304,8 @@ Compare elementwise and set prediction using: - [x] Add a local ttbar launcher for paired elementwise and set-output training. - [x] Add reusable seeded comparison scenarios with local, Tallinn, LUMI, and Flatiron hardware profiles. +- [x] Keep partial validation batches in distributed runs so small `nvalid` + samples do not produce zero per-rank batches. - [ ] Run a short CLD-hits training comparison against the elementwise baseline. - [x] Document initial correctness, timing, memory, and scaling measurements here; add physics accuracy after training. diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index 4e07d1e5c..4bc7d05e2 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -17,10 +17,10 @@ training: parameters: lr: 0.001 num_steps: 20000 - val_freq: 2000 - checkpoint_freq: 2000 - nvalid: 100 - ntest: 100 + val_freq: 1000 + checkpoint_freq: 1000 + nvalid: 512 + ntest: 512 sampler_mode: interleaved-shards validation_diagnostics_batches: 4 pad_to_multiple_elements: 128 diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index 792e6aee0..59cbe3f62 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -754,7 +754,11 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, prefetch_factor=config.prefetch_factor, # pin_memory=use_cuda, # pin_memory_device="cuda:{}".format(rank) if use_cuda else "", - drop_last=True, + # Training uses fixed-size batches, but a bounded validation + # sample can be smaller than one per-rank batch (for example, + # nvalid=100 with 8 ranks and batch_size=64). Keep that partial + # validation batch so every rank participates in evaluation. + drop_last=split == "train", worker_init_fn=set_worker_sharing_strategy, generator=loader_generator, persistent_workers=config.num_workers > 0, diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py index 0bdfbf618..9ff795e38 100644 --- a/tests/test_dataloader.py +++ b/tests/test_dataloader.py @@ -25,6 +25,67 @@ def setUp(self): def tearDown(self): shutil.rmtree(self.tempdir) + @patch("mlpf.model.PFDataset.PFDataset") + def test_validation_keeps_partial_per_rank_batch(self, MockPFDataset): + """A small nvalid must still produce a batch on every distributed rank.""" + mock_pf_instance = MockPFDataset.return_value + mock_pf_instance.ds = MockDictDataset( + size=100, + keys=("X", "ytarget", "genmet"), + shapes=((1, 2), (1, 2), (1,)), + ) + + config = MLPFConfig.model_validate( + { + "dataset": "cms", + "data_dir": "/tmp/dummy_data", + "model": {"type": "attention", "attention": {"num_convs": 2}}, + "conv_type": "attention", + "train_dataset": { + "cms": { + "physical": { + "batch_size": 1, + "samples": {"sample1": {"version": "1.0.0", "splits": ["split1"]}}, + } + } + }, + "valid_dataset": { + "cms": { + "physical": { + "batch_size": 1, + "samples": {"sample1": {"version": "1.0.0", "splits": ["split1"]}}, + } + } + }, + "ntrain": 100, + "nvalid": 100, + "num_workers": 1, + "prefetch_factor": 2, + "sort_data": False, + "pad_to_multiple_elements": None, + "gpu_batch_multiplier": 64, + "sampler_mode": "interleaved-shards", + } + ) + + loaders, _ = get_interleaved_dataloaders( + world_size=8, + rank=2, + config=config, + use_cuda=False, + use_ray=False, + shuffle_train=False, + ) + + train_data_loader = loaders["train"].data_loader.data_loaders[0] + valid_loader = loaders["valid"] + valid_data_loader = valid_loader.data_loaders[0] + + self.assertTrue(train_data_loader.drop_last) + self.assertFalse(valid_data_loader.drop_last) + self.assertEqual(len(valid_loader), 1) + self.assertEqual(next(iter(valid_loader)).X.shape[0], 13) + @patch("mlpf.model.PFDataset.PFDataset") def test_restoration(self, MockPFDataset): """Ensures that the dataloader state is correctly saved and restored.""" From 1f41c27b3b22fdd05ca059054c0021e1fd4a0945 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 4 Sep 2026 14:05:11 +0300 Subject: [PATCH 09/29] format --- mlpf/model/PFDataset.py | 8 +-- mlpf/model/inference.py | 9 +-- mlpf/model/mlpf.py | 4 +- mlpf/model/set_decoder.py | 33 +++------- mlpf/model/set_losses.py | 71 ++++++---------------- mlpf/model/training.py | 23 ++----- mlpf/model/validation_metrics.py | 101 +++++++------------------------ mlpf/training_scenarios.py | 74 ++++++---------------- mlpf/training_submission.py | 24 ++------ tests/test_set_prediction.py | 32 +++------- tests/test_training_scenarios.py | 4 +- tests/test_validation_metrics.py | 28 +++------ 12 files changed, 100 insertions(+), 311 deletions(-) diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index 59cbe3f62..9b4bb22dc 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -721,13 +721,9 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, sampler_mode = DatasetSamplerMode(config.sampler_mode) _logger.info(f"{split}_dataset sampler_mode={sampler_mode.value} shuffle={shuffle}") if world_size > 1 and sampler_mode == DatasetSamplerMode.INTERLEAVED_SHARDS: - sampler = DistributedInterleavedShardSampler( - dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed - ) + sampler = DistributedInterleavedShardSampler(dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed) elif world_size > 1: - sampler = DistributedShardConsecutiveSampler( - dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed - ) + sampler = DistributedShardConsecutiveSampler(dataset, world_size=world_size, rank=rank, shuffle=shuffle, seed=config.seed) elif sampler_mode == DatasetSamplerMode.INTERLEAVED_SHARDS: sampler = InterleavedShardSampler(dataset, shuffle=shuffle, seed=config.seed) else: diff --git a/mlpf/model/inference.py b/mlpf/model/inference.py index aa354bf3f..fe59667a4 100644 --- a/mlpf/model/inference.py +++ b/mlpf/model/inference.py @@ -56,9 +56,7 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m ytarget = unpack_target(batch.ytarget_set.to(torch.float32), model_module) ytarget["pt"] = torch.exp(ytarget["pt"]) ytarget["energy"] = torch.exp(ytarget["energy"]) - ytarget["momentum"] = torch.stack( - [ytarget["pt"], ytarget["eta"], ytarget["sin_phi"], ytarget["cos_phi"], ytarget["energy"]], dim=-1 - ) + ytarget["momentum"] = torch.stack([ytarget["pt"], ytarget["eta"], ytarget["sin_phi"], ytarget["cos_phi"], ytarget["energy"]], dim=-1) ytarget["p4"] = torch.stack([ytarget["pt"], ytarget["eta"], ytarget["phi"], ytarget["energy"]], dim=-1) else: batch.ytarget[..., 2] = batch.ytarget_pt_orig @@ -98,10 +96,7 @@ def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_m for flat_arr, typ in [(ytarget, "target"), (ycand, "cand"), (ypred, "pred")]: collection_mask = collection_masks[typ] counts = collection_mask.sum(dim=1).cpu().numpy() - values = { - key: value[collection_mask].detach().cpu().float().contiguous().numpy() - for key, value in flat_arr.items() - } + values = {key: value[collection_mask].detach().cpu().float().contiguous().numpy() for key, value in flat_arr.items()} awk_arr = awkward.Array(values) awkvals[typ] = awkward.unflatten(awk_arr, counts) Xs = awkward.unflatten(awkward.from_numpy(X), input_counts) diff --git a/mlpf/model/mlpf.py b/mlpf/model/mlpf.py index 2506e44de..794155617 100644 --- a/mlpf/model/mlpf.py +++ b/mlpf/model/mlpf.py @@ -1469,7 +1469,9 @@ def __init__( ) _logger.info("output_mode={}".format(self.output_mode.value)) _logger.info("set_decoder parameters: {}".format(count_parameters(self.set_decoder) if self.set_decoder is not None else 0)) - _logger.info("nn_binary_particle parameters: {}".format(count_parameters(self.nn_binary_particle) if self.nn_binary_particle is not None else 0)) + _logger.info( + "nn_binary_particle parameters: {}".format(count_parameters(self.nn_binary_particle) if self.nn_binary_particle is not None else 0) + ) _logger.info("nn_pid parameters: {}".format(count_parameters(self.nn_pid) if self.nn_pid is not None else 0)) _logger.info("nn_pt parameters: {}".format(count_parameters(self.nn_pt) if self.nn_pt is not None else 0)) _logger.info("nn_eta parameters: {}".format(count_parameters(self.nn_eta) if self.nn_eta is not None else 0)) diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py index 33f46f321..2f1bc66fa 100644 --- a/mlpf/model/set_decoder.py +++ b/mlpf/model/set_decoder.py @@ -16,13 +16,9 @@ def __init__(self, embedding_dim, num_heads, ffn_dim, dropout=0.0): super().__init__() self.query_norm = nn.LayerNorm(embedding_dim) self.memory_norm = nn.LayerNorm(embedding_dim) - self.cross_attention = nn.MultiheadAttention( - embedding_dim, num_heads, dropout=dropout, batch_first=True - ) + self.cross_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) self.self_norm = nn.LayerNorm(embedding_dim) - self.self_attention = nn.MultiheadAttention( - embedding_dim, num_heads, dropout=dropout, batch_first=True - ) + self.self_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) self.ffn_norm = nn.LayerNorm(embedding_dim) self.ffn = nn.Sequential( nn.Linear(embedding_dim, ffn_dim), @@ -37,13 +33,9 @@ def forward(self, slots, memory, memory_mask): normalized_memory = self.memory_norm(memory) cross_outputs = [] for event_idx in range(memory.shape[0]): - event_memory = normalized_memory[ - event_idx : event_idx + 1, memory_mask[event_idx] - ] + event_memory = normalized_memory[event_idx : event_idx + 1, memory_mask[event_idx]] if event_memory.shape[1] == 0: - cross_outputs.append( - torch.zeros_like(cross_queries[event_idx : event_idx + 1]) - ) + cross_outputs.append(torch.zeros_like(cross_queries[event_idx : event_idx + 1])) continue event_output, _ = self.cross_attention( cross_queries[event_idx : event_idx + 1], @@ -55,9 +47,7 @@ def forward(self, slots, memory, memory_mask): slots = slots + torch.cat(cross_outputs, dim=0) normalized_slots = self.self_norm(slots) - self_output, _ = self.self_attention( - normalized_slots, normalized_slots, normalized_slots, need_weights=False - ) + self_output, _ = self.self_attention(normalized_slots, normalized_slots, normalized_slots, need_weights=False) slots = slots + self_output return slots + self.ffn(self.ffn_norm(slots)) @@ -68,19 +58,14 @@ class ParticleSetDecoder(nn.Module): def __init__(self, embedding_dim, num_classes, config): super().__init__() if embedding_dim % config.num_heads != 0: - raise ValueError( - f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}" - ) + raise ValueError(f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}") self.num_slots = config.num_slots self.queries = nn.Parameter(torch.empty(1, config.num_slots, embedding_dim)) nn.init.trunc_normal_(self.queries, std=0.02) ffn_dim = int(config.ffn_multiplier * embedding_dim) self.layers = nn.ModuleList( - ParticleSetDecoderLayer( - embedding_dim, config.num_heads, ffn_dim, config.dropout - ) - for _ in range(config.num_layers) + ParticleSetDecoderLayer(embedding_dim, config.num_heads, ffn_dim, config.dropout) for _ in range(config.num_layers) ) self.output_norm = nn.LayerNorm(embedding_dim) self.presence_head = nn.Linear(embedding_dim, 2) @@ -97,8 +82,6 @@ def forward(self, memory, memory_mask): pid = self.pid_head(slots) momentum = self.momentum_head(slots) phi_direction = F.normalize(momentum[..., 2:4], dim=-1, eps=1e-6) - momentum = torch.cat( - [momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1 - ) + momentum = torch.cat([momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1) pileup = torch.zeros_like(presence) return presence, pid, momentum, pileup diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py index 4d61aa6cb..e64d55e01 100644 --- a/mlpf/model/set_losses.py +++ b/mlpf/model/set_losses.py @@ -23,14 +23,10 @@ def _pairwise_matching_cost(target, prediction, weights): target_cls = target["cls_id"].long() presence_cost = -F.log_softmax(prediction["cls_binary"].float(), dim=-1)[:, 1:2] - pid_cost = -F.log_softmax(prediction["cls_id_onehot"].float(), dim=-1)[ - :, target_cls - ] + pid_cost = -F.log_softmax(prediction["cls_id_onehot"].float(), dim=-1)[:, target_cls] def l1_cost(feature): - return torch.abs( - prediction[feature].float()[:, None] - target[feature].float()[None, :] - ) + return torch.abs(prediction[feature].float()[:, None] - target[feature].float()[None, :]) pred_direction = F.normalize( torch.stack([prediction["sin_phi"], prediction["cos_phi"]], dim=-1).float(), @@ -64,20 +60,14 @@ def hungarian_match(targets, predictions, target_mask, weights=None): valid = target_mask[event_idx].bool() num_targets = int(valid.sum().item()) if num_targets > num_slots: - raise ValueError( - f"Event {event_idx} has {num_targets} targets but the decoder has only {num_slots} slots" - ) + raise ValueError(f"Event {event_idx} has {num_targets} targets but the decoder has only {num_slots} slots") if num_targets == 0: - empty = torch.empty( - 0, dtype=torch.long, device=predictions["cls_binary"].device - ) + empty = torch.empty(0, dtype=torch.long, device=predictions["cls_binary"].device) matches.append((empty, empty)) continue event_targets = {key: value[event_idx][valid] for key, value in targets.items()} - event_predictions = { - key: value[event_idx] for key, value in predictions.items() - } + event_predictions = {key: value[event_idx] for key, value in predictions.items()} cost = _pairwise_matching_cost(event_targets, event_predictions, weights) slot_indices, target_indices = linear_sum_assignment(cost.float().cpu().numpy()) matches.append( @@ -101,9 +91,7 @@ def set_event_loss( matches = hungarian_match(targets, predictions, target_mask, matcher_weights) device = predictions["cls_binary"].device - presence_targets = torch.zeros( - predictions["cls_binary"].shape[:2], dtype=torch.long, device=device - ) + presence_targets = torch.zeros(predictions["cls_binary"].shape[:2], dtype=torch.long, device=device) matched_predictions = {key: [] for key in ("cls_id_onehot", *REGRESSION_FEATURES)} matched_targets = {key: [] for key in ("cls_id", *REGRESSION_FEATURES)} @@ -115,13 +103,9 @@ def set_event_loss( for key in matched_predictions: matched_predictions[key].append(predictions[key][event_idx, slot_indices]) for key in matched_targets: - matched_targets[key].append( - targets[key][event_idx, valid_targets][target_indices] - ) + matched_targets[key].append(targets[key][event_idx, valid_targets][target_indices]) - presence_class_weights = predictions["cls_binary"].new_tensor( - [no_object_weight, 1.0] - ) + presence_class_weights = predictions["cls_binary"].new_tensor([no_object_weight, 1.0]) losses = { "Classification_binary": 10.0 * F.cross_entropy( @@ -143,27 +127,15 @@ def set_event_loss( losses[f"Regression_{feature}"] = zero return losses, matches - matched_predictions = { - key: torch.cat(value, dim=0) for key, value in matched_predictions.items() - } - matched_targets = { - key: torch.cat(value, dim=0) for key, value in matched_targets.items() - } - losses["Classification"] = F.cross_entropy( - matched_predictions["cls_id_onehot"], matched_targets["cls_id"] - ) + matched_predictions = {key: torch.cat(value, dim=0) for key, value in matched_predictions.items()} + matched_targets = {key: torch.cat(value, dim=0) for key, value in matched_targets.items()} + losses["Classification"] = F.cross_entropy(matched_predictions["cls_id_onehot"], matched_targets["cls_id"]) - sqrt_target_pt = torch.sqrt( - torch.exp(matched_targets["pt"].float()).clamp_min(1e-6) - ) + sqrt_target_pt = torch.sqrt(torch.exp(matched_targets["pt"].float()).clamp_min(1e-6)) for feature in REGRESSION_FEATURES: prediction = torch.nan_to_num(matched_predictions[feature].float()) - per_particle = regression_weights[feature] * F.mse_loss( - prediction, matched_targets[feature].float(), reduction="none" - ) - losses[f"Regression_{feature}"] = ( - per_particle * sqrt_target_pt - ).sum() / num_matched + per_particle = regression_weights[feature] * F.mse_loss(prediction, matched_targets[feature].float(), reduction="none") + losses[f"Regression_{feature}"] = (per_particle * sqrt_target_pt).sum() / num_matched return losses, matches @@ -179,15 +151,9 @@ def set_mlpf_loss( """Compute the set-prediction objective with the standard task names.""" if batch.target_mask is None: - raise ValueError( - "Set prediction requires batch.ytarget_set and batch.target_mask" - ) + raise ValueError("Set prediction requires batch.ytarget_set and batch.target_mask") - effective_regression_weights = ( - regression_weights - if task_loss_weighter is None - else {feature: 1.0 for feature in REGRESSION_FEATURES} - ) + effective_regression_weights = regression_weights if task_loss_weighter is None else {feature: 1.0 for feature in REGRESSION_FEATURES} losses, _ = set_event_loss( targets, predictions, @@ -213,8 +179,5 @@ def set_mlpf_loss( detached_losses = {key: value.detach() for key, value in losses.items()} if diagnostics is not None: - diagnostics = { - name: {task: value.detach() for task, value in values.items()} - for name, values in diagnostics.items() - } + diagnostics = {name: {task: value.detach() for task, value in values.items()} for name, values in diagnostics.items()} return loss_opt, detached_losses, diagnostics diff --git a/mlpf/model/training.py b/mlpf/model/training.py index aac1b2069..025f74ede 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -343,14 +343,10 @@ def model_step(batch, model, loss_fn, regression_weights): model_module = model.module if hasattr(model, "module") else model if model_module.output_mode == OutputMode.SET: ytarget = unpack_target(batch.ytarget_set, model_module) - loss_opt, losses_detached, task_loss_diagnostics = set_mlpf_loss( - ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) - ) + loss_opt, losses_detached, task_loss_diagnostics = set_mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) else: ytarget = unpack_target(batch.ytarget, model_module) - loss_opt, losses_detached, task_loss_diagnostics = loss_fn( - ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) - ) + loss_opt, losses_detached, task_loss_diagnostics = loss_fn(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) return loss_opt, losses_detached, task_loss_diagnostics, ypred_raw, ypred, ytarget @@ -459,14 +455,10 @@ def train_step( model_module = model.module if hasattr(model, "module") else model if model_module.output_mode == OutputMode.SET: ytarget = unpack_target(batch.ytarget_set, model_module) - loss_opt, loss, task_loss_diagnostics = set_mlpf_loss( - ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) - ) + loss_opt, loss, task_loss_diagnostics = set_mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) else: ytarget = unpack_target(batch.ytarget, model_module) - loss_opt, loss, task_loss_diagnostics = mlpf_loss( - ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model) - ) + loss_opt, loss, task_loss_diagnostics = mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) phase_start = _record_phase_time_if_enabled(diagnostics.get("time", {}), "loss", phase_start, device_type, log_this_step) if log_this_step: _collect_step_memory(rank, "after_loss", diagnostics) @@ -798,12 +790,7 @@ def evaluate( ) # Save validation plots for first batch - if ( - model_module.output_mode == OutputMode.ELEMENTWISE - and (rank == 0 or rank == "cpu") - and ival == 0 - and config.make_plots - ): + if model_module.output_mode == OutputMode.ELEMENTWISE and (rank == 0 or rank == "cpu") and ival == 0 and config.make_plots: validation_plots(batch, ypred_raw, ytarget, ypred, tensorboard_writer, step, outdir) # Accumulate losses diff --git a/mlpf/model/validation_metrics.py b/mlpf/model/validation_metrics.py index 1acfeb5ba..01dd7616c 100644 --- a/mlpf/model/validation_metrics.py +++ b/mlpf/model/validation_metrics.py @@ -39,12 +39,8 @@ def validation_particle_collections(batch, predictions, output_mode): # Legacy/custom collaters may omit the cached absolute values. The # elementwise targets store log(target/input), so reconstruct them # without requiring the input dataset to be regenerated. - targets["pt"] = torch.exp(targets["pt"].clamp(-20.0, 20.0)) * batch.X[ - ..., 1 - ].to(torch.float32) - targets["energy"] = torch.exp( - targets["energy"].clamp(-20.0, 20.0) - ) * batch.X[..., 5].to(torch.float32) + targets["pt"] = torch.exp(targets["pt"].clamp(-20.0, 20.0)) * batch.X[..., 1].to(torch.float32) + targets["energy"] = torch.exp(targets["energy"].clamp(-20.0, 20.0)) * batch.X[..., 5].to(torch.float32) target_mask = batch.mask.bool() & (targets["cls_id"] != 0) prediction_mask = batch.mask.bool() & (predictions["cls_id"] != 0) @@ -104,12 +100,8 @@ def _clean_kinematics(collection, mask): } for name, (minimum, maximum) in limits.items(): value = collection[name][mask].detach().to(device="cpu", dtype=torch.float32) - selected[name] = torch.nan_to_num( - value, nan=0.0, posinf=maximum, neginf=minimum - ).clamp(minimum, maximum) - selected["cls_id"] = ( - collection["cls_id"][mask].detach().to(device="cpu", dtype=torch.long) - ) + selected[name] = torch.nan_to_num(value, nan=0.0, posinf=maximum, neginf=minimum).clamp(minimum, maximum) + selected["cls_id"] = collection["cls_id"][mask].detach().to(device="cpu", dtype=torch.long) selected["phi"] = torch.atan2(selected["sin_phi"], selected["cos_phi"]) return selected @@ -119,14 +111,8 @@ def _pairwise_geometry(targets, predictions): delta_phi = predictions["phi"][:, None] - targets["phi"][None, :] delta_phi = torch.remainder(delta_phi + math.pi, 2.0 * math.pi) - math.pi delta_r = torch.sqrt(delta_eta.square() + delta_phi.square()) - log_pt_ratio = torch.abs( - torch.log(predictions["pt"].clamp_min(1.0e-8))[:, None] - - torch.log(targets["pt"].clamp_min(1.0e-8))[None, :] - ) - relative_pt = ( - torch.abs(predictions["pt"][:, None] - targets["pt"][None, :]) - / targets["pt"].clamp_min(1.0e-8)[None, :] - ) + log_pt_ratio = torch.abs(torch.log(predictions["pt"].clamp_min(1.0e-8))[:, None] - torch.log(targets["pt"].clamp_min(1.0e-8))[None, :]) + relative_pt = torch.abs(predictions["pt"][:, None] - targets["pt"][None, :]) / targets["pt"].clamp_min(1.0e-8)[None, :] return delta_eta, delta_phi, delta_r, log_pt_ratio, relative_pt @@ -142,26 +128,14 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): matched_target_indices = torch.empty(0, dtype=torch.long) delta_eta = delta_phi = delta_r = log_pt_ratio = pairwise_relative_pt = None if num_targets and num_predictions: - delta_eta, delta_phi, delta_r, log_pt_ratio, pairwise_relative_pt = ( - _pairwise_geometry(targets, predictions) - ) + delta_eta, delta_phi, delta_r, log_pt_ratio, pairwise_relative_pt = _pairwise_geometry(targets, predictions) cost = (delta_r / MATCH_DR).square() + (log_pt_ratio / MATCH_LOG_PT).square() prediction_indices, target_indices = linear_sum_assignment(cost.numpy()) - matched_prediction_indices = torch.as_tensor( - prediction_indices, dtype=torch.long - ) + matched_prediction_indices = torch.as_tensor(prediction_indices, dtype=torch.long) matched_target_indices = torch.as_tensor(target_indices, dtype=torch.long) - matched_dr = ( - delta_r[matched_prediction_indices, matched_target_indices] - if delta_r is not None - else torch.empty(0) - ) - relative_pt = ( - pairwise_relative_pt[matched_prediction_indices, matched_target_indices] - if pairwise_relative_pt is not None - else torch.empty(0) - ) + matched_dr = delta_r[matched_prediction_indices, matched_target_indices] if delta_r is not None else torch.empty(0) + relative_pt = pairwise_relative_pt[matched_prediction_indices, matched_target_indices] if pairwise_relative_pt is not None else torch.empty(0) accepted = (matched_dr < MATCH_DR) & (relative_pt < MATCH_REL_PT) num_accepted = int(accepted.sum()) @@ -195,9 +169,7 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): _add(metrics, "matching/f1", 2 * num_accepted, num_targets + num_predictions) if delta_r is not None: - close_to_any_target = ( - (delta_r < MATCH_DR) & (pairwise_relative_pt < MATCH_REL_PT) - ).any(dim=1) + close_to_any_target = ((delta_r < MATCH_DR) & (pairwise_relative_pt < MATCH_REL_PT)).any(dim=1) num_duplicates = max(int(close_to_any_target.sum()) - num_accepted, 0) else: num_duplicates = 0 @@ -210,33 +182,22 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): accepted_prediction_pt = predictions["pt"][accepted_prediction_indices] accepted_target_energy = targets["energy"][accepted_target_indices] accepted_prediction_energy = predictions["energy"][accepted_prediction_indices] - accepted_relative_pt = torch.abs( - accepted_prediction_pt - accepted_target_pt - ) / accepted_target_pt.clamp_min(1.0e-8) - accepted_relative_energy = torch.abs( - accepted_prediction_energy - accepted_target_energy - ) / accepted_target_energy.clamp_min(1.0e-8) - pid_correct = ( - predictions["cls_id"][accepted_prediction_indices] - == targets["cls_id"][accepted_target_indices] - ) + accepted_relative_pt = torch.abs(accepted_prediction_pt - accepted_target_pt) / accepted_target_pt.clamp_min(1.0e-8) + accepted_relative_energy = torch.abs(accepted_prediction_energy - accepted_target_energy) / accepted_target_energy.clamp_min(1.0e-8) + pid_correct = predictions["cls_id"][accepted_prediction_indices] == targets["cls_id"][accepted_target_indices] _add(metrics, "matched/pid_accuracy", pid_correct.sum(), num_accepted) _add(metrics, "matched/delta_r_mean", matched_dr[accepted].sum(), num_accepted) _add( metrics, "matched/delta_eta_abs_mean", - delta_eta[matched_prediction_indices, matched_target_indices][accepted] - .abs() - .sum(), + delta_eta[matched_prediction_indices, matched_target_indices][accepted].abs().sum(), num_accepted, ) _add( metrics, "matched/delta_phi_abs_mean", - delta_phi[matched_prediction_indices, matched_target_indices][accepted] - .abs() - .sum(), + delta_phi[matched_prediction_indices, matched_target_indices][accepted].abs().sum(), num_accepted, ) _add( @@ -284,9 +245,7 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): pt_denominator = target_scalar_pt.clamp_min(1.0e-8) energy_residual = (prediction_energy - target_energy) / energy_denominator scalar_pt_residual = (prediction_scalar_pt - target_scalar_pt) / pt_denominator - _add( - metrics, "event/energy_response_mean", prediction_energy / energy_denominator, 1 - ) + _add(metrics, "event/energy_response_mean", prediction_energy / energy_denominator, 1) _add(metrics, "event/energy_relative_abs_error", energy_residual.abs(), 1) _add( metrics, @@ -296,20 +255,10 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): ) _add(metrics, "event/scalar_pt_relative_abs_error", scalar_pt_residual.abs(), 1) - target_px = ( - targets["pt"].to(torch.float64) * torch.cos(targets["phi"].to(torch.float64)) - ).sum() - target_py = ( - targets["pt"].to(torch.float64) * torch.sin(targets["phi"].to(torch.float64)) - ).sum() - prediction_px = ( - predictions["pt"].to(torch.float64) - * torch.cos(predictions["phi"].to(torch.float64)) - ).sum() - prediction_py = ( - predictions["pt"].to(torch.float64) - * torch.sin(predictions["phi"].to(torch.float64)) - ).sum() + target_px = (targets["pt"].to(torch.float64) * torch.cos(targets["phi"].to(torch.float64))).sum() + target_py = (targets["pt"].to(torch.float64) * torch.sin(targets["phi"].to(torch.float64))).sum() + prediction_px = (predictions["pt"].to(torch.float64) * torch.cos(predictions["phi"].to(torch.float64))).sum() + prediction_py = (predictions["pt"].to(torch.float64) * torch.sin(predictions["phi"].to(torch.float64))).sum() vector_pt_error = torch.hypot(prediction_px - target_px, prediction_py - target_py) target_met = torch.hypot(target_px, target_py) prediction_met = torch.hypot(prediction_px, prediction_py) @@ -317,9 +266,7 @@ def _accumulate_event_metrics(metrics, targets, predictions, num_classes): _add(metrics, "event/met_abs_error", torch.abs(prediction_met - target_met), 1) -def compute_validation_particle_metrics( - targets, target_mask, predictions, prediction_mask, num_classes -): +def compute_validation_particle_metrics(targets, target_mask, predictions, prediction_mask, num_classes): """Compute additive, scheme-independent particle metrics for one batch.""" metrics = _empty_metrics(num_classes) @@ -332,7 +279,5 @@ def compute_validation_particle_metrics( {name: value[event_idx] for name, value in predictions.items()}, prediction_mask[event_idx], ) - _accumulate_event_metrics( - metrics, event_targets, event_predictions, num_classes - ) + _accumulate_event_metrics(metrics, event_targets, event_predictions, num_classes) return {name: tuple(values) for name, values in metrics.items()} diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index 78d19fa77..cafc5af06 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -40,9 +40,7 @@ class ScenarioVariant(BaseModel): def reject_derived_overrides(self): invalid = DERIVED_KEYS.intersection(self.overrides) if invalid: - raise ValueError( - f"Variant overrides must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Variant overrides must not set derived keys: {sorted(invalid)}") return self @@ -56,9 +54,7 @@ class ScenarioTraining(BaseModel): def reject_derived_parameters(self): invalid = DERIVED_KEYS.intersection(self.parameters) if invalid: - raise ValueError( - f"Scenario parameters must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Scenario parameters must not set derived keys: {sorted(invalid)}") return self @@ -72,9 +68,7 @@ class TrainingScenario(BaseModel): seeds: list[int] = Field(min_length=1) training: ScenarioTraining common_overrides: dict[str, Any] = Field(default_factory=dict) - allowed_variant_differences: list[str] = Field( - default_factory=lambda: ["model.output_mode", "model.set_decoder"] - ) + allowed_variant_differences: list[str] = Field(default_factory=lambda: ["model.output_mode", "model.set_decoder"]) @model_validator(mode="after") def validate_scenario(self): @@ -86,9 +80,7 @@ def validate_scenario(self): raise ValueError("Scenario seeds must be non-negative") invalid = DERIVED_KEYS.intersection(self.common_overrides) if invalid: - raise ValueError( - f"Common overrides must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Common overrides must not set derived keys: {sorted(invalid)}") return self @@ -132,10 +124,7 @@ class PlatformProfile(BaseModel): def validate_runtime_overrides(self): invalid = set(self.runtime_overrides).difference(PLATFORM_OVERRIDE_KEYS) if invalid: - raise ValueError( - "Platform profiles may only set runtime-specific overrides; " - f"invalid keys: {sorted(invalid)}" - ) + raise ValueError("Platform profiles may only set runtime-specific overrides; " f"invalid keys: {sorted(invalid)}") return self @@ -166,13 +155,8 @@ def load_training_scenario(path): def load_platform_profile(path): profile = PlatformProfile.model_validate(_read_yaml(path)) profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) - profile.experiments_dir = os.path.expandvars( - os.path.expanduser(profile.experiments_dir) - ) - profile.environment = { - key: os.path.expandvars(os.path.expanduser(value)) - for key, value in profile.environment.items() - } + profile.experiments_dir = os.path.expandvars(os.path.expanduser(profile.experiments_dir)) + profile.environment = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.environment.items()} return profile @@ -229,8 +213,7 @@ def _training_batch_size(config): batch_sizes = {dataset.batch_size for dataset in physical_datasets} if len(batch_sizes) != 1: raise ValueError( - "Automatic global-batch resolution requires every physical training dataset " - f"to use the same batch size, got {sorted(batch_sizes)}" + "Automatic global-batch resolution requires every physical training dataset " f"to use the same batch size, got {sorted(batch_sizes)}" ) return next(iter(batch_sizes)) @@ -256,16 +239,12 @@ def resolve_scenario_job( extra_overrides=None, ): if variant_name not in scenario.variants: - raise ValueError( - f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}" - ) + raise ValueError(f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}") variant = scenario.variants[variant_name] extra_overrides = extra_overrides or {} invalid = DERIVED_KEYS.intersection(extra_overrides) if invalid: - raise ValueError( - f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}" - ) + raise ValueError(f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}") settings = _merge_settings(scenario, platform, variant, extra_overrides) settings["seed"] = seed selected_spec = str(spec_file or scenario.spec_file) @@ -279,19 +258,14 @@ def resolve_scenario_job( extra_args=_settings_as_extra_args(settings), ) - target_global_batch = ( - global_batch_size - if global_batch_size is not None - else scenario.training.global_batch_size - ) + target_global_batch = global_batch_size if global_batch_size is not None else scenario.training.global_batch_size if target_global_batch <= 0: raise ValueError("global_batch_size must be positive") dataset_batch_size = _training_batch_size(config) divisor = platform.gpus * dataset_batch_size if target_global_batch % divisor: raise ValueError( - f"global_batch_size={target_global_batch} is not divisible by " - f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" + f"global_batch_size={target_global_batch} is not divisible by " f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" ) multiplier = target_global_batch // divisor settings["gpu_batch_multiplier"] = multiplier @@ -331,9 +305,7 @@ def _flatten(value, prefix=""): def _difference_allowed(path, allowed_paths): - return any( - path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths - ) + return any(path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths) def validate_variant_invariants(jobs, allowed_paths): @@ -345,17 +317,11 @@ def validate_variant_invariants(jobs, allowed_paths): differences = { path: (reference.get(path), candidate.get(path)) for path in sorted(set(reference) | set(candidate)) - if reference.get(path) != candidate.get(path) - and not _difference_allowed(path, allowed_paths) + if reference.get(path) != candidate.get(path) and not _difference_allowed(path, allowed_paths) } if differences: - details = ", ".join( - f"{path}: {values[0]!r} != {values[1]!r}" - for path, values in differences.items() - ) - raise ValueError( - f"Scenario variants differ outside allowed fields: {details}" - ) + details = ", ".join(f"{path}: {values[0]!r} != {values[1]!r}" for path, values in differences.items()) + raise ValueError(f"Scenario variants differ outside allowed fields: {details}") def resolve_scenario_jobs( @@ -442,9 +408,7 @@ def _experiment_path(platform, job, timestamp=None): def run_scenario_job(job, scenario, platform, spec_file, *, dry_run=False): - experiment_dir = _experiment_path( - platform, job, timestamp="TIMESTAMP" if dry_run else None - ) + experiment_dir = _experiment_path(platform, job, timestamp="TIMESTAMP" if dry_run else None) command = _pipeline_command(job, scenario, platform, spec_file, experiment_dir) print(shlex.join(command), flush=True) if dry_run: @@ -480,9 +444,7 @@ def _parse_set_overrides(values): def _validate_slurm_allocation(platform): allocated = os.environ.get("SLURM_GPUS_PER_NODE") if allocated and allocated.isdigit() and int(allocated) != platform.gpus: - raise ValueError( - f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}" - ) + raise ValueError(f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}") def main(argv=None): diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py index 4c8cd9ec7..3e131238c 100644 --- a/mlpf/training_submission.py +++ b/mlpf/training_submission.py @@ -44,13 +44,8 @@ def resolve_flatiron_profile_path(reference, repo_root): def available_choices(repo_root, site="flatiron"): - scenarios = sorted( - path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml") - ) - accelerators = sorted( - path.stem.removeprefix(f"{site}_") - for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml") - ) + scenarios = sorted(path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml")) + accelerators = sorted(path.stem.removeprefix(f"{site}_") for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml")) return scenarios, accelerators @@ -92,9 +87,7 @@ def build_slurm_submission( scenario.seeds = [seed] profile = load_platform_profile(profile_path) if profile.slurm is None: - raise ValueError( - f"Platform profile {profile.name!r} has no Slurm configuration" - ) + raise ValueError(f"Platform profile {profile.name!r} has no Slurm configuration") spec_file = Path(scenario.spec_file) if not spec_file.is_absolute(): @@ -105,9 +98,7 @@ def build_slurm_submission( slurm = profile.slurm logs_dir = repo_root / "logs_slurm" - worker = ( - Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") - ) + worker = Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") command = [ "sbatch", "--time", @@ -170,9 +161,7 @@ def main(argv=None, *, site="flatiron"): action="store_true", help="Print the sbatch command without submitting", ) - parser.add_argument( - "--list", action="store_true", help="List available scenarios and accelerators" - ) + parser.add_argument("--list", action="store_true", help="List available scenarios and accelerators") args = parser.parse_args(argv) repo_root = Path(__file__).resolve().parents[1] @@ -194,8 +183,7 @@ def main(argv=None, *, site="flatiron"): worker=_worker_for_site(repo_root, site), ) print( - f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" - + shlex.join(command), + f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" + shlex.join(command), flush=True, ) if args.dry_run: diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py index 96815509c..c00dabe14 100644 --- a/tests/test_set_prediction.py +++ b/tests/test_set_prediction.py @@ -8,9 +8,7 @@ from mlpf.model.utils import unpack_predictions, unpack_target -REGRESSION_WEIGHTS = { - feature: 1.0 for feature in ("pt", "eta", "sin_phi", "cos_phi", "energy") -} +REGRESSION_WEIGHTS = {feature: 1.0 for feature in ("pt", "eta", "sin_phi", "cos_phi", "energy")} def make_config(num_slots=4): @@ -124,9 +122,7 @@ def test_set_model_output_axis_is_num_slots(): assert pid.shape == (2, 4, config.num_classes) assert momentum.shape == (2, 4, 5) assert pileup.shape == (2, 4, 2) - torch.testing.assert_close( - torch.linalg.vector_norm(momentum[..., 2:4], dim=-1), torch.ones(2, 4) - ) + torch.testing.assert_close(torch.linalg.vector_norm(momentum[..., 2:4], dim=-1), torch.ones(2, 4)) def test_attention_set_model_has_no_unused_elementwise_parameters(): @@ -141,9 +137,7 @@ def test_attention_set_model_has_no_unused_elementwise_parameters(): assert model.classification_norm is None assert model.regression_norm is None - assert [ - name for name, parameter in model.named_parameters() if parameter.grad is None - ] == [] + assert [name for name, parameter in model.named_parameters() if parameter.grad is None] == [] def test_hungarian_match_finds_permuted_particles(): @@ -151,9 +145,7 @@ def test_hungarian_match_finds_permuted_particles(): targets = unpack_target(ytarget_tensor, None) predictions = { "cls_binary": torch.tensor([[[-5.0, 5.0], [-5.0, 5.0]]]), - "cls_id_onehot": torch.tensor( - [[[-5.0, -5.0, 5.0, -5.0, -5.0, -5.0], [-5.0, 5.0, -5.0, -5.0, -5.0, -5.0]]] - ), + "cls_id_onehot": torch.tensor([[[-5.0, -5.0, 5.0, -5.0, -5.0, -5.0], [-5.0, 5.0, -5.0, -5.0, -5.0, -5.0]]]), "pt": targets["pt"].flip(1), "eta": targets["eta"].flip(1), "sin_phi": targets["sin_phi"].flip(1), @@ -182,17 +174,13 @@ def test_set_loss_is_target_permutation_invariant(): "energy": torch.randn(1, 4, requires_grad=True), } targets = unpack_target(target_tensor, None) - losses, _ = set_event_loss( - targets, predictions, batch.target_mask, REGRESSION_WEIGHTS - ) + losses, _ = set_event_loss(targets, predictions, batch.target_mask, REGRESSION_WEIGHTS) permutation = torch.tensor([1, 0]) permuted_tensor = target_tensor[:, permutation] permuted_targets = unpack_target(permuted_tensor, None) permuted_mask = permuted_tensor[..., 0] != 0 - permuted_losses, _ = set_event_loss( - permuted_targets, predictions, permuted_mask, REGRESSION_WEIGHTS - ) + permuted_losses, _ = set_event_loss(permuted_targets, predictions, permuted_mask, REGRESSION_WEIGHTS) for key in losses: torch.testing.assert_close(losses[key], permuted_losses[key]) @@ -231,9 +219,7 @@ def test_set_loss_supports_an_event_without_targets(): "energy": torch.randn(1, 4, requires_grad=True), } - losses, matches = set_event_loss( - targets, predictions, batch.target_mask, REGRESSION_WEIGHTS - ) + losses, matches = set_event_loss(targets, predictions, batch.target_mask, REGRESSION_WEIGHTS) loss = sum(losses.values()) loss.backward() @@ -270,9 +256,7 @@ def test_set_model_10k_inputs_forward_backward(): raw_predictions = model(batch.X, batch.mask) predictions = unpack_predictions(raw_predictions) targets = unpack_target(batch.ytarget_set, model) - losses, _ = set_event_loss( - targets, predictions, batch.target_mask, REGRESSION_WEIGHTS - ) + losses, _ = set_event_loss(targets, predictions, batch.target_mask, REGRESSION_WEIGHTS) loss = sum(losses.values()) loss.backward() diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index a97c85372..9dfd48e41 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -156,6 +156,4 @@ def test_experiments_are_grouped_under_the_scenario_directory(): path = _experiment_path(platform, job, timestamp="TIMESTAMP") - assert path == Path( - "experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP" - ) + assert path == Path("experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP") diff --git a/tests/test_validation_metrics.py b/tests/test_validation_metrics.py index 71ed74f25..37fc19ee5 100644 --- a/tests/test_validation_metrics.py +++ b/tests/test_validation_metrics.py @@ -82,15 +82,11 @@ def test_common_particle_metrics_are_prediction_permutation_invariant(): energy=[24.0, 12.0], ) masks = (torch.ones(1, 2, dtype=torch.bool),) * 2 - original = compute_validation_particle_metrics( - targets, masks[0], predictions, masks[1], num_classes=6 - ) + original = compute_validation_particle_metrics(targets, masks[0], predictions, masks[1], num_classes=6) permutation = torch.tensor([1, 0]) permuted = {name: value[:, permutation] for name, value in predictions.items()} - reordered = compute_validation_particle_metrics( - targets, masks[0], permuted, masks[1], num_classes=6 - ) + reordered = compute_validation_particle_metrics(targets, masks[0], permuted, masks[1], num_classes=6) assert original == reordered @@ -120,19 +116,11 @@ def test_validation_collections_restore_same_physical_targets_for_both_modes(): energy=[12.0, 24.0, 0.0], ) - element_collections = validation_particle_collections( - element_batch, predictions, OutputMode.ELEMENTWISE - ) - set_collections = validation_particle_collections( - set_batch, predictions, OutputMode.SET - ) + element_collections = validation_particle_collections(element_batch, predictions, OutputMode.ELEMENTWISE) + set_collections = validation_particle_collections(set_batch, predictions, OutputMode.SET) - torch.testing.assert_close( - element_collections[0]["pt"][:, :2], set_collections[0]["pt"] - ) - torch.testing.assert_close( - element_collections[0]["energy"][:, :2], set_collections[0]["energy"] - ) + torch.testing.assert_close(element_collections[0]["pt"][:, :2], set_collections[0]["pt"]) + torch.testing.assert_close(element_collections[0]["energy"][:, :2], set_collections[0]["energy"]) assert element_collections[1].sum() == set_collections[1].sum() == 2 assert element_collections[3].sum() == set_collections[3].sum() == 2 @@ -156,9 +144,7 @@ def test_elementwise_validation_reconstructs_targets_without_cached_values(): energy=[12.0, 18.0], ) - targets, _, _, _ = validation_particle_collections( - batch, predictions, OutputMode.ELEMENTWISE - ) + targets, _, _, _ = validation_particle_collections(batch, predictions, OutputMode.ELEMENTWISE) torch.testing.assert_close(targets["pt"], torch.tensor([[10.0, 8.0]])) torch.testing.assert_close(targets["energy"], torch.tensor([[12.0, 18.0]])) From f5c57c351d2e4551a1461d2f7b7913a88f768c78 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sat, 5 Sep 2026 17:15:56 +0300 Subject: [PATCH 10/29] Improve set-based particle training --- DOING.md | 9 +- .../scenarios/cld_hits_output_comparison.yaml | 25 ++ mlpf/conf.py | 34 +++ mlpf/model/mlpf.py | 25 +- mlpf/model/set_decoder.py | 281 ++++++++++++++++-- mlpf/model/set_losses.py | 85 ++++-- mlpf/model/training.py | 32 +- tests/test_set_prediction.py | 113 ++++++- 8 files changed, 532 insertions(+), 72 deletions(-) diff --git a/DOING.md b/DOING.md index 08a8fc612..c280bee33 100644 --- a/DOING.md +++ b/DOING.md @@ -267,6 +267,8 @@ Compare elementwise and set prediction using: - [x] Add a hit-dataset set-mode example to `particleflow_spec.yaml`. - [x] Refactor/reuse the backbone encoder without changing legacy forward behavior. - [x] Implement learned fixed queries and two decoder layers. +- [x] Add detector-seeded, input-conditioned queries with local eta-phi + cross-attention and iterative direction refinement. - [ ] Implement memory-efficient packed cross-attention plus a CPU test fallback. - [x] Implement presence, PID, and absolute-momentum heads. - [x] Assert target-slot overflow instead of truncating; add aggregate logging later. @@ -274,9 +276,12 @@ Compare elementwise and set prediction using: ### Matching and loss - [x] Implement per-event Hungarian matching. -- [x] Implement matching costs with cyclic phi handling; expose them in model configuration later. +- [x] Implement configurable, dimensionless matching costs using delta-R and + relative log-momentum scales. - [x] Implement matched presence, PID, and regression losses. - [x] Decide and test event/particle normalization and no-particle weighting. +- [x] Add configurable no-object weighting, cardinality supervision, and + auxiliary losses for intermediate decoder layers. - [x] Integrate calibrated task-loss weighting for set mode. - [ ] Log target count, active-slot count, matched cost, and unmatched-slot statistics. @@ -287,6 +292,8 @@ Compare elementwise and set prediction using: closure metrics during validation. - [ ] Update validation diagnostics for independent axes. - [x] Update `predict_particles` for set outputs and absolute inverse transforms. +- [x] Decouple set-slot presence selection from PID and expose the inference + probability threshold in configuration. - [x] Update parquet inference serialization to use separate input, target, and prediction counts. - [ ] Update particle, jet, and MET metrics for set outputs. diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index 4bc7d05e2..b1c0aa536 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -7,6 +7,31 @@ variants: 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] diff --git a/mlpf/conf.py b/mlpf/conf.py index 16e67ea26..cef13e7e9 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -44,6 +44,11 @@ class OutputMode(Enum): SET = "set" +class SetQueryInit(Enum): + LEARNED = "learned" + INPUT_CONDITIONED = "input-conditioned" + + class DatasetSamplerMode(Enum): SHARD_CONSECUTIVE = "shard-consecutive" INTERLEAVED_SHARDS = "interleaved-shards" @@ -591,6 +596,21 @@ class HitFeatureEngineeringConfig(BaseModel): calorimeter_neighborhood: bool = True +class SetMatcherConfig(BaseModel): + """Dimensionless costs used by set-output Hungarian matching.""" + + model_config = ConfigDict(extra="forbid") + + presence: float = Field(default=1.0, ge=0.0) + pid: float = Field(default=1.0, ge=0.0) + geometry: float = Field(default=1.0, ge=0.0) + pt: float = Field(default=1.0, ge=0.0) + energy: float = Field(default=0.0, ge=0.0) + dr_scale: float = Field(default=0.1, gt=0.0) + log_pt_scale: float = Field(default=0.6931471805599453, gt=0.0) + log_energy_scale: float = Field(default=0.6931471805599453, gt=0.0) + + class SetDecoderConfig(BaseModel): """Configuration for permutation-invariant particle-set prediction.""" @@ -601,6 +621,20 @@ class SetDecoderConfig(BaseModel): num_heads: int = Field(default=8, gt=0) ffn_multiplier: float = Field(default=4.0, gt=0.0) dropout: float = Field(default=0.0, ge=0.0, lt=1.0) + query_init: SetQueryInit = SetQueryInit.LEARNED + local_attention_radius: Optional[float] = Field(default=None, gt=0.0) + tracker_query_fraction: float = Field(default=0.5, ge=0.0, le=1.0) + presence_threshold: float = Field(default=0.5, gt=0.0, lt=1.0) + no_object_weight: float = Field(default=1.0, gt=0.0) + cardinality_loss_weight: float = Field(default=0.0, ge=0.0) + auxiliary_loss_weight: float = Field(default=0.0, ge=0.0) + matcher: SetMatcherConfig = Field(default_factory=SetMatcherConfig) + + @model_validator(mode="after") + def validate_query_locality(self): + if self.local_attention_radius is not None and self.query_init != SetQueryInit.INPUT_CONDITIONED: + raise ValueError("set decoder local_attention_radius requires query_init='input-conditioned'") + return self class ModelArchitectureConfig(BaseModel): diff --git a/mlpf/model/mlpf.py b/mlpf/model/mlpf.py index 794155617..5732a1c93 100644 --- a/mlpf/model/mlpf.py +++ b/mlpf/model/mlpf.py @@ -1778,7 +1778,7 @@ def final_norm_reg(self): # @torch.compile def forward(self, X_features, mask): if self.output_mode == OutputMode.SET: - return self.set_decoder(self.encode_backbone(X_features, mask), mask) + return self.set_decoder(self.encode_backbone(X_features, mask), mask, X_features) X_features = self._engineer_input_features(X_features, mask) if self.use_split_backbone: @@ -1855,16 +1855,23 @@ def predict_particles(self, X_features, mask): ypred = unpack_predictions(tuple(ypred_raw)) ypred["ispu"] = torch.softmax(ypred["ispu"], axis=-1)[:, :, -1] - # By default, use standard argmax - pred_cls = torch.argmax(ypred_raw[0], axis=-1) + if self.output_mode == OutputMode.SET: + presence_probability = torch.softmax(ypred_raw[0], dim=-1)[..., 1] + active = presence_probability >= self.config.set_decoder.presence_threshold + # A present set slot must represent a physical PID. The no-particle + # class is owned exclusively by the presence head. + physical_pid = torch.argmax(ypred_raw[1][..., 1:], dim=-1) + 1 + ypred["cls_id"] = torch.where(active, physical_pid, torch.zeros_like(physical_pid)) + else: + active = torch.argmax(ypred_raw[0], dim=-1).bool() # Zero out predictions for non-particles - ypred["cls_id"][pred_cls == 0] = 0 - ypred["pt"][pred_cls == 0] = 0 - ypred["energy"][pred_cls == 0] = 0 - ypred["eta"][pred_cls == 0] = 0 - ypred["sin_phi"][pred_cls == 0] = 0 - ypred["cos_phi"][pred_cls == 0] = 0 + ypred["cls_id"][~active] = 0 + ypred["pt"][~active] = 0 + ypred["energy"][~active] = 0 + ypred["eta"][~active] = 0 + ypred["sin_phi"][~active] = 0 + ypred["cos_phi"][~active] = 0 return ypred diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py index 2f1bc66fa..b5c199289 100644 --- a/mlpf/model/set_decoder.py +++ b/mlpf/model/set_decoder.py @@ -1,24 +1,28 @@ +import math + import torch from torch import nn from torch.nn import functional as F -class ParticleSetDecoderLayer(nn.Module): - """Pre-norm particle-query decoder layer. +def _wrapped_delta_phi(left, right): + return torch.remainder(left - right + math.pi, 2.0 * math.pi) - math.pi + - Cross-attention is evaluated event by event on only the valid input embeddings. - With ``need_weights=False``, PyTorch dispatches through scaled-dot-product - attention and can select a fused FlashAttention kernel on CUDA without retaining - the slot-by-input attention matrix. - """ +class ParticleSetDecoderLayer(nn.Module): + """Pre-norm particle-query decoder layer with optional local cross-attention.""" def __init__(self, embedding_dim, num_heads, ffn_dim, dropout=0.0): super().__init__() self.query_norm = nn.LayerNorm(embedding_dim) self.memory_norm = nn.LayerNorm(embedding_dim) - self.cross_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) + self.cross_attention = nn.MultiheadAttention( + embedding_dim, num_heads, dropout=dropout, batch_first=True + ) self.self_norm = nn.LayerNorm(embedding_dim) - self.self_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) + self.self_attention = nn.MultiheadAttention( + embedding_dim, num_heads, dropout=dropout, batch_first=True + ) self.ffn_norm = nn.LayerNorm(embedding_dim) self.ffn = nn.Sequential( nn.Linear(embedding_dim, ffn_dim), @@ -28,60 +32,279 @@ def __init__(self, embedding_dim, num_heads, ffn_dim, dropout=0.0): nn.Dropout(dropout), ) - def forward(self, slots, memory, memory_mask): + def forward( + self, + slots, + memory, + memory_mask, + query_references=None, + query_reference_mask=None, + memory_positions=None, + local_attention_radius=None, + ): cross_queries = self.query_norm(slots) normalized_memory = self.memory_norm(memory) cross_outputs = [] for event_idx in range(memory.shape[0]): - event_memory = normalized_memory[event_idx : event_idx + 1, memory_mask[event_idx]] + valid_memory = memory_mask[event_idx] + event_memory = normalized_memory[event_idx : event_idx + 1, valid_memory] if event_memory.shape[1] == 0: - cross_outputs.append(torch.zeros_like(cross_queries[event_idx : event_idx + 1])) + cross_outputs.append( + torch.zeros_like(cross_queries[event_idx : event_idx + 1]) + ) continue + + attention_mask = None + if local_attention_radius is not None: + event_positions = memory_positions[event_idx, valid_memory] + reference = query_references[event_idx] + delta_eta = reference[:, None, 0] - event_positions[None, :, 0] + delta_phi = _wrapped_delta_phi( + reference[:, None, 1], event_positions[None, :, 1] + ) + attention_mask = ( + delta_eta.square() + delta_phi.square() > local_attention_radius**2 + ) + if query_reference_mask is not None: + attention_mask[~query_reference_mask[event_idx]] = False + + # A sparse or malformed event must never produce a fully masked + # attention row. Fall back to its nearest valid input. + fully_masked = attention_mask.all(dim=1) + if fully_masked.any(): + distance = delta_eta.square() + delta_phi.square() + nearest = distance[fully_masked].argmin(dim=1) + attention_mask[fully_masked] = True + attention_mask[fully_masked, nearest] = False + event_output, _ = self.cross_attention( cross_queries[event_idx : event_idx + 1], event_memory, event_memory, + attn_mask=attention_mask, need_weights=False, ) cross_outputs.append(event_output) slots = slots + torch.cat(cross_outputs, dim=0) normalized_slots = self.self_norm(slots) - self_output, _ = self.self_attention(normalized_slots, normalized_slots, normalized_slots, need_weights=False) + self_output, _ = self.self_attention( + normalized_slots, normalized_slots, normalized_slots, need_weights=False + ) slots = slots + self_output return slots + self.ffn(self.ffn_norm(slots)) class ParticleSetDecoder(nn.Module): - """Decode a fixed bank of learned queries into an unordered particle set.""" + """Decode learned or detector-seeded queries into an unordered particle set.""" def __init__(self, embedding_dim, num_classes, config): super().__init__() if embedding_dim % config.num_heads != 0: - raise ValueError(f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}") + raise ValueError( + f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}" + ) self.num_slots = config.num_slots + self.query_init = getattr(config.query_init, "value", config.query_init) + self.local_attention_radius = config.local_attention_radius + self.tracker_query_fraction = config.tracker_query_fraction + self.use_auxiliary_losses = config.auxiliary_loss_weight > 0 self.queries = nn.Parameter(torch.empty(1, config.num_slots, embedding_dim)) nn.init.trunc_normal_(self.queries, std=0.02) ffn_dim = int(config.ffn_multiplier * embedding_dim) self.layers = nn.ModuleList( - ParticleSetDecoderLayer(embedding_dim, config.num_heads, ffn_dim, config.dropout) for _ in range(config.num_layers) + ParticleSetDecoderLayer( + embedding_dim, config.num_heads, ffn_dim, config.dropout + ) + for _ in range(config.num_layers) ) self.output_norm = nn.LayerNorm(embedding_dim) self.presence_head = nn.Linear(embedding_dim, 2) self.pid_head = nn.Linear(embedding_dim, num_classes) - self.momentum_head = nn.Linear(embedding_dim, 5) - - def forward(self, memory, memory_mask): - slots = self.queries.expand(memory.shape[0], -1, -1) - for layer in self.layers: - slots = layer(slots, memory, memory_mask.bool()) - slots = self.output_norm(slots) - - presence = self.presence_head(slots) - pid = self.pid_head(slots) - momentum = self.momentum_head(slots) - phi_direction = F.normalize(momentum[..., 2:4], dim=-1, eps=1e-6) - momentum = torch.cat([momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1) + + if self.query_init == "input-conditioned": + self.seed_projection = nn.Linear(embedding_dim, embedding_dim) + self.reference_embedding = nn.Sequential( + nn.Linear(3, embedding_dim), + nn.GELU(), + nn.Linear(embedding_dim, embedding_dim), + ) + self.reference_delta_heads = nn.ModuleList( + nn.Linear(embedding_dim, 2) for _ in self.layers + ) + self.scale_head = nn.Linear(embedding_dim, 2) + self.momentum_head = None + else: + self.seed_projection = None + self.reference_embedding = None + self.reference_delta_heads = None + self.scale_head = None + self.momentum_head = nn.Linear(embedding_dim, 5) + + # Populated on every forward pass. The main four-tensor return signature + # remains unchanged for inference and elementwise compatibility. + self.auxiliary_outputs = [] + + @staticmethod + def _take_topk(scores, candidates, count): + count = min(count, int(candidates.sum().item())) + if count == 0: + return torch.empty(0, dtype=torch.long, device=scores.device) + ranked = scores.masked_fill(~candidates, -torch.inf) + return torch.topk(ranked, count, sorted=True).indices + + def _input_conditioned_queries(self, memory, memory_mask, input_features): + if input_features is None or input_features.shape[-1] < 6: + raise ValueError( + "Input-conditioned set queries require raw input features through energy" + ) + + batch_size = memory.shape[0] + slots = self.queries.expand(batch_size, -1, -1).clone() + references = memory.new_zeros( + (batch_size, self.num_slots, 2), dtype=torch.float32 + ) + reference_mask = torch.zeros( + (batch_size, self.num_slots), dtype=torch.bool, device=memory.device + ) + num_tracker_slots = round(self.num_slots * self.tracker_query_fraction) + + element_type = input_features[..., 0] + proposal_score = torch.log1p( + input_features[..., 1].float().abs() + ) + torch.log1p(input_features[..., 5].float().abs()) + proposal_score = torch.nan_to_num( + proposal_score, nan=-torch.inf, posinf=1.0e6, neginf=-torch.inf + ) + input_eta = torch.nan_to_num( + input_features[..., 2].float(), nan=0.0, posinf=10.0, neginf=-10.0 + ).clamp(-10.0, 10.0) + input_phi = torch.atan2( + input_features[..., 3].float(), input_features[..., 4].float() + ) + + for event_idx in range(batch_size): + valid = memory_mask[event_idx].bool() + chosen_mask = torch.zeros_like(valid) + tracker = valid & (element_type[event_idx] == 1) + calorimeter = valid & (element_type[event_idx] == 2) + + tracker_indices = self._take_topk( + proposal_score[event_idx], tracker, num_tracker_slots + ) + chosen_mask[tracker_indices] = True + calo_slots = self.num_slots - len(tracker_indices) + calo_indices = self._take_topk( + proposal_score[event_idx], calorimeter & ~chosen_mask, calo_slots + ) + chosen_mask[calo_indices] = True + selected = torch.cat([tracker_indices, calo_indices]) + + remaining = self.num_slots - len(selected) + if remaining: + fallback = self._take_topk( + proposal_score[event_idx], valid & ~chosen_mask, remaining + ) + selected = torch.cat([selected, fallback]) + + num_selected = len(selected) + if num_selected == 0: + continue + reference = torch.stack( + [input_eta[event_idx, selected], input_phi[event_idx, selected]], dim=-1 + ) + position_features = torch.stack( + [ + reference[:, 0], + torch.sin(reference[:, 1]), + torch.cos(reference[:, 1]), + ], + dim=-1, + ) + slots[event_idx, :num_selected] = ( + slots[event_idx, :num_selected] + + self.seed_projection(memory[event_idx, selected]) + + self.reference_embedding(position_features).to(memory.dtype) + ) + references[event_idx, :num_selected] = reference + reference_mask[event_idx, :num_selected] = True + return slots, references, reference_mask + + def _predict(self, slots, references=None): + normalized_slots = self.output_norm(slots) + presence = self.presence_head(normalized_slots) + pid = self.pid_head(normalized_slots) + if references is None: + momentum = self.momentum_head(normalized_slots) + phi_direction = F.normalize(momentum[..., 2:4], dim=-1, eps=1e-6) + momentum = torch.cat( + [momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1 + ) + else: + scales = self.scale_head(normalized_slots) + momentum = torch.stack( + [ + scales[..., 0], + references[..., 0], + torch.sin(references[..., 1]), + torch.cos(references[..., 1]), + scales[..., 1], + ], + dim=-1, + ) pileup = torch.zeros_like(presence) return presence, pid, momentum, pileup + + def forward(self, memory, memory_mask, input_features=None): + memory_mask = memory_mask.bool() + references = reference_mask = memory_positions = None + if self.query_init == "input-conditioned": + slots, references, reference_mask = self._input_conditioned_queries( + memory, memory_mask, input_features + ) + memory_positions = torch.stack( + [ + torch.nan_to_num( + input_features[..., 2].float(), + nan=0.0, + posinf=10.0, + neginf=-10.0, + ).clamp(-10.0, 10.0), + torch.atan2( + input_features[..., 3].float(), input_features[..., 4].float() + ), + ], + dim=-1, + ) + else: + slots = self.queries.expand(memory.shape[0], -1, -1) + + outputs = [] + for layer_index, layer in enumerate(self.layers): + slots = layer( + slots, + memory, + memory_mask, + query_references=references, + query_reference_mask=reference_mask, + memory_positions=memory_positions, + local_attention_radius=self.local_attention_radius, + ) + if references is not None: + normalized_slots = self.output_norm(slots) + delta = torch.tanh( + self.reference_delta_heads[layer_index](normalized_slots) + ) + step_size = self.local_attention_radius or 1.0 + eta = references[..., 0] + step_size * delta[..., 0] + phi = references[..., 1] + step_size * delta[..., 1] + references = torch.stack( + [eta, torch.atan2(torch.sin(phi), torch.cos(phi))], dim=-1 + ) + is_final_layer = layer_index == len(self.layers) - 1 + if self.use_auxiliary_losses or is_final_layer: + outputs.append(self._predict(slots, references)) + + self.auxiliary_outputs = outputs[:-1] if self.use_auxiliary_losses else [] + return outputs[-1] diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py index e64d55e01..b5d14270e 100644 --- a/mlpf/model/set_losses.py +++ b/mlpf/model/set_losses.py @@ -1,3 +1,4 @@ +import math from dataclasses import dataclass import torch @@ -12,10 +13,12 @@ class SetMatcherWeights: presence: float = 1.0 pid: float = 1.0 + geometry: float = 1.0 pt: float = 1.0 - eta: float = 1.0 - phi: float = 1.0 - energy: float = 1.0 + energy: float = 0.0 + dr_scale: float = 0.1 + log_pt_scale: float = math.log(2.0) + log_energy_scale: float = math.log(2.0) def _pairwise_matching_cost(target, prediction, weights): @@ -25,28 +28,22 @@ def _pairwise_matching_cost(target, prediction, weights): presence_cost = -F.log_softmax(prediction["cls_binary"].float(), dim=-1)[:, 1:2] pid_cost = -F.log_softmax(prediction["cls_id_onehot"].float(), dim=-1)[:, target_cls] - def l1_cost(feature): - return torch.abs(prediction[feature].float()[:, None] - target[feature].float()[None, :]) + pred_phi = torch.atan2(prediction["sin_phi"].float(), prediction["cos_phi"].float()) + target_phi = torch.atan2(target["sin_phi"].float(), target["cos_phi"].float()) + delta_phi = pred_phi[:, None] - target_phi[None, :] + delta_phi = torch.atan2(torch.sin(delta_phi), torch.cos(delta_phi)) + delta_eta = prediction["eta"].float()[:, None] - target["eta"].float()[None, :] + delta_r = torch.sqrt(delta_eta.square() + delta_phi.square() + 1e-12) - pred_direction = F.normalize( - torch.stack([prediction["sin_phi"], prediction["cos_phi"]], dim=-1).float(), - dim=-1, - eps=1e-6, - ) - target_direction = F.normalize( - torch.stack([target["sin_phi"], target["cos_phi"]], dim=-1).float(), - dim=-1, - eps=1e-6, - ) - phi_cost = 1.0 - pred_direction @ target_direction.transpose(0, 1) + log_pt_cost = torch.abs(prediction["pt"].float()[:, None] - target["pt"].float()[None, :]) + log_energy_cost = torch.abs(prediction["energy"].float()[:, None] - target["energy"].float()[None, :]) return ( weights.presence * presence_cost + weights.pid * pid_cost - + weights.pt * l1_cost("pt") - + weights.eta * l1_cost("eta") - + weights.phi * phi_cost - + weights.energy * l1_cost("energy") + + weights.geometry * delta_r / weights.dr_scale + + weights.pt * log_pt_cost / weights.log_pt_scale + + weights.energy * log_energy_cost / weights.log_energy_scale ).detach() @@ -85,11 +82,14 @@ def set_event_loss( target_mask, regression_weights, matcher_weights=None, - no_object_weight=0.1, + no_object_weight=1.0, + cardinality_loss_weight=0.0, + matches=None, ): """Permutation-invariant particle-set loss for a padded event batch.""" - matches = hungarian_match(targets, predictions, target_mask, matcher_weights) + if matches is None: + matches = hungarian_match(targets, predictions, target_mask, matcher_weights) device = predictions["cls_binary"].device presence_targets = torch.zeros(predictions["cls_binary"].shape[:2], dtype=torch.long, device=device) @@ -107,13 +107,16 @@ def set_event_loss( presence_class_weights = predictions["cls_binary"].new_tensor([no_object_weight, 1.0]) losses = { - "Classification_binary": 10.0 - * F.cross_entropy( + "Classification_binary": F.cross_entropy( predictions["cls_binary"].reshape(-1, 2), presence_targets.reshape(-1), weight=presence_class_weights, ) } + if cardinality_loss_weight > 0: + predicted_count = F.softmax(predictions["cls_binary"].float(), dim=-1)[..., 1].sum(dim=1) + target_count = target_mask.sum(dim=1).to(dtype=predicted_count.dtype) + losses["Cardinality"] = cardinality_loss_weight * F.smooth_l1_loss(predicted_count, target_count) num_matched = int(presence_targets.sum().item()) if num_matched == 0: @@ -146,7 +149,10 @@ def set_mlpf_loss( regression_weights, task_loss_weighter=None, matcher_weights=None, - no_object_weight=0.1, + no_object_weight=1.0, + cardinality_loss_weight=0.0, + auxiliary_predictions=None, + auxiliary_loss_weight=0.0, ): """Compute the set-prediction objective with the standard task names.""" @@ -154,22 +160,43 @@ def set_mlpf_loss( raise ValueError("Set prediction requires batch.ytarget_set and batch.target_mask") effective_regression_weights = regression_weights if task_loss_weighter is None else {feature: 1.0 for feature in REGRESSION_FEATURES} - losses, _ = set_event_loss( + losses, matches = set_event_loss( targets, predictions, batch.target_mask, effective_regression_weights, matcher_weights=matcher_weights, no_object_weight=no_object_weight, + cardinality_loss_weight=cardinality_loss_weight, ) + task_losses = {task: losses[task] for task in LOSS_TASKS} if task_loss_weighter is None: - loss_opt = sum(losses.values()) + loss_opt = sum(task_losses.values()) diagnostics = None else: # Keep the same task names so the existing one-time calibration can be # evaluated for set mode rather than introducing a second mechanism. - assert tuple(losses) == LOSS_TASKS - loss_opt, diagnostics = task_loss_weighter(losses) + loss_opt, diagnostics = task_loss_weighter(task_losses) + + if "Cardinality" in losses: + loss_opt = loss_opt + losses["Cardinality"] + + if auxiliary_predictions and auxiliary_loss_weight > 0: + auxiliary_losses = [] + for auxiliary_prediction in auxiliary_predictions: + layer_losses, _ = set_event_loss( + targets, + auxiliary_prediction, + batch.target_mask, + effective_regression_weights, + matcher_weights=matcher_weights, + no_object_weight=no_object_weight, + cardinality_loss_weight=cardinality_loss_weight, + matches=matches, + ) + auxiliary_losses.append(sum(layer_losses.values())) + losses["Auxiliary"] = auxiliary_loss_weight * torch.stack(auxiliary_losses).mean() + loss_opt = loss_opt + losses["Auxiliary"] losses["Total"] = loss_opt if not torch.isfinite(loss_opt): diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 025f74ede..6f88d8053 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -87,7 +87,7 @@ mlpf_loss, particle_loss, ) -from mlpf.model.set_losses import set_mlpf_loss +from mlpf.model.set_losses import SetMatcherWeights, set_mlpf_loss from mlpf.model.validation_metrics import compute_validation_particle_metrics, validation_particle_collections from mlpf.utils import create_comet_experiment from mlpf.conf import INPUT_TYPE_LABELS, MLPFConfig, OutputMode, SOURCE_LABELS @@ -218,6 +218,18 @@ def _get_task_loss_weighter(model): return getattr(model_module, "task_loss_weighter", None) +def _set_loss_kwargs(model_module): + config = model_module.config.set_decoder + auxiliary_predictions = [unpack_predictions(prediction) for prediction in model_module.set_decoder.auxiliary_outputs] + return { + "matcher_weights": SetMatcherWeights(**config.matcher.model_dump()), + "no_object_weight": config.no_object_weight, + "cardinality_loss_weight": config.cardinality_loss_weight, + "auxiliary_predictions": auxiliary_predictions, + "auxiliary_loss_weight": config.auxiliary_loss_weight, + } + + def _format_task_diagnostic(task_diagnostic): if not task_diagnostic: return "" @@ -343,7 +355,14 @@ def model_step(batch, model, loss_fn, regression_weights): model_module = model.module if hasattr(model, "module") else model if model_module.output_mode == OutputMode.SET: ytarget = unpack_target(batch.ytarget_set, model_module) - loss_opt, losses_detached, task_loss_diagnostics = set_mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + loss_opt, losses_detached, task_loss_diagnostics = set_mlpf_loss( + ytarget, + ypred, + batch, + regression_weights, + _get_task_loss_weighter(model), + **_set_loss_kwargs(model_module), + ) else: ytarget = unpack_target(batch.ytarget, model_module) loss_opt, losses_detached, task_loss_diagnostics = loss_fn(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) @@ -455,7 +474,14 @@ def train_step( model_module = model.module if hasattr(model, "module") else model if model_module.output_mode == OutputMode.SET: ytarget = unpack_target(batch.ytarget_set, model_module) - loss_opt, loss, task_loss_diagnostics = set_mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + loss_opt, loss, task_loss_diagnostics = set_mlpf_loss( + ytarget, + ypred, + batch, + regression_weights, + _get_task_loss_weighter(model), + **_set_loss_kwargs(model_module), + ) else: ytarget = unpack_target(batch.ytarget, model_module) loss_opt, loss, task_loss_diagnostics = mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py index c00dabe14..732412b44 100644 --- a/tests/test_set_prediction.py +++ b/tests/test_set_prediction.py @@ -11,7 +11,7 @@ REGRESSION_WEIGHTS = {feature: 1.0 for feature in ("pt", "eta", "sin_phi", "cos_phi", "energy")} -def make_config(num_slots=4): +def make_config(num_slots=4, **set_decoder_overrides): return MLPFConfig.model_validate( { "dataset": "cld_hits", @@ -31,6 +31,7 @@ def make_config(num_slots=4): "num_slots": num_slots, "num_layers": 2, "num_heads": 2, + **set_decoder_overrides, }, "hit_feature_engineering": {"enabled": False}, }, @@ -92,6 +93,13 @@ def test_set_config_populates_decoder_defaults(): assert config.model.set_decoder is not None assert config.model.set_decoder.num_slots == 256 + assert config.model.set_decoder.no_object_weight == 1.0 + assert config.model.set_decoder.matcher.dr_scale == 0.1 + + +def test_set_config_rejects_local_attention_for_global_queries(): + with pytest.raises(ValueError, match="local_attention_radius requires"): + make_config(local_attention_radius=0.4) def test_set_config_rejects_non_hit_datasets(): @@ -125,6 +133,60 @@ def test_set_model_output_axis_is_num_slots(): torch.testing.assert_close(torch.linalg.vector_norm(momentum[..., 2:4], dim=-1), torch.ones(2, 4)) +def test_input_conditioned_set_decoder_forward_backward(): + config = make_config( + num_slots=6, + query_init="input-conditioned", + local_attention_radius=0.4, + tracker_query_fraction=0.5, + num_layers=3, + auxiliary_loss_weight=0.25, + ) + model = MLPF(config) + X = torch.randn(2, 12, config.input_dim) + X[..., 0] = torch.tensor([1, 2] * 6) + X[..., 1] = X[..., 1].abs() + 0.1 + X[..., 5] = X[..., 5].abs() + 0.1 + mask = torch.ones(2, 12, dtype=torch.bool) + mask[1, 9:] = False + + predictions = model(X, mask) + assert predictions[2].shape == (2, 6, 5) + assert len(model.set_decoder.auxiliary_outputs) == 2 + assert all(torch.isfinite(output).all() for prediction in predictions for output in [prediction]) + + auxiliary = [output for prediction in model.set_decoder.auxiliary_outputs for output in prediction] + sum(output.square().mean() for output in [*predictions, *auxiliary]).backward() + assert [name for name, parameter in model.named_parameters() if parameter.grad is None] == [] + + +def test_model_step_applies_configured_cardinality_and_auxiliary_losses(): + from mlpf.model.training import model_step + + config = make_config( + num_slots=6, + query_init="input-conditioned", + local_attention_radius=0.4, + num_layers=3, + cardinality_loss_weight=0.05, + auxiliary_loss_weight=0.25, + ) + model = MLPF(config) + X = torch.randn(1, 10, config.input_dim) + X[..., 0] = torch.tensor([1, 2] * 5) + X[..., 1] = X[..., 1].abs() + 0.1 + X[..., 5] = X[..., 5].abs() + 0.1 + batch = PFBatch(X=X, ytarget_set=make_target_tensor(num_targets=2)) + + loss, losses, _, _, _, _ = model_step(batch, model, None, REGRESSION_WEIGHTS) + loss.backward() + + assert torch.isfinite(loss) + assert losses["Cardinality"] > 0 + assert losses["Auxiliary"] > 0 + assert model.set_decoder.reference_delta_heads[0].weight.grad is not None + + def test_attention_set_model_has_no_unused_elementwise_parameters(): config = make_attention_config() model = MLPF(config) @@ -230,6 +292,38 @@ def test_set_loss_supports_an_event_without_targets(): assert all(prediction.grad is not None for prediction in predictions.values()) +def test_cardinality_loss_penalizes_excess_present_slots(): + target_tensor = make_target_tensor(num_targets=2) + targets = unpack_target(target_tensor, None) + target_mask = torch.ones(1, 2, dtype=torch.bool) + predictions = { + "cls_binary": torch.tensor([[[0.0, 5.0], [0.0, 5.0], [5.0, 0.0], [5.0, 0.0]]]), + "cls_id_onehot": torch.zeros(1, 4, 6), + "pt": torch.zeros(1, 4), + "eta": torch.zeros(1, 4), + "sin_phi": torch.zeros(1, 4), + "cos_phi": torch.ones(1, 4), + "energy": torch.zeros(1, 4), + } + calibrated_losses, _ = set_event_loss( + targets, + predictions, + target_mask, + REGRESSION_WEIGHTS, + cardinality_loss_weight=1.0, + ) + predictions["cls_binary"] = torch.tensor([[[0.0, 5.0]] * 4]) + excess_losses, _ = set_event_loss( + targets, + predictions, + target_mask, + REGRESSION_WEIGHTS, + cardinality_loss_weight=1.0, + ) + + assert calibrated_losses["Cardinality"] < excess_losses["Cardinality"] + + def test_predict_particles_restores_absolute_set_kinematics(): model = MLPF(make_config(num_slots=3)).eval() X = torch.ones(1, 6, 15) @@ -242,6 +336,23 @@ def test_predict_particles_restores_absolute_set_kinematics(): assert torch.all(prediction["energy"] >= 0) +def test_set_presence_threshold_controls_inference_selection(): + model = MLPF(make_config(num_slots=2, presence_threshold=0.9)).eval() + presence = torch.tensor([[[0.0, 2.0], [0.0, 4.0]]]) + pid = torch.full((1, 2, model.num_classes), -5.0) + pid[0, 0, 2] = 5.0 + pid[0, 1, 3] = 5.0 + momentum = torch.zeros(1, 2, 5) + momentum[..., 3] = 1.0 + model.forward = lambda _features, _mask: (presence, pid, momentum, torch.zeros_like(presence)) + + prediction = model.predict_particles(torch.ones(1, 3, 15), torch.ones(1, 3, dtype=torch.bool)) + + assert prediction["cls_id"].tolist() == [[0, 3]] + assert prediction["pt"][0, 0] == 0 + assert prediction["pt"][0, 1] > 0 + + def test_set_model_10k_inputs_forward_backward(): torch.manual_seed(7) config = make_config(num_slots=256) From 44a3dbaf74be4c33918b69e788abf9e105c8b8f3 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sat, 5 Sep 2026 17:16:34 +0300 Subject: [PATCH 11/29] update batches --- configs/training/scenarios/cld_hits_output_comparison.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index b1c0aa536..667911669 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -42,8 +42,8 @@ training: parameters: lr: 0.001 num_steps: 20000 - val_freq: 1000 - checkpoint_freq: 1000 + val_freq: 2000 + checkpoint_freq: 2000 nvalid: 512 ntest: 512 sampler_mode: interleaved-shards From 09fbdfb27591fe641dafbf43e4bdf5733d7a9e1d Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sat, 5 Sep 2026 18:24:13 +0300 Subject: [PATCH 12/29] add studies readme --- notebooks/studies/README.md | 163 ++++++++++++++++++++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 notebooks/studies/README.md diff --git a/notebooks/studies/README.md b/notebooks/studies/README.md new file mode 100644 index 000000000..72f5c6580 --- /dev/null +++ b/notebooks/studies/README.md @@ -0,0 +1,163 @@ +# Reproducible notebook studies + +Use this directory for dated, reproducible analyses and presentations. A study +should remain renderable after its source experiment directory has been moved or +deleted. + +## Directory layout + +Name a study `YYYYMMDD_short_description` and use this layout: + +```text +YYYYMMDD_short_description/ +├── README.md +├── study_name.ipynb +├── render.sh +├── inputs/ +│ └── / +│ ├── history/ +│ └── tensorboard/ +│ ├── train/ +│ └── valid/ +└── output/ + ├── study_name.executed.ipynb + ├── study_name.slides.html + ├── study_name.slides.pdf + └── generated figures +``` + +Use descriptive names such as `elementwise`, `set`, or `pf_baseline` inside +`inputs/`; do not retain timestamp-heavy experiment directory names as the only +description of a run. + +## Archive the inputs + +Copy the inputs needed to reproduce the notebook into the study directory. Do +not use symlinks or make the notebook read mutable paths under `experiments/`, +scratch storage, or a user's home directory. + +For a training comparison, the useful archive normally includes: + +- history JSON files used for curves and tables; +- the rank-zero training log used for timing, memory, or loss calibration; +- TensorBoard event files from both the training and validation writers, even + when the current notebook reads only one of them; +- the training configuration, hyperparameters, scenario manifest, and resolved + `particleflow_spec.yaml`; +- the final validation plots or compact source data used in the presentation. + +Copy files without modifying their contents. Renaming a rank-zero log to +`train.log` is fine when documented by the directory structure. Preserve the +complete history and final plot directory rather than selecting only individual +points or panels; this leaves enough context for later follow-up plots. + +Keep training and validation TensorBoard event files in separate directories, +preferably `tensorboard/train/` and `tensorboard/valid/`. Do not combine the two +writers in one directory: they can use overlapping tag names, and TensorBoard or +`EventAccumulator` may otherwise merge them into a misleading scalar history. + +Avoid copying large derived artifacts that are not required by the notebook, +especially checkpoints, per-checkpoint prediction parquet files, and full local +dataset caches. If one of these is essential, include only the necessary subset +and explain the choice and size in the study README. + +A study may reuse a baseline already archived by an earlier study instead of +duplicating it. Point only to that dated study, and document the dependency in +the new README. New campaign inputs must still be archived under the new study. + +## Make the notebook portable + +Resolve paths from the study directory and write generated files only under +`output/`. A typical setup cell is: + +```python +from pathlib import Path + +try: + STUDY_DIR = Path.cwd() + if not (STUDY_DIR / "README.md").is_file(): + STUDY_DIR = Path("notebooks/studies/YYYYMMDD_short_description").resolve() +except NameError: + STUDY_DIR = Path.cwd() + +INPUT_DIR = STUDY_DIR / "inputs" +OUTPUT_DIR = STUDY_DIR / "output" +OUTPUT_DIR.mkdir(exist_ok=True) +``` + +Do not embed absolute paths. Keep exploratory code out of the final slide export +with notebook tags such as `hide-input`, while retaining it in the source and +executed notebooks. + +## Document provenance and limitations + +The study README should record: + +- the question being studied and the compared runs; +- exact run identifiers, timestamps, random seeds, dataset names and versions; +- hardware, batch size, training length, and the selected validation step; +- what is present under `inputs/` and what was intentionally excluded; +- dependencies on earlier archived studies; +- commands needed to render the notebook or regenerate validation inputs; +- known comparison limitations, incomplete jobs, or differences in sample size. + +Keep important caveats visible in the slides as well as in the README. + +## Render and verify + +Provide an executable `render.sh` that can be called from the repository root. +It should execute the source notebook, export HTML slides, and export a PDF when +Chrome or Chromium is available. + +Before considering the study complete: + +1. Confirm that the source notebook is valid JSON. +2. Search the notebook for references to `experiments/`, absolute paths, and + temporary locations. +3. Execute `render.sh` using only the archived input paths. +4. Check that every expected figure, the executed notebook, HTML, and PDF were + regenerated successfully. +5. Open or render the PDF and inspect the first, representative middle, and last + pages for clipping, blank pages, or missing images. +6. Review the final archive size and verify that no checkpoint, dataset cache, or + unintended prediction dump was included. +7. Confirm that every archived training run has nonempty, separately stored + TensorBoard event files for both training and validation. + +## Upload to the private study bucket + +The off-machine archive is the private Hugging Face bucket +[`jpata/particleflow-studies`](https://huggingface.co/buckets/jpata/particleflow-studies). +Access requires a Hugging Face account authorized for the bucket and an +authenticated `hf` CLI. Keep each dated directory at the same top-level name in +the bucket, and upload this README along with the studies. + +From the repository root, prepare and review a non-deleting plan for the dated +study directory, apply it, and copy this README separately: + +```bash +.venv/bin/hf auth whoami +.venv/bin/hf buckets sync \ + notebooks/studies/YYYYMMDD_short_description \ + hf://buckets/jpata/particleflow-studies/YYYYMMDD_short_description \ + --no-delete \ + --plan /tmp/YYYYMMDD_short_description-sync.jsonl +.venv/bin/hf buckets sync \ + --apply /tmp/YYYYMMDD_short_description-sync.jsonl +.venv/bin/hf buckets cp \ + notebooks/studies/README.md \ + hf://buckets/jpata/particleflow-studies/README.md +``` + +After uploading, run the same command with a new `--plan` path. A complete sync +should report zero uploads, downloads, and deletes. Sync dated directories +individually so sibling transfer archives, such as `.zip` or `.tar.gz` files, are +not uploaded accidentally. Do not use `--delete` unless remote removal is +intentional and the complete plan has been reviewed. + +For a standalone transfer archive, run from the repository root: + +```bash +tar -czf YYYYMMDD_short_description.tar.gz \ + notebooks/studies/YYYYMMDD_short_description +``` From be8b0b5cf1e87abc09c16a305784098eb7b95362 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sat, 5 Sep 2026 19:54:06 +0300 Subject: [PATCH 13/29] format --- mlpf/model/set_decoder.py | 101 ++++++++++---------------------------- 1 file changed, 25 insertions(+), 76 deletions(-) diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py index b5c199289..2d9fce902 100644 --- a/mlpf/model/set_decoder.py +++ b/mlpf/model/set_decoder.py @@ -16,13 +16,9 @@ def __init__(self, embedding_dim, num_heads, ffn_dim, dropout=0.0): super().__init__() self.query_norm = nn.LayerNorm(embedding_dim) self.memory_norm = nn.LayerNorm(embedding_dim) - self.cross_attention = nn.MultiheadAttention( - embedding_dim, num_heads, dropout=dropout, batch_first=True - ) + self.cross_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) self.self_norm = nn.LayerNorm(embedding_dim) - self.self_attention = nn.MultiheadAttention( - embedding_dim, num_heads, dropout=dropout, batch_first=True - ) + self.self_attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True) self.ffn_norm = nn.LayerNorm(embedding_dim) self.ffn = nn.Sequential( nn.Linear(embedding_dim, ffn_dim), @@ -49,9 +45,7 @@ def forward( valid_memory = memory_mask[event_idx] event_memory = normalized_memory[event_idx : event_idx + 1, valid_memory] if event_memory.shape[1] == 0: - cross_outputs.append( - torch.zeros_like(cross_queries[event_idx : event_idx + 1]) - ) + cross_outputs.append(torch.zeros_like(cross_queries[event_idx : event_idx + 1])) continue attention_mask = None @@ -59,12 +53,8 @@ def forward( event_positions = memory_positions[event_idx, valid_memory] reference = query_references[event_idx] delta_eta = reference[:, None, 0] - event_positions[None, :, 0] - delta_phi = _wrapped_delta_phi( - reference[:, None, 1], event_positions[None, :, 1] - ) - attention_mask = ( - delta_eta.square() + delta_phi.square() > local_attention_radius**2 - ) + delta_phi = _wrapped_delta_phi(reference[:, None, 1], event_positions[None, :, 1]) + attention_mask = delta_eta.square() + delta_phi.square() > local_attention_radius**2 if query_reference_mask is not None: attention_mask[~query_reference_mask[event_idx]] = False @@ -88,9 +78,7 @@ def forward( slots = slots + torch.cat(cross_outputs, dim=0) normalized_slots = self.self_norm(slots) - self_output, _ = self.self_attention( - normalized_slots, normalized_slots, normalized_slots, need_weights=False - ) + self_output, _ = self.self_attention(normalized_slots, normalized_slots, normalized_slots, need_weights=False) slots = slots + self_output return slots + self.ffn(self.ffn_norm(slots)) @@ -101,9 +89,7 @@ class ParticleSetDecoder(nn.Module): def __init__(self, embedding_dim, num_classes, config): super().__init__() if embedding_dim % config.num_heads != 0: - raise ValueError( - f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}" - ) + raise ValueError(f"Set decoder embedding_dim={embedding_dim} must be divisible by num_heads={config.num_heads}") self.num_slots = config.num_slots self.query_init = getattr(config.query_init, "value", config.query_init) @@ -114,10 +100,7 @@ def __init__(self, embedding_dim, num_classes, config): nn.init.trunc_normal_(self.queries, std=0.02) ffn_dim = int(config.ffn_multiplier * embedding_dim) self.layers = nn.ModuleList( - ParticleSetDecoderLayer( - embedding_dim, config.num_heads, ffn_dim, config.dropout - ) - for _ in range(config.num_layers) + ParticleSetDecoderLayer(embedding_dim, config.num_heads, ffn_dim, config.dropout) for _ in range(config.num_layers) ) self.output_norm = nn.LayerNorm(embedding_dim) self.presence_head = nn.Linear(embedding_dim, 2) @@ -130,9 +113,7 @@ def __init__(self, embedding_dim, num_classes, config): nn.GELU(), nn.Linear(embedding_dim, embedding_dim), ) - self.reference_delta_heads = nn.ModuleList( - nn.Linear(embedding_dim, 2) for _ in self.layers - ) + self.reference_delta_heads = nn.ModuleList(nn.Linear(embedding_dim, 2) for _ in self.layers) self.scale_head = nn.Linear(embedding_dim, 2) self.momentum_head = None else: @@ -156,33 +137,19 @@ def _take_topk(scores, candidates, count): def _input_conditioned_queries(self, memory, memory_mask, input_features): if input_features is None or input_features.shape[-1] < 6: - raise ValueError( - "Input-conditioned set queries require raw input features through energy" - ) + raise ValueError("Input-conditioned set queries require raw input features through energy") batch_size = memory.shape[0] slots = self.queries.expand(batch_size, -1, -1).clone() - references = memory.new_zeros( - (batch_size, self.num_slots, 2), dtype=torch.float32 - ) - reference_mask = torch.zeros( - (batch_size, self.num_slots), dtype=torch.bool, device=memory.device - ) + references = memory.new_zeros((batch_size, self.num_slots, 2), dtype=torch.float32) + reference_mask = torch.zeros((batch_size, self.num_slots), dtype=torch.bool, device=memory.device) num_tracker_slots = round(self.num_slots * self.tracker_query_fraction) element_type = input_features[..., 0] - proposal_score = torch.log1p( - input_features[..., 1].float().abs() - ) + torch.log1p(input_features[..., 5].float().abs()) - proposal_score = torch.nan_to_num( - proposal_score, nan=-torch.inf, posinf=1.0e6, neginf=-torch.inf - ) - input_eta = torch.nan_to_num( - input_features[..., 2].float(), nan=0.0, posinf=10.0, neginf=-10.0 - ).clamp(-10.0, 10.0) - input_phi = torch.atan2( - input_features[..., 3].float(), input_features[..., 4].float() - ) + proposal_score = torch.log1p(input_features[..., 1].float().abs()) + torch.log1p(input_features[..., 5].float().abs()) + proposal_score = torch.nan_to_num(proposal_score, nan=-torch.inf, posinf=1.0e6, neginf=-torch.inf) + input_eta = torch.nan_to_num(input_features[..., 2].float(), nan=0.0, posinf=10.0, neginf=-10.0).clamp(-10.0, 10.0) + input_phi = torch.atan2(input_features[..., 3].float(), input_features[..., 4].float()) for event_idx in range(batch_size): valid = memory_mask[event_idx].bool() @@ -190,30 +157,22 @@ def _input_conditioned_queries(self, memory, memory_mask, input_features): tracker = valid & (element_type[event_idx] == 1) calorimeter = valid & (element_type[event_idx] == 2) - tracker_indices = self._take_topk( - proposal_score[event_idx], tracker, num_tracker_slots - ) + tracker_indices = self._take_topk(proposal_score[event_idx], tracker, num_tracker_slots) chosen_mask[tracker_indices] = True calo_slots = self.num_slots - len(tracker_indices) - calo_indices = self._take_topk( - proposal_score[event_idx], calorimeter & ~chosen_mask, calo_slots - ) + calo_indices = self._take_topk(proposal_score[event_idx], calorimeter & ~chosen_mask, calo_slots) chosen_mask[calo_indices] = True selected = torch.cat([tracker_indices, calo_indices]) remaining = self.num_slots - len(selected) if remaining: - fallback = self._take_topk( - proposal_score[event_idx], valid & ~chosen_mask, remaining - ) + fallback = self._take_topk(proposal_score[event_idx], valid & ~chosen_mask, remaining) selected = torch.cat([selected, fallback]) num_selected = len(selected) if num_selected == 0: continue - reference = torch.stack( - [input_eta[event_idx, selected], input_phi[event_idx, selected]], dim=-1 - ) + reference = torch.stack([input_eta[event_idx, selected], input_phi[event_idx, selected]], dim=-1) position_features = torch.stack( [ reference[:, 0], @@ -238,9 +197,7 @@ def _predict(self, slots, references=None): if references is None: momentum = self.momentum_head(normalized_slots) phi_direction = F.normalize(momentum[..., 2:4], dim=-1, eps=1e-6) - momentum = torch.cat( - [momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1 - ) + momentum = torch.cat([momentum[..., :2], phi_direction, momentum[..., 4:5]], dim=-1) else: scales = self.scale_head(normalized_slots) momentum = torch.stack( @@ -260,9 +217,7 @@ def forward(self, memory, memory_mask, input_features=None): memory_mask = memory_mask.bool() references = reference_mask = memory_positions = None if self.query_init == "input-conditioned": - slots, references, reference_mask = self._input_conditioned_queries( - memory, memory_mask, input_features - ) + slots, references, reference_mask = self._input_conditioned_queries(memory, memory_mask, input_features) memory_positions = torch.stack( [ torch.nan_to_num( @@ -271,9 +226,7 @@ def forward(self, memory, memory_mask, input_features=None): posinf=10.0, neginf=-10.0, ).clamp(-10.0, 10.0), - torch.atan2( - input_features[..., 3].float(), input_features[..., 4].float() - ), + torch.atan2(input_features[..., 3].float(), input_features[..., 4].float()), ], dim=-1, ) @@ -293,15 +246,11 @@ def forward(self, memory, memory_mask, input_features=None): ) if references is not None: normalized_slots = self.output_norm(slots) - delta = torch.tanh( - self.reference_delta_heads[layer_index](normalized_slots) - ) + delta = torch.tanh(self.reference_delta_heads[layer_index](normalized_slots)) step_size = self.local_attention_radius or 1.0 eta = references[..., 0] + step_size * delta[..., 0] phi = references[..., 1] + step_size * delta[..., 1] - references = torch.stack( - [eta, torch.atan2(torch.sin(phi), torch.cos(phi))], dim=-1 - ) + references = torch.stack([eta, torch.atan2(torch.sin(phi), torch.cos(phi))], dim=-1) is_final_layer = layer_index == len(self.layers) - 1 if self.use_auxiliary_losses or is_final_layer: outputs.append(self._predict(slots, references)) From 6199628886efaa7792ce8aa500327022d7f65457 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sat, 5 Sep 2026 13:26:54 -0400 Subject: [PATCH 14/29] use h200 --- .../platforms/{flatiron_b200.yaml => flatiron_h200.yaml} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename configs/training/platforms/{flatiron_b200.yaml => flatiron_h200.yaml} (85%) diff --git a/configs/training/platforms/flatiron_b200.yaml b/configs/training/platforms/flatiron_h200.yaml similarity index 85% rename from configs/training/platforms/flatiron_b200.yaml rename to configs/training/platforms/flatiron_h200.yaml index 0fc04efd3..e10bb3f85 100644 --- a/configs/training/platforms/flatiron_b200.yaml +++ b/configs/training/platforms/flatiron_h200.yaml @@ -1,4 +1,4 @@ -name: flatiron_b200 +name: flatiron_h200 gpus: 8 data_dir: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds experiments_dir: /mnt/home/${USER}/particleflow/experiments @@ -8,8 +8,8 @@ runtime_overrides: prefetch_factor: 2 model.attention.use_flash_attn_varlen: false slurm: - partition: gpu - constraint: b200 + partition: gpuxl + constraint: h200 time: "12:00:00" nodes: 1 tasks_per_node: 1 From 1b40de0a1987ec5a677b769117ddea11a2722819 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sun, 6 Sep 2026 08:52:35 +0300 Subject: [PATCH 15/29] Add hit backbone comparison scenario --- .../cld_hits_backbone_comparison.yaml | 46 +++++++++++++++++++ mlpf/training_scenarios.py | 2 + tests/test_training_scenarios.py | 21 ++++++++- tests/test_training_submission.py | 3 +- 4 files changed, 70 insertions(+), 2 deletions(-) create mode 100644 configs/training/scenarios/cld_hits_backbone_comparison.yaml diff --git a/configs/training/scenarios/cld_hits_backbone_comparison.yaml b/configs/training/scenarios/cld_hits_backbone_comparison.yaml new file mode 100644 index 000000000..09221f770 --- /dev/null +++ b/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 diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index cafc5af06..d849e0638 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -175,6 +175,8 @@ def _temporary_environment(values): def _serialize_cli_value(value): + if value is None: + return "null" if isinstance(value, bool): return "true" if value else "false" if isinstance(value, list): diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index 9dfd48e41..530a133ff 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -17,6 +17,7 @@ ROOT = Path(__file__).resolve().parents[1] SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" +BACKBONE_SCENARIO = ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" PLATFORMS = ROOT / "configs/training/platforms" @@ -41,12 +42,30 @@ def test_comparison_scenario_resolves_both_output_modes_with_same_seed(): assert all(job.resolved_config.seed == 12345 for job in jobs) +def test_backbone_comparison_scenario_keeps_elementwise_output_and_depth_fixed(): + scenario = load_training_scenario(BACKBONE_SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + + assert [job.variant_name for job in jobs] == ["attention", "heptv2"] + assert [job.resolved_config.model.type.value for job in jobs] == ["attention", "heptv2"] + assert {job.resolved_config.model.output_mode.value for job in jobs} == {"elementwise"} + assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} + assert jobs[1].resolved_config.model.heptv2.block_size == 128 + + @pytest.mark.parametrize( ("profile_name", "expected_multiplier"), [ ("flatiron_h100.yaml", 64), ("flatiron_a100.yaml", 128), - ("flatiron_b200.yaml", 64), + ("flatiron_h200.yaml", 64), ("tallinn_l40.yaml", 256), ("lumi_mi250x.yaml", 64), ], diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index 2b6a759a4..e27c7f25d 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -16,7 +16,8 @@ def test_picker_discovers_scenarios_and_accelerators(): scenarios, accelerators = available_choices(ROOT) assert "cld_hits_output_comparison" in scenarios - assert {"a100", "b200", "h100"}.issubset(accelerators) + assert "cld_hits_backbone_comparison" in scenarios + assert {"a100", "h100", "h200"}.issubset(accelerators) def test_h100_submission_is_derived_from_scenario_and_profile(): From 2ba46293986f40e0b26c63df8694484e6054e24e Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Sun, 6 Sep 2026 09:30:06 +0300 Subject: [PATCH 16/29] Fix unused HEPTv2 parameters in DDP --- mlpf/model/heptv2.py | 6 +++++- mlpf/standalone/train.py | 6 +++++- tests/test_hept_layers.py | 20 ++++++++++++++++++++ 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/mlpf/model/heptv2.py b/mlpf/model/heptv2.py index c0a36293e..f08d7eeab 100644 --- a/mlpf/model/heptv2.py +++ b/mlpf/model/heptv2.py @@ -455,7 +455,12 @@ def __init__(self, name=None, embedding_dim=128, num_heads=16, width=512, dropou self.w_k = nn.Linear(embedding_dim, embedding_dim, bias=False) self.w_v = nn.Linear(embedding_dim, embedding_dim, bias=False) + # HEPTv2 hashes detector coordinates directly and no longer uses the + # learned relative-position projection from HEPT v1. Keep the frozen + # tensors in the state dict for compatibility with existing HEPTv2 + # checkpoints and optimizer parameter groups. self.w_rpe = nn.Linear(self.num_w_per_dist * (self.coords_dim - 1), self.num_heads * self.dim_per_head) + self.w_rpe.requires_grad_(False) self.pe_func = PELearned(input_channel=self.coords_dim, h_dim=embedding_dim) if pe_type == "learned" else None @@ -552,7 +557,6 @@ def forward(self, x, mask, X_features, return_attn=False): k, v, coords=coords, - w_rpe=self.w_rpe, regions_h=regions_h, region_indices=region_indices, raw_size=raw_size, diff --git a/mlpf/standalone/train.py b/mlpf/standalone/train.py index 9fa4ea45a..497fceac5 100644 --- a/mlpf/standalone/train.py +++ b/mlpf/standalone/train.py @@ -452,7 +452,12 @@ def __init__(self, name=None, embedding_dim=128, num_heads=16, width=512, dropou self.w_k = nn.Linear(embedding_dim, embedding_dim, bias=False) self.w_v = nn.Linear(embedding_dim, embedding_dim, bias=False) + # HEPTv2 hashes detector coordinates directly and no longer uses the + # learned relative-position projection from HEPT v1. Keep the frozen + # tensors in the state dict for compatibility with existing HEPTv2 + # checkpoints and optimizer parameter groups. self.w_rpe = nn.Linear(self.num_w_per_dist * (self.coords_dim - 1), self.num_heads * self.dim_per_head) + self.w_rpe.requires_grad_(False) self.pe_func = PELearned(input_channel=self.coords_dim, h_dim=embedding_dim) if pe_type == "learned" else None @@ -549,7 +554,6 @@ def forward(self, x, mask, X_features, return_attn=False): k, v, coords=coords, - w_rpe=self.w_rpe, regions_h=regions_h, region_indices=region_indices, raw_size=raw_size, diff --git a/tests/test_hept_layers.py b/tests/test_hept_layers.py index e267777a3..69e1b4685 100644 --- a/tests/test_hept_layers.py +++ b/tests/test_hept_layers.py @@ -88,6 +88,26 @@ def test_hept_layer_backward_has_finite_gradients(layer_cls): assert all(torch.isfinite(grad).all() for grad in grads) +def test_heptv2_layer_backward_reaches_all_trainable_parameters(): + torch.manual_seed(2) + layer = _make_layer(HEPTv2Layer) + layer.train() + + batch_size, seq_len, embedding_dim = 2, 16, 32 + x = torch.randn(batch_size, seq_len, embedding_dim, requires_grad=True) + mask = torch.ones(batch_size, seq_len, dtype=torch.bool) + mask[0, 10:] = False + mask[1, 14:] = False + features = _make_x_features(batch_size, seq_len) + + layer(x, mask, features)[mask].pow(2).mean().backward() + + missing_gradients = [name for name, param in layer.named_parameters() if param.requires_grad and param.grad is None] + assert missing_gradients == [] + assert not layer.w_rpe.weight.requires_grad + assert not layer.w_rpe.bias.requires_grad + + @pytest.mark.parametrize("layer_cls", [HEPTLayer, HEPTv2Layer]) def test_hept_layer_mask_edge_cases_keep_at_least_one_valid_token(layer_cls): torch.manual_seed(3) From 52f6846f1d0301b2ed5c54dc781cfe3f6ff2a955 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Mon, 7 Sep 2026 15:47:47 +0300 Subject: [PATCH 17/29] Add particle-origin query alignment loss --- .../scenarios/cld_hits_output_comparison.yaml | 7 + mlpf/conf.py | 2 + mlpf/model/set_decoder.py | 11 ++ mlpf/model/set_losses.py | 150 ++++++++++++++++++ mlpf/model/training.py | 4 + tests/test_set_prediction.py | 73 ++++++++- 6 files changed, 246 insertions(+), 1 deletion(-) diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index 667911669..029b28957 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -21,6 +21,13 @@ variants: model.set_decoder.no_object_weight: 1.0 model.set_decoder.cardinality_loss_weight: 0.05 model.set_decoder.auxiliary_loss_weight: 0.25 + # Align each matched query with the pooled encoder hits carrying the same + # event-local particle_number. The symmetric direction makes all other + # slots negatives, discouraging multiple queries for one particle. + # A random 55-particle/256-query assignment has an unweighted loss near + # five, so 0.02 starts this term at roughly 0.1 in the total objective. + model.set_decoder.query_origin_loss_weight: 0.02 + model.set_decoder.query_origin_temperature: 0.1 # 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. diff --git a/mlpf/conf.py b/mlpf/conf.py index cef13e7e9..8121d1e19 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -628,6 +628,8 @@ class SetDecoderConfig(BaseModel): no_object_weight: float = Field(default=1.0, gt=0.0) cardinality_loss_weight: float = Field(default=0.0, ge=0.0) auxiliary_loss_weight: float = Field(default=0.0, ge=0.0) + query_origin_loss_weight: float = Field(default=0.0, ge=0.0) + query_origin_temperature: float = Field(default=0.1, gt=0.0) matcher: SetMatcherConfig = Field(default_factory=SetMatcherConfig) @model_validator(mode="after") diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py index 2d9fce902..27527705a 100644 --- a/mlpf/model/set_decoder.py +++ b/mlpf/model/set_decoder.py @@ -96,6 +96,7 @@ def __init__(self, embedding_dim, num_classes, config): self.local_attention_radius = config.local_attention_radius self.tracker_query_fraction = config.tracker_query_fraction self.use_auxiliary_losses = config.auxiliary_loss_weight > 0 + self.use_query_origin_loss = config.query_origin_loss_weight > 0 self.queries = nn.Parameter(torch.empty(1, config.num_slots, embedding_dim)) nn.init.trunc_normal_(self.queries, std=0.02) ffn_dim = int(config.ffn_multiplier * embedding_dim) @@ -126,6 +127,8 @@ def __init__(self, embedding_dim, num_classes, config): # Populated on every forward pass. The main four-tensor return signature # remains unchanged for inference and elementwise compatibility. self.auxiliary_outputs = [] + self.query_origin_query_embeddings = None + self.query_origin_memory_embeddings = None @staticmethod def _take_topk(scores, candidates, count): @@ -214,6 +217,8 @@ def _predict(self, slots, references=None): return presence, pid, momentum, pileup def forward(self, memory, memory_mask, input_features=None): + self.query_origin_query_embeddings = None + self.query_origin_memory_embeddings = None memory_mask = memory_mask.bool() references = reference_mask = memory_positions = None if self.query_init == "input-conditioned": @@ -256,4 +261,10 @@ def forward(self, memory, memory_mask, input_features=None): outputs.append(self._predict(slots, references)) self.auxiliary_outputs = outputs[:-1] if self.use_auxiliary_losses else [] + if self.use_query_origin_loss: + # Keep references to the existing activations rather than materializing + # a query-by-hit ownership tensor. The loss pools hits by truth particle + # in O(num_hits * embedding_dim) time and memory. + self.query_origin_query_embeddings = slots + self.query_origin_memory_embeddings = memory return outputs[-1] diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py index b5d14270e..26eb0ee86 100644 --- a/mlpf/model/set_losses.py +++ b/mlpf/model/set_losses.py @@ -6,9 +6,13 @@ from torch.nn import functional as F from mlpf.logger import _logger +from mlpf.conf import Y_FEATURES from mlpf.model.losses import LOSS_TASKS, REGRESSION_FEATURES +PARTICLE_NUMBER_INDEX = Y_FEATURES.index("particle_number") + + @dataclass(frozen=True) class SetMatcherWeights: presence: float = 1.0 @@ -76,6 +80,130 @@ def hungarian_match(targets, predictions, target_mask, weights=None): return matches +def query_origin_contrastive_loss( + query_embeddings, + memory_embeddings, + input_particle_numbers, + target_particle_numbers, + input_mask, + target_mask, + matches, + temperature=0.1, +): + """Align queries to truth-linked hit groups with a symmetric contrastive loss. + + ``particle_number`` is used only as an event-local grouping label. Hit + embeddings are accumulated directly into target-particle prototypes, avoiding + the O(num_queries * num_hits) ownership tensor that a dense mask loss would + require. The particle-to-query direction includes every query as a negative, + so duplicate queries aligned to the same particle are explicitly penalized. + """ + + if temperature <= 0: + raise ValueError("query-origin temperature must be positive") + if query_embeddings is None or memory_embeddings is None: + raise ValueError("query-origin loss requires decoder query and memory embeddings") + + zero = query_embeddings.reshape(-1)[0].float() * 0.0 + memory_embeddings.reshape(-1)[0].float() * 0.0 + # Perform reductions and similarities in FP32 even under BF16 autocast. + with torch.autocast(device_type=query_embeddings.device.type, enabled=False): + batch_size, num_inputs, embedding_dim = memory_embeddings.shape + num_targets = target_particle_numbers.shape[1] + if batch_size == 0 or num_inputs == 0 or num_targets == 0: + return zero + + hit_numbers = input_particle_numbers.long() + target_numbers = target_particle_numbers.long() + valid_hits = input_mask.bool() & (hit_numbers > 0) + valid_numbered_targets = target_mask.bool() & (target_numbers > 0) + if not valid_hits.any() or not valid_numbered_targets.any(): + return zero + + # particle_number is event-local. A composite key lets one searchsorted + # map every hit in the batch to its compact target row without a Python + # event loop or a dense query-by-hit ownership tensor. + max_particle_number = torch.maximum( + hit_numbers.masked_fill(~valid_hits, 0).max(), + target_numbers.masked_fill(~valid_numbered_targets, 0).max(), + ) + key_stride = max_particle_number + 1 + event_offsets = torch.arange(batch_size, device=query_embeddings.device, dtype=torch.long) * key_stride + hit_keys = (hit_numbers + event_offsets[:, None])[valid_hits] + target_keys = (target_numbers + event_offsets[:, None])[valid_numbered_targets] + target_flat_indices = torch.arange( + batch_size * num_targets, device=query_embeddings.device, dtype=torch.long + ).reshape(batch_size, num_targets)[valid_numbered_targets] + + sorted_target_keys, target_key_order = torch.sort(target_keys) + sorted_target_flat_indices = target_flat_indices[target_key_order] + positions = torch.searchsorted(sorted_target_keys, hit_keys) + in_range = positions < len(sorted_target_keys) + safe_positions = positions.clamp_max(len(sorted_target_keys) - 1) + associated = in_range & (sorted_target_keys[safe_positions] == hit_keys) + if not associated.any(): + return zero + + group_flat_indices = sorted_target_flat_indices[safe_positions[associated]] + # Select associated hits before promoting BF16 activations to FP32 so a + # padded batch never acquires a full-size FP32 memory copy. + valid_memory = memory_embeddings[valid_hits][associated].float() + prototype_sums = valid_memory.new_zeros((batch_size * num_targets, embedding_dim)) + prototype_sums.index_add_(0, group_flat_indices, valid_memory) + prototype_counts = valid_memory.new_zeros(batch_size * num_targets) + prototype_counts.index_add_( + 0, + group_flat_indices, + torch.ones_like(group_flat_indices, dtype=valid_memory.dtype), + ) + prototype_sums = prototype_sums.reshape(batch_size, num_targets, embedding_dim) + prototype_counts = prototype_counts.reshape(batch_size, num_targets) + valid_prototypes = prototype_counts > 0 + prototypes = prototype_sums / prototype_counts.clamp_min(1.0)[..., None] + prototypes = F.normalize(prototypes, dim=-1, eps=1.0e-6) + queries = F.normalize(query_embeddings.float(), dim=-1, eps=1.0e-6) + similarities = torch.bmm(queries, prototypes.transpose(1, 2)) / temperature + similarities = similarities.masked_fill(~valid_prototypes[:, None, :], torch.finfo(similarities.dtype).min) + + pair_batches = [] + pair_slots = [] + pair_targets = [] + for event_idx, (slot_indices, target_indices) in enumerate(matches): + if len(slot_indices): + # Hungarian target indices address the compact valid-target view, + # whereas the batched prototype tensor retains padded positions. + valid_target_positions = torch.nonzero(target_mask[event_idx], as_tuple=False).squeeze(1) + pair_batches.append(torch.full_like(slot_indices, event_idx)) + pair_slots.append(slot_indices) + pair_targets.append(valid_target_positions[target_indices]) + if not pair_slots: + return zero + + pair_batches = torch.cat(pair_batches) + pair_slots = torch.cat(pair_slots) + pair_targets = torch.cat(pair_targets) + has_prototype = valid_prototypes[pair_batches, pair_targets] + pair_batches = pair_batches[has_prototype] + pair_slots = pair_slots[has_prototype] + pair_targets = pair_targets[has_prototype] + if len(pair_slots) == 0: + return zero + + # Query -> particle aligns each matched query with its originating hit + # group. Particle -> query makes that group select exactly one query; + # all unmatched and duplicate queries participate as negatives. + query_to_particle = F.cross_entropy( + similarities[pair_batches, pair_slots], + pair_targets, + reduction="sum", + ) + particle_to_query = F.cross_entropy( + similarities[pair_batches, :, pair_targets], + pair_slots, + reduction="sum", + ) + return (query_to_particle + particle_to_query) / (2 * len(pair_slots)) + + def set_event_loss( targets, predictions, @@ -153,6 +281,10 @@ def set_mlpf_loss( cardinality_loss_weight=0.0, auxiliary_predictions=None, auxiliary_loss_weight=0.0, + query_embeddings=None, + memory_embeddings=None, + query_origin_loss_weight=0.0, + query_origin_temperature=0.1, ): """Compute the set-prediction objective with the standard task names.""" @@ -181,6 +313,24 @@ def set_mlpf_loss( if "Cardinality" in losses: loss_opt = loss_opt + losses["Cardinality"] + if query_origin_loss_weight > 0: + if batch.ytarget is None: + raise ValueError("query-origin loss requires per-hit ytarget particle_number labels") + if "particle_number" not in targets: + raise ValueError("query-origin loss requires particle_number in set targets") + origin_loss = query_origin_contrastive_loss( + query_embeddings, + memory_embeddings, + batch.ytarget[..., PARTICLE_NUMBER_INDEX], + targets["particle_number"], + batch.mask, + batch.target_mask, + matches, + temperature=query_origin_temperature, + ) + losses["Query_origin"] = query_origin_loss_weight * origin_loss + loss_opt = loss_opt + losses["Query_origin"] + if auxiliary_predictions and auxiliary_loss_weight > 0: auxiliary_losses = [] for auxiliary_prediction in auxiliary_predictions: diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 6f88d8053..5c2613b72 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -227,6 +227,10 @@ def _set_loss_kwargs(model_module): "cardinality_loss_weight": config.cardinality_loss_weight, "auxiliary_predictions": auxiliary_predictions, "auxiliary_loss_weight": config.auxiliary_loss_weight, + "query_embeddings": model_module.set_decoder.query_origin_query_embeddings, + "memory_embeddings": model_module.set_decoder.query_origin_memory_embeddings, + "query_origin_loss_weight": config.query_origin_loss_weight, + "query_origin_temperature": config.query_origin_temperature, } diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py index 732412b44..e1cddf487 100644 --- a/tests/test_set_prediction.py +++ b/tests/test_set_prediction.py @@ -4,7 +4,7 @@ from mlpf.conf import MLPFConfig from mlpf.model.PFDataset import PFBatch from mlpf.model.mlpf import MLPF -from mlpf.model.set_losses import hungarian_match, set_event_loss +from mlpf.model.set_losses import hungarian_match, query_origin_contrastive_loss, set_event_loss from mlpf.model.utils import unpack_predictions, unpack_target @@ -94,6 +94,8 @@ def test_set_config_populates_decoder_defaults(): assert config.model.set_decoder is not None assert config.model.set_decoder.num_slots == 256 assert config.model.set_decoder.no_object_weight == 1.0 + assert config.model.set_decoder.query_origin_loss_weight == 0.0 + assert config.model.set_decoder.query_origin_temperature == 0.1 assert config.model.set_decoder.matcher.dr_scale == 0.1 @@ -187,6 +189,75 @@ def test_model_step_applies_configured_cardinality_and_auxiliary_losses(): assert model.set_decoder.reference_delta_heads[0].weight.grad is not None +def test_query_origin_loss_is_number_invariant_and_penalizes_duplicate_queries(): + memory = torch.tensor([[[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]], requires_grad=True) + aligned_queries = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]]], requires_grad=True) + duplicate_queries = aligned_queries.detach().clone() + duplicate_queries[0, 2] = torch.tensor([1.0, 0.0]) + hit_numbers = torch.tensor([[11, 11, 29, 29]]) + target_numbers = torch.tensor([[11, 0, 29]]) + input_mask = torch.ones(1, 4, dtype=torch.bool) + target_mask = torch.tensor([[True, False, True]]) + matches = [(torch.tensor([0, 1]), torch.tensor([0, 1]))] + + aligned_loss = query_origin_contrastive_loss( + aligned_queries, + memory, + hit_numbers, + target_numbers, + input_mask, + target_mask, + matches, + ) + duplicate_loss = query_origin_contrastive_loss( + duplicate_queries, + memory, + hit_numbers, + target_numbers, + input_mask, + target_mask, + matches, + ) + renumbered_loss = query_origin_contrastive_loss( + aligned_queries, + memory, + torch.tensor([[103, 103, 7, 7]]), + torch.tensor([[103, 0, 7]]), + input_mask, + target_mask, + matches, + ) + + assert aligned_loss < duplicate_loss + torch.testing.assert_close(aligned_loss, renumbered_loss) + aligned_loss.backward() + assert torch.isfinite(aligned_queries.grad).all() + assert torch.isfinite(memory.grad).all() + + +def test_model_step_applies_query_origin_loss(): + from mlpf.model.training import model_step + + config = make_config(num_slots=4, query_origin_loss_weight=0.1) + model = MLPF(config) + X = torch.randn(1, 8, config.input_dim) + X[..., 0] = 1 + X[..., 1] = X[..., 1].abs() + 0.1 + X[..., 5] = X[..., 5].abs() + 0.1 + ytarget = torch.zeros(1, 8, 14) + ytarget[0, :4, 13] = 1 + ytarget[0, 4:, 13] = 2 + batch = PFBatch(X=X, ytarget=ytarget, ytarget_set=make_target_tensor(num_targets=2)) + + loss, losses, _, _, _, _ = model_step(batch, model, None, REGRESSION_WEIGHTS) + loss.backward() + + assert torch.isfinite(loss) + assert losses["Query_origin"] > 0 + assert model.set_decoder.queries.grad is not None + assert torch.isfinite(model.set_decoder.queries.grad).all() + + def test_attention_set_model_has_no_unused_elementwise_parameters(): config = make_attention_config() model = MLPF(config) From ec4f81c4bff0aefab5323997ed156f84df5d03b0 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Tue, 8 Sep 2026 18:54:34 +0300 Subject: [PATCH 18/29] Revert particle-origin query alignment loss --- .../scenarios/cld_hits_output_comparison.yaml | 7 - mlpf/conf.py | 2 - mlpf/model/set_decoder.py | 11 -- mlpf/model/set_losses.py | 150 ------------------ mlpf/model/training.py | 4 - tests/test_set_prediction.py | 73 +-------- 6 files changed, 1 insertion(+), 246 deletions(-) diff --git a/configs/training/scenarios/cld_hits_output_comparison.yaml b/configs/training/scenarios/cld_hits_output_comparison.yaml index 029b28957..667911669 100644 --- a/configs/training/scenarios/cld_hits_output_comparison.yaml +++ b/configs/training/scenarios/cld_hits_output_comparison.yaml @@ -21,13 +21,6 @@ variants: model.set_decoder.no_object_weight: 1.0 model.set_decoder.cardinality_loss_weight: 0.05 model.set_decoder.auxiliary_loss_weight: 0.25 - # Align each matched query with the pooled encoder hits carrying the same - # event-local particle_number. The symmetric direction makes all other - # slots negatives, discouraging multiple queries for one particle. - # A random 55-particle/256-query assignment has an unweighted loss near - # five, so 0.02 starts this term at roughly 0.1 in the total objective. - model.set_decoder.query_origin_loss_weight: 0.02 - model.set_decoder.query_origin_temperature: 0.1 # 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. diff --git a/mlpf/conf.py b/mlpf/conf.py index 8121d1e19..cef13e7e9 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -628,8 +628,6 @@ class SetDecoderConfig(BaseModel): no_object_weight: float = Field(default=1.0, gt=0.0) cardinality_loss_weight: float = Field(default=0.0, ge=0.0) auxiliary_loss_weight: float = Field(default=0.0, ge=0.0) - query_origin_loss_weight: float = Field(default=0.0, ge=0.0) - query_origin_temperature: float = Field(default=0.1, gt=0.0) matcher: SetMatcherConfig = Field(default_factory=SetMatcherConfig) @model_validator(mode="after") diff --git a/mlpf/model/set_decoder.py b/mlpf/model/set_decoder.py index 27527705a..2d9fce902 100644 --- a/mlpf/model/set_decoder.py +++ b/mlpf/model/set_decoder.py @@ -96,7 +96,6 @@ def __init__(self, embedding_dim, num_classes, config): self.local_attention_radius = config.local_attention_radius self.tracker_query_fraction = config.tracker_query_fraction self.use_auxiliary_losses = config.auxiliary_loss_weight > 0 - self.use_query_origin_loss = config.query_origin_loss_weight > 0 self.queries = nn.Parameter(torch.empty(1, config.num_slots, embedding_dim)) nn.init.trunc_normal_(self.queries, std=0.02) ffn_dim = int(config.ffn_multiplier * embedding_dim) @@ -127,8 +126,6 @@ def __init__(self, embedding_dim, num_classes, config): # Populated on every forward pass. The main four-tensor return signature # remains unchanged for inference and elementwise compatibility. self.auxiliary_outputs = [] - self.query_origin_query_embeddings = None - self.query_origin_memory_embeddings = None @staticmethod def _take_topk(scores, candidates, count): @@ -217,8 +214,6 @@ def _predict(self, slots, references=None): return presence, pid, momentum, pileup def forward(self, memory, memory_mask, input_features=None): - self.query_origin_query_embeddings = None - self.query_origin_memory_embeddings = None memory_mask = memory_mask.bool() references = reference_mask = memory_positions = None if self.query_init == "input-conditioned": @@ -261,10 +256,4 @@ def forward(self, memory, memory_mask, input_features=None): outputs.append(self._predict(slots, references)) self.auxiliary_outputs = outputs[:-1] if self.use_auxiliary_losses else [] - if self.use_query_origin_loss: - # Keep references to the existing activations rather than materializing - # a query-by-hit ownership tensor. The loss pools hits by truth particle - # in O(num_hits * embedding_dim) time and memory. - self.query_origin_query_embeddings = slots - self.query_origin_memory_embeddings = memory return outputs[-1] diff --git a/mlpf/model/set_losses.py b/mlpf/model/set_losses.py index 26eb0ee86..b5d14270e 100644 --- a/mlpf/model/set_losses.py +++ b/mlpf/model/set_losses.py @@ -6,13 +6,9 @@ from torch.nn import functional as F from mlpf.logger import _logger -from mlpf.conf import Y_FEATURES from mlpf.model.losses import LOSS_TASKS, REGRESSION_FEATURES -PARTICLE_NUMBER_INDEX = Y_FEATURES.index("particle_number") - - @dataclass(frozen=True) class SetMatcherWeights: presence: float = 1.0 @@ -80,130 +76,6 @@ def hungarian_match(targets, predictions, target_mask, weights=None): return matches -def query_origin_contrastive_loss( - query_embeddings, - memory_embeddings, - input_particle_numbers, - target_particle_numbers, - input_mask, - target_mask, - matches, - temperature=0.1, -): - """Align queries to truth-linked hit groups with a symmetric contrastive loss. - - ``particle_number`` is used only as an event-local grouping label. Hit - embeddings are accumulated directly into target-particle prototypes, avoiding - the O(num_queries * num_hits) ownership tensor that a dense mask loss would - require. The particle-to-query direction includes every query as a negative, - so duplicate queries aligned to the same particle are explicitly penalized. - """ - - if temperature <= 0: - raise ValueError("query-origin temperature must be positive") - if query_embeddings is None or memory_embeddings is None: - raise ValueError("query-origin loss requires decoder query and memory embeddings") - - zero = query_embeddings.reshape(-1)[0].float() * 0.0 + memory_embeddings.reshape(-1)[0].float() * 0.0 - # Perform reductions and similarities in FP32 even under BF16 autocast. - with torch.autocast(device_type=query_embeddings.device.type, enabled=False): - batch_size, num_inputs, embedding_dim = memory_embeddings.shape - num_targets = target_particle_numbers.shape[1] - if batch_size == 0 or num_inputs == 0 or num_targets == 0: - return zero - - hit_numbers = input_particle_numbers.long() - target_numbers = target_particle_numbers.long() - valid_hits = input_mask.bool() & (hit_numbers > 0) - valid_numbered_targets = target_mask.bool() & (target_numbers > 0) - if not valid_hits.any() or not valid_numbered_targets.any(): - return zero - - # particle_number is event-local. A composite key lets one searchsorted - # map every hit in the batch to its compact target row without a Python - # event loop or a dense query-by-hit ownership tensor. - max_particle_number = torch.maximum( - hit_numbers.masked_fill(~valid_hits, 0).max(), - target_numbers.masked_fill(~valid_numbered_targets, 0).max(), - ) - key_stride = max_particle_number + 1 - event_offsets = torch.arange(batch_size, device=query_embeddings.device, dtype=torch.long) * key_stride - hit_keys = (hit_numbers + event_offsets[:, None])[valid_hits] - target_keys = (target_numbers + event_offsets[:, None])[valid_numbered_targets] - target_flat_indices = torch.arange( - batch_size * num_targets, device=query_embeddings.device, dtype=torch.long - ).reshape(batch_size, num_targets)[valid_numbered_targets] - - sorted_target_keys, target_key_order = torch.sort(target_keys) - sorted_target_flat_indices = target_flat_indices[target_key_order] - positions = torch.searchsorted(sorted_target_keys, hit_keys) - in_range = positions < len(sorted_target_keys) - safe_positions = positions.clamp_max(len(sorted_target_keys) - 1) - associated = in_range & (sorted_target_keys[safe_positions] == hit_keys) - if not associated.any(): - return zero - - group_flat_indices = sorted_target_flat_indices[safe_positions[associated]] - # Select associated hits before promoting BF16 activations to FP32 so a - # padded batch never acquires a full-size FP32 memory copy. - valid_memory = memory_embeddings[valid_hits][associated].float() - prototype_sums = valid_memory.new_zeros((batch_size * num_targets, embedding_dim)) - prototype_sums.index_add_(0, group_flat_indices, valid_memory) - prototype_counts = valid_memory.new_zeros(batch_size * num_targets) - prototype_counts.index_add_( - 0, - group_flat_indices, - torch.ones_like(group_flat_indices, dtype=valid_memory.dtype), - ) - prototype_sums = prototype_sums.reshape(batch_size, num_targets, embedding_dim) - prototype_counts = prototype_counts.reshape(batch_size, num_targets) - valid_prototypes = prototype_counts > 0 - prototypes = prototype_sums / prototype_counts.clamp_min(1.0)[..., None] - prototypes = F.normalize(prototypes, dim=-1, eps=1.0e-6) - queries = F.normalize(query_embeddings.float(), dim=-1, eps=1.0e-6) - similarities = torch.bmm(queries, prototypes.transpose(1, 2)) / temperature - similarities = similarities.masked_fill(~valid_prototypes[:, None, :], torch.finfo(similarities.dtype).min) - - pair_batches = [] - pair_slots = [] - pair_targets = [] - for event_idx, (slot_indices, target_indices) in enumerate(matches): - if len(slot_indices): - # Hungarian target indices address the compact valid-target view, - # whereas the batched prototype tensor retains padded positions. - valid_target_positions = torch.nonzero(target_mask[event_idx], as_tuple=False).squeeze(1) - pair_batches.append(torch.full_like(slot_indices, event_idx)) - pair_slots.append(slot_indices) - pair_targets.append(valid_target_positions[target_indices]) - if not pair_slots: - return zero - - pair_batches = torch.cat(pair_batches) - pair_slots = torch.cat(pair_slots) - pair_targets = torch.cat(pair_targets) - has_prototype = valid_prototypes[pair_batches, pair_targets] - pair_batches = pair_batches[has_prototype] - pair_slots = pair_slots[has_prototype] - pair_targets = pair_targets[has_prototype] - if len(pair_slots) == 0: - return zero - - # Query -> particle aligns each matched query with its originating hit - # group. Particle -> query makes that group select exactly one query; - # all unmatched and duplicate queries participate as negatives. - query_to_particle = F.cross_entropy( - similarities[pair_batches, pair_slots], - pair_targets, - reduction="sum", - ) - particle_to_query = F.cross_entropy( - similarities[pair_batches, :, pair_targets], - pair_slots, - reduction="sum", - ) - return (query_to_particle + particle_to_query) / (2 * len(pair_slots)) - - def set_event_loss( targets, predictions, @@ -281,10 +153,6 @@ def set_mlpf_loss( cardinality_loss_weight=0.0, auxiliary_predictions=None, auxiliary_loss_weight=0.0, - query_embeddings=None, - memory_embeddings=None, - query_origin_loss_weight=0.0, - query_origin_temperature=0.1, ): """Compute the set-prediction objective with the standard task names.""" @@ -313,24 +181,6 @@ def set_mlpf_loss( if "Cardinality" in losses: loss_opt = loss_opt + losses["Cardinality"] - if query_origin_loss_weight > 0: - if batch.ytarget is None: - raise ValueError("query-origin loss requires per-hit ytarget particle_number labels") - if "particle_number" not in targets: - raise ValueError("query-origin loss requires particle_number in set targets") - origin_loss = query_origin_contrastive_loss( - query_embeddings, - memory_embeddings, - batch.ytarget[..., PARTICLE_NUMBER_INDEX], - targets["particle_number"], - batch.mask, - batch.target_mask, - matches, - temperature=query_origin_temperature, - ) - losses["Query_origin"] = query_origin_loss_weight * origin_loss - loss_opt = loss_opt + losses["Query_origin"] - if auxiliary_predictions and auxiliary_loss_weight > 0: auxiliary_losses = [] for auxiliary_prediction in auxiliary_predictions: diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 5c2613b72..6f88d8053 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -227,10 +227,6 @@ def _set_loss_kwargs(model_module): "cardinality_loss_weight": config.cardinality_loss_weight, "auxiliary_predictions": auxiliary_predictions, "auxiliary_loss_weight": config.auxiliary_loss_weight, - "query_embeddings": model_module.set_decoder.query_origin_query_embeddings, - "memory_embeddings": model_module.set_decoder.query_origin_memory_embeddings, - "query_origin_loss_weight": config.query_origin_loss_weight, - "query_origin_temperature": config.query_origin_temperature, } diff --git a/tests/test_set_prediction.py b/tests/test_set_prediction.py index e1cddf487..732412b44 100644 --- a/tests/test_set_prediction.py +++ b/tests/test_set_prediction.py @@ -4,7 +4,7 @@ from mlpf.conf import MLPFConfig from mlpf.model.PFDataset import PFBatch from mlpf.model.mlpf import MLPF -from mlpf.model.set_losses import hungarian_match, query_origin_contrastive_loss, set_event_loss +from mlpf.model.set_losses import hungarian_match, set_event_loss from mlpf.model.utils import unpack_predictions, unpack_target @@ -94,8 +94,6 @@ def test_set_config_populates_decoder_defaults(): assert config.model.set_decoder is not None assert config.model.set_decoder.num_slots == 256 assert config.model.set_decoder.no_object_weight == 1.0 - assert config.model.set_decoder.query_origin_loss_weight == 0.0 - assert config.model.set_decoder.query_origin_temperature == 0.1 assert config.model.set_decoder.matcher.dr_scale == 0.1 @@ -189,75 +187,6 @@ def test_model_step_applies_configured_cardinality_and_auxiliary_losses(): assert model.set_decoder.reference_delta_heads[0].weight.grad is not None -def test_query_origin_loss_is_number_invariant_and_penalizes_duplicate_queries(): - memory = torch.tensor([[[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]], requires_grad=True) - aligned_queries = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]]], requires_grad=True) - duplicate_queries = aligned_queries.detach().clone() - duplicate_queries[0, 2] = torch.tensor([1.0, 0.0]) - hit_numbers = torch.tensor([[11, 11, 29, 29]]) - target_numbers = torch.tensor([[11, 0, 29]]) - input_mask = torch.ones(1, 4, dtype=torch.bool) - target_mask = torch.tensor([[True, False, True]]) - matches = [(torch.tensor([0, 1]), torch.tensor([0, 1]))] - - aligned_loss = query_origin_contrastive_loss( - aligned_queries, - memory, - hit_numbers, - target_numbers, - input_mask, - target_mask, - matches, - ) - duplicate_loss = query_origin_contrastive_loss( - duplicate_queries, - memory, - hit_numbers, - target_numbers, - input_mask, - target_mask, - matches, - ) - renumbered_loss = query_origin_contrastive_loss( - aligned_queries, - memory, - torch.tensor([[103, 103, 7, 7]]), - torch.tensor([[103, 0, 7]]), - input_mask, - target_mask, - matches, - ) - - assert aligned_loss < duplicate_loss - torch.testing.assert_close(aligned_loss, renumbered_loss) - aligned_loss.backward() - assert torch.isfinite(aligned_queries.grad).all() - assert torch.isfinite(memory.grad).all() - - -def test_model_step_applies_query_origin_loss(): - from mlpf.model.training import model_step - - config = make_config(num_slots=4, query_origin_loss_weight=0.1) - model = MLPF(config) - X = torch.randn(1, 8, config.input_dim) - X[..., 0] = 1 - X[..., 1] = X[..., 1].abs() + 0.1 - X[..., 5] = X[..., 5].abs() + 0.1 - ytarget = torch.zeros(1, 8, 14) - ytarget[0, :4, 13] = 1 - ytarget[0, 4:, 13] = 2 - batch = PFBatch(X=X, ytarget=ytarget, ytarget_set=make_target_tensor(num_targets=2)) - - loss, losses, _, _, _, _ = model_step(batch, model, None, REGRESSION_WEIGHTS) - loss.backward() - - assert torch.isfinite(loss) - assert losses["Query_origin"] > 0 - assert model.set_decoder.queries.grad is not None - assert torch.isfinite(model.set_decoder.queries.grad).all() - - def test_attention_set_model_has_no_unused_elementwise_parameters(): config = make_attention_config() model = MLPF(config) From 125a5a5803e796e7d1ad1c672b14e872ab968ab6 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Tue, 8 Sep 2026 23:33:03 +0300 Subject: [PATCH 19/29] Improve hit-training evaluation and experiment setup --- configs/training/platforms/flatiron_h100.yaml | 2 +- configs/training/platforms/flatiron_h200.yaml | 2 +- .../scenarios/cld_pf_hits_comparison.yaml | 74 +++++++++++++ mlpf/conf.py | 8 ++ mlpf/jet_utils.py | 104 +++++++++++++----- mlpf/model/inference.py | 3 +- mlpf/model/losses.py | 27 +++-- mlpf/model/training.py | 52 +++++++-- mlpf/plotting/plot_utils.py | 70 ++++++------ particleflow_spec.yaml | 2 + tests/test_jet_utils.py | 87 +++++++++++++++ tests/test_standard_loss.py | 26 ++++- tests/test_training_scenarios.py | 25 +++++ tests/test_training_submission.py | 1 + 14 files changed, 405 insertions(+), 78 deletions(-) create mode 100644 configs/training/scenarios/cld_pf_hits_comparison.yaml create mode 100644 tests/test_jet_utils.py diff --git a/configs/training/platforms/flatiron_h100.yaml b/configs/training/platforms/flatiron_h100.yaml index bfaac20af..ca3357579 100644 --- a/configs/training/platforms/flatiron_h100.yaml +++ b/configs/training/platforms/flatiron_h100.yaml @@ -10,7 +10,7 @@ runtime_overrides: slurm: partition: gpu constraint: h100 - time: "12:00:00" + time: "48:00:00" nodes: 1 tasks_per_node: 1 cpus_per_task: 64 diff --git a/configs/training/platforms/flatiron_h200.yaml b/configs/training/platforms/flatiron_h200.yaml index e10bb3f85..873acf247 100644 --- a/configs/training/platforms/flatiron_h200.yaml +++ b/configs/training/platforms/flatiron_h200.yaml @@ -10,7 +10,7 @@ runtime_overrides: slurm: partition: gpuxl constraint: h200 - time: "12:00:00" + time: "48:00:00" nodes: 1 tasks_per_node: 1 cpus_per_task: 64 diff --git a/configs/training/scenarios/cld_pf_hits_comparison.yaml b/configs/training/scenarios/cld_pf_hits_comparison.yaml new file mode 100644 index 000000000..f85e45262 --- /dev/null +++ b/configs/training/scenarios/cld_pf_hits_comparison.yaml @@ -0,0 +1,74 @@ +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 + set_hits: + 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 + 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.backbone.num_tracker_layers: 2 + model.backbone.num_calo_layers: 2 + model.backbone.num_common_layers: 2 + 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 + # Only the set-based hit variant has a set decoder. + - model.output_mode + - model.set_decoder diff --git a/mlpf/conf.py b/mlpf/conf.py index cef13e7e9..5437ace99 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -426,6 +426,7 @@ def get_names(cls): "r": 0.4, "ptcut": 3.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, Dataset.CLIC.value: { "algo": "ee_genkt_algorithm", @@ -433,6 +434,7 @@ def get_names(cls): "p": -1.0, "ptcut": 5.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, Dataset.CLD.value: { "algo": "ee_genkt_algorithm", @@ -440,6 +442,7 @@ def get_names(cls): "p": -1.0, "ptcut": 5.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, Dataset.CLIC_HITS.value: { "algo": "ee_genkt_algorithm", @@ -447,6 +450,7 @@ def get_names(cls): "p": -1.0, "ptcut": 5.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, Dataset.CLD_HITS.value: { "algo": "ee_genkt_algorithm", @@ -454,6 +458,7 @@ def get_names(cls): "p": -1.0, "ptcut": 5.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, } @@ -651,6 +656,9 @@ class ModelArchitectureConfig(BaseModel): trainable: str = "all" task_queries: bool = True output_mode: OutputMode = OutputMode.ELEMENTWISE + # None keeps the standard binary cross-entropy. A non-negative value uses + # focal loss with this gamma for the elementwise particle-presence head. + binary_classification_focal_gamma: Optional[float] = Field(default=None, ge=0.0) backbone: Optional[BackboneConfig] = None hit_feature_engineering: HitFeatureEngineeringConfig = Field(default_factory=HitFeatureEngineeringConfig) set_decoder: Optional[SetDecoderConfig] = None diff --git a/mlpf/jet_utils.py b/mlpf/jet_utils.py index 22af783ee..293f83941 100644 --- a/mlpf/jet_utils.py +++ b/mlpf/jet_utils.py @@ -3,6 +3,7 @@ import numba import awkward import vector +from scipy.optimize import linear_sum_assignment @numba.njit @@ -18,34 +19,83 @@ def deltar(eta1, phi1, eta2, phi2): return np.sqrt(deta**2 + dphi**2) -@numba.njit def _match_jets_event(j1_eta, j1_phi, j2_eta, j2_phi, deltaR_cut): - jet_inds_1 = np.empty(len(j1_eta), dtype=np.int64) - jet_inds_2 = np.empty(len(j1_eta), dtype=np.int64) - num_matches = 0 - - # loop over the first jet collection - for ij1 in range(len(j1_eta)): - # compute deltaR from this jet to all jets in the other collection - min_idx_dr = -1 - min_dr = np.inf - - # loop over the other jet collection - for ij2 in range(len(j2_eta)): - # Workaround for https://github.com/scikit-hep/vector/issues/303 - # dr = j1[ij1].deltaR(j2[ij2]) - dr = deltar(j1_eta[ij1], j1_phi[ij1], j2_eta[ij2], j2_phi[ij2]) - if dr < min_dr: - min_idx_dr = ij2 - min_dr = dr - - # has to be closer than the deltaR_cut - if min_idx_dr >= 0 and min_dr < deltaR_cut: - jet_inds_1[num_matches] = ij1 - jet_inds_2[num_matches] = min_idx_dr - num_matches += 1 - - return jet_inds_1[:num_matches], jet_inds_2[:num_matches] + """Return a maximum-cardinality, minimum-deltaR one-to-one assignment.""" + + if deltaR_cut <= 0: + raise ValueError("deltaR_cut must be positive") + + num_jets_1 = len(j1_eta) + num_jets_2 = len(j2_eta) + if num_jets_1 == 0 or num_jets_2 == 0: + empty = np.empty(0, dtype=np.int64) + return empty, empty + + delta_eta = j1_eta[:, None] - j2_eta[None, :] + delta_phi = j1_phi[:, None] - j2_phi[None, :] + delta_phi = np.arctan2(np.sin(delta_phi), np.cos(delta_phi)) + delta_r = np.sqrt(delta_eta**2 + delta_phi**2) + valid = np.isfinite(delta_r) & (delta_r < deltaR_cut) + + # Give every jet in the first collection its own dummy unmatched column. + # The unmatched cost dominates the sum of all valid normalized distances, + # so the assignment first maximizes cardinality and then minimizes deltaR. + max_pairs = min(num_jets_1, num_jets_2) + unmatched_cost = float(max_pairs + 1) + invalid_cost = unmatched_cost * float(num_jets_1 + 1) + cost = np.full((num_jets_1, num_jets_2 + num_jets_1), invalid_cost, dtype=np.float64) + cost[:, :num_jets_2] = np.where(valid, delta_r / deltaR_cut, invalid_cost) + cost[np.arange(num_jets_1), num_jets_2 + np.arange(num_jets_1)] = unmatched_cost + + jet_inds_1, columns = linear_sum_assignment(cost) + real_match = columns < num_jets_2 + jet_inds_1 = jet_inds_1[real_match] + jet_inds_2 = columns[real_match] + accepted = valid[jet_inds_1, jet_inds_2] + return jet_inds_1[accepted], jet_inds_2[accepted] + + +def jet_matching_metrics(response_ratios, num_reference_jets, num_candidate_jets, response_rel_pt_cut=0.5): + """Summarize unique angular matches and response-qualified matches. + + ``response_ratios`` contains candidate/reference pT for the angularly + matched pairs returned by :func:`match_jets`. A response-qualified match + additionally satisfies ``abs(candidate/reference - 1) < response_rel_pt_cut``. + """ + + if response_rel_pt_cut <= 0: + raise ValueError("response_rel_pt_cut must be positive") + + ratios = np.asarray(response_ratios, dtype=np.float64).reshape(-1) + num_reference_jets = int(num_reference_jets) + num_candidate_jets = int(num_candidate_jets) + num_angular_matches = len(ratios) + if num_angular_matches > min(num_reference_jets, num_candidate_jets): + raise ValueError("one-to-one angular matches cannot exceed either jet collection") + + num_response_matches = int(np.sum(np.isfinite(ratios) & (np.abs(ratios - 1.0) < response_rel_pt_cut))) + + def fraction(numerator, denominator): + return float(numerator / denominator) if denominator else float("nan") + + metrics = { + "num_reference_jets": num_reference_jets, + "num_candidate_jets": num_candidate_jets, + "num_angular_matches": num_angular_matches, + "num_response_qualified_matches": num_response_matches, + "response_rel_pt_cut": float(response_rel_pt_cut), + "angular_recall": fraction(num_angular_matches, num_reference_jets), + "angular_precision": fraction(num_angular_matches, num_candidate_jets), + "angular_f1": fraction(2 * num_angular_matches, num_reference_jets + num_candidate_jets), + "angular_fake_rate": fraction(num_candidate_jets - num_angular_matches, num_candidate_jets), + "response_qualified_recall": fraction(num_response_matches, num_reference_jets), + "response_qualified_precision": fraction(num_response_matches, num_candidate_jets), + "response_qualified_f1": fraction(2 * num_response_matches, num_reference_jets + num_candidate_jets), + } + # Keep the historical key readable by existing dashboards. Its semantics + # are now the one-to-one angular recall rather than target-wise nearest-neighbor recall. + metrics["match_frac"] = metrics["angular_recall"] + return metrics def match_jets(jets1, jets2, deltaR_cut): diff --git a/mlpf/model/inference.py b/mlpf/model/inference.py index fe59667a4..0238d4f6d 100644 --- a/mlpf/model/inference.py +++ b/mlpf/model/inference.py @@ -32,7 +32,7 @@ from mlpf.logger import _logger from mlpf.model.utils import unpack_target -from mlpf.conf import OutputMode +from mlpf.conf import JET_CONFIG, OutputMode def predict_one_batch(conv_type, model, i, batch, rank, jetdef, jet_ptcut, jet_match_dr, outpath, dir_name, sample): @@ -222,6 +222,7 @@ def make_plots(outpath, sample, dataset, dir_name="", num_test_events=None, base sample=sample, dataset=ds_name, baseline_yvals=baseline_yvals, + response_rel_pt_cut=JET_CONFIG[ds_name]["match_rel_pt"], ) _logger.info("Plotted jet ratio") diff --git a/mlpf/model/losses.py b/mlpf/model/losses.py index 5d00854a9..134f1c00d 100644 --- a/mlpf/model/losses.py +++ b/mlpf/model/losses.py @@ -134,13 +134,16 @@ def sliced_wasserstein_loss(y_pred, y_true, num_projections=200): return ret -def classification_loss(y, ypred): +def classification_loss(y, ypred, binary_focal_gamma=None): """Compute per-element particle-presence and particle-ID losses.""" cls_id = y["cls_id"] num_elements = cls_id.numel() is_particle = cls_id != 0 - binary = 10.0 * F.cross_entropy(ypred["cls_binary"], is_particle.long()) + if binary_focal_gamma is None: + binary = 10.0 * F.cross_entropy(ypred["cls_binary"], is_particle.long()) + else: + binary = 10.0 * FocalLoss(gamma=binary_focal_gamma)(ypred["cls_binary"], is_particle.long()) pid_per_element = FocalLoss(gamma=2.0, reduction="none")(ypred["cls_id_onehot"], cls_id) pid_per_element = torch.where(is_particle, pid_per_element, torch.zeros_like(pid_per_element)) @@ -173,14 +176,14 @@ def regression_loss(y, ypred, input_pt, regression_weights): return losses -def particle_loss(y, ypred, input_pt, regression_weights): +def particle_loss(y, ypred, input_pt, regression_weights, *, binary_focal_gamma=None): """Compute classification and regression losses over flattened particles.""" - losses = classification_loss(y, ypred) + losses = classification_loss(y, ypred, binary_focal_gamma=binary_focal_gamma) losses.update(regression_loss(y, ypred, input_pt, regression_weights)) return losses -def event_loss(y, ypred, batch, regression_weights): +def event_loss(y, ypred, batch, regression_weights, *, binary_focal_gamma=None): """Compute losses for complete padded event batches. The standard loss currently contains only independent particle terms. @@ -200,18 +203,24 @@ def event_loss(y, ypred, batch, regression_weights): } input_pt = batch.X[..., 1][valid] - return particle_loss(particle_targets, particle_predictions, input_pt, regression_weights) + return particle_loss( + particle_targets, + particle_predictions, + input_pt, + regression_weights, + binary_focal_gamma=binary_focal_gamma, + ) -def mlpf_loss(y, ypred, batch, regression_weights, task_loss_weighter=None): +def mlpf_loss(y, ypred, batch, regression_weights, task_loss_weighter=None, *, binary_focal_gamma=None): """Compute the standard MLPF objective for a batch of events.""" if task_loss_weighter is None: - loss = event_loss(y, ypred, batch, regression_weights) + loss = event_loss(y, ypred, batch, regression_weights, binary_focal_gamma=binary_focal_gamma) loss_opt = sum(loss.values()) task_loss_diagnostics = None else: unweighted_regression_weights = {feature: 1.0 for feature in REGRESSION_FEATURES} - loss = event_loss(y, ypred, batch, unweighted_regression_weights) + loss = event_loss(y, ypred, batch, unweighted_regression_weights, binary_focal_gamma=binary_focal_gamma) loss_opt, task_loss_diagnostics = task_loss_weighter(loss) loss["Total"] = loss_opt diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 6f88d8053..fe7a5ba94 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -94,6 +94,23 @@ from mlpf.jet_utils import get_jet_config UNIT_REGRESSION_WEIGHTS = {feature: 1.0 for feature in REGRESSION_FEATURES} +JET_VALIDATION_METRICS = ( + "med", + "iqr", + "num_reference_jets", + "num_candidate_jets", + "num_angular_matches", + "num_response_qualified_matches", + "response_rel_pt_cut", + "angular_recall", + "angular_precision", + "angular_f1", + "angular_fake_rate", + "response_qualified_recall", + "response_qualified_precision", + "response_qualified_f1", + "match_frac", +) def seed_everything(seed): @@ -154,7 +171,7 @@ def _add_accumulator(accum, key, value, count=1.0): accum[key][1] += torch.as_tensor(float(count), device=value.device, dtype=torch.float32) -def _accumulate_domain_losses_and_stats(batch, ytarget, ypred, regression_weights, accum): +def _accumulate_domain_losses_and_stats(batch, ytarget, ypred, regression_weights, accum, *, binary_focal_gamma=None): domain_labels = _event_domain_labels(batch) if domain_labels is None: return @@ -175,7 +192,13 @@ def _accumulate_domain_losses_and_stats(batch, ytarget, ypred, regression_weight "cls_id_onehot": ypred["cls_id_onehot"][valid], **{feature: ypred[feature][valid] for feature in REGRESSION_FEATURES}, } - losses = particle_loss(particle_targets, particle_predictions, batch.X[..., 1][valid], regression_weights) + losses = particle_loss( + particle_targets, + particle_predictions, + batch.X[..., 1][valid], + regression_weights, + binary_focal_gamma=binary_focal_gamma, + ) for loss_name, loss_value in losses.items(): _add_accumulator(accum, f"diagnostic/loss/{label}/{loss_name}", loss_value) @@ -365,7 +388,14 @@ def model_step(batch, model, loss_fn, regression_weights): ) else: ytarget = unpack_target(batch.ytarget, model_module) - loss_opt, losses_detached, task_loss_diagnostics = loss_fn(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + loss_opt, losses_detached, task_loss_diagnostics = loss_fn( + ytarget, + ypred, + batch, + regression_weights, + _get_task_loss_weighter(model), + binary_focal_gamma=model_module.config.binary_classification_focal_gamma, + ) return loss_opt, losses_detached, task_loss_diagnostics, ypred_raw, ypred, ytarget @@ -484,7 +514,14 @@ def train_step( ) else: ytarget = unpack_target(batch.ytarget, model_module) - loss_opt, loss, task_loss_diagnostics = mlpf_loss(ytarget, ypred, batch, regression_weights, _get_task_loss_weighter(model)) + loss_opt, loss, task_loss_diagnostics = mlpf_loss( + ytarget, + ypred, + batch, + regression_weights, + _get_task_loss_weighter(model), + binary_focal_gamma=model_module.config.binary_classification_focal_gamma, + ) phase_start = _record_phase_time_if_enabled(diagnostics.get("time", {}), "loss", phase_start, device_type, log_this_step) if log_this_step: _collect_step_memory(rank, "after_loss", diagnostics) @@ -813,6 +850,7 @@ def evaluate( ypred, UNIT_REGRESSION_WEIGHTS, diagnostic_accum, + binary_focal_gamma=model_module.config.binary_classification_focal_gamma, ) # Save validation plots for first batch @@ -1035,7 +1073,7 @@ def _run_validation_cycle( plot_metrics = make_plots(outdir, sample, config.dataset, testdir_name, config.ntest) plot_metrics_sample[sample] = plot_metrics # Log key jet metrics to TensorBoard and CometML - for k in ["med", "iqr", "match_frac"]: + for k in JET_VALIDATION_METRICS: metric_name = f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/{k}" metric_value = plot_metrics["jet_ratio"]["jet_ratio_target_to_pred_pt"][k] tensorboard_writer_valid.add_scalar(metric_name, metric_value, step) @@ -1072,12 +1110,12 @@ def _run_validation_cycle( "valid_loader_state_dict": valid_loader.state_dict(), } for sample in plot_metrics_sample.keys(): - for metric in ["iqr", "match_frac"]: + for metric in JET_VALIDATION_METRICS: metric_name = f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/{metric}" metrics[metric_name] = plot_metrics_sample[sample]["jet_ratio"]["jet_ratio_target_to_pred_pt"][metric] metrics[f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/combined"] = ( metrics[f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/iqr"] - - metrics[f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/match_frac"] + - metrics[f"step/{sample}/jet_ratio/jet_ratio_target_to_pred_pt/response_qualified_f1"] ) save_checkpoint(Path(temp_checkpoint_dir) / "checkpoint.pth", model, optimizer, extra_state) ray.train.report(metrics, checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir)) diff --git a/mlpf/plotting/plot_utils.py b/mlpf/plotting/plot_utils.py index 299d8c0c1..64db72cee 100644 --- a/mlpf/plotting/plot_utils.py +++ b/mlpf/plotting/plot_utils.py @@ -13,6 +13,8 @@ import sys import vector +from mlpf.jet_utils import jet_matching_metrics + SAMPLE_LABEL_CMS = { "TTbar_14TeV_TuneCUETP8M1_cfi": r"$\mathrm{t}\bar{\mathrm{t}}$+PU events", "ZTT_All_hadronic_14TeV_TuneCUETP8M1_cfi": r"$Z\rightarrow \tau \tau$+PU events", @@ -239,6 +241,24 @@ def med_iqr(arr): return p50, p75 - p25 +def jet_response_metrics(yvals, reference, candidate, response_rel_pt_cut=0.5): + """Return response resolution and auditable one-to-one matching metrics.""" + + response_ratios = yvals[f"jet_ratio_{reference}_to_{candidate}_pt"] + median, iqr = med_iqr(response_ratios) + metrics = { + "med": median, + "iqr": iqr, + **jet_matching_metrics( + response_ratios, + awkward.count(yvals[f"jets_{reference}_pt"], axis=None), + awkward.count(yvals[f"jets_{candidate}_pt"], axis=None), + response_rel_pt_cut=response_rel_pt_cut, + ), + } + return metrics + + def get_eff(df, pid): v0 = np.sum(df == pid) return v0 / len(df), np.sqrt(v0) / len(df) @@ -299,7 +319,6 @@ def cld_label(ax): "cld": cld_label, "clic_hits": clic_label, "cld_hits": cld_label, - "clic_hits": clic_label, } @@ -692,6 +711,7 @@ def plot_jet_ratio( dataset=None, sample=None, baseline_yvals=None, + response_rel_pt_cut=0.5, ): baseline_yvals = baseline_yvals if baseline_yvals is not None else yvals plt.figure() @@ -702,11 +722,9 @@ def plot_jet_ratio( ret_dict = {} p = med_iqr(yvals["jet_ratio_gen_to_target_pt"]) - ret_dict["jet_ratio_gen_to_target_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(yvals["jet_ratio_gen_to_target_pt"]) / awkward.count(yvals["jets_gen_pt"]), - } + ret_dict["jet_ratio_gen_to_target_pt"] = jet_response_metrics( + yvals, "gen", "target", response_rel_pt_cut + ) plt.hist( yvals["jet_ratio_gen_to_target_pt"], bins=bins, @@ -716,11 +734,9 @@ def plot_jet_ratio( ) p = med_iqr(baseline_yvals["jet_ratio_gen_to_cand_pt"]) - ret_dict["jet_ratio_gen_to_cand_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(baseline_yvals["jet_ratio_gen_to_cand_pt"]) / awkward.count(baseline_yvals["jets_gen_pt"]), - } + ret_dict["jet_ratio_gen_to_cand_pt"] = jet_response_metrics( + baseline_yvals, "gen", "cand", response_rel_pt_cut + ) plt.hist( baseline_yvals["jet_ratio_gen_to_cand_pt"], bins=bins, @@ -730,11 +746,9 @@ def plot_jet_ratio( ) p = med_iqr(yvals["jet_ratio_gen_to_pred_pt"]) - ret_dict["jet_ratio_gen_to_pred_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(yvals["jet_ratio_gen_to_pred_pt"]) / awkward.count(yvals["jets_gen_pt"]), - } + ret_dict["jet_ratio_gen_to_pred_pt"] = jet_response_metrics( + yvals, "gen", "pred", response_rel_pt_cut + ) plt.hist( yvals["jet_ratio_gen_to_pred_pt"], bins=bins, @@ -744,11 +758,9 @@ def plot_jet_ratio( ) p = med_iqr(yvals["jet_ratio_gen_to_pred_nopu_pt"]) - ret_dict["jet_ratio_gen_to_pred_nopu_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(yvals["jet_ratio_gen_to_pred_nopu_pt"]) / awkward.count(yvals["jets_gen_pt"]), - } + ret_dict["jet_ratio_gen_to_pred_nopu_pt"] = jet_response_metrics( + yvals, "gen", "pred_nopu", response_rel_pt_cut + ) plt.hist( yvals["jet_ratio_gen_to_pred_nopu_pt"], bins=bins, @@ -784,11 +796,9 @@ def plot_jet_ratio( ax = plt.axes() p = med_iqr(baseline_yvals["jet_ratio_target_to_cand_pt"]) - ret_dict["jet_ratio_target_to_cand_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(baseline_yvals["jet_ratio_target_to_cand_pt"]) / awkward.count(baseline_yvals["jets_target_pt"]), - } + ret_dict["jet_ratio_target_to_cand_pt"] = jet_response_metrics( + baseline_yvals, "target", "cand", response_rel_pt_cut + ) plt.plot([], []) plt.hist( baseline_yvals["jet_ratio_target_to_cand_pt"], @@ -798,11 +808,9 @@ def plot_jet_ratio( label="PF $({:.2f}\pm{:.2f})$".format(p[0], p[1]), ) p = med_iqr(yvals["jet_ratio_target_to_pred_pt"]) - ret_dict["jet_ratio_target_to_pred_pt"] = { - "med": p[0], - "iqr": p[1], - "match_frac": awkward.count(yvals["jet_ratio_target_to_pred_pt"]) / awkward.count(yvals["jets_target_pt"]), - } + ret_dict["jet_ratio_target_to_pred_pt"] = jet_response_metrics( + yvals, "target", "pred", response_rel_pt_cut + ) plt.hist( yvals["jet_ratio_target_to_pred_pt"], bins=bins, diff --git a/particleflow_spec.yaml b/particleflow_spec.yaml index 63dfc6cf4..e3fae9e46 100644 --- a/particleflow_spec.yaml +++ b/particleflow_spec.yaml @@ -604,6 +604,7 @@ models: type: "attention" input_encoding: "split" output_mode: "elementwise" + binary_classification_focal_gamma: 2.0 attention: num_convs: 3 head_dim: 16 @@ -739,6 +740,7 @@ models: architecture: type: "attention" input_encoding: "split" + binary_classification_focal_gamma: 2.0 attention: num_convs: 3 head_dim: 16 diff --git a/tests/test_jet_utils.py b/tests/test_jet_utils.py new file mode 100644 index 000000000..4f118fe74 --- /dev/null +++ b/tests/test_jet_utils.py @@ -0,0 +1,87 @@ +import awkward as ak +import numpy as np +import pytest + +from mlpf.jet_utils import jet_matching_metrics, match_jets +from mlpf.plotting.plot_utils import jet_response_metrics + + +def make_jets(eta_events, phi_events): + return ak.zip({"eta": eta_events, "phi": phi_events}) + + +def test_match_jets_is_one_to_one_and_maximizes_angular_matches(): + reference = make_jets([[0.0, 0.08]], [[0.0, 0.0]]) + candidate = make_jets([[0.04, 0.12]], [[0.0, 0.0]]) + + reference_indices, candidate_indices = match_jets(reference, candidate, 0.1) + + assert reference_indices == [[0, 1]] + assert candidate_indices == [[0, 1]] + + +def test_match_jets_does_not_reuse_a_candidate(): + reference = make_jets([[0.0, 0.05]], [[0.0, 0.0]]) + candidate = make_jets([[0.025]], [[0.0]]) + + reference_indices, candidate_indices = match_jets(reference, candidate, 0.1) + + assert len(reference_indices[0]) == 1 + assert candidate_indices == [[0]] + + +def test_match_jets_wraps_phi_and_handles_empty_events(): + reference = make_jets([[0.0], []], [[np.pi - 0.01], []]) + candidate = make_jets([[0.0], []], [[-np.pi + 0.01], []]) + + reference_indices, candidate_indices = match_jets(reference, candidate, 0.1) + + assert reference_indices == [[0], []] + assert candidate_indices == [[0], []] + + +def test_jet_matching_metrics_separate_angular_and_response_quality(): + metrics = jet_matching_metrics( + response_ratios=[1.0, 1.4, 1.6], + num_reference_jets=4, + num_candidate_jets=5, + response_rel_pt_cut=0.5, + ) + + assert metrics["num_reference_jets"] == 4 + assert metrics["num_candidate_jets"] == 5 + assert metrics["num_angular_matches"] == 3 + assert metrics["num_response_qualified_matches"] == 2 + assert metrics["response_rel_pt_cut"] == 0.5 + assert metrics["angular_recall"] == pytest.approx(3 / 4) + assert metrics["angular_precision"] == pytest.approx(3 / 5) + assert metrics["angular_f1"] == pytest.approx(6 / 9) + assert metrics["angular_fake_rate"] == pytest.approx(2 / 5) + assert metrics["response_qualified_recall"] == pytest.approx(2 / 4) + assert metrics["response_qualified_precision"] == pytest.approx(2 / 5) + assert metrics["response_qualified_f1"] == pytest.approx(4 / 9) + assert metrics["match_frac"] == metrics["angular_recall"] + + +def test_jet_matching_metrics_reject_invalid_counts_and_cuts(): + with pytest.raises(ValueError, match="cannot exceed"): + jet_matching_metrics([1.0, 1.0], 1, 2) + with pytest.raises(ValueError, match="must be positive"): + jet_matching_metrics([], 0, 0, response_rel_pt_cut=0) + + +def test_jet_response_metrics_uses_both_collection_denominators(): + yvals = { + "jet_ratio_target_to_pred_pt": np.asarray([1.0, 1.6]), + "jets_target_pt": ak.Array([[10.0, 20.0, 30.0]]), + "jets_pred_pt": ak.Array([[10.0, 32.0, 8.0, 6.0]]), + } + + metrics = jet_response_metrics(yvals, "target", "pred", response_rel_pt_cut=0.5) + + assert metrics["med"] == pytest.approx(1.3) + assert metrics["iqr"] == pytest.approx(0.3) + assert metrics["angular_recall"] == pytest.approx(2 / 3) + assert metrics["angular_fake_rate"] == pytest.approx(2 / 4) + assert metrics["response_qualified_recall"] == pytest.approx(1 / 3) + assert metrics["response_qualified_precision"] == pytest.approx(1 / 4) diff --git a/tests/test_standard_loss.py b/tests/test_standard_loss.py index 3cccf4971..d12ce5717 100644 --- a/tests/test_standard_loss.py +++ b/tests/test_standard_loss.py @@ -5,7 +5,15 @@ import math import torch -from mlpf.model.losses import REGRESSION_FEATURES, event_loss, mlpf_loss, particle_loss, regression_loss +from mlpf.model.losses import ( + FocalLoss, + REGRESSION_FEATURES, + classification_loss, + event_loss, + mlpf_loss, + particle_loss, + regression_loss, +) from mlpf.model.PFDataset import PFBatch @@ -48,6 +56,22 @@ def get_mock_data(batch_size=2, seq_len=10, num_classes=6): return batch, y, ypred +def test_binary_particle_classification_can_use_focal_loss(): + y = {"cls_id": torch.tensor([0, 0, 0, 1])} + logits = torch.tensor([[5.0, -5.0], [4.0, -4.0], [3.0, -3.0], [0.5, -0.5]]) + ypred = { + "cls_binary": logits, + "cls_id_onehot": torch.zeros(4, 2), + } + + focal = classification_loss(y, ypred, binary_focal_gamma=2.0)["Classification_binary"] + expected = 10.0 * FocalLoss(gamma=2.0)(logits, (y["cls_id"] != 0).long()) + cross_entropy = classification_loss(y, ypred)["Classification_binary"] + + torch.testing.assert_close(focal, expected) + assert focal < cross_entropy + + def test_mlpf_loss_standard(): batch, y, ypred = get_mock_data() loss_opt, losses, _ = mlpf_loss(y, ypred, batch, REGRESSION_WEIGHTS) diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index 530a133ff..25857edd6 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -18,6 +18,7 @@ ROOT = Path(__file__).resolve().parents[1] SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" BACKBONE_SCENARIO = ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" +PF_HITS_SCENARIO = ROOT / "configs/training/scenarios/cld_pf_hits_comparison.yaml" PLATFORMS = ROOT / "configs/training/platforms" @@ -60,6 +61,30 @@ def test_backbone_comparison_scenario_keeps_elementwise_output_and_depth_fixed() assert jobs[1].resolved_config.model.heptv2.block_size == 128 +def test_pf_hits_comparison_scenario_resolves_three_40k_variants(): + scenario = load_training_scenario(PF_HITS_SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + + assert [job.variant_name for job in jobs] == ["pf", "elementwise_hits", "set_hits"] + assert [job.model_name for job in jobs] == ["pyg-cld-v1", "pyg-cld-hits-v1", "pyg-cld-hits-set-v1"] + assert [job.resolved_config.dataset.value for job in jobs] == ["cld", "cld_hits", "cld_hits"] + assert [job.resolved_config.model.output_mode.value for job in jobs] == ["elementwise", "elementwise", "set"] + assert [job.resolved_config.model.binary_classification_focal_gamma for job in jobs] == [None, 2.0, 2.0] + assert {job.resolved_config.num_steps for job in jobs} == {40000} + assert {job.resolved_config.val_freq for job in jobs} == {5000} + assert {job.resolved_config.checkpoint_freq for job in jobs} == {5000} + assert {job.resolved_config.lr for job in jobs} == {0.001} + assert {job.global_batch_size for job in jobs} == {8} + assert {job.seed for job in jobs} == {12345} + + @pytest.mark.parametrize( ("profile_name", "expected_multiplier"), [ diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index e27c7f25d..c3ec343eb 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -17,6 +17,7 @@ def test_picker_discovers_scenarios_and_accelerators(): assert "cld_hits_output_comparison" in scenarios assert "cld_hits_backbone_comparison" in scenarios + assert "cld_pf_hits_comparison" in scenarios assert {"a100", "h100", "h200"}.issubset(accelerators) From 24e8205bae0ed3fece9e69ff4b8518b6713c25f4 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Tue, 8 Sep 2026 23:45:35 +0300 Subject: [PATCH 20/29] Fix PF backbone in hit comparison scenario --- .../scenarios/cld_pf_hits_comparison.yaml | 17 ++++++++++++++--- tests/test_training_scenarios.py | 9 +++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/configs/training/scenarios/cld_pf_hits_comparison.yaml b/configs/training/scenarios/cld_pf_hits_comparison.yaml index f85e45262..ec438522a 100644 --- a/configs/training/scenarios/cld_pf_hits_comparison.yaml +++ b/configs/training/scenarios/cld_pf_hits_comparison.yaml @@ -7,9 +7,18 @@ variants: 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 @@ -52,9 +61,6 @@ common_overrides: model.task_queries: false 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 model.attention.use_jagged_attention: true allowed_variant_differences: @@ -69,6 +75,11 @@ allowed_variant_differences: - 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 diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index 25857edd6..da6c5cb3a 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -77,6 +77,15 @@ def test_pf_hits_comparison_scenario_resolves_three_40k_variants(): assert [job.resolved_config.dataset.value for job in jobs] == ["cld", "cld_hits", "cld_hits"] assert [job.resolved_config.model.output_mode.value for job in jobs] == ["elementwise", "elementwise", "set"] assert [job.resolved_config.model.binary_classification_focal_gamma for job in jobs] == [None, 2.0, 2.0] + assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} + assert [ + ( + job.resolved_config.model.backbone.num_tracker_layers, + job.resolved_config.model.backbone.num_calo_layers, + job.resolved_config.model.backbone.num_common_layers, + ) + for job in jobs + ] == [(None, None, None), (2, 2, 2), (2, 2, 2)] assert {job.resolved_config.num_steps for job in jobs} == {40000} assert {job.resolved_config.val_freq for job in jobs} == {5000} assert {job.resolved_config.checkpoint_freq for job in jobs} == {5000} From 7a7dd02764d1d763ec72de83d6fb91116f3b0019 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 14:53:18 +0300 Subject: [PATCH 21/29] format --- mlpf/plotting/plot_utils.py | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/mlpf/plotting/plot_utils.py b/mlpf/plotting/plot_utils.py index 64db72cee..3eb28cdfe 100644 --- a/mlpf/plotting/plot_utils.py +++ b/mlpf/plotting/plot_utils.py @@ -722,9 +722,7 @@ def plot_jet_ratio( ret_dict = {} p = med_iqr(yvals["jet_ratio_gen_to_target_pt"]) - ret_dict["jet_ratio_gen_to_target_pt"] = jet_response_metrics( - yvals, "gen", "target", response_rel_pt_cut - ) + ret_dict["jet_ratio_gen_to_target_pt"] = jet_response_metrics(yvals, "gen", "target", response_rel_pt_cut) plt.hist( yvals["jet_ratio_gen_to_target_pt"], bins=bins, @@ -734,9 +732,7 @@ def plot_jet_ratio( ) p = med_iqr(baseline_yvals["jet_ratio_gen_to_cand_pt"]) - ret_dict["jet_ratio_gen_to_cand_pt"] = jet_response_metrics( - baseline_yvals, "gen", "cand", response_rel_pt_cut - ) + ret_dict["jet_ratio_gen_to_cand_pt"] = jet_response_metrics(baseline_yvals, "gen", "cand", response_rel_pt_cut) plt.hist( baseline_yvals["jet_ratio_gen_to_cand_pt"], bins=bins, @@ -746,9 +742,7 @@ def plot_jet_ratio( ) p = med_iqr(yvals["jet_ratio_gen_to_pred_pt"]) - ret_dict["jet_ratio_gen_to_pred_pt"] = jet_response_metrics( - yvals, "gen", "pred", response_rel_pt_cut - ) + ret_dict["jet_ratio_gen_to_pred_pt"] = jet_response_metrics(yvals, "gen", "pred", response_rel_pt_cut) plt.hist( yvals["jet_ratio_gen_to_pred_pt"], bins=bins, @@ -758,9 +752,7 @@ def plot_jet_ratio( ) p = med_iqr(yvals["jet_ratio_gen_to_pred_nopu_pt"]) - ret_dict["jet_ratio_gen_to_pred_nopu_pt"] = jet_response_metrics( - yvals, "gen", "pred_nopu", response_rel_pt_cut - ) + ret_dict["jet_ratio_gen_to_pred_nopu_pt"] = jet_response_metrics(yvals, "gen", "pred_nopu", response_rel_pt_cut) plt.hist( yvals["jet_ratio_gen_to_pred_nopu_pt"], bins=bins, @@ -796,9 +788,7 @@ def plot_jet_ratio( ax = plt.axes() p = med_iqr(baseline_yvals["jet_ratio_target_to_cand_pt"]) - ret_dict["jet_ratio_target_to_cand_pt"] = jet_response_metrics( - baseline_yvals, "target", "cand", response_rel_pt_cut - ) + ret_dict["jet_ratio_target_to_cand_pt"] = jet_response_metrics(baseline_yvals, "target", "cand", response_rel_pt_cut) plt.plot([], []) plt.hist( baseline_yvals["jet_ratio_target_to_cand_pt"], @@ -808,9 +798,7 @@ def plot_jet_ratio( label="PF $({:.2f}\pm{:.2f})$".format(p[0], p[1]), ) p = med_iqr(yvals["jet_ratio_target_to_pred_pt"]) - ret_dict["jet_ratio_target_to_pred_pt"] = jet_response_metrics( - yvals, "target", "pred", response_rel_pt_cut - ) + ret_dict["jet_ratio_target_to_pred_pt"] = jet_response_metrics(yvals, "target", "pred", response_rel_pt_cut) plt.hist( yvals["jet_ratio_target_to_pred_pt"], bins=bins, From 5304040664470f7f01ab787bac208814d6034654 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 15:04:37 +0300 Subject: [PATCH 22/29] CLIC and CLD trainings --- configs/training/platforms/flatiron_a100.yaml | 4 +- configs/training/platforms/flatiron_h100.yaml | 4 +- configs/training/platforms/flatiron_h200.yaml | 4 +- configs/training/platforms/local.yaml | 4 +- configs/training/platforms/tallinn_l40.yaml | 4 +- mlpf/training_scenarios.py | 47 +++++++--- particleflow_spec.yaml | 23 ++++- scripts/local/make_local_available_spec.py | 1 + scripts/local/train_scenario.sh | 10 ++- tests/test_training_scenarios.py | 89 +++++++++++++++++++ tests/test_training_submission.py | 11 +++ 11 files changed, 180 insertions(+), 21 deletions(-) diff --git a/configs/training/platforms/flatiron_a100.yaml b/configs/training/platforms/flatiron_a100.yaml index e6d60ae19..99ee214d2 100644 --- a/configs/training/platforms/flatiron_a100.yaml +++ b/configs/training/platforms/flatiron_a100.yaml @@ -1,6 +1,8 @@ name: flatiron_a100 gpus: 4 -data_dir: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds +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 diff --git a/configs/training/platforms/flatiron_h100.yaml b/configs/training/platforms/flatiron_h100.yaml index ca3357579..47c6e5312 100644 --- a/configs/training/platforms/flatiron_h100.yaml +++ b/configs/training/platforms/flatiron_h100.yaml @@ -1,6 +1,8 @@ name: flatiron_h100 gpus: 8 -data_dir: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds +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 diff --git a/configs/training/platforms/flatiron_h200.yaml b/configs/training/platforms/flatiron_h200.yaml index 873acf247..45ec87bb5 100644 --- a/configs/training/platforms/flatiron_h200.yaml +++ b/configs/training/platforms/flatiron_h200.yaml @@ -1,6 +1,8 @@ name: flatiron_h200 gpus: 8 -data_dir: /mnt/ceph/users/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds +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 diff --git a/configs/training/platforms/local.yaml b/configs/training/platforms/local.yaml index be1f813c3..9ac9b7661 100644 --- a/configs/training/platforms/local.yaml +++ b/configs/training/platforms/local.yaml @@ -1,6 +1,8 @@ name: local gpus: 1 -data_dir: /mnt/work/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds +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 diff --git a/configs/training/platforms/tallinn_l40.yaml b/configs/training/platforms/tallinn_l40.yaml index b841a0ee9..323460865 100644 --- a/configs/training/platforms/tallinn_l40.yaml +++ b/configs/training/platforms/tallinn_l40.yaml @@ -1,6 +1,8 @@ name: tallinn_l40 gpus: 2 -data_dir: /local/${USER}/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds +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 diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index d849e0638..fb4ea69fc 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -34,6 +34,8 @@ class ScenarioVariant(BaseModel): model_config = ConfigDict(extra="forbid") model_name: str + # Defaults to the scenario-level production; set it to compare detectors. + production_name: str | None = None overrides: dict[str, Any] = Field(default_factory=dict) @model_validator(mode="after") @@ -83,6 +85,9 @@ def validate_scenario(self): raise ValueError(f"Common overrides must not set derived keys: {sorted(invalid)}") return self + def variant_production(self, variant_name): + return self.variants[variant_name].production_name or self.production_name + class SlurmProfile(BaseModel): model_config = ConfigDict(extra="forbid") @@ -114,7 +119,9 @@ class PlatformProfile(BaseModel): name: str gpus: int = Field(gt=0) - data_dir: str + # Either one TFDS directory shared by every production, or a mapping from + # production name to that production's TFDS directory. + data_dir: str | dict[str, str] experiments_dir: str environment: dict[str, str] = Field(default_factory=dict) runtime_overrides: dict[str, Any] = Field(default_factory=dict) @@ -125,8 +132,19 @@ def validate_runtime_overrides(self): invalid = set(self.runtime_overrides).difference(PLATFORM_OVERRIDE_KEYS) if invalid: raise ValueError("Platform profiles may only set runtime-specific overrides; " f"invalid keys: {sorted(invalid)}") + if isinstance(self.data_dir, dict) and not self.data_dir: + raise ValueError("Platform data_dir mapping must name at least one production") return self + def data_dir_for(self, production_name): + if isinstance(self.data_dir, str): + return self.data_dir + if production_name not in self.data_dir: + raise ValueError( + f"Platform profile {self.name!r} has no data_dir for production {production_name!r}; " f"available: {sorted(self.data_dir)}" + ) + return self.data_dir[production_name] + class ResolvedScenarioJob(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -135,6 +153,8 @@ class ResolvedScenarioJob(BaseModel): platform_name: str variant_name: str model_name: str + production_name: str + data_dir: str seed: int global_batch_size: int per_gpu_batch_size: int @@ -154,7 +174,10 @@ def load_training_scenario(path): def load_platform_profile(path): profile = PlatformProfile.model_validate(_read_yaml(path)) - profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) + if isinstance(profile.data_dir, str): + profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) + else: + profile.data_dir = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.data_dir.items()} profile.experiments_dir = os.path.expandvars(os.path.expanduser(profile.experiments_dir)) profile.environment = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.environment.items()} return profile @@ -193,12 +216,12 @@ def _settings_as_extra_args(settings): return args -def _config_args(profile, settings): +def _config_args(profile, settings, data_dir): return SimpleNamespace( train=True, test=True, pipeline=False, - data_dir=profile.data_dir, + data_dir=data_dir, gpus=profile.gpus, compile=settings.get("compile"), comet=settings.get("comet"), @@ -250,13 +273,15 @@ def resolve_scenario_job( settings = _merge_settings(scenario, platform, variant, extra_overrides) settings["seed"] = seed selected_spec = str(spec_file or scenario.spec_file) + production_name = scenario.variant_production(variant_name) + data_dir = platform.data_dir_for(production_name) with _temporary_environment(platform.environment): config = MLPFConfig.from_spec( selected_spec, variant.model_name, - scenario.production_name, - args=_config_args(platform, settings), + production_name, + args=_config_args(platform, settings, data_dir), extra_args=_settings_as_extra_args(settings), ) @@ -276,8 +301,8 @@ def resolve_scenario_job( config = MLPFConfig.from_spec( selected_spec, variant.model_name, - scenario.production_name, - args=_config_args(platform, settings), + production_name, + args=_config_args(platform, settings, data_dir), extra_args=_settings_as_extra_args(settings), ) @@ -286,6 +311,8 @@ def resolve_scenario_job( platform_name=platform.name, variant_name=variant_name, model_name=variant.model_name, + production_name=production_name, + data_dir=data_dir, seed=seed, global_batch_size=target_global_batch, per_gpu_batch_size=dataset_batch_size * multiplier, @@ -366,9 +393,9 @@ def _pipeline_command(job, scenario, platform, spec_file, experiment_dir): "--model-name", job.model_name, "--production-name", - scenario.production_name, + job.production_name, "--data-dir", - platform.data_dir, + job.data_dir, "--experiment-dir", str(experiment_dir), "train", diff --git a/particleflow_spec.yaml b/particleflow_spec.yaml index e3fae9e46..8f6138862 100644 --- a/particleflow_spec.yaml +++ b/particleflow_spec.yaml @@ -728,7 +728,7 @@ models: version: "3.2.1" splits: ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] - pyg-clic-hits-v1: + pyg-clic-hits-v1: &pyg_clic_hits_v1 <<: *defaults dataset: clic_hits gpu_batch_multiplier: 16 @@ -737,9 +737,10 @@ models: batch_size: 1 lr: 0.0001 - architecture: + architecture: &pyg_clic_hits_architecture type: "attention" input_encoding: "split" + output_mode: "elementwise" binary_classification_focal_gamma: 2.0 attention: num_convs: 3 @@ -785,3 +786,21 @@ models: - name: "clic_edm_qq_hits" version: "3.2.1" splits: ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] + + pyg-clic-hits-set-v1: + <<: *pyg_clic_hits_v1 + gpu_batch_multiplier: 16 + + hyperparameters: + batch_size: 1 + lr: 0.0001 + + architecture: + <<: *pyg_clic_hits_architecture + output_mode: "set" + set_decoder: + num_slots: 256 + num_layers: 2 + num_heads: 8 + ffn_multiplier: 4.0 + dropout: 0.0 diff --git a/scripts/local/make_local_available_spec.py b/scripts/local/make_local_available_spec.py index 481b49153..656d5c230 100755 --- a/scripts/local/make_local_available_spec.py +++ b/scripts/local/make_local_available_spec.py @@ -33,6 +33,7 @@ def main(): "pyg-cld-hits-v1": ("cld_hits", "cld_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-cld-hits-set-v1": ("cld_hits", "cld_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-clic-hits-v1": ("clic_hits", "clic_edm_ttbar_hits", args.hit_version, args.hit_splits), + "pyg-clic-hits-set-v1": ("clic_hits", "clic_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-cld-v1": ("cld", "cld_edm_ttbar_pf", args.pf_version, args.pf_splits), "pyg-clic-v1": ("clic", "clic_edm_ttbar_pf", args.pf_version, args.pf_splits), } diff --git a/scripts/local/train_scenario.sh b/scripts/local/train_scenario.sh index 67d859e8f..fdbd2666a 100755 --- a/scripts/local/train_scenario.sh +++ b/scripts/local/train_scenario.sh @@ -44,7 +44,6 @@ export PF_SITE=local PLATFORM_FILE=${PLATFORM_FILE:-configs/training/platforms/local.yaml} SPEC_FILE=${SPEC_FILE:-$(uv run python3 scripts/get_param.py "$SCENARIO_FILE" spec_file particleflow_spec.yaml)} -PRODUCTION_NAME=$(uv run python3 scripts/get_param.py "$SCENARIO_FILE" production_name) USE_LOCAL_AVAILABLE_SPEC=${USE_LOCAL_AVAILABLE_SPEC:-true} LOCAL_SPEC_FILE=${LOCAL_SPEC_FILE:-/tmp/particleflow_local_available_spec.yaml} SEED=${SEED:-} @@ -66,6 +65,9 @@ PREFETCH_FACTOR=${PREFETCH_FACTOR:-4} VALIDATION_DIAGNOSTICS_BATCHES=${VALIDATION_DIAGNOSTICS_BATCHES:-4} EXPERIMENTS_DIR=${EXPERIMENTS_DIR:-experiments} PAD_TO_MULTIPLE_ELEMENTS=${PAD_TO_MULTIPLE_ELEMENTS:-128} +# The platform profile maps each production (cld, clic) to its TFDS directory. +# Set DATA_DIR to force one directory for every variant instead. +DATA_DIR=${DATA_DIR:-} read -r -a HIT_SPLIT_LIST <<< "$HIT_SPLITS" if [[ "$USE_LOCAL_AVAILABLE_SPEC" == "true" ]]; then @@ -76,14 +78,11 @@ if [[ "$USE_LOCAL_AVAILABLE_SPEC" == "true" ]]; then SPEC_FILE="$LOCAL_SPEC_FILE" fi -DATA_DIR=${DATA_DIR:-$(uv run python3 scripts/get_param.py "$SPEC_FILE" productions."$PRODUCTION_NAME".workspace_dir)/tfds/} - RUN_ARGS=( --scenario "$SCENARIO_FILE" --platform "$PLATFORM_FILE" --spec-file "$SPEC_FILE" --global-batch-size "$GLOBAL_BATCH_SIZE" - --data-dir "$DATA_DIR" --experiments-dir "$EXPERIMENTS_DIR" --set "data_config=$DATA_CONFIG" --set "num_steps=$NUM_STEPS" @@ -96,6 +95,9 @@ RUN_ARGS=( --set "validation_diagnostics_batches=$VALIDATION_DIAGNOSTICS_BATCHES" --set "pad_to_multiple_elements=$PAD_TO_MULTIPLE_ELEMENTS" ) +if [[ -n "$DATA_DIR" ]]; then + RUN_ARGS+=(--data-dir "$DATA_DIR") +fi if [[ -n "$SEED" ]]; then RUN_ARGS+=(--seed "$SEED") fi diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index da6c5cb3a..d3eda3a6d 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -19,6 +19,7 @@ SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" BACKBONE_SCENARIO = ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" PF_HITS_SCENARIO = ROOT / "configs/training/scenarios/cld_pf_hits_comparison.yaml" +CLIC_CLD_SCENARIO = ROOT / "configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml" PLATFORMS = ROOT / "configs/training/platforms" @@ -94,6 +95,94 @@ def test_pf_hits_comparison_scenario_resolves_three_40k_variants(): assert {job.seed for job in jobs} == {12345} +def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): + scenario = load_training_scenario(CLIC_CLD_SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + + assert [job.variant_name for job in jobs] == ["cld_pf", "cld_set_hits", "clic_pf", "clic_set_hits"] + assert [job.model_name for job in jobs] == ["pyg-cld-v1", "pyg-cld-hits-set-v1", "pyg-clic-v1", "pyg-clic-hits-set-v1"] + assert [job.production_name for job in jobs] == ["cld", "cld", "clic", "clic"] + assert [job.data_dir for job in jobs] == [ + platform.data_dir["cld"], + platform.data_dir["cld"], + platform.data_dir["clic"], + platform.data_dir["clic"], + ] + assert [job.resolved_config.data_dir for job in jobs] == [job.data_dir for job in jobs] + assert [job.resolved_config.dataset.value for job in jobs] == ["cld", "cld_hits", "clic", "clic_hits"] + assert [job.resolved_config.model.output_mode.value for job in jobs] == ["elementwise", "set", "elementwise", "set"] + # The set-based hit models run twice the backbone depth of the PF models. + assert [job.resolved_config.model.backbone.num_convs for job in jobs] == [6, 12, 6, 12] + assert [ + ( + job.resolved_config.model.backbone.num_tracker_layers, + job.resolved_config.model.backbone.num_calo_layers, + job.resolved_config.model.backbone.num_common_layers, + ) + for job in jobs + ] == [(None, None, None), (4, 4, 4), (None, None, None), (4, 4, 4)] + assert {job.resolved_config.model.set_decoder.num_layers for job in jobs if job.resolved_config.model.set_decoder} == {8} + assert {job.resolved_config.num_steps for job in jobs} == {50000} + assert {job.resolved_config.val_freq for job in jobs} == {5000} + assert {job.resolved_config.lr for job in jobs} == {0.001} + assert {job.seed for job in jobs} == {12345} + + +def test_platform_data_dir_mapping_requires_the_variant_production(): + scenario = load_training_scenario(CLIC_CLD_SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + platform.data_dir = {"cld": platform.data_dir["cld"]} + + with pytest.raises(ValueError, match="no data_dir for production 'clic'"): + resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + + platform.data_dir = "/tmp/shared_tfds" + jobs = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + ) + assert {job.data_dir for job in jobs} == {"/tmp/shared_tfds"} + + +def test_cli_dry_run_uses_per_production_data_dir(capsys): + from mlpf.training_scenarios import main + + main( + [ + "--scenario", + str(CLIC_CLD_SCENARIO), + "--platform", + str(PLATFORMS / "local.yaml"), + "--spec-file", + str(ROOT / "particleflow_spec.yaml"), + "--global-batch-size", + "8", + "--variant", + "clic_set_hits", + "--dry-run", + ] + ) + + command = capsys.readouterr().out + assert "--production-name clic" in command + assert "--data-dir /mnt/work/mlpf/clic/v1.2.5_key4hep_2025-05-29/tfds" in command + assert "--model.backbone.num_convs 12" in command + + @pytest.mark.parametrize( ("profile_name", "expected_multiplier"), [ diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index c3ec343eb..2fe6a0a04 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -18,9 +18,20 @@ def test_picker_discovers_scenarios_and_accelerators(): assert "cld_hits_output_comparison" in scenarios assert "cld_hits_backbone_comparison" in scenarios assert "cld_pf_hits_comparison" in scenarios + assert "clic_cld_pf_set_hits_comparison" in scenarios assert {"a100", "h100", "h200"}.issubset(accelerators) +def test_multi_production_scenario_submits_one_array_task_per_variant(): + scenario = resolve_scenario_path("clic_cld_pf_set_hits_comparison", ROOT) + profile = resolve_flatiron_profile_path("h100", ROOT) + + command, jobs = build_slurm_submission(scenario, profile, ROOT) + + assert [job.production_name for job in jobs] == ["cld", "cld", "clic", "clic"] + assert command[command.index("--array") + 1] == "0-3" + + def test_h100_submission_is_derived_from_scenario_and_profile(): scenario = resolve_scenario_path("cld_hits_output_comparison", ROOT) profile = resolve_flatiron_profile_path("h100", ROOT) From d5e1561e52ff844581584dff1f20894ae1f77ab5 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 15:06:51 +0300 Subject: [PATCH 23/29] add missing --- .../clic_cld_pf_set_hits_comparison.yaml | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml diff --git a/configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml b/configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml new file mode 100644 index 000000000..06ebfed3b --- /dev/null +++ b/configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml @@ -0,0 +1,95 @@ +name: clic_cld_pf_set_hits_comparison +spec_file: particleflow_spec.yaml +# Scenario-level default; the CLIC variants override it below. +production_name: cld + +# Four jobs per seed: the learned PF model on tracks/clusters and the set-based +# model on raw detector hits, each for CLD and for CLIC. +variants: + cld_pf: + model_name: pyg-cld-v1 + cld_set_hits: + model_name: pyg-cld-hits-set-v1 + overrides: &set_hits_overrides + # Twice the elementwise backbone depth of the PF model (12 layers instead + # of 6): four layers per detector branch plus four common layers. + model.backbone.num_convs: 12 + model.backbone.num_tracker_layers: 4 + model.backbone.num_calo_layers: 4 + model.backbone.num_common_layers: 4 + # Start slots from energetic detector inputs so their initial direction is + # physical, then refine within a local eta-phi neighborhood. The decoder is + # also twice as deep as in cld_pf_hits_comparison (8 layers instead of 4). + 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: 8 + 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 + clic_pf: + model_name: pyg-clic-v1 + production_name: clic + clic_set_hits: + model_name: pyg-clic-hits-set-v1 + production_name: clic + overrides: *set_hits_overrides + +# Add seeds here; launchers derive the four 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: 50000 + 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 and + # the deliberately deeper hit backbone above. + model.task_queries: false + model.backbone.mode: shared + model.backbone.num_convs: 6 + model.attention.use_jagged_attention: true + +allowed_variant_differences: + # The PF runs read tracks/clusters while the hit runs read detector hits, and + # the CLD and CLIC productions live in different TFDS directories. + - data_dir + - 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 set training uses focal presence classification; PF keeps CE. + - model.binary_classification_focal_gamma + # The set-based hit models run a 2x deeper backbone with detector-specific + # branches. The PF models use all six layers as common layers. + - model.backbone.num_convs + - model.backbone.num_tracker_layers + - model.backbone.num_calo_layers + - model.backbone.num_common_layers + # Only the set-based hit variants have a set decoder. + - model.output_mode + - model.set_decoder From 6d2c9cbd9f0a1aa53ef261adcdcbb5810fdea5c9 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 15:45:55 +0300 Subject: [PATCH 24/29] format --- mlpf/model/PFDataset.py | 4 ++-- mlpf/model/training.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index 4a10f9c1a..e79141996 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -754,7 +754,7 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, rank_index = int(rank) if isinstance(rank, int) else 0 split_offset = 0 if split == "train" else 1000 loader_generator.manual_seed(config.seed + 10_000 * rank_index + split_offset + len(loaders[split])) - + worker_kwargs = {} if config.num_workers > 0: worker_kwargs = { @@ -762,7 +762,7 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, "worker_init_fn": set_worker_sharing_strategy, "persistent_workers": True, } - + loader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, diff --git a/mlpf/model/training.py b/mlpf/model/training.py index bcb08bda1..06d6585b4 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -1444,7 +1444,7 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi test_loader_generator = torch.Generator() rank_index = int(rank) if isinstance(rank, int) else 0 test_loader_generator.manual_seed(config.seed + 10_000 * rank_index + 2000) - + worker_kwargs = {} if config.num_workers > 0: worker_kwargs = { @@ -1452,7 +1452,7 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi "worker_init_fn": set_worker_sharing_strategy, "persistent_workers": True, } - + test_loader = torch.utils.data.DataLoader( ds, batch_size=batch_size, From 668551989da3885298bf8e1b5ae0c5c148bed78d Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 16:02:11 +0300 Subject: [PATCH 25/29] fixes --- mlpf/conf.py | 1 + mlpf/model/PFDataset.py | 5 ++++- particleflow_spec.yaml | 2 +- scripts/local/make_local_available_spec.py | 5 ++++- snakemake_jobs/cld/gen/gen_gun_e_10gev.sh | 4 ++-- snakemake_jobs/cld/gen/gen_gun_mu_10gev.sh | 4 ++-- snakemake_jobs/cld/gen/gen_gun_pi_10gev.sh | 4 ++-- snakemake_jobs/cld/gen/gen_qq.sh | 4 ++-- snakemake_jobs/cld/gen/gen_ttbar.sh | 4 ++-- snakemake_jobs/cld/gen/gen_ww_fullhad.sh | 4 ++-- snakemake_jobs/cld/gen/gen_z_qq.sh | 4 ++-- snakemake_jobs/cld/gen/gen_z_tautau.sh | 4 ++-- snakemake_jobs/cld/gen/gen_zh_tautau_240.sh | 4 ++-- snakemake_jobs/cld/gen/gen_zh_tautau_365.sh | 4 ++-- snakemake_jobs/cld/gen/gen_zz.sh | 4 ++-- snakemake_jobs/cld/gen/gen_zz_tautau_240.sh | 4 ++-- snakemake_jobs/cld/gen/gen_zz_tautau_365.sh | 4 ++-- snakemake_jobs/cld/post/post_gun_e_10gev.sh | 2 +- snakemake_jobs/cld/post/post_gun_mu_10gev.sh | 2 +- snakemake_jobs/cld/post/post_gun_pi_10gev.sh | 2 +- snakemake_jobs/cld/post/post_qq.sh | 2 +- snakemake_jobs/cld/post/post_ttbar.sh | 2 +- snakemake_jobs/cld/post/post_ww_fullhad.sh | 2 +- snakemake_jobs/cld/post/post_z_qq.sh | 2 +- snakemake_jobs/cld/post/post_z_tautau.sh | 2 +- snakemake_jobs/cld/post/post_zh_tautau_240.sh | 2 +- snakemake_jobs/cld/post/post_zh_tautau_365.sh | 2 +- snakemake_jobs/cld/post/post_zz.sh | 2 +- snakemake_jobs/cld/post/post_zz_tautau_240.sh | 2 +- snakemake_jobs/cld/post/post_zz_tautau_365.sh | 2 +- snakemake_jobs/cld/tfds/tfds_qq.sh | 2 +- snakemake_jobs/cld/tfds/tfds_ttbar.sh | 2 +- snakemake_jobs/cld/tfds/tfds_ww_fullhad.sh | 2 +- snakemake_jobs/cld/tfds/tfds_zz.sh | 2 +- snakemake_jobs/cld/tfds_hit/tfds_hit_qq.sh | 2 +- snakemake_jobs/cld/tfds_hit/tfds_hit_ttbar.sh | 2 +- snakemake_jobs/cld/tfds_hit/tfds_hit_ww_fullhad.sh | 2 +- snakemake_jobs/cld/tfds_hit/tfds_hit_zz.sh | 2 +- snakemake_jobs/cld/train/train_pyg-cld-hits-v1_cld.sh | 2 +- snakemake_jobs/cld/train/train_pyg-cld-v1_cld.sh | 2 +- snakemake_jobs/idea/gen/gen_gun_e_10gev.sh | 4 ++-- snakemake_jobs/idea/gen/gen_gun_mu_10gev.sh | 4 ++-- snakemake_jobs/idea/gen/gen_gun_pi_10gev.sh | 4 ++-- snakemake_jobs/idea/gen/gen_qq.sh | 4 ++-- snakemake_jobs/idea/gen/gen_ttbar.sh | 4 ++-- snakemake_jobs/idea/gen/gen_ww_fullhad.sh | 4 ++-- snakemake_jobs/idea/post/post_gun_e_10gev.sh | 2 +- snakemake_jobs/idea/post/post_gun_mu_10gev.sh | 2 +- snakemake_jobs/idea/post/post_gun_pi_10gev.sh | 2 +- snakemake_jobs/idea/post/post_qq.sh | 2 +- snakemake_jobs/idea/post/post_ttbar.sh | 2 +- snakemake_jobs/idea/post/post_ww_fullhad.sh | 2 +- snakemake_jobs/idea/tfds/tfds_qq.sh | 2 +- snakemake_jobs/idea/tfds/tfds_ttbar.sh | 2 +- snakemake_jobs/idea/tfds/tfds_ww_fullhad.sh | 2 +- snakemake_jobs/idea/train/train_pyg-idea-pipeline-v1_idea.sh | 2 +- 56 files changed, 81 insertions(+), 74 deletions(-) diff --git a/mlpf/conf.py b/mlpf/conf.py index 76403b177..f6db34e50 100644 --- a/mlpf/conf.py +++ b/mlpf/conf.py @@ -581,6 +581,7 @@ def get_names(cls): "p": -1.0, "ptcut": 5.0, "match_dr": 0.1, + "match_rel_pt": 0.5, }, Dataset.CLIC_HITS.value: { "algo": "ee_genkt_algorithm", diff --git a/mlpf/model/PFDataset.py b/mlpf/model/PFDataset.py index e79141996..e435efae9 100644 --- a/mlpf/model/PFDataset.py +++ b/mlpf/model/PFDataset.py @@ -769,9 +769,12 @@ def get_interleaved_dataloaders(world_size, rank, config: MLPFConfig, use_cuda, collate_fn=Collater(per_particle_keys, ["genmet", "source_id", "input_type_id"]), sampler=sampler, num_workers=config.num_workers, + # Training uses fixed-size batches, but a bounded validation + # sample can be smaller than one per-rank batch (for example, + # nvalid=100 with 8 ranks and batch_size=64). Keep that partial + # validation batch so every rank participates in evaluation. drop_last=split == "train", generator=loader_generator, - drop_last=True, **worker_kwargs, ) diff --git a/particleflow_spec.yaml b/particleflow_spec.yaml index a10b3bb22..1929a2312 100644 --- a/particleflow_spec.yaml +++ b/particleflow_spec.yaml @@ -24,7 +24,7 @@ project: storage_root: "/local/joosep/mlpf" scratch_root: "/scratch/local/joosep" tmpdir: "/scratch/local/joosep/tmp" - project_root: "/home/joosep/particleflow-dev" + project_root: "/home/joosep/particleflow" cmssw_dir: "/scratch/persistent/joosep/CMSSW_15_0_5" gpu_partition: "gpu" cpu_partition: "main" diff --git a/scripts/local/make_local_available_spec.py b/scripts/local/make_local_available_spec.py index 656d5c230..05b901fa8 100755 --- a/scripts/local/make_local_available_spec.py +++ b/scripts/local/make_local_available_spec.py @@ -14,13 +14,15 @@ def set_ttbar_only(model_config, dataset_key, sample_name, version, splits): def main(): - parser = argparse.ArgumentParser(description="Restrict CLD/CLIC models to locally available ttbar datasets.") + parser = argparse.ArgumentParser(description="Restrict CLD/CLIC/IDEA models to locally available ttbar datasets.") parser.add_argument("input_spec", type=Path) parser.add_argument("output_spec", type=Path) parser.add_argument("--hit-version", default="3.2.1") parser.add_argument("--hit-splits", nargs="+", default=["1"]) parser.add_argument("--pf-version", default="3.2.0") parser.add_argument("--pf-splits", nargs="+", default=[str(i) for i in range(1, 11)]) + parser.add_argument("--idea-version", default="0.1.0") + parser.add_argument("--idea-splits", nargs="+", default=["1"]) args = parser.parse_args() input_spec = args.input_spec @@ -36,6 +38,7 @@ def main(): "pyg-clic-hits-set-v1": ("clic_hits", "clic_edm_ttbar_hits", args.hit_version, args.hit_splits), "pyg-cld-v1": ("cld", "cld_edm_ttbar_pf", args.pf_version, args.pf_splits), "pyg-clic-v1": ("clic", "clic_edm_ttbar_pf", args.pf_version, args.pf_splits), + "pyg-idea-pipeline-v1": ("idea", "idea_edm_ttbar_pf", args.idea_version, args.idea_splits), } for model_name, (dataset_key, sample_name, version, splits) in local_datasets.items(): if model_name in spec["models"]: diff --git a/snakemake_jobs/cld/gen/gen_gun_e_10gev.sh b/snakemake_jobs/cld/gen/gen_gun_e_10gev.sh index 334ad81bc..162f1cffc 100755 --- a/snakemake_jobs/cld/gen/gen_gun_e_10gev.sh +++ b/snakemake_jobs/cld/gen/gen_gun_e_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_e_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_e_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh gun_e_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_gun_mu_10gev.sh b/snakemake_jobs/cld/gen/gen_gun_mu_10gev.sh index 2e8f8a883..d40e3a5f8 100755 --- a/snakemake_jobs/cld/gen/gen_gun_mu_10gev.sh +++ b/snakemake_jobs/cld/gen/gen_gun_mu_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_mu_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_mu_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh gun_mu_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_gun_pi_10gev.sh b/snakemake_jobs/cld/gen/gen_gun_pi_10gev.sh index 0123e800a..72dd9256c 100755 --- a/snakemake_jobs/cld/gen/gen_gun_pi_10gev.sh +++ b/snakemake_jobs/cld/gen/gen_gun_pi_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_pi_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/gun_pi_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh gun_pi_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_qq.sh b/snakemake_jobs/cld/gen/gen_qq.sh index 175a8a89c..40c5f1aae 100755 --- a/snakemake_jobs/cld/gen/gen_qq.sh +++ b/snakemake_jobs/cld/gen/gen_qq.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_qq_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_qq_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_qq_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_ttbar.sh b/snakemake_jobs/cld/gen/gen_ttbar.sh index 878683404..7e1651eb2 100755 --- a/snakemake_jobs/cld/gen/gen_ttbar.sh +++ b/snakemake_jobs/cld/gen/gen_ttbar.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ttbar_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ttbar_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ttbar_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_ww_fullhad.sh b/snakemake_jobs/cld/gen/gen_ww_fullhad.sh index b57ffb14a..3f63a8bfb 100755 --- a/snakemake_jobs/cld/gen/gen_ww_fullhad.sh +++ b/snakemake_jobs/cld/gen/gen_ww_fullhad.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_WW_fullhad_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_WW_fullhad_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_WW_fullhad_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_z_qq.sh b/snakemake_jobs/cld/gen/gen_z_qq.sh index 40cfe721c..22e9d7dad 100755 --- a/snakemake_jobs/cld/gen/gen_z_qq.sh +++ b/snakemake_jobs/cld/gen/gen_z_qq.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_qq_ecm91/root/reco_p8_ee_Z_qq_ecm91_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_qq_ecm91/root/reco_p8_ee_Z_qq_ecm91_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_Z_qq_ecm91_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_Z_qq_ecm91_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_Z_qq_ecm91 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_qq_ecm91/root/reco_p8_ee_Z_qq_ecm91_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_z_tautau.sh b/snakemake_jobs/cld/gen/gen_z_tautau.sh index 163ec9d38..f494a3a59 100755 --- a/snakemake_jobs/cld/gen/gen_z_tautau.sh +++ b/snakemake_jobs/cld/gen/gen_z_tautau.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_tautau_ecm91/root/reco_p8_ee_Z_tautau_ecm91_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_tautau_ecm91/root/reco_p8_ee_Z_tautau_ecm91_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_Z_tautau_ecm91_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_Z_tautau_ecm91_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_Z_tautau_ecm91 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_Z_tautau_ecm91/root/reco_p8_ee_Z_tautau_ecm91_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_zh_tautau_240.sh b/snakemake_jobs/cld/gen/gen_zh_tautau_240.sh index 1001b63f2..09b9c685a 100755 --- a/snakemake_jobs/cld/gen/gen_zh_tautau_240.sh +++ b/snakemake_jobs/cld/gen/gen_zh_tautau_240.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm240/root/reco_p8_ee_ZH_Htautau_ecm240_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm240/root/reco_p8_ee_ZH_Htautau_ecm240_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZH_Htautau_ecm240_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZH_Htautau_ecm240_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ZH_Htautau_ecm240 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm240/root/reco_p8_ee_ZH_Htautau_ecm240_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_zh_tautau_365.sh b/snakemake_jobs/cld/gen/gen_zh_tautau_365.sh index c949d5196..d583e2a43 100755 --- a/snakemake_jobs/cld/gen/gen_zh_tautau_365.sh +++ b/snakemake_jobs/cld/gen/gen_zh_tautau_365.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm365/root/reco_p8_ee_ZH_Htautau_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm365/root/reco_p8_ee_ZH_Htautau_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZH_Htautau_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZH_Htautau_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ZH_Htautau_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZH_Htautau_ecm365/root/reco_p8_ee_ZH_Htautau_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_zz.sh b/snakemake_jobs/cld/gen/gen_zz.sh index 7729158a2..c76e0a7c6 100755 --- a/snakemake_jobs/cld/gen/gen_zz.sh +++ b/snakemake_jobs/cld/gen/gen_zz.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_ecm365/root/reco_p8_ee_ZZ_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_ecm365/root/reco_p8_ee_ZZ_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ZZ_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_ecm365/root/reco_p8_ee_ZZ_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_zz_tautau_240.sh b/snakemake_jobs/cld/gen/gen_zz_tautau_240.sh index 381163630..31f26a9b3 100755 --- a/snakemake_jobs/cld/gen/gen_zz_tautau_240.sh +++ b/snakemake_jobs/cld/gen/gen_zz_tautau_240.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm240/root/reco_p8_ee_ZZ_tautau_ecm240_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm240/root/reco_p8_ee_ZZ_tautau_ecm240_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_tautau_ecm240_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_tautau_ecm240_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ZZ_tautau_ecm240 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm240/root/reco_p8_ee_ZZ_tautau_ecm240_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/gen/gen_zz_tautau_365.sh b/snakemake_jobs/cld/gen/gen_zz_tautau_365.sh index 20e03ba15..055d22cb3 100755 --- a/snakemake_jobs/cld/gen/gen_zz_tautau_365.sh +++ b/snakemake_jobs/cld/gen/gen_zz_tautau_365.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm365/root/reco_p8_ee_ZZ_tautau_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm365/root/reco_p8_ee_ZZ_tautau_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_tautau_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 + export OUTDIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/cld/CLDConfig && export WORKDIR="/scratch/local/joosep/p8_ee_ZZ_tautau_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 bash mlpf/data/key4hep/gen/cld/run_sim.sh p8_ee_ZZ_tautau_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/gen/p8_ee_ZZ_tautau_ecm365/root/reco_p8_ee_ZZ_tautau_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/cld/post/post_gun_e_10gev.sh b/snakemake_jobs/cld/post/post_gun_e_10gev.sh index a1ba9b882..635cf6879 100755 --- a/snakemake_jobs/cld/post/post_gun_e_10gev.sh +++ b/snakemake_jobs/cld/post/post_gun_e_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_gun_mu_10gev.sh b/snakemake_jobs/cld/post/post_gun_mu_10gev.sh index 48c624bd6..7d9dce6c2 100755 --- a/snakemake_jobs/cld/post/post_gun_mu_10gev.sh +++ b/snakemake_jobs/cld/post/post_gun_mu_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_gun_pi_10gev.sh b/snakemake_jobs/cld/post/post_gun_pi_10gev.sh index 80807028e..94f4efb12 100755 --- a/snakemake_jobs/cld/post/post_gun_pi_10gev.sh +++ b/snakemake_jobs/cld/post/post_gun_pi_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_qq.sh b/snakemake_jobs/cld/post/post_qq.sh index 0a3ceea15..eb2242df7 100755 --- a/snakemake_jobs/cld/post/post_qq.sh +++ b/snakemake_jobs/cld/post/post_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_ttbar.sh b/snakemake_jobs/cld/post/post_ttbar.sh index 894f8a56f..ee9398b77 100755 --- a/snakemake_jobs/cld/post/post_ttbar.sh +++ b/snakemake_jobs/cld/post/post_ttbar.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_ww_fullhad.sh b/snakemake_jobs/cld/post/post_ww_fullhad.sh index a1d0d5f70..1bb7c2a35 100755 --- a/snakemake_jobs/cld/post/post_ww_fullhad.sh +++ b/snakemake_jobs/cld/post/post_ww_fullhad.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_z_qq.sh b/snakemake_jobs/cld/post/post_z_qq.sh index 93703576b..a4f620f84 100755 --- a/snakemake_jobs/cld/post/post_z_qq.sh +++ b/snakemake_jobs/cld/post/post_z_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_z_tautau.sh b/snakemake_jobs/cld/post/post_z_tautau.sh index ff4815a56..c04f1e278 100755 --- a/snakemake_jobs/cld/post/post_z_tautau.sh +++ b/snakemake_jobs/cld/post/post_z_tautau.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_zh_tautau_240.sh b/snakemake_jobs/cld/post/post_zh_tautau_240.sh index 3d6e11e0e..9160fa778 100755 --- a/snakemake_jobs/cld/post/post_zh_tautau_240.sh +++ b/snakemake_jobs/cld/post/post_zh_tautau_240.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_zh_tautau_365.sh b/snakemake_jobs/cld/post/post_zh_tautau_365.sh index 84be8b6aa..fd72fcf2e 100755 --- a/snakemake_jobs/cld/post/post_zh_tautau_365.sh +++ b/snakemake_jobs/cld/post/post_zh_tautau_365.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_zz.sh b/snakemake_jobs/cld/post/post_zz.sh index 8e0a84881..65160227f 100755 --- a/snakemake_jobs/cld/post/post_zz.sh +++ b/snakemake_jobs/cld/post/post_zz.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_zz_tautau_240.sh b/snakemake_jobs/cld/post/post_zz_tautau_240.sh index bd5712d9e..e80b09fce 100755 --- a/snakemake_jobs/cld/post/post_zz_tautau_240.sh +++ b/snakemake_jobs/cld/post/post_zz_tautau_240.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/post/post_zz_tautau_365.sh b/snakemake_jobs/cld/post/post_zz_tautau_365.sh index 14a273dd0..b1643a6ed 100755 --- a/snakemake_jobs/cld/post/post_zz_tautau_365.sh +++ b/snakemake_jobs/cld/post/post_zz_tautau_365.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/cld/tfds/tfds_qq.sh b/snakemake_jobs/cld/tfds/tfds_qq.sh index 3fba9fedc..fd54ee0e9 100755 --- a/snakemake_jobs/cld/tfds/tfds_qq.sh +++ b/snakemake_jobs/cld/tfds/tfds_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=qq_tfds_$config_id diff --git a/snakemake_jobs/cld/tfds/tfds_ttbar.sh b/snakemake_jobs/cld/tfds/tfds_ttbar.sh index f582fc54f..783f29602 100755 --- a/snakemake_jobs/cld/tfds/tfds_ttbar.sh +++ b/snakemake_jobs/cld/tfds/tfds_ttbar.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ttbar_tfds_$config_id diff --git a/snakemake_jobs/cld/tfds/tfds_ww_fullhad.sh b/snakemake_jobs/cld/tfds/tfds_ww_fullhad.sh index 05a485d19..139456666 100755 --- a/snakemake_jobs/cld/tfds/tfds_ww_fullhad.sh +++ b/snakemake_jobs/cld/tfds/tfds_ww_fullhad.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ww_fullhad_tfds_$config_id diff --git a/snakemake_jobs/cld/tfds/tfds_zz.sh b/snakemake_jobs/cld/tfds/tfds_zz.sh index 8ecfbb58f..daf283183 100755 --- a/snakemake_jobs/cld/tfds/tfds_zz.sh +++ b/snakemake_jobs/cld/tfds/tfds_zz.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=zz_tfds_$config_id diff --git a/snakemake_jobs/cld/tfds_hit/tfds_hit_qq.sh b/snakemake_jobs/cld/tfds_hit/tfds_hit_qq.sh index 92846b549..6796d8010 100755 --- a/snakemake_jobs/cld/tfds_hit/tfds_hit_qq.sh +++ b/snakemake_jobs/cld/tfds_hit/tfds_hit_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=qq_tfds_hit_$config_id diff --git a/snakemake_jobs/cld/tfds_hit/tfds_hit_ttbar.sh b/snakemake_jobs/cld/tfds_hit/tfds_hit_ttbar.sh index 9ec9c2d26..f6642a2ba 100755 --- a/snakemake_jobs/cld/tfds_hit/tfds_hit_ttbar.sh +++ b/snakemake_jobs/cld/tfds_hit/tfds_hit_ttbar.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ttbar_tfds_hit_$config_id diff --git a/snakemake_jobs/cld/tfds_hit/tfds_hit_ww_fullhad.sh b/snakemake_jobs/cld/tfds_hit/tfds_hit_ww_fullhad.sh index 3e88dd0dd..39465bd1e 100755 --- a/snakemake_jobs/cld/tfds_hit/tfds_hit_ww_fullhad.sh +++ b/snakemake_jobs/cld/tfds_hit/tfds_hit_ww_fullhad.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ww_fullhad_tfds_hit_$config_id diff --git a/snakemake_jobs/cld/tfds_hit/tfds_hit_zz.sh b/snakemake_jobs/cld/tfds_hit/tfds_hit_zz.sh index 8550b1a66..36dd6c4fd 100755 --- a/snakemake_jobs/cld/tfds_hit/tfds_hit_zz.sh +++ b/snakemake_jobs/cld/tfds_hit/tfds_hit_zz.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=zz_tfds_hit_$config_id diff --git a/snakemake_jobs/cld/train/train_pyg-cld-hits-v1_cld.sh b/snakemake_jobs/cld/train/train_pyg-cld-hits-v1_cld.sh index ae624881f..e7dce70d8 100755 --- a/snakemake_jobs/cld/train/train_pyg-cld-hits-v1_cld.sh +++ b/snakemake_jobs/cld/train/train_pyg-cld-hits-v1_cld.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH export TFDS_DATA_DIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds diff --git a/snakemake_jobs/cld/train/train_pyg-cld-v1_cld.sh b/snakemake_jobs/cld/train/train_pyg-cld-v1_cld.sh index ded60f14a..b7db278e2 100755 --- a/snakemake_jobs/cld/train/train_pyg-cld-v1_cld.sh +++ b/snakemake_jobs/cld/train/train_pyg-cld-v1_cld.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH export TFDS_DATA_DIR=/local/joosep/mlpf/cld/v1.2.5_key4hep_2025-05-29/tfds diff --git a/snakemake_jobs/idea/gen/gen_gun_e_10gev.sh b/snakemake_jobs/idea/gen/gen_gun_e_10gev.sh index f1915f3d4..9f9cc4d95 100755 --- a/snakemake_jobs/idea/gen/gen_gun_e_10gev.sh +++ b/snakemake_jobs/idea/gen/gen_gun_e_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_e_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_e_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh gun_e_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_e_10gev/root/reco_gun_e_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/gen/gen_gun_mu_10gev.sh b/snakemake_jobs/idea/gen/gen_gun_mu_10gev.sh index f603fbaa0..6caf63510 100755 --- a/snakemake_jobs/idea/gen/gen_gun_mu_10gev.sh +++ b/snakemake_jobs/idea/gen/gen_gun_mu_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_mu_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_mu_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh gun_mu_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_mu_10gev/root/reco_gun_mu_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/gen/gen_gun_pi_10gev.sh b/snakemake_jobs/idea/gen/gen_gun_pi_10gev.sh index a85c0c7d5..11861e546 100755 --- a/snakemake_jobs/idea/gen/gen_gun_pi_10gev.sh +++ b/snakemake_jobs/idea/gen/gen_gun_pi_10gev.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_pi_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/gun_pi_10gev_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh gun_pi_10gev $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/gun_pi_10gev/root/reco_gun_pi_10gev_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/gen/gen_qq.sh b/snakemake_jobs/idea/gen/gen_qq.sh index 6c7b797ec..6ae364fc5 100755 --- a/snakemake_jobs/idea/gen/gen_qq.sh +++ b/snakemake_jobs/idea/gen/gen_qq.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_qq_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_qq_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh p8_ee_qq_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_qq_ecm365/root/reco_p8_ee_qq_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/gen/gen_ttbar.sh b/snakemake_jobs/idea/gen/gen_ttbar.sh index 27d2bbe39..4ad3a4114 100755 --- a/snakemake_jobs/idea/gen/gen_ttbar.sh +++ b/snakemake_jobs/idea/gen/gen_ttbar.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_ttbar_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_ttbar_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh p8_ee_ttbar_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_ttbar_ecm365/root/reco_p8_ee_ttbar_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/gen/gen_ww_fullhad.sh b/snakemake_jobs/idea/gen/gen_ww_fullhad.sh index 1b66d1114..dc729183e 100755 --- a/snakemake_jobs/idea/gen/gen_ww_fullhad.sh +++ b/snakemake_jobs/idea/gen/gen_ww_fullhad.sh @@ -11,14 +11,14 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow start_seed=$1 for (( i=0; i<1; i++ )); do seed=$((start_seed + i)) if [ ! -f /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root ]; then echo "Generating /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root" - export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow-dev/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_WW_fullhad_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 + export OUTDIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/ && export CONFIG_DIR=/home/joosep/particleflow/mlpf/data/key4hep/gen/idea && export WORKDIR="/scratch/local/joosep/p8_ee_WW_fullhad_ecm365_${seed}_${SLURM_JOB_ID:-manual-$$}" && export NEV=100 && export PROGRESS_INTERVAL=60 bash mlpf/data/key4hep/gen/idea/run_sim.sh p8_ee_WW_fullhad_ecm365 $seed nopu else echo "Skipping /local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/gen/p8_ee_WW_fullhad_ecm365/root/reco_p8_ee_WW_fullhad_ecm365_${seed}.root, already exists" diff --git a/snakemake_jobs/idea/post/post_gun_e_10gev.sh b/snakemake_jobs/idea/post/post_gun_e_10gev.sh index 3315b4662..e42850fbc 100755 --- a/snakemake_jobs/idea/post/post_gun_e_10gev.sh +++ b/snakemake_jobs/idea/post/post_gun_e_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/post/post_gun_mu_10gev.sh b/snakemake_jobs/idea/post/post_gun_mu_10gev.sh index b9954a056..6e8b85f19 100755 --- a/snakemake_jobs/idea/post/post_gun_mu_10gev.sh +++ b/snakemake_jobs/idea/post/post_gun_mu_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/post/post_gun_pi_10gev.sh b/snakemake_jobs/idea/post/post_gun_pi_10gev.sh index 422514859..426794f95 100755 --- a/snakemake_jobs/idea/post/post_gun_pi_10gev.sh +++ b/snakemake_jobs/idea/post/post_gun_pi_10gev.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/post/post_qq.sh b/snakemake_jobs/idea/post/post_qq.sh index 132a1fdb2..2d97cbab3 100755 --- a/snakemake_jobs/idea/post/post_qq.sh +++ b/snakemake_jobs/idea/post/post_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/post/post_ttbar.sh b/snakemake_jobs/idea/post/post_ttbar.sh index 2c9751c60..608ff2e26 100755 --- a/snakemake_jobs/idea/post/post_ttbar.sh +++ b/snakemake_jobs/idea/post/post_ttbar.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/post/post_ww_fullhad.sh b/snakemake_jobs/idea/post/post_ww_fullhad.sh index 8ae75e7a4..4d6596388 100755 --- a/snakemake_jobs/idea/post/post_ww_fullhad.sh +++ b/snakemake_jobs/idea/post/post_ww_fullhad.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH start_seed=$1 diff --git a/snakemake_jobs/idea/tfds/tfds_qq.sh b/snakemake_jobs/idea/tfds/tfds_qq.sh index 87429335a..ed93916e6 100755 --- a/snakemake_jobs/idea/tfds/tfds_qq.sh +++ b/snakemake_jobs/idea/tfds/tfds_qq.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=qq_tfds_$config_id diff --git a/snakemake_jobs/idea/tfds/tfds_ttbar.sh b/snakemake_jobs/idea/tfds/tfds_ttbar.sh index 236bc33bd..6cb8e2a7e 100755 --- a/snakemake_jobs/idea/tfds/tfds_ttbar.sh +++ b/snakemake_jobs/idea/tfds/tfds_ttbar.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ttbar_tfds_$config_id diff --git a/snakemake_jobs/idea/tfds/tfds_ww_fullhad.sh b/snakemake_jobs/idea/tfds/tfds_ww_fullhad.sh index c11e3a4ee..338f5a533 100755 --- a/snakemake_jobs/idea/tfds/tfds_ww_fullhad.sh +++ b/snakemake_jobs/idea/tfds/tfds_ww_fullhad.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow config_id=$1 tfds_id=ww_fullhad_tfds_$config_id diff --git a/snakemake_jobs/idea/train/train_pyg-idea-pipeline-v1_idea.sh b/snakemake_jobs/idea/train/train_pyg-idea-pipeline-v1_idea.sh index dba42c6af..abc387a12 100755 --- a/snakemake_jobs/idea/train/train_pyg-idea-pipeline-v1_idea.sh +++ b/snakemake_jobs/idea/train/train_pyg-idea-pipeline-v1_idea.sh @@ -11,7 +11,7 @@ export TEMPDIR=/scratch/local/joosep/tmp export TEMP=/scratch/local/joosep/tmp export TMP=/scratch/local/joosep/tmp mkdir -p $TMPDIR -cd /home/joosep/particleflow-dev +cd /home/joosep/particleflow export PYTHONPATH=$(pwd):$PYTHONPATH export TFDS_DATA_DIR=/local/joosep/mlpf/idea/IDEA_o1_v03_fccconfig_a05a3a9/tfds From 59a7613b6f7481703c03c7e9a51da450e19bc1d2 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 11 Sep 2026 21:39:12 +0300 Subject: [PATCH 26/29] fix shutdown noise --- mlpf/model/training.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mlpf/model/training.py b/mlpf/model/training.py index 06d6585b4..a66539422 100644 --- a/mlpf/model/training.py +++ b/mlpf/model/training.py @@ -1447,10 +1447,12 @@ def run_test(rank, world_size, config: MLPFConfig, outdir, model, sample, testdi worker_kwargs = {} if config.num_workers > 0: + # This loader is consumed once and then discarded. Persistent workers + # would only defer their shutdown to DataLoader destruction, which can + # race with multiprocessing queue cleanup in multi-GPU runs. worker_kwargs = { "prefetch_factor": config.prefetch_factor, "worker_init_fn": set_worker_sharing_strategy, - "persistent_workers": True, } test_loader = torch.utils.data.DataLoader( From ef8329de4a4d7c415ca0ef63c8e179b7194f952b Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Mon, 14 Sep 2026 10:40:54 +0300 Subject: [PATCH 27/29] Add continuation for unfinished scenario jobs --- mlpf/training_scenarios.py | 220 +++++++++++++++++++++++----- mlpf/training_submission.py | 67 ++++++++- scripts/flatiron/run_uv_scenario.sh | 8 + scripts/lumi/run_scenario.sh | 8 + scripts/tallinn/run_scenario.sh | 8 + tests/test_training_scenarios.py | 120 +++++++++++++-- tests/test_training_submission.py | 58 ++++++++ 7 files changed, 432 insertions(+), 57 deletions(-) diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index fb4ea69fc..6078b1adf 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -5,6 +5,7 @@ import datetime import json import os +import re import shlex import subprocess from pathlib import Path @@ -42,7 +43,9 @@ class ScenarioVariant(BaseModel): def reject_derived_overrides(self): invalid = DERIVED_KEYS.intersection(self.overrides) if invalid: - raise ValueError(f"Variant overrides must not set derived keys: {sorted(invalid)}") + raise ValueError( + f"Variant overrides must not set derived keys: {sorted(invalid)}" + ) return self @@ -56,7 +59,9 @@ class ScenarioTraining(BaseModel): def reject_derived_parameters(self): invalid = DERIVED_KEYS.intersection(self.parameters) if invalid: - raise ValueError(f"Scenario parameters must not set derived keys: {sorted(invalid)}") + raise ValueError( + f"Scenario parameters must not set derived keys: {sorted(invalid)}" + ) return self @@ -70,7 +75,9 @@ class TrainingScenario(BaseModel): seeds: list[int] = Field(min_length=1) training: ScenarioTraining common_overrides: dict[str, Any] = Field(default_factory=dict) - allowed_variant_differences: list[str] = Field(default_factory=lambda: ["model.output_mode", "model.set_decoder"]) + allowed_variant_differences: list[str] = Field( + default_factory=lambda: ["model.output_mode", "model.set_decoder"] + ) @model_validator(mode="after") def validate_scenario(self): @@ -82,7 +89,9 @@ def validate_scenario(self): raise ValueError("Scenario seeds must be non-negative") invalid = DERIVED_KEYS.intersection(self.common_overrides) if invalid: - raise ValueError(f"Common overrides must not set derived keys: {sorted(invalid)}") + raise ValueError( + f"Common overrides must not set derived keys: {sorted(invalid)}" + ) return self def variant_production(self, variant_name): @@ -131,9 +140,14 @@ class PlatformProfile(BaseModel): def validate_runtime_overrides(self): invalid = set(self.runtime_overrides).difference(PLATFORM_OVERRIDE_KEYS) if invalid: - raise ValueError("Platform profiles may only set runtime-specific overrides; " f"invalid keys: {sorted(invalid)}") + raise ValueError( + "Platform profiles may only set runtime-specific overrides; " + f"invalid keys: {sorted(invalid)}" + ) if isinstance(self.data_dir, dict) and not self.data_dir: - raise ValueError("Platform data_dir mapping must name at least one production") + raise ValueError( + "Platform data_dir mapping must name at least one production" + ) return self def data_dir_for(self, production_name): @@ -141,7 +155,8 @@ def data_dir_for(self, production_name): return self.data_dir if production_name not in self.data_dir: raise ValueError( - f"Platform profile {self.name!r} has no data_dir for production {production_name!r}; " f"available: {sorted(self.data_dir)}" + f"Platform profile {self.name!r} has no data_dir for production {production_name!r}; " + f"available: {sorted(self.data_dir)}" ) return self.data_dir[production_name] @@ -163,6 +178,14 @@ class ResolvedScenarioJob(BaseModel): resolved_config: MLPFConfig +class ScenarioContinuation(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + experiment_dir: Path + checkpoint: Path | None + step: int = 0 + + def _read_yaml(path): with Path(path).open() as handle: return yaml.safe_load(handle) @@ -177,9 +200,17 @@ def load_platform_profile(path): if isinstance(profile.data_dir, str): profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) else: - profile.data_dir = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.data_dir.items()} - profile.experiments_dir = os.path.expandvars(os.path.expanduser(profile.experiments_dir)) - profile.environment = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.environment.items()} + profile.data_dir = { + key: os.path.expandvars(os.path.expanduser(value)) + for key, value in profile.data_dir.items() + } + profile.experiments_dir = os.path.expandvars( + os.path.expanduser(profile.experiments_dir) + ) + profile.environment = { + key: os.path.expandvars(os.path.expanduser(value)) + for key, value in profile.environment.items() + } return profile @@ -238,7 +269,8 @@ def _training_batch_size(config): batch_sizes = {dataset.batch_size for dataset in physical_datasets} if len(batch_sizes) != 1: raise ValueError( - "Automatic global-batch resolution requires every physical training dataset " f"to use the same batch size, got {sorted(batch_sizes)}" + "Automatic global-batch resolution requires every physical training dataset " + f"to use the same batch size, got {sorted(batch_sizes)}" ) return next(iter(batch_sizes)) @@ -264,12 +296,16 @@ def resolve_scenario_job( extra_overrides=None, ): if variant_name not in scenario.variants: - raise ValueError(f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}") + raise ValueError( + f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}" + ) variant = scenario.variants[variant_name] extra_overrides = extra_overrides or {} invalid = DERIVED_KEYS.intersection(extra_overrides) if invalid: - raise ValueError(f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}") + raise ValueError( + f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}" + ) settings = _merge_settings(scenario, platform, variant, extra_overrides) settings["seed"] = seed selected_spec = str(spec_file or scenario.spec_file) @@ -285,14 +321,19 @@ def resolve_scenario_job( extra_args=_settings_as_extra_args(settings), ) - target_global_batch = global_batch_size if global_batch_size is not None else scenario.training.global_batch_size + target_global_batch = ( + global_batch_size + if global_batch_size is not None + else scenario.training.global_batch_size + ) if target_global_batch <= 0: raise ValueError("global_batch_size must be positive") dataset_batch_size = _training_batch_size(config) divisor = platform.gpus * dataset_batch_size if target_global_batch % divisor: raise ValueError( - f"global_batch_size={target_global_batch} is not divisible by " f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" + f"global_batch_size={target_global_batch} is not divisible by " + f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" ) multiplier = target_global_batch // divisor settings["gpu_batch_multiplier"] = multiplier @@ -334,7 +375,9 @@ def _flatten(value, prefix=""): def _difference_allowed(path, allowed_paths): - return any(path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths) + return any( + path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths + ) def validate_variant_invariants(jobs, allowed_paths): @@ -346,11 +389,17 @@ def validate_variant_invariants(jobs, allowed_paths): differences = { path: (reference.get(path), candidate.get(path)) for path in sorted(set(reference) | set(candidate)) - if reference.get(path) != candidate.get(path) and not _difference_allowed(path, allowed_paths) + if reference.get(path) != candidate.get(path) + and not _difference_allowed(path, allowed_paths) } if differences: - details = ", ".join(f"{path}: {values[0]!r} != {values[1]!r}" for path, values in differences.items()) - raise ValueError(f"Scenario variants differ outside allowed fields: {details}") + details = ", ".join( + f"{path}: {values[0]!r} != {values[1]!r}" + for path, values in differences.items() + ) + raise ValueError( + f"Scenario variants differ outside allowed fields: {details}" + ) def resolve_scenario_jobs( @@ -436,23 +485,101 @@ def _experiment_path(platform, job, timestamp=None): return Path(platform.experiments_dir) / job.scenario_name / experiment_name -def run_scenario_job(job, scenario, platform, spec_file, *, dry_run=False): - experiment_dir = _experiment_path(platform, job, timestamp="TIMESTAMP" if dry_run else None) +def _checkpoint_step(path): + match = re.fullmatch(r"checkpoint-(\d+)(?:-.*)?\.pth", Path(path).name) + return int(match.group(1)) if match else None + + +def find_scenario_continuation(job, platform): + """Find the most advanced compatible run for one scenario job.""" + scenario_dir = Path(platform.experiments_dir) / job.scenario_name + desired_config = job.resolved_config.model_dump(mode="json") + candidates = [] + for manifest_path in scenario_dir.glob("*/scenario-manifest.json"): + try: + with manifest_path.open() as handle: + manifest = json.load(handle) + except (OSError, json.JSONDecodeError): + continue + saved_job = manifest.get("job", {}) + if ( + saved_job.get("scenario_name") != job.scenario_name + or saved_job.get("platform_name") != job.platform_name + or saved_job.get("variant_name") != job.variant_name + or saved_job.get("seed") != job.seed + or manifest.get("resolved_config") != desired_config + ): + continue + + experiment_dir = manifest_path.parent + checkpoints = [] + for checkpoint_path in (experiment_dir / "checkpoints").glob( + "checkpoint-*.pth" + ): + step = _checkpoint_step(checkpoint_path) + if step is not None: + checkpoints.append((step, checkpoint_path)) + step, checkpoint = max(checkpoints, default=(0, None), key=lambda item: item[0]) + candidates.append( + (step, manifest_path.stat().st_mtime, experiment_dir, checkpoint) + ) + + if not candidates: + return None + step, _, experiment_dir, checkpoint = max( + candidates, key=lambda item: (item[0], item[1]) + ) + return ScenarioContinuation( + experiment_dir=experiment_dir, checkpoint=checkpoint, step=step + ) + + +def run_scenario_job( + job, scenario, platform, spec_file, *, dry_run=False, continue_run=False +): + continuation = find_scenario_continuation(job, platform) if continue_run else None + if continue_run and continuation is None: + raise ValueError( + f"No compatible prior run found for {job.variant_name} seed {job.seed}" + ) + experiment_dir = ( + continuation.experiment_dir + if continuation is not None + else _experiment_path(platform, job, timestamp="TIMESTAMP" if dry_run else None) + ) + if continuation is not None and continuation.checkpoint is not None: + job.settings["load"] = str(continuation.checkpoint) command = _pipeline_command(job, scenario, platform, spec_file, experiment_dir) print(shlex.join(command), flush=True) if dry_run: return - experiment_dir.mkdir(parents=True, exist_ok=False) - manifest = { - "scenario": scenario.model_dump(mode="json"), - "platform": platform.model_dump(mode="json"), - "job": job.model_dump(mode="json", exclude={"resolved_config"}), - "resolved_config": job.resolved_config.model_dump(mode="json"), - "command": command, - "git_revision": _git_revision(), - } - with (experiment_dir / "scenario-manifest.json").open("w") as handle: + manifest_path = experiment_dir / "scenario-manifest.json" + if continuation is None: + experiment_dir.mkdir(parents=True, exist_ok=False) + manifest = { + "scenario": scenario.model_dump(mode="json"), + "platform": platform.model_dump(mode="json"), + "job": job.model_dump(mode="json", exclude={"resolved_config"}), + "resolved_config": job.resolved_config.model_dump(mode="json"), + "command": command, + "git_revision": _git_revision(), + } + else: + with manifest_path.open() as handle: + manifest = json.load(handle) + manifest.setdefault("continuations", []).append( + { + "checkpoint": str(continuation.checkpoint) + if continuation.checkpoint + else None, + "step": continuation.step, + "command": command, + "git_revision": _git_revision(), + "submitted_at": datetime.datetime.now().isoformat(), + } + ) + with manifest_path.open("w") as handle: json.dump(manifest, handle, indent=2) environment = os.environ.copy() @@ -473,7 +600,9 @@ def _parse_set_overrides(values): def _validate_slurm_allocation(platform): allocated = os.environ.get("SLURM_GPUS_PER_NODE") if allocated and allocated.isdigit() and int(allocated) != platform.gpus: - raise ValueError(f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}") + raise ValueError( + f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}" + ) def main(argv=None): @@ -490,6 +619,12 @@ def main(argv=None): parser.add_argument("--experiments-dir") parser.add_argument("--set", action="append", default=[], metavar="KEY=VALUE") parser.add_argument("--dry-run", action="store_true") + parser.add_argument( + "--continue", + dest="continue_run", + action="store_true", + help="Resume this unfinished scenario job", + ) args = parser.parse_args(argv) scenario = load_training_scenario(args.scenario) @@ -524,7 +659,26 @@ def main(argv=None): raise ValueError("No jobs matched the requested variant") for job in jobs: - run_scenario_job(job, scenario, platform, spec_file, dry_run=args.dry_run) + if args.continue_run: + continuation = find_scenario_continuation(job, platform) + if continuation is None: + raise ValueError( + f"No compatible prior run found for {job.variant_name} seed {job.seed}" + ) + target_step = job.resolved_config.num_steps + if continuation.step >= target_step: + print( + f"Skipping completed job {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step}" + ) + continue + run_scenario_job( + job, + scenario, + platform, + spec_file, + dry_run=args.dry_run, + continue_run=args.continue_run, + ) if __name__ == "__main__": diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py index 3e131238c..475ab45ad 100644 --- a/mlpf/training_submission.py +++ b/mlpf/training_submission.py @@ -6,6 +6,7 @@ from pathlib import Path from mlpf.training_scenarios import ( + find_scenario_continuation, load_platform_profile, load_training_scenario, resolve_scenario_jobs, @@ -44,8 +45,13 @@ def resolve_flatiron_profile_path(reference, repo_root): def available_choices(repo_root, site="flatiron"): - scenarios = sorted(path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml")) - accelerators = sorted(path.stem.removeprefix(f"{site}_") for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml")) + scenarios = sorted( + path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml") + ) + accelerators = sorted( + path.stem.removeprefix(f"{site}_") + for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml") + ) return scenarios, accelerators @@ -79,6 +85,7 @@ def build_slurm_submission( *, seed=None, worker=None, + continue_run=False, ): scenario = load_training_scenario(scenario_path) if seed is not None: @@ -87,7 +94,9 @@ def build_slurm_submission( scenario.seeds = [seed] profile = load_platform_profile(profile_path) if profile.slurm is None: - raise ValueError(f"Platform profile {profile.name!r} has no Slurm configuration") + raise ValueError( + f"Platform profile {profile.name!r} has no Slurm configuration" + ) spec_file = Path(scenario.spec_file) if not spec_file.is_absolute(): @@ -96,9 +105,37 @@ def build_slurm_submission( if not jobs: raise ValueError("Scenario did not resolve to any jobs") + selected_indices = list(range(len(jobs))) + if continue_run: + selected_indices = [] + for index, job in enumerate(jobs): + continuation = find_scenario_continuation(job, profile) + if continuation is None: + print( + f"Skipping {job.variant_name} seed {job.seed}: no compatible prior run" + ) + continue + target_step = job.resolved_config.num_steps + if continuation.step >= target_step: + print( + f"Skipping {job.variant_name} seed {job.seed}: complete at step {continuation.step}/{target_step}" + ) + continue + source = ( + continuation.checkpoint.name if continuation.checkpoint else "start" + ) + print( + f"Continuing {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step} from {source}" + ) + selected_indices.append(index) + if not selected_indices: + raise ValueError("No unfinished compatible scenario jobs found") + slurm = profile.slurm logs_dir = repo_root / "logs_slurm" - worker = Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") + worker = ( + Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") + ) command = [ "sbatch", "--time", @@ -126,7 +163,9 @@ def build_slurm_submission( command.extend( [ "--array", - f"0-{len(jobs) - 1}", + ",".join(str(index) for index in selected_indices) + if continue_run + else f"0-{len(jobs) - 1}", "--job-name", scenario.name, "--output", @@ -144,7 +183,9 @@ def build_slurm_submission( ) if seed is not None: command.extend(["--seed", str(seed)]) - return command, jobs + if continue_run: + command.append("--continue") + return command, [jobs[index] for index in selected_indices] def main(argv=None, *, site="flatiron"): @@ -161,7 +202,15 @@ def main(argv=None, *, site="flatiron"): action="store_true", help="Print the sbatch command without submitting", ) - parser.add_argument("--list", action="store_true", help="List available scenarios and accelerators") + parser.add_argument( + "--list", action="store_true", help="List available scenarios and accelerators" + ) + parser.add_argument( + "--continue", + dest="continue_run", + action="store_true", + help="Resubmit only compatible jobs whose latest checkpoint is below num_steps", + ) args = parser.parse_args(argv) repo_root = Path(__file__).resolve().parents[1] @@ -181,9 +230,11 @@ def main(argv=None, *, site="flatiron"): repo_root, seed=args.seed, worker=_worker_for_site(repo_root, site), + continue_run=args.continue_run, ) print( - f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" + shlex.join(command), + f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" + + shlex.join(command), flush=True, ) if args.dry_run: diff --git a/scripts/flatiron/run_uv_scenario.sh b/scripts/flatiron/run_uv_scenario.sh index 5ae40bf6f..b92850a0a 100755 --- a/scripts/flatiron/run_uv_scenario.sh +++ b/scripts/flatiron/run_uv_scenario.sh @@ -7,6 +7,7 @@ shift 2 TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} SEED_OVERRIDE=${SEED:-} +CONTINUE_RUN=0 REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} while [[ $# -gt 0 ]]; do case "$1" in @@ -22,6 +23,10 @@ while [[ $# -gt 0 ]]; do REPO_ROOT=${2:?--repo-root requires a value} shift 2 ;; + --continue) + CONTINUE_RUN=1 + shift + ;; *) echo "Unknown argument: $1" >&2 exit 2 @@ -58,5 +63,8 @@ RUN_ARGS=( if [[ -n "$SEED_OVERRIDE" ]]; then RUN_ARGS+=(--seed "$SEED_OVERRIDE") fi +if [[ "$CONTINUE_RUN" == 1 ]]; then + RUN_ARGS+=(--continue) +fi uv run python3 scripts/training/run_scenario.py "${RUN_ARGS[@]}" diff --git a/scripts/lumi/run_scenario.sh b/scripts/lumi/run_scenario.sh index efae3292b..21b71d9f6 100755 --- a/scripts/lumi/run_scenario.sh +++ b/scripts/lumi/run_scenario.sh @@ -7,6 +7,7 @@ shift 2 TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} SEED_OVERRIDE=${SEED:-} +CONTINUE_RUN=0 REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} while [[ $# -gt 0 ]]; do case "$1" in @@ -22,6 +23,10 @@ while [[ $# -gt 0 ]]; do REPO_ROOT=${2:?--repo-root requires a value} shift 2 ;; + --continue) + CONTINUE_RUN=1 + shift + ;; *) echo "Unknown argument: $1" >&2 exit 2 @@ -63,6 +68,9 @@ RUN_ARGS=( if [[ -n "$SEED_OVERRIDE" ]]; then RUN_ARGS+=(--seed "$SEED_OVERRIDE") fi +if [[ "$CONTINUE_RUN" == 1 ]]; then + RUN_ARGS+=(--continue) +fi singularity exec \ -B /scratch/project_465001293 \ diff --git a/scripts/tallinn/run_scenario.sh b/scripts/tallinn/run_scenario.sh index 20c5c78d1..4db803d4a 100755 --- a/scripts/tallinn/run_scenario.sh +++ b/scripts/tallinn/run_scenario.sh @@ -7,6 +7,7 @@ shift 2 TASK_INDEX=${SLURM_ARRAY_TASK_ID:-0} SEED_OVERRIDE=${SEED:-} +CONTINUE_RUN=0 REPO_ROOT=${MLPF_REPO_ROOT:-${SLURM_SUBMIT_DIR:-$PWD}} while [[ $# -gt 0 ]]; do case "$1" in @@ -22,6 +23,10 @@ while [[ $# -gt 0 ]]; do REPO_ROOT=${2:?--repo-root requires a value} shift 2 ;; + --continue) + CONTINUE_RUN=1 + shift + ;; *) echo "Unknown argument: $1" >&2 exit 2 @@ -54,5 +59,8 @@ RUN_ARGS=( if [[ -n "$SEED_OVERRIDE" ]]; then RUN_ARGS+=(--seed "$SEED_OVERRIDE") fi +if [[ "$CONTINUE_RUN" == 1 ]]; then + RUN_ARGS+=(--continue) +fi exec uv run python3 scripts/training/run_scenario.py "${RUN_ARGS[@]}" diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index d3eda3a6d..d8c9efcf9 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -1,3 +1,4 @@ +import json from copy import deepcopy from pathlib import Path @@ -8,6 +9,7 @@ ScenarioVariant, ScenarioTraining, _experiment_path, + find_scenario_continuation, load_platform_profile, load_training_scenario, resolve_scenario_jobs, @@ -17,9 +19,13 @@ ROOT = Path(__file__).resolve().parents[1] SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" -BACKBONE_SCENARIO = ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" +BACKBONE_SCENARIO = ( + ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" +) PF_HITS_SCENARIO = ROOT / "configs/training/scenarios/cld_pf_hits_comparison.yaml" -CLIC_CLD_SCENARIO = ROOT / "configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml" +CLIC_CLD_SCENARIO = ( + ROOT / "configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml" +) PLATFORMS = ROOT / "configs/training/platforms" @@ -56,8 +62,13 @@ def test_backbone_comparison_scenario_keeps_elementwise_output_and_depth_fixed() ) assert [job.variant_name for job in jobs] == ["attention", "heptv2"] - assert [job.resolved_config.model.type.value for job in jobs] == ["attention", "heptv2"] - assert {job.resolved_config.model.output_mode.value for job in jobs} == {"elementwise"} + assert [job.resolved_config.model.type.value for job in jobs] == [ + "attention", + "heptv2", + ] + assert {job.resolved_config.model.output_mode.value for job in jobs} == { + "elementwise" + } assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} assert jobs[1].resolved_config.model.heptv2.block_size == 128 @@ -74,10 +85,24 @@ def test_pf_hits_comparison_scenario_resolves_three_40k_variants(): ) assert [job.variant_name for job in jobs] == ["pf", "elementwise_hits", "set_hits"] - assert [job.model_name for job in jobs] == ["pyg-cld-v1", "pyg-cld-hits-v1", "pyg-cld-hits-set-v1"] - assert [job.resolved_config.dataset.value for job in jobs] == ["cld", "cld_hits", "cld_hits"] - assert [job.resolved_config.model.output_mode.value for job in jobs] == ["elementwise", "elementwise", "set"] - assert [job.resolved_config.model.binary_classification_focal_gamma for job in jobs] == [None, 2.0, 2.0] + assert [job.model_name for job in jobs] == [ + "pyg-cld-v1", + "pyg-cld-hits-v1", + "pyg-cld-hits-set-v1", + ] + assert [job.resolved_config.dataset.value for job in jobs] == [ + "cld", + "cld_hits", + "cld_hits", + ] + assert [job.resolved_config.model.output_mode.value for job in jobs] == [ + "elementwise", + "elementwise", + "set", + ] + assert [ + job.resolved_config.model.binary_classification_focal_gamma for job in jobs + ] == [None, 2.0, 2.0] assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} assert [ ( @@ -106,8 +131,18 @@ def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): global_batch_size=8, ) - assert [job.variant_name for job in jobs] == ["cld_pf", "cld_set_hits", "clic_pf", "clic_set_hits"] - assert [job.model_name for job in jobs] == ["pyg-cld-v1", "pyg-cld-hits-set-v1", "pyg-clic-v1", "pyg-clic-hits-set-v1"] + assert [job.variant_name for job in jobs] == [ + "cld_pf", + "cld_set_hits", + "clic_pf", + "clic_set_hits", + ] + assert [job.model_name for job in jobs] == [ + "pyg-cld-v1", + "pyg-cld-hits-set-v1", + "pyg-clic-v1", + "pyg-clic-hits-set-v1", + ] assert [job.production_name for job in jobs] == ["cld", "cld", "clic", "clic"] assert [job.data_dir for job in jobs] == [ platform.data_dir["cld"], @@ -115,11 +150,28 @@ def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): platform.data_dir["clic"], platform.data_dir["clic"], ] - assert [job.resolved_config.data_dir for job in jobs] == [job.data_dir for job in jobs] - assert [job.resolved_config.dataset.value for job in jobs] == ["cld", "cld_hits", "clic", "clic_hits"] - assert [job.resolved_config.model.output_mode.value for job in jobs] == ["elementwise", "set", "elementwise", "set"] + assert [job.resolved_config.data_dir for job in jobs] == [ + job.data_dir for job in jobs + ] + assert [job.resolved_config.dataset.value for job in jobs] == [ + "cld", + "cld_hits", + "clic", + "clic_hits", + ] + assert [job.resolved_config.model.output_mode.value for job in jobs] == [ + "elementwise", + "set", + "elementwise", + "set", + ] # The set-based hit models run twice the backbone depth of the PF models. - assert [job.resolved_config.model.backbone.num_convs for job in jobs] == [6, 12, 6, 12] + assert [job.resolved_config.model.backbone.num_convs for job in jobs] == [ + 6, + 12, + 6, + 12, + ] assert [ ( job.resolved_config.model.backbone.num_tracker_layers, @@ -128,7 +180,11 @@ def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): ) for job in jobs ] == [(None, None, None), (4, 4, 4), (None, None, None), (4, 4, 4)] - assert {job.resolved_config.model.set_decoder.num_layers for job in jobs if job.resolved_config.model.set_decoder} == {8} + assert { + job.resolved_config.model.set_decoder.num_layers + for job in jobs + if job.resolved_config.model.set_decoder + } == {8} assert {job.resolved_config.num_steps for job in jobs} == {50000} assert {job.resolved_config.val_freq for job in jobs} == {5000} assert {job.resolved_config.lr for job in jobs} == {0.001} @@ -298,4 +354,36 @@ def test_experiments_are_grouped_under_the_scenario_directory(): path = _experiment_path(platform, job, timestamp="TIMESTAMP") - assert path == Path("experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP") + assert path == Path( + "experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP" + ) + + +def test_continuation_chooses_most_advanced_compatible_checkpoint(tmp_path): + scenario = load_training_scenario(SCENARIO) + platform = load_platform_profile(PLATFORMS / "local.yaml") + platform.experiments_dir = str(tmp_path) + job = resolve_scenario_jobs( + scenario, + platform, + spec_file=ROOT / "particleflow_spec.yaml", + global_batch_size=8, + )[0] + + for suffix, step in [("older", 10000), ("newer", 5000), ("best", 15000)]: + run_dir = ( + tmp_path / scenario.name / f"{job.variant_name}_seed{job.seed}_{suffix}" + ) + (run_dir / "checkpoints").mkdir(parents=True) + manifest = { + "job": job.model_dump(mode="json", exclude={"resolved_config"}), + "resolved_config": job.resolved_config.model_dump(mode="json"), + } + (run_dir / "scenario-manifest.json").write_text(json.dumps(manifest)) + (run_dir / "checkpoints" / f"checkpoint-{step}.pth").touch() + + continuation = find_scenario_continuation(job, platform) + + assert continuation.step == 15000 + assert continuation.checkpoint.name == "checkpoint-15000.pth" + assert continuation.experiment_dir.name.endswith("_best") diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index 2fe6a0a04..d8a468334 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -1,5 +1,8 @@ +import json from pathlib import Path +import yaml + from mlpf.training_submission import ( available_choices, build_slurm_submission, @@ -7,6 +10,11 @@ resolve_platform_profile_path, resolve_scenario_path, ) +from mlpf.training_scenarios import ( + load_platform_profile, + load_training_scenario, + resolve_scenario_jobs, +) ROOT = Path(__file__).resolve().parents[1] @@ -109,3 +117,53 @@ def test_picker_discovers_site_specific_accelerators(): assert tallinn_accelerators == ["l40"] assert lumi_accelerators == ["mi250x"] + + +def _write_scenario_run(experiments_dir, scenario, profile, job, step): + run_dir = ( + experiments_dir / scenario.name / f"{job.variant_name}_seed{job.seed}_test" + ) + checkpoint_dir = run_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True) + manifest = { + "job": job.model_dump(mode="json", exclude={"resolved_config"}), + "resolved_config": job.resolved_config.model_dump(mode="json"), + } + (run_dir / "scenario-manifest.json").write_text(json.dumps(manifest)) + (checkpoint_dir / f"checkpoint-{step}.pth").touch() + + +def test_continue_submission_selects_only_unfinished_original_array_indices(tmp_path): + scenario_path = resolve_scenario_path("clic_cld_pf_set_hits_comparison", ROOT) + original_profile_path = resolve_flatiron_profile_path("h100", ROOT) + profile_data = yaml.safe_load(original_profile_path.read_text()) + profile_data["experiments_dir"] = str(tmp_path / "experiments") + profile_path = tmp_path / "flatiron_h100.yaml" + profile_path.write_text(yaml.safe_dump(profile_data)) + + scenario = load_training_scenario(scenario_path) + profile = load_platform_profile(profile_path) + jobs = resolve_scenario_jobs(scenario, profile, spec_file=ROOT / scenario.spec_file) + for index, job in enumerate(jobs): + _write_scenario_run( + Path(profile.experiments_dir), + scenario, + profile, + job, + 40000 if index in {1, 3} else 50000, + ) + + command, selected_jobs = build_slurm_submission( + scenario_path, + profile_path, + ROOT, + worker=ROOT / "scripts/flatiron/run_uv_scenario.sh", + continue_run=True, + ) + + assert [job.variant_name for job in selected_jobs] == [ + "cld_set_hits", + "clic_set_hits", + ] + assert command[command.index("--array") + 1] == "1,3" + assert command[-1] == "--continue" From 1f77e9f0165c56184569553481811934f7849fe1 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Thu, 17 Sep 2026 10:13:04 +0300 Subject: [PATCH 28/29] update visualization script --- scripts/visualize_key4hep.py | 285 +++++++++++++++++++++++++---------- 1 file changed, 207 insertions(+), 78 deletions(-) diff --git a/scripts/visualize_key4hep.py b/scripts/visualize_key4hep.py index 1ab6de8d8..3c81c73b9 100644 --- a/scripts/visualize_key4hep.py +++ b/scripts/visualize_key4hep.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Render perspective Key4HEP event displays from EDM4hep ROOT files. +"""Render consistently oriented Key4HEP event displays from EDM4hep ROOT files. The display overlays reconstructed tracks, calorimeter clusters, detector hits, and stable generator-level particles. Detector type is inferred from the EDM4hep @@ -28,6 +28,12 @@ } +def _production_suffix(root_file: str | Path) -> str | None: + """Return the trailing numerical ROOT-file suffix used as the Pythia seed.""" + suffix = Path(root_file).stem.rsplit("_", 1)[-1] + return suffix if suffix.isdigit() else None + + @dataclass(frozen=True) class DetectorConfig: key: str @@ -46,6 +52,7 @@ class DetectorConfig: cluster_size_scale: float cluster_size_max: float cluster_alpha: float + magnetic_field_tesla: float DETECTORS = { @@ -79,6 +86,40 @@ class DetectorConfig: cluster_size_scale=2.0, cluster_size_max=25.0, cluster_alpha=0.9, + magnetic_field_tesla=2.0, + ), + "clic": DetectorConfig( + key="clic", + title="CLIC", + track_collection="SiTracks_Refitted", + track_label="Reconstructed tracks", + cluster_collection="PandoraClusters", + hit_collections=( + ("VXDTrackerHits", "Tracker hits", "#d62728"), + ("VXDEndcapTrackerHits", "Tracker hits", "#d62728"), + ("ITrackerHits", "Tracker hits", "#d62728"), + ("OTrackerHits", "Tracker hits", "#d62728"), + ("ITrackerEndcapHits", "Tracker hits", "#d62728"), + ("OTrackerEndcapHits", "Tracker hits", "#d62728"), + ("ECALBarrel", "ECAL hits", "#1f77b4"), + ("ECALEndcap", "ECAL hits", "#1f77b4"), + ("ECALOther", "ECAL hits", "#1f77b4"), + ("HCALBarrel", "HCAL hits", "#2ca02c"), + ("HCALEndcap", "HCAL hits", "#2ca02c"), + ("HCALOther", "HCAL hits", "#2ca02c"), + ("MUON", "Muon hits", "#ff7f0e"), + ), + track_radius=1600.0, + track_half_z=2300.0, + particle_barrel_radius=1750.0, + particle_endcap_z=2300.0, + particle_max_length=3300.0, + plot_limit=3600.0, + cluster_size_base=5.0, + cluster_size_scale=2.0, + cluster_size_max=25.0, + cluster_alpha=0.9, + magnetic_field_tesla=4.0, ), "idea": DetectorConfig( key="idea", @@ -106,6 +147,7 @@ class DetectorConfig: cluster_size_scale=0.8, cluster_size_max=10.0, cluster_alpha=0.45, + magnetic_field_tesla=2.0, ), } @@ -118,7 +160,15 @@ def _detector_config(tree, detector: str = "auto") -> DetectorConfig: f"{detector.upper()} collections are not present: expected " f"{config.track_collection} and {config.cluster_collection}" ) return config - matches = [config for config in DETECTORS.values() if config.track_collection in tree and config.cluster_collection in tree] + # CLD and CLIC share the main track and cluster collection names in these + # productions. ECALOther is specific to the CLIC detector model. + if "ECALOther" in tree: + return DETECTORS["clic"] + matches = [ + config + for config in DETECTORS.values() + if config.key != "clic" and config.track_collection in tree and config.cluster_collection in tree + ] if len(matches) != 1: found = ", ".join(config.key for config in matches) or "none" raise ValueError(f"could not infer detector uniquely (matched: {found}); use --detector") @@ -306,8 +356,10 @@ def render_event( max_hits: int = 800, detector: str = "auto", plot_limit: float | None = None, + show_particles: bool = True, + target_only: bool = False, ) -> str: - """Render one CLD or IDEA event as a perspective 3D PNG.""" + """Render one CLD, CLIC, or IDEA event in the transverse x-y plane.""" tree = _open_root(root_file)["events"] if not 0 <= event < tree.num_entries: raise IndexError(f"event {event} is outside [0, {tree.num_entries})") @@ -319,86 +371,144 @@ def render_event( rng = np.random.default_rng(event) def project(x, y, z): - """Fast perspective projection matching a conventional 3D camera.""" - azimuth, elevation = np.deg2rad(36), np.deg2rad(19) - horizontal = np.cos(azimuth) * x - np.sin(azimuth) * y - depth_axis = np.sin(azimuth) * x + np.cos(azimuth) * y - vertical = np.cos(elevation) * z - np.sin(elevation) * depth_axis - depth = np.sin(elevation) * z + np.cos(elevation) * depth_axis - perspective = 1.0 / np.clip(1.0 - depth / 12000.0, 0.55, 1.55) - return horizontal * perspective, vertical * perspective - - for collection, label, color in config.hit_collections: - if collection not in tree: - continue - x = _event(tree, f"{collection}/{collection}.position.x", event) - y = _event(tree, f"{collection}/{collection}.position.y", event) - z = _event(tree, f"{collection}/{collection}.position.z", event) - if len(x) > max_hits: - idx = np.sort(rng.choice(len(x), max_hits, replace=False)) - x, y, z = x[idx], y[idx], z[idx] - sx, sy = project(x, y, z) - ax.scatter(sx, sy, s=4.0, color=color, alpha=0.5, edgecolors="none", rasterized=True, label=label if label not in shown_labels else None) - shown_labels.add(label) - - tx, ty, tz = _track_trajectories(tree, event, config) - sx, sy = project(tx, ty, tz) - ax.plot(sx, sy, color="#ef4444", linewidth=0.8, alpha=0.78, label=config.track_label) + """Project along z with +x left and +y up in every rendered view.""" + del z + return -np.asarray(x), np.asarray(y) - cluster = config.cluster_collection - cx = _event(tree, f"{cluster}/{cluster}.position.x", event) - cy = _event(tree, f"{cluster}/{cluster}.position.y", event) - cz = _event(tree, f"{cluster}/{cluster}.position.z", event) - energy = _event(tree, f"{cluster}/{cluster}.energy", event) - sx, sy = project(cx, cy, cz) - ax.scatter( - sx, - sy, - s=np.clip( - config.cluster_size_base + config.cluster_size_scale * np.sqrt(np.maximum(energy, 0)), - config.cluster_size_base, - config.cluster_size_max, - ), - c=energy, - cmap="viridis", - alpha=config.cluster_alpha, - edgecolors="none", - rasterized=True, - label="Calorimeter clusters", - ) - - status = _event(tree, "MCParticles/MCParticles.generatorStatus", event) - px = _event(tree, "MCParticles/MCParticles.momentum.x", event) - py = _event(tree, "MCParticles/MCParticles.momentum.y", event) - pz = _event(tree, "MCParticles/MCParticles.momentum.z", event) - pdg = np.abs(_event(tree, "MCParticles/MCParticles.PDG", event)).astype(int) - charge = _event(tree, "MCParticles/MCParticles.charge", event) - mass = _event(tree, "MCParticles/MCParticles.mass", event) - particle_energy = np.sqrt(px * px + py * py + pz * pz + mass * mass) - keep = status == 1 - px, py, pz, pdg, charge, particle_energy = (v[keep] for v in (px, py, pz, pdg, charge, particle_energy)) - particle_kind = np.where(np.isin(pdg, [11, 13, 22]), pdg, np.where(np.abs(charge) > 0, 211, 130)) - for code, (name, color) in PARTICLE_STYLES.items(): - selected = particle_kind == code - particle_x, particle_y, particle_z = [], [], [] - for vx, vy, vz, particle_e, particle_charge in zip(px[selected], py[selected], pz[selected], particle_energy[selected], charge[selected]): - norm = np.sqrt(vx * vx + vy * vy + vz * vz) - if norm == 0: + if not target_only: + for collection, label, color in config.hit_collections: + if collection not in tree: continue - length = _particle_display_length(vx, vy, vz, particle_e, abs(particle_charge) < 0.5, config) - scale = length / norm - particle_x.extend([0, scale * vx, np.nan]) - particle_y.extend([0, scale * vy, np.nan]) - particle_z.extend([0, scale * vz, np.nan]) - if particle_x: - sx, sy = project(np.asarray(particle_x), np.asarray(particle_y), np.asarray(particle_z)) - ax.plot(sx, sy, color=color, linewidth=0.85, linestyle="--", alpha=0.62, label=f"Particle: {name}") + x = _event(tree, f"{collection}/{collection}.position.x", event) + y = _event(tree, f"{collection}/{collection}.position.y", event) + z = _event(tree, f"{collection}/{collection}.position.z", event) + if len(x) > max_hits: + idx = np.sort(rng.choice(len(x), max_hits, replace=False)) + x, y, z = x[idx], y[idx], z[idx] + sx, sy = project(x, y, z) + ax.scatter(sx, sy, s=4.0, color=color, alpha=0.5, edgecolors="none", rasterized=True, label=label if label not in shown_labels else None) + shown_labels.add(label) + + tx, ty, tz = _track_trajectories(tree, event, config) + sx, sy = project(tx, ty, tz) + ax.plot(sx, sy, color="#111827", linewidth=0.9, alpha=0.82, label=config.track_label) + + cluster = config.cluster_collection + cx = _event(tree, f"{cluster}/{cluster}.position.x", event) + cy = _event(tree, f"{cluster}/{cluster}.position.y", event) + cz = _event(tree, f"{cluster}/{cluster}.position.z", event) + energy = _event(tree, f"{cluster}/{cluster}.energy", event) + sx, sy = project(cx, cy, cz) + ax.scatter( + sx, + sy, + s=np.clip( + config.cluster_size_base + config.cluster_size_scale * np.sqrt(np.maximum(energy, 0)), + config.cluster_size_base, + config.cluster_size_max, + ), + c=energy, + cmap="viridis", + alpha=config.cluster_alpha, + edgecolors="none", + rasterized=True, + label="Calorimeter clusters", + ) + + if show_particles or target_only: + status = _event(tree, "MCParticles/MCParticles.generatorStatus", event) + px = _event(tree, "MCParticles/MCParticles.momentum.x", event) + py = _event(tree, "MCParticles/MCParticles.momentum.y", event) + pz = _event(tree, "MCParticles/MCParticles.momentum.z", event) + pdg = np.abs(_event(tree, "MCParticles/MCParticles.PDG", event)).astype(int) + charge = _event(tree, "MCParticles/MCParticles.charge", event) + mass = _event(tree, "MCParticles/MCParticles.mass", event) + particle_energy = np.sqrt(px * px + py * py + pz * pz + mass * mass) + # Visible status-1 particles are a compact proxy for the MLPF target + # population. The full postprocessing additionally accounts for + # detector association and merging; neutrinos are never visible. + keep = (status == 1) & ~np.isin(pdg, [12, 14, 16]) + px, py, pz, pdg, charge, particle_energy = (v[keep] for v in (px, py, pz, pdg, charge, particle_energy)) + particle_kind = np.where(np.isin(pdg, [11, 13, 22]), pdg, np.where(np.abs(charge) > 0, 211, 130)) + for code, (name, color) in PARTICLE_STYLES.items(): + selected = particle_kind == code + particle_x, particle_y, particle_z = [], [], [] + endpoint_x, endpoint_y, endpoint_z, endpoint_energy = [], [], [], [] + for vx, vy, vz, particle_e, particle_charge in zip(px[selected], py[selected], pz[selected], particle_energy[selected], charge[selected]): + norm = np.sqrt(vx * vx + vy * vy + vz * vz) + if norm == 0: + continue + if target_only: + energy_fraction = np.clip(np.log1p(max(particle_e, 0.0)) / np.log1p(100.0), 0.0, 1.0) + length = display_limit * (0.28 + 0.58 * energy_fraction) + else: + length = _particle_display_length(vx, vy, vz, particle_e, abs(particle_charge) < 0.5, config) + transverse_momentum = np.hypot(vx, vy) + if target_only and abs(particle_charge) >= 0.5 and transverse_momentum > 1e-6: + # Helical propagation in the detector's axial solenoidal + # field. Radius is in mm for pT in GeV and B in tesla. + signed_radius = ( + transverse_momentum * 1000.0 / (0.3 * config.magnetic_field_tesla * particle_charge) + ) + tan_lambda = vz / transverse_momentum + transverse_arc = length / np.sqrt(1.0 + tan_lambda * tan_lambda) + arc = np.linspace(0.0, transverse_arc, 36) + phi = np.arctan2(vy, vx) + angle = phi - arc / signed_radius + path_x = signed_radius * np.sin(phi) - signed_radius * np.sin(angle) + path_y = -signed_radius * np.cos(phi) + signed_radius * np.cos(angle) + path_z = arc * tan_lambda + else: + scale = length / norm + path_x = np.asarray([0.0, scale * vx]) + path_y = np.asarray([0.0, scale * vy]) + path_z = np.asarray([0.0, scale * vz]) + particle_x.extend(path_x.tolist() + [np.nan]) + particle_y.extend(path_y.tolist() + [np.nan]) + particle_z.extend(path_z.tolist() + [np.nan]) + endpoint_x.append(path_x[-1]) + endpoint_y.append(path_y[-1]) + endpoint_z.append(path_z[-1]) + endpoint_energy.append(particle_e) + if particle_x: + sx, sy = project(np.asarray(particle_x), np.asarray(particle_y), np.asarray(particle_z)) + ax.plot( + sx, + sy, + color=color, + linewidth=1.25 if target_only else 0.85, + linestyle="-" if target_only else "--", + alpha=0.78 if target_only else 0.62, + label=name if target_only else f"Particle: {name}", + ) + if target_only: + ex, ey = project(np.asarray(endpoint_x), np.asarray(endpoint_y), np.asarray(endpoint_z)) + ax.scatter( + ex, + ey, + s=np.clip(5 + 1.5 * np.sqrt(np.asarray(endpoint_energy)), 5, 22), + color=color, + alpha=0.85, + edgecolors="none", + ) ax.set_aspect("equal", adjustable="box") ax.set_xlim(-display_limit, display_limit) ax.set_ylim(-display_limit, display_limit) - ax.set_title(f"{config.title} — event {event}", fontsize=15) + seed = _production_suffix(root_file) + seed_label = f" — seed {seed}" if seed is not None else "" + title_suffix = " — visible status-1 target proxy" if target_only else "" + ax.set_title(f"{config.title}{seed_label} — event {event}{title_suffix}", fontsize=15) ax.legend(loc="upper left", fontsize=7.5, ncol=2, frameon=True, framealpha=0.9) + # Fixed camera-orientation marker: the beam axis is perpendicular to the + # image, +x points left and +y points up. The circled dot denotes +z out of + # the screen (toward the viewer). + axis_origin = (0.91, 0.10) + ax.annotate("", xy=(0.82, 0.10), xytext=axis_origin, xycoords="axes fraction", arrowprops={"arrowstyle": "->", "color": "#475569", "lw": 1.2}) + ax.annotate("", xy=(0.91, 0.19), xytext=axis_origin, xycoords="axes fraction", arrowprops={"arrowstyle": "->", "color": "#475569", "lw": 1.2}) + ax.text(0.805, 0.085, "+x", transform=ax.transAxes, fontsize=8, color="#475569", ha="right", va="top") + ax.text(0.925, 0.195, "+y", transform=ax.transAxes, fontsize=8, color="#475569", ha="left", va="bottom") + ax.text(0.91, 0.065, r"$\odot\ +z$", transform=ax.transAxes, fontsize=8, color="#475569", ha="center", va="top") ax.axis("off") fig.savefig(output, dpi=150, facecolor="white", bbox_inches="tight", pad_inches=0.02) plt.close(fig) @@ -425,6 +535,8 @@ def main() -> None: parser.add_argument("--output-dir", type=Path, default=Path("event_displays")) parser.add_argument("--detector", choices=("auto", *DETECTORS), default="auto") parser.add_argument("--max-hits", type=int, default=800, help="maximum displayed hits per collection") + parser.add_argument("--no-particles", action="store_true", help="omit stable generator-particle guide lines") + parser.add_argument("--target-only", action="store_true", help="render only visible status-1 MC particles as a target proxy") parser.add_argument("--debug", action="store_true", help="also render track/hit and cluster/hit association checks") args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) @@ -433,13 +545,30 @@ def main() -> None: with _open_root(root_file) as source: config = _detector_config(source["events"], args.detector) inputs.append((root_file, config)) + production_suffixes = [_production_suffix(root_file) for root_file, _ in inputs] + numeric_suffixes = [suffix for suffix in production_suffixes if suffix is not None] + if len(inputs) > 1 and len(numeric_suffixes) == len(inputs) and len(set(numeric_suffixes)) != 1: + raise ValueError( + "comparison inputs must have the same trailing numerical suffix " + f"(got: {', '.join(numeric_suffixes)})" + ) comparison_limit = max(config.plot_limit for _, config in inputs) if len(inputs) > 1 else None for event in args.events: event_images = [] for root_file, config in inputs: - output = args.output_dir / f"{config.key}_event_{event}.png" - detector = render_event(root_file, event, output, args.max_hits, args.detector, comparison_limit) + suffix = "_targets" if args.target_only else "" + output = args.output_dir / f"{config.key}_event_{event}{suffix}.png" + detector = render_event( + root_file, + event, + output, + args.max_hits, + args.detector, + comparison_limit, + show_particles=not args.no_particles, + target_only=args.target_only, + ) event_images.append((output, detector)) print(output) if args.debug: From d8797cd95671f6795b1e659880cd8463697f9f9f Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Thu, 17 Sep 2026 10:13:23 +0300 Subject: [PATCH 29/29] format --- mlpf/training_scenarios.py | 122 ++++++++---------------------- mlpf/training_submission.py | 44 +++-------- scripts/visualize_key4hep.py | 13 +--- tests/test_training_scenarios.py | 34 ++------- tests/test_training_submission.py | 4 +- 5 files changed, 53 insertions(+), 164 deletions(-) diff --git a/mlpf/training_scenarios.py b/mlpf/training_scenarios.py index 6078b1adf..ca1f45447 100644 --- a/mlpf/training_scenarios.py +++ b/mlpf/training_scenarios.py @@ -43,9 +43,7 @@ class ScenarioVariant(BaseModel): def reject_derived_overrides(self): invalid = DERIVED_KEYS.intersection(self.overrides) if invalid: - raise ValueError( - f"Variant overrides must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Variant overrides must not set derived keys: {sorted(invalid)}") return self @@ -59,9 +57,7 @@ class ScenarioTraining(BaseModel): def reject_derived_parameters(self): invalid = DERIVED_KEYS.intersection(self.parameters) if invalid: - raise ValueError( - f"Scenario parameters must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Scenario parameters must not set derived keys: {sorted(invalid)}") return self @@ -75,9 +71,7 @@ class TrainingScenario(BaseModel): seeds: list[int] = Field(min_length=1) training: ScenarioTraining common_overrides: dict[str, Any] = Field(default_factory=dict) - allowed_variant_differences: list[str] = Field( - default_factory=lambda: ["model.output_mode", "model.set_decoder"] - ) + allowed_variant_differences: list[str] = Field(default_factory=lambda: ["model.output_mode", "model.set_decoder"]) @model_validator(mode="after") def validate_scenario(self): @@ -89,9 +83,7 @@ def validate_scenario(self): raise ValueError("Scenario seeds must be non-negative") invalid = DERIVED_KEYS.intersection(self.common_overrides) if invalid: - raise ValueError( - f"Common overrides must not set derived keys: {sorted(invalid)}" - ) + raise ValueError(f"Common overrides must not set derived keys: {sorted(invalid)}") return self def variant_production(self, variant_name): @@ -140,14 +132,9 @@ class PlatformProfile(BaseModel): def validate_runtime_overrides(self): invalid = set(self.runtime_overrides).difference(PLATFORM_OVERRIDE_KEYS) if invalid: - raise ValueError( - "Platform profiles may only set runtime-specific overrides; " - f"invalid keys: {sorted(invalid)}" - ) + raise ValueError("Platform profiles may only set runtime-specific overrides; " f"invalid keys: {sorted(invalid)}") if isinstance(self.data_dir, dict) and not self.data_dir: - raise ValueError( - "Platform data_dir mapping must name at least one production" - ) + raise ValueError("Platform data_dir mapping must name at least one production") return self def data_dir_for(self, production_name): @@ -155,8 +142,7 @@ def data_dir_for(self, production_name): return self.data_dir if production_name not in self.data_dir: raise ValueError( - f"Platform profile {self.name!r} has no data_dir for production {production_name!r}; " - f"available: {sorted(self.data_dir)}" + f"Platform profile {self.name!r} has no data_dir for production {production_name!r}; " f"available: {sorted(self.data_dir)}" ) return self.data_dir[production_name] @@ -200,17 +186,9 @@ def load_platform_profile(path): if isinstance(profile.data_dir, str): profile.data_dir = os.path.expandvars(os.path.expanduser(profile.data_dir)) else: - profile.data_dir = { - key: os.path.expandvars(os.path.expanduser(value)) - for key, value in profile.data_dir.items() - } - profile.experiments_dir = os.path.expandvars( - os.path.expanduser(profile.experiments_dir) - ) - profile.environment = { - key: os.path.expandvars(os.path.expanduser(value)) - for key, value in profile.environment.items() - } + profile.data_dir = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.data_dir.items()} + profile.experiments_dir = os.path.expandvars(os.path.expanduser(profile.experiments_dir)) + profile.environment = {key: os.path.expandvars(os.path.expanduser(value)) for key, value in profile.environment.items()} return profile @@ -269,8 +247,7 @@ def _training_batch_size(config): batch_sizes = {dataset.batch_size for dataset in physical_datasets} if len(batch_sizes) != 1: raise ValueError( - "Automatic global-batch resolution requires every physical training dataset " - f"to use the same batch size, got {sorted(batch_sizes)}" + "Automatic global-batch resolution requires every physical training dataset " f"to use the same batch size, got {sorted(batch_sizes)}" ) return next(iter(batch_sizes)) @@ -296,16 +273,12 @@ def resolve_scenario_job( extra_overrides=None, ): if variant_name not in scenario.variants: - raise ValueError( - f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}" - ) + raise ValueError(f"Unknown variant {variant_name!r}; choose from {sorted(scenario.variants)}") variant = scenario.variants[variant_name] extra_overrides = extra_overrides or {} invalid = DERIVED_KEYS.intersection(extra_overrides) if invalid: - raise ValueError( - f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}" - ) + raise ValueError(f"Use the dedicated runner options for derived settings, not --set: {sorted(invalid)}") settings = _merge_settings(scenario, platform, variant, extra_overrides) settings["seed"] = seed selected_spec = str(spec_file or scenario.spec_file) @@ -321,19 +294,14 @@ def resolve_scenario_job( extra_args=_settings_as_extra_args(settings), ) - target_global_batch = ( - global_batch_size - if global_batch_size is not None - else scenario.training.global_batch_size - ) + target_global_batch = global_batch_size if global_batch_size is not None else scenario.training.global_batch_size if target_global_batch <= 0: raise ValueError("global_batch_size must be positive") dataset_batch_size = _training_batch_size(config) divisor = platform.gpus * dataset_batch_size if target_global_batch % divisor: raise ValueError( - f"global_batch_size={target_global_batch} is not divisible by " - f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" + f"global_batch_size={target_global_batch} is not divisible by " f"gpus={platform.gpus} * dataset_batch_size={dataset_batch_size}" ) multiplier = target_global_batch // divisor settings["gpu_batch_multiplier"] = multiplier @@ -375,9 +343,7 @@ def _flatten(value, prefix=""): def _difference_allowed(path, allowed_paths): - return any( - path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths - ) + return any(path == allowed or path.startswith(f"{allowed}.") for allowed in allowed_paths) def validate_variant_invariants(jobs, allowed_paths): @@ -389,17 +355,11 @@ def validate_variant_invariants(jobs, allowed_paths): differences = { path: (reference.get(path), candidate.get(path)) for path in sorted(set(reference) | set(candidate)) - if reference.get(path) != candidate.get(path) - and not _difference_allowed(path, allowed_paths) + if reference.get(path) != candidate.get(path) and not _difference_allowed(path, allowed_paths) } if differences: - details = ", ".join( - f"{path}: {values[0]!r} != {values[1]!r}" - for path, values in differences.items() - ) - raise ValueError( - f"Scenario variants differ outside allowed fields: {details}" - ) + details = ", ".join(f"{path}: {values[0]!r} != {values[1]!r}" for path, values in differences.items()) + raise ValueError(f"Scenario variants differ outside allowed fields: {details}") def resolve_scenario_jobs( @@ -513,39 +473,25 @@ def find_scenario_continuation(job, platform): experiment_dir = manifest_path.parent checkpoints = [] - for checkpoint_path in (experiment_dir / "checkpoints").glob( - "checkpoint-*.pth" - ): + for checkpoint_path in (experiment_dir / "checkpoints").glob("checkpoint-*.pth"): step = _checkpoint_step(checkpoint_path) if step is not None: checkpoints.append((step, checkpoint_path)) step, checkpoint = max(checkpoints, default=(0, None), key=lambda item: item[0]) - candidates.append( - (step, manifest_path.stat().st_mtime, experiment_dir, checkpoint) - ) + candidates.append((step, manifest_path.stat().st_mtime, experiment_dir, checkpoint)) if not candidates: return None - step, _, experiment_dir, checkpoint = max( - candidates, key=lambda item: (item[0], item[1]) - ) - return ScenarioContinuation( - experiment_dir=experiment_dir, checkpoint=checkpoint, step=step - ) + step, _, experiment_dir, checkpoint = max(candidates, key=lambda item: (item[0], item[1])) + return ScenarioContinuation(experiment_dir=experiment_dir, checkpoint=checkpoint, step=step) -def run_scenario_job( - job, scenario, platform, spec_file, *, dry_run=False, continue_run=False -): +def run_scenario_job(job, scenario, platform, spec_file, *, dry_run=False, continue_run=False): continuation = find_scenario_continuation(job, platform) if continue_run else None if continue_run and continuation is None: - raise ValueError( - f"No compatible prior run found for {job.variant_name} seed {job.seed}" - ) + raise ValueError(f"No compatible prior run found for {job.variant_name} seed {job.seed}") experiment_dir = ( - continuation.experiment_dir - if continuation is not None - else _experiment_path(platform, job, timestamp="TIMESTAMP" if dry_run else None) + continuation.experiment_dir if continuation is not None else _experiment_path(platform, job, timestamp="TIMESTAMP" if dry_run else None) ) if continuation is not None and continuation.checkpoint is not None: job.settings["load"] = str(continuation.checkpoint) @@ -570,9 +516,7 @@ def run_scenario_job( manifest = json.load(handle) manifest.setdefault("continuations", []).append( { - "checkpoint": str(continuation.checkpoint) - if continuation.checkpoint - else None, + "checkpoint": str(continuation.checkpoint) if continuation.checkpoint else None, "step": continuation.step, "command": command, "git_revision": _git_revision(), @@ -600,9 +544,7 @@ def _parse_set_overrides(values): def _validate_slurm_allocation(platform): allocated = os.environ.get("SLURM_GPUS_PER_NODE") if allocated and allocated.isdigit() and int(allocated) != platform.gpus: - raise ValueError( - f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}" - ) + raise ValueError(f"Platform profile requests {platform.gpus} GPUs but Slurm allocated {allocated}") def main(argv=None): @@ -662,14 +604,10 @@ def main(argv=None): if args.continue_run: continuation = find_scenario_continuation(job, platform) if continuation is None: - raise ValueError( - f"No compatible prior run found for {job.variant_name} seed {job.seed}" - ) + raise ValueError(f"No compatible prior run found for {job.variant_name} seed {job.seed}") target_step = job.resolved_config.num_steps if continuation.step >= target_step: - print( - f"Skipping completed job {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step}" - ) + print(f"Skipping completed job {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step}") continue run_scenario_job( job, diff --git a/mlpf/training_submission.py b/mlpf/training_submission.py index 475ab45ad..8e6331cc8 100644 --- a/mlpf/training_submission.py +++ b/mlpf/training_submission.py @@ -45,13 +45,8 @@ def resolve_flatiron_profile_path(reference, repo_root): def available_choices(repo_root, site="flatiron"): - scenarios = sorted( - path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml") - ) - accelerators = sorted( - path.stem.removeprefix(f"{site}_") - for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml") - ) + scenarios = sorted(path.stem for path in (repo_root / "configs/training/scenarios").glob("*.yaml")) + accelerators = sorted(path.stem.removeprefix(f"{site}_") for path in (repo_root / "configs/training/platforms").glob(f"{site}_*.yaml")) return scenarios, accelerators @@ -94,9 +89,7 @@ def build_slurm_submission( scenario.seeds = [seed] profile = load_platform_profile(profile_path) if profile.slurm is None: - raise ValueError( - f"Platform profile {profile.name!r} has no Slurm configuration" - ) + raise ValueError(f"Platform profile {profile.name!r} has no Slurm configuration") spec_file = Path(scenario.spec_file) if not spec_file.is_absolute(): @@ -111,31 +104,21 @@ def build_slurm_submission( for index, job in enumerate(jobs): continuation = find_scenario_continuation(job, profile) if continuation is None: - print( - f"Skipping {job.variant_name} seed {job.seed}: no compatible prior run" - ) + print(f"Skipping {job.variant_name} seed {job.seed}: no compatible prior run") continue target_step = job.resolved_config.num_steps if continuation.step >= target_step: - print( - f"Skipping {job.variant_name} seed {job.seed}: complete at step {continuation.step}/{target_step}" - ) + print(f"Skipping {job.variant_name} seed {job.seed}: complete at step {continuation.step}/{target_step}") continue - source = ( - continuation.checkpoint.name if continuation.checkpoint else "start" - ) - print( - f"Continuing {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step} from {source}" - ) + source = continuation.checkpoint.name if continuation.checkpoint else "start" + print(f"Continuing {job.variant_name} seed {job.seed}: step {continuation.step}/{target_step} from {source}") selected_indices.append(index) if not selected_indices: raise ValueError("No unfinished compatible scenario jobs found") slurm = profile.slurm logs_dir = repo_root / "logs_slurm" - worker = ( - Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") - ) + worker = Path(worker) if worker is not None else _worker_for_site(repo_root, "flatiron") command = [ "sbatch", "--time", @@ -163,9 +146,7 @@ def build_slurm_submission( command.extend( [ "--array", - ",".join(str(index) for index in selected_indices) - if continue_run - else f"0-{len(jobs) - 1}", + ",".join(str(index) for index in selected_indices) if continue_run else f"0-{len(jobs) - 1}", "--job-name", scenario.name, "--output", @@ -202,9 +183,7 @@ def main(argv=None, *, site="flatiron"): action="store_true", help="Print the sbatch command without submitting", ) - parser.add_argument( - "--list", action="store_true", help="List available scenarios and accelerators" - ) + parser.add_argument("--list", action="store_true", help="List available scenarios and accelerators") parser.add_argument( "--continue", dest="continue_run", @@ -233,8 +212,7 @@ def main(argv=None, *, site="flatiron"): continue_run=args.continue_run, ) print( - f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" - + shlex.join(command), + f"Submitting {len(jobs)} jobs for {args.scenario} on {args.accelerator}:\n" + shlex.join(command), flush=True, ) if args.dry_run: diff --git a/scripts/visualize_key4hep.py b/scripts/visualize_key4hep.py index 3c81c73b9..0d5e0462a 100644 --- a/scripts/visualize_key4hep.py +++ b/scripts/visualize_key4hep.py @@ -165,9 +165,7 @@ def _detector_config(tree, detector: str = "auto") -> DetectorConfig: if "ECALOther" in tree: return DETECTORS["clic"] matches = [ - config - for config in DETECTORS.values() - if config.key != "clic" and config.track_collection in tree and config.cluster_collection in tree + config for config in DETECTORS.values() if config.key != "clic" and config.track_collection in tree and config.cluster_collection in tree ] if len(matches) != 1: found = ", ".join(config.key for config in matches) or "none" @@ -447,9 +445,7 @@ def project(x, y, z): if target_only and abs(particle_charge) >= 0.5 and transverse_momentum > 1e-6: # Helical propagation in the detector's axial solenoidal # field. Radius is in mm for pT in GeV and B in tesla. - signed_radius = ( - transverse_momentum * 1000.0 / (0.3 * config.magnetic_field_tesla * particle_charge) - ) + signed_radius = transverse_momentum * 1000.0 / (0.3 * config.magnetic_field_tesla * particle_charge) tan_lambda = vz / transverse_momentum transverse_arc = length / np.sqrt(1.0 + tan_lambda * tan_lambda) arc = np.linspace(0.0, transverse_arc, 36) @@ -548,10 +544,7 @@ def main() -> None: production_suffixes = [_production_suffix(root_file) for root_file, _ in inputs] numeric_suffixes = [suffix for suffix in production_suffixes if suffix is not None] if len(inputs) > 1 and len(numeric_suffixes) == len(inputs) and len(set(numeric_suffixes)) != 1: - raise ValueError( - "comparison inputs must have the same trailing numerical suffix " - f"(got: {', '.join(numeric_suffixes)})" - ) + raise ValueError("comparison inputs must have the same trailing numerical suffix " f"(got: {', '.join(numeric_suffixes)})") comparison_limit = max(config.plot_limit for _, config in inputs) if len(inputs) > 1 else None for event in args.events: diff --git a/tests/test_training_scenarios.py b/tests/test_training_scenarios.py index d8c9efcf9..2c5674103 100644 --- a/tests/test_training_scenarios.py +++ b/tests/test_training_scenarios.py @@ -19,13 +19,9 @@ ROOT = Path(__file__).resolve().parents[1] SCENARIO = ROOT / "configs/training/scenarios/cld_hits_output_comparison.yaml" -BACKBONE_SCENARIO = ( - ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" -) +BACKBONE_SCENARIO = ROOT / "configs/training/scenarios/cld_hits_backbone_comparison.yaml" PF_HITS_SCENARIO = ROOT / "configs/training/scenarios/cld_pf_hits_comparison.yaml" -CLIC_CLD_SCENARIO = ( - ROOT / "configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml" -) +CLIC_CLD_SCENARIO = ROOT / "configs/training/scenarios/clic_cld_pf_set_hits_comparison.yaml" PLATFORMS = ROOT / "configs/training/platforms" @@ -66,9 +62,7 @@ def test_backbone_comparison_scenario_keeps_elementwise_output_and_depth_fixed() "attention", "heptv2", ] - assert {job.resolved_config.model.output_mode.value for job in jobs} == { - "elementwise" - } + assert {job.resolved_config.model.output_mode.value for job in jobs} == {"elementwise"} assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} assert jobs[1].resolved_config.model.heptv2.block_size == 128 @@ -100,9 +94,7 @@ def test_pf_hits_comparison_scenario_resolves_three_40k_variants(): "elementwise", "set", ] - assert [ - job.resolved_config.model.binary_classification_focal_gamma for job in jobs - ] == [None, 2.0, 2.0] + assert [job.resolved_config.model.binary_classification_focal_gamma for job in jobs] == [None, 2.0, 2.0] assert {job.resolved_config.model.backbone.num_convs for job in jobs} == {6} assert [ ( @@ -150,9 +142,7 @@ def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): platform.data_dir["clic"], platform.data_dir["clic"], ] - assert [job.resolved_config.data_dir for job in jobs] == [ - job.data_dir for job in jobs - ] + assert [job.resolved_config.data_dir for job in jobs] == [job.data_dir for job in jobs] assert [job.resolved_config.dataset.value for job in jobs] == [ "cld", "cld_hits", @@ -180,11 +170,7 @@ def test_clic_cld_scenario_resolves_pf_and_set_hits_per_detector(): ) for job in jobs ] == [(None, None, None), (4, 4, 4), (None, None, None), (4, 4, 4)] - assert { - job.resolved_config.model.set_decoder.num_layers - for job in jobs - if job.resolved_config.model.set_decoder - } == {8} + assert {job.resolved_config.model.set_decoder.num_layers for job in jobs if job.resolved_config.model.set_decoder} == {8} assert {job.resolved_config.num_steps for job in jobs} == {50000} assert {job.resolved_config.val_freq for job in jobs} == {5000} assert {job.resolved_config.lr for job in jobs} == {0.001} @@ -354,9 +340,7 @@ def test_experiments_are_grouped_under_the_scenario_directory(): path = _experiment_path(platform, job, timestamp="TIMESTAMP") - assert path == Path( - "experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP" - ) + assert path == Path("experiments/cld_hits_output_comparison/elementwise_seed12345_TIMESTAMP") def test_continuation_chooses_most_advanced_compatible_checkpoint(tmp_path): @@ -371,9 +355,7 @@ def test_continuation_chooses_most_advanced_compatible_checkpoint(tmp_path): )[0] for suffix, step in [("older", 10000), ("newer", 5000), ("best", 15000)]: - run_dir = ( - tmp_path / scenario.name / f"{job.variant_name}_seed{job.seed}_{suffix}" - ) + run_dir = tmp_path / scenario.name / f"{job.variant_name}_seed{job.seed}_{suffix}" (run_dir / "checkpoints").mkdir(parents=True) manifest = { "job": job.model_dump(mode="json", exclude={"resolved_config"}), diff --git a/tests/test_training_submission.py b/tests/test_training_submission.py index d8a468334..c1ac3edfd 100644 --- a/tests/test_training_submission.py +++ b/tests/test_training_submission.py @@ -120,9 +120,7 @@ def test_picker_discovers_site_specific_accelerators(): def _write_scenario_run(experiments_dir, scenario, profile, job, step): - run_dir = ( - experiments_dir / scenario.name / f"{job.variant_name}_seed{job.seed}_test" - ) + run_dir = experiments_dir / scenario.name / f"{job.variant_name}_seed{job.seed}_test" checkpoint_dir = run_dir / "checkpoints" checkpoint_dir.mkdir(parents=True) manifest = {