Skip to content

[AutoParallel] GPT support shared parameters #10783

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 7 commits into from
Closed

[AutoParallel] GPT support shared parameters #10783

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
90db0e7
init shared parameters
waliwali777 Jul 10, 2025
18b121b
update shared_parameters
waliwali777 Jul 10, 2025
080e701
fix status
waliwali777 Jul 15, 2025
4b3b571
update shard param placements
waliwali777 Jul 16, 2025
e4f88f4
update shard param placements
waliwali777 Jul 16, 2025
49ba28c
fix lint
waliwali777 Jul 17, 2025
20c3b84
update loss base
waliwali777 Jul 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions paddlenlp/transformers/gpt/modeling_auto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -658,10 +658,10 @@ def __init__(
config.hidden_size,
)
self.word_embeddings.weight = dist.shard_tensor(
self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Replicate()]
self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(0)]
Copy link
Contributor

@liym27 liym27 Jul 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

疑问：这里为什么要行切？ embedding 算子是不支持行切的，所以实际上还是要 Allgather 的
如果要行切，是否能替换成 c_embedding

Copy link
Contributor Author

@waliwali777 waliwali777 Jul 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

因为 lmhead 层为了支持 parallel_cross_centropy，所以 weight 必须是 行切；因为共享参数需要切分状态是一致的，所以参数同步时会引入一些通信。
经过测试 lmhead weight 和 embedding weight 不同切分状态的组合，发现同为行切时，此时性能最佳；虽然embedding 计算时会引入 allgather，但该该通信总耗时是最少的
替换成 c_embedding 是可行的，可以之后 PR 中再支持

)
self.position_embeddings.weight = dist.shard_tensor(
self.position_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)]
self.position_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(0)]
)
self.dropout = nn.Dropout(config.hidden_dropout_prob)

Expand Down Expand Up @@ -699,6 +699,7 @@ def forward(self, input_ids, position_ids=None, inputs_embeddings=None):
# The 'with' block ensures the correct seed context is used
with seed_guard_context(current_seed):
embeddings = self.dropout(embeddings)
embeddings = dist.reshard(embeddings, get_mesh(), [dist.Shard(0), dist.Replicate()])
return embeddings


Expand Down
25 changes: 24 additions & 1 deletion paddlenlp/transformers/gpt/modeling_auto_pp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,10 @@ def manual_model_split(model, stage_idx, group, mode, pp_degree):

layer_lists = model.layers

shared_params_names = {"gpt_shared_weight": ["embedding_0.w_0.dist", "gptlm_head_auto_0.w_0.dist"]}

shared_mp = build_shared_param_map(model, shared_params_names)

def _build_stage(model, stage_idx, group):
new_model = None
if stage_idx == 0:
Expand All @@ -151,7 +155,7 @@ def _build_stage(model, stage_idx, group):
new_model = GPTChunk(
layer_lists[stage_idx * chunk_size : (stage_idx + 1) * chunk_size], is_first=False, is_last=False
)
stage = PipelineStage(new_model, stage_idx, chunk_num, group=group)
stage = PipelineStage(new_model, stage_idx, chunk_num, group=group, shared_parameters=shared_mp)
return stage

stages = []
Expand All @@ -161,6 +165,25 @@ def _build_stage(model, stage_idx, group):
return stages


def build_shared_param_map(model, shared_params_names):
shared_mp = []
for key, pair in shared_params_names.items():
assert len(pair) == 2, "Only exactly two parameters are supported for sharing."
ori_name = pair[0]
sync_name = pair[1]
ori_param = get_param_from_name(ori_name, model)
sync_param = get_param_from_name(sync_name, model)
shared_mp.append({"params": [ori_param, sync_param]})
return shared_mp


def get_param_from_name(param_name, model):
for param in model.parameters():
if param.name == param_name:
return param
raise ValueError(f"{param_name} not found in model parameters")


def get_gpt_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group):
assert mode in ["VPP", "1F1B", "FThenB"]
stages = manual_model_split(model, group.rank, group, mode, pp_degree)
Expand Down
16 changes: 8 additions & 8 deletions scripts/distribute/ci_case_auto.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2506,11 +2506,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2() {
ips=-1
mem=-1
echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
loss_base=10.55853653 # output of dropout is different after supporting spmd
loss_base=10.55727577 # output of dropout is different after supporting spmd
ips_base=-1
mem_base=-1
if [ $IS_A100 -ne 0 ];then
loss_base=10.56019211 # after add dropout spmd
loss_base=10.56668472 # after add dropout spmd
fi
check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
echo "=========== $FUNCNAME run end ==========="
Expand Down Expand Up @@ -2578,11 +2578,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2-MP2() {
ips=-1
mem=-1
echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
loss_base=10.5657959 # output of dropout is different after supporting spmd
loss_base=10.49585533 # output of dropout is different after supporting spmd
Copy link
Contributor

@liym27 liym27 Jul 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

用了共享参数后，loss 差异这么大吗

Copy link
Contributor Author

@waliwali777 waliwali777 Jul 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

我这边是发现隐藏层层数较少时，共享参数 + 其他切分状态修改， loss 受影响比较大

ips_base=-1
mem_base=-1
if [ $IS_A100 -ne 0 ];then
loss_base=10.5760107 # after add dropout spmd
loss_base=10.51038742 # after add dropout spmd
fi
check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
echo "=========== $FUNCNAME run end ==========="
Expand Down Expand Up @@ -2651,11 +2651,11 @@ function llm_gpt_dygraph_auto_bs8_fp32_DP2-MP2-PP2() {
mem=-1
echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
# loss_base=10.59993172 # note: need to debug
loss_base=10.57174778 # output of dropout is different after supporting spmd
loss_base=10.49603939 # output of dropout is different after supporting spmd
ips_base=-1
mem_base=-1
if [ $IS_A100 -ne 0 ];then
loss_base=10.57701015 # after add dropout spmd
loss_base=10.51580238 # after add dropout spmd
fi
check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
echo "=========== $FUNCNAME run end ==========="
Expand Down Expand Up @@ -2724,11 +2724,11 @@ function llm_gpt_dygraph_auto_bs8_fp16_DP2-MP2-PP2() {
mem=-1
echo "result: loss=$loss ips=$ips mem=$mem loss_md5=$loss_md5"
# loss_base=10.58456802 # note: need to debug
loss_base=10.57304478
loss_base=10.49809837
ips_base=-1
mem_base=-1
if [ $IS_A100 -ne 0 ];then
loss_base=10.57861042 # after add dropout spmd
loss_base=10.51762962 # after add dropout spmd
fi
check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
echo "=========== $FUNCNAME run end ==========="
Expand Down
Loading