Merge pull request #100 from laekov/doc-fix

Fix document for megatron

fmoe/megatron/balance.py

@@ -46,6 +46,9 @@ def generate_megatron_gate_hook(layer_idx, num_expert_global):
 
 
 def add_balance_log(model, writer, iteration):
+    r"""
+    Note that this function does not work with pipeline parallelism
+    """
     from megatron import is_last_rank
 
     while hasattr(model, 'module'):

fmoe/megatron/layers.py

@@ -162,9 +162,6 @@ def fmoefy(
     they are trained in data-parallel mode. This can be useful when testing on
     small models that do not require high training throughput or large parameter
     capacity.
-    Note that pipeline parallel is not supported yet. When distributed experts
-    are enabled, their communicator should be Megatron's
-    tensor_model_parall_comm x data_parallel_comm, which is not created.
     """
     from megatron import get_args