[zyhu-hu]

Summary

Updates the README for the public release of the STRATA v1 weights and the SCREAM training data, fixes two bugs found while running the released checkpoint from public files only, and makes all rollout examples use the recommended settings.

README

• Header: one row of badges (paper, project page, model, data, license) instead of a link list.
• Environment variables: one table of the variables each task needs (forecast, training, wandb, s3:// via an rclone remote); .env is optional.
• Data: the three nvidia/STRATA-SCREAM-* Zarr repositories with sizes, the store layout, reading a store over hf://, and the directory names the training configs expect under ZARR_ROOT. Partially downloaded stores read missing chunks as zeros.
• Auxiliary files: which files come from nvidia/strata, and which are generated with data_prep/scrip_generation/generate_all.sh or copied from data/.
• Running the STRATA v1 checkpoint:
  • valid initial times; the sdecadal run used October 2020 SSTs with a 1994 atmospheric initial condition, so forecasts are evaluated against SCREAM output, not observations;
  • one forecast command with the recommended settings (--n-steps 2 as a smoke test);
  • batch size and output size notes, and the output store layout;
  • SLURM and Perlmutter pointers.
• Training:
  • launch commands (train.py, torchrun, srun, experiment_manager.py);
  • the config the released checkpoint was trained with;
  • a note that fine-tuning from strata_v1.pth currently needs torch.compile off.
• License section: code Apache-2.0, weights OpenMDW-1.1, SCREAM data CC BY 4.0.
• Removed the "will be shared soon" wording and repeated explanations.

Fixes

• --output-steps N > 1 frame labels (screamcast/zarr_writer.py, scripts/ace/run_screamcast_nudged.py): the writers saved the forecasts after steps 1, N+1, 2N+1, … but the store's step coordinate labels frame k as lead (k+1)·N, so every saved frame was (N−1)·dt earlier than its label.
  • A new helper, output_frame_index(), now saves leads N, 2N, …, and both writers use it.
  • --output-steps must be ≥ 1 and divide --n-steps; this is checked before the model loads.
  • --output-steps 1 (the default) is unchanged.
• Precip log (run_screamcast_nudged.py): the per-step log converted the m/s fluxes with 86400 only, so it printed m/day under an mm/day label (1000× too small). It now uses 1000 * 86400, like the other scripts in the repo.
• slurm/submit_inference.sh sets PYTHONPATH from the submit directory. Without it, every rank failed to import data_catalog.
• CHANGELOG entry for the above.

Rollout settings

The README forecast, slurm/submit_inference.sh and perlmutter_scripts/strata_v1_rollout.sh now all use the recommended settings: tile jitter on (4 px, cycle16), --no-halo-adjoint-cross-face, and SHT omega lmax 3.

Testing

• Forecast from public downloads only (weights from nvidia/strata, initial condition read anonymously from hf://datasets/nvidia/STRATA-SCREAM-sdecadal, grid files from generate_all.sh) on 4 × H100 80 GB:
  • strict checkpoint load and finite outputs;
  • errors at 10 min to 1 h well below persistence;
  • torch.compile on and off both run.
• pytest tests/test_run_screamcast_nudged_write.py tests/test_regional_averages.py tests/test_zarr_writer.py: 28 passed.
  • The new test checks that every saved frame holds the forecast for the lead its step label names, for N = 1, 2, 3 and 6.
  • It fails on the old indexing.
• make lint passes; both shell scripts pass bash -n; the PYTHONPATH fix was checked in the container (import fails without it, works with it).
• Not run: the full pytest suite, tests/distributed, and an end-to-end forecast with the exact final flag set (jitter on together with --no-halo-adjoint-cross-face). Each setting was run separately.

Review

An independent review by four reviewers with no prior context (correctness, tests with a mutation check, claims audit, blind), with each serious finding re-checked by a separate verifier, found no blockers. The confirmed issues are fixed in this PR:

• training: validation store variables, and the torch.compile caveat for fine-tuning;
• the SLURM PYTHONPATH;
• the Perlmutter script description;
• the sdecadal date description;
• the CHANGELOG entry;
• --output-steps validation.

Mutations of the new frame-index function are caught by the tests. The script's legacy and pressure-level writer wiring is not tested directly.

🤖 Generated with Claude Code