Add option to use flash attention #222

Warvito · 2023-02-04T14:34:39Z

Fixes #210

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito · 2023-02-04T14:49:47Z

@ericspod After adding the option to use the efficient memory to the model I started to get the following error when running the unit tests:

======================================================================
ERROR: test_script_conditioned_2d_models (tests.test_diffusion_model_unet.TestDiffusionModelUNet2D)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/media/walter/Storage/Projects/GenerativeModels/tests/test_diffusion_model_unet.py", line 402, in test_script_conditioned_2d_models
    test_script_save(net, torch.rand((1, 1, 16, 16)), torch.randint(0, 1000, (1,)).long(), torch.rand((1, 1, 3)))
  File "/media/walter/Storage/Projects/GenerativeModels/tests/utils.py", line 723, in test_script_save
    convert_to_torchscript(
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/monai/networks/utils.py", line 597, in convert_to_torchscript
    script_module = torch.jit.script(model, **kwargs)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 476, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 542, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 393, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 894, in compile_unbound_method
    create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 393, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_recursive.py", line 863, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/_script.py", line 1340, in script
    ast = get_jit_def(obj, obj.__name__)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/frontend.py", line 293, in get_jit_def
    return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/frontend.py", line 331, in build_def
    param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types)
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/torch/jit/frontend.py", line 366, in build_param_list
    raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/media/walter/Storage/Projects/GenerativeModels/venv/lib/python3.10/site-packages/xformers/ops/fmha/__init__.py", line 125
    scale: Optional[float] = None,
    *,
    op: Optional[AttentionOp] = None,
                                ~~~~ <--- HERE
) -> torch.Tensor:
    """Implements the memory-efficient attention mechanism following
'CrossAttention._memory_efficient_attention_xformers' is being compiled since it was called from 'CrossAttention.forward'
  File "/media/walter/Storage/Projects/GenerativeModels/generative/networks/nets/diffusion_model_unet.py", line 164
    
        if self.use_flash_attention:
            x = self._memory_efficient_attention_xformers(query, key, value)
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        else:
            x = self._attention(query, key, value)

It looks like a problem with the dependency and torchscript. How would you suggest to deal with this issue?

ericspod · 2023-02-06T15:00:33Z

Where you call memory_efficient_attention you may have to provide all arguments that it requires and do so positionally, so not like keywords like it is and explicitly providing the default values for those you don't want to provide values for here.

…tion-to-the-diffusion-unet # Conflicts: # generative/networks/nets/diffusion_model_unet.py # requirements-dev.txt

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Add option to use flash attention

1756a2c

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito linked an issue Feb 4, 2023 that may be closed by this pull request

Add option to use memory efficient attention to the Diffusion UNet #210

Closed

Warvito added 2 commits February 4, 2023 14:36

Fix docstring

258aa66

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Add xformers==0.0.16 to requirements-dev.txt

9b31ca8

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

OeslleLucena force-pushed the main branch from eeeca8f to b0e3937 Compare February 8, 2023 11:02

Warvito added 2 commits February 26, 2023 19:12

Merge branch 'main' into 210-add-option-to-use-memory-efficient-atten…

ef0f743

…tion-to-the-diffusion-unet # Conflicts: # generative/networks/nets/diffusion_model_unet.py # requirements-dev.txt

Fixed torchscript error

617fffa

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito marked this pull request as ready for review February 26, 2023 19:19

Warvito added the need reviewer This PR need a reviewer label Feb 27, 2023

Warvito merged commit 811ac1e into main Feb 27, 2023

Warvito deleted the 210-add-option-to-use-memory-efficient-attention-to-the-diffusion-unet branch March 1, 2023 13:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add option to use flash attention #222

Add option to use flash attention #222

Uh oh!

Warvito commented Feb 4, 2023

Uh oh!

Warvito commented Feb 4, 2023

Uh oh!

ericspod commented Feb 6, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Add option to use flash attention #222

Add option to use flash attention #222

Uh oh!

Conversation

Warvito commented Feb 4, 2023

Uh oh!

Warvito commented Feb 4, 2023

Uh oh!

ericspod commented Feb 6, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants