AnimateDiffPipeline performance regression on CPU.

[jiqing-feng]

Describe the bug

regression PR: #11098 .

The attn_processor will turn the hidden_states into channel last (NHWC) layout, which will have great benefit on CPU when running matmul.

But the op: .contiguous() will turn the tensor into NCHW layout as the tensor orginal layout, it will make the matmul slower than the NHWC layout (converted by attn processor). Besides, the op .contiguous() also costs too much time on CPU if the tensor layout is bad (just like in this case).

Reproduction

numactl -C 0-31 --membind 0 python test.py

from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
from transformers import set_seed
import torch
import time

SEED = 42
device = "cpu"
model_dtype = torch.float16
WARM_UP = 4
RUN = 4

set_seed(SEED)

print("\nLoading AnimateDiff-Lightning model...")
model_id = "ByteDance/AnimateDiff-Lightning"
step = 4
ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
base = "emilianJR/epiCRealism"

adapter = MotionAdapter().to(device, model_dtype)
adapter.load_state_dict(load_file(hf_hub_download(model_id, ckpt), device=device))
pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=model_dtype).to(device)
pipe.scheduler = EulerDiscreteScheduler.from_config(
 pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear"
)

def run_inference():
 set_seed(SEED)
 with torch.no_grad():
 output = pipe(
 prompt="An astronaut riding a green horse",
 guidance_scale=1.0,
 num_inference_steps=4,
 ).frames[0]
 return output

# Warm up
print(f"\nWarming up ({WARM_UP} iterations)...")
for i in range(WARM_UP):
 run_inference()
 print(f" Warm-up {i+1}/{WARM_UP} done")

# Benchmark
print(f"\nRunning benchmark ({RUN} iterations)...")
elapsed_times = []
for i in range(RUN):
 start = time.perf_counter()
 output = run_inference()
 end = time.perf_counter()
 elapsed = (end - start) * 1000 # ms
 elapsed_times.append(elapsed)
 print(f" Run {i+1}/{RUN}: {elapsed:.2f} ms")

# Statistics
avg_time = sum(elapsed_times) / len(elapsed_times)
min_time = min(elapsed_times)
max_time = max(elapsed_times)
print(f"\n{'='*50}")
print(f"Results ({RUN} runs):")
print(f" Average: {avg_time:.2f} ms")
print(f" Min: {min_time:.2f} ms")
print(f" Max: {max_time:.2f} ms")
print(f"{'='*50}")

The pipeline latency has 50% performance regression after the regression PR.

Since the PR is targeted to fix the DDP issue, I think we can check if DDP before using .contiguous(). WDYT? @sayakpaul

Hi @jinc7461 . Could you please provide the script to reproduce the error, and give me some advice to check before using .contiguous() ? Thanks!

cc @sywangyi

Logs

System Info

torch 2.11.0.dev20260113+cpu
platform: Intel Xeon 6

Who can help?

No response

[sayakpaul]

Thanks for pointing this out. Running these pipelines on CPU is something unique to us but it's great to see the possibility and practice.

I would suggest you to open a PR in this case.

[jiqing-feng]

Thanks for pointing this out. Running these pipelines on CPU is something unique to us but it's great to see the possibility and practice.

I would suggest you to open a PR in this case.

Yes, I will open a PR to work around it once I know which case breaks if not using .contiguous()

[jinc7461]

The script to reproduce the error is mentioned in #3809.
I believe this can be resolved by checking if the model is an instance of DistributedDataParallel (DDP) as follows in src/diffusers/models/resnet.py:

from torch.nn.parallel import DistributedDataParallel as DDP

if isinstance(model, DDP):
 xxxx

[jiqing-feng]

Thanks! I will reproduce it and check how to fix it.

[jiqing-feng]

Fixed in #12977. Please review it. Thanks!