PEER - Pytorch
Project description
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}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
peer_pytorch-0.0.7.tar.gz
(264.9 kB
view hashes)
Built Distribution
Close
Hashes for peer_pytorch-0.0.7-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4a5fba3e82830310920e3bc5298b7bd85316a35459263bf0c0466a949149d8cd |
|
MD5 | 8d982c1bbe1f5e95c84b98f55a43c5c0 |
|
BLAKE2b-256 | 09b60eb126aa8a7b342e2b4fc3c3f83e5803f9b9e216476f7d64a53206fb5f6c |