[experiment] one shot decoding instead of framewise in Cog VAE

[a-r-r-o-w]

What does this PR do?

Small experiment to see if we can perform decoding at once instead of batching frame-wise, while applying quantization to see if the overall memory required remains approximately the same.

Here are the results:

dtype	max_memory	max_memory_reserved
bfloat16	54.291	71.307
int8wo	54.284	71.312
int8dq	54.283	71.308
fp6_e3m2	54.285	71.307

So it seems like this doesn't work but good to have it verified, unless I've done something wrong. This is due to needing to materialize the bfloat16 tensors for intermediate computations, which is done whether or not we use quantization.

code

import argparse
import gc

import numpy as np
import torch
from diffusers.utils import load_video
from diffusers import AutoencoderKLCogVideoX
from torchao.quantization import int8_weight_only, int8_dynamic_activation_int8_weight, quantize_, fpx_weight_only


DTYPE_CONVERTER = {
    "float16": lambda module: module.to(dtype=torch.float16),
    "bfloat16": lambda module: module.to(dtype=torch.bfloat16),
    "int8wo": lambda module: quantize_(module, int8_weight_only()),
    "int8dq": lambda module: quantize_(module, int8_dynamic_activation_int8_weight()),
    "fp6_e3m2": lambda module: quantize_(module, fpx_weight_only(3, 2)),
}


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float16", "bfloat16", "int8wo", "int8dq", "fp6_e3m2"])
    parser.add_argument("--disable_framewise_batching", action="store_true")
    return parser.parse_args()


def reset_memory():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.reset_accumulated_memory_stats()


def print_memory():
    memory = torch.cuda.memory_allocated() / 1024**3
    max_memory = torch.cuda.max_memory_allocated() / 1024**3
    max_reserved = torch.cuda.max_memory_reserved() / 1024**3
    print(f"{memory=:.3f}")
    print(f"{max_memory=:.3f}")
    print(f"{max_reserved=:.3f}")


@torch.no_grad()
def main(args):
    reset_memory()
    vae = AutoencoderKLCogVideoX.from_pretrained(
        "THUDM/CogVideoX-2b", subfolder="vae", torch_dtype=torch.bfloat16
    ).to("cuda")

    if args.disable_framewise_batching:
        vae.disable_framewise_batching()

    vae_return = DTYPE_CONVERTER[args.dtype](vae)

    if vae_return:
        vae = vae_return

    print("===== Before inference =====")
    print_memory()
    
    video = load_video("inputs/hiker.mp4")
    video = np.array(video).transpose(3, 0, 1, 2) / 127.5 - 1
    video = torch.from_numpy(video).unsqueeze(0).to("cuda", dtype=vae.dtype)
    
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    
    start.record()
    encoded = vae.encode(video, return_dict=False)[0].sample()
    decoded = (vae.decode(encoded, return_dict=False)[0].float() + 1) * 255
    end.record()
    torch.cuda.synchronize()
    time_elapsed = start.elapsed_time(end) / 1000

    decoded = torch.clamp(decoded, 0, 255)
    video = decoded.permute(0, 2, 3, 4, 1)[0].cpu().numpy().astype(np.uint8)
    print(decoded.shape)

    print("===== After inference =====")
    print("Inference time:", round(time_elapsed, 3))
    print_memory()


if __name__ == "__main__":
    args = get_args()
    main(args)

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.

@sayakpaul

[a-r-r-o-w]

Closed since doesn't need to be merged. We can still discuss and re-open if we find something wrong with what I've done and actually see a memory improvement.

[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.

[sayakpaul]

How about inference speed? I think one-shot decoding especially when compiled should lead to a reasonable amount of speedup (given we can completely avoid recompilation). Can we verify that? I think that should be the main goal.

Additionally, the memory measurement seems to be contradictory with what’s been reported in diffusers-torchao under this section:

Additional fp8dqrow, fp6_e3m2, fp5_e2m2 and fp4_e2m1 benchmarks
H100

Is it because we are doing decoding one-shot? What am I missing out on?

[sayakpaul]

I will keep this open for a bit as I want to give it a try. Hope that is okay.

[a-r-r-o-w]

Is it because we are doing decoding one-shot? What am I missing out on?

Yes, we need to materialize a tensor that is 12-13 times as larger (or maybe 6 times as large in the intermediate layers (not sure exactly)) than before because we're not doing framewise-batched decoding.

How about inference speed? I think one-shot decoding especially when compiled should lead to a reasonable amount of speedup (given we can completely avoid recompilation). Can we verify that? I think that should be the main goal.

I thought it was to see if quantization + one-shot decoding could use similar amount of memory as bf16 precision frame-wise decoding, but I see what you mean. I'll give compilation a test

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[a-r-r-o-w]

@sayakpaul Should this be closed?

[sayakpaul]

Yeah this can be closed.