feat: add 1F1B schedule #96
Merged
Conversation
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
JYMiracle305
force-pushed
the
add_1F1B
branch
3 times, most recently
from
496bbfd to
7108a12
Compare
JYMiracle305
force-pushed
the
add_1F1B
branch
2 times, most recently
from
3726518 to
9af4751
Compare
Contributor
Author
|
新增超参数virtual_pipeline_parallel（vpp_size），表示PP场景对stage进行虚拟切分的块数，PP场景将模型切分成pp_size * vpp_size块，分配到对应的设备上；重构后统一不同调度策略对上层的接口，构造调度器PipelineParallelScheduler时根据不同策略填充任务Task表，任务表中保存子任务（关联chunk、microbatch和当前属于正/反向），训练时上层调用StepMicroBatches，StepMicroBatches内部遍历任务表。 virtual_pipeline_parallel为1时，调度表示如下： virtual_pipeline_parallel大于1时，调度表示如下： 多机训练gpt2，配置 DDP=2，TP=2(SP=ON)，PP=2(VPP=2) 多机训练LLaMA3，配置 DDP=2，TP=2(SP=ON)，PP=2(VPP=2) |
JYMiracle305
force-pushed
the
add_1F1B
branch
from
9af4751 to
a413a6e
Compare
JYMiracle305
requested review from
Chamberlain0w0 and
kilinchange
and removed request for
kilinchange
kilinchange
requested changes
Dec 24, 2025
example/gpt2/net.cc
Outdated
| std::vector<std::shared_ptr<infini_train::Tensor>> | ||
| GPT2::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) { | ||
| int pp_rank = nn::parallel::pp_rank; | ||
| void GPT2::BuildChunks() { |
Contributor
There was a problem hiding this comment.
针对 Transformer 模型的话，BuildChunks 也可以合并，gpt2/llama 仅是一个 pos_emb 的区别，加个 if 判断就可以
JYMiracle305
force-pushed
the
add_1F1B
branch
3 times, most recently
from
0f5628b to
aeb8ee0
Compare
kilinchange
requested changes
Dec 25, 2025
JYMiracle305
force-pushed
the
add_1F1B
branch
2 times, most recently
from
f8b086c to
c22da40
Compare
JYMiracle305
force-pushed
the
add_1F1B
branch
from
c22da40 to
213a164
Compare
feat: decouple DDP and PP for 1F1B pipeline in 3d parallelism
kilinchange
requested changes
Jan 4, 2026
kilinchange
approved these changes
Jan 5, 2026
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
No description provided.