Multi-stage freeze/unfreeze callback

[ryspark]

Adds a MultiStageFineTuning callback for flexible, epoch-based freeze/unfreeze schedules.

• Define which layers to freeze/unfreeze at specific epochs with freeze_selectors and unfreeze_selectors.
• Unfreeze can add new param groups at scaled LR (base_lr / unfreeze_lr_factor).
• Optionally scale existing group LRs (scale_existing_groups) when deeper layers are unfrozen.
• Keeps ReduceLROnPlateau.min_lrs aligned with optimizer groups.
• Each stage runs once and overrides previous freezes for resume-safe fine-tuning.
• Test coverage for the new callback since it's more complex than the previous FreezeUnfreeze.

Should be merged after #221 and this helios PR (the tests will fail until this one is merged due to a missing import from helios). The only new changes are to freeze_unfreeze.py and the corresponding test file, plus moe module was moved to helios so deleted here.

Example usage:

from rslearn.train.callbacks.freeze_unfreeze import MultiStageFineTuning

cb = MultiStageFineTuning([
    # Stage 0: Train head only
    {"at_epoch": 0,
     "freeze_selectors": ["backbone"],
     "unfreeze_selectors": ["head"]},

    # Stage 1: Unfreeze LoRA adapters at lower LR, scale head LR down
    # 'unfreeze_selectors' lets us unfreeze specific layers specified by 'freeze_selectors'
    {"at_epoch": 3,
     "freeze_selectors": ["backbone"],
     "unfreeze_selectors": ["lora_adapter", "head"],
     "unfreeze_lr_factor": 10.0,
     "scale_existing_groups": 0.5},

    # Stage 2: Train full model
    {"at_epoch": 6,
     "freeze_selectors": [],
     "unfreeze_selectors": ["backbone", "head"],  # unfreeze everything
     "unfreeze_lr_factor": 1.0},
])