From 01a33c68cb5e4e61bd9a4e3628693c002ba1dff8 Mon Sep 17 00:00:00 2001 From: xmarre Date: Fri, 6 Mar 2026 06:26:28 +0100 Subject: [PATCH] Disable cudagraph for proj layers --- README.md | 26 +++++++++- torch_compile_lora_safe.py | 99 ++++++++++++++++++++++++++++++++------ 2 files changed, 109 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 9715c3c..5e7ca9e 100644 --- a/README.md +++ b/README.md @@ -63,18 +63,32 @@ Strict Flux names targeted by default: - `single_blocks.{i}.linear1` - `single_blocks.{i}.linear2` +For cudagraph-capable Flux runs (`backend = cudagraphs`, or `backend = inductor` +with `disable_cudagraphs = false`), the target set is narrowed to the +cudagraph-safe qkv subset only: + +- `double_blocks.{i}.img_attn.qkv` +- `double_blocks.{i}.txt_attn.qkv` + +In that mode, inferred Flux targets and full-model fallback are intentionally +disabled to avoid widening back into unstable replay paths. + Optional widening controls: - `allow_flux_inferred_targets` (`false` by default) - `false`: use only strict names above. - `true`: allow inferred Flux leaf targets (including `img_mlp/txt_mlp` pattern matches) for compatibility across variants. + - Ignored for cudagraph-capable Flux runs, which stay on the strict qkv-only + set. - `fallback_to_full_model_if_no_targets` (`false` by default) - `false`: if transformer-only target discovery finds nothing, skip compile instead of silently widening scope. - `true`: if discovery finds nothing, fall back to compiling `diffusion_model`. + - Ignored for cudagraph-capable Flux runs, which skip compile instead of + widening scope. For non-Flux transformer-style models, it compiles known block containers such as: @@ -128,23 +142,30 @@ If no known transformer compile targets are found, behavior is controlled by - Not applied to the `cudagraphs` backend. - Recommended if you hit `cudaMallocAsync` / cudagraph instability in Inductor-based runs. + - Set to `false` only when you specifically want Inductor cudagraphs. - `compile_transformer_only` (`BOOLEAN`) - `false` (default): compile the entire `diffusion_model`. - `true`: compile discovered transformer targets only. - Flux behavior (default strict mode): targets only `qkv/proj` attention leaves plus `linear1/linear2` leaves. + - For cudagraph-capable Flux runs, discovery is narrowed further to + `img_attn.qkv` and `txt_attn.qkv` only. - `allow_flux_inferred_targets` (`BOOLEAN`, default `false`) - `false` (recommended): strict Flux target set only. - `true`: allow inferred Flux leaf targets (`img_mlp/txt_mlp`-style matches) when preferred names differ. + - Ignored for cudagraph-capable Flux runs, which stay on the strict qkv-only + set. - `fallback_to_full_model_if_no_targets` (`BOOLEAN`, default `false`) - `false` (recommended): if transformer-only discovery finds nothing, skip compile rather than widening scope. - `true`: if transformer-only discovery finds nothing, compile `diffusion_model`. + - Ignored for cudagraph-capable Flux runs, which skip compile instead of + widening scope. --- @@ -153,7 +174,8 @@ If no known transformer compile targets are found, behavior is controlled by If you're unsure where to begin: - Backend: `inductor` -- `disable_cudagraphs`: `true` (recommended default for stability) +- `disable_cudagraphs`: `true` (recommended default for stability). Set to + `false` only when you specifically want Inductor cudagraphs. - Mode: `default` - `fullgraph`: `false` - `dynamic`: `false` @@ -176,6 +198,8 @@ Then iterate one setting at a time based on stability and speed. - If a backend fails, switch to `inductor` first. - If you encounter `cudaMallocAsync`/cudagraph issues, keep `disable_cudagraphs = true`. +- For cudagraph-enabled Flux validation, start with qkv-only targeting and do + not widen to `proj` unless logs prove the path is stable. - Guard filtering: when ComfyUI's `skip_torch_compile_dict` is available, this node passes it as `options['guard_filter_fn']` to reduce guard churn on diff --git a/torch_compile_lora_safe.py b/torch_compile_lora_safe.py index c8791ff..b5e997e 100644 --- a/torch_compile_lora_safe.py +++ b/torch_compile_lora_safe.py @@ -7,6 +7,11 @@ import torch.nn as nn _LOG = logging.getLogger(__name__) +_DEFAULT_FORCE_PARAMETER_STATIC_SHAPES = getattr( + torch._dynamo.config, "force_parameter_static_shapes", None +) +_DEFAULT_GRAPH_PARTITION = getattr(torch._inductor.config, "graph_partition", None) + try: from comfy_api.torch_helpers import set_torch_compile_wrapper as _set_torch_compile_wrapper except Exception: @@ -93,8 +98,10 @@ class TorchCompileModel_LoRASafe: "img_attn.proj", "txt_attn.qkv", "txt_attn.proj", - "img_mlp", - "txt_mlp", + ) + _FLUX_DOUBLE_CUDAGRAPH_SAFE_SUFFIXES = ( + "img_attn.qkv", + "txt_attn.qkv", ) _FLUX_SINGLE_PREFERRED_SUFFIXES = ("linear1", "linear2") @@ -191,6 +198,12 @@ class TorchCompileModel_LoRASafe: return False return any(True for _ in module.parameters(recurse=False)) + @staticmethod + def _is_flux_model(diffusion_model: nn.Module) -> bool: + return hasattr(diffusion_model, "double_blocks") and hasattr( + diffusion_model, "single_blocks" + ) + @classmethod def _try_preferred_suffixes(cls, block: nn.Module, suffixes): targets = [] @@ -208,9 +221,17 @@ class TorchCompileModel_LoRASafe: @classmethod def _discover_flux_double_block_suffixes( - cls, block: nn.Module, allow_inferred_targets: bool = False + cls, + block: nn.Module, + allow_inferred_targets: bool = False, + cudagraph_capable: bool = False, ): - preferred = cls._try_preferred_suffixes(block, cls._FLUX_DOUBLE_PREFERRED_SUFFIXES) + preferred_suffixes = ( + cls._FLUX_DOUBLE_CUDAGRAPH_SAFE_SUFFIXES + if cudagraph_capable + else cls._FLUX_DOUBLE_PREFERRED_SUFFIXES + ) + preferred = cls._try_preferred_suffixes(block, preferred_suffixes) if preferred or not allow_inferred_targets: return preferred @@ -231,8 +252,14 @@ class TorchCompileModel_LoRASafe: @classmethod def _discover_flux_single_block_suffixes( - cls, block: nn.Module, allow_inferred_targets: bool = False + cls, + block: nn.Module, + allow_inferred_targets: bool = False, + cudagraph_capable: bool = False, ): + if cudagraph_capable: + return [] + preferred = cls._try_preferred_suffixes(block, cls._FLUX_SINGLE_PREFERRED_SUFFIXES) if preferred or not allow_inferred_targets: return preferred @@ -247,25 +274,34 @@ class TorchCompileModel_LoRASafe: return list(dict.fromkeys(inferred)) @classmethod - def _discover_compile_keys(cls, diffusion_model, allow_flux_inferred_targets: bool = False): + def _discover_compile_keys( + cls, + diffusion_model, + allow_flux_inferred_targets: bool = False, + cudagraph_capable: bool = False, + ): keys = [] # Flux special-case: compile only leaf tensor-heavy modules, not full block forward. - if hasattr(diffusion_model, "double_blocks") and hasattr(diffusion_model, "single_blocks"): + if cls._is_flux_model(diffusion_model): double_blocks = getattr(diffusion_model, "double_blocks") single_blocks = getattr(diffusion_model, "single_blocks") if isinstance(double_blocks, (nn.ModuleList, list, tuple)): for i, block in enumerate(double_blocks): for suffix in cls._discover_flux_double_block_suffixes( - block, allow_inferred_targets=allow_flux_inferred_targets + block, + allow_inferred_targets=allow_flux_inferred_targets, + cudagraph_capable=cudagraph_capable, ): keys.append(f"diffusion_model.double_blocks.{i}.{suffix}") if isinstance(single_blocks, (nn.ModuleList, list, tuple)): for i, block in enumerate(single_blocks): for suffix in cls._discover_flux_single_block_suffixes( - block, allow_inferred_targets=allow_flux_inferred_targets + block, + allow_inferred_targets=allow_flux_inferred_targets, + cudagraph_capable=cudagraph_capable, ): keys.append(f"diffusion_model.single_blocks.{i}.{suffix}") @@ -345,6 +381,19 @@ class TorchCompileModel_LoRASafe: self._configure_torch_logs(debug_torch_logs) + cudagraph_capable = backend == "cudagraphs" or ( + backend == "inductor" and not disable_cudagraphs + ) + + if _DEFAULT_FORCE_PARAMETER_STATIC_SHAPES is not None: + torch._dynamo.config.force_parameter_static_shapes = ( + False if cudagraph_capable else _DEFAULT_FORCE_PARAMETER_STATIC_SHAPES + ) + if _DEFAULT_GRAPH_PARTITION is not None: + torch._inductor.config.graph_partition = ( + True if cudagraph_capable else _DEFAULT_GRAPH_PARTITION + ) + # This flag is process-global, so set it on every call to match the # current node input instead of leaking state from earlier runs. if backend in ("inductor", "cudagraphs"): @@ -363,25 +412,42 @@ class TorchCompileModel_LoRASafe: ) if compile_transformer_only: + is_flux_model = self._is_flux_model(dm) compile_keys = self._discover_compile_keys( - dm, allow_flux_inferred_targets=allow_flux_inferred_targets + dm, + allow_flux_inferred_targets=( + allow_flux_inferred_targets and not cudagraph_capable + ), + cudagraph_capable=cudagraph_capable, ) if not compile_keys: - if fallback_to_full_model_if_no_targets: + if fallback_to_full_model_if_no_targets and not ( + cudagraph_capable and is_flux_model + ): _LOG.warning( "TorchCompileModel_LoRASafe: no known transformer compile targets found; " "falling back to whole diffusion model compile." ) compile_keys = ["diffusion_model"] else: - _LOG.warning( - "TorchCompileModel_LoRASafe: no known transformer compile targets found; " - "skipping compile instead of falling back to whole model." - ) + if cudagraph_capable and is_flux_model: + _LOG.warning( + "TorchCompileModel_LoRASafe: no cudagraph-safe Flux transformer " + "compile targets found; skipping compile instead of widening target scope." + ) + else: + _LOG.warning( + "TorchCompileModel_LoRASafe: no known transformer compile targets found; " + "skipping compile instead of falling back to whole model." + ) return (m,) else: compile_keys = ["diffusion_model"] + if cudagraph_capable and compile_transformer_only and is_flux_model: + _LOG.warning("Flux cudagraph targets: %s", compile_keys[:16]) + _LOG.warning("Flux cudagraph target count: %d", len(compile_keys)) + if _set_torch_compile_wrapper is not None: try: _set_torch_compile_wrapper(model=m, keys=compile_keys, **compile_kw) @@ -407,3 +473,6 @@ class TorchCompileModel_LoRASafe: NODE_CLASS_MAPPINGS = { "TorchCompileModel_LoRASafe": TorchCompileModel_LoRASafe, } + + +