[dont merge]. Diff 0.2 -> 0.3

.ci/docker/build.sh

@@ -13,14 +13,20 @@ shift
 echo "Building ${IMAGE_NAME} Docker image"
 
 OS=ubuntu
-OS_VERSION=20.04
 CLANG_VERSION=""
-PYTHON_VERSION=3.11
+PYTHON_VERSION=3.12
 MINICONDA_VERSION=24.3.0-0
 
 case "${IMAGE_NAME}" in
   torchtitan-ubuntu-20.04-clang12)
+    OS_VERSION=20.04
     CLANG_VERSION=12
+    BASE_IMAGE=nvidia/cuda:12.4.1-cudnn-runtime-ubuntu${OS_VERSION}
+    ;;
+  torchtitan-rocm-ubuntu-22.04-clang12)
+    OS_VERSION=22.04
+    CLANG_VERSION=12
+    BASE_IMAGE=rocm/dev-ubuntu-${OS_VERSION}:latest
     ;;
   *)
     echo "Invalid image name ${IMAGE_NAME}"
@@ -30,6 +36,7 @@ esac
 docker build \
   --no-cache \
   --progress=plain \
+  --build-arg "BASE_IMAGE=${BASE_IMAGE}" \
   --build-arg "OS_VERSION=${OS_VERSION}" \
   --build-arg "CLANG_VERSION=${CLANG_VERSION}" \
   --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
@@ -38,3 +45,4 @@ docker build \
   -f "${OS}"/Dockerfile \
   "$@" \
   .
+

.ci/docker/common/install_conda.sh

@@ -42,6 +42,7 @@ install_pip_dependencies() {
   pip_install -r /opt/conda/requirements-dev.txt
   pip_install -r /opt/conda/requirements.txt
   pip_install -r /opt/conda/requirements-flux.txt
+  pip_install -r /opt/conda/requirements-vlm.txt
   popd
 }
 

.ci/docker/requirements-dev.txt

@@ -2,5 +2,4 @@ expecttest==0.1.6
 pytest==7.3.2
 pytest-cov
 pre-commit
-tomli-w >= 1.1.0
 transformers

.ci/docker/requirements-vlm.txt

@@ -0,0 +1,2 @@
+einops
+pillow

.ci/docker/requirements.txt

@@ -1,10 +1,12 @@
 torchdata >= 0.8.0
 datasets >= 3.6.0
 tomli >= 1.1.0 ; python_version < "3.11"
+tomli-w >= 1.1.0
 tensorboard
 tabulate
 wandb
 fsspec
 tyro
 tokenizers >= 0.15.0
 safetensors
+psutil

.ci/docker/ubuntu/Dockerfile

@@ -1,6 +1,6 @@
-ARG OS_VERSION
+ARG BASE_IMAGE
 
-FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu${OS_VERSION}
+FROM ${BASE_IMAGE}
 
 ARG OS_VERSION
 
@@ -32,10 +32,11 @@ ENV PATH /opt/conda/envs/py_$PYTHON_VERSION/bin:/opt/conda/bin:$PATH
 COPY requirements-dev.txt /opt/conda/
 COPY requirements.txt /opt/conda/
 COPY requirements-flux.txt /opt/conda/
+COPY requirements-vlm.txt /opt/conda/
 COPY conda-env-ci.txt /opt/conda/
 COPY ./common/install_conda.sh install_conda.sh
 COPY ./common/utils.sh utils.sh
-RUN bash ./install_conda.sh && rm install_conda.sh utils.sh /opt/conda/requirements-dev.txt /opt/conda/requirements.txt /opt/conda/requirements-flux.txt /opt/conda/conda-env-ci.txt
+RUN bash ./install_conda.sh && rm install_conda.sh utils.sh /opt/conda/requirements-dev.txt /opt/conda/requirements.txt /opt/conda/requirements-flux.txt /opt/conda/requirements-vlm.txt /opt/conda/conda-env-ci.txt
 
 USER ci-user
 CMD ["bash"]

.github/CODEOWNERS

@@ -10,4 +10,4 @@
 /torchtitan/experiments/
 
 # codeowners for experiments/forge
-/torchtitan/experiments/forge/* @ebsmothers @pbontrager @joecummings @allenwang28 @tianyu-l @wwwjn
+/torchtitan/experiments/forge/* @ebsmothers @pbontrager @joecummings @allenwang28 @tianyu-l @wwwjn @fegin