megatron-core 0.13.0rc3__tar.gz → 0.13.0rc4__tar.gz

{megatron_core-0.13.0rc3/megatron_core.egg-info → megatron_core-0.13.0rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.13.0rc3
+Version: 0.13.0rc4
 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.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/T5/t5_spec.py

@@ -16,6 +16,8 @@ from megatron.core.transformer.transformer_block import TransformerBlockSubmodul
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/bert/bert_layer_specs.py

@@ -12,6 +12,8 @@ from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEDotProductAttention,
         TELayerNormColumnParallelLinear,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/common/embeddings/rotary_pos_embedding.py

@@ -315,5 +315,5 @@ class MultimodalRotaryEmbedding(nn.Module):
         if parallel_state.get_context_parallel_world_size() > 1:
             # slice rotary_pos_emb along sequence dimension and select the parition of the current
             # CP rank
-            emb = get_pos_emb_on_this_cp_rank(emb, 1)
+            emb = get_pos_emb_on_this_cp_rank(emb, 0, parallel_state.get_context_parallel_group())
         return emb

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/gpt/gpt_layer_specs.py

@@ -34,6 +34,8 @@ from megatron.core.transformer.transformer_layer import (
 )
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
@@ -42,6 +44,8 @@ except ImportError:
     HAVE_TE = False
 
 try:
+    import nvidia_kitchen  # pylint: disable=unused-import
+
     from megatron.core.extensions.kitchen import KitchenSpecProvider
 
     HAVE_KITCHEN = True

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py

@@ -29,6 +29,8 @@ from megatron.core.transformer.transformer_layer import (
 from megatron.core.utils import is_te_min_version
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEDotProductAttention,
         TELayerNormColumnParallelLinear,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/gpt/moe_module_specs.py

@@ -9,6 +9,8 @@ from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
 from megatron.core.transformer.spec_utils import ModuleSpec
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
     HAVE_TE = True

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/retro/decoder_spec.py

@@ -40,6 +40,8 @@ except ImportError:
     HAVE_APEX = False
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/models/retro/encoder_spec.py

@@ -21,6 +21,8 @@ from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 13
 PATCH = 0
-PRE_RELEASE = 'rc3'
+PRE_RELEASE = 'rc4'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/cuda_graphs.py

@@ -23,6 +23,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
     from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, fp8_autocast
     from transformer_engine.pytorch.graph import restore_fp8_tensors, save_fp8_tensors
     from transformer_engine.pytorch.graph import set_capture_end as te_set_capture_end

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/heterogeneous/linear_replacements.py

@@ -16,6 +16,8 @@ from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import divide
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear
 
     HAVE_TE = True
@@ -67,6 +69,7 @@ if HAVE_TE:
             )
 
         def forward(self, x, **kwargs):
+            """Forward of TELayerNormColumnParallelLinearGathered"""
             out, bias = super().forward(x)
             assert bias is None, "bias should be None since we set skip_bias_add=False"
 
@@ -100,6 +103,7 @@ class ColumnParallelLinearGathered(ColumnParallelLinear):
         runtime_gather_output: bool | None = None,
         **kwargs,
     ):
+        """Forward of ColumnParallelLinearGathered"""
         out, bias = super().forward(input_, weight, runtime_gather_output)
         assert bias is None, "bias should be None since we set skip_bias_add=False"
 

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/moe/experts.py

@@ -40,6 +40,7 @@ from megatron.core.transformer.utils import (
 )
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
 
     from megatron.core.extensions.transformer_engine import Fp8Padding, Fp8Unpadding
 

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/moe/moe_layer.py

@@ -21,6 +21,8 @@ from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_config import TransformerConfig
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import te_checkpoint
 
     HAVE_TE = True

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/moe/moe_utils.py

@@ -10,6 +10,8 @@ from megatron.core.process_groups_config import ModelCommProcessGroups
 from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         fused_permute,
         fused_permute_with_probs,

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/multi_token_prediction.py

@@ -32,6 +32,8 @@ SUPPORTED_ATTN_MASK = [
 ]
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
     HAVE_TE = True

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4}/megatron/core/transformer/transformer_block.py

@@ -28,27 +28,38 @@ from megatron.core.transformer.utils import sharded_state_dict_default
 from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
 
 try:
+    import transformer_engine.pytorch as te  # pylint: disable=unused-import
+
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+try:
+    import apex  # pylint: disable=unused-import
+
+    HAVE_APEX = True
+except ImportError:
+    HAVE_APEX = False
+
+get_cpu_offload_context = None
+te_checkpoint = None
+
+if HAVE_TE:
     from megatron.core.extensions.transformer_engine import (
         TENorm,
         get_cpu_offload_context,
         te_checkpoint,
     )
 
-    HAVE_TE = True
     LayerNormImpl = TENorm
-except ImportError:
-    HAVE_TE = False
-    get_cpu_offload_context = None
-
-    try:
-        import apex  # pylint: disable=unused-import
 
-        LayerNormImpl = FusedLayerNorm
+elif HAVE_APEX:
+    LayerNormImpl = FusedLayerNorm
 
-    except ImportError:
-        from megatron.core.transformer.torch_norm import WrappedTorchNorm
+else:
+    from megatron.core.transformer.torch_norm import WrappedTorchNorm
 
-        LayerNormImpl = WrappedTorchNorm
+    LayerNormImpl = WrappedTorchNorm
 
 
 def get_num_layers_to_build(config: TransformerConfig, vp_stage: Optional[int] = None) -> int:

{megatron_core-0.13.0rc3 → megatron_core-0.13.0rc4/megatron_core.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.13.0rc3
+Version: 0.13.0rc4
 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>