Skip to content

v2.4.4 - v2.2.7 Duplicate File I/O Elimination — Loss and Restoration

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 16 May 22:07
· 274 commits to main since this release
v2.4.4
08db6e8
EN 中文

Overview

Documents the fact that the "first LoRA duplicate file read elimination" feature added in v2.2.7 (2026-01-07) was unintentionally lost in the v2.3.0 (2026-01-22) AWQ modulation layer LoRA fix commit, and subsequently restored.

1. Feature Added in v2.2.7

Commit: 1ff4db9
Release: v2.2.7 — LoRA Loading Speed Improvement: Duplicate File Read Elimination

Code Added

# OPTIMIZATION: Cache first LoRA state dict for reuse in processing loop
_cached_first_lora_state_dict = None
if lora_configs:
    first_lora_path_or_dict, first_lora_strength = lora_configs[0]
    first_lora_state_dict = _load_lora_state_dict(first_lora_path_or_dict)
    _cached_first_lora_state_dict = first_lora_state_dict  # Cache for reuse

# 1. Aggregate weights from all LoRAs
for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
    lora_name = lora_path_or_dict if isinstance(lora_path_or_dict, str) else "dict"
    # OPTIMIZATION: Reuse cached first LoRA state dict to avoid duplicate file I/O
    if idx == 0 and _cached_first_lora_state_dict is not None:
        lora_state_dict = _cached_first_lora_state_dict
    else:
        lora_state_dict = _load_lora_state_dict(lora_path_or_dict)

Impact and Significance

Why the first LoRA was read twice

compose_loras_v2() loaded the file in two separate stages:

  1. Debug logging stage (lines 1199-1212)

    first_lora_state_dict = _load_lora_state_dict(first_lora_path_or_dict)
for key in first_lora_state_dict.keys():
    parsed_res = _classify_and_map_key(key)
    logger.info(...)
# ← first_lora_state_dict is discarded as a local variable

  2. Actual processing loop stage (lines 1217-1223)

    for lora_path_or_dict, strength in lora_configs:
    lora_state_dict = _load_lora_state_dict(lora_path_or_dict)  # ← Same file re-read

The state_dict loaded for debug logging (hundreds of MB to several GB) went out of scope and became eligible for garbage collection. The actual processing loop then re-read the same file from disk, causing entirely wasted work.

Quantified waste

Operation Before fix After fix Reduction
File I/O 2 times 1 time 50%
Binary parse & tensor deserialization 2 times 1 time 50%
Key classification (_classify_and_map_key) 2×N times 1×N times 50%

N = number of keys in the LoRA file (typically hundreds to thousands)

Concrete performance impact

  • 500MB LoRA file: ~1-3 seconds saved (depending on disk speed)
  • 2GB LoRA file: ~5-10 seconds saved
  • User perception: Reduced wait time from prompt execution to generation start

Why this matters

  1. Single LoRA is the most common use case

    • Many workflows apply only one LoRA
    • In that case, the "first LoRA" is the only LoRA, so waste is 100%

  2. File I/O is a synchronous block

    • comfy.utils.load_torch_file() runs synchronously; downstream processing halts until it completes
    • The UI thread is blocked, making the ComfyUI console appear "frozen"

  3. Deserialization is CPU-intensive

    • Binary parsing in safetensors / torch.load places heavy CPU load
    • Especially noticeable with large LoRAs (FLUX, SDXL, Qwen-class)

  4. Key classification is a chain of regex matching

    • _classify_and_map_key() runs ~30 regex patterns against each key
    • Cost scales linearly with the number of keys

Cache safety

  • _cached_first_lora_state_dict is a local variable inside compose_loras_v2
  • Automatically discarded when the function exits; no state leaks across calls
  • state_dict is treated as read-only; debug logging does not mutate it
  • The idx == 0 condition ensures the second and subsequent LoRAs are unaffected

2. Preserved in v2.2.8

Commit: c113968
Release: v2.2.8 — Fix for Issue #44: Console Slowdown When Loading Unsupported LoRA Formats

Issue #44 (massive log output for unsupported LoRA formats) was fixed, but the _cached_first_lora_state_dict cache mechanism was preserved as-is.

git show c113968:nunchaku_code/lora_qwen.py | Select-String "cached_first_lora"
→ _cached_first_lora_state_dict = None
→ _cached_first_lora_state_dict = first_lora_state_dict
→ if idx == 0 and _cached_first_lora_state_dict is not None:

3. Lost in v2.3.0

Commit: 6498929
Release: v2.3.0 — Implement AWQ modulation layer LoRA fix with Runtime Monkey Patch
Change scope: nunchaku_code/lora_qwen.py | 865 ++++++++++++++--------------- | 266 insertions(+), 599 deletions(-)

Lost code (confirmed via git diff)

-        # OPTIMIZATION: Cache first LoRA state dict for reuse in processing loop
-        _cached_first_lora_state_dict = None
-        if lora_configs:
-            first_lora_path_or_dict, first_lora_strength = lora_configs[0]
-            first_lora_state_dict = _load_lora_state_dict(first_lora_path_or_dict)
-            _cached_first_lora_state_dict = first_lora_state_dict  # Cache for reuse

-        # 1. Aggregate weights from all LoRAs
-        for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
-            lora_name = lora_path_or_dict if isinstance(lora_path_or_dict, str) else "dict"
-            # OPTIMIZATION: Reuse cached first LoRA state dict to avoid duplicate file I/O
-            if idx == 0 and _cached_first_lora_state_dict is not None:
-                lora_state_dict = _cached_first_lora_state_dict
-            else:
-                lora_state_dict = _load_lora_state_dict(lora_path_or_dict)
+    # --- 4. Main Loading Loop ---
+    for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
+        if abs(strength) < 1e-5:
+            continue

Cause of loss

A migration omission during the sectional restructuring of compose_loras_v2().

v2.2.8 structure:

Cache init → Debug logging → Loop (idx == 0 cache reuse)

v2.3.0 structure:

# --- 1. Format Detection ---
# --- 2. Early Exit ---
# --- 3. Debug Inspection (First LoRA) ---
# --- 4. Main Loading Loop ---

In Section 3, first_lora_state_dict is loaded but remains a local variable inaccessible from Section 4. Section 4 calls _load_lora_state_dict_robust(lora_path_or_dict) unconditionally, so the first LoRA is read twice.

enumerate(lora_configs) remained, but all three elements of the cache variable (declaration, save, and reuse) disappeared.

4. Code at v2.4.3

# --- 3. Debug Inspection (First LoRA) ---
if lora_configs:
    first_lora_path_or_dict, first_lora_strength = lora_configs[0]
    first_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict)
    _first_detection = _detect_lora_format(first_lora_state_dict)

# --- 4. Main Loading Loop ---
for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
    ...
    lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict)  # Second read at idx==0

Duplicate file I/O had reappeared.

5. Compatibility with AWQ Monkey Patch

The AWQ modulation layer LoRA fix (_awq_lora_forward) operates inside _apply_lora_to_module. This is at the tail end of the compose_loras_v2 flow: file load → weight processing → module application.

File-load optimization (caching) sits at the entry point of this flow. The two are independent layers and can coexist:

_cached_first_lora_state_dict = None  # Added
if lora_configs:
    first_lora_path_or_dict, first_lora_strength = lora_configs[0]
    first_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict)
    _cached_first_lora_state_dict = first_lora_state_dict  # Added
    _first_detection = _detect_lora_format(first_lora_state_dict)

for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):
    ...
    if idx == 0 and _cached_first_lora_state_dict is not None:  # Added
        lora_state_dict = _cached_first_lora_state_dict
    else:
        lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict)

This change has zero impact on AWQ monkey-patch code (the should_apply_awq_mod check in Section 2, the AWQ branch inside _apply_lora_to_module). The cache mechanism and the AWQ patch are independent layers.

6. Facts of the Migration Omission

Changes actually made in commit 6498929

nunchaku_code/lora_qwen.py | 865 ++++++++++++++--------------- | 266 insertions(+), 599 deletions(-)
  • 599 lines deleted: A massive rewrite of the compose_loras_v2 function structure.
  • 266 lines added: Mostly AWQ monkey-patch code (_awq_lora_forward, DEFAULT_APPLY_AWQ_MOD_ENV, etc.).

How the structure changed

v2.2.8 and earlier (sequential processing):

def compose_loras_v2(...):
    # Initialize cache variable
    _cached_first_lora_state_dict = None
    
    # Load first LoRA for debug logging
    if lora_configs:
        first_lora = _load_lora_state_dict(...)
        _cached_first_lora_state_dict = first_lora  # Save to cache
    
    # Main loop (reuse cache)
    for idx, (lora_path, strength) in enumerate(lora_configs):
        if idx == 0 and _cached_first_lora_state_dict is not None:
            lora_state_dict = _cached_first_lora_state_dict  # Reuse
        else:
            lora_state_dict = _load_lora_state_dict(lora_path)

v2.3.0 and later (sectioned):

def compose_loras_v2(...):
    # --- 1. Z-Image / NextDiT Handling ---
    # --- 2. Environment Variable / Argument Logic for AWQ Mod ---
    
    # --- 3. Debug Inspection (First LoRA) ---
    if lora_configs:
        first_lora = _load_lora_state_dict_robust(...)
        # ← Loaded here, but not saved to a cache variable
    
    # --- 4. Main Loading Loop ---
    for idx, (lora_path, strength) in enumerate(lora_configs):
        lora_state_dict = _load_lora_state_dict_robust(lora_path)
        # ← Re-read at idx==0. No cache variable exists, so no branch is possible.

What disappeared (all three elements)

Element v2.2.8 v2.3.0
Cache variable declaration _cached_first_lora_state_dict = None Deleted
Save to cache _cached_first_lora_state_dict = first_lora_state_dict Deleted
Cache reuse if idx == 0 and _cached_first_lora_state_dict is not None: Deleted

What remained

for idx, (lora_path_or_dict, strength) in enumerate(lora_configs):

enumerate remained. It looks like a trace of the optimization, but because the idx == 0 branch logic is gone, enumerate functions only as a counter and is never used for cache reuse.

Facts from the commit message

  • Commit message: Implement AWQ modulation layer LoRA fix with Runtime Monkey Patch
  • The body contains no mention of "cache deletion," "optimization removal," or "file I/O changes."
  • The purpose was to add the AWQ monkey patch; there was no intent to discard the existing file I/O optimization.

Conclusion

When restructuring compose_loras_v2 into sections in commit 6498929, the developer kept enumerate(lora_configs) but forgot to port the three elements of _cached_first_lora_state_dict (declaration, save, and reuse) into the new structure. The commit's purpose was to add the AWQ monkey patch; removal of the file I/O optimization was not the intent.

7. Restoration Facts

File modified

nunchaku_code/lora_qwen.py

Code added

Section 3 (Debug Inspection):

    # --- 3. Debug Inspection (First LoRA) ---
    # OPTIMIZATION: Cache first LoRA state dict for reuse in processing loop
    _cached_first_lora_state_dict = None
    if lora_configs:
        first_lora_path_or_dict, first_lora_strength = lora_configs[0]
        first_lora_state_dict = _load_lora_state_dict_robust(first_lora_path_or_dict)
        _cached_first_lora_state_dict = first_lora_state_dict  # Cache for reuse

Section 4 (Main Loading Loop):

        # OPTIMIZATION: Reuse cached first LoRA state dict to avoid duplicate file I/O
        if idx == 0 and _cached_first_lora_state_dict is not None:
            lora_state_dict = _cached_first_lora_state_dict
        else:
            lora_state_dict = _load_lora_state_dict_robust(lora_path_or_dict)

Significance

Restored the v2.2.7 "first LoRA duplicate read elimination" feature in full compatibility with the AWQ modulation layer monkey patch (v2.3.0).

  • Section 3 saves the loaded first_lora_state_dict into _cached_first_lora_state_dict
  • Section 4 reuses that cache at idx == 0, avoiding a second _load_lora_state_dict_robust call
  • Achieves a 50% reduction in file I/O, binary parsing, and key classification

This fix has zero impact on the AWQ monkey patch (_awq_lora_forward, the AWQ branch inside _apply_lora_to_module). The cache mechanism lives in the file-loading layer; the AWQ patch lives in the module-application layer. They are completely independent.

Assets 2
Loading