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.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for Expert Parallel (EP) + FSDP2 + Multi-LoRA SFT for DeepSeek-V4. It implements a new TargetParameterLoraManager to manage LoRA target parameters (especially for MoE experts), integrates it with the MultiLora and MultiLoraTransformersModel classes, and updates the native FSDP strategy to handle these parametrizations. The review feedback suggests supporting both 2D PEFT and 3D formats when loading checkpoints in set_state_dict to allow loading standard PEFT checkpoints, and fixing a NameError in the test file where a commented-out function is called in the __main__ block.
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.
| if __name__ == "__main__": | ||
| test_target_parameter_state_dict_loads_with_peft() No newline at end of file |
There was a problem hiding this comment.
The test file calls test_target_parameter_state_dict_loads_with_peft() in the __main__ block, but that function is commented out on lines 138-168. This will cause a NameError when running the file directly. We should call test_multilora_state_dict_round_trips_target_parameters() instead.
| if __name__ == "__main__": | |
| test_target_parameter_state_dict_loads_with_peft() | |
| if __name__ == "__main__": | |
| test_multilora_state_dict_round_trips_target_parameters() |
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：
