(seq_len, hidden_size, num_heads, batch_size=8, qkv=None)
| 2 | import torch |
| 3 | |
| 4 | def test_mea(seq_len, hidden_size, num_heads, batch_size=8, qkv=None): |
| 5 | torch.cuda.manual_seed(0) |
| 6 | q, k, v = qkv.clone().cuda().requires_grad_(True) |
| 7 | # profile time and memory |
| 8 | torch.cuda.reset_peak_memory_stats() |
| 9 | with torch.autograd.profiler.profile(use_cuda=True, profile_memory=True) as prof: |
| 10 | out = memory_efficient_attention(q, k, v) |
| 11 | out.sum().backward() |
| 12 | print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10)) |
| 13 | |
| 14 | print(torch.cuda.max_memory_allocated() / 1024 ** 3) |
| 15 | # clean cuda cache |
| 16 | del q, k, v |
| 17 | torch.cuda.empty_cache() |
| 18 | return out |
| 19 | |
| 20 | def attention(query, key, value): |
| 21 | scale = 1 / query.shape[-1] ** 0.5 |
nothing calls this directly
no test coverage detected