⚡️ Speed up function zero_module by 75%

[codeflash-ai]

📄 75% (0.75x) speedup for zero_module in invokeai/backend/flux/controlnet/zero_module.py

⏱️ Runtime : 2.22 milliseconds → 1.26 milliseconds (best of 203 runs)

📝 Explanation and details

The optimization achieves a 75% speedup by replacing torch.nn.init.zeros_(p) with p.zero_() and wrapping the operation in torch.no_grad().

Key optimizations:

• Direct tensor method vs. init function: p.zero_() is a direct tensor operation that zeros the parameter in-place, while torch.nn.init.zeros_(p) goes through PyTorch's initialization framework with additional function call overhead
• Gradient context optimization: torch.no_grad() prevents PyTorch from tracking operations for autograd, reducing memory overhead and computation when zeroing parameters
• Parameter list caching: Pre-collecting parameters with list(module.parameters()) avoids repeated generator calls within the loop

Performance impact by test case:

• Large modules with many parameters see the biggest gains (104-226% faster for sequential models with many layers)
• Standard neural network layers (Linear, Conv2d) show consistent 15-22% improvements
• Modules with no parameters show slight regression (~7-12% slower) due to the added list creation overhead, but this is negligible in practice

Hot path benefits: Based on the function reference, zero_module is called during ControlNet initialization to create zero-initialized linear layers for controlnet_blocks and controlnet_single_blocks. Since ControlNet models can have dozens of these blocks (matching the depth of the base FLUX model), this optimization significantly reduces model initialization time - a critical performance factor for ML inference pipelines where models may be loaded/reloaded frequently.

The optimization is most effective for modules with multiple parameters, making it ideal for the neural network layers typically used in ControlNet architectures.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 33 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from typing import TypeVar

# imports
import pytest  # used for our unit tests
# function to test
import torch
import torch.nn as nn
from invokeai.backend.flux.controlnet.zero_module import zero_module

T = TypeVar("T", bound=torch.nn.Module)
from invokeai.backend.flux.controlnet.zero_module import zero_module

# unit tests

# ----------- BASIC TEST CASES -----------

def test_zero_linear_layer_weights_and_bias():
    # Test that a simple Linear layer's weights and bias are zeroed
    layer = nn.Linear(4, 2)
    zero_module(layer) # 14.6μs -> 12.4μs (18.1% faster)

def test_zero_conv2d_layer_weights_and_bias():
    # Test that a Conv2d layer's weights and bias are zeroed
    layer = nn.Conv2d(3, 6, 5)
    zero_module(layer) # 14.0μs -> 11.7μs (19.6% faster)

def test_zero_module_returns_same_object():
    # Test that zero_module returns the same module object, not a copy
    layer = nn.Linear(2, 2)
    codeflash_output = zero_module(layer); returned = codeflash_output # 13.3μs -> 11.5μs (16.2% faster)

def test_zero_sequential_module():
    # Test that all parameters in a Sequential module are zeroed
    model = nn.Sequential(nn.Linear(3, 3), nn.ReLU(), nn.Linear(3, 1))
    zero_module(model) # 23.0μs -> 16.3μs (40.9% faster)
    for param in model.parameters():
        pass

# ----------- EDGE TEST CASES -----------

def test_zero_module_on_module_with_no_parameters():
    # Test that zero_module works on modules with no parameters (should not fail)
    class DummyModule(nn.Module):
        def __init__(self):
            super().__init__()
    dummy = DummyModule()
    zero_module(dummy) # 4.42μs -> 4.80μs (7.84% slower)
    # There are no parameters, so nothing to check

def test_zero_module_on_shared_parameters():
    # Test that shared parameters are zeroed only once
    shared = nn.Parameter(torch.randn(5, 5))
    class SharedModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.p1 = shared
            self.p2 = shared
    mod = SharedModule()
    zero_module(mod) # 16.9μs -> 16.8μs (0.894% faster)

def test_zero_module_on_non_contiguous_parameters():
    # Test that non-contiguous parameters are zeroed
    layer = nn.Linear(10, 10)
    # Make weight non-contiguous by transposing
    layer.weight.data = layer.weight.data.t()
    zero_module(layer) # 14.3μs -> 11.8μs (21.4% faster)

def test_zero_module_on_parameter_with_requires_grad_false():
    # Test that parameters with requires_grad=False are zeroed
    layer = nn.Linear(3, 3)
    layer.weight.requires_grad = False
    zero_module(layer) # 13.1μs -> 11.7μs (11.6% faster)

def test_zero_module_on_module_with_buffers():
    # Test that buffers are NOT zeroed (only parameters)
    class BufferModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.param = nn.Parameter(torch.randn(2, 2))
            self.register_buffer('buf', torch.ones(2, 2))
    mod = BufferModule()
    zero_module(mod) # 15.5μs -> 15.7μs (1.54% slower)

def test_zero_module_on_parameter_with_nan_inf():
    # Test that parameters with NaN or Inf values are zeroed
    layer = nn.Linear(2, 2)
    layer.weight.data.fill_(float('nan'))
    layer.bias.data.fill_(float('inf'))
    zero_module(layer) # 13.6μs -> 11.6μs (17.2% faster)

def test_zero_module_on_empty_parameter():
    # Test that an empty parameter tensor is handled
    class EmptyParamModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.empty = nn.Parameter(torch.empty(0))
    mod = EmptyParamModule()
    zero_module(mod) # 17.9μs -> 17.2μs (3.94% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_zero_module_on_large_linear_layer():
    # Test zero_module on a large Linear layer
    # 1000 x 1000 = 1,000,000 floats = ~4MB
    layer = nn.Linear(1000, 1000)
    zero_module(layer) # 170μs -> 164μs (3.96% faster)

def test_zero_module_on_large_sequential_model():
    # Test zero_module on a large Sequential model with many layers
    layers = [nn.Linear(100, 100) for _ in range(10)]
    model = nn.Sequential(*layers)
    zero_module(model) # 81.8μs -> 40.2μs (104% faster)
    for param in model.parameters():
        pass

def test_zero_module_on_large_conv2d_layer():
    # Test zero_module on a large Conv2d layer
    layer = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
    zero_module(layer) # 15.5μs -> 13.2μs (17.1% faster)

def test_zero_module_on_module_with_many_parameters():
    # Test zero_module on a module with many parameters
    class ManyParamsModule(nn.Module):
        def __init__(self):
            super().__init__()
            for i in range(50):
                setattr(self, f'param_{i}', nn.Parameter(torch.randn(10, 10)))
    mod = ManyParamsModule()
    zero_module(mod) # 145μs -> 45.3μs (221% faster)
    for name, param in mod.named_parameters():
        pass

def test_zero_module_on_nested_modules():
    # Test zero_module on deeply nested modules
    class NestedModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.seq = nn.Sequential(
                nn.Linear(10, 10),
                nn.Sequential(
                    nn.Linear(10, 10),
                    nn.Sequential(
                        nn.Linear(10, 10)
                    )
                )
            )
    mod = NestedModule()
    zero_module(mod) # 34.1μs -> 22.2μs (53.4% faster)
    for param in mod.parameters():
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import TypeVar

# imports
import pytest  # used for our unit tests
# function to test
# (from invokeai/backend/flux/controlnet/zero_module.py)
import torch
import torch.nn as nn
from invokeai.backend.flux.controlnet.zero_module import zero_module

T = TypeVar("T", bound=torch.nn.Module)
from invokeai.backend.flux.controlnet.zero_module import zero_module

# unit tests

# ----------------------- Basic Test Cases -----------------------

def test_zero_module_linear_basic():
    # Test zeroing a simple Linear layer
    linear = nn.Linear(3, 2)
    zero_module(linear) # 16.5μs -> 13.7μs (20.6% faster)
    # All parameters (weight and bias) should be zero
    for p in linear.parameters():
        pass

def test_zero_module_conv2d_basic():
    # Test zeroing a simple Conv2d layer
    conv = nn.Conv2d(1, 2, kernel_size=3)
    zero_module(conv) # 14.3μs -> 12.1μs (18.1% faster)
    for p in conv.parameters():
        pass

def test_zero_module_sequential_basic():
    # Test zeroing a Sequential module with multiple layers
    model = nn.Sequential(
        nn.Linear(4, 4),
        nn.ReLU(),
        nn.Linear(4, 2)
    )
    zero_module(model) # 23.9μs -> 16.4μs (45.6% faster)
    # Only Linear layers have parameters
    for module in model:
        if hasattr(module, 'weight'):
            pass
        if hasattr(module, 'bias') and module.bias is not None:
            pass

def test_zero_module_module_with_no_parameters():
    # Test zeroing a module with no parameters (e.g., nn.ReLU)
    relu = nn.ReLU()
    zero_module(relu) # 4.11μs -> 4.70μs (12.5% slower)
    # No parameters to check

def test_zero_module_multiple_calls():
    # Test that calling zero_module multiple times is idempotent
    linear = nn.Linear(5, 5)
    zero_module(linear) # 13.9μs -> 11.6μs (19.3% faster)
    zero_module(linear) # 7.93μs -> 5.92μs (34.0% faster)
    for p in linear.parameters():
        pass

# ----------------------- Edge Test Cases -----------------------

def test_zero_module_shared_parameters():
    # Test zeroing a module with shared parameters
    linear = nn.Linear(3, 3)
    # Share weight with another Linear
    linear2 = nn.Linear(3, 3)
    linear2.weight = linear.weight
    zero_module(linear) # 13.4μs -> 11.0μs (22.4% faster)

def test_zero_module_parameter_requires_grad_false():
    # Test zeroing a module with parameters that do not require grad
    linear = nn.Linear(2, 2)
    for p in linear.parameters():
        p.requires_grad = False
    zero_module(linear) # 10.8μs -> 8.83μs (22.2% faster)
    for p in linear.parameters():
        pass

def test_zero_module_custom_module_with_buffers():
    # Test zeroing a custom module with buffers (buffers should NOT be zeroed)
    class CustomModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(2, 2)
            self.register_buffer('buf', torch.ones(2, 2))
    mod = CustomModule()
    zero_module(mod) # 15.5μs -> 13.5μs (15.1% faster)
    for p in mod.parameters():
        pass

def test_zero_module_empty_module():
    # Test zeroing an empty nn.Module (no parameters, no buffers)
    class EmptyModule(nn.Module):
        pass
    mod = EmptyModule()
    zero_module(mod) # 4.36μs -> 4.69μs (7.06% slower)

def test_zero_module_parameter_with_nan_inf():
    # Test zeroing a module whose parameters contain NaN or Inf
    linear = nn.Linear(2, 2)
    with torch.no_grad():
        linear.weight.fill_(float('nan'))
        linear.bias.fill_(float('inf'))
    zero_module(linear) # 13.8μs -> 11.5μs (20.1% faster)
    for p in linear.parameters():
        pass

def test_zero_module_parameter_with_zero_shape():
    # Test zeroing a module with a parameter of zero shape
    class ZeroParamModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = nn.Parameter(torch.empty(0))
    mod = ZeroParamModule()
    zero_module(mod) # 17.0μs -> 17.6μs (3.23% slower)

# ----------------------- Large Scale Test Cases -----------------------

def test_zero_module_large_linear():
    # Test zeroing a large Linear layer (within 100MB limit)
    in_features = 1000
    out_features = 1000
    linear = nn.Linear(in_features, out_features)
    zero_module(linear) # 169μs -> 163μs (3.68% faster)
    for p in linear.parameters():
        pass

def test_zero_module_large_sequential():
    # Test zeroing a Sequential module with many layers
    layers = [nn.Linear(10, 10) for _ in range(100)]
    model = nn.Sequential(*layers)
    zero_module(model) # 623μs -> 191μs (226% faster)
    for layer in layers:
        for p in layer.parameters():
            pass

def test_zero_module_large_conv2d():
    # Test zeroing a large Conv2d layer (within 100MB)
    conv = nn.Conv2d(32, 32, kernel_size=16)
    zero_module(conv) # 39.4μs -> 36.6μs (7.56% faster)
    for p in conv.parameters():
        pass

def test_zero_module_large_custom_module():
    # Test zeroing a custom module with many submodules
    class BigModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.linears = nn.ModuleList([nn.Linear(20, 20) for _ in range(50)])
            self.convs = nn.ModuleList([nn.Conv2d(3, 3, 3) for _ in range(20)])
        def forward(self, x):
            pass
    mod = BigModule()
    zero_module(mod) # 452μs -> 147μs (207% faster)
    for linear in mod.linears:
        for p in linear.parameters():
            pass
    for conv in mod.convs:
        for p in conv.parameters():
            pass

def test_zero_module_large_shared_parameter():
    # Test zeroing a module with a large shared parameter
    shared_weight = nn.Parameter(torch.ones(1000, 1000))
    class SharedModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.weight1 = shared_weight
            self.weight2 = shared_weight
    mod = SharedModule()
    zero_module(mod) # 169μs -> 167μs (0.813% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-zero_module-mhx2mq0d and push.