Add functional tests for NVFP4 precision
- 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
Assessment
This issue has not been assessed yet.