[megatron] 12/n towards Kimi K2.6: harden AdapterStore swaps for Tinker API multi-LoRA session churn - #2063
Conversation
…deleted-tenant live state Three fixes for the session-churn crash seen on 2026-08-18 (swap_to_adapter -> CUDA invalid argument at grad copy_, then silent state inheritance on the retry run): - _snapshot/_restore skip grad_data copies while the DDP grad buffers are offloaded (storage().resize_(0) leaves a stale view; Megatron zero-fills on reload anyway, and grads are only offloaded post-step). Snapshot records zeros instead of reading freed memory. - create() no longer adopts the live GPU state when current_id was cleared by delete(): a new _live_dirty flag distinguishes "live is pristine" (true first create) from "live mirrors a deleted tenant" (must seed from pristine, swap_to restores it). - register_pristine() synchronizes after its non_blocking D2H snapshot so create()'s CPU-side _copy_slot can't race the in-flight DMA. Unit tests fake the DDP buffers via monkeypatched _iter_buffers and cover the exact production ordering (create new model -> expiry-delete of current -> swap with grads offloaded). Co-authored-by: Cursor <cursoragent@cursor.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses two production issues in AdapterStore for Megatron: it prevents crashes during swap_to() when DDP gradient buffers are offloaded by skipping gradient copies when storage size is zero, and introduces a _live_dirty flag to prevent newly created adapters from adopting the live GPU state of a deleted adapter. Unit tests have been added to cover these scenarios. The review feedback suggests defensively guarding against buf.grad_data being None or missing in _grad_data_live to prevent potential AttributeError crashes.
| gd = buf.grad_data | ||
| return not (gd.is_cuda and gd.untyped_storage().size() == 0) |
There was a problem hiding this comment.
Defensively guard against buf.grad_data being None. While Megatron currently offloads gradients by resizing the storage to 0, future versions or alternative DDP implementations might set grad_data to None directly. Checking for None prevents potential AttributeError crashes.
| gd = buf.grad_data | |
| return not (gd.is_cuda and gd.untyped_storage().size() == 0) | |
| gd = getattr(buf, "grad_data", None) | |
| if gd is None: | |
| return False | |
| return not (gd.is_cuda and gd.untyped_storage().size() == 0) |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit e9a3a7c. Configure here.
| # else: grad storage is offloaded/freed. zero_grad_buffer() / | ||
| # restore_grad_buffers() reallocate it zero-filled before the next | ||
| # forward_backward, and a slot swapped in while offloaded carries | ||
| # post-step (zero) grads anyway — skipping loses nothing. |
There was a problem hiding this comment.
Offloaded swap drops pending adapter grads
High Severity
_restore skips writing a slot's grad_data while DDP grad storage is offloaded, then swap_to still marks that adapter current. A later forward_backward backloads grads as zeros and the swap is a no-op, so pending CPU grads never land. forward() already swaps with the optimizer still offloaded, so a tenant that accumulated, got swapped away, then came back via sample/forward can silently drop interrupted forward_backward grads.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit e9a3a7c. Configure here.
|
hey @casper-hansen i think this should be solved by #2000, closing this for now, feel free to re-open if #2000 looks insufficient |
2 participants
Part of the Kimi K2.6/K2.7 series (previous: #2062). Standalone: fixes apply to the existing multi-LoRA
AdapterStoreon main. Pairs naturally with #2062 (offload) and #2026 (merge_lora=false sync), which make the fixed code paths much hotter.What
Three fixes to
AdapterStore, plus a CPU-side unit-test suite (test_adapter_store_swap.py) that fakes the DDP buffers via a monkeypatched_iter_buffersand replays the exact production ordering (create new model -> expiry-delete of the current one -> swap with grads offloaded):_snapshot/_restoreskipgrad_datacopies while the DDP grad buffers are offloaded:storage().resize_(0)leaves a stale view behind, and reading/writing it is UB. Megatron zero-fills grad buffers on reload anyway, and grads are only ever offloaded post-step, so recording zeros is faithful.create()no longer adopts the live GPU state whencurrent_idwas cleared bydelete(): a new_live_dirtyflag distinguishes "live is pristine" (true first create) from "live still mirrors a deleted tenant" (must seed from the pristine snapshot;swap_torestores it).register_pristine()synchronizes after itsnon_blockingD2H snapshot socreate()'s CPU-side_copy_slotcannot race the in-flight DMA.Why
Long-running Tinker API service with multiple LoRA sessions churning (clients create models, train, hit session expiry, reconnect) on our Kimi K2.7 B300 deployment. On 2026-08-18 we hit the failure chain this PR fixes: a
swap_to_adaptercrashed with CUDAinvalid argumentat the gradcopy_(grad buffers had been offloaded post-step, leaving resized-to-zero storages), the run restarted, and the retry then silently inherited the previous tenant's live adapter state becausecreate()treated the leftover live weights as pristine after the expirydelete()had clearedcurrent_id. The second bug is the nastier one — no crash, just a new session fine-tuning on top of a dead tenant's weights.Made with Cursor