lucidrains / peer-pytorch

Pytorch implementation of the PEER block from the paper, Mixture of A Million Experts, by Xu Owen He at Deepmind

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}
}

peer-pytorch's Issues

where are the 1M pretrained loras

can we have the 1M pretrained Lora weights only if the base model is the same? or 1M lora with corresponding output layers

todo

• batchnorm and look into the causality issue from pkm paper
• think about the GLU extension w/ triton (author suggested it as potential follow up work) if public reports good results
• improvise a mixture of attention module with this scheme

Usage with x-transformers

PEER looks like an interesting approach and thanks for implementing so cleanly! I do have a quick question though about recommended usage with x-transformers. Would something like this be a good way of using it?

import torch
from PEER_pytorch import PEER
from x_transformers import ContinuousTransformerWrapper, Encoder

peer = PEER(
    dim = 512,
    heads = 8,                   
    num_experts = 1_000_000,     
    num_experts_per_head = 16,   
    dim_key = 128,
    pre_rmsnorm = True
).cuda()


pre_peer = ContinuousTransformerWrapper(
    max_seq_len = 1024,
    attn_layers = Encoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)

post_peer = ContinuousTransformerWrapper(
    max_seq_len = 1024,
    attn_layers = Encoder(
        dim = 512,
        depth = 6,
        heads = 8
    )
)


x = torch.randn(2, 1024, 512).cuda()

out = pre_peer(x)
out = peer(out) + out
out = post_peer(out)`