Skip to content
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.

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

fix(shard.py): fix isp unpack data indexes err in rotary emb #316

Merged
merged 4 commits into from
Sep 6, 2024

Conversation

huangting4201
Copy link
Collaborator

@huangting4201 huangting4201 commented Sep 5, 2024

bug修复1

问题描述:在当前的代码实现中,unpack data场景下会pop出indexes字段,而在进行rotary_emb计算的位置编码时需要用到该字段,如果不传的话,默认值是offsets=0,这在非序列并行下是没有问题的。但在isp算法,序列并行场景下,seq维度是切分的,这样的话,如果没有indexes字段与切分后的seq字段进行一一对应,会导致rotary_emb计算出现问题。而这在mtp/msp/fsp算法下都没有问题,这是因为这些算法的序列并行在进行linear计算前会把seq维度进行allgather,所以在rotary_emb计算时seq维度是全量的,因此直接使用默认值offsets=0就没有问题。

代码修复:主要在core/parallel/shard.py文件中,首先unpack data保留indexes字段,然后基于sp并行进行切分,同时在2D seq parallel时支持负载均衡式的切分indexes,与seq切分对齐。

bug修复2

问题描述:在dp4cp4并行配置下,跑真实训练数据,大概五百多步后会出现Nan grad norm的情况。对比实验发现,dp4并行,dp4tp4并行,与dp4hp4并行场景下均不会出现该问题。

代码修复:排查 ring attn 开源代码,更新 update_out_and_lse 函数中使用的计算算子,提高数值稳定性。修复该代码后,再进行loss曲线对比:
4a804fcd-37ee-4e5c-b0a7-b3f06c9c9747

7e7baa44-6d0b-4ce7-aaf6-c1c26209f90b

新特性1

因测试发现VocabSequenceParallelCrossEntropyLoss计算出的loss有问题,所以目前改成了对head output进行all2all处理后,仍使用TP并行计算loss的方式。复用parallel_output开关,若为True,则走FlashCrossEntropyLoss方法;若为False,则走nn.CrossEntropyLoss方法。loss曲线如下:
35a72af5-e4b7-4254-999a-b5e702d9ae26

新特性2

适配evaluate模块,支持isp序列并行。因valid data没有indexes字段,所以手动加上,使得其支持isp序列并行。

@sunpengsdu sunpengsdu merged commit 3dfb540 into develop Sep 6, 2024
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants