fix(packaging): AUT-1323 align dev package and TE revision

examples/moe_recipes/deepseek_v4_flash/gb200/mxfp8_SL4K_128GPU_TP1PP1EP64.yaml

@@ -52,6 +52,7 @@ DEPENDENCIES:
     WORKDIR /home/
     RUN git clone https://github.com/Dao-AILab/fast-hadamard-transform.git && \
         cd fast-hadamard-transform && \
+        git checkout f134af63deb2df17e1171a9ec1ea4a7d8604d5ca && \
         pip install --no-build-isolation .
 
     # Emerging-Optimizers (Muon)
@@ -64,6 +65,7 @@ DEPENDENCIES:
     WORKDIR /opt/
     RUN git clone --branch nv_dev https://github.com/deepseek-ai/FlashMLA.git && \
         cd FlashMLA && \
+        git checkout b7643bd && \
         FLASH_MLA_DISABLE_SM90=1 \
         NVCC_THREADS=16 \
         CFLAGS="-I/usr/local/cuda/include/cccl" \

megatron/post_training/utils.py

@@ -3,10 +3,7 @@
 import inspect
 import os
 
-import modelopt
-import modelopt.torch.quantization as mtq
 import torch
-from modelopt.torch.quantization.utils import is_quantized
 from packaging.version import Version
 
 from megatron.core import parallel_state
@@ -15,6 +12,7 @@
 
 def modelopt_version_higher_than(target_version: str):
     """Check if Model-Optimizer is greater than this version."""
+    import modelopt
     info = "rank {:3}/{:3} checking if nvidia-modelopt-{} is higher than {}".format(
         torch.distributed.get_rank(),
         torch.distributed.get_world_size(),
@@ -27,6 +25,7 @@ def modelopt_version_higher_than(target_version: str):
 
 def modelopt_version_at_least(target_version: str):
     """Check if Model-Optimizer is greater or equal than this version."""
+    import modelopt
     info = "rank {:3}/{:3} checking if nvidia-modelopt-{} is at least {}".format(
         torch.distributed.get_rank(),
         torch.distributed.get_world_size(),
@@ -101,6 +100,8 @@ def _empty_like_if_meta(tensor: torch.Tensor, *, device: torch.device):
 
 
 def print_distributed_quant_summary(model, msg=""):
+    import modelopt.torch.quantization as mtq
+    from modelopt.torch.quantization.utils import is_quantized
     from megatron.core import parallel_state
     from megatron.core.utils import unwrap_model
     from megatron.training import print_rank_0

pyproject.toml

@@ -8,7 +8,14 @@ build-backend = "setuptools.build_meta"
 include-package-data = true
 
 [tool.setuptools.packages.find]
-include = ["megatron.core", "megatron.core.*", "megatron.training", "megatron.training.*"]
+include = [
+    "megatron.core",
+    "megatron.core.*",
+    "megatron.training",
+    "megatron.training.*",
+    "megatron.post_training",
+    "megatron.post_training.*",
+]
 
 [tool.setuptools.dynamic]
 version = { attr = "megatron.core.package_info.__version__" }
@@ -84,7 +91,7 @@ mlm = [
 
 dev = [
     "nvidia-modelopt[torch]; sys_platform != 'darwin'",
-    "transformer-engine[pytorch,core_cu13]>=2.9.0a0,<2.12.0",
+    "transformer-engine[pytorch,core_cu13]>=2.18.0,<2.19.0",
     "nvidia-resiliency-ext==0.6.0",
     "tqdm",
     "einops~=0.8",
@@ -198,7 +205,7 @@ flash_mla = [
     { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" },
 ]
 nvidia-cudnn-frontend = { git = "https://github.com/NVIDIA/cudnn-frontend.git", rev = "0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }
-transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "f031cf87bd054c7558b887df7bed93975456667f" }
+transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "cfda38340769c432c12db3c0df574f8e6ce46c8c" }
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "17ae86b64d7f75653351664f5d8c9e466faede00" }
 emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.3.0" }
 fast-hadamard-transform = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git", rev = "f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }