Skip to content

Commit

Permalink
Merge pull request #42 from laekov/v0.2.0-pre-release
Browse files Browse the repository at this point in the history
Checkout version number to V0.2.0
  • Loading branch information
laekov authored May 31, 2021
2 parents d205aae + 411e57f commit c96f886
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 8 deletions.
22 changes: 16 additions & 6 deletions examples/megatron/README.md
Original file line number Diff line number Diff line change
@@ -1,15 +1,25 @@
FastMoE currently works with both `v2.0` and `v2.1` release of
FastMoE works with different versions of
[Megatron-LM](https://github.com/nvidia/megatron-lm).
See `fmoe/megatron/utils.py` for arguments of FastMoE.

Patches which you can find in this directory are used to easily enable MoE in
different versions of Megatron-LM for training Bert. The usage is the same in
other training scripts.
An example patch is provided for `v2.2` release.
The patch can be directly applied to add FastMoE support if you are using
Megatron-LM v2.2.
Otherwise, you may need to manually enable FastMoE in your codebase.
The patch includes the following modifications.

The patch works in the following way.
### Add arguments to Megatron's argparser

In `megatron/arguments.py`, add `_add_fmoe_args` to the parser.

### Patch checkpoint

In `megatron/training.py`, replace `load_checkpoint` and `save_checkpoint` by
functions with the same name in `fmoe.megatron.checkpointing`.

### Building the model in FastMoE style

In `pretrain_bert.py`, the `fmoe.megatron.fmoefy` function is used as an
In `megatron/training.py`, the `fmoe.megatron.fmoefy` function is used as an
entrance to one-key introduce FastMoE layer to replace the MLP layers in the
transformer language models.

Expand Down
13 changes: 11 additions & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,15 @@
cxx_flags = []
ext_libs = []

authors = [
'Jiaao He',
'Jiezhong Qiu',
'Aohan Zeng',
'Tiago Antunes',
'Jinjun Peng',
'Qin Li',
]

if os.environ.get('USE_NCCL', '0') == '1':
cxx_flags.append('-DFMOE_USE_NCCL')
ext_libs.append('nccl')
Expand All @@ -14,9 +23,9 @@
if __name__ == '__main__':
setuptools.setup(
name='fastmoe',
version='0.1.2',
version='0.2.0',
description='An efficient Mixture-of-Experts system for PyTorch',
author='Jiaao He, Jiezhong Qiu and Aohan Zeng',
author=', '.join(authors),
author_email='[email protected]',
license='Apache-2',
url='https://github.com/laekov/fastmoe',
Expand Down

0 comments on commit c96f886

Please sign in to comment.