From b8bc1696707657565173329f71b9f2f530a53fca Mon Sep 17 00:00:00 2001 From: Eric Tang Date: Fri, 7 Aug 2026 00:14:15 +0000 Subject: [PATCH] x --- .../workers/megatron/adapter_store.py | 118 ++++++++++- .../workers/megatron/megatron_worker.py | 26 +++ .../skyrl_train/workers/worker_dispatch.py | 17 +- .../test_multi_lora_megatron_colocated.py | 200 ++++++++++++++++++ 4 files changed, 356 insertions(+), 5 deletions(-) create mode 100644 tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py diff --git a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py b/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py index ec7f82b538..a61cad361e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py +++ b/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py @@ -37,10 +37,26 @@ def _iter_buffers(model_chunks) -> Iterable[Tuple[int, int, Any]]: def _new_pinned_like(t: torch.Tensor) -> torch.Tensor: - """Allocate a pinned-CPU tensor with the same shape/dtype as t.""" + """Allocate a pinned-CPU tensor with the same shape/dtype as t. + + Safe to call on an offloaded buffer: only shape/dtype metadata is read, + which survives ``storage().resize_(0)``. + """ return torch.empty_like(t, device="cpu").pin_memory() +def _is_resident(t: Optional[torch.Tensor]) -> bool: + """True when ``t`` still owns storage we can copy to/from. + + Megatron's offload path frees GPU buffers with ``storage().resize_(0)`` + (see ``_ParamAndGradBuffer.offload_to_cpu``) and leaves the tensor object + — full shape and all — pointing at a zero-sized storage. Copying from it + raises ``cudaErrorInvalidValue``, so every read/write of a DDP buffer has + to be gated on this. + """ + return t is not None and t.untyped_storage().size() > 0 + + def _expected_lora_param_check(model_chunks) -> None: """Sanity-check: every trainable param under DDP buffers is a LoRA adapter param. @@ -149,6 +165,10 @@ def __init__(self) -> None: self._pristine: Optional[AdapterSlot] = None self._current_id: Optional[str] = None self._signature: Optional[LoraSignature] = None + # True while the DDP grad buffers are offloaded, i.e. each adapter's + # grads live only in its CPU slot. Set by park_grads, cleared by + # unpark_grads. + self._grads_parked: bool = False @property def current_id(self) -> Optional[str]: @@ -223,12 +243,34 @@ def _allocate_empty_slot(self, model_chunks, optimizer) -> AdapterSlot: slot.cpu_param_group_state.append(group_state) return slot + @staticmethod + def _require_param_residency(buf, mc_idx: int, buf_idx: int) -> None: + """Fail loudly when a swap is attempted with model params offloaded. + + Callers must backload the model before swapping (the dispatch does + this via ``_ensure_on_gpu(..., need_model=True)``). Grads are allowed + to be offloaded — see :meth:`park_grads` — but params are not, because + the CPU mirror Megatron keeps for them (``param_data_cpu``) is not + per-adapter and would hand the next backload the wrong tenant's + weights. + """ + if not _is_resident(buf.param_data): + raise RuntimeError( + f"AdapterStore: DDP buffer {mc_idx}/{buf_idx} param_data is offloaded; " + f"backload the model before swapping adapters." + ) + @torch.no_grad() def _snapshot(self, slot: AdapterSlot, model_chunks, optimizer) -> None: """Copy live GPU state into `slot` (CPU).""" for mc_idx, buf_idx, buf in _iter_buffers(model_chunks): + self._require_param_residency(buf, mc_idx, buf_idx) slot.cpu_param_data[mc_idx][buf_idx].copy_(buf.param_data, non_blocking=True) - slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True) + # Skip grads while the grad buffers are offloaded: their GPU + # storage is freed, and the slot's copy — parked by park_grads() + # just before the offload — is already the authoritative one. + if _is_resident(buf.grad_data): + slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True) for opt_idx, _opt in enumerate(iter_opts(optimizer)): groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or [] for g, group in enumerate(groups): @@ -254,8 +296,13 @@ def _snapshot(self, slot: AdapterSlot, model_chunks, optimizer) -> None: def _restore(self, slot: AdapterSlot, model_chunks, optimizer) -> None: """Copy `slot` (CPU) into live GPU state.""" for mc_idx, buf_idx, buf in _iter_buffers(model_chunks): + self._require_param_residency(buf, mc_idx, buf_idx) buf.param_data.copy_(slot.cpu_param_data[mc_idx][buf_idx], non_blocking=True) - buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True) + # Offloaded grad buffers have no storage to restore into; the + # incoming adapter's grads stay parked in its slot and are + # re-materialised by unpark_grads() on the next backload. + if _is_resident(buf.grad_data): + buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True) for opt_idx, _opt in enumerate(iter_opts(optimizer)): groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or [] for g, group in enumerate(groups): @@ -372,6 +419,66 @@ def _copy_slot(self, src: AdapterSlot, dst: AdapterSlot) -> None: else: dst_pg[k] = v + # ------------------------------------------------------------------ + # Grad parking across CPU offload + # ------------------------------------------------------------------ + + @torch.no_grad() + def park_grads(self, model_chunks) -> None: + """Save the live adapter's grads into its slot before a grad offload. + + Megatron frees ``grad_data`` on offload and zero-fills it on reload, + so grads accumulated by a ``forward_backward`` that hasn't reached its + ``optim_step`` yet are destroyed outright. Under colocation another + tenant's request (a sample, a forward) offloads in exactly that gap, + so the grads have to be parked per-adapter instead. + + Call immediately before the strategy's offload. No-op when the grad + buffers are already offloaded or no adapter is live. + """ + if self._current_id is None or self._current_id not in self._slots: + return + slot = self._slots[self._current_id] + parked = False + for mc_idx, buf_idx, buf in _iter_buffers(model_chunks): + if not _is_resident(buf.grad_data): + continue + slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True) + parked = True + if parked: + torch.cuda.current_stream().synchronize() + self._grads_parked = True + + @torch.no_grad() + def unpark_grads(self, model_chunks) -> None: + """Re-materialise the live adapter's grads after a grad backload. + + Mirror of :meth:`park_grads`. Restores the adapter that is live *now*, + which need not be the one that was live at park time — a swap in the + offload window only moves CPU slots around, and this is where the + result lands back on the GPU. + + Call immediately after the strategy's backload. No-op unless grads + were parked and the buffers are resident again. + """ + if not self._grads_parked: + return + if self._current_id is None or self._current_id not in self._slots: + # Live adapter was deleted while offloaded: nothing to restore, + # and the reloaded buffers are already zeroed. + self._grads_parked = False + return + slot = self._slots[self._current_id] + restored = False + for mc_idx, buf_idx, buf in _iter_buffers(model_chunks): + if not _is_resident(buf.grad_data): + continue + buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True) + restored = True + if restored: + torch.cuda.current_stream().synchronize() + self._grads_parked = False + @torch.no_grad() def delete(self, model_id: str) -> None: """Drop the slot for `model_id`. @@ -402,6 +509,11 @@ def swap_to(self, model_id: str, model_chunks, optimizer) -> None: agree on the live adapter before the next collective. TP/PP/EP groups do not need barriers because the swap is identical-shape on all ranks within those groups (LoRA signature is fixed). + + Model params must be GPU-resident; the DDP grad buffers and the + DistributedOptimizer state need not be (colocation offloads them + between requests). Offloaded grads are left parked in their slots — + see :meth:`park_grads`. """ if model_id not in self._slots: raise KeyError(f"AdapterStore: unknown adapter '{model_id}'") diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 40fa6fe181..239c71aebf 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -1552,6 +1552,32 @@ def swap_to_adapter(self, model_id: str) -> None: return # FFT path: no-op self.adapter_store.swap_to(model_id, self.actor_module, self.optimizer) + def offload_to_cpu(self, offload_optimizer: bool = True, offload_model: bool = True): + """Offload worker state, parking the live adapter's grads first. + + The optimizer half of the offload frees the DDP grad buffers outright + (Megatron's reload zero-fills them), which would silently drop grads a + tenant accumulated in a ``forward_backward`` whose ``optim_step`` + hasn't arrived yet — under colocation another tenant's sample lands in + exactly that gap. Parking moves them into the adapter's CPU slot, + where they survive both the offload and any swap that happens while + offloaded. + """ + if offload_optimizer and self.adapter_store is not None and self.actor_module is not None: + self.adapter_store.park_grads(self.actor_module) + super().offload_to_cpu(offload_optimizer=offload_optimizer, offload_model=offload_model) + + def backload_to_gpu(self, backload_optimizer: bool = True, backload_model: bool = True): + """Backload worker state, re-materialising the live adapter's grads. + + Counterpart to :meth:`offload_to_cpu`. Restores whichever adapter is + live at this point, which may differ from the one parked if a swap + happened while offloaded. + """ + super().backload_to_gpu(backload_optimizer=backload_optimizer, backload_model=backload_model) + if backload_optimizer and self.adapter_store is not None and self.actor_module is not None: + self.adapter_store.unpark_grads(self.actor_module) + def adapter_store_state(self) -> dict: """Diagnostic: return current_id + registered model_ids. Cheap; useful for tests.""" diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 52b8f71a3c..e6800559a3 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/skyrl/backends/skyrl_train/workers/worker_dispatch.py @@ -86,11 +86,17 @@ def ensure_active_adapter(self, role: str, model_id: Optional[str]) -> None: No-op when ``model_id is None`` (single-tenant / FFT path) or when the workers don't have an AdapterStore (non-LoRA strategies). - Must be called *after* ``_ensure_on_gpu(role, ...)`` so the model - and optimizer storages are live before we tensor.copy_() into them. + The swap copies the DDP param buffers, so the model must be resident + before it runs. Callers generally arrive here right after their own + ``_ensure_on_gpu``; we repeat it (a no-op against dispatch-local + state) so paths that only need the optimizer — ``set_lr``, say — + can't hand the AdapterStore freed param storage. The optimizer and + grad buffers are allowed to stay offloaded: the store copies the + optimizer's CPU tensors in place and leaves parked grads alone. """ if model_id is None or role not in self._actor_groups: return + self._ensure_on_gpu(role, need_optimizer=False, need_model=True) ray.get(self._actor_groups[role].async_run_ray_method("pass_through", "swap_to_adapter", model_id)) def register_adapter(self, role: str, model_id: str) -> None: @@ -417,7 +423,14 @@ def optim_step(self, model: str, model_id: Optional[str] = None) -> Optional[flo """Run optimizer step. For single-tenant training, the model should already be on GPU from forward_backward. For multi-tenant LoRA training, ``model_id`` is used to ensure the correct adapter is used. + + The residency check is not redundant with ``forward_backward``'s under + multi-tenancy: another tenant's request (e.g. a sample, which offloads + the optimizer) can land between this model's forward_backward and its + optim_step, so the state this step writes to may have been offloaded + in the gap. """ + self._ensure_on_gpu(model, need_optimizer=True, need_model=True) self.ensure_active_adapter(model, model_id) refs = self._actor_groups[model].async_run_ray_method("pass_through", "optim_step") grad_norms = ray.get(refs) diff --git a/tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py b/tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py new file mode 100644 index 0000000000..2d91b83e26 --- /dev/null +++ b/tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py @@ -0,0 +1,200 @@ +"""Multi-LoRA tests for the *colocated* Megatron path (``colocate_all=True``). + +Sibling of ``test_multi_lora_megatron.py``, which pins +``trainer.placement.colocate_all=False``. Colocation is the SkyRL-Train +default (and what the Tinker API assumes when the operator passes no +override), and it changes the memory lifecycle underneath every adapter +swap: + + * ``WorkerDispatch._prepare_for_weight_sync`` offloads the optimizer and + the DDP grad buffers before syncing weights to vLLM. + * ``forward``/``forward_from_staged`` only backload the *model*, so the + optimizer + grad buffers stay offloaded across a pure forward. + * Megatron's ``offload_grad_buffers`` frees ``grad_data`` outright + (``storage().resize_(0)``) and ``restore_grad_buffers`` zero-fills it. + +So on this path a swap runs against partially-offloaded state, and any +grads the live adapter accumulated but has not yet consumed in +``optim_step`` are destroyed by the offload. Both are exercised below. + +Run with: + uv run --extra tinker --extra megatron --with pytest --with pytest-timeout \\ + pytest -s tests/tinker/skyrl_train/test_multi_lora_megatron_colocated.py +""" + +from __future__ import annotations + +import pytest + +cuda_available = False +try: # pragma: no cover - import guard + import torch + + cuda_available = bool(torch.cuda.is_available() and torch.cuda.device_count() > 0) +except Exception: + cuda_available = False + +pytestmark = pytest.mark.skipif(not cuda_available, reason="multi-LoRA Megatron tests require at least one CUDA GPU") + +tinker = pytest.importorskip("tinker") +from tinker import types as tinker_types # noqa: E402 + +from tests.tinker.skyrl_train.test_multi_lora_megatron import ( # noqa: E402 + BASE_MODEL, + TINKER_API_KEY, + _api_server, + _make_datum, + _sample_greedy, +) + +TEST_PORT = 8021 + +# Same tiny config as the non-colocated module, but colocated: the policy +# worker and the single vLLM engine share one GPU, so the dispatch offloads +# policy state to CPU around every sample. +BACKEND_CONFIG = { + "strategy": "megatron", + "trainer.placement.policy_num_gpus_per_node": 1, + "trainer.placement.policy_num_nodes": 1, + "trainer.placement.colocate_all": True, + "trainer.policy.megatron_config.tensor_model_parallel_size": 1, + "trainer.policy.megatron_config.pipeline_model_parallel_size": 1, + "trainer.policy.megatron_config.lora_config.merge_lora": False, + "trainer.policy.model.lora.max_loras": 4, + "trainer.policy.model.lora.max_cpu_loras": 4, + "generator.inference_engine.num_engines": 1, + "generator.inference_engine.tensor_parallel_size": 1, + "generator.inference_engine.gpu_memory_utilization": 0.3, +} + + +@pytest.fixture(scope="module") +def server(): + with _api_server(TEST_PORT, BACKEND_CONFIG) as proc: + yield proc + + +@pytest.fixture +def service_client(server): + return tinker.ServiceClient(base_url=f"http://0.0.0.0:{TEST_PORT}/", api_key=TINKER_API_KEY) + + +def test_sample_non_live_adapter_colocated(service_client): + """Minimal repro: sampling an adapter that isn't the live one. + + After ``create_lora_training_client`` for A then B, A is the live + adapter on the workers and (under colocation) the optimizer + grad + buffers are offloaded. ``save_weights_for_sampler(B)`` then swaps + A -> B while ``grad_data`` has been freed, so the AdapterStore's + snapshot copies from a zero-sized CUDA storage: + + cpu_grad_data[...].copy_(buf.grad_data) + torch.AcceleratorError: CUDA error: invalid argument + + No training required — the offload after ``init_model`` is enough. + """ + a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + b = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + tok = a.get_tokenizer() + + tokens = _sample_greedy(b, name="colo_b0", tok=tok, prompt="Question: 2+2?\nAnswer:") + assert len(tokens) > 0 + + +def test_pending_grads_survive_other_tenant_sample(service_client): + """A's un-consumed grads must survive another tenant's sample. + + The Tinker request stream interleaves tenants at request granularity, so + ``A.forward_backward`` / ``A.optim_step`` can straddle a ``B.sample``. + Under colocation that sample offloads the grad buffers, and Megatron's + grad offload *drops* the grads (reload zero-fills). If they aren't parked + in A's adapter slot, A's ``optim_step`` applies an all-zero gradient. + + Measured against an undisturbed control adapter C that runs the same + single step from pristine on the same data: A's post-step loss has to + land on C's. Asserting only ``post < pre`` on A is not enough — Adam + still applies weight decay to a zero gradient, which drifts the loss by + ~1e-6 and would satisfy a bare inequality while the step itself was a + no-op (the real step moves it by ~1e-2). + """ + a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + b = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + c = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + tok = a.get_tokenizer() + data = [_make_datum(tok, "Question: 1+1?\nAnswer:", " 2")] + lr = tinker_types.AdamParams(learning_rate=1e-2) + + def _loss(out): + return sum(sum(o["elementwise_loss"].data) for o in out.loss_fn_outputs) + + # A accumulates grads but does not step yet. + pre_a = _loss(a.forward_backward(data, "cross_entropy").result()) + + # B samples in the gap: offloads optimizer + grad buffers, swaps A -> B. + _sample_greedy(b, name="colo_pending_b", tok=tok, prompt="Question: 2+2?\nAnswer:") + + # A steps on the grads it accumulated before B intervened. + a.optim_step(lr).result() + post_a = _loss(a.forward_backward(data, "cross_entropy").result()) + + # C: same step, no sample in the gap. + pre_c = _loss(c.forward_backward(data, "cross_entropy").result()) + c.optim_step(lr).result() + post_c = _loss(c.forward_backward(data, "cross_entropy").result()) + + improvement_c = pre_c - post_c + print( + f"\n[pending_grads] A: pre={pre_a!r} post={post_a!r} Δ={post_a - pre_a:.6e}\n" + f"[pending_grads] C: pre={pre_c!r} post={post_c!r} Δ={post_c - pre_c:.6e}" + ) + + assert pre_a == pre_c, f"A and C were not equally pristine: {pre_a!r} vs {pre_c!r}" + assert improvement_c > 0, f"control adapter C did not learn (pre={pre_c!r}, post={post_c!r}); test is inconclusive" + # Tolerance scales with the control's own step, so this stays meaningful + # whatever the tiny model happens to do. + assert abs(post_a - post_c) <= 0.05 * improvement_c, ( + f"A's step diverged from the undisturbed control: A post={post_a!r}, C post={post_c!r} " + f"(C improved by {improvement_c:.6e}). A's pending grads were dropped by the grad-buffer " + "offload that B's sample triggered." + ) + + +def test_two_adapters_train_and_sample_colocated(service_client): + """Colocated analog of ``test_two_adapters_sample_independently``. + + Full interleaving of training and sampling across two tenants: every + sample round-trips through offload -> swap -> broadcast -> offload, and + every training request backloads and swaps back. + """ + a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + b = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8) + tok = a.get_tokenizer() + data = [_make_datum(tok, "Question: 1+1?\nAnswer:", " 2")] + sample_prompt = "Question: 2+2?\nAnswer:" + + for _ in range(2): + a.forward_backward(data, "cross_entropy").result() + a.optim_step(tinker_types.AdamParams(learning_rate=1e-2)).result() + tokens_a_first = _sample_greedy(a, name="colo_a1", tok=tok, prompt=sample_prompt) + + b.forward_backward(data, "cross_entropy").result() + b.optim_step(tinker_types.AdamParams(learning_rate=5e-2)).result() + tokens_b = _sample_greedy(b, name="colo_b1", tok=tok, prompt=sample_prompt) + + a.forward_backward(data, "cross_entropy").result() + a.optim_step(tinker_types.AdamParams(learning_rate=1e-2)).result() + tokens_a_continued = _sample_greedy(a, name="colo_a2", tok=tok, prompt=sample_prompt) + print( + f"\n[colocated_independently] tokens_a_first={tokens_a_first}\n" + f"[colocated_independently] tokens_b ={tokens_b}\n" + f"[colocated_independently] tokens_a_cont ={tokens_a_continued}" + ) + + assert tokens_a_continued != tokens_a_first, ( + f"A's tokens did not change after one more optim_step (A={tokens_a_continued!r}, " + f"prior={tokens_a_first!r}). A's optimizer state may have been wiped by B." + ) + assert tokens_a_continued != tokens_b, ( + f"A's continued sample matches B's sample: A={tokens_a_continued!r}, B={tokens_b!r}. " + "B's adapter sync may have clobbered A's slot on vLLM." + )