Documentation

URL or Section:

• src/MaxText/configs/base.yml (lines 91-98)
• src/MaxText/pyconfig.py (line 287)
• src/MaxText/layers/quantizations.py (lines 642-692)
• docs/explanations/quantization.md (lines 51-56)

Describe the problem:

The inline comments in base.yml documenting the quantization config option contain incorrect option names that don't match the actual implementation in the codebase. This could cause users to use invalid configuration values.

Issues found:

1. CRITICAL - Wrong option name:

   • base.yml documented 'nanoo_fp8'
   • Actual valid option in code: 'fp8_nanoo' (see pyconfig.py line 287 and quantizations.py line 682)
   • Users following the comments would use an invalid option that doesn't exist!

2. Missing option:

   • base.yml did not mention 'fp8_gpu'
   • This option exists in pyconfig.py valid_quant_methods (line 287) and has implementation in quantizations.py (line 673)

3. Misleading descriptions:

   • base.yml described 'fp8' as "for 8-bit floating-point GeMMs on NVIDIA GPUs"
   • But 'fp8_gpu' is the actual NVIDIA-specific option
   • 'fp8' is the generic 8-bit floating-point quantization

4. Inconsistency across documentation sources:

   • pyconfig.py: defines valid options as ("", "int8", "fp8", "fp8_full", "fp8_gpu", "fp8_nanoo")
   • quantizations.py: implements all these options in the case statement
   • quantization.md: documents some of these options
   • base.yml: had different names and missing options

Expected behavior:
The comments in base.yml should exactly match:

1. The valid option names defined in pyconfig.py
2. The case statements implemented in quantizations.py
3. The documentation in quantization.md

Impact:

• Users following the base.yml comments might try to use 'nanoo_fp8' which would fail validation, causing confusion and errors
• Misleading description like "fp8 is for GPU" could cause users to incorrectly choose 'fp8' when they specifically need GPU optimization ('fp8_gpu'), or vice versa, leading to suboptimal performance or unexpected behavior

Additional Context

AI GDE / Kakao