PEER-pytorch 0.2.1
Last updated:
0 purchases
Description:
PEERpytorch 0.2.1
PEER - Pytorch
Pytorch implementation of the PEER block from the Deepmind paper, Mixture of A Million Experts, by Xu Owen He.
Install
$ pip install PEER-pytorch
Usage
import torch
from PEER_pytorch import PEER
peer = PEER(
dim = 512,
heads = 8, # tested up to 32 - (hk = heads * num_experts_per_head (16))
num_experts = 1_000_000, # he chose 1 million
num_experts_per_head = 16, # he settled on 16, but was 32 in PKM paper
dim_key = 128,
pre_rmsnorm = True
).cuda()
x = torch.randn(2, 1024, 512).cuda()
out = peer(x) + x
assert x.shape == out.shape
Citations
@inproceedings{He2024MixtureOA,
title = {Mixture of A Million Experts},
author = {Xu Owen He},
year = {2024},
url = {https://api.semanticscholar.org/CorpusID:271038610}
}
@article{Csordas2023ApproximatingTF,
title = {Approximating Two-Layer Feedforward Networks for Efficient Transformers},
author = {R'obert Csord'as and Kazuki Irie and J{\"u}rgen Schmidhuber},
journal = {ArXiv},
year = {2023},
volume = {abs/2310.10837},
url = {https://api.semanticscholar.org/CorpusID:264172384}
}
License:
For personal and professional use. You cannot resell or redistribute these repositories in their original state.
Customer Reviews
There are no reviews.
