search

HazyResearch / zoology

Understand and test language model architectures on synthetic tasks.
Apache License 2.0
164 stars 28 forks source link

Broken link #9

Closed thomasj02 closed 11 months ago

thomasj02 commented 11 months ago

The second zoology blog post has this text:

we compare the PyTorch reference implementation provided here, with the CUDA code and observe the following speedups! Note that the PyTorch code runs out of memory at N=32768 and BS=4. Experiments are run on an 80GB A100 GPU.

Unfortunately that link doesn't seem to work anymore!

On a related note, is the CUDA code available?

simran-arora commented 11 months ago

Hi they got removed in a refactor -- here you go: https://github.com/HazyResearch/zoology/tree/d0c2fa46005c638b525f29fdeb9fc6ece06fdd28/based_refs

We will release the CUDA code with full release! The structure of ```linear_attn_ref.py''' at the link above shows the python forward and backward pass for the CUDA algorithm we use.