ci: check megatron.training imports in installation test

.github/workflows/install-test.yml

@@ -77,6 +77,12 @@ jobs:
           package-name: megatron.core
           python-binary: ${{ env.UV_PROJECT_ENVIRONMENT }}/bin/python
 
+      - name: Check imports for megatron.training
+        uses: ./FW-CI-templates/.github/actions/check-imports
+        with:
+          package-name: megatron.training
+          python-binary: ${{ env.UV_PROJECT_ENVIRONMENT }}/bin/python
+
   uv-test-pytorch:
     needs: [pre-flight]
     if: |

megatron/training/config/container.py

@@ -6,7 +6,12 @@
 from dataclasses import is_dataclass
 from typing import Any, Type, TypeVar
 
-import yaml
+try:
+    import yaml
+
+    HAVE_YAML = True
+except ImportError:
+    HAVE_YAML = False
 
 from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
 from megatron.core.msc_utils import MultiStorageClientFeature
@@ -94,6 +99,12 @@ def from_yaml(cls: Type[T], yaml_path: str, mode: InstantiationMode = Instantiat
         Returns:
             A new instance of this class initialized with the YAML file values
         """
+        if not HAVE_YAML:
+            raise ImportError(
+                "PyYAML is required to load a config from YAML. "
+                "Install via `pip install pyyaml`."
+            )
+
         from omegaconf import OmegaConf
 
         if MultiStorageClientFeature.is_enabled():
@@ -197,6 +208,12 @@ def to_yaml(self, yaml_path: str) -> None:
         Args:
             yaml_path: Path where to save the YAML file.
         """
+        if not HAVE_YAML:
+            raise ImportError(
+                "PyYAML is required to save a config to YAML. "
+                "Install via `pip install pyyaml`."
+            )
+
         config_dict = self.to_dict()
 
         with safe_yaml_representers():
@@ -212,6 +229,12 @@ def print_yaml(self) -> None:
         """
         Print the config container to the console in YAML format.
         """
+        if not HAVE_YAML:
+            raise ImportError(
+                "PyYAML is required to print a config as YAML. "
+                "Install via `pip install pyyaml`."
+            )
+
         config_dict = self.to_dict()
         with safe_yaml_representers():
             print(yaml.safe_dump(config_dict, default_flow_style=False))

megatron/training/config/yaml_utils.py

@@ -6,7 +6,12 @@
 from contextlib import contextmanager
 from typing import Generator
 
-import yaml
+try:
+    import yaml
+
+    HAVE_YAML = True
+except ImportError:
+    HAVE_YAML = False
 
 
 @contextmanager
@@ -22,6 +27,12 @@ def safe_yaml_representers() -> Generator[None, None, None]:
         with safe_yaml_representers():
             yaml_str = yaml.safe_dump(my_complex_object)
     """
+    if not HAVE_YAML:
+        raise ImportError(
+            "PyYAML is required to register YAML representers. "
+            "Install via `pip install pyyaml`."
+        )
+
     # Save original representers
     original_representers = yaml.SafeDumper.yaml_representers.copy()
     original_multi_representers = yaml.SafeDumper.yaml_multi_representers.copy()

megatron/training/distillation/logits_saver.py

@@ -35,7 +35,13 @@
 
 import torch
 import torch.distributed as dist
-import zstandard
+
+try:
+    import zstandard
+
+    HAVE_ZSTANDARD = True
+except ImportError:
+    HAVE_ZSTANDARD = False
 
 from megatron.core import parallel_state
 from megatron.core.models.common.language_module.language_module import LanguageModule
@@ -579,6 +585,12 @@ def _write_batched_tar(
             # NOTE: MSC is not enabled in the async saving process by default.
             MultiStorageClientFeature.enable()
 
+        if not HAVE_ZSTANDARD:
+            raise ImportError(
+                "zstandard is required to write batched logit tars. "
+                "Install via `pip install zstandard`."
+            )
+
         storage_makedirs(os.path.dirname(tar_path), exist_ok=True)
         write_path = tar_path if is_remote_storage_path(tar_path) else f"{tar_path}.tmp"
         compressor = zstandard.ZstdCompressor(level=3)

megatron/training/distillation/utils_logits.py

@@ -24,7 +24,13 @@
 import torch
 import torch.distributed as dist
 from torch.utils.data import get_worker_info
-import zstandard
+
+try:
+    import zstandard
+
+    HAVE_ZSTANDARD = True
+except ImportError:
+    HAVE_ZSTANDARD = False
 
 from megatron.core.msc_utils import MultiStorageClientFeature
 from megatron.training import get_args
@@ -353,6 +359,11 @@ def iter_logprobs_tar_entries(
 
 def decode_logprobs_payload(data: bytes) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
     """Decode one zstd-compressed cached-logits payload."""
+    if not HAVE_ZSTANDARD:
+        raise ImportError(
+            "zstandard is required to decode cached-logits payloads. "
+            "Install via `pip install zstandard`."
+        )
     data = zstandard.ZstdDecompressor().decompress(data)
     tensors = torch.load(io.BytesIO(data), weights_only=True)
     indices_list = [

megatron/training/yaml_arguments.py

@@ -9,7 +9,13 @@
 import re
 import torch
 import types
-import yaml
+
+try:
+    import yaml
+
+    HAVE_YAML = True
+except ImportError:
+    HAVE_YAML = False
 
 from itertools import chain, starmap
 from types import SimpleNamespace
@@ -28,8 +34,9 @@ def env_constructor(loader, node):
         assert os.environ.get(group) is not None, f"environment variable {group} in yaml not found"
         value = value.replace(f"${{{group}}}", os.environ.get(group))
     return value
-yaml.add_implicit_resolver("!pathex", env_pattern)
-yaml.add_constructor("!pathex", env_constructor)
+if HAVE_YAML:
+    yaml.add_implicit_resolver("!pathex", env_pattern)
+    yaml.add_constructor("!pathex", env_constructor)
 
 
 str_dtype_to_torch = {
@@ -428,6 +435,11 @@ def squared_relu(x):
 
 def load_yaml(yaml_path):
     print(f"warning using experimental yaml arguments feature, argparse arguments will be ignored")
+    if not HAVE_YAML:
+        raise ImportError(
+            "PyYAML is required to load YAML arguments. "
+            "Install via `pip install pyyaml`."
+        )
     with open(yaml_path, "r") as f:
         config = yaml.safe_load(f)
         # Convert to nested namespace