[tinker][megatron] Multi-LoRA Megatron + Tinker API RL Training

[erictang000]

Summary

Multi-tenant LoRA training and per-adapter sampling for the SkyRL-Train Megatron + FSDP backends, exposed via the Tinker API. Multiple rl_loop clients can run concurrently against a single Tinker server, each owning their own LoRA adapter; the backend swaps the live adapter into Megatron's optimizer/buffer state on demand and registers per-tenant adapters on vLLM by model_id.

What's in the diff

Six files, ~238 LOC net.

Backend / dispatch (skyrl/backends/skyrl_train_backend.py, workers/worker_dispatch.py):

• Drop the v1 single-tenant gates in sample() and save_sampler_checkpoint(). The batched-sample path now validates every model_id in a mixed batch is a known policy, and resolves the inference-engine model name per request (model_id for LoRA tenants, resolve_policy_model_name(cfg) as the FFT / single-tenant fallback). The engine's find_batchable_sample legitimately mixes adapters in a single sample call now that multi-LoRA is live.
• WorkerDispatch.save_weights_for_sampler takes an optional model_id. Before broadcasting it calls ensure_active_adapter("policy", model_id) to swap the requested adapter into the optimizer/model on every worker — without this we'd export some other tenant's LoRA weights to vLLM. model_id is forwarded into _broadcast_to_inference_engines so the worker registers the adapter on vLLM under that name.
• model_id=None everywhere preserves the legacy single-tenant code path.

Per-tenant LoRA save+load (workers/megatron/megatron_worker.py, workers/fsdp/fsdp_worker.py):

• Both _save_lora_adapters_and_sync and broadcast_to_inference_engines take an optional lora_name / model_id. With model_id set the adapter writes into a per-tenant subdir of lora_sync_path and registers on vLLM under that name (via RemoteInferenceClient.load_lora_adapter). With model_id=None the legacy SKYRL_LORA_ADAPTER_NAME shared-path behavior is unchanged.
• Megatron path additionally requires merge_lora=False so vLLM serves each tenant's adapter by name rather than baking it into the base weights.

Config knobs (skyrl/train/config/ppo_base_config.yaml):

• New trainer.policy.model.lora.max_loras (default 1) — maps to vLLM's max_loras, the number of adapters that can be active in a single GPU batch.
• New trainer.policy.model.lora.max_cpu_loras (default null → vLLM defaults to max_loras) — the total LoRA capacity in vLLM's CPU LRU cache.

Tests (tests/tinker/skyrl_train/test_multi_lora_megatron.py):

• Existing train-only tests (test_two_adapters_train_independently, test_per_adapter_step_isolation, test_rank_mismatch_rejected, test_delete_then_train_remaining) continue to pass under the per-tenant config (merge_lora=False + max_loras=4).
• Removed the train-only test_sample_with_two_adapters_errors — the gate it covered is gone.
• Added test_per_adapter_sample_isolation: two pristine adapters following an identical training trajectory must produce bit-identical greedy samples through vLLM at every step. Catches: vLLM serving the wrong adapter for a model_id, load_lora_adapter("B", ...) clobbering adapter A's slot, optimizer-state regressions that surface as different greedy tokens.
• Added test_two_adapters_sample_independently: A trains, B trains in between (registering its own LoRA on vLLM), A trains one more step. Asserts (1) A's tokens advance after the resumed step (A's optimizer state survived B's intervention end-to-end through sampling) and (2) A's final tokens differ from B's (B's adapter sync did not clobber A's slot).
• New _sample_greedy(tc, name, tok, prompt) helper: save_weights_and_get_sampling_client(name=...) then a deterministic temperature=0 sample.

Operator contract

To run multi-tenant LoRA on Megatron, set:

trainer:
  policy:
    megatron_config:
      lora_config:
        merge_lora: false
    model:
      lora:
        max_loras: <max concurrent adapters>
        max_cpu_loras: <total adapter capacity>

max_cpu_loras should be ≥ the number of distinct adapters you expect to serve concurrently, otherwise vLLM evicts an adapter mid-run and the next sample 404s. There is no auto-rehydrate yet — operators size for the workload. The FFT and single-tenant LoRA paths are unchanged; existing configs need no modification.

Test plan

• pytest tests/tinker/skyrl_train/test_multi_lora_megatron.py (6 tests, GPU-gated, requires 2 GPUs for Megatron policy + vLLM)
• All existing single-tenant test_api.py paths
• End-to-end smoke: two rl_loop clients with distinct run_names sharing one Tinker server; verify both converge on their respective tasks without contamination

[gemini-code-assist]

Code Review

This pull request implements multi-tenant LoRA training and sampling for the Megatron backend, enabling multiple adapters to share a single base model. It introduces an AdapterStore to manage swapping LoRA weights and optimizer states (including gradients and Adam moments) between GPU and pinned CPU memory. The RemoteInferenceClient and various training components have been updated to require explicit model identification for routing generation requests to specific adapters. Additionally, configuration for concurrent LoRA capacity in vLLM is now supported. Feedback suggests using more specific exception types for API-level validation errors in the inference client.

[gemini-code-assist]

The _require_model function raises a ValueError if the model is missing. While this is good for validation, consider if a more specific exception type or a custom error class would be more appropriate for API-level validation errors to distinguish them from internal logic errors.

[erictang000]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements multi-tenant LoRA support for FSDP and Megatron backends, enabling concurrent training and sampling of multiple adapters. Changes include per-adapter synchronization paths, updated worker dispatch logic, and support for multi-adapter batches in the sampling backend. Feedback highlights a potential path traversal vulnerability when constructing sync paths from model_id and notes that the adapter name is missing in the weight update logic for non-remote inference clients.

[erictang000]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements multi-tenant LoRA support for the Megatron backend in SkyRL, allowing multiple LoRA adapters to be trained and sampled concurrently against a single server. Key changes include updating the weight synchronization and inference engine logic to handle per-tenant adapter directories and vLLM registration, adding necessary configuration options, and updating documentation and tests to reflect the new multi-tenancy capabilities. My review identified two issues: the shutil.rmtree call in megatron_worker.py incorrectly uses ignore_errors=True while attempting to catch exceptions, and the error message in skyrl_train_backend.py incorrectly references 'fsdp' instead of 'fsdp2'.