Seeking Help: OutOfOrderError with DiffusionModel

Model Support
1

Hi everyone,

I’m encountering a perplexing OutOfOrderError when trying to collect layer activations during model generation. My goal is to extract the outputs from specific layers while a model is running.

Here’s a simplified version of my code:

with model.generate(
    prompts,
    **generate_kwargs,
    device=device,
    seed=seed,
):
    with model.iter[0:50] as step:
        for name, layer in layers.items():
            print(step, name)
            activations.append(layer.output) # also tried with output.save()
    
    outputs.extend(model.output.images)

return activations, outputs

The error I’m getting is:

OutOfOrderError: Value was missed for model.unet.mid_block.attentions.0.transformer_blocks.0.attn1.to_q.output.i0. Did you call an Envoy out of order?

I suspect this might be happening because the order in which I’m iterating through my layers dictionary isn’t the same as the order in which the layers are executed during the forward pass. The model.iter context seems to be designed to handle this, but I’m clearly doing something wrong.

Is there a standard or recommended way to correctly collect activations from a list of specific layers without causing an OutOfOrderError? Do I need to get the layers in the correct execution order first? If so, how can I do that? Any advice or examples would be greatly appreciated!

Thanks in advance for your help! :folded_hands:

2

Hi @mamiglia - happy to see you using our DiffusionModel class!

Can you share more details about your tracing setup:

  • nnsight version
  • model name
  • layers you are looping over

Thanks.

3

Hey @AdamBelfki , here’s the details:

  • nnsight version 0.5.1
  • model=stable-diffusion-v1-5/stable-diffusion-v1-5

Here’s a reduced version of the code I use:

layers = OrderedDict()
regex = r"blocks?\.(1\.)?attentions\.0.*attn2\.to_[kqv]"

for name, module in model.unet.named_modules():
    if re.search(regex, name):
        layers[name] = module

So layers.keys() returns

'model.unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q', 
'model.unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k', 
'model.unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v', 
'model.unet.up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q', 
'model.unet.up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k', 
'model.unet.up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v', 
'model.unet.mid_block.attentions.0.transformer_blocks.0.attn2.to_q', 
'model.unet.mid_block.attentions.0.transformer_blocks.0.attn2.to_k', 
'model.unet.mid_block.attentions.0.transformer_blocks.0.attn2.to_v'

So as you can see the problem is that mid_blocks is extracted after the up_blocks, so they are called in the wrong order. I could sort them manually, but I’d like to avoid this.

4

bump @AdamBelfki

would it make sense to use empty invokes here?

5

Also @mamiglia if you just want to collect the activation, you should use the nnsight 0.5 new cache feature: NNsight 0.5: Official Walkthrough — nnsight
that would also solve your problem

6

Hey @mamiglia . This is my suggested workflow for you using the cache.

import re
from collections import OrderedDict

import torch

from nnsight.modeling.diffusion import DiffusionModel

id = "CompVis/stable-diffusion-v1-4"
model = DiffusionModel(id, dispatch=True).to('cuda:0').to(torch.bfloat16)


layers = OrderedDict()
regex = r"blocks?\.(1\.)?attentions\.0.*attn2\.to_[kqv]"

for name, module in model.unet.named_modules():
    if re.search(regex, name):
        layers[name] = module
        
with model.generate("world",num_inference_steps=50, seed=40) as tracer:
    
    cache = tracer.cache(layers)
    
    output = tracer.result.save()
7

I’m curious why you use an OrderedDict instead of a list here, does it change anything?