Support weights-only checkpoint loading

[j316chuck]

What Tinker supports

Tinker 0.23 intentionally exposes two checkpoint-loading APIs:

• TrainingClient.load_state(...) sends optimizer=False: restore model weights and reset optimizer state.
• TrainingClient.load_state_with_optimizer(...) sends optimizer=True: restore model weights and optimizer state.

Weights-only loading is needed for warm starts and stage transitions where parameters should transfer but stale Adam momentum and scheduler state should not.

What SkyRL is missing

SkyRL currently loses this distinction:

1. LoadWeightsRequest does not expose the optimizer flag.
2. The load_weights endpoint creates LoadWeightsInput without it.
3. SkyRLTrainBackend.load_checkpoint hardcodes load_optimizer_states=True and load_lr_scheduler_states=True.

As a result, load_state() and load_state_with_optimizer() both restore optimizer and scheduler state on the SkyRL backend.

Feature Request

• Add the load_optimizer boolean to SkyRL load-weight request/input types.
• Thread it through the API and engine to load_checkpoint.
• Restore optimizer and scheduler state only when load_optimizer=True.
• Add a regression test that first creates nonzero Adam moments, then verifies:
  • load_state() resets exp_avg, exp_avg_sq, and scheduler state.
  • load_state_with_optimizer() preserves them.

[j316chuck]

One additional resume-path feature request:

SkyRL does not expose a server-side create_model_from_state_with_optimizer operation. Tinker’s create_training_client_from_state_with_optimizer_async is a client-side convenience workflow that:

1. reads checkpoint metadata,
2. creates the LoRA training client/model,
3. calls load_state_with_optimizer_async, and
4. awaits the load future.

There could also be a startup-latency speed up in a server-side combined operation. Checkpoint lookup, download, and staging can run concurrently with Megatron model/optimizer construction and vLLM engine startup when their inputs are already known. Applying weights and optimizer state must still synchronize after model creation but overlapping the independent initialization and checkpoint-I/O phases could materially reduce resumed-job startup time.