# Task 15: Profile and Optimize GNN4Colliders Performance Implement a focused, measurement-driven performance optimization pass for GNN4Colliders. This task builds on the stable and parity-tested implementation from: ```text Task 7: metadata-aware caching, splits, batching, and GraphBatch Task 8: ROOT-GNN model and transfer/fine-tuning Task 9: classification tasks, losses, metrics, and output handling Task 10: training and evaluation lifecycle Task 11: checkpointing and resume Task 12: evaluation, prediction, and output serialization Task 13: Hydra configuration and CLI Task 14: DDP, Perlmutter, and Slurm execution ``` Before making changes, read: ```text AGENTS.md README.md docs/architecture.md docs/migration.md ``` Then inspect: ```text src/gnn4colliders/data/ src/gnn4colliders/features/ src/gnn4colliders/graphs/ src/gnn4colliders/models/root_gnn/ src/gnn4colliders/tasks/ src/gnn4colliders/training/ src/gnn4colliders/inference/ src/gnn4colliders/distributed/ src/gnn4colliders/cli/ configs/ scripts/slurm/ tests/ ``` Also inspect any existing benchmark/profiling utilities in the repository. Treat the currently validated scientific behavior as fixed. Do not change model mathematics, task semantics, data selection, feature definitions, graph topology, event weighting, checkpoint semantics, or inference outputs merely to obtain better benchmark numbers. --- # Goal Measure where time and memory are spent in the end-to-end ROOT-GNN workflow, identify the dominant bottlenecks, and implement only optimizations justified by measurements. Profile the pipeline: ```text ROOT / Awkward I/O -> EventSample construction -> shared feature construction -> graph construction -> cache read/write -> sampling / batching -> GraphBatch -> host-to-device transfer -> ROOT-GNN forward -> task loss -> backward -> optimizer step -> DDP synchronization where enabled ``` The task should answer: ```text Where is wall-clock time spent? Where is CPU time spent? Where is GPU time spent? Where are synchronization stalls? Is the GPU waiting for data? What consumes the most CPU/GPU memory? Which optimizations materially improve throughput? Which attempted optimizations do not help? Do optimized paths preserve existing parity? ``` Performance changes must be evidence-driven. --- # 1. Add a benchmark/profiling area Add a dedicated benchmark area. Prefer: ```text benchmarks/ README.md benchmark_preprocessing.py benchmark_dataloader.py benchmark_training.py benchmark_inference.py ``` or a similarly small structure. Do not put manual performance benchmarks in the normal unit-test suite. Benchmark programs should be runnable directly from the repository. --- # 2. Benchmark reproducibility Benchmarks must be reproducible. Use: * explicit seeds * fixed benchmark inputs * fixed iteration counts * warmup iterations where needed * explicit device * explicit batch size * explicit worker count * explicit graph/event sizes Do not rely on wall-clock-generated random data. Record enough information in benchmark output to interpret results. --- # 3. Benchmark output Prefer concise structured benchmark output. At minimum report relevant values such as: ```text events / second graphs / second batches / second training steps / second milliseconds / batch seconds / epoch ``` Where relevant also report: ```text mean median min/max standard deviation ``` Do not report excessive precision. --- # 4. Benchmark metadata Each benchmark run should report or record relevant execution metadata. For example: ```text Python version PyTorch version DGL version device GPU model if available world size batch size num_workers graph/event size dtype ``` Keep this lightweight. Do not add a benchmark database or tracking service. --- # 5. Establish a baseline first Before changing production performance code, establish baseline measurements. Record baseline numbers in: ```text benchmarks/README.md ``` or a concise performance document. Do not rewrite large parts of the pipeline before measuring the current implementation. --- # 6. Separate setup from steady-state timing Do not mix one-time setup costs into steady-state throughput unless intentionally measuring startup. For example distinguish: ```text dataset construction cache construction first batch steady-state cached batches model initialization first CUDA execution steady-state forward/backward ``` CUDA initialization and kernel compilation/warmup should not distort the steady-state benchmark. --- # 7. Correct CUDA timing For CUDA benchmarks, account for asynchronous execution. Use appropriate synchronization around timed regions, such as: ```python torch.cuda.synchronize() ``` or CUDA events. Do not use naïve wall-clock timing around asynchronous CUDA operations and treat it as accurate GPU execution time. --- # 8. Profiling tools Use standard tooling first. Prefer: ```text cProfile / pstats PyTorch profiler torch.utils.benchmark CUDA memory statistics ``` where appropriate. Do not add a heavyweight profiling dependency unless necessary. --- # 9. PyTorch profiler Add a documented way to profile a short training section with: ```python torch.profiler ``` Capture enough iterations to distinguish: ```text CPU preprocessing host-to-device copies DGL operations matrix operations backward optimizer work distributed synchronization ``` Do not enable profiler overhead by default during ordinary training. --- # 10. Profile feature construction Measure shared feature construction independently. Profile: * event/object transformation * energy calculation * scaling * stacking/concatenation * dtype conversions * avoidable allocations Do not change feature values or ordering. --- # 11. Feature-construction optimization Only after profiling, consider safe improvements such as: ```text reducing repeated conversions reducing unnecessary copies vectorizing Python loops preallocating outputs avoiding repeated scale-array construction ``` Preserve exact schema and ordering. Add parity tests for any changed implementation. --- # 12. Profile graph construction Measure graph construction separately. Include: ```text edge-index construction fully connected no-self-loop topology edge feature calculation phi wrapping DGL graph creation ndata/edata population ``` Test representative node counts. Remember that fully connected graph edge count scales as: ```text N * (N - 1) ``` Do not accidentally change topology while optimizing. --- # 13. Edge-index construction Inspect whether edge index generation is repeated unnecessarily. If safe, consider efficient deterministic generation of: ```text src dst ``` for the fully connected directed no-self-loop graph. Preserve the validated edge ordering if parity tests depend on it. Do not substitute an ordering with the same mathematical graph unless ordering is explicitly known to be irrelevant. --- # 14. Edge-feature optimization Profile: ```text deta dphi dR ``` calculation. Look for: ```text repeated indexing temporary arrays Python loops unnecessary device/host conversions ``` Optimize only if measurable. Preserve phi wrapping exactly. --- # 15. Graph construction caching If graph topology is repeatedly constructed for identical node counts, evaluate whether caching reusable topology indices provides a measurable benefit. For example, a bounded cache keyed by: ```text number of nodes device ``` may be useful. Do not implement unbounded caching. Do not cache event-specific edge features as if they were topology. Only keep this optimization if benchmarks show a real benefit. --- # 16. Profile cache behavior Measure: ```text cold cache creation warm cache loading cache deserialization filesystem throughput ``` Distinguish CPU processing cost from filesystem cost. Do not infer cache performance from total training time alone. --- # 17. Cache format changes require evidence Do not replace the Task 7 cache format solely because another format seems theoretically faster. If cache serialization/deserialization is a proven bottleneck, evaluate a narrowly scoped improvement. Any format change must preserve: ```text schema versioning preprocessing fingerprinting metadata sample identity compatibility validation ``` Do not break existing caches silently. If a format changes, bump the cache schema version. --- # 18. Profile DataLoader throughput Measure DataLoader performance independently from model execution. Benchmark combinations of: ```text batch_size num_workers persistent_workers prefetch_factor pin_memory ``` where relevant. Use a bounded, meaningful search. Do not brute-force dozens of configurations without a hypothesis. --- # 19. DataLoader worker count Determine whether additional workers improve actual throughput. On Perlmutter this may depend strongly on: ```text CPU allocation filesystem batching implementation cache format ``` Do not make a large `num_workers` value the global default based on a laptop or a single node. Keep defaults conservative unless benchmark evidence supports changing them. --- # 20. Persistent workers Evaluate: ```python persistent_workers=True ``` for multi-epoch training when `num_workers > 0`. Keep it only if: * compatible with the dataset implementation * does not introduce stale state * improves measured epoch throughput Do not enable it when `num_workers=0`. --- # 21. Prefetching Evaluate DataLoader prefetch behavior. Do not set very large prefetch values by default. Measure both throughput and host-memory cost. Expose performance-related loader settings through configuration only when useful. --- # 22. Pinned memory For CUDA workflows, benchmark: ```python pin_memory=True ``` where compatible with the custom batch structure. Do not assume DGL graph objects behave identically to ordinary tensor-only batches. Test actual host-to-device timing. --- # 23. Non-blocking device transfers If pinned memory is effective, evaluate: ```python tensor.to(device, non_blocking=True) ``` for ordinary tensors. Do not mark operations non-blocking without satisfying the required memory conditions. Keep transfer logic centralized in `GraphBatch.to(...)` or the existing device-transfer boundary. --- # 24. Profile GraphBatch transfer Measure separately: ```text DGL graph transfer labels transfer global-feature transfer weight transfer other numeric metadata transfer ``` Do not transfer strings such as `sample_id` to GPU. Look for repeated transfers of data that is never consumed by the model/task. --- # 25. Avoid unnecessary GPU metadata Audit which `GraphBatch` fields need to reside on GPU. Typically: ```text graph features global features labels during training/evaluation weights required by the loss ``` need device access. Metadata such as: ```text sample_id source file tree name entry index fold ``` should remain on CPU unless actually required by GPU computation. Do not discard metadata; simply avoid pointless GPU transfer. --- # 26. Profile model forward Profile `EdgeNetwork` forward separately from the DataLoader. Measure: ```text node encoder edge encoder global encoder each message-passing step global pooling/update decoder classifier ``` where profiler attribution allows. Determine whether time is dominated by: ```text DGL message passing MLPs graph pooling memory movement Python overhead ``` Do not optimize blindly. --- # 27. DGL graph-local state Audit Task 8's use of: ```python graph.local_scope() ``` and temporary `ndata` / `edata`. Ensure temporary graph-state safety is not creating obviously unnecessary copying. Do not remove `local_scope()` merely for speed if doing so leaks mutation between forward calls. Correctness remains higher priority. --- # 28. Reduce repeated graph-data lookups If profiling shows meaningful Python/DGL lookup overhead, safely reduce repeated access to: ```text graph.ndata[...] graph.edata[...] ``` inside tight loops. Do not duplicate large tensors merely to save trivial dictionary lookup time. Require benchmark evidence. --- # 29. MLP optimization Profile the Task 8 MLP blocks. Consider only semantics-preserving changes such as: ```text removing avoidable Python overhead avoiding repeated module construction using efficient contiguous tensors where needed ``` Do not change: ```text layer ordering activation functions LayerNorm placement dropout semantics hidden dimensions ``` without explicitly leaving Task 15 scope. --- # 30. Tensor contiguity Profile whether non-contiguous tensors cause material overhead in hot paths. Use `.contiguous()` only where it measurably helps or is required. Do not scatter unnecessary contiguous copies through the model. --- # 31. In-place operations Do not introduce in-place operations merely for theoretical memory savings if they make autograd or parity fragile. Only use in-place operations where clearly safe and beneficial. Add tests around any altered computation path. --- # 32. Profile backward Measure forward and backward separately. Determine whether training is dominated by: ```text forward graph operations backward graph operations MLP gradients optimizer update DDP all-reduce ``` This matters before selecting optimization targets. --- # 33. Optimizer overhead Measure optimizer step overhead. Do not replace the scientifically/configurationally active optimizer merely for speed. If PyTorch exposes a semantics-compatible implementation option such as: ```text foreach fused ``` benchmark it separately. Only use it if compatibility and platform support are clear. --- # 34. `zero_grad(set_to_none=True)` Benchmark and consider: ```python optimizer.zero_grad(set_to_none=True) ``` if not already used. This is often a safe performance improvement, but verify: * existing code does not rely on zero tensors instead of `None` * training parity remains intact Document the choice. --- # 35. Mixed precision Now that correctness and DDP are established, evaluate mixed precision as an optional performance feature. Do not make it mandatory. For modern supported CUDA hardware, evaluate: ```text bf16 ``` before automatically choosing fp16, where the platform supports it. The actual supported precision mode should be determined from the current execution environment. --- # 36. AMP configuration If mixed precision is implemented, expose it explicitly. Conceptually: ```yaml trainer: precision: float32 ``` with supported alternatives such as: ```text bfloat16 float16 ``` only if tested. Do not silently enable AMP based on CUDA availability. --- # 37. AMP task/loss behavior Keep numerically sensitive operations in suitable precision. Do not force Task 9 metric computation into reduced precision. Metrics should continue to operate on stable detached outputs. Check weighted losses carefully under negative or large event weights. --- # 38. AMP parity For optional mixed precision, do not require bitwise parity with float32. Instead validate: ```text finite loss finite gradients reasonable numerical agreement stable task metrics ``` using explicitly justified tolerances. Float32 remains the reference correctness path. --- # 39. AMP performance threshold Do not keep a mixed-precision implementation merely because it technically works. Record whether it improves: ```text throughput GPU memory time per epoch ``` on the target GPU environment. Document the measured benefit. --- # 40. `torch.compile` Evaluate `torch.compile` only after establishing ordinary eager-mode profiles. Do not restructure the entire model around compilation. Treat compilation as optional. Benchmark: ```text first-call compile cost steady-state throughput compatibility with DGL compatibility with DDP ``` If the active DGL graph workflow does not benefit or causes graph breaks, document that and do not force it. --- # 41. Compilation configuration If retained, expose compilation explicitly. For example: ```yaml trainer: compile: enabled: false ``` Do not enable it by default without clear target-environment evidence. --- # 42. No custom CUDA kernels in this task Do not write custom CUDA/C++ extensions. Do not add Triton kernels. Do not add CuPy implementations. Use existing PyTorch/DGL primitives first. Custom kernels require a separate task with dedicated correctness/performance justification. --- # 43. No Numba unless profiling proves need Do not add Numba merely because preprocessing is CPU-side. Use NumPy/Awkward/vectorized operations first. Only introduce a new compiled dependency if there is a demonstrated bottleneck that cannot be addressed cleanly otherwise. Prefer dependency restraint. --- # 44. DDP profiling Profile distributed training separately from single-GPU training. Measure: ```text computation all-reduce / communication data loading rank imbalance synchronization ``` Use PyTorch profiler or other standard mechanisms where practical. Do not assume poor scaling is caused by NCCL before measuring. --- # 45. DDP scaling efficiency For available Perlmutter hardware, benchmark at least conceptually: ```text 1 GPU 2 GPUs 4 GPUs ``` and, if practical: ```text multiple nodes ``` Report: ```text throughput speedup scaling efficiency ``` For example: ```text speedup(N) = throughput(N) / throughput(1) efficiency(N) = speedup(N) / N ``` Do not treat perfect linear scaling as a correctness requirement. --- # 46. Fair DDP comparisons Be explicit about whether scaling benchmarks use: ```text fixed per-GPU batch size ``` or: ```text fixed global batch size ``` Do not compare inconsistent workloads without saying so. Prefer fixed per-GPU batch for throughput scaling measurements, while documenting the resulting larger global batch. --- # 47. Rank imbalance Profile rank step times. Determine whether some ranks systematically receive more expensive graph batches because of varying node/edge counts. If graph-size imbalance is material, document it. Do not immediately redesign the sampler in this task unless a simple semantics-preserving improvement is available. --- # 48. Graph-size-aware batching If profiling proves that highly variable graph size is a major performance problem, evaluate a simple size-aware batching strategy. Examples might include grouping events by approximate: ```text node count edge count ``` before batching. However, this can change sample ordering and stochastic training behavior. Do not make it the default unless: * semantics are understood * reproducibility remains explicit * parity expectations are updated appropriately * throughput improves materially Prefer documenting this as a follow-up if it becomes a significant algorithmic/data-loader change. --- # 49. Padding efficiency If Task 7 retains an active padding mode, measure: ```text real nodes vs padded nodes real edges vs padded edges ``` and quantify wasted work. Do not remove a compatibility-required padding mode merely because it is inefficient. If a more efficient mode can be optional, benchmark and document it separately. --- # 50. Memory profiling Measure memory at important stages. For CPU where practical: ```text dataset/cache loading batch construction output accumulation ``` For CUDA measure: ```text allocated memory reserved memory peak memory ``` using PyTorch-supported APIs. Record peak training and inference memory for representative benchmarks. --- # 51. Detect obvious retained tensors Inspect the training and inference loops for accidentally retained computation graphs. Examples: ```text storing non-detached loss tensors storing GPU logits for entire epochs keeping graph references in history ``` Fix such issues if found. Add regression tests when practical. --- # 52. Epoch metric accumulation memory Task 10/14 may collect logits/targets for full-split metrics. Profile the memory cost. Keep the current correct behavior as baseline. If this becomes a bottleneck, consider moving accumulated tensors to CPU as soon as possible. Do not redesign ROC AUC as an approximate streaming metric in this task. --- # 53. Inference memory Profile `PredictionResult` accumulation. Ensure predictions are detached and moved to CPU. Do not keep all results on GPU. If large-result memory is identified as a serious issue, document streaming/sharded output as a future task rather than over-expanding Task 15. --- # 54. Inference throughput Benchmark: ```text batch size DataLoader workers device transfer forward postprocessing output accumulation ``` separately. Do not include ROOT/NPZ serialization time in model throughput numbers unless explicitly measuring end-to-end inference. --- # 55. Output serialization benchmark Separately measure: ```text NPZ writing ROOT score writing ``` for representative output sizes. Do not optimize writers unless serialization is a meaningful bottleneck in actual inference. --- # 56. CPU-to-GPU overlap If profiling shows the GPU waiting significantly for data and pinned/non-blocking transfers are functioning, evaluate whether the existing DataLoader naturally overlaps preprocessing with GPU execution. Do not create a custom CUDA stream/prefetch framework unless simpler DataLoader improvements are insufficient. If advanced prefetching appears necessary, document it as a follow-up. --- # 57. Avoid premature micro-optimizations Do not spend significant code complexity optimizing components contributing negligible runtime. Prioritize the top measured bottlenecks. Prefer: ```text 10% simple improvement ``` over: ```text 1% improvement with substantial architectural complexity ``` unless the latter is scientifically/operationally important. --- # 58. Benchmark before/after every retained optimization For each production optimization retained in the final diff, record: ```text baseline optimized relative improvement benchmark scenario device/environment ``` Do not claim improvements without before/after numbers. --- # 59. Remove unsuccessful experiments Do not leave speculative optimization code disabled throughout the repository. If an experiment does not provide a meaningful benefit: * revert it * document the benchmark result briefly if useful Keep production code simple. --- # 60. Performance configuration Only add config options for optimizations that are actually retained. Possible examples: ```yaml data: pin_memory: true persistent_workers: true prefetch_factor: 2 trainer: precision: float32 compile: enabled: false ``` Do not turn every internal implementation detail into YAML. --- # 61. Conservative defaults Do not change baseline defaults merely because a particular Perlmutter benchmark was faster. Defaults should remain: ```text portable predictable correct ``` Performance profiles may override settings for the target machine. --- # 62. Perlmutter performance profile If useful, extend: ```text configs/environment/perlmutter.yaml ``` or add a dedicated performance-oriented trainer/data config. For example, a profile could set appropriate: ```text num_workers pin_memory persistent_workers precision ``` based on actual measurements. Do not hardcode values without benchmark evidence. --- # 63. Debug/local profile Keep local/debug configs lightweight. Do not apply aggressive Perlmutter worker/GPU settings to ordinary development. --- # 64. Scientific parity after optimization Run existing parity tests after every meaningful production-code optimization. At minimum preserve: ```text feature values edge topology/features model float32 outputs loss semantics metric semantics checkpoint semantics inference sample alignment ``` Performance changes must not invalidate these contracts. --- # 65. Deterministic model parity For float32 eager-mode optimizations, compare fixed-weight model outputs before and after changes. Use established parity tolerances. Do not weaken tolerances merely to allow an unnecessary optimization. --- # 66. Training-step parity Where optimizer/training internals change, such as: ```text zero_grad(set_to_none=True) fused/foreach optimizer path ``` test at least one deterministic training step against the baseline. Confirm parameter updates remain equivalent within justified tolerances. --- # 67. Batch-size invariance Performance changes to batching/device movement must not break Task 12 batch-size invariance. Run the existing inference tests. --- # 68. DDP correctness after optimization Run Task 14 distributed tests after retained changes. Do not optimize single-GPU behavior at the cost of incorrect distributed training. In particular verify: ```text global loss metrics checkpoint writing fine-tuning ``` still work under DDP. --- # 69. No algorithmic model redesign Do not change: ```text message-passing equations number of processing steps aggregation function pooling hidden dimensions classifier structure ``` for performance. Those would be new model experiments, not implementation optimization. --- # 70. No scientific preprocessing redesign Do not change: ```text selected objects feature definitions feature ordering energy calculation phi wrapping graph connectivity fold/split semantics event weights ``` for speed. --- # 71. Benchmark command examples Document commands in: ```text benchmarks/README.md ``` For example: ```bash uv run python benchmarks/benchmark_preprocessing.py uv run python benchmarks/benchmark_dataloader.py uv run python benchmarks/benchmark_training.py --device cpu uv run python benchmarks/benchmark_training.py --device cuda uv run python benchmarks/benchmark_inference.py --device cuda ``` Use the actual final argument interface. --- # 72. Optional CLI benchmark command Do not add: ```bash gnn4colliders benchmark ``` unless it provides clear value. Standalone benchmark scripts are sufficient for Task 15. Keep the user-facing scientific CLI focused. --- # 73. Perlmutter benchmark scripts If helpful, add a small Slurm benchmark launcher under: ```text scripts/slurm/ ``` such as: ```text benchmark_single_gpu.sh benchmark_multi_gpu.sh ``` Keep them thin. Do not duplicate benchmark logic in shell. --- # 74. Benchmark artifacts Do not commit large profiler traces or benchmark output files. Add generated profiling artifacts to `.gitignore` where appropriate. Examples: ```text *.pt.trace.json profiles/ benchmark-results/ ``` Do not ignore source benchmark scripts. --- # 75. Unit tests Performance scripts themselves do not need extensive unit tests. However, add regression tests for production-code changes. Examples: ```text optimized graph topology cache returns correct indices non-blocking batch transfer preserves fields AMP configuration validation works compile configuration does not affect default path ``` Focus testing on correctness, not timing thresholds. --- # 76. Do not add timing assertions to ordinary tests Do not add brittle tests such as: ```python assert runtime < 0.1 ``` to the unit/integration suite. Performance depends on hardware and CI load. Benchmarks measure speed; tests verify semantics. --- # 77. Benchmark smoke tests If useful, add only lightweight import/smoke coverage verifying benchmark scripts can initialize their argument/config setup. Do not run meaningful performance workloads in CI. --- # 78. Documentation Update: ```text docs/architecture.md ``` only where retained optimizations alter implementation architecture. For example document: ```text pinned/non-blocking transfer boundary optional precision mode optional compile mode topology cache ``` Do not fill architecture documentation with benchmark tables. --- # 79. Performance documentation Add: ```text docs/performance.md ``` or use: ```text benchmarks/README.md ``` to record: * benchmark methodology * representative hardware * baseline results * optimized results * retained optimizations * rejected optimizations * known bottlenecks * recommended Perlmutter settings Keep results clearly tied to hardware/configuration. --- # 80. README update Add a short performance section to README. Point to the detailed benchmark documentation. Include practical recommended invocation/config examples where validated. Do not advertise performance numbers without stating the hardware/configuration used. --- # 81. AGENTS.md Add concise durable performance rules if absent: ```text Profile before optimizing. Do not change scientific semantics for performance. Benchmark retained optimizations before and after. Keep performance features optional unless portability is proven. Do not add timing thresholds to ordinary tests. ``` Avoid duplicating detailed benchmark documentation. --- # 82. Dependency restraint Prefer existing dependencies and Python/PyTorch tooling. Do not add: ```text Numba CuPy Triton DeepSpeed custom CUDA extensions external profilers as mandatory dependencies ``` without strong measured justification. If an external profiling tool is useful manually, document it rather than making it a runtime dependency. --- # 83. Benchmark ROOT-GNN fine-tuning Include at least one representative fine-tuning benchmark. Measure: ```text frozen backbone unfrozen backbone ``` if both are real workflows. This may reveal substantially different backward/optimizer costs. --- # 84. Benchmark multiclass pretraining Include the active multiclass pretraining path. Use the standard active output size and representative model configuration. Do not benchmark only the smaller binary fine-tuning model. --- # 85. Frozen-backbone optimization When the fine-tuning backbone is frozen, verify unnecessary autograd work is not being performed. Parameters with: ```python requires_grad = False ``` should not accumulate gradients. Do not add special-case detached forward logic unless it is correct and measurably beneficial. Remember that gradients may still be required through backbone activations depending on what is trainable downstream; reason carefully before detaching anything. --- # 86. Inference-mode correctness Ensure inference uses: ```python torch.inference_mode() ``` where already appropriate. If Task 12 uses `no_grad()` and changing to `inference_mode()` is compatible, benchmark the difference. Keep the change only if safe. --- # 87. Avoid repeated model/device setup Audit CLI/training/inference paths for repeated: ```text model.to(device) checkpoint reload task reconstruction ``` inside batch/epoch loops. Fix obvious repeated setup if found. Add regression coverage where useful. --- # 88. Avoid repeated configuration parsing Configuration resolution should happen at application startup, not per batch. Do not micro-optimize Hydra itself unless profiling somehow proves it is in the training hot path, which it should not be. --- # 89. Avoid repeated filesystem checks in hot loops Audit for repeated: ```text Path.exists() cache metadata reads checkpoint discovery ``` inside batch loops. Move invariant filesystem work outside hot execution loops. Do not weaken correctness checks at application boundaries. --- # 90. Performance logging If useful, add optional per-epoch throughput reporting to the trainer. For example: ```text samples/sec batches/sec epoch duration ``` Keep this lightweight. Do not synchronize CUDA unnecessarily every batch just to produce timing logs. Use coarse epoch-level wall time for ordinary logging. --- # 91. Distributed throughput logging Under DDP, report global throughput correctly. For example: ```text global samples processed / epoch wall-clock time ``` Do not multiply throughput incorrectly if sampler padding/duplicate handling is present. Only rank 0 should report ordinary performance summaries. --- # 92. Warmup exclusions For GPU benchmark measurements, exclude initial warmup where appropriate. For actual end-to-end startup benchmark, include it explicitly as a separate number. Do not mix the two. --- # 93. Synchronization overhead Avoid introducing unnecessary: ```python torch.cuda.synchronize() dist.barrier() ``` into production hot loops. Synchronization may be needed for profiling but should not remain in normal execution unless required for correctness. --- # 94. Memory cleanup Do not add routine: ```python torch.cuda.empty_cache() ``` inside training loops as a supposed optimization. Use it only for explicit diagnostic/phase-boundary reasons if needed. Let PyTorch's allocator manage normal execution. --- # 95. Benchmark reporting At the end of the task, summarize results in a compact table. For each retained change report something like: ```text Optimization | Scenario | Before | After | Improvement | Memory impact ``` Also list tested changes that were rejected. Do not cherry-pick only successful results. --- # 96. Prioritize simplicity If two implementations provide similar performance, keep the simpler implementation. Do not sacrifice maintainability for marginal gains. This repository is intended to support future model families such as ROOT-Transformer, so shared performance infrastructure should remain architecture-neutral where reasonable. --- # 97. Future model compatibility Do not introduce graph-specific assumptions into shared: ```text data training distributed inference ``` layers solely for ROOT-GNN optimization. Graph-specific optimizations belong under: ```text graphs/ models/root_gnn/ ``` where possible. Keep the future flow viable: ```text EventSample -> shared features ├── GraphSample -> ROOT-GNN └── SequenceSample -> ROOT-Transformer ``` --- # 98. Validation Run the correctness suite before and after performance changes: ```bash uv run pytest ``` Run focused parity tests: ```bash uv run pytest tests/parity -v ``` Run distributed tests: ```bash uv run pytest tests/unit/distributed -v uv run pytest tests/integration -k distributed -v ``` Run lint/format checks: ```bash uv run ruff check src tests benchmarks uv run ruff format --check src tests benchmarks ``` Run representative benchmarks. At minimum: ```bash uv run python benchmarks/benchmark_preprocessing.py uv run python benchmarks/benchmark_dataloader.py uv run python benchmarks/benchmark_training.py --device cpu ``` On a CUDA environment also run the appropriate: ```bash uv run python benchmarks/benchmark_training.py --device cuda uv run python benchmarks/benchmark_inference.py --device cuda ``` On Perlmutter, run the validated single-GPU benchmark and multi-GPU benchmark where available. Inspect: ```bash git status git diff ``` Verify: * legacy source remains unchanged * scientific parity is preserved * no experimental architecture changes were introduced * no unmeasured speculative optimization remains * optional performance paths do not change default semantics * no profiler output artifacts are accidentally tracked * no personal Perlmutter paths/accounts were added * no unrelated changes are included --- # Completion criteria Task 15 is complete when: 1. Reproducible benchmark scripts exist. 2. Baseline preprocessing performance is measured. 3. Baseline DataLoader performance is measured. 4. Baseline model inference performance is measured. 5. Baseline training-step performance is measured. 6. At least representative CPU behavior is characterized. 7. Representative GPU behavior is characterized where the environment permits. 8. PyTorch profiling identifies the main runtime bottlenecks. 9. Peak memory behavior is measured for representative training/inference. 10. Retained optimizations have before/after benchmark evidence. 11. Unsuccessful optimization experiments are removed from production code. 12. Feature-construction parity remains intact. 13. Graph topology/edge-feature parity remains intact. 14. Model float32 parity remains intact. 15. Loss/metric semantics remain intact. 16. Inference output/sample alignment remains intact. 17. Checkpoint compatibility remains intact. 18. Single-process training remains correct. 19. DDP training remains correct. 20. Distributed global metrics remain correct. 21. Optional DataLoader performance improvements are implemented if beneficial. 22. Optional host-to-device transfer improvements are implemented if beneficial. 23. Optional optimizer implementation improvements are implemented only if semantics remain valid. 24. Mixed precision is characterized and optionally supported if beneficial. 25. `torch.compile` is characterized and retained only if beneficial and compatible. 26. Recommended Perlmutter performance settings are documented. 27. Performance results clearly identify hardware and configuration. 28. Full tests pass. 29. Lint/format checks pass. 30. Legacy code remains untouched. --- # Completion report Report: 1. files created 2. files modified 3. benchmark structure 4. benchmark methodology 5. benchmark hardware/environment 6. baseline feature-construction performance 7. baseline graph-construction performance 8. baseline cache performance 9. baseline DataLoader throughput 10. baseline CPU training performance 11. baseline GPU training performance, if measured 12. baseline inference performance 13. baseline peak memory usage 14. profiler-identified bottlenecks 15. retained feature/data optimizations 16. retained graph optimizations 17. retained DataLoader optimizations 18. retained device-transfer optimizations 19. retained model/training optimizations 20. AMP/bfloat16 results 21. `torch.compile` results 22. single-GPU before/after results 23. DDP scaling results 24. memory before/after results 25. rejected optimization experiments and why 26. parity/correctness results 27. recommended local defaults 28. recommended Perlmutter settings 29. remaining bottlenecks 30. suggested future performance work 31. validation commands and results After validation succeeds, create one Git commit containing only Task 15 changes. Use: ```text perf: profile and optimize ROOT-GNN execution ``` Before committing, inspect the final diff and ensure no unrelated files or generated profiling artifacts are included.