Questions about finetuning Aurora on 9 new variables (~3 days / epoch on an H100 95 GB)

Hello, thank you for the amazing work! I’m trying to finetune Aurora on nine additional surface variables and training is extremely slow:
≈ 8 min per step ⇒ ~3 days per epoch on an NVIDIA H100 (95 GB), despite batch_size=1 and gradient accumulation over 8 steps.

Training setup

Questions

1. Aurora‑Small “not recommended”
   Docs say to avoid the Small model.

   • What is the underlying reason? Is it the Architecture, stability, or just accuracy?
   • Can it still converge with more epochs?

2. Loss on denormalised values
   Predictions are denormalised before loss computation.

   • Would computing loss in normalised space (and perhaps dropping per‑variable weighting) improve gradient stability?
   • Any best‑practice guidance?

3. Variable scale imbalance
   New variables span orders of magnitude (MAE/MSE from ~10⁶ to < 0.01).

   • I used the paper’s two‑stage weighting: all weights = 1 (initial run, on a sub-set of the data for 10 epochs) → Calculate the weight (1/MAE) for the new variables.
   • Should I also recompute weights for the original variables?
   • Are other approaches (e.g. log‑scaling) better suited for Aurora?

4. Regional fine‑tuning
   Has anyone successfully trained on a geographic subset (e.g. Europe only)?

   • The model accepts any input shape divisible by patch_size (default = 4), so in theory this should work.
   • Are there caveats or recommended settings for regional training?

5. Memory usage seemingly independent of variable count
   GPU memory stays roughly constant no matter how many variables I finetune on.
   • Even when I removed the original surface + atmospheric variables and trained only on my 9 new variables, memory usage was unchanged. Is this expected?
   • Is it recommended to finetune exclusively on new variables, or must the original set always be included? *(The atmospheric variables for 2024 at 4 timesteps/day already occupy ~300 GB as a single NetCDF!)
   • Any tips for reducing memory footprint?

Start with the Aurora fine-tuning documentation and the paper, then reproduce the reported setup: aurora-base and aurora-small, batch size 1, gradient accumulation, and nine added variables. Trace the training, loss, regional-input, and memory behavior described in the issue; done means documenting confirmed causes, supported settings, and recommended fine-tuning guidance.