diff --git a/scripts/prototype/test_nvfp4_moe.py b/scripts/prototype/test_nvfp4_moe.py index 3b273a1a59..cc39ceef41 100644 --- a/scripts/prototype/test_nvfp4_moe.py +++ b/scripts/prototype/test_nvfp4_moe.py @@ -1,11 +1,16 @@ """Minimal OLMoE inference example with optional NVFP4 quantization for expert weights.""" +import gc +import inspect +import subprocess +import tempfile from contextlib import nullcontext import fire import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.models.olmoe.modeling_olmoe import OlmoeExperts +from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer from torchao.prototype.mx_formats.inference_workflow import ( NVFP4DynamicActivationNVFP4WeightConfig, @@ -14,7 +19,7 @@ from torchao.quantization import FqnToConfig, quantize_ -def main(recipe: str = "bf16"): +def main(recipe: str = "bf16", run_lm_eval: bool = False): print(f"{recipe=}") model_id = "allenai/OLMoE-1B-7B-0924" @@ -66,6 +71,49 @@ def main(recipe: str = "bf16"): print(tokenizer.decode(output[0], skip_special_tokens=True)) + if run_lm_eval: + if recipe == "nvfp4": + source = inspect.getsource(TorchAoHfQuantizer.get_weight_conversions) + if "gate_up_proj" not in source: + raise RuntimeError( + "Your version of `transformers` does not support NVFP4 MoE serialization. " + "Please install a version that includes " + "https://github.com/huggingface/transformers/pull/45609" + ) + + with tempfile.TemporaryDirectory() as output_dir: + print(f"\nSaving model to {output_dir}...") + + if recipe != "bf16": + from transformers import TorchAoConfig + + torchao_config = TorchAoConfig(quant_type=config) + model.config.quantization_config = torchao_config + model.hf_quantizer = TorchAoHfQuantizer(torchao_config) + + model.save_pretrained(output_dir, safe_serialization=False) + tokenizer.save_pretrained(output_dir) + + del model + gc.collect() + torch.cuda.empty_cache() + + lm_eval_cmd = [ + "lm_eval", + "--model", + "hf", + "--model_args", + f"pretrained={output_dir}", + "--tasks", + "wikitext", + "--num_fewshot", + "0", + "--batch_size", + "16", + ] + print(f"Running: {' '.join(lm_eval_cmd)}") + subprocess.run(lm_eval_cmd, check=True) + if __name__ == "__main__": fire.Fire(main)