Enable load_weights_only in checkpointing

[j316chuck]

Summary

• Preserve Tinker's optimizer flag through the load-weights API, engine, and JAX/SkyRL-Train backends.
• Restore optimizer and scheduler state only for load_state_with_optimizer(); weights-only loads leave the live optimizer untouched.
• Keep Tinker's public optimizer field for SDK compatibility while using load_optimizer internally.

Fixes #1993.

Testing

uv run --isolated --extra dev --extra tinker --extra jax pytest tests/tinker/test_api_validation.py tests/tinker/test_engine.py tests/tinker/test_api.py::test_training_workflow -q
uv run --isolated --extra skyrl-train --extra dev pytest tests/backends/skyrl_train/test_checkpoint_loading.py -q

The API, engine, SDK workflow, and backend checkpoint tests pass. In a live H200 deployment, a cold one-step job saved a checkpoint and a second job on the same direct endpoint loaded it with load_optimizer=False, completed training, and saved a new checkpoint. The weights-only /load_weights call took 0.841s.

[gemini-code-assist]

Code Review

This pull request introduces an "optimizer" boolean flag to the checkpoint loading process across multiple backends (JAX, Ray JAX, and SkyRL Train) and exposes it via the Tinker API, enabling optional restoration of optimizer and scheduler states. The feedback highlights two main areas of improvement: first, the "optimizer" field in "LoadWeightsRequest" and "LoadWeightsInput" should default to "False" to preserve backward compatibility with older clients; second, a safety check should be added in the JAX backend to prevent a potential "KeyError" if "optimizer_state" is missing from the checkpoint data.

[gemini-code-assist]

Adding optimizer as a required boolean field in LoadWeightsRequest without a default value breaks backward compatibility with older clients or SDKs that do not send this field in their payload. Providing a default value (e.g., False) ensures that the API remains backward-compatible and robust against missing fields.

Suggested change

	optimizer: bool
	optimizer: bool = False

[kalectory]

clean