Support multi-LoRA training with EP + FSDP2#236
Conversation
![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 implements support for DeepSeek-V4 EP Multi-LoRA target parameters in MultiLoraTransformersModel, allowing multiple target-parameter LoRA adapters to reside in memory while activating only one at a time. The feedback highlights several critical issues in the target-parameter manager: a shape mismatch error in reset_slot when expert parallel is enabled due to unsharded initial weights, significant memory overhead from cloning the entire target parameter instead of just storing its ndim, and a broadcasting shape mismatch when computing delta weights for 2D parameters. Additionally, a potential NameError was identified in the new cookbook when resuming from checkpoints if the adapter list is empty.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…om parametrizations
…lete parameter restoration
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for EP + FSDP2 + Multi-LoRA SFT, including a DeepSeek-V4 cookbook and a new TargetParameterLoraManager to manage multi-LoRA target parameters. The review feedback highlights three key issues: a shape mismatch when loading full checkpoints in sharded EP environments, a performance bottleneck from dynamically registering and removing PyTorch parametrizations on every forward pass, and an incorrect gather dimension for twinkle_lora parameters in the expert state dict gather logic.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| @contextmanager | ||
| def activate(self, slot_name: str | None, disable_lora: bool = False): | ||
| if disable_lora or slot_name is None or slot_name not in self.lora_A: | ||
| yield | ||
| return | ||
|
|
||
| module = self.record.module | ||
| param_name = self.record.parameter_name | ||
| already_parametrized = nn.utils.parametrize.is_parametrized(module, param_name) | ||
| if not already_parametrized: | ||
| # LoRA weights change after EP + FSDP sharding, so they must be computed dynamically and not be pre-fixed. | ||
| # delta_weight = self.get_delta_weight(slot_name) # lora_weight = B @ A * scaling | ||
| requires_grad_before = self.base_parameter.requires_grad | ||
| nn.utils.parametrize.register_parametrization( | ||
| self.record.module, | ||
| self.record.parameter_name, | ||
| LoraParameterProxy(self, slot_name), | ||
| ) | ||
| module.parametrizations[param_name].original.requires_grad_(requires_grad_before) | ||
| try: | ||
| with nn.utils.parametrize.cached(): | ||
| yield | ||
| finally: | ||
| if not already_parametrized: | ||
| nn.utils.parametrize.remove_parametrizations( | ||
| self.record.module, | ||
| self.record.parameter_name, | ||
| leave_parametrized=False, | ||
| ) |
There was a problem hiding this comment.
Performance Bottleneck: Dynamic Parametrization Registration
Registering and removing PyTorch parametrizations (register_parametrization and remove_parametrizations) on every single forward/backward pass introduces significant Python runtime overhead, which can severely degrade training throughput (steps per second).
Suggested Improvement
Register the parametrization once during wrapper initialization (or during patch), and have the LoraParameterProxy dynamically query the active slot name from the wrapper. The activate context manager can then simply toggle the active slot name on the wrapper, which is a fast O(1) state change with zero overhead.
Here is how the refactored code would look:
class LoraParameterProxy(nn.Module):
def __init__(self, LoraWrapper):
super().__init__()
self.LoraWrapper = LoraWrapper
def forward(self, weight: torch.Tensor) -> torch.Tensor:
slot_name = self.LoraWrapper.active_slot
if slot_name is None or self.LoraWrapper.disable_adapters:
return weight
delta_weight = self.LoraWrapper.get_delta_weight(slot_name)
return weight + delta_weightAnd in TargetParameterLoraWrapper.__init__:
self.active_slot: str | None = None
self.disable_adapters = False
self._init_slots()
nn.utils.parametrize.register_parametrization(
self.record.module,
self.record.parameter_name,
LoraParameterProxy(self),
)And in TargetParameterLoraWrapper.activate:
@contextmanager
def activate(self, slot_name: str | None, disable_lora: bool = False):
old_slot = self.active_slot
old_disable = self.disable_adapters
self.active_slot = slot_name
self.disable_adapters = disable_lora
try:
with nn.utils.parametrize.cached():
yield
finally:
self.active_slot = old_slot
self.disable_adapters = old_disable| def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): | ||
| _ensure_dummy_zmq() | ||
| from twinkle.model.transformers.strategy.native_fsdp import _ep_expert_state_dict_gather_dim | ||
|
|
||
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_A.weight") == 0 | ||
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_A.weight") == 0 | ||
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_B.weight") == 1 | ||
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_B.weight") == 1 |
There was a problem hiding this comment.
Correctness Bug: Incorrect Gather Dimension for _twinkle_lora_ Parameters
In src/twinkle/model/transformers/strategy/native_fsdp.py, _ep_expert_state_dict_gather_dim returns 1 for any parameter containing lora_B:
def _ep_expert_state_dict_gather_dim(name: str) -> int:
if 'lora_B' in name:
return 1
return 0While this is correct for PEFT's native flattened target_parameters representation, it is incorrect for TargetParameterLoraWrapper's parameters (which contain _twinkle_lora_). In TargetParameterLoraWrapper, lora_B is stored as a 3D tensor of shape (num_experts, out_features, r), meaning the expert dimension is 0, not 1.
This will cause silent corruption or shape mismatches during checkpoint saving/loading or FSDP broadcasting of multi-LoRA expert adapters.
Solution
Update _ep_expert_state_dict_gather_dim in src/twinkle/model/transformers/strategy/native_fsdp.py to return 0 if _twinkle_lora_ is in the parameter name:
def _ep_expert_state_dict_gather_dim(name: str) -> int:
if '_twinkle_lora_' in name:
return 0
if 'lora_B' in name:
return 1
return 0We should also add a test case to verify this behavior.
| def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): | |
| _ensure_dummy_zmq() | |
| from twinkle.model.transformers.strategy.native_fsdp import _ep_expert_state_dict_gather_dim | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_A.weight") == 0 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_A.weight") == 0 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_B.weight") == 1 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_B.weight") == 1 | |
| def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): | |
| _ensure_dummy_zmq() | |
| from twinkle.model.transformers.strategy.native_fsdp import _ep_expert_state_dict_gather_dim | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_A.weight") == 0 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_A.weight") == 0 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_B.weight") == 1 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_B.weight") == 1 | |
| assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts._twinkle_lora_gate_up_proj.lora_B.lora_0.weight") == 0 |
3 participants
PR type
PR information
Background
Twinkle currently supports single-adapter EP + LoRA training on packed MoE expert weights (gate_up_proj / down_proj) via PEFT's target_parameters interface. The MultiLoRA framework enables multi-tenant adapter deployment but only supports target_modules-based LoRA (attached at nn.Module layer level), not target_parameters (raw Parameter tensors). PEFT does not natively support multiple adapters on target_parameters, creating a gap for multi-tenant LoRA in EP scenarios.
This PR
This PR introduces multi-LoRA training under EP + FSDP2 by extending MultiLoRA with a target_parameters multi-slot path, enabling direct attachment of tenant adapters to packed MoE expert weights. Key changes include physical slot allocation and tenant mapping, FSDP2 sharding compatibility, and preserved single-tenant activation semantics. This unifies MultiLoRA support across both LoRA attachment paradigms, enabling efficient multi-tenant fine-tuning of MoE models under EP + FSDP2.
Experiment results
Training loss curves for two tenants on DeepSeek-V4-Flash：
