megatron-core 0.16.0rc0.dev127378__tar.gz → 0.16.0rc0.dev128858__tar.gz

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev127378
+Version: 0.16.0rc0.dev128858
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/datasets/megatron_tokenizer.py

@@ -1,11 +1,14 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 import json
+import logging
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Any
 
 import numpy
 
+logger = logging.getLogger(__name__)
+
 
 class MegatronLegacyTokenizer(ABC):
     """Abstract class for tokenizer
@@ -20,6 +23,12 @@ class MegatronLegacyTokenizer(ABC):
     """
 
     def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any):
+        # Deprecation warning
+        logger.warning(
+            "You’re using the legacy tokenizer system, which is deprecated "
+            "and will be removed in a future release. Please migrate to the new tokenizer system "
+            "(`megatron.core.tokenizers.MegatronTokenizer`)."
+        )
         self.unique_identifiers = OrderedDict()
         self.unique_identifiers["class"] = type(self).__name__
         self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths)

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/fp8_utils.py

@@ -10,6 +10,12 @@ from typing import List, Optional
 import torch
 
 from megatron.core.enums import Fp8Recipe
+from megatron.core.tensor_parallel import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    gather_from_sequence_parallel_region,
+    reduce_scatter_to_sequence_parallel_region,
+)
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import get_te_version, is_te_min_version
 
@@ -112,6 +118,27 @@ def get_fp8_align_size(fp8_recipe: Fp8Recipe) -> int:
         return 16
 
 
+def is_column_parallel_linear(module):
+    """Returns whether the given module is a ColumnParallelLinear layer."""
+    if HAVE_TE and (
+        isinstance(module, TEColumnParallelLinear)
+        or isinstance(module, TELayerNormColumnParallelLinear)
+    ):
+        return True
+    elif isinstance(module, ColumnParallelLinear):
+        return True
+    return False
+
+
+def is_row_parallel_linear(module):
+    """Returns whether the given module is a RowParallelLinear layer."""
+    if HAVE_TE and isinstance(module, TERowParallelLinear):
+        return True
+    elif isinstance(module, RowParallelLinear):
+        return True
+    return False
+
+
 """
 The code below abstracts the functionalities needed for implementing "--fp8-param-gather" into
 several functions. It provides different implementations for each function based on different
@@ -587,6 +614,18 @@ if HAVE_TE:
             if not FP8GlobalStateManager.is_fp8_enabled():
                 return original_forward(input_tensor, *args, **kwargs)
 
+            # With sequence parallelism we need to all-gather before padding
+            # and reduce-scatter after unpadding
+            if is_sequence_parallel := getattr(module, "sequence_parallel", False):
+                if is_column_parallel_linear(module):
+                    input_tensor = gather_from_sequence_parallel_region(
+                        input_tensor, group=module.tp_group
+                    )
+
+                # Disable sequence parallelism on the module because we are handling the
+                # all-gather and reduce-scatter externally
+                module.sequence_parallel = False
+
             seq_len, batch_size, hidden_size = input_tensor.shape
             # Reshape to (S, B*H) to pad sequence dimension
             input_2d = input_tensor.reshape(seq_len, -1)
@@ -612,6 +651,16 @@ if HAVE_TE:
             unpadded_output_2d = _unpad_func(output_2d, [seq_len])
             unpadded_output = unpadded_output_2d.reshape(seq_len, batch_size, output_hidden_size)
 
+            if is_sequence_parallel:
+                # Reduce-scatter after unpadding
+                if is_row_parallel_linear(module):
+                    unpadded_output = reduce_scatter_to_sequence_parallel_region(
+                        unpadded_output, group=module.tp_group
+                    )
+
+                # Reset sequence parallelism flag on the module
+                module.sequence_parallel = True
+
             if other_outputs:
                 return (unpadded_output,) + other_outputs
             else:

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 16
 PATCH = 0
-PRE_RELEASE = 'rc0.dev127378'
+PRE_RELEASE = 'rc0.dev128858'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/safe_globals.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 from argparse import Namespace
+from io import BytesIO
 from pathlib import PosixPath
 from types import SimpleNamespace
 
@@ -26,6 +27,7 @@ SAFE_GLOBALS = [
     RerunDiagnostic,
     RerunMode,
     RerunState,
+    BytesIO,
 ]
 
 

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/transformer/dot_product_attention.py

@@ -126,6 +126,8 @@ class DotProductAttention(MegatronModule):
                     )
                 ),
             )
+            if config.perform_initialization:
+                self.softmax_offset = config.init_method(self.softmax_offset)
         else:
             raise ValueError("Softmax type not supported")
 

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron/core/transformer/moe/router.py

@@ -66,6 +66,8 @@ class Router(ABC, MegatronModule):
         """Reset the router parameters."""
         if self.config.perform_initialization:
             self.config.init_method(self.weight)
+            if self.bias is not None:
+                self.config.init_method(self.bias)
         self.weight.data = self.weight.data.to(dtype=self.config.params_dtype)
         setattr(self.weight, 'sequence_parallel', self.config.sequence_parallel)
         if self.bias is not None:

{megatron_core-0.16.0rc0.dev127378 → megatron_core-0.16.0rc0.dev128858}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev127378
+Version: 0.16.0rc0.dev128858
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>