Fix (brevitas_examples/llm): more checks for FX-related args (#1441)

---------

Co-authored-by: Pablo Monteagudo Lago <44771380+pablomlago@users.noreply.github.com>

Author: Giuseppe5

src/brevitas_examples/llm/llm_args.py

@@ -485,6 +485,10 @@ def create_args_parser() -> ArgumentParser:
     return parser
 
 
+def fx_required(args: Namespace):
+    return args.weight_equalization or args.act_equalization == 'fx' or args.rotation == 'fx' or args.ln_affine_merge or args.convert_layernorm_to_rmsnorm or args.quant_sdpa == 'fx'
+
+
 def validate(args: Namespace, extra_args: Optional[List[str]] = None) -> None:
     if args.optimize_rotations:
         assert args.rotation in ['fx', 'fused_no_fx'], f"Rotations can only be optimized if --rotation=fx or --rotation=fused_no_fx"
@@ -494,6 +498,12 @@ def validate(args: Namespace, extra_args: Optional[List[str]] = None) -> None:
         assert args.ln_affine_merge, 'Graph rotation requires to merge LN/RMS norm affine parameters'
         assert args.replace_rmsnorm, 'Graph rotation requires to replace HF RMSNorm with PyTorch ones (torch 2.4+ require)'
         assert args.convert_layernorm_to_rmsnorm, 'Graph rotation requires to replace LayerNorm with RMSNorm'
+        # FX is not compatible with few-shot evaluation
+        assert args.few_shot_eval is None, "FX is not compatible with few shot evaluation, use fused_no_fx"
+    # Otherwise we might end up tracing through dynamo twice and other weird errors.
+    # Fused_no_fx takes care of all the rotations-related transformations
+    if fx_required(args):
+        assert args.rotation != "fused_no_fx", "fused_no_fx is incompatible with any option that requires FX tracing"
     elif args.rotation == 'fused_no_fx':
         assert not args.convert_layernorm_to_rmsnorm, 'LayerNorm is automatically replaced with RMSNorm when running with --rotation=fused_no_fx. Remove the flag --convert-layernorm-to-rmsnorm'
         assert args.replace_rmsnorm, 'Graph rotation requires to replace HF RMSNorm with PyTorch ones (torch 2.4+ require)'

src/brevitas_examples/llm/main.py

@@ -38,6 +38,7 @@
 from brevitas_examples.common.parse_utils import parse_args
 from brevitas_examples.llm.gguf_export.export import save_quantized_as_gguf
 from brevitas_examples.llm.llm_args import create_args_parser
+from brevitas_examples.llm.llm_args import fx_required
 from brevitas_examples.llm.llm_args import validate
 from brevitas_examples.llm.llm_quant.awq.pre_quant import apply_awq
 from brevitas_examples.llm.llm_quant.bias_corr import apply_bias_correction
@@ -201,10 +202,6 @@ def model_export(model, tokenizer, ref_input, args, config=None):
         ds.save(export_path, io_report_callback=None)
 
 
-def fx_required(args):
-    return args.weight_equalization or args.act_equalization == 'fx' or args.rotation == 'fx' or args.ln_affine_merge or args.convert_layernorm_to_rmsnorm or args.quant_sdpa == 'fx'
-
-
 # Recursive function to unwrap equalized layers
 def find_equalized_layer(layer):
     if hasattr(layer, 'layer'):