@a-r-r-o-w
I'm encountering an error when running test code with gradient checkpointing enabled in my PR.

When backward() is called, the following function is executed:

x = self._gradient_checkpointing_func(resnet, x, feat_cache, feat_idx)

This function calls the forward method of the ResidualBlock. The problem is that the forward method include feat_idx += 1 (feat_idx[0] += 1), which eventually causes feat_idx > len(feat_cache), resulting in an IndexError.

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_wan.py

Lines 239 to 279 in 3be6706

    
           def forward(self, x, feat_cache=None, feat_idx=[0]): 
        
               # Apply shortcut connection 
        
               h = self.conv_shortcut(x) 
        
               # First normalization and activation 
        
               x = self.norm1(x) 
        
               x = self.nonlinearity(x) 
        
               if feat_cache is not None: 
        
                   idx = feat_idx[0] 
        
                   cache_x = x[:, :, -CACHE_T:, :, :].clone() 
        
                   if cache_x.shape[2] < 2 and feat_cache[idx] is not None: 
        
                       cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) 
        
                   x = self.conv1(x, feat_cache[idx]) 
        
                   feat_cache[idx] = cache_x 
        
                   feat_idx[0] += 1 
        
               else: 
        
                   x = self.conv1(x) 
        
               # Second normalization and activation 
        
               x = self.norm2(x) 
        
               x = self.nonlinearity(x) 
        
               # Dropout 
        
               x = self.dropout(x) 
        
               if feat_cache is not None: 
        
                   idx = feat_idx[0] 
        
                   cache_x = x[:, :, -CACHE_T:, :, :].clone() 
        
                   if cache_x.shape[2] < 2 and feat_cache[idx] is not None: 
        
                       cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) 
        
                   x = self.conv2(x, feat_cache[idx]) 
        
                   feat_cache[idx] = cache_x 
        
                   feat_idx[0] += 1 
        
               else: 
        
                   x = self.conv2(x) 
        
               # Add residual connection 
        
               return x + h

I'm not sure how to resolve this issue and would appreciate some guidance. Is there a recommended way to handle the feat_idx incrementing when using gradient checkpointing?

Thank you for your help!

AutoencoderKLWan - support grandient_checkpointing #11071

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions