diff --git a/.github/workflows/install-test.yml b/.github/workflows/install-test.yml index f340e5aa2d8..3505937cd92 100644 --- a/.github/workflows/install-test.yml +++ b/.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: | diff --git a/megatron/training/config/container.py b/megatron/training/config/container.py index c13f73f52e9..bb1e4f2cf5d 100644 --- a/megatron/training/config/container.py +++ b/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)) diff --git a/megatron/training/config/yaml_utils.py b/megatron/training/config/yaml_utils.py index f088a8ba484..0d26801b6e2 100644 --- a/megatron/training/config/yaml_utils.py +++ b/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() diff --git a/megatron/training/distillation/logits_saver.py b/megatron/training/distillation/logits_saver.py index 33b035e44e2..36e7aacbcba 100644 --- a/megatron/training/distillation/logits_saver.py +++ b/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) diff --git a/megatron/training/distillation/utils_logits.py b/megatron/training/distillation/utils_logits.py index d2441a5a1d0..dd22b2b8b14 100644 --- a/megatron/training/distillation/utils_logits.py +++ b/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 = [ diff --git a/megatron/training/yaml_arguments.py b/megatron/training/yaml_arguments.py index d44f4d31822..93a2a7abc73 100644 --- a/megatron/training/yaml_arguments.py +++ b/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