FSDP Concurrent Multi-LoRA Training

[atemaguer]

Adds concurrent multi-tenant LoRA training to the SkyRL-Train FSDP Tinker backend. Multiple Tinker clients can create_model against the same server, each with an independent LoRA adapter sharing one frozen base model. Mixed-adapter requests execute in the same FSDP forward/backward pass.

Architecture

• New resident MultiLoRAManager allocates a fixed number of adapter slots before FSDP wrapping. Each MultiLoRALinear owns independent A/B parameters per slot while the frozen base layer remains shared and FSDP-sharded.
• A row-level adapter_indices tensor routes each sample to its adapter. Tokens are sorted by active adapter and processed with torch.nn.functional.grouped_mm; only active adapter banks are materialized for each invocation.
• One AdamW optimizer retains independent per-parameter state for every slot. Batched optimizer requests construct slot-specific parameter groups, preserving each adapter's learning rate, betas, epsilon, weight decay, gradients, and step state.
• Per-adapter training checkpoints include LoRA parameters, optimizer state, optimizer hyperparameters, seed, and signature metadata. Checkpoints can be restored into a different resident slot.

MoE support

• Fused expert banks are detected structurally and wrapped with MultiLoRAExperts. Routing uses composite (expert, adapter) groups for the expert gate, up, and down projections.
• Frozen expert GEMMs and LoRA A/B GEMMs use ragged grouped MM. Shared expert dense projections use the same row-level adapter routing.
• PEFT-compatible export expands fused expert banks into per-expert LoRA keys for vLLM loading.
• Unaligned dense outputs such as Qwen 3.5's scalar shared-expert gate are padded to the grouped-MM alignment and sliced back to their original width.

API surface

• trainer.policy.model.lora.implementation=concurrent opts into the resident FSDP path.
• max_lora_adapters controls trainer-side resident capacity; the existing max_loras and max_cpu_loras settings continue to control vLLM serving capacity.
• Additional policy model_ids must use the same (rank, alpha) signature as the first adapter. Registration, deletion, forward/backward, optimizer steps, checkpoint save/load, and sampler export are adapter-aware.
• The Tinker engine batches compatible optimizer requests, and the SkyRL-Train backend keeps mixed model IDs together for concurrent FSDP execution.

Constraints

• The grouped-MM path requires BF16 CUDA tensors on compute capability 8.0 or newer, with aligned input and rank dimensions.
• Concurrent FSDP LoRA currently requires sequence_parallel_size=1 and remove_microbatch_padding=False.
• Multimodal checkpoints are supported with policy.language_model_only=True.

Files

• New: skyrl/backends/skyrl_train/workers/fsdp/multi_lora.py, skyrl/benchmarks/bench_fsdp_multi_lora.py, and focused FSDP/Tinker multi-LoRA tests.
• Modified: the SkyRL-Train backend, FSDP worker/model wrapper, worker dispatch, training-batch plumbing, Tinker engine batching, and LoRA configuration.
• Removed: the checked-in Modal test integration; production GPU launchers used during development remain disposable and outside the repository.

Verification

• Focused suite: 27 passed, 2 skipped across the FSDP multi-LoRA layer and Tinker routing tests.
• Code quality: Ruff, Black, hardcoded-secret scan, and git diff --check pass.
• CUDA parity: grouped-MM forward, input gradients, and adapter gradients matched an independent per-token/per-expert reference on H100.
• Production training: Qwen/Qwen3.5-35B-A3B, 4x H100 80GB, two concurrent rank-32 adapters, three optimizer steps. Steps 2 and 3 executed as true mixed-adapter batches; losses decreased from 0.7536 to 0.00038 and 0.7907 to 0.00240. A distributed training checkpoint saved successfully.
• Remaining production validation: the colocated vLLM engine core failed during startup before adapter export/broadcast, so the Qwen 3.5 sampler-sync path was not validated by that run. Existing unit/integration coverage still exercises PEFT export and sampler routing.

[gemini-code-assist]

Code Review

This pull request introduces a concurrent multi-adapter LoRA implementation for FSDP2 policy paths, allowing multiple independent LoRA adapters to reside in memory and enabling batched optimizer steps and mixed-adapter forward/backward passes. The review feedback highlights several critical improvements for correctness and robustness: correcting the mathematical bounds for kaiming initialization on sharded DTensors, converting 3D weight tensors to lists of 2D tensors for torch._grouped_mm compatibility, and adding optimizer-presence guards across the FSDP worker and multi-LoRA manager to prevent crashes and support loading checkpoint weights in inference-only mode.

[gemini-code-assist]

Initializing the local shard of a sharded DTensor directly using kaiming_uniform_ is mathematically incorrect because kaiming_uniform_ calculates the fan-in and fan-out based on the shape of the tensor passed to it. If the parameter is sharded along the input or output dimension, the calculated fan-in/fan-out will be divided by the shard count, leading to incorrect initialization bounds and potential training instability.\n\nSince kaiming_uniform_ with a=math.sqrt(5) simplifies exactly to a uniform distribution with bounds [-1/sqrt(in_features), 1/sqrt(in_features)], we can calculate the bounds using the global shape of self.lora_A.weight to ensure correct and robust initialization regardless of the sharding layout.

    @torch.no_grad()\n    def reset_parameters(self, seed: Optional[int] = None) -> None:\n        if self.lora_A.weight.is_meta:\n            return\n        local_a = _local_tensor(self.lora_A.weight)\n        local_b = _local_tensor(self.lora_B.weight)\n        bound = 1.0 / math.sqrt(self.lora_A.weight.shape[1])\n        if seed is None:\n            local_a.uniform_(-bound, bound)\n        else:\n            devices = [local_a.device] if local_a.is_cuda else []\n            with torch.random.fork_rng(devices=devices):\n                torch.manual_seed(seed + _shard_seed_offset(self.lora_A.weight))\n                local_a.uniform_(-bound, bound)\n        nn.init.zeros_(local_b)

[gemini-code-assist]

The standard torch._grouped_mm API in PyTorch expects the second argument (weights) to be a list of 2D tensors (Tensor[]), not a single 3D tensor. Passing a 3D tensor directly will raise a TypeError on a real GPU. Converting the 3D weight tensors to lists of 2D tensors using list(lora_a.transpose(1, 2)) ensures compatibility with both the mock/fake implementation and the real PyTorch C++ kernel.

        intermediate = torch._grouped_mm(sorted_inputs, list(lora_a.transpose(1, 2)), offs=group_offsets)\n        sorted_delta = torch._grouped_mm(intermediate, list(lora_b.transpose(1, 2)), offs=group_offsets)

[gemini-code-assist]

In inference-only mode, the optimizer is not initialized (self.optimizer is None). However, we still need to load the adapter weights for inference. Making the optimizer argument optional in load_training_state allows loading weights without requiring an active optimizer.

    def load_training_state(\n        self,\n        model_id: str,\n        state: dict[str, object],\n        optimizer: Optional[torch.optim.Optimizer] = None,\n    ) -> None:

[gemini-code-assist]

Wrap the optimizer state restoration in an if optimizer is not None: block to support loading weights in inference-only mode where no optimizer is present.

            if optimizer is not None:\n                optimizer.state.pop(parameter, None)\n                saved_param_state = saved_optimizer_state.get(name)\n                if saved_param_state is None:\n                    continue\n                if not isinstance(saved_param_state, dict):\n                    raise ValueError(f"Concurrent LoRA checkpoint optimizer state for '{name}' is invalid")\n                restored_state: dict[str, object] = {}\n                for key, value in saved_param_state.items():\n                    if isinstance(value, torch.Tensor) and tuple(value.shape) == tuple(parameter.shape):\n                        restored_tensor = torch.zeros_like(parameter, dtype=value.dtype)\n                        _copy_from_full_tensor(restored_tensor, value)\n                        restored_state[key] = restored_tensor\n                    elif isinstance(value, torch.Tensor):\n                        restored_state[key] = value.detach().clone()\n                    else:\n                        restored_state[key] = deepcopy(value)\n                optimizer.state[parameter] = restored_state

[gemini-code-assist]

In inference-only mode, self.optimizer is None. If load_optimizer_states is False, we should still be allowed to load the checkpoint weights. Guard the optimizer check so it only raises an error if load_optimizer_states is True.

        if self.optimizer is None and load_optimizer_states:\n            raise RuntimeError("Optimizer is not initialized but load_optimizer_states is True")

[gemini-code-assist]

If self.optimizer is None (e.g., in inference-only mode), accessing self.optimizer.param_groups will raise an AttributeError. Guard the parameter group update with if self.optimizer is not None:.

        optimizer_hparams = self._multi_lora_manager.optimizer_hparams_for(model_id)\n        if optimizer_hparams and self.optimizer is not None:\n            for param_group in self.optimizer.param_groups:\n                param_group["lr"] = optimizer_hparams["learning_rate"]\n                param_group["betas"] = (\n                    optimizer_hparams["beta1"],\n                    optimizer_hparams["beta2"],\n                )\n                param_group["eps"] = optimizer_hparams["eps"]\n                param_group["weight_decay"] = optimizer_hparams["weight_decay"]