feat: Add LoRA (Low-Rank Adaptation) support for efficient model fine-tuning#108
Open
chen2021673 wants to merge 7 commits intomasterfrom
Open
feat: Add LoRA (Low-Rank Adaptation) support for efficient model fine-tuning#108chen2021673 wants to merge 7 commits intomasterfrom
chen2021673 wants to merge 7 commits intomasterfrom
Conversation
- Add LoRA module infrastructure with configurable rank, alpha, dropout - Implement LoRALinear wrapper for seamless integration with Linear layers - Support tensor parallelism via LoRAParallelLinear - Add LoRAModel utility for managing multiple LoRA layers - Integrate LoRA configuration and utilities - Add GPT2 example demonstrating LoRA fine-tuning - Include comprehensive usage documentation and test suite Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Refactor LoRA config construction with proper target module parsing - Add GetLoRAModel for in-place LoRA layer injection - Fix DDP reducer to correctly handle LoRA parameters
- Fix RowParallel/ColumnParallel LoRA input handling to match base module behavior - Add shape-based defensive checks for TP/SP consistency - Move TP/SP communication helper function declarations to utils.h - Move getter implementations from header to .cc file - Add unit test for SaveLoRAWeights/LoadLoRAWeights functionality Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Chamberlain0w0
requested changes
Mar 2, 2026
| // LoRA A: [rank, in_features] - replicated across TP ranks (implemented as Linear) | ||
| // LoRA B: [out_features_per_partition, rank] - sharded like base weight (implemented as ColumnParallelLinear with | ||
| // gather_output) | ||
| class LoRAColumnParallelLinear : public nn::CloneableModule<LoRAColumnParallelLinear> { |
Contributor
There was a problem hiding this comment.
感觉是不是可以继承自原 ColumnParallelLinear，篇幅上可以省一些基类的成员定义和 getter
Contributor
Author
There was a problem hiding this comment.
这里不建议继承，因为LoRA 是在原 Linear 上叠加增量方法（典型decorator，LoRA 不管并行细节，把通信留给 base），不是一种新的 ColumnParallelLinear，此时组合优于继承。继承会让基类和 base_module_ 各自维护一套 weight / flags，容易不一致。
| continue; | ||
| } | ||
|
|
||
| if (type == Linear::kType) { |
Contributor
There was a problem hiding this comment.
这个文件里从这里开始有比较多这种三个 if 判断，但实际上就是一个 class name 的差异的代码，感觉可以采取一些更优雅的写法
Contributor
Author
There was a problem hiding this comment.
这里最多只能再抽一个公共函数减少每个if分支里的内容，因为type是运行期确定的，不能再用模板简化分支数量了。但这里逻辑应该不会再有增量了，我认为可以接受
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Added LoRA (Low-Rank Adaptation) support for parameter-efficient fine-tuning. This feature significantly reduces the number of trainable parameters through low-rank decomposition, enabling efficient fine-tuning of large models.
Changes
New Features
LoRA Infrastructure (
infini_train/include/nn/lora/):lora_config.h/cc- LoRA configuration (rank, alpha, dropout)lora_linear.h/cc- LoRA linear layer wrapperlora_model.h/cc- Multi-LoRA layer managementlora_parallel_linear.h/cc- Tensor parallelism supportlora_utils.h/cc- Utility functionsTests:
test/lora/test_lora.cc- Unit testsDocumentation:
docs/lora_usage.md- Usage documentationExamples:
example/gpt2/main.cc- Added LoRA training exampleBuild:
CMakeLists.txt- Added test_lora build targetTest Result
精度：
性能：
llama3 运行结果对比：