Rename land_sea_mask.nc to land_mask.nc

[sgpearse]

Fix #146

[azrael417]

Hello, thanks for filing the PR. I have pushed some changes to your fork. I realized that there is an asymmetry between the loading and storing of these invariants and their stats. I asked Claude to fix these and pushed to your PR branch, I hope that is OK. Would you mind reviewing this and see if everything works for you on this branch now?

[sgpearse]

Thanks @azrael417. Here's what I'm seeing when trying to test it.

Setup: fresh venv, earth2studio[fcn3]==0.16.0, this branch installed via pip install git+@fix-land-mask-filename, torch 2.6.0+cu124, torch-harmonics built from main (needed for the fused param and renamed precompute_latitudes this branch requires).

Result: from_pretrained() downloads weights fine (including the corrected land_mask.nc), but fails at load_state_dict with:

Missing key(s): "model.normalization_means", "model.normalization_stds", "model.encoder.conv.weight", "model.decoder.conv.weight"
Unexpected key(s): "model.atmo_encoder.conv.weight", "model.surf_encoder.conv.weight", "model.atmo_decoder.conv.weight", "model.surf_decoder.conv.weight"

And many channel-size mismatches throughout (e.g. [1,677,677,360] vs [1,44,44,360]).

I confirmed this reproduces identically on unmodified main with the same checkpoint, so it's not something introduced by the changes in this PR — it looks like nvidia/fourcastnet3's best_ckpt_mp0.tar was saved under the older split atmo/surf encoder-decoder architecture, and current main has since consolidated to a single encoder/decoder.

I think the landmask/normalization fix itself is fine, but I can't exercise it end-to-end with the currently published checkpoint. Could you point me to a checkpoint that matches current main, or a commit I should pin makani to that still matches best_ckpt_mp0.tar's architecture? Want to give this a real test rather than just eyeballing the diff.

[azrael417]

Thanks for reporting this, we will fix this asap.

[azrael417]

This issue is tracked here: #148. Merging this now.

[azrael417]

Can you please check the branch belonging to this MR: #149
I verified that it can load the HF model correctly, but I need to wait for my colleague to validate the full architecture. But please have a look if that works. We will merge it as soon as everything is validated.

[sgpearse]

I will test tomorrow and let you know. Thanks!

[sgpearse]

Confirmed the branch loads the HF checkpoint correctly on our end with L40 and CUDA 12.7. The land_sea_mask and encoder/decoder split do load without error.

Inference is still in progress on our side. I've hit GPU memory contention from another process during decode and not from the model itself. I can re-test with a clean GPU to confirm end-to-end if needed. Thank you for the fix!

[azrael417]

That is good to hear. We did an inference run from within makani and the results look reasonable. Let me know if you encounter any more issues. Just open a new issue in this case. If it is a persistent OOM, then it is maybe on the inference side and maybe open one against earth2 studio.