about flash attention v1 and v2 speed

[daixiangzi]

When I use different flash versions, the speed is almost the same. From my understanding, shouldn't v2 be faster than v1?
PyTorch version:
v1:'torch2.0.1+cu117'
v2:'2.2.0a0+git18e1a37 from source install on today'
this my test code

   rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    distributed.init_process_group(backend="nccl")
    torch.cuda.set_device(local_rank)
    backbone = ViT_S_16().cuda()
    backbone = torch.nn.parallel.DistributedDataParallel(
        module=backbone,
        bucket_cap_mb=32,
        find_unused_parameters=True,
        static_graph=True)
    auto_scaler = torch.cuda.amp.grad_scaler.GradScaler(growth_interval=200)
    ce = torch.nn.CrossEntropyLoss().cuda()
    opt = torch.optim.AdamW(backbone.parameters(), 0.001)
    t = time.time()
    
    for i in tqdm(range(1000)):
        img = torch.randn(512,3,224,224,device=local_rank)
        local_labels = torch.zeros(size=(512,), dtype=torch.long, device=local_rank)
        with torch.cuda.amp.autocast(True):
                result = backbone(img)
                loss = ce(result, local_labels)
        auto_scaler.scale(loss).backward()
        auto_scaler.unscale_(opt)
        auto_scaler.step(opt)
        auto_scaler.update()
        opt.zero_grad()
    print((time.time()-t)/1000)

[bdhirsh]

cc @drisspg

[drisspg]

There is not alot to go off of here @daixiangzi, can you try isolating the SDPA component of your model? Just because SDPA is ~2x faster doesn't mean that end to end your model will run 2x faster. There are likely other bottlenecks of your model. I would suggest trying the torch profiler to get a better understanding of the runtime characteristics of your PyTorch code:

https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html

[daixiangzi]

There is not alot to go off of here @daixiangzi, can you try isolating the SDPA component of your model? Just because SDPA is ~2x faster doesn't mean that end to end your model will run 2x faster. There are likely other bottlenecks of your model. I would suggest trying the torch profiler to get a better understanding of the runtime characteristics of your PyTorch code:

https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html

thank for you reply's, I will try to it

[daixiangzi]

v1:

v2:

[daixiangzi]

host total seems that v2 is higher than v1

[daixiangzi]

with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA],schedule=torch.profiler.schedule(skip_first=3, wait=1, warmup=1, active=1, repeat=1),use_cuda=True,on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/')) as prof:

        for i in tqdm(range(100)):
            img = torch.randn(512,3,224,224,device=local_rank)
            local_labels = torch.zeros(size=(512,), dtype=torch.long, device=local_rank)
            with torch.cuda.amp.autocast(True):
                        result = backbone(img)
                        loss = ce(result, local_labels)
            auto_scaler.scale(loss).backward()
            auto_scaler.unscale_(opt)
            auto_scaler.step(opt)
            auto_scaler.update()
            opt.zero_grad()
            prof.step()

my profile code

[daixiangzi]

sorry,my code have some bugs,i fixed bug, flashv2 greater than flash v1 on MultiheadAttention layer