Merge pull request #42 from laekov/v0.2.0-pre-release

Checkout version number to V0.2.0

examples/megatron/README.md

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

setup.py

@@ -6,6 +6,15 @@ import os
 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')
@@ -14,9 +23,9 @@ if os.environ.get('USE_NCCL', '0') == '1':
 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='hja20@mails.tsinghua.edu.cn',
         license='Apache-2',
         url='https://github.com/laekov/fastmoe',