fix(lora): backload optimizer before adapter swap

[pcmoritz]

The adapter swap copy_()s the live param buffer, grad buffer, fp32 main params and Adam state to/from CPU slots, so all of them must be resident. Megatron's grad offload frees grad_data in place (storage().resize_(0)) rather than moving it, so swapping while the optimizer half is offloaded reads a dangling device pointer and fails with a bare "CUDA error: invalid argument".

• ensure_active_adapter() now backloads with need_optimizer=True itself instead of trusting the caller, since paths like forward and save_weights_for_sampler deliberately keep the optimizer off GPU.
• Move ensure_active_adapter() ahead of _prepare_for_weight_sync(), which is what pushes those buffers back off GPU.
• Add _check_resident() in AdapterStore so a freed storage raises an actionable error instead of the opaque CUDA failure.

[gemini-code-assist]

Code Review

This pull request introduces residency checks for Megatron DDP buffers during adapter swaps to prevent cryptic CUDA errors caused by accessing offloaded storage. It also ensures that the model and optimizer are resident on the GPU before swapping adapters, and reorders the weight synchronization sequence. The review feedback suggests two valuable improvements: caching the active adapter ID in WorkerDispatch to avoid redundant GPU allocations and transfers on every forward pass, and adding a defensive None check in _check_resident to prevent potential AttributeError crashes.

[gemini-code-assist]

Forcing the optimizer to GPU (need_optimizer=True) on every call to ensure_active_adapter defeats the offloading optimization during forward passes (e.g., in forward and forward_from_staged), where the optimizer is deliberately kept off GPU to save memory. Since adapter swaps are only needed when the requested model_id is different from the currently active adapter, we can cache the active adapter ID in WorkerDispatch and return early if no swap is required. This avoids redundant GPU memory allocation and expensive CPU-GPU transfers on every forward pass.

Suggested change

	if model_id is None or role not in self._actor_groups:
	return
	self._ensure_on_gpu(role, need_optimizer=True, need_model=True)
	ray.get(self._actor_groups[role].async_run_ray_method("pass_through", "swap_to_adapter", model_id))
	if model_id is None or role not in self._actor_groups:
	return
	if not hasattr(self, "_active_adapter"):
	self._active_adapter = {}
	if self._active_adapter.get(role) == model_id:
	return
	self._ensure_on_gpu(role, need_optimizer=True, need_model=True)
	ray.get(self._actor_groups[role].async_run_ray_method("pass_through", "swap_to_adapter", model_id))
	self._active_adapter[role] = model_id

[gemini-code-assist]

To prevent potential AttributeError crashes if t is None (for example, if certain gradient or parameter buffers are not initialized or are optional in some configurations), we should add a defensive None check at the beginning of _check_resident.

Suggested change

	def _check_resident(t: torch.Tensor, what: str) -> None:
	"""Fail loudly if `t`'s storage has been freed by an offload.
	Megatron offloads DDP buffers in place: `offload_grad_buffers()` /
	`buffer.offload_to_cpu()` call `storage().resize_(0)` and leave the tensor
	view intact ("accessing them during offload is undefined behavior").
	copy_()ing such a tensor hands cudaMemcpyAsync a dangling device pointer,
	which surfaces as a bare `CUDA error: invalid argument` from whichever
	line touched it. Say what actually went wrong instead.
	"""
	if t.numel() > 0 and t.untyped_storage().size() == 0:
	raise RuntimeError(
	f"AdapterStore: {what} is offloaded (storage freed); backload the "
	"policy model *and* optimizer/grad buffers before swapping adapters."
	)
	def _check_resident(t: Optional[torch.Tensor], what: str) -> None:
	"""Fail loudly if `t`'s storage has been freed by an offload.
	Megatron offloads DDP buffers in place: `offload_grad_buffers()` /
	`buffer.offload_to_cpu()` call `storage().resize_(0)` and leave the tensor
	view intact ("accessing them during offload is undefined behavior").
	copy_()ing such a tensor hands cudaMemcpyAsync a dangling device pointer,
	which surfaces as a bare `CUDA error: invalid argument` from whichever
	line touched it. Say what actually went wrong instead.
	"""
	if t is None:
	return
	if t.numel() > 0 and t.untyped_storage().size() == 0:
	raise RuntimeError(
	f"AdapterStore: {what} is offloaded (storage freed); backload the "
	"policy model *and* optimizer/grad buffers before swapping adapters."
	)

[erictang000]

if we have the call order

self.forward("ref", ...)
self.forward(model="policy",...,model_id="lora1")

won't we unnecessarily run the lora1 forward with the optimizer state on GPU now? can we avoid that in some way even if the grad buffer/optim state should technically be small for lora

[erictang000]

closed by #2000