同步最新代码

megatron/core/models/gpt/gpt_layer_specs.py

 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
+import warnings
 from typing import Optional
 
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.dot_product_attention import DotProductAttention
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
-from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
-from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
 from megatron.core.transformer.multi_latent_attention import (
     MLASelfAttention,
     MLASelfAttentionSubmodules,
@@ -26,12 +26,10 @@ from megatron.core.utils import is_te_min_version
 
 try:
     from megatron.core.extensions.transformer_engine import (
-        TEColumnParallelGroupedLinear,
         TEColumnParallelLinear,
         TEDotProductAttention,
         TELayerNormColumnParallelLinear,
         TENorm,
-        TERowParallelGroupedLinear,
         TERowParallelLinear,
     )
 
@@ -47,8 +45,6 @@ try:
     HAVE_APEX = True
     LNImpl = FusedLayerNorm
 except ImportError:
-    import warnings
-
     from megatron.core.transformer.torch_norm import WrappedTorchNorm
 
     warnings.warn('Apex is not installed. Falling back to Torch Norm')
@@ -60,7 +56,8 @@ def get_gpt_layer_with_transformer_engine_spec(
     moe_grouped_gemm: Optional[bool] = False,
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
-    fp8: Optional[str] = None,
+    fp8: Optional[str] = None,  # pylint: disable=unused-arguments
+    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ) -> ModuleSpec:
     """Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
 
@@ -69,13 +66,24 @@ def get_gpt_layer_with_transformer_engine_spec(
         num_experts (int, optional): Number of experts. Defaults to None.
         moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
-        fp8 (str, optional): Flag to decide the linear layer spec for MoE. Defaults to None.
+        fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
+        moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
+                                                      Defaults to False.
 
     Returns:
         ModuleSpec: Module specification with TE modules
     """
+    if fp8 is not None:
+        warnings.warn(
+            'The fp8 argument in "get_gpt_layer_with_transformer_engine_spec" has been deprecated'
+            ' and will be removed soon. Please update your code accordingly.'
+        )
+
     mlp = _get_mlp_module_spec(
-        use_te=True, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, fp8=fp8
+        use_te=True,
+        num_experts=num_experts,
+        moe_grouped_gemm=moe_grouped_gemm,
+        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
     )
 
     if multi_latent_attention:
@@ -138,6 +146,8 @@ def get_gpt_layer_local_spec(
     moe_grouped_gemm: Optional[bool] = False,
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
+    fp8: Optional[str] = None,  # pylint: disable=unused-arguments
+    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ) -> ModuleSpec:
     """Use this spec for an implementation using only modules in Megatron-Core.
 
@@ -146,13 +156,24 @@ def get_gpt_layer_local_spec(
         num_experts (int, optional): Number of experts. Defaults to None.
         moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
+        fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
+        moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
+                                                      Defaults to False.
 
     Returns:
         ModuleSpec: Module specification with Megatron-Core modules
     """
+    if fp8 is not None:
+        warnings.warn(
+            'The fp8 argument in "get_gpt_layer_local_spec" has been deprecated'
+            ' and will be removed soon. Please update your code accordingly.'
+        )
 
     mlp = _get_mlp_module_spec(
-        use_te=False, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
+        use_te=False,
+        num_experts=num_experts,
+        moe_grouped_gemm=moe_grouped_gemm,
+        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
     )
 
     if multi_latent_attention:
@@ -213,63 +234,33 @@ def _get_mlp_module_spec(
     use_te: Optional[bool] = True,
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
-    fp8: Optional[str] = None,
+    fp8: Optional[str] = None,  # pylint: disable=unused-arguments
+    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ) -> ModuleSpec:
-    """Helper function to get module spec for MLP"""
-    if num_experts is not None:
-        moe_spec = _get_moe_module_spec(
-            use_te=True, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, fp8=fp8
+    """Helper function to get module spec for MLP/MoE"""
+    if fp8 is not None:
+        warnings.warn(
+            'The fp8 argument in "_get_mlp_module_spec" has been deprecated'
+            ' and will be removed soon. Please update your code accordingly.'
         )
-        return moe_spec
-
-    return ModuleSpec(
-        module=MLP,
-        submodules=MLPSubmodules(
-            linear_fc1=TELayerNormColumnParallelLinear if use_te else ColumnParallelLinear,
-            linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
-        ),
-    )
 
-
-def _get_moe_module_spec(
-    use_te: Optional[bool] = True,
-    num_experts: Optional[int] = None,
-    moe_grouped_gemm: Optional[bool] = False,
-    fp8: Optional[str] = None,
-) -> ModuleSpec:
-    """Helper function to get module spec for MoE"""
     if num_experts is None:
-        return None
-    if use_te and moe_grouped_gemm:
-        linear_fc1 = TEColumnParallelGroupedLinear
-        linear_fc2 = TERowParallelGroupedLinear
-    elif use_te and fp8:
-        linear_fc1 = TEColumnParallelLinear
-        linear_fc2 = TERowParallelLinear
-    else:
-        linear_fc1 = ColumnParallelLinear
-        linear_fc2 = RowParallelLinear
-
-    use_te_grouped_gemm = use_te and TEColumnParallelGroupedLinear is not None
-
-    return ModuleSpec(
-        module=MoELayer,
-        submodules=MoESubmodules(
-            experts=(
-                MLPSubmodules(linear_fc1=linear_fc1, linear_fc2=linear_fc2)
-                if not moe_grouped_gemm or use_te_grouped_gemm
-                else None
-            ),
-            shared_experts=ModuleSpec(
-                module=SharedExpertMLP,
-                params={"gate": False},
-                submodules=MLPSubmodules(
-                    linear_fc1=TEColumnParallelLinear if use_te else ColumnParallelLinear,
-                    linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
-                ),
+        # Dense MLP w/ or w/o TE modules.
+        return ModuleSpec(
+            module=MLP,
+            submodules=MLPSubmodules(
+                linear_fc1=TELayerNormColumnParallelLinear if use_te else ColumnParallelLinear,
+                linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
             ),
-        ),
-    )
+        )
+    else:
+        # Mixture of experts with modules in megatron core.
+        return get_moe_module_spec(
+            use_te=use_te,
+            num_experts=num_experts,
+            moe_grouped_gemm=moe_grouped_gemm,
+            moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        )
 
 
 def get_gpt_decoder_block_spec(
@@ -288,7 +279,7 @@ def get_gpt_decoder_block_spec(
             moe_grouped_gemm=False,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            fp8=config.fp8,
+            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         )
         if use_transformer_engine
         else get_gpt_layer_local_spec(
@@ -296,6 +287,7 @@ def get_gpt_decoder_block_spec(
             moe_grouped_gemm=False,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
+            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         )
     )
     moe_layer_spec = (
@@ -304,7 +296,7 @@ def get_gpt_decoder_block_spec(
             moe_grouped_gemm=config.moe_grouped_gemm,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            fp8=config.fp8,
+            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         )
         if use_transformer_engine
         else get_gpt_layer_local_spec(
@@ -312,6 +304,7 @@ def get_gpt_decoder_block_spec(
             moe_grouped_gemm=config.moe_grouped_gemm,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
+            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         )
     )
 

megatron/core/models/gpt/moe_module_specs.py

+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import warnings
+from typing import Optional
+
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.mlp import MLPSubmodules
+from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP, TEGroupedMLP
+from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.utils import get_te_version, is_te_min_version
+
+try:
+    from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelGroupedLinear,
+        TEColumnParallelLinear,
+        TERowParallelGroupedLinear,
+        TERowParallelLinear,
+    )
+
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+
+def get_moe_module_spec(
+    use_te: Optional[bool] = True,
+    num_experts: Optional[int] = None,
+    moe_grouped_gemm: Optional[bool] = False,
+    moe_use_legacy_grouped_gemm: Optional[bool] = False,
+) -> ModuleSpec:
+    """Helper function to get module spec for MoE"""
+    assert num_experts is not None
+
+    mlp = MLPSubmodules(
+        linear_fc1=TEColumnParallelLinear if use_te else ColumnParallelLinear,
+        linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
+    )
+
+    # experts spec
+    if moe_grouped_gemm:
+        ## use GroupedMLP
+        if use_te and TEColumnParallelGroupedLinear is not None and not moe_use_legacy_grouped_gemm:
+            ## use TEGroupedLinear
+            expert_module = TEGroupedMLP
+            expert_submodule = MLPSubmodules(
+                linear_fc1=TEColumnParallelGroupedLinear, linear_fc2=TERowParallelGroupedLinear
+            )
+        else:
+            ## use legacy GroupedMLP
+            expert_module = GroupedMLP
+            expert_submodule = None
+            warnings.warn(
+                'The legacy GroupedMLP will be deprecated in Megatron-Core v0.12.0. '
+                'Please update the TransformerEngine to version>=1.7.0 and use TEGroupedMLP.'
+            )
+    else:
+        ## use SequentialMLP
+        expert_module = SequentialMLP
+        if use_te and not is_te_min_version("1.7.0.dev0"):
+            warnings.warn(
+                "Only transformer-engine>=1.7.0 supports MoE experts, "
+                f"but your version is {get_te_version()}. Use local linear implementation instead."
+            )
+            expert_submodule = MLPSubmodules(
+                linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
+            )
+        else:
+            expert_submodule = mlp
+
+    experts = ModuleSpec(module=expert_module, submodules=expert_submodule)
+
+    # shared experts spec
+    shared_experts = ModuleSpec(module=SharedExpertMLP, params={"gate": False}, submodules=mlp)
+
+    # MoE module spec
+    moe_module_spec = ModuleSpec(
+        module=MoELayer, submodules=MoESubmodules(experts=experts, shared_experts=shared_experts)
+    )
+    return moe_module_spec

megatron/core/models/multimodal/llava_model.py

@@ -36,6 +36,7 @@ IGNORE_INDEX = -100  # ID for labels that should be ignored.
 # Image token index can be tokenizer dependent so the default value does not work in all cases.
 DEFAULT_IMAGE_TOKEN_INDEX = -200
 IMAGE_TOKEN = "<image>"
+VIDEO_TOKEN = "<video>"
 
 
 # Note: This is under development and may be missing features.
@@ -828,7 +829,7 @@ class LLaVAModel(MegatronModule):
             ).contiguous()  # [b, text_seq_len, h_language]
 
         # Assume 1 tile per image if the number of tiles is not provided.
-        if num_image_tiles is None:
+        if num_image_tiles is None and images is not None:
             num_image_tiles = torch.ones(images.shape[0], dtype=torch.int, device=input_ids.device)
 
         combined_embeddings, new_labels, new_loss_mask = self._preprocess_data(