fix(megatron): refresh masters after weights-only load

[bvolpato]

A Megatron weights-only checkpoint load updates the model or LoRA tensors but leaves optimizer master weights at their previous values. The next optimizer step copies those stale masters back over the restored checkpoint. This affects Tinker load_state() / load_weights(optimizer=False) and complements the selective-resume controls in #2010.

Keep the live optimizer available to the strategy while the existing restore flags control optimizer history. After a weights-only load, refresh masters from the restored model. Use Megatron's native reload and fill the gaps in the pinned HybridDeviceOptimizer and Transformer Engine precision-aware paths, preserving Adam moments, counters, scheduler state, and parameter aliases. Full optimizer restores retain the checkpoint's master precision.

Reproduction and validation

uv run --isolated --extra skyrl-train --extra dev pytest \
  tests/backends/skyrl_train/distributed/test_megatron_checkpoint_weights_only.py \
  tests/backends/skyrl_train/test_checkpoint_loading.py -q

• The actual original Worker/Strategy fail both full-model and LoRA weights-only cases: the model loads 5.0, while its master remains 0.9. Both full-state controls pass.
• 15 CPU tests pass with the fix, covering a subsequent real Adam update, preserved history, full-state precision, LoRA, chained/empty shards, CPU offload, non-distributed offload, and precision-aware master/remainder storage. Distributed checkpoint I/O and unavailable GPU dependencies are mocked.
• Added a two-GPU full-training/LoRA checkpoint roundtrip. It reloads an earlier checkpoint and runs a real backward/optimizer window at zero learning rate, checking that master copy-back preserves the loaded logprobs:

uv run --isolated --extra megatron --extra dev pytest \
  tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_weights_only_checkpoint.py -s

The GPU test has not run on this host. Keep this draft pending distributed validation. Ruff, Black, and gitleaks hooks passed.

Downsides

Weights-only loading now performs the necessary local model-to-master copies. CPU-offloaded optimizers also refresh their host copies, and quantized parameters may require temporary dequantization. This cost occurs during checkpoint loading rather than each training step. Compatibility handling follows the currently pinned Megatron/Transformer Engine layouts and should be rechecked on dependency upgrades.

Risk and rollback

The main risk is incorrect master/shard mapping in a GPU optimizer variant; the CPU invariants and pending GPU roundtrip target that boundary. FSDP still honors its explicit optimizer restore flag. To avoid the affected path during rollback, restore optimizer state as well as weights, then revert the change.