You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Is your feature request related to a problem? Please describe.
In the most updated NVIDIA/Megatron-LM repo, several hints suggest that it is possible to train using nvfp4 with Transformer Engine support:
in megatron/core/enums.py:
class Fp4Recipe(str, enum.Enum):
"""FP4 recipe names: nvfp4."""
nvfp4 = "nvfp4"
in megatron/training/augments.py
FP4 related arguments
group.add_argument('--fp4-format', default=None,
choices=['e2m1'],
help='Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass',
dest='fp4')
group.add_argument('--fp4-recipe', default='nvfp4',
choices=['nvfp4'],
help='Which fp4 recipe to use for FP4 tensors in the forward and backward pass',
dest='fp4_recipe')
group.add_argument('--fp4-param-gather', action='store_true',
help='Keep the compute param in fp4 (do not use any other intermediate '
'dtype) and perform the param all-gather in fp4.',
dest='fp4_param')
Explicitly, it requires a minimum version of te 2.7.0dev, as in fp4_utils.py:
if HAVE_TE:
if is_te_min_version("2.7.0.dev0"):
try:
from transformer_engine.pytorch.tensor.nvfp4_tensor import (
NVFP4Tensor as FP4_TENSOR_CLASS,
)
HAVE_TE_FP4_TENSOR_CLASS = True
However, when I check the most updated main branch or release 2.7 branch, such files do not exists yet.
Describe the solution you'd like
Is the proper te 2.7.0.dev0 branch released to public yet? If it is indeed, could you please kindly direct me to the correct branch/tag?
Activity
ptrendx commentedon Sep 17, 2025
Please see PR #2177
zobeideThePlayer commentedon Sep 18, 2025
Awesome! Thank you!