fix(megatron): snapshot precision-aware optimizer state

[hershg]

• Resolve precision-aware distributed-optimizer shards to their live FP32 master tensors before multi-LoRA snapshots. This fixes the None main-parameter failure observed during pristine registration in XID 1050929.
• Snapshot and restore FP32 masters, Adam state, and dynamic step counters for both standard and CPU-offloaded optimizer layouts.

Testing

uv run --isolated --extra dev --with ray pytest tests/backends/skyrl_train/workers/test_adapter_store_lifecycle.py -q
uv run --isolated --extra dev pre-commit run --all-files --config .pre-commit-config.yaml

Six focused tests pass; the full pre-commit suite is green.

[greptile-apps]

Confidence Score: 4/5

The PR is safe to merge with two non-blocking CI reproducibility and diagnostics issues.

The runtime changes have focused lifecycle and concurrency coverage; the accepted concerns are confined to the DeepGEMM CI invocation using ambient packages and not persisting its output.

Files Needing Attention: ci/gpu_ci_run_h100.sh

Important Files Changed

Filename	Overview
skyrl/backends/skyrl_train/workers/megatron/adapter_store.py	Adds precision-aware FP32-master resolution and snapshots master parameters, optimizer tensor state, and dynamic counters.
skyrl/tinker/api.py	Integrates in-memory forwarded futures, selected SQLite locking, sampling caches, and threaded protobuf serialization.
skyrl/tinker/external_future_store.py	Introduces negative process-epoch request IDs, concurrent waiting, forwarding-task tracking, and TTL-based reclamation.
skyrl/tinker/extra/external_inference.py	Allows external inference results to complete in-memory futures while retaining the legacy database fallback.
skyrl/tinker/extra/skyrl_train_inference_forwarding.py	Routes SkyRL-Train inference results into the new store and aligns HTTP pool timeout behavior.
pyproject.toml	Adds a source-built DeepGEMM dependency and updates pinned Megatron Bridge/Core fork revisions.
ci/gpu_ci_run_h100.sh	Adds the DeepGEMM test to a Megatron CI invocation that lacks the required isolated uv environment and persistent output logging.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[LoRA adapter swap] --> B[Resolve optimizer shard]
  B --> C[FP32 master tensor]
  C --> D[Pinned CPU snapshot]
  D --> E[Restore selected adapter]
  F[Forwarded sample request] --> G[ExternalFutureStore]
  G --> H[Inference forwarding task]
  H --> I[Completed or failed result]
  I --> J[retrieve_future response]
  J --> K[TTL reclamation]

Loading

_{Reviews (1): Last reviewed commit: "fix(megatron): snapshot precision-aware ..." | Re-trigger Greptile}

[greptile-apps]

DeepGEMM CI lacks isolation

The newly covered DeepGEMM test runs under uv run without --isolated, making dependency resolution and imports sensitive to packages already installed on the CI worker rather than only the locked Megatron environment.

Context Used: CLAUDE.md (source)

[greptile-apps]

CI output is not persisted

The changed Megatron test invocation does not redirect its output to a persistent log file as required by the repository guidance, making complete diagnostics harder to retain and inspect after an H100 test failure.

Context Used: CLAUDE.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[gemini-code-assist]

Code Review

This pull request introduces an in-memory ExternalFutureStore to handle forwarded sample futures without writing them to the database, optimizing the SQLite write path. It also adds database write serialization locks for SQLite to prevent connection pool exhaustion, caches sampling models, and validates sampler checkpoints once. Additionally, DeepGEMM is now built against the selected PyTorch runtime as an external package. In the Megatron adapter store, a critical bug was identified in _iter_optimizer_state_groups where the optimizer state is not re-retrieved after resolving the FP32 master parameter, which would cause the optimizer states to be lost during snapshot and restore.

[gemini-code-assist]

In _iter_optimizer_state_groups, when main_param is initially None (common in precision-aware setups), state_key is set to shard_param and state is retrieved using optimizer.state.get(state_key, {}). However, if main_param is subsequently resolved via param_to_inner_param, the actual optimizer state is keyed by main_param (the FP32 master parameter), not shard_param. Because state is not re-retrieved using the resolved main_param, the returned state remains an empty dictionary {}. This causes the optimizer states (like exp_avg, exp_avg_sq, and step) to be completely lost/ignored during snapshot and restore for these parameters.

Suggested change

	for param_idx, (main_param, shard_param) in enumerate(zip(main_group, shard_group)):
	state_key = main_param if main_param is not None else shard_param
	state = optimizer.state.get(state_key, {})
	if main_param is None:
	main_param = getattr(optimizer, "param_to_inner_param", {}).get(shard_param)
	if main_param is None:
	main_param = state.get("master_param")
	if not isinstance(main_param, torch.Tensor):
	raise RuntimeError(
	"AdapterStore: precision-aware optimizer has no FP32 master tensor "
	f"for group {group_idx} parameter {param_idx}"
	)
	state_group.append((main_param, state))
	for param_idx, (main_param, shard_param) in enumerate(zip(main_group, shard_group)):
	state_key = main_param if main_param is not None else shard_param
	state = optimizer.state.get(state_key, {})
	if main_param is None:
	main_param = getattr(optimizer, "param_to_inner_param", {}).get(shard_param)
	if main_param is not None:
	state = optimizer.state.get(main_param, {})
	else:
	main_param = state.get("master_param")
	if not isinstance(main_param, torch.Tensor):
	raise RuntimeError(
	"AdapterStore: precision-aware optimizer has no FP32 master tensor "
	f"for group {group_idx} parameter {param_idx}"
	)
	state_group.append((main_param, state))

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit d0e4a8a. Configure here.}

[cursor]

Wrong tensor used as master

High Severity

When a precision-aware shard has no entry in shard_fp32_from_float16_groups, _iter_optimizer_state_groups treats param_to_inner_param as the FP32 master and never consults state["master_param"] or param_to_fp32_param. On HybridDeviceOptimizer, that map is the same-dtype inner copy, so snapshots skip the real master (also excluded from cpu_opt_state). Adapter swaps then leave stale FP32 masters in place and corrupt later optim_steps under CPU-offload.

 

^{Reviewed by Cursor Bugbot for commit d0e4a8a. Configure here.}