Add Transformer Engine support to Praxis#36

Open

ashors1 wants to merge 16 commits intogoogle:mainfrom

nvjax-svc-0:patch/te_support

Commits on Aug 26, 2024

Adding TE support

mingxu1067
authored and
terrykong
committed
Fix missing vars wiht PP.

mingxu1067
authored and
terrykong
committed
Add checkpoint_policy checker for fused attn + dropout

zlsh80826
authored and
terrykong
committed
Support more TE configurations

zlsh80826
authored and
terrykong
committed
Change the gated activations orders

zlsh80826
authored and
terrykong
committed
Remove RoPE restriction from DPA module

zlsh80826
authored and
terrykong
committed
Add rotary_pos_emb_group dispatch

zlsh80826
authored and
terrykong
committed
Fix the missing .cls

zlsh80826
authored and
terrykong
committed
Fixed the unexpected input sharding pattern when TE enabled.

mingxu1067
authored and
terrykong
committed
Addind a comment to get_input_bld

mingxu1067
authored and
terrykong
committed
Leverage the origianl input bld when ENABLE_SP=False

mingxu1067
authored and
terrykong
committed
Use causal_padding instead of padding

zlsh80826
authored and
terrykong
committed
Fix rope convert

zlsh80826
authored and
terrykong
committed
Add test_te_helper.py

zlsh80826
authored and
terrykong
committed
Add LoRA support in TE (#11)

hemildesai
authored and
terrykong
committed

Commits on Dec 6, 2024

Fix missing mutalbes and carry_init of TE + FP8 + PP

mingxu1067
authored and
ashors1
committed