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

Open

apoorvtintin wants to merge 34 commits intoapple:mainfrom apoorvtintin:mainline_upstream_do_not_merge

+707-148

Commits on Dec 10, 2024

Support fine grained activation sharding. (patrick-toulme#21 )

patrick-toulme
authored and
apivovarov
committed
boilerplate code for neuron

apoorvtintin
authored and
apivovarov
committed
Intermediate commit for others

apoorvtintin
authored and
apivovarov
committed
fuji 70B v2 runs

apoorvtintin
authored and
apivovarov
committed
add flash attention and remat

apoorvtintin
authored and
apivovarov
committed
Maybe shard bug fixes

patrick-toulme
authored and
apivovarov
committed
Support fine grained embedding table sharding.

patrick-toulme
authored and
apivovarov
committed
rotary replaced

apoorvtintin
authored and
apivovarov
committed
use custom remat function

apoorvtintin
authored and
apivovarov
committed
clean flash attention, enable GQA since it works

apoorvtintin
authored and
apivovarov
committed
fix evaler

apoorvtintin
authored and
apivovarov
committed
fixed non shared lm head

apoorvtintin
authored and
apivovarov
committed
fix remat

apoorvtintin
authored and
apivovarov
committed
Fix Pytest skipping.

lipovsek-aws
authored and
apivovarov
committed
fix neuron imports to make unit tests run (patrick-toulme#27 )

lipovsek-aws
authored and
apivovarov
committed
flash attention & input data sharding test
apoorvtintin
committed

Commits on Dec 11, 2024

Commits on Dec 19, 2024

import neuron_attention earlier (patrick-toulme#38 )
apoorvtintin
authored

Commits on Dec 27, 2024

import jax_neuronx to enable buffer donation
apoorvtintin
committed

Commits on Jan 15, 2025

Use different remat strategy on 8 nodes
apoorvtintin
committed