[tinker][fsdp] Multi-LoRA adapter store + Tinker API - #1956
Conversation
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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.
| 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, | |
| ) |
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
FSDPAdapterStoreholds a pinned-CPU snapshot permodel_id, plus a pristine template used to initialize new adapters.(rank, alpha, target_modules, world_size)signature and rejects trainable non-LoRA parameters.API surface
create_modelprimes Adam state, captures the pristine adapter, and registers the initialmodel_id.create_modelcalls register CPU-backed adapter slots through the existing Tinker/Megatron adapter lifecycle.forward,forward_backward,optim_step, checkpoint operations, and sampler synchronization use the existingWorkerDispatch.ensure_active_adapterpath.Scope
max_lorasandmax_cpu_lorascontinue to configure vLLM serving.Verification
git diff --checkpass.9 passedin 8m47s.