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.

Fixes #1993.

Testing

uv run --extra tinker --extra dev pytest tests/tinker/test_api_validation.py::test_load_weights_request_preserves_optimizer_choice tests/tinker/test_engine.py::test_process_load_weights_forwards_optimizer_choice
uv run --extra tinker --extra skyrl-train --extra ray --extra dev pytest tests/backends/skyrl_train/test_checkpoint_loading.py

All seven focused API, engine, and checkpoint-loading cases pass.

[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