Skip to content

Add seq parallelism for attention and MoE MLP #1328

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 53 commits into from
Closed

Add seq parallelism for attention and MoE MLP #1328

Changes from all commits
Commits
Show all changes
53 commits
Select commit Hold shift + click to select a range
7b8f711
[exp] seq exp sharding
ZhiyuLi-goog Dec 29, 2024
802313a
update
ZhiyuLi-goog Dec 29, 2024
d8d4595
update
ZhiyuLi-goog Dec 29, 2024
b7f3225
merge for sp
suexu1025 Feb 27, 2025
7ed5fd1
fix merge parts
suexu1025 Feb 27, 2025
a1c6973
update merge confict base config
suexu1025 Feb 27, 2025
3f2d278
update to fix sharding mismatch
suexu1025 Feb 28, 2025
3e06ebb
update sub_seq for masks
suexu1025 Feb 27, 2025
d23d27b
update sharding axis
suexu1025 Feb 27, 2025
924ce77
update with reshape
suexu1025 Feb 28, 2025
b62812d
solve merge conflict
suexu1025 Mar 1, 2025
746f4a3
update for generate sharding
suexu1025 Feb 28, 2025
a6d345c
enable compute_axis configurable in mixtral model
suexu1025 Mar 4, 2025
e06c3d6
address output_logits sharding
suexu1025 Mar 5, 2025
65a64d4
clean up
suexu1025 Mar 5, 2025
23cd85f
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 5, 2025
10a9d82
update
suexu1025 Mar 5, 2025
0cca6df
update
suexu1025 Mar 6, 2025
cd005f3
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 6, 2025
ebae8e0
fix tests
suexu1025 Mar 6, 2025
2e0c459
added contition for non-sharded kernel for cp during inference only
suexu1025 Mar 6, 2025
37c843e
update
suexu1025 Mar 6, 2025
b63c63b
bug fix
suexu1025 Mar 7, 2025
9b32dc0
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 7, 2025
82d7fc3
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 7, 2025
4007e7c
fix tests
suexu1025 Mar 7, 2025
72f2a90
adddress comment
suexu1025 Mar 7, 2025
8da48f5
update
suexu1025 Mar 7, 2025
8a43dd5
address comments
suexu1025 Mar 7, 2025
56deeda
address comments
suexu1025 Mar 7, 2025
1c6be59
revert
suexu1025 Mar 7, 2025
bd0e199
address lint
suexu1025 Mar 7, 2025
44d646f
reformat for lint
suexu1025 Mar 7, 2025
5172068
update MOE test
suexu1025 Mar 7, 2025
d6787c3
add comment to explain grouping in generate_mask for moe model
suexu1025 Mar 7, 2025
f964acd
address the comments
suexu1025 Mar 8, 2025
930d77b
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 8, 2025
c5174de
update to fix tests
suexu1025 Mar 8, 2025
5c3fe75
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 8, 2025
b86e035
seperate yml for inference
suexu1025 Mar 8, 2025
cf8f0ec
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 10, 2025
e96340e
update to address training perf difference
suexu1025 Mar 11, 2025
7446563
update
suexu1025 Mar 11, 2025
3b5346f
revert back mask_shape for tests
suexu1025 Mar 11, 2025
c7ec0a4
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 11, 2025
e3d56c0
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 11, 2025
b7dcb1e
added back reshape and clean up merge changes
suexu1025 Mar 11, 2025
c859b4c
address comment to remove reshape
suexu1025 Mar 11, 2025
7d91629
update with different softmaxt score for inference/training for mask_…
suexu1025 Mar 11, 2025
1a0fdb3
lint
suexu1025 Mar 11, 2025
4d94d95
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 11, 2025
9a60574
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 11, 2025
9a12bd2
Merge branch 'main' into qinwen/sharding_merge_main
suexu1025 Mar 11, 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

These merge commits were added into this branch cleanly.

There are no new changes to show.