Skip to content

Coalesce TorchTrainDataset rollout loads - #799

Draft
jder wants to merge 1 commit into
mainfrom
agent/coalesce-torch-train-dataset-loads
Draft

Coalesce TorchTrainDataset rollout loads#799
jder wants to merge 1 commit into
mainfrom
agent/coalesce-torch-train-dataset-loads

Conversation

@jder

@jder jder commented Jul 16, 2026

Copy link
Copy Markdown
Member

Summary

  • Confirmed TorchTrainDataset was materializing each autoregressive rollout step independently.
  • For shared prognostic sources, the step-N label range is the same data as the step-(N+1) input range, so non-dask loading duplicated adjacent time reads.
  • Coalesce each sample's rollout time range into one selected prognostic tensor when input and label sources are shared, then slice that tensor into per-step examples.
  • Keep split-source paths for cross-resolution training, while still selecting each input, label, and boundary range once per sample.
  • Add a regression test that verifies shared-source training selects one steps + 1 prognostic range and preserves the overlapping label/input values.

Validation

  • uv run pytest tests/test_datasets.py -k 'train_dataset_coalesces_shared_prognostic_loads or train_dataset_normalize_pre_fill or train_dataset_no_input_change'
  • uv run pytest tests/test_datasets.py -k 'loader__data_shape__across_schedules or train_dataset_coalesces_shared_prognostic_loads'
  • uvx ruff format --check src/samudra/datasets.py tests/test_datasets.py
  • uvx ruff check src/samudra/datasets.py tests/test_datasets.py
  • uvx pre-commit run --files src/samudra/datasets.py tests/test_datasets.py

Note: uv run pytest tests/test_datasets.py reached the unrelated samudra_multi inference setup and errored because flash_perceiver is not installed for auto flash mode in this environment; the training dataset tests before that point passed (117 passed, 12 errors).

Refs #798

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

3 participants