FROM public.ecr.aws/w6p6i9i7/aws-efa-nccl-rdma:22.03-pt-py3

###################################################
## Install Megatron
RUN apt update && apt install ninja-build python3-pybind11 python3-dev -y 
RUN git clone -b v3.0.2 https://github.com/NVIDIA/Megatron-LM.git /megatron