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:
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:
AGENTS.md
README.md
docs/architecture.md
docs/migration.md
Then inspect:
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:
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:
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:
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:
events / second
graphs / second
batches / second
training steps / second
milliseconds / batch
seconds / epoch
Where relevant also report:
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:
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:
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:
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:
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:
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:
torch.profiler
Capture enough iterations to distinguish:
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:
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:
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:
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:
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:
deta
dphi
dR
calculation.
Look for:
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:
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:
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:
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:
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:
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:
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:
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:
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:
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:
graph features
global features
labels during training/evaluation
weights required by the loss
need device access.
Metadata such as:
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:
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:
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:
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:
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:
removing avoidable Python overhead
avoiding repeated module construction
using efficient contiguous tensors where needed
Do not change:
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:
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:
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:
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:
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:
trainer:
precision: float32
with supported alternatives such as:
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:
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:
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:
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:
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:
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:
1 GPU
2 GPUs
4 GPUs
and, if practical:
multiple nodes
Report:
throughput
speedup
scaling efficiency
For example:
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:
fixed per-GPU batch size
or:
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:
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:
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:
dataset/cache loading
batch construction
output accumulation
For CUDA measure:
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:
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:
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:
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:
10% simple improvement
over:
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:
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:
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:
portable
predictable
correct
Performance profiles may override settings for the target machine.
62. Perlmutter performance profile
If useful, extend:
configs/environment/perlmutter.yaml
or add a dedicated performance-oriented trainer/data config.
For example, a profile could set appropriate:
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:
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:
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:
global loss
metrics
checkpoint writing
fine-tuning
still work under DDP.
69. No algorithmic model redesign
Do not change:
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:
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:
benchmarks/README.md
For example:
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:
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:
scripts/slurm/
such as:
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:
*.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:
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:
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:
docs/architecture.md
only where retained optimizations alter implementation architecture.
For example document:
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:
docs/performance.md
or use:
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:
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:
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:
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:
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:
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:
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:
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:
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:
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:
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:
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:
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:
data
training
distributed
inference
layers solely for ROOT-GNN optimization.
Graph-specific optimizations belong under:
graphs/
models/root_gnn/
where possible.
Keep the future flow viable:
EventSample
->
shared features
├── GraphSample -> ROOT-GNN
└── SequenceSample -> ROOT-Transformer
98. Validation
Run the correctness suite before and after performance changes:
uv run pytest
Run focused parity tests:
uv run pytest tests/parity -v
Run distributed tests:
uv run pytest tests/unit/distributed -v
uv run pytest tests/integration -k distributed -v
Run lint/format checks:
uv run ruff check src tests benchmarks
uv run ruff format --check src tests benchmarks
Run representative benchmarks.
At minimum:
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:
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:
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:
- Reproducible benchmark scripts exist.
- Baseline preprocessing performance is measured.
- Baseline DataLoader performance is measured.
- Baseline model inference performance is measured.
- Baseline training-step performance is measured.
- At least representative CPU behavior is characterized.
- Representative GPU behavior is characterized where the environment permits.
- PyTorch profiling identifies the main runtime bottlenecks.
- Peak memory behavior is measured for representative training/inference.
- Retained optimizations have before/after benchmark evidence.
- Unsuccessful optimization experiments are removed from production code.
- Feature-construction parity remains intact.
- Graph topology/edge-feature parity remains intact.
- Model float32 parity remains intact.
- Loss/metric semantics remain intact.
- Inference output/sample alignment remains intact.
- Checkpoint compatibility remains intact.
- Single-process training remains correct.
- DDP training remains correct.
- Distributed global metrics remain correct.
- Optional DataLoader performance improvements are implemented if beneficial.
- Optional host-to-device transfer improvements are implemented if beneficial.
- Optional optimizer implementation improvements are implemented only if semantics remain valid.
- Mixed precision is characterized and optionally supported if beneficial.
torch.compileis characterized and retained only if beneficial and compatible.- Recommended Perlmutter performance settings are documented.
- Performance results clearly identify hardware and configuration.
- Full tests pass.
- Lint/format checks pass.
- Legacy code remains untouched.
Completion report
Report:
- files created
- files modified
- benchmark structure
- benchmark methodology
- benchmark hardware/environment
- baseline feature-construction performance
- baseline graph-construction performance
- baseline cache performance
- baseline DataLoader throughput
- baseline CPU training performance
- baseline GPU training performance, if measured
- baseline inference performance
- baseline peak memory usage
- profiler-identified bottlenecks
- retained feature/data optimizations
- retained graph optimizations
- retained DataLoader optimizations
- retained device-transfer optimizations
- retained model/training optimizations
- AMP/bfloat16 results
torch.compileresults- single-GPU before/after results
- DDP scaling results
- memory before/after results
- rejected optimization experiments and why
- parity/correctness results
- recommended local defaults
- recommended Perlmutter settings
- remaining bottlenecks
- suggested future performance work
- validation commands and results
After validation succeeds, create one Git commit containing only Task 15 changes.
Use:
perf: profile and optimize ROOT-GNN execution
Before committing, inspect the final diff and ensure no unrelated files or generated profiling artifacts are included.