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.

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

Jump to bottom

[DO-NOT-MERGE] PR encompassing all changes needed to support neuron on Axlearn #886

Open
wants to merge 32 commits into
base: main
Choose a base branch
from
Open

[DO-NOT-MERGE] PR encompassing all changes needed to support neuron on Axlearn #886

wants to merge 32 commits into from

Conversation

apoorvtintin
Copy link

@apoorvtintin apoorvtintin commented Dec 11, 2024 •
edited
Loading

This PR contains all of the changes needed to run fuji family of models on TRN2/1. This PR is actively being broken down to create smaller and manageable PRs that can be reviewed/merged to Axlearn.

Smaller PRs created

patrick-toulme and others added 28 commits December 10, 2024 13:11
@patrick-toulme @apivovarov
Support fine grained activation sharding. (#21)
783b5ab
@apoorvtintin @apivovarov
boilerplate code for neuron
2bd3956
@apoorvtintin @apivovarov
Intermediate commit for others
608b63d
@apoorvtintin @apivovarov
fuji 70B v2 runs
0e39d18
@apoorvtintin @apivovarov
add flash attention and remat
1445b4d
@patrick-toulme @apivovarov
Maybe shard bug fixes
6c70c6a
@patrick-toulme @apivovarov
Support fine grained embedding table sharding.
c8e79df
@apoorvtintin @apivovarov
rotary replaced
eb4edf6
@apoorvtintin @apivovarov
use custom remat function
01a5ff5
@apoorvtintin @apivovarov
clean flash attention, enable GQA since it works
b83d68c
@apoorvtintin @apivovarov
fix evaler
8062384
@apoorvtintin @apivovarov
fixed non shared lm head
ce67d13
@apoorvtintin @apivovarov
fix remat
4cd7d09
@lipovsek-aws @apivovarov
Fix Pytest skipping.
d6bae35
@lipovsek-aws @apivovarov
fix neuron imports to make unit tests run (#27)
8ee79cc 
* fix neuron attention to make unit tests run

* remove top level iport from flash attention utils

* revert LNC change
@apoorvtintin
flash attention & input data sharding test
8de3b9c
@apoorvtintin
clean up flash attention
3c03930
@patrick-toulme
Logits in FP32
28b97a3
@lipovsek-aws
fix regressions and messages (#29)
a80b674 
* fix regressions and messages

* tpu_health_check.py import

* not defined variables

* full name for neuron
@apoorvtintin
make full model default
482b7fa
@apoorvtintin
fix host array test
e9e3279
@aws-zhehongb
duplicate the kernel if we cannot shard on num_heads (#30)
d1d7da8 
Merge fix for heads % 2 != 0
@apoorvtintin
enable more neuron attention test
dfa013e
@apoorvtintin
set num_layers to original value for 7B
adadc7c
@apoorvtintin
add fsdp tp meshes for all fuji models
01ebc3f
@apoorvtintin
revert train_batch size change
041a54e
@apoorvtintin
put sharding change for embedding under a flag
c855cd6
@apoorvtintin
fixif input sharding is none for evaler
4a78272
@apoorvtintin apoorvtintin requested a review from a team as a code owner December 19, 2024 00:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ruomingp ruomingp Awaiting requested review from ruomingp ruomingp is a code owner

@markblee markblee Awaiting requested review from markblee markblee is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@apoorvtintin @patrick-toulme @lipovsek-aws @aws-zhehongb @indhub