Update

[ghstack-poisoned]

Author: vmoens

.github/unittest/llm/scripts_llm/environment.yml renamed to .github/unittest/linux_libs/scripts_llm/environment.yml

@@ -23,4 +23,14 @@ lib_dir="${env_dir}/lib"
 
 conda deactivate && conda activate ./env
 
-pytest test/llm -vvv --instafail --durations 600 --capture no --error-for-skips
+python -c "import transformers, datasets"
+
+pytest test/test_rlhf.py --instafail -v --durations 200 --capture no --error-for-skips
+
+python examples/rlhf/train_rlhf.py \
+  sys.device=cuda:0 sys.ref_device=cuda:0 \
+  model.name_or_path=gpt2 train.max_epochs=2 \
+  data.batch_size=2 train.ppo.ppo_batch_size=2 \
+  train.ppo.ppo_num_epochs=1 reward_model.name_or_path= \
+  train.ppo.episode_length=8 train.ppo.num_rollouts_per_epoch=4 \
+  data.block_size=110 io.logger=csv

.github/unittest/llm/scripts_llm/install.sh renamed to .github/unittest/linux_libs/scripts_llm/install.sh

@@ -6,13 +6,10 @@
 # Do not install PyTorch and torchvision here, otherwise they also get cached.
 
 set -e
-export DEBIAN_FRONTEND=noninteractive
-export TZ=UTC
-apt-get update
-apt-get install -yq --no-install-recommends git cmake
+apt-get update && apt-get upgrade -y && apt-get install -y git cmake
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
-apt-get install -yq --no-install-recommends wget \
+apt-get install -y wget \
     gcc \
     g++ \
     unzip \
@@ -30,10 +27,7 @@ apt-get install -yq --no-install-recommends wget \
     libgles2
 
 # Upgrade specific package
-apt-get install -yq --no-install-recommends --only-upgrade libstdc++6
-
-apt-get clean
-rm -rf /var/lib/apt/lists/*
+apt-get upgrade -y libstdc++6
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 root_dir="$(git rev-parse --show-toplevel)"

.github/unittest/llm/scripts_llm/post_process.sh renamed to .github/unittest/linux_libs/scripts_llm/post_process.sh

@@ -21,15 +21,14 @@ permissions:
 
 jobs:
   unittests:
-    if: ${{ github.event_name == 'push' || (github.event_name == 'pull_request' && contains(join(github.event.pull_request.labels.*.name, ', '), 'llm/')) }}
     strategy:
       matrix:
         python_version: ["3.9"]
         cuda_arch_version: ["12.8"]
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       repository: pytorch/rl
-      runner: "linux.g6.4xlarge.experimental.nvidia.gpu"
+      runner: "linux.g5.4xlarge.nvidia.gpu"
       # gpu-arch-type: cuda
       # gpu-arch-version: "11.7"
       docker-image: "nvidia/cudagl:11.4.0-base"
@@ -51,7 +50,7 @@ jobs:
         export TF_CPP_MIN_LOG_LEVEL=0
         export TD_GET_DEFAULTS_TO_NONE=1
 
-        bash .github/unittest/llm/scripts_llm/setup_env.sh
-        bash .github/unittest/llm/scripts_llm/install.sh
-        bash .github/unittest/llm/scripts_llm/run_test.sh
-        bash .github/unittest/llm/scripts_llm/post_process.sh
+        bash .github/unittest/linux_libs/scripts_llm/setup_env.sh
+        bash .github/unittest/linux_libs/scripts_llm/install.sh
+        bash .github/unittest/linux_libs/scripts_llm/run_test.sh
+        bash .github/unittest/linux_libs/scripts_llm/post_process.sh