[zhou9402]

Ports miles' low-precision path: pure-torch quantization oracle, MXFP8 trainer-side export quantizer + --extra-high-precision-layers, offline HF→MXFP8 converter, NVFP4 primitives + QDQ fallback + fake-QAT TE hook, NVFP4 Megatron→HF export (modelopt_fp4 layout), offline HF→NVFP4 converter, and vllm-nvfp4-inloop.patch for in-loop NVFP4 weight updates.

Source: radixark/miles@ba28991834. Port is additive (new modules + new add_*_arguments() groups); slime-origin files are not restructured so future [Sync] Mirror Slime PRs stay clean. One commit per sub-PR; review each commit separately.

Commits:

• 2b405c2 Add pure-torch low-precision quantization reference oracle
• a644f74 Add MXFP8 trainer-side export quantizer and --extra-high-precision-layers flags
• a333ffb Add offline HF->MXFP8 checkpoint converter and metadata-only rollout schema
• 4e4d08c Add NVFP4 primitives, pure-torch QDQ fallback, and fake-QAT TE hook
• 49d62ef Add NVFP4 Megatron->HF export quantizer (modelopt_fp4 layout)
• ab377b7 Add offline HF->NVFP4 checkpoint converter
• e66e155 Add vLLM patch for in-loop NVFP4 (modelopt_fp4) weight updates

Verification: every quantizer is checked against the pure-torch oracle on CPU; tests/test_mxfp8_te_vs_oracle.py (real TE path) needs Blackwell. Patch applies in Dockerfile order; docker build pending. On rebase, drop this branch's copy of vime/utils/hf_utils/{weight_mapping,config}.py in favor of F04's.

Merge order of the miles-port series: F10 → F02 → F06 → F11a → F11b → F05 → F09 → F04 → F03 → F07 → F08 → F01. train.py / vime/utils/arguments.py are touched by several branches; later ones are rebased after each merge.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #35444838 | 📁 Comparing e66e155 against latest (d7e40fe)

  🔍 Preview build  

58 files changed · ± 58 modified

± Modified

• genindex.html
• library.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 48 more...

[gemini-code-assist]

Code Review

This pull request introduces comprehensive support for low-precision quantization formats (MXFP8 and NVFP4) optimized for Blackwell architectures, including offline conversion tools, pure-torch reference implementations, fake QAT fallbacks, and integration with the Megatron-to-HF conversion pipeline. It also adds a vLLM patch to support in-loop NVFP4 weight updates during training. The review feedback highlights a discrepancy in default environment variable parsing between the fake QAT config and reference utilities that could cause numerical divergence, and recommends using safe dictionary access on quantization_config to prevent potential KeyError exceptions.

[gemini-code-assist]

There is a discrepancy in the default value of NVTE_NVFP4_4OVER6_E4M3_USE_256 when it is not set in the environment. In vime/utils/nvfp4.py, it defaults to "all" (which resolves to 256 if 4over6 is enabled), whereas in nvfp4_qdq_torch.py, _env_flag defaults to "" (which resolves to 448). This inconsistency can cause numerical divergence between the training-time fake QAT and the offline/live quantizers.

We should update _env_flag to accept a default value so that NVTE_NVFP4_4OVER6_E4M3_USE_256 can default to "all" consistently.

Suggested change

	def _env_flag(name: str) -> bool:
	return os.getenv(name, "").strip().lower() in ("1", "true", "yes", "weights", "all")
	@lru_cache(maxsize=1)
	def current_nvfp4_qdq_config() -> NVFP4QDQConfig:
	"""QDQ config resolved once per process from NVTE_* env vars.
	Mirrors the INT4 path, which reads its env once at hook install time.
	Tests must call ``current_nvfp4_qdq_config.cache_clear()`` after changing
	the env (see reset_nvfp4_qdq_config_for_tests).
	"""
	if _env_flag("NVTE_USE_FAST_MATH"):
	raise RuntimeError("NVFP4 fake QAT requires NVTE_USE_FAST_MATH=0 (exact RN numerics).")
	use_4over6 = _env_flag("NVTE_NVFP4_4OVER6")
	e4m3_max = 448
	if use_4over6 and _env_flag("NVTE_NVFP4_4OVER6_E4M3_USE_256"):
	e4m3_max = 256
	def _env_flag(name: str, default: str = "") -> bool:
	return os.getenv(name, default).strip().lower() in ("1", "true", "yes", "weights", "all")
	@lru_cache(maxsize=1)
	def current_nvfp4_qdq_config() -> NVFP4QDQConfig:
	"""QDQ config resolved once per process from NVTE_* env vars.
	Mirrors the INT4 path, which reads its env once at hook install time.
	Tests must call ``current_nvfp4_qdq_config.cache_clear()`` after changing
	the env (see reset_nvfp4_qdq_config_for_tests).
	"""
	if _env_flag("NVTE_USE_FAST_MATH"):
	raise RuntimeError("NVFP4 fake QAT requires NVTE_USE_FAST_MATH=0 (exact RN numerics).")
	use_4over6 = _env_flag("NVTE_NVFP4_4OVER6")
	e4m3_max = 448
	if use_4over6 and _env_flag("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all"):
	e4m3_max = 256

[gemini-code-assist]

Accessing quantization_config["quant_method"] directly can raise a KeyError if the key is missing. It is safer to use .get("quant_method") to handle missing keys gracefully.

Suggested change

	if quantization_config["quant_method"] == "mxfp8":
	from .quantizer_mxfp8 import quantize_params_mxfp8
	return quantize_params_mxfp8(args, megatron_name, converted_named_params, quantization_config)
	# NVFP4 checkpoints: quant_algo "NVFP4" with quant_method "modelopt" /
	# "modelopt_fp4", or a bare "nvfp4" quant_method.
	if (
	quantization_config.get("quant_algo") == "NVFP4"
	or quantization_config["quant_method"] in ("nvfp4", "NVFP4", "modelopt_fp4")
	):
	from .quantizer_nvfp4 import quantize_params_nvfp4
	return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config)
	if quantization_config.get("quant_method") == "mxfp8":
	from .quantizer_mxfp8 import quantize_params_mxfp8
	return quantize_params_mxfp8(args, megatron_name, converted_named_params, quantization_config)
	# NVFP4 checkpoints: quant_algo "NVFP4" with quant_method "modelopt" /
	# "modelopt_fp4", or a bare "nvfp4" quant_method.
	if (
	quantization_config.get("quant_algo") == "NVFP4"
	or quantization_config.get("quant_method") in ("nvfp4", "NVFP4", "modelopt_fp4")
	):
	from .quantizer_nvfp4 import quantize_params_nvfp4
	return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config)