Mesh Transformer Jax

A haiku library using the new(ly documented) xmap operator in Jax for model parallelism of transformers.

See enwik8_example.py for an example of using this to implement an autoregressive language model.

Benchmarks

On a TPU v3-8 (see tpuv38_example.py):

~2.7B model

Initialized in 121.842s
Total parameters: 2722382080
Compiled in 49.0534s
it: 0, loss: 20.311113357543945
<snip>
it: 90, loss: 3.987450361251831
100 steps in 109.385s
effective flops (not including attn): 2.4466e+14

~4.8B model

Initialized in 101.016s
Total parameters: 4836720896
Compiled in 52.7404s
it: 0, loss: 4.632925987243652
<snip>
it: 40, loss: 3.2406811714172363
50 steps in 102.559s
effective flops (not including attn): 2.31803e+14

10B model

Initialized in 152.762s
Total parameters: 10073579776
Compiled in 92.6539s
it: 0, loss: 5.3125
<snip>
it: 40, loss: 3.65625
50 steps in 100.235s
effective flops (not including attn): 2.46988e+14

Model parallel transformers in Jax and Haiku

Related tags

Overview

Mesh Transformer Jax

Benchmarks

~2.7B model

~4.8B model

10B model

TODO

Owner

Ben Wang

Implementation of the federated dual coordinate descent (FedDCD) method.

Pixray is an image generation system

EM-POSE 3D Human Pose Estimation from Sparse Electromagnetic Trackers.

Contrastive Fact Verification

[NeurIPS 2021] “Improving Contrastive Learning on Imbalanced Data via Open-World Sampling”,

MLSpace: Hassle-free machine learning & deep learning development

Post-training Quantization for Neural Networks with Provable Guarantees

Official PyTorch implementation for FastDPM, a fast sampling algorithm for diffusion probabilistic models

Aesara is a Python library that allows one to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays.

This is the code for our KILT leaderboard submission to the T-REx and zsRE tasks. It includes code for training a DPR model then continuing training with RAG.

NeRD: Neural Reflectance Decomposition from Image Collections

PyTorch trainer and model for Sequence Classification

This is the official repository of Music Playlist Title Generation: A Machine-Translation Approach.

A PyTorch port of the Neural 3D Mesh Renderer

Motion Planner Augmented Reinforcement Learning for Robot Manipulation in Obstructed Environments (CoRL 2020)

3D AffordanceNet is a 3D point cloud benchmark consisting of 23k shapes from 23 semantic object categories, annotated with 56k affordance annotations and covering 18 visual affordance categories.

GemNet model in PyTorch, as proposed in "GemNet: Universal Directional Graph Neural Networks for Molecules" (NeurIPS 2021)

Image to Image translation, image generataton, few shot learning

[CVPR 2022] Semi-Supervised Semantic Segmentation Using Unreliable Pseudo-Labels

3D cascade RCNN for object detection on point cloud