Inference Pipeline Optimization - hpy_splits

[javak87]

Description

Inference Pipeline Optimization

Tokenizer CPU time was dominated by hpy_splits: a Python loop over every HEALPix cell (12 * 4**hl) with a tiny torch.argsort / cat / split per cell. This replaces that with a vectorized NumPy path while preserving the same token index lists.

• hpy_cell_splits: Uses one np.lexsort((theta, cell_id)) to group points by cell and sort them by latitude within each cell. Occupancy is computed using np.bincount(..., minlength=12 * 4**hl). It returns packed (idxs_ord, counts) instead of a Python list of per-cell arrays.
• hpy_splits: Scatters the packed indices into one padded buffer, splits once into tokens, and then slices the tokens back per cell. The same padding (index 0) and offset_step behavior are preserved.
• tokenize_space / tokenize_spacetime: Unchanged; they continue to consume the same nested token lists.
• Tests: tests/test_hpy_splits.py freezes the original implementations and checks torch.equal on tokens for padded and unpadded cases, offset_step, tied coordinates, remainder tokens, and hl=5.

hpy_cell_splits is only used by hpy_splits in this tree.

Inference Timing

Environment: Jupiter, 1 node, mini-epoch 8 from cw6a4szu

The following command was used for both branches. test_config.output.num_samples=0 is required to avoid OOM.

../WeatherGenerator-private/hpc/launch-slurm.py --stage inference --from-run-id cw6a4szu \
  --mini-epoch 8 --nodes 1 --time 15:00 \
  --options test_config.output.num_samples=0 test_config.forecast.num_steps=8 \
  test_config.samples_per_mini_epoch=4 test_config.validation_noise_levels=[] \
  data_loading.num_workers=1 data_loading.num_workers_validation=1 \
  test_config.compute_full_validation_loss=False

Branch	Runs (s/it)	Average (s/it)
develop-ssl-diffusion-v1	u6xdzoq8 60.50, xhohdobd 58.61, mjijni1u 58.84	59.15
This branch	j6dsiw2c 40.81, n0f2bw1i 41.32, d3gonkm4 41.74	41.29

About 30% faster. Validation loss on j6dsiw2c: 0.0329427495598793 (unchanged vs baseline).

Note: These timings only show up if the job tree is running javad/optimize-hpy_splits branch on private repo. launch-slurm.py --from-run-id cw6a4szu copies Python from that training run, not from your checkout, so launching from the optimized branch alone does not change the tokenizer unless those files are overlaid (or otherwise copied) into the new job directory.

Issue Number

#2778

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

FastEvaluation

• I have updated the public documentation if necessary

[clessig]

I tested the code and get:

Old: 17:39:13 -> 17:47:51
New: 17:28:39 -> 17:37:13

Different branch, though: clessig/preop_082026/inference_no_targets

[javak87]

I tested the code and get:

Old: 17:39:13 -> 17:47:51 New: 17:28:39 -> 17:37:13

Different branch, though: clessig/preop_082026/inference_no_targets

Thanks for the test.
Please make sure you are in this private repo branch: javad/optimize-hpy_splits

[javak87]

Close the PR due to incorrect timing!
Open the new PR with correct timing:
#2791