NVIDIA / NVIDIA/Megatron-LM

Add functional tests for NVFP4 precision

Open
#4,161 2 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

We need functional tests for NVFP4 native (packed) weights to ensure end-to-end training correctness in CI, similar to the existing FP8 native weight functional tests (e.g., gpt3_weekly_mcore_tp2_pp2_current_scaling_native_fp8_tp_pp_sp_tp_overlap).

***Background***
The fp4_native_weights PR adds support for NVFP4 packed weights in the DDP pipeline (--fp4-param-gather). Currently only unit tests exist (tests/unit_tests/test_fp4_param.py). There are no functional tests under tests/functional_tests/test_cases/ covering FP4.

FP8 has multiple functional tests with --fp8-param-gather: true that serve as a good template:

gpt/gpt3_weekly_mcore_tp2_pp2_current_scaling_native_fp8_tp_pp_sp_tp_overlap
gpt/gpt3_weekly_mcore_tp4_cp2_current_scaling_native_fp8_tp_sp_cp_tp_overlap

***Proposed tests***
Add two new functional test cases:

1. NVFP4 baseline (without --fp4-param-gather)
Standard NVFP4 quantization via TE, no native weight packing in DDP
Config: --fp4 e2m1 --fp4-recipe nvfp4 --use-distributed-optimizer --tensor-model-parallel-size 2 --pipeline-model-parallel-size 1 --sequence-parallel
Purpose: establishes a baseline loss curve for FP4 training
2. NVFP4 with native param gather (--fp4-param-gather)
Same as above but with --fp4-param-gather: true
Purpose: verifies that native NVFP4 packed storage + all-gather produces the same loss curve as the non-packed baseline
Should also enable --overlap-param-gather and --overlap-grad-reduce to exercise the overlap paths
Both tests should:
Track lm loss, iteration-time, and mem-allocated-bytes metrics
Run with --bf16, --use-distributed-optimizer, --transformer-impl transformer_engine
Use a small model (e.g., 4-12 layers, hidden 256-512) to keep runtime reasonable

***Hardware requirement***
NVFP4 GEMM is supported on Blackwell and newer. The --fp4-param-gather tests specifically need TE >= 2.7.0.dev0. These should be tagged as weekly tests (similar to the FP8 native weight tests) since they require specific hardware.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.