Add include_fc and use_combined_linear argument in the SABlock

[KumoLiu]

Fixes #7991
Fixes #7992

Description

Add include_fc and use_combined_linear argument in the SABlock.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[KumoLiu]

RelPosEmbedding.DECOMPOSED can not convert to torchscript, leave it in this PR. Error shows as below.

'__torch__.torch.nn.modules.container.ParameterList (of Python compilation unit at: 0x55c00fc6db20)' object has no attribute or method '__len__'. Did you forget to initialize an attribute in __init__()?:
  File "/workspace/Code/MONAI/monai/networks/blocks/attention_utils.py", line 98

[KumoLiu]

Remaining parts need to be addressed.
#7977 (comment) -- addressed.
#7977 (comment) -- addressed.
#7977 (comment) -- addressed.

[KumoLiu]

/build

[ericspod]

For reference, I had attempted to test the similarity between the SABlock in core versus that in GenerativeModels and didn't have any luck narrowing down where the differences were. I have this at the end of test_selfattension.py which defines some failing tests that should pass if I haven't missed anything:

from tests.utils import SkipIfBeforePyTorchVersion, assert_allclose
from monai.utils import set_determinism

generative, has_generative = optional_import("generative")
xops, has_xformers = optional_import("xformers.ops")


class TestComparison(unittest.TestCase):
    @parameterized.expand([["cuda:0", True], ["cuda:0", False], ["cpu", False]])
    @skipUnless(has_einops, "Requires einops")
    @skipUnless(has_generative, "Requires generative")
    @skipUnless(has_xformers, "Requires xformers")
    @SkipIfBeforePyTorchVersion((2, 0))
    def test_generative_vs_core(self, _device, use_flash_attention):
        device = torch.device(_device)
        input_shape = (2, 512, 360)
        input_param = {
            "hidden_size": 360,
            "num_heads": 4,
            "dropout_rate": 0,
            "use_flash_attention": use_flash_attention,
        }

        set_determinism(0)
        net_gen = generative.networks.blocks.SABlock(**input_param).to(device)
        set_determinism(0)
        net_monai = SABlock(**input_param).to(device)
        set_determinism(0)
        input = torch.randn(input_shape).to(device)

        net_monai.load_state_dict(net_gen.state_dict())

        with eval_mode(net_gen, net_monai):
            set_determinism(0)
            r1 = net_gen(input)
            set_determinism(0)
            r2 = net_monai(input)
            assert_allclose(r1.detach().cpu().numpy(), r2.detach().cpu().numpy())

I'm not suggesting we add this test but we should come back to where the differences are coming from.

[ericspod]

Hi @KumoLiu I had a few comments about the tests that need to be addressed to test the actual cases you want. I'm otherwise good with things.

[guopengf]

Hi @KumoLiu, shall we also add those parameters (use_flash_attention, include_fc and use_combined_linear) to monai.networks.nets.controlnet and monai.networks.nets.autoencoderkl? Those two networks also use attention layers.

[KumoLiu]

For reference, I had attempted to test the similarity between the SABlock in core versus that in GenerativeModels and didn't have any luck narrowing down where the differences were.

I actually tested this locally and confirmed that they return the same result. It seems that the issue with your script might be due to differences in the initial weights. To address this, I copied and reused the state dict. I didn't include this test in this pr because we aim to avoid introducing generative.

For your reference, here is the test code I used:

import torch
from copy import deepcopy
from generative.networks.nets.diffusion_model_unet import CrossAttention, AttentionBlock
from monai.utils import set_determinism
from monai.networks.blocks import CrossAttentionBlock, SpatialAttentionBlock

set_determinism(0)
use_flash_attention = True

cross_block_gen = CrossAttention(query_dim=512, num_attention_heads=16, num_head_channels=32, use_flash_attention=use_flash_attention).to('cuda')
cross_block_core = CrossAttentionBlock(hidden_size=512, num_heads=16, dim_head=32, use_flash_attention=use_flash_attention).to('cuda')
state_dict_gen_cross = cross_block_gen.state_dict()
state_dict_core_cross = deepcopy(state_dict_gen_cross)
state_dict_core_cross["out_proj.weight"] = state_dict_gen_cross["to_out.0.weight"]
state_dict_core_cross["out_proj.bias"] = state_dict_gen_cross["to_out.0.bias"]
del state_dict_core_cross["to_out.0.weight"]
del state_dict_core_cross["to_out.0.bias"]

atten_block_gen = AttentionBlock(spatial_dims=2, num_channels=512, num_head_channels=32, use_flash_attention=use_flash_attention).to('cuda')
atten_block_core = SpatialAttentionBlock(spatial_dims=2, num_channels=512, num_head_channels=32, use_combined_linear=False, include_fc=False, use_flash_attention=use_flash_attention).to('cuda')
state_dict_gen_atten = atten_block_gen.state_dict()
state_dict_core_atten = atten_block_core.state_dict()
for key in state_dict_core_atten.keys():
    if 'attn' in key and "proj" not in key:
        state_dict_core_atten[key] = state_dict_gen_atten[key[5:]]
    if key in state_dict_gen_atten:
        state_dict_core_atten[key] = state_dict_gen_atten[key]
state_dict_core_atten["attn.out_proj.weight"] = state_dict_gen_atten["proj_attn.weight"]
state_dict_core_atten["attn.out_proj.bias"] = state_dict_gen_atten["proj_attn.bias"]


input_cross = torch.rand(1, 256, 512).to("cuda")
input_cross2 = deepcopy(input_cross)
cross_block_gen.load_state_dict(state_dict_gen_cross)

out_gen = cross_block_gen(input_cross)
cross_block_core.load_state_dict(state_dict_core_cross)
out_core = cross_block_core(input_cross2)

print('diff in cross', (out_gen-out_core).abs().sum())

input_atten = torch.rand(1, 512, 16, 16).to("cuda")
input_atten2 = deepcopy(input_atten)
atten_block_gen.load_state_dict(state_dict_gen_atten)

out_gen = atten_block_gen(input_atten)
atten_block_core.load_state_dict(state_dict_core_atten)
out_core = atten_block_core(input_atten2)

print('diff in atten', (out_gen-out_core).abs().sum())

Result:

diff in cross tensor(0., device='cuda:0', grad_fn=<SumBackward0>)
diff in atten tensor(0., device='cuda:0', grad_fn=<SumBackward0>)

And for this one :

We may also want to test with and without CUDA, I've noticed some differences on CPU between using flash attention and not.

Yes, we also include these tests. The gpu pipeline have been removed to blossom, but it will stilled be tested.

[KumoLiu]

Hi @KumoLiu, shall we also add those parameters (use_flash_attention, include_fc and use_combined_linear) to monai.networks.nets.controlnet and monai.networks.nets.autoencoderkl? Those two networks also use attention layers.

Hi @guopengf, added in the latest commit.

[KumoLiu]

/build