Repository: pytorch
Issue: [torch.export] Error When Trying To Express Dynamism For Transformer Model of SD3
Labels: oncall: pt2,oncall: export

### 🐛 Describe the bug

**Brief Description:**
I am trying to export the transformer model of Stable Diffusion 3 using `torch.export.export_for_training`. The error occurs when trying to express dynamism for the feature map height and width. The reproducer and traceback are given below.

**Reproducible Code:**
```
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", text_encoder_3=None, tokenizer_3=None)

unet_kwargs = {}
unet_kwargs["hidden_states"] = torch.ones((2, 16, 64, 64))
unet_kwargs["timestep"] = torch.from_numpy(np.array([1, 2], dtype=np.float32))
unet_kwargs["encoder_hidden_states"] = torch.ones((2, 154, 4096))
unet_kwargs["pooled_projections"] = torch.ones((2, 2048))

#Feature map height and width are dynamic
fm_height = torch.export.Dim('fm_height', min=16)
fm_width = torch.export.Dim('fm_width', min=16)

#iterate through the unet kwargs and set only hidden state kwarg to dynamic
dynamic_shapes = {key: (None if key != "hidden_states" else {2: fm_height, 3: fm_width}) for key in unet_kwargs.keys()}
transformer = torch.export.export_for_training(pipe.transformer.eval(), args=(), kwargs=(unet_kwargs), dynamic_shapes=dynamic_shapes).module()

```

**Error Traceback:**
```
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0] Error while creating guard:
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0] Name: ''
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Source: shape_env
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Create Function: SHAPE_ENV
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Guard Types: None
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Code List: None
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Object Weakref: None
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     Guarded Class Weakref: None
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0] Traceback (most recent call last):
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_guards.py", line 281, in create
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     return self.create_fn(builder, self)
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 1836, in SHAPE_ENV
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     guards = output_graph.shape_env.produce_guards(
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 4178, in produce_guards
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]     raise ConstraintViolationError(
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0] torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (fm_height, fm_width)! For more information, run with TORCH_LOGS="+dynamic".
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_height = L['hidden_states'].size()[2] in the specified range 16 <= fm_height <= 9223372036854775806 satisfy the generated guard Ne(((-(L['hidden_states'].size()[2]//2))//2) + 96, 0).
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_height = L['hidden_states'].size()[2] in the specified range 16 <= fm_height <= 9223372036854775806 satisfy the generated guard (L['hidden_states'].size()[2]//2) + ((-(L['hidden_states'].size()[2]//2))//2) + 96 <= 192.
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_width = L['hidden_states'].size()[3] in the specified range 16 <= fm_width <= 9223372036854775806 satisfy the generated guard Ne(((-(L['hidden_states'].size()[3]//2))//2) + 96, 0).
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_width = L['hidden_states'].size()[3] in the specified range 16 <= fm_width <= 9223372036854775806 satisfy the generated guard (L['hidden_states'].size()[3]//2) + ((-(L['hidden_states'].size()[3]//2))//2) + 96 <= 192.
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_width = L['hidden_states'].size()[3] in the specified range 16 <= fm_width <= 9223372036854775806 satisfy the generated guard Ne(294912, 1536*((L['hidden_states'].size()[3]//2))).
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_height = L['hidden_states'].size()[2] in the specified range 16 <= fm_height <= 9223372036854775806 satisfy the generated guard Ne(Mod((L['hidden_states'].size()[2]//2), 2*((L['hidden_states'].size()[2]//2))), 0).
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_width = L['hidden_states'].size()[3] in the specified range 16 <= fm_width <= 9223372036854775806 satisfy the generated guard Ne(Mod((L['hidden_states'].size()[3]//2), 2*((L['hidden_states'].size()[3]//2))), 0).
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_height = L['hidden_states'].size()[2] in the specified range 16 <= fm_height <= 9223372036854775806 satisfy the generated guard 16 <= L['hidden_states'].size()[2] and L['hidden_states'].size()[2] <= 385
E0115 12:29:57.483000 144841 torch/_guards.py:283] [8/0]   - Not all values of fm_width = L['hidden_states'].size()[3] in the specified range 16 <= fm_width <= 9223372036854775806 satisfy the generated guard 16 <= L['hidden_states'].size()[3] and L['hidden_states'].size()[3] <= 385
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0] Created at:
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 615, in transform
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]     tracer = InstructionTranslator(
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2670, in __init__
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]     output=OutputGraph(
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 317, in __init__
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]     self.init_ambient_guards()
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]   File "/home/user/Downloads/ov_notebooks_sd3/openvino_notebooks/.venv/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 463, in init_ambient_guards
E0115 12:29:57.485000 144841 torch/_guards.py:285] [8/0]     self.guards.add(ShapeEnvSource().make_guard(GuardBuilder.SHAPE_ENV))


torch/export/__init__.py:154, in export_for_training(mod, args, kwargs, dynamic_shapes, strict, preserve_module_call_signature)
    148 if isinstance(mod, torch.jit.ScriptModule):
    149     raise ValueError(
    150         "Exporting a ScriptModule is not supported. "
    151         "Maybe try converting your ScriptModule to an ExportedProgram "
    152         "using `TS2EPConverter(mod, args, kwargs).convert()` instead."
    153     )
--> 154 return _export_for_training(
    155     mod,
    156     args,
    157     kwargs,
    158     dynamic_shapes,
    159     strict=strict,
    160     preserve_module_call_signature=preserve_module_call_signature,
    161 )

torch/export/_trace.py:1017, in _log_export_wrapper.<locals>.wrapper(*args, **kwargs)
   1010     else:
   1011         log_export_usage(
   1012             event="export.error.unclassified",
   1013             type=error_type,
   1014             message=str(e),
   1015             flags=_EXPORT_FLAGS,
   1016         )
-> 1017     raise e
   1018 finally:
   1019     _EXPORT_FLAGS = None

torch/export/_trace.py:990, in _log_export_wrapper.<locals>.wrapper(*args, **kwargs)
    988 try:
    989     start = time.time()
--> 990     ep = fn(*args, **kwargs)
    991     end = time.time()
    992     log_export_usage(
    993         event="export.time",
    994         metrics=end - start,
    995         flags=_EXPORT_FLAGS,
    996         **get_ep_stats(ep),
    997     )

torch/export/exported_program.py:114, in _disable_prexisiting_fake_mode.<locals>.wrapper(*args, **kwargs)
    111 @functools.wraps(fn)
    112 def wrapper(*args, **kwargs):
    113     with unset_fake_temporarily():
--> 114         return fn(*args, **kwargs)

torch/export/_trace.py:1746, in _export_for_training(mod, args, kwargs, dynamic_shapes, strict, preserve_module_call_signature)
   1727 (
   1728     args,
   1729     kwargs,
   (...)
   1732     dynamic_shapes,
   1733 ) = _process_export_inputs(mod, args, kwargs, dynamic_shapes)
   1735 export_func = (
   1736     functools.partial(
   1737         _strict_export_lower_to_aten_ir,
   (...)
   1744     )
   1745 )
-> 1746 export_artifact = export_func(  # type: ignore[operator]
   1747     mod=mod,
   1748     args=args,
   1749     kwargs=kwargs,
   1750     dynamic_shapes=dynamic_shapes,
   1751     preserve_module_call_signature=preserve_module_call_signature,
   1752     pre_dispatch=False,
   1753     original_state_dict=original_state_dict,
   1754     orig_in_spec=orig_in_spec,
   1755     allow_complex_guards_as_runtime_asserts=False,
   1756     _is_torch_jit_trace=False,
   1757 )
   1759 export_graph_signature = export_artifact.aten.sig
   1761 forward_arg_names = _get_forward_arg_names(mod, args, kwargs)

torch/export/_trace.py:1252, in _strict_export_lower_to_aten_ir(mod, args, kwargs, dynamic_shapes, preserve_module_call_signature, pre_dispatch, original_state_dict, orig_in_spec, allow_complex_guards_as_runtime_asserts, _is_torch_jit_trace, lower_to_aten_callback)
   1239 def _strict_export_lower_to_aten_ir(
   1240     mod: torch.nn.Module,
   1241     args: Tuple[Any, ...],
   (...)
   1250     lower_to_aten_callback: Callable,
   1251 ) -> ExportArtifact:
-> 1252     gm_torch_level = _export_to_torch_ir(
   1253         mod,
   1254         args,
   1255         kwargs,
   1256         dynamic_shapes,
   1257         preserve_module_call_signature=preserve_module_call_signature,
   1258         restore_fqn=False,  # don't need to restore because we will do it later
   1259         allow_complex_guards_as_runtime_asserts=allow_complex_guards_as_runtime_asserts,
   1260         _log_export_usage=False,
   1261     )
   1263     # We detect the fake_mode by looking at gm_torch_level's placeholders, this is the fake_mode created in dynamo.
   1264     (
   1265         fake_args,
   1266         fake_kwargs,
   1267         dynamo_fake_mode,
   1268     ) = _extract_fake_inputs(gm_torch_level, args, kwargs)

torch/export/_trace.py:560, in _export_to_torch_ir(f, args, kwargs, dynamic_shapes, preserve_module_call_signature, disable_constraint_solver, allow_complex_guards_as_runtime_asserts, restore_fqn, _log_export_usage, same_signature)
    556     module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]] = {}
    557     with _wrap_submodules(
    558         f, preserve_module_call_signature, module_call_specs
    559     ), _ignore_backend_decomps():
--> 560         gm_torch_level, _ = torch._dynamo.export(
    561             f,
    562             dynamic_shapes=transformed_dynamic_shapes,  # type: ignore[arg-type]
    563             tracing_mode="symbolic",
    564             disable_constraint_solver=disable_constraint_solver,
    565             # currently the following 2 flags are tied together for export purposes,
    566             # but untangle for sake of dynamo export api
    567             prefer_deferred_runtime_asserts_over_guards=True,
    568             allow_complex_guards_as_runtime_asserts=allow_complex_guards_as_runtime_asserts,
    569             _log_export_usage=_log_export_usage,
    570             same_signature=same_signature,
    571         )(
    572             *args,
    573             **kwargs,
    574         )
    575 except (ConstraintViolationError, ValueRangeError) as e:
    576     raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904

torch/_dynamo/eval_frame.py:1448, in export.<locals>.inner(*args, **kwargs)
   1446 dim_constraints.solve()
   1447 forced_specializations = dim_constraints.forced_specializations()
-> 1448 msg = dim_constraints.prettify_results(
   1449     original_signature,
   1450     dynamic_shapes,
   1451     constraint_violation_error,
   1452     forced_specializations,
   1453 )
   1454 if constraint_violation_error:
   1455     constraint_violation_error.args = (
   1456         constraint_violation_error.args[0] + msg,
   1457     )

torch/fx/experimental/symbolic_shapes.py:2248, in DimConstraints.prettify_results(self, original_signature, dynamic_shapes, constraint_violation_error, forced_specializations)
   2245     for s, val in forced_specializations.items():
   2246         buf += f"  - solving the guards generated for {s} resulted in a specialized value of {val}.\n"
-> 2248 self._process_derived_dim_roots(results, name_to_dim)
   2250 dims = []
   2251 others = []

torch/fx/experimental/symbolic_shapes.py:2064, in DimConstraints._process_derived_dim_roots(self, results, name_to_dim)
   2062 # create result & dim
   2063 results[str(root)] = {"min": min_, "max": max_}
-> 2064 name_to_dim[str(root)] = Dim(str(root), min=min_, max=max_)
   2065 # remove old root min/max bounds
   2066 c.pop("min", None)

torch/export/dynamic_shapes.py:227, in Dim(name, min, max)
    225 _min = 0 if min is None else min
    226 _max = int_oo if max is None else max
--> 227 assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
    228 assert name.isidentifier(), f"Dim name must be a valid identifier, got {name}"
    229 dim = _Dim(name, (int,), {"min": _min, "max": _max})

AssertionError: Cannot create Dim with inconsistent min=-4, max=-96
```

### Versions

```
Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.26.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov  6 2024, 20:22:13) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.8.0-51-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 560.28.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               36
On-line CPU(s) list:                  0-35
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   2
Core(s) per socket:                   18
Socket(s):                            1
Stepping:                             7
CPU max MHz:                          4800.0000
CPU min MHz:                          1200.0000
BogoMIPS:                             6000.00
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts vnmi avx512_vnni md_clear flush_l1d arch_capabilities
Virtualization:                       VT-x
L1d cache:                            576 KiB (18 instances)
L1i cache:                            576 KiB (18 instances)
L2 cache:                             18 MiB (18 instances)
L3 cache:                             24.8 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-35
Vulnerability Gather data sampling:   Mitigation; Microcode
Vulnerability Itlb multihit:          KVM: Mitigation: VMX disabled
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy[/](https://file+.vscode-resource.vscode-cdn.net/)swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] flake8==7.1.1
[pip3] mypy==1.12.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] onnx==1.17.0
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.1
[pip3] torchvision==0.20.1+cpu
[pip3] triton==3.1.0
[conda] Could not collect
```

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4