()
| 44 | return out.transpose(1,2) |
| 45 | |
| 46 | def test_mixin(): |
| 47 | with torch.no_grad(): |
| 48 | from sat.model import BaseModel, AutoModel |
| 49 | from sat.model.mixins import MemoryEfficientAttentionMixin, TransposedMemoryEfficientAttentionMixin |
| 50 | model = BaseModel(args=BaseModel.get_args( |
| 51 | max_sequence_length=5000, |
| 52 | )) |
| 53 | model = model.cuda().eval().half() |
| 54 | x = torch.tensor([range(4096)], device='cuda') |
| 55 | with torch.autograd.profiler.profile(use_cuda=True, profile_memory=True) as prof: |
| 56 | a = model(input_ids=x, position_ids=x, attention_mask=None) |
| 57 | print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10)) |
| 58 | model.add_mixin('mea', TransposedMemoryEfficientAttentionMixin()) |
| 59 | with torch.autograd.profiler.profile(use_cuda=True, profile_memory=True) as prof: |
| 60 | b = model(input_ids=x, position_ids=x, attention_mask=None) |
| 61 | print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10)) |
| 62 | print(((a[0]-b[0]).abs()/(a[0].abs().mean())).max()) |
| 63 | |
| 64 | if __name__ == "__main__": |
| 65 | # seq_len = 2048 |
no test coverage detected