import logging

import warnings

# May help avoid undefined symbol errors https://pytorch.org/cppdocs/notes/faq.html#undefined-symbol-errors-from-pytorch-aten

import torch

__all__ = ["amp", "fp16_utils", "optimizers", "normalization", "transformer"]

if torch.distributed.is_available():

from . import parallel

__all__.append("parallel")

from . import amp

from . import fp16_utils

# For optimizers and normalization there is no Python fallback.

# Absence of cuda backend is a hard error.

# I would like the errors from importing fused_adam_cuda or fused_layer_norm_cuda

# to be triggered lazily, because if someone has installed with --cpp_ext and --cuda_ext

# so they expect those backends to be available, but for some reason they actually aren't

# available (for example because they built improperly in a way that isn't revealed until

# load time) the error message is timely and visible.

from . import optimizers

from . import normalization

# Logging utilities for apex.transformer module

class RankInfoFormatter(logging.Formatter):

def format(self, record):

from apex.transformer.parallel_state import get_rank_info

record.rank_info = get_rank_info()

return super().format(record)

_library_root_logger = logging.getLogger(__name__)

handler = logging.StreamHandler()

handler.setFormatter(RankInfoFormatter("%(asctime)s - PID:%(process)d - rank:%(rank_info)s - %(filename)s:%(lineno)d - %(levelname)s - %(message)s", "%y-%m-%d %H:%M:%S"))

_library_root_logger.addHandler(handler)

_library_root_logger.propagate = False

def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:

cudnn_available = torch.backends.cudnn.is_available()

cudnn_version = torch.backends.cudnn.version() if cudnn_available else None

if not (cudnn_available and (cudnn_version >= required_cudnn_version)):

warnings.warn(

f"`{global_option}` depends on cuDNN {required_cudnn_version} or later, "

f"but {'cuDNN is not available' if not cudnn_available else cudnn_version}"

)

return False

return True

class DeprecatedFeatureWarning(FutureWarning):

pass

def deprecated_warning(msg: str) -> None:

if (

not torch.distributed.is_available

or not torch.distributed.is_initialized()

or (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0)

):

warnings.warn(msg, DeprecatedFeatureWarning)