[fix][megatron] Make multi-LoRA adapter swaps safe against colocation offload

[erictang000]

[fix][megatron] Make multi-LoRA adapter swaps safe against colocation offload

Problem

Running a multi-tenant Tinker script that both trains and samples crashes on the Megatron backend:

ray::MegatronPolicyWorkerBase.swap_to_adapter()
  File ".../workers/megatron/adapter_store.py", line 417, in swap_to
    self._snapshot(current_slot, model_chunks, optimizer)
  File ".../workers/megatron/adapter_store.py", line 231, in _snapshot
    slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
torch.AcceleratorError: CUDA error: invalid argument

Multi-tenant scripts that only train are fine; the failure needs a sample.

Root cause

Under colocate_all (the SkyRL-Train default, and what the Tinker API assumes when the operator passes no override), the dispatch offloads policy state between requests. Megatron's grad-buffer offload does not move grads to CPU — it frees them with storage().resize_(0) and zero-fills on reload. The tensor object keeps its full shape, so nothing looks wrong until a copy dereferences the zero-sized storage:

save_weights_for_sampler(B)
  → _prepare_for_weight_sync()  → _offload(offload_optimizer=True) → grad_data storage freed
  → ensure_active_adapter(B)    → swap A→B → cpu_grad_data.copy_(buf.grad_data)  💥

Train-only scripts survive because forward_backward backloads everything (need_optimizer=True). Two more instances of the same shape were sitting next to it:

• forward / forward_from_staged backload only the model, so a swap during a plain forward runs against the same freed grad buffers.
• optim_step did no residency check at all, so a swap there could hit freed param storage once another tenant's sample had offloaded the model.

Underneath the crash was a silent correctness bug: because the offload drops grads rather than saving them, grads a tenant accumulated in a forward_backward whose optim_step hasn't arrived yet are destroyed when another tenant's request lands in the gap — which is exactly what Tinker's request-granularity interleaving produces. That tenant's optim_step then applies an all-zero gradient and silently does nothing.

Fix

adapter_store.py

• _is_resident() gates every DDP-buffer read/write. Grad copies are skipped while the buffers are offloaded; params must be resident, and a swap without them now raises a clear RuntimeError instead of a CUDA error.
• New park_grads() / unpark_grads(): the live adapter's grads are saved into its own CPU slot immediately before an offload and re-materialised on backload. Unparking restores whichever adapter is live at that moment, which need not be the one parked — a swap inside the offload window only moves CPU slots around, and this is where the result lands back on the GPU.

The optimizer needs no special handling: its offload genuinely moves tensors to CPU, so the store's copies degrade to CPU→CPU and the values ride back to GPU on the next backload.

megatron_worker.py — MegatronPolicyWorkerBase brackets offload_to_cpu / backload_to_gpu with park/unpark.

worker_dispatch.py — ensure_active_adapter guarantees model residency before swapping (covers set_lr, which only asks for the optimizer); optim_step ensures residency, which is not redundant with forward_backward's under multi-tenancy.

Test plan

New tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py — sibling of test_multi_lora_megatron.py with colocate_all=True. The existing module pins it to False, which is why none of this was covered. CI already globs tests/tinker/skyrl_train/, so the module is picked up with no workflow change.

test	on main	with this PR
test_sample_non_live_adapter_colocated	the traceback above, with no training at all	pass
test_pending_grads_survive_other_tenant_sample	A's optim_step is a no-op (Δ=0.0) vs the control's −4.6e-02	pass, A bit-identical to control
test_two_adapters_train_and_sample_colocated	—	pass

The grad test measures A against an undisturbed control adapter rather than asserting post < pre: a bare inequality passes even with grads dropped, because Adam still applies weight decay to a zero gradient and drifts the loss by ~1e-6. Sensitivity was confirmed by disabling the parking hook and re-running.

Also run:

• tests/tinker/skyrl_train/test_multi_lora_megatron.py — 6 existing non-colocated tests, no regression (9/9 with the new module in one session).
• tests/backends/skyrl_train/gpu/gpu_ci/test_worker_dispatch_offload.py -k "training_switch or gpu_state_tracking or only_one_model" — 3 passed, covering the FFT colocated path touched in optim_step.

Not run: tests/backends/skyrl_train/distributed/test_worker_dispatch.py hangs in Ray worker registration on the dev host — identically on a clean tree with these changes stashed, so pre-existing, but it does leave the dispatch unit tests unverified locally.

[gemini-code-assist]

Code Review

This pull request adds support for multi-LoRA training and adapter swapping on the colocated Megatron path. It introduces gradient 'parking' and 'unparking' mechanisms to preserve accumulated gradients across CPU offloads, ensures proper GPU residency of models and optimizers during swaps and optimizer steps, and adds a comprehensive test suite to verify these colocated behaviors. I have no feedback to provide as there are no review comments.