[core] CogVideoX memory optimizations in VAE encode

[a-r-r-o-w]

(cherry picked from commit bf890bc)

What does this PR do?

Adds VAE encode memory optimizations from #9333 separately for a cleaner history of additions and because this is needed in #9302.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@DN6 @yiyixuxu

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

Can this be configurable? Why change to a different hardcoded value?

[a-r-r-o-w]

This can be configurable to multiples of temporal_compression_ratio which is 4. We can configure this using vae.num_sample_frames_batch_size = 16 for the moment, and I think it's best to do it that way. 8, here, corresponds to 1 second of video because Cog was trained on 8 FPS videos. I would leave the configuration option out of enable_tiling and keep it for power users who have looked through the code and understand its usage tbh

[DN6]

LGTM. Just a comment about whether num_sample_frames_batch_size can be made configurable?

[YanzuoLu]

This optimization doesn't seem to be the right thing to do I think maybe?
As the gather_norm is enabled for all the vae encoders in CogVideoX(both 1.0 and 1.5), it means that the normalization here is identically applied on the whole batched data [B,C,T,H,W] with full sequence of time frames.
If we encode in framewise like 8 frames per forward call, the mean and std obtained for calculating group_norm is applied on a sub-sequence of time frames, which should be not expected.
And I also see that in CogVideoX1.5, the decoder is also with gather_norm, it should indicate that framewise decoding is also not a perfect practice that can align with sat implementation.
I want to know if my understanding is wrong? @a-r-r-o-w

[a-r-r-o-w]

gather_norm is just the SpatialNorm3D layer in the VAE. It contains one GroupNorm and two Conv3D layers. Since GroupNorm only operates on the channel dimension in groups of fixed size, you can chunk inference across all the other dimensions (in this case, we are chunking about the frame dimension). The Conv3D is a bit more tricky to understand because of the intermediate conv_cache padding, but it is also independent of the frame dimension as long as we correctly handle padding (otherwise conv across temporal dimension may lead to lower or higher frames than expected depending on whether you are doing encoding or decoding).

If there was any operation in which the channel dimension was involved but was not independent of frame dimension (for example, if there were any attention layers), then framewise encoding/decoding would not have been possible. This is the case with Mochi VAE, for example. The encoder contains attention layer, so you can't chunk across frame dimension and perform framewise encoding, but since decoder is comprised only of convolutions (with appropriate padding) and norms, computation can be done independent of the frame dimension.

I hope that makes sense haha!

[YanzuoLu]

Yes I know what you mean about conv3d because I'm just adapting the diffusers implementation with real context parallel to support my distillation training.
But my question is just about this GroupNorm.
The hidden_states are at dimensions like [Batch, Channel, TimeFrame, Height, Width].
So when we apply GroupNorm operation, let's say the group size is 32, it takes [1, 32, T, H, W] to calculate the mean and std for normalization as in here.
Such that if the vae training process is at context parallel with gather norm, it is taking the whole TimeFrame sequence to calculate it.
But if not, only the TimeFrame sequence in current GPU would be used, let's say we have cp_size=4, then the tensors for calculating mean and std would bd [1, 32, T/4, H, W], such it now depends on different position of time frames

[YanzuoLu]

If the gather_norm is enabled, it means that the mean and std need to take consideration of the whole TimeFrame sequence, this will disable us to conduct framewise encoding/decoding.

If the gather_norm is disabled, we need to carefully align with the training process to have a specific X frames per forward.

Is that right? I have been confused about this for a long time...

[a-r-r-o-w]

Maybe I can give a simpler example with nn.Linear to explain how doing it framewise, or all frames at once, is equivalent. Consider this code:

B, S = 4, 100
model = nn.Linear(256, 1024)
input = torch.randn((B, S, 256))

# Normal forward
output1 = model(input)

# Chunking across batch dimension with total 256 // 4 = 64 forward passes
output2 = torch.cat([model(x) for x in input.split(4, dim=0)], dim=0)

# Chunking across sequence length dimension with total 100 // 20 = 5 forward passes
output3 = torch.cat([model(x) for x in input.split(20, dim=1)], dim=1)

# Collapsing across single dimension with total (256 * 100) // 50 = 512 forward passes
input = input.flatten(0, 1)
output4 = torch.cat([model(x) for x in input.split(50, dim=0)], dim=0)
output4 = output4.unflatten(0, (B, -1))

The outputs in all these cases are the exact same. Due to numerical precision issues and rounding and order of operations, there is maybe a small difference at the order of 1e-7 to 1e-15 or lower, which is negligible.

This kind of chunking across dimension is possible for any kind of layer when done correctly. As group norm only affects small number of groups on the channel dimension, computation is fully independent of all the other dimensions (this is why you can perform spatial tiling as well as framewise decoding). It is really equivalent and there is no mistake in doing this.

For your context parallel implementation, the simplest thing you could do for parallelizing is collapse your tensor to a [-1, channels] shape tensor and then evenly distribute across all processes. It should roughly be the same as any other way of implementing context parallel. I hope this explanation helps clear the doubts with groupnorm a bit

On another note, we are trying to look into distillation of CogVideoX too internally, so if you're interested in collaborating on the training and further future work, we would definitely love to join you! If this is of interest, let me know and I will setup a Slack channel where we can communicate together

[YanzuoLu]

Of course I would be happy to do so!!! I think I have been in a channel with @sayakpaul about our Hyper-SD.
Maybe you could send me an new channel invitation email if you want to have new one?

[YanzuoLu]

nonono...
Please refer to the pseudo code in GroupNorm paper.
The mean and std is calculated on the whole sample.

[YanzuoLu]

It means that it is not like linear layers to be independent of the other dimensions I think!