Add Transformer Engine support to Paxml#46

Open

ashors1 wants to merge 12 commits intogoogle:mainfrom nvjax-svc-0:patch/te_support

+155-50

Commits on Sep 30, 2024

add TE support

ashors1
authored and
terrykong
committed
Adding dropout support when enabling TE.

mingxu1067
authored and
terrykong
committed
Set deterministic=True for inference.

mingxu1067
authored and
terrykong
committed
Fix the excluded list for excluded_for_learner

zlsh80826
authored and
terrykong
committed
Adapting to TE/JAX/Custom_partitioning.

mingxu1067
authored and
terrykong
committed
Adding TE-compatiable PipelinedTransformer

mingxu1067
authored and
terrykong
committed
Apply OWG to TE's FP8 meta

mingxu1067
authored and
terrykong
committed
Remove Praxis related setup (Moving to Praxis TE/Patch)

mingxu1067
authored and
terrykong
committed
Fix missing DEFAULT_INIT_MUTABLE_LIST

mingxu1067
authored and
terrykong
committed
Revert mutable kwarg in abstract_init_with_metadata in init checkpoint rule

hemildesai
authored and
terrykong
committed
Support FM32 to OWG parameters.

mingxu1067
authored and
terrykong
committed

Commits on Oct 29, 2024

fix OOM with TE
ashors1
committed