[tinker][fsdp] Concurrent Multi-LoRA training

[atemaguer]

Stacked on #1956. Please review the single-resident FSDP adapter store there first. Once #1956 lands, this PR contains only the concurrent-residency and grouped-GEMM delta.

Adds concurrent multi-tenant LoRA training to the SkyRL-Train FSDP Tinker backend. Multiple Tinker clients retain independent resident LoRA adapters and can execute in one mixed-adapter FSDP forward/backward pass.

Architecture

• MultiLoRAManager allocates max_lora_adapters resident slots before FSDP wrapping. Each MultiLoRALinear owns independent A/B parameters per slot while sharing the frozen FSDP base layer.
• A row-level adapter_indices tensor routes each sample. Active tokens and adapter banks are packed for torch.nn.functional.grouped_mm.
• One AdamW optimizer retains independent state for every slot. Batched optimizer requests use slot-specific parameter groups and preserve each adapter's hyperparameters, gradients, and step state.
• Per-adapter checkpoints contain LoRA parameters, optimizer state, optimizer hyperparameters, seed, and signature metadata and can be restored into a different resident slot.

MoE support

• Fused expert banks are wrapped with MultiLoRAExperts; routing uses composite (expert, adapter) groups for gate, up, and down projections.
• Frozen expert and LoRA A/B projections 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.
• Unaligned outputs such as Qwen 3.5's scalar shared-expert gate are padded for grouped MM and sliced back to their original width.

API and constraints

• trainer.policy.model.lora.implementation=concurrent opts into this path.
• max_lora_adapters controls trainer-side resident capacity. vLLM's existing max_loras and max_cpu_loras remain serving-side settings.
• Registration, deletion, forward/backward, optimizer steps, checkpoint save/load, and sampler export are adapter-aware.
• All resident adapters must share the initial (rank, alpha) signature.
• Grouped MM requires BF16 CUDA tensors on compute capability 8.0 or newer, aligned input/rank dimensions, sequence_parallel_size=1, and remove_microbatch_padding=False.

Throughput

The checked-in skyrl/benchmarks/bench_fsdp_multi_lora_tinker.py benchmark compares two workloads:

• Underfilled: one 512-token sequence per adapter. Concurrent execution can combine otherwise-small tenant batches.
• Saturated: each adapter independently submits max_tokens_per_microbatch / 512 sequences. A single tenant therefore fills a complete microbatch without other adapters.

Both implementations ran end-to-end through the Tinker API on separate H100 workers using Qwen/Qwen3.5-4B, rank 32, five warmup steps, and ten measured steps. Each step includes request scheduling, forward, backward, adapter swap/routing, and optimizer execution. The clients submit requests concurrently behind a barrier, and the benchmark records the engine's observed request coalescing sizes. Throughput uses median total step latency; the optimizer learning rate is zero to execute a stable repeated workload without weight drift.

Underfilled tenants

max_tokens_per_microbatch	Active adapters	Single resident	Concurrent grouped MM	Speedup
1024	2	551.6 tok/s	428.7 tok/s	0.78x
1024	4	512.6 tok/s	597.9 tok/s	1.17x
2048	2	551.8 tok/s	620.5 tok/s	1.12x
2048	4	512.6 tok/s	831.0 tok/s	1.62x
4096	2	552.2 tok/s	690.7 tok/s	1.25x
4096	4	512.4 tok/s	1223.6 tok/s	2.39x

Saturated tenants

max_tokens_per_microbatch	Active adapters	Single resident	Concurrent grouped MM	Speedup
1024	2	780.7 tok/s	603.6 tok/s	0.77x
1024	4	816.6 tok/s	634.9 tok/s	0.78x
2048	2	1559.1 tok/s	1194.5 tok/s	0.77x
2048	4	1632.1 tok/s	1256.8 tok/s	0.77x
4096	2	3100.5 tok/s	2388.0 tok/s	0.77x
4096	4	3243.3 tok/s	2984.0 tok/s	0.92x

The results support the compute-bound hypothesis: grouped MM does not improve throughput when each tenant already fills the token cap. Its benefit is coalescing underfilled tenants into larger work. For four underfilled adapters at a 4096-token cap, median forward/backward falls from 3.143s to 1.381s and optimizer time from 0.854s to 0.290s, producing the 2.39x end-to-end gain. With four saturated adapters at the same cap, grouped forward/backward is slower (5.203s versus 4.198s); batched optimization remains faster (0.290s versus 0.857s), but net throughput is 0.92x.

Reproduce one token cap on two GPUs (one command per GPU), then compare:

uv run --extra tinker --extra fsdp python -m skyrl.benchmarks.bench_fsdp_multi_lora_tinker \
  --implementation single --model Qwen/Qwen3.5-4B --rank 32 \
  --adapter-counts 2,4 --workloads underfilled,saturated \
  --sequence-length 512 --max-tokens-per-microbatch 4096 \
  --warmup-steps 5 --measured-steps 10 --output single-4096.json

uv run --extra tinker --extra fsdp python -m skyrl.benchmarks.bench_fsdp_multi_lora_tinker \
  --implementation concurrent --model Qwen/Qwen3.5-4B --rank 32 \
  --adapter-counts 2,4 --workloads underfilled,saturated \
  --sequence-length 512 --max-tokens-per-microbatch 4096 \
  --warmup-steps 5 --measured-steps 10 --output concurrent-4096.json

uv run python -m skyrl.benchmarks.bench_fsdp_multi_lora_tinker \
  --compare single-4096.json concurrent-4096.json --output comparison-4096.json

Repeat with token caps 1024 and 2048 for the complete table.

End-to-end validation

Real Tinker cookbook GSM8K SFT and RL workloads ran concurrently against Qwen/Qwen3.5-35B-A3B with four FSDP trainer H100s, two tensor-parallel vLLM H100s, and two rank-32 adapters:

• SFT completed 3/3 optimizer steps.
• RL completed three sampling rounds, 13 requests, four sampler syncs, and two trainable updates; the final zero-advantage batch was correctly skipped.
• RL rewards were 0.15625 and 0.0625; gradient norms were 290.65 and 638.80.
• SFT -> RL -> SFT sampling probes each returned two 128-token sequences with exactly 128 logprobs.
• The six-GPU workload exited cleanly.

Verification

• Focused local suite: 66 passed, 3 skipped.
• Benchmark and routing suite: 13 passed.
• Single-resident base production suite: 9 passed on two H100s.
• Grouped-MM CUDA forward, input-gradient, and adapter-gradient parity passed on H100.

[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"]