Skip to content

nvfp4 training with hints Megatron #2171

Closed

Description

@zobeideThePlayer

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

  1. ptrendx commented on Sep 17, 2025

    @ptrendx
    Member

    Please see PR #2177

  2. zobeideThePlayer commented on Sep 18, 2025

    @zobeideThePlayer
    ContributorAuthor

    Please see PR #2177

    Awesome! Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions