merge megatron_v0.11.0

dcu_megatron/adaptor/megatron_adaptor.py

@@ -165,13 +165,14 @@ class CoreAdaptation(MegatronAdaptationABC):
 
     def patch_tensor_parallel(self):
         from ..core.tensor_parallel.cross_entropy import VocabParallelCrossEntropy
-        from ..core.tensor_parallel import vocab_parallel_embedding_forward, vocab_parallel_embedding_init
+        from ..core.tensor_parallel import vocab_parallel_embedding_forward, vocab_parallel_embedding_init_wrapper
 
         # VocabParallelEmbedding
         MegatronAdaptation.register('megatron.core.tensor_parallel.layers.VocabParallelEmbedding.forward',
                                     vocab_parallel_embedding_forward)
         MegatronAdaptation.register('megatron.core.tensor_parallel.layers.VocabParallelEmbedding.__init__',
-                                    vocab_parallel_embedding_init)
+                                    vocab_parallel_embedding_init_wrapper,
+                                    apply_wrapper=True)
 
         # VocabParallelCrossEntropy
         MegatronAdaptation.register('megatron.core.tensor_parallel.cross_entropy.VocabParallelCrossEntropy.calculate_predicted_logits',
@@ -233,13 +234,20 @@ class LegacyAdaptation(MegatronAdaptationABC):
         self.patch_legacy_models()
 
     def patch_legacy_models(self):
-        from ..legacy.model.transformer import ParallelMLPPatch, ParallelAttentionPatch
+        from ..legacy.model.transformer import (
+            parallel_mlp_init_wrapper,
+            ParallelAttentionPatch,
+            parallel_attention_init_wrapper
+        )
         from ..legacy.model.utils import get_norm
 
         # ParallecMLP
         MegatronAdaptation.register('megatron.legacy.model.transformer.ParallelMLP.__init__',
-                                    ParallelMLPPatch.__init__)
-
+                                    parallel_mlp_init_wrapper,
+                                    apply_wrapper=True)
+        MegatronAdaptation.register('megatron.legacy.model.transformer.ParallelAttention.__init__',
+                                    parallel_attention_init_wrapper,
+                                    apply_wrapper=True)
         MegatronAdaptation.register('megatron.legacy.model.transformer.ParallelAttention.forward',
                                     ParallelAttentionPatch.forward)
 

dcu_megatron/core/tensor_parallel/layers.py

 import os
+import copy
 import socket
 import warnings
 from functools import wraps
 from typing import Callable, List, Optional
 
-try:
+if int(os.getenv("USE_FLUX_OVERLAP", "0"))
+    try:
         import flux
-except ImportError:
+        from dcu_megatron.core.utils import is_flux_min_version
+    except ImportError:
         raise ImportError("flux is NOT installed")
 
 import torch
@@ -53,6 +56,7 @@ from megatron.core.tensor_parallel.layers import (
 from dcu_megatron.core.utils import is_flux_min_version
 
 
+
 _grad_accum_fusion_available = True
 try:
     import fused_weight_gradient_mlp_cuda
@@ -60,65 +64,30 @@ except ImportError:
     _grad_accum_fusion_available = False
 
 
-def vocab_parallel_embedding_init(
-    self,
-    num_embeddings: int,
-    embedding_dim: int,
-    *,
-    init_method: Callable,
-    reduce_scatter_embeddings: bool = False,
-    config: ModelParallelConfig,
-    skip_weight_param_allocation: bool = False
-):
-    super(VocabParallelEmbedding, self).__init__()
-    # Keep the input dimensions.
-    self.num_embeddings = num_embeddings
-    self.embedding_dim = embedding_dim
-    self.reduce_scatter_embeddings = reduce_scatter_embeddings
-    self.tensor_model_parallel_size = get_tensor_model_parallel_world_size()
-    # Divide the weight matrix along the vocaburaly dimension.
-    (self.vocab_start_index, self.vocab_end_index) = (
-        VocabUtility.vocab_range_from_global_vocab_size(
-            self.num_embeddings,
-            get_tensor_model_parallel_rank(),
-            self.tensor_model_parallel_size,
-        )
-    )
-    self.num_embeddings_per_partition = self.vocab_end_index - self.vocab_start_index
-    self.deterministic_mode = config.deterministic_mode
-
-    # Allocate weights and initialize.
-    if not skip_weight_param_allocation:
-        if config.use_cpu_initialization:
-            self.weight = Parameter(
-                torch.empty(
-                    self.num_embeddings_per_partition, self.embedding_dim, dtype=config.params_dtype
-                )
-            )
-            if config.perform_initialization:
-                _initialize_affine_weight_cpu(
-                    self.weight,
-                    self.num_embeddings,
-                    self.embedding_dim,
-                    self.num_embeddings_per_partition,
-                    0,
-                    init_method,
-                    params_dtype=config.params_dtype,
-                )
-        else:
-            self.weight = Parameter(
-                torch.empty(
-                    self.num_embeddings_per_partition,
-                    self.embedding_dim,
-                    device=torch.cuda.current_device(),
-                    dtype=config.params_dtype,
-                )
-            )
-            if config.perform_initialization:
-                _initialize_affine_weight_gpu(self.weight, init_method, partition_dim=0, stride=1)
-    else:
+def vocab_parallel_embedding_init_wrapper(fn):
+    @wraps(fn)
+    def wrapper(self,
+                *args,
+                skip_weight_param_allocation: bool = False,
+                **kwargs
+        ):
+
+        if (
+            skip_weight_param_allocation
+            and "config" in kwargs
+            and hasattr(kwargs["config"], "perform_initialization")
+        ):
+            config = copy.deepcopy(kwargs["config"])
+            config.perform_initialization = False
+            kwargs["config"] = config
+
+        fn(self, *args, **kwargs)
+
+        if skip_weight_param_allocation:
             self.weight = None
 
+    return wrapper
+
 
 @torch.compile(mode='max-autotune-no-cudagraphs')
 def vocab_parallel_embedding_forward(self, input_, weight=None):

dcu_megatron/core/transformer/mtp/multi_token_predictor.py

@@ -16,6 +16,7 @@ from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross
 
 from megatron.core.transformer import ModuleSpec, TransformerConfig, build_module
 from ...tensor_parallel.random import CheckpointWithoutOutput
+from ...tensor_parallel import FluxColumnParallelLinear
 
 
 @dataclass

dcu_megatron/legacy/model/rms_norm.py

-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
-
 import torch
-from torch import nn
-
-class RMSNorm(torch.nn.Module):
-
-    def __init__(self,
-                 dim: int,
-                 eps: float = 1e-6,
-                 sequence_parallel: bool = False,
-                 config: dict = None):
-        """RMS Normaliation module
-
-        Args:
-            dim (int): The width of input, i.e. hidden size
-            eps (float): epsilon to use for the norm, default to 1e-6
-            sequence_parallel (bool): Set to true if sequence parallelism is being used,
-              this marks the weights as needing to be allreduced.
-        """
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-        setattr(self.weight, 'sequence_parallel', sequence_parallel)
-
-    def _norm(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    @torch.compile(mode="max-autotune-no-cudagraphs")
-    def forward(self, x):
-        output = self._norm(x.float()).type_as(x)
-        return output * self.weight
-
-import torch
-
-from typing import Any, Callable, Dict, Optional, Tuple, Union
-import lightop # rmsnorm_forward,rmsnorm_backward
+from typing import Optional
+import lightop
 
 from functools import partial
 from megatron.core.utils import is_torch_min_version

dcu_megatron/legacy/model/transformer.py

 import torch
 import torch.nn.functional as F
 
+from functools import wraps
+
 from megatron.training import get_args
 from megatron.core import tensor_parallel
 from megatron.legacy.model.enums import AttnType
 from megatron.core.models.common.embeddings import apply_rotary_pos_emb
 from megatron.legacy.model.module import MegatronModule
-from megatron.legacy.model.transformer import ParallelMLP
-from megatron.legacy.model.utils import (
-    erf_gelu,
-    openai_gelu,
-)
 
 try:
     from einops import rearrange
 except ImportError:
     rearrange = None
 
-class ParallelMLPPatch(MegatronModule):
-    """MLP.
+try: # 使用定长fa
+    from flash_attn import flash_attn_func
+except ImportError:
+    flash_attn_func = None
 
-    MLP will take the input with h hidden state, project it to 4*h
-    hidden dimension, perform nonlinear transformation, and project the
-    state back into h hidden dimension.
-    """
-    def __init__(self, config, is_expert=False):
-        super(ParallelMLP, self).__init__()
-        args = get_args()
+try:
+    from einops import rearrange
+except ImportError:
+    rearrange = None
 
-        self.add_bias = config.add_bias_linear
-
-        ffn_hidden_size = config.ffn_hidden_size
-        if config.gated_linear_unit:
-            ffn_hidden_size *= 2
-
-        # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
-        self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
-            config.hidden_size,
-            ffn_hidden_size,
-            config=config,
-            init_method=config.init_method,
-            bias=self.add_bias,
-            gather_output=False,
-            skip_bias_add=True,
-            is_expert=is_expert,
-        )
 
-        self.bias_gelu_fusion = False
-        self.activation_func = None
-        self.swiglu = args.swiglu
+def parallel_mlp_init_wrapper(fn):
+    @wraps(fn)
+    def wrapper(self, *args, **kwargs):
+        fn(self, *args, **kwargs)
 
-        if args.openai_gelu:
-            self.activation_func = openai_gelu
-        elif args.onnx_safe:
-            self.activation_func = erf_gelu
-        elif args.swiglu:
+        args = get_args()
+        if args.swiglu:
             @torch.compile(mode="max-autotune-no-cudagraphs")
             def swiglu(x):
                 x = torch.chunk(x, 2, dim=-1)
                 return F.silu(x[0]) * x[1]
             self.activation_func = swiglu
-        elif args.squared_relu:
-            def squared_relu(x):
-                return torch.pow(F.relu(x), 2)
-            self.activation_func = squared_relu
-        else:
-            self.bias_gelu_fusion = args.bias_gelu_fusion
-            self.activation_func = F.gelu
-
-        # Project back to h.
-        self.dense_4h_to_h = tensor_parallel.RowParallelLinear(
-            config.ffn_hidden_size,
-            config.hidden_size,
-            config=config,
-            init_method=config.output_layer_init_method,
-            bias=self.add_bias,
-            skip_bias_add=True,
-            input_is_parallel=True,
-            is_expert=is_expert,
+
+    return wrapper
+
+
+class FlashFixedSelfAttention(torch.nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0,
+                 device=None, dtype=None):
+        super().__init__()
+        assert flash_attn_func is not None, ('Please install FlashAttention first, '
+                                                      'e.g., with pip install flash-attn')
+        assert rearrange is not None, 'Please install einops first, e.g., with pip install einops'
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+        self.flash_attn_func = flash_attn_func
+
+    def forward(self, q, k, v):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            q, k, v: The tensor containing the query, key, and value. (B, S, H, D)
+        """
+
+        assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q,k,v)))
+        assert all((i.is_cuda for i in (q,k,v)))
+
+        output = self.flash_attn_func(q, k, v, dropout_p=self.dropout_p, softmax_scale=self.softmax_scale, causal=self.causal)
+        # [b,s,a,dim]
+        return output
+
+
+def parallel_attention_init_wrapper(fn):
+    @wraps(fn)
+    def wrapper(self, *args, **kwargs):
+        fn(self, *args, **kwargs)
+
+        if self.use_flash_attn:
+            self.core_attention_flash = FlashFixedSelfAttention(
+                causal=True, attention_dropout=self.config.attention_dropout
             )
 
+    return wrapper
+
 
 class ParallelAttentionPatch(MegatronModule):
     """Parallel self-attention layer abstract class.
@@ -87,6 +97,7 @@ class ParallelAttentionPatch(MegatronModule):
     Self-attention layer takes input with size [s, b, h]
     and returns output of the same size.
     """
+
     def forward(self, hidden_states, attention_mask,
                 encoder_output=None, inference_params=None,
                 rotary_pos_emb=None):

dcu_megatron/legacy/model/utils.py

 from megatron.training import get_args
-from megatron.legacy.model import LayerNorm
-from .rms_norm import RMSNorm, LightopRMSNorm
+from megatron.legacy.model import LayerNorm, RMSNorm
+from .rms_norm import LightopRMSNorm
 
 
 def get_norm(config):