Disable corrector for the first N epochs

[elynnwu]

Adds a corrector_disabled_epochs option to corrector configs that skips the corrector for train-mode steps during the first N training epochs. It is always applied in eval mode (validation and inference), so only training is affected.

Some runs get frozen precipitation stuck at exactly zero (e.g. job): the force-positive clamp zeroes the sparse frozen precip channel whenever an early update makes the raw prediction negative, and gradients through the clamp are zero there, so it never recovers. Disabling the corrector for the first epoch lets the loss see the raw prediction so gradients can pull the channel positive before the clamp returns.

corrector_disabled_epochs is a field on CorrectorConfigABC; its @final get_corrector wraps the built corrector in a new EpochScheduledCorrector (which skips an arbitrary wrapped corrector in train mode during the disabled epochs) when the value is > 0. Concrete configs just implement _get_corrector and are unaware of scheduling. The flag is valid only on the corrector config — CorrectorSelector rejects it and delegates — so there's one place to set it and no double-wrapping. The disabled flag is persisted in checkpoint state so mid-epoch resume keeps the interrupted epoch's behavior. Defaults to 0, so existing configs and checkpoints are unaffected.

Config

corrector:
  conserve_dry_air: true
  force_positive_names:
  - frozen_precipitation_rate
  corrector_disabled_epochs: 1   # inside `config:` for selector-based correctors

Changes

• CorrectorConfigABC: corrector_disabled_epochs field + @final get_corrector wrapping via new abstract _get_corrector; CorrectorABC gains default no-op train/set_epoch/get_state/load_state; new EpochScheduledCorrector wrapper.
• CorrectorSelector rejects corrector_disabled_epochs (must be set on the wrapped config).
• Atmosphere/Ocean/Ice configs implement _get_corrector; step layer + Stepper forward set_epoch/train/get_state/load_state to the corrector.

[elynnwu]

These two jobs are identical except the new job disabled corrector for the first epoch: wandb. The new job does not have the same zero frozen precipitation issue. Note that I was able to reproduce the zero frozen precip behavior in 3 experiments, all with n384 (one with equal loss, two with ACE2 loss weighting on Beaker and Perlmutter).

[mcgibbon]

LGTM, but consider the corrector refactor sooner rather than later. I can also act on it if you like, a Claude agent pointed to your issue and this PR should be able to handle it.