Skip to content

[tinker][fsdp] Multi-LoRA adapter store + Tinker API - #1956

Open
atemaguer wants to merge 2 commits into
NovaSky-AI:mainfrom
atemaguer:fsdp-multilora-adapter-store
Open

[tinker][fsdp] Multi-LoRA adapter store + Tinker API#1956
atemaguer wants to merge 2 commits into
NovaSky-AI:mainfrom
atemaguer:fsdp-multilora-adapter-store

Conversation

@atemaguer

@atemaguer atemaguer commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

First PR in a two-part series. The concurrent resident/grouped-GEMM implementation follows in #1938.

Adds multi-tenant LoRA training to the SkyRL-Train FSDP Tinker backend with one PEFT adapter resident on GPU at a time. Multiple Tinker clients share one frozen base model while each client retains independent LoRA weights, gradients, and optimizer state.

Architecture

  • New FSDPAdapterStore holds a pinned-CPU snapshot per model_id, plus a pristine template used to initialize new adapters.
  • Each snapshot contains the worker's local FSDP parameter shards, unconsumed gradients, and Adam state. Adapter swaps copy local shards only; they do not gather full parameters.
  • Swaps are bracketed by distributed barriers when the process group is initialized and synchronize CUDA copies before the next request.
  • The store validates a fixed (rank, alpha, target_modules, world_size) signature and rejects trainable non-LoRA parameters.

API surface

  • The first LoRA create_model primes Adam state, captures the pristine adapter, and registers the initial model_id.
  • Later policy create_model calls register CPU-backed adapter slots through the existing Tinker/Megatron adapter lifecycle.
  • forward, forward_backward, optim_step, checkpoint operations, and sampler synchronization use the existing WorkerDispatch.ensure_active_adapter path.
  • Deleting one of multiple adapters releases only that slot; deleting the last model retains the existing full-runtime teardown behavior.

Scope

  • This PR intentionally keeps exactly one trainer adapter resident at a time, matching the current Megatron backend.
  • Resident concurrent adapters, mixed-adapter batches, grouped GEMMs, batched optimizer steps, and MoE support remain in [tinker][fsdp] Concurrent Multi-LoRA training #1938.
  • No new trainer capacity setting is introduced; adapter capacity is bounded by host memory. Existing max_loras and max_cpu_loras continue to configure vLLM serving.

Verification

  • CPU state-isolation tests cover independent weights, gradients, Adam moments/steps, signature validation, parameter-layout validation, deletion, and swap recovery.
  • Ruff, Black, and git diff --check pass.
  • Two-H100 Modal Tinker run: 9 passed in 8m47s.
  • Real FSDP forward/backward and optimizer updates passed.
  • Second-adapter registration and training completed without rebuilding the worker group.
  • Alternating identical updates remained bit-identical across two adapters.
  • Per-adapter LoRA export, vLLM load, and sampling completed successfully.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a CPU-backed adapter storage mechanism (FSDPAdapterStore) for the FSDP LoRA worker, enabling multi-tenant LoRA training by swapping local parameter shards, gradients, and optimizer states between GPU and pinned CPU memory. It integrates this store into the FSDP worker and training backend, and adds comprehensive unit and integration tests. However, a critical runtime issue was identified in adapter_store.py where torch.empty_like is called with the unsupported pin_memory argument, which will raise a TypeError and must be corrected.

def _cpu_copy(tensor: torch.Tensor) -> torch.Tensor:
local = _local_tensor(tensor).detach()
pin_memory = local.device.type == "cuda" or local.is_pinned()
result = torch.empty_like(local, device="cpu", pin_memory=pin_memory)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In PyTorch, torch.empty_like does not support the pin_memory argument. Calling it with pin_memory=pin_memory will raise a TypeError at runtime. To allocate pinned CPU memory directly, use torch.empty with the tensor's shape, dtype, and layout instead.

Suggested change
result = torch.empty_like(local, device="cpu", pin_memory=pin_memory)
result = torch.empty(
local.shape,
dtype=local.dtype,
layout=local.layout,
device="cpu",
pin_memory=pin_memory,
)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant