__init__.py — Megatron 引擎模块初始化

文件概述

导出 MegatronEngine 和 MegatronEngineWithLMHead。包含一个 HACK：在非 CUDA 环境下设置 TORCH_CUDA_ARCH_LIST 避免 JIT 编译错误。

if not is_cuda_available and "TORCH_CUDA_ARCH_LIST" not in os.environ:
    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0"

from .transformer_impl import MegatronEngine, MegatronEngineWithLMHead

if not is_cuda_available:
    del os.environ["TORCH_CUDA_ARCH_LIST"]

小结

处理了 Megatron 模块在 CPU 环境导入时的兼容性问题。