1、如何使用按照Megatron环境

1、先安装pytorch，注意版本匹配

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

2、安装Megatron-core(原来的megatron-lm 1.1.5)

安装包地址，然后再终端直接 pip install

https://files.pythonhosted.org/packages/ff/17/cf9ab8e7aec4ab89e697e43e52d9801c8b788b79de5c0c810f154d7c0a2f/megatron_core-0.9.0-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

api文档

megatron-core的API文档

3、安装apex，同样的方法

要求NVIDIA cuDNN v8.5 或更高版本

然后再从源安装apex

https://github.com/NVIDIA/apex.git
cd apex

 官网说：Each apex.contrib module requires one or more install options other than --cpp_ext and --cuda_ext. Note that contrib modules do not necessarily support stable PyTorch releases.

不一定支持稳定的 PyTorch 版本，很可能会因为某些版本的torch不支持导致安装失败，注意torch的版本。

一定要注意cuda和torch兼容

为了避免torch和cuda版本兼容的时候还报错

Failed to build installable wheels for some pyproject.toml based projects (apex)，

建议纯 Python 构建，方法是

pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./

4、安装TransformerEngine

git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine
git checkout $te_commit 
git submodule init && git submodule update 
NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi pip install .