From c5935ed19e617fd692390117e204f0161d6e3c79 Mon Sep 17 00:00:00 2001
From: IvanKobzarev <ivan.kobzarev@gmail.com>
Date: Mon, 11 May 2026 06:32:28 -0700
Subject: [PATCH] Skip _check_forward_args validation during make_fx tracing

During make_fx tracing, the forward args are FakeTensors which have
different device/repr than the meta tensors captured during AutoParallel
tracing. This causes _check_forward_args to raise spurious validation
errors. Detect this by checking for meta tensors or FakeTensors in
the args and skip validation in that case.

Authored by Claude.
---
 autoparallel/input_validation.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/autoparallel/input_validation.py b/autoparallel/input_validation.py
index 7990ef1c..b5d4d803 100644
--- a/autoparallel/input_validation.py
+++ b/autoparallel/input_validation.py
@@ -71,6 +71,15 @@ def _check_forward_args(args, expected_inputs, dynamic_dims=frozenset()):
         expected_inputs: Expected shapes from _compute_expected_inputs.
         dynamic_dims: Set of (arg_index, dim) pairs for dynamic dimensions.
     """
+    # Skip validation during make_fx tracing -- FakeTensors have different
+    # device/repr than the meta tensors used during AutoParallel tracing.
+    if any(
+        a.is_meta or getattr(a, "fake_mode", None) is not None
+        for a in args
+        if isinstance(a, torch.Tensor)
+    ):
+        return
+
     if len(args) != len(expected_inputs):
         raise ValueError(
             f"AutoParallel: expected {len(expected_inputs)} arguments "