hook up hf + moe + nvfp4 script to lm_eval

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)