diff --git a/.github/workflows/_run_test.yml b/.github/workflows/_run_test.yml
index faf1848cad..d75a6a7178 100644
--- a/.github/workflows/_run_test.yml
+++ b/.github/workflows/_run_test.yml
@@ -64,6 +64,24 @@ jobs:
         - name: Docker pull image
           run: |
             docker pull nemoci.azurecr.io/nemo_reinforcer_container:${{ github.run_id }}
+        
+        # NOTE: under certain circumstances, the checkout action cannot clean up the workspace properly, so
+        # this workaround is needed to ensure that the workspace is clean by removing all files created by root.
+        #
+        # The error observed looked like this from the checkout action:
+        #      Run actions/checkout@v4
+        #      ...
+        #      Deleting the contents of '/home/azureuser/actions-runner/_work/reinforcer/reinforcer'
+        #      Error: File was unable to be removed Error: EACCES: permission denied, rmdir '/home/azureuser/actions-runner/_work/reinforcer/reinforcer/docs/_build/doctest'
+        - name: Forcefully clean up the repository
+          run: |
+            docker run --rm -u root \
+              -v /home/azureuser/actions-runner/_work/reinforcer/reinforcer:/home/azureuser/actions-runner/_work/reinforcer/reinforcer \
+              nemoci.azurecr.io/nemo_reinforcer_container:${{ github.run_id }} \
+              bash -x -c "ls -lah /home/azureuser/actions-runner/_work/reinforcer/reinforcer && shopt -s dotglob && rm -rf /home/azureuser/actions-runner/_work/reinforcer/reinforcer/*"
+
+        - name: Checkout repository
+          uses: actions/checkout@v4
 
         - name: Start container
           run: |
@@ -71,8 +89,9 @@ jobs:
               --env TRANSFORMERS_OFFLINE=0 \
               --env HYDRA_FULL_ERROR=1 \
               --env HF_HOME=/home/TestData/reinforcer/hf_home \
-              --env REINFORCER_CI_DIR=/home/TestData/reinforcer \
-              --env REINFORCER_REPO_DIR=/opt/NeMo-Reinforcer \
+              --env REINFORCER_REPO_DIR=/opt/reinforcer \
+              --volume $PWD:/opt/reinforcer \
+              --volume /mnt/datadrive/TestData/reinforcer/datasets:/opt/reinforcer/datasets:ro \
               --volume /mnt/datadrive/TestData/reinforcer/checkpoints:/home/TestData/reinforcer/checkpoints:ro \
               --volume /mnt/datadrive/TestData/reinforcer/hf_home/hub:/home/TestData/reinforcer/hf_home/hub \
               nemoci.azurecr.io/nemo_reinforcer_container:${{ github.run_id }} \
@@ -94,6 +113,9 @@ jobs:
             set -e
 
             cmd=$(cat <<"RUN_TEST_EOF"
+            # This is needed since we create virtualenvs in the workspace, so this allows it to be cleaned up if necessary
+            umask 000
+
             nvidia-smi
 
             # In case git commands need to be run inside Reinforcer
diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index d79d0cf61f..cbb5e8132b 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -102,6 +102,20 @@ jobs:
           pre-commit install
           pre-commit run --all-files --show-diff-on-failure --color=always
   
+  sphinx-build:
+    name: Sphinx build
+    needs: [pre-flight]
+    runs-on: ubuntu-latest
+    if: ${{ needs.pre-flight.outputs.run_ci == 'true' }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: build docs
+        run: |
+          pip install uv
+          cd docs/
+          uv run --extra docs sphinx-build . _build/html
+
   build-container:
     if: ${{ needs.pre-flight.outputs.run_ci == 'true' }}
     needs: [pre-flight]
@@ -115,6 +129,20 @@ jobs:
         MAX_JOBS=32
         REINFORCER_COMMIT=${{ github.sha }}
 
+  sphinx-doctest:
+    name: Sphinx doctest
+    needs: [build-container, pre-flight]
+    uses: ./.github/workflows/_run_test.yml
+    if: ${{ needs.pre-flight.outputs.run_ci == 'true' }}
+    with:
+      RUNNER: self-hosted-azure
+      TIMEOUT: 10
+      SCRIPT: |
+        cd ${REINFORCER_REPO_DIR}/docs
+        uv run --extra docs sphinx-build -b doctest . _build/doctest
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
   unit-tests:
     name: Unit tests
     needs: [build-container, pre-flight]
@@ -125,6 +153,27 @@ jobs:
       TIMEOUT: 10
       SCRIPT: |
         cd ${REINFORCER_REPO_DIR}
-        uv run bash -x ./tests/run_unit.sh
+        uv run --extra test bash -x ./tests/run_unit.sh
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
+  functional-tests:
+    name: ${{ matrix.test_case }}
+    needs: [build-container, pre-flight]
+    uses: ./.github/workflows/_run_test.yml
+    if: ${{ needs.pre-flight.outputs.run_ci == 'true' }}
+    strategy:
+      matrix:
+        test_case:
+          - sft.sh
+          - grpo.sh
+    with:
+      # TODO: For now, allow these to fail since the checks are not robust.
+      OPTIONAL: true
+      RUNNER: self-hosted-azure
+      TIMEOUT: 8
+      SCRIPT: |
+        cd ${REINFORCER_REPO_DIR}
+        uv run bash ./tests/functional/${{ matrix.test_case }}
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
\ No newline at end of file
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 3ac272200e..0abb6ebd9e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,18 +1,41 @@
 ARG BASE_IMAGE=anyscale/ray:2.43.0-py312-cu125
-FROM ${BASE_IMAGE}
+FROM ${BASE_IMAGE} AS base
+# base is just ray + uv with minimal installs so it is a very lightweight container
 
-WORKDIR /opt/NeMo-Reinforcer
+# It is more convenient for users to run as root
+USER root
 
-RUN sudo apt-get update && sudo apt-get install -y jq
+RUN apt-get update && sudo apt-get install -y jq
 
+RUN pip install uv
 RUN echo "unset RAY_RUNTIME_ENV_HOOK" >> /home/ray/.bashrc
 
-COPY pyproject.toml .
+FROM base AS hermetic
+# hermetic creates a virtual environment with the default dependencies pre-installed for convenience
 
-RUN pip install uv && \
-    uv venv -p python3.12 && \
-    uv pip install -r pyproject.toml --extra dev --extra test
+COPY --chown=ray --chmod=755 pyproject.toml /opt/reinforcer/pyproject.toml
+RUN chmod 755 /home/ray/.cache
+WORKDIR /opt/reinforcer
+RUN uv venv .venv
+# uv sync has a more reliable resolver than simple uv pip install which can fail
+RUN uv sync --extra test --extra dev --extra docs --no-install-project
 
-COPY . .
+ENV VIRTUAL_ENV=/opt/reinforcer/.venv
+ENV PATH="/opt/reinforcer/.venv/bin:$PATH"
+# The ray images automatically activate the anaconda venv. We will
+# comment this out of the .bashrc to give the same UX between docker
+# and other clusters like slurm.
+RUN <<"EOF"
+cp ~/.bashrc ~/.bashrc.backup  # backup existing .bashrc
 
-RUN uv pip install -e .
+# Comment out the conda initialize block
+sed -i '/# >>> conda initialize >>>/,/# <<< conda initialize <<</ { /^[^#]/ s/^/# / }' ~/.bashrc
+
+# Comment out any line that explicitly exports the anaconda3 PATH
+sed -i '/export PATH=\$HOME\/anaconda3\/bin:\$PATH/ s/^/# /' ~/.bashrc
+EOF
+
+COPY --chown=ray --chmod=755 . /opt/reinforcer
+RUN uv pip install --no-deps --editable /opt/reinforcer
+
+FROM base
\ No newline at end of file
diff --git a/docs/cluster.md b/docs/cluster.md
index 5d2f5cb4b8..be34d2af1b 100644
--- a/docs/cluster.md
+++ b/docs/cluster.md
@@ -21,7 +21,7 @@ export UV_CACHE_DIR=/path/that/all/workers/can/access/uv_cache
 
 ```sh
 # Run from the root of NeMo-Reinforcer repo
-NUM_ACTOR_NODES=1  # Total nodes requested are $NUM_ACTOR_NODES + 1 (+1 for head node)
+NUM_ACTOR_NODES=1  # Total nodes requested (head is colocated on ray-worker-0)
 
 COMMAND="bash -c 'uv pip install -e .; uv run ./examples/run_grpo.py'" \
 RAY_DEDUP_LOGS=0 \
@@ -55,7 +55,7 @@ tail -f 1980204-logs/ray-driver.log
 To run interactively, launch the same command as the [Batched Job Submission](#batched-job-submission) except omit the `COMMAND` line:
 ```sh
 # Run from the root of NeMo-Reinforcer repo
-NUM_ACTOR_NODES=1  # Total nodes requested are $NUM_ACTOR_NODES + 1 (+1 for head node)
+NUM_ACTOR_NODES=1  # Total nodes requested (head is colocated on ray-worker-0)
 
 RAY_DEDUP_LOGS=0 \
 UV_CACHE_DIR=YOUR_UV_CACHE_DIR \
diff --git a/docs/docker.md b/docs/docker.md
index e270372666..15b897bef2 100644
--- a/docs/docker.md
+++ b/docs/docker.md
@@ -1,7 +1,22 @@
-# Building Docker Image
+# Building Docker Images
 
-## Docker Build
+### Base Image
+If you only need the base image with ray + uv, you can build it like so:
 ```sh
 cd docker/
-docker buildx build -t nemo-reinforcer -f Dockerfile .
+docker buildx build -t reinforcer -f Dockerfile ..
 ```
+
+This is **our recommendation** as it is a small image and allows you to specify your python dependencies at runtime.
+
+### Hermetic Image
+We also provide a way to build the docker image with all of default dependencies to get started.
+```sh
+cd docker/
+docker buildx build --target hermetic -t reinforcer -f Dockerfile ..
+```
+
+This image sets up the python environment for you, so you do not have to use `uv` if you don't need
+any other packages.
+
+This image is useful in situations where you may not have network connectivity to re-download packages.
diff --git a/docs/testing.md b/docs/testing.md
index ec184ab0c0..610067961c 100644
--- a/docs/testing.md
+++ b/docs/testing.md
@@ -1,24 +1,70 @@
-# Testing NeMo-Reinforcer
+# Testing Reinforcer
 
 ## Unit Tests
 
+:::{important}
+Unit tests require 2 GPUs to test the full suite.
+:::
+
 ```sh
+# Install the project and the test dependencies
 uv pip install -e '.[test]'
+
+# Run the unit tests using local GPUs
 uv run bash tests/run_unit.sh 
 ```
 
-### Run Unit Tests Hermetic
+:::{note}
+Tests can also be run on SLURM with `ray.sub`, but note that some tests will be skipped
+due to no GPUs being located on the head node. To run the full suite of tests, please
+launch on a regular GPU allocation.
+:::
 
-If your local environment does not have all the necessary dependencies (e.g., `gcc`, `nvcc`)
-or there is concern that something in your environment may be misconfigured, you can also run
-the tests in docker with this script:
+### Running Unit Tests in a Hermetic Environment
+
+For environments lacking necessary dependencies (e.g., `gcc`, `nvcc`)
+or where environmental configuration may be problematic, tests can be run
+in docker with this script:
 
 ```sh
 CONTAINER=... bash tests/run_unit_in_docker.sh
 ```
 
-The `CONTAINER` can be built by following the instructions [here](docker.md).
+The required `CONTAINER` can be built by following the instructions in the [docker documentation](docker.md).
 
 ## Functional tests
 
-TBD
+:::{important}
+Functional tests may require multiple GPUs to run. See each script to understand the requirements.
+:::
+
+Functional tests are located under `tests/functional/`.
+
+```sh
+# Install the project and the test dependencies
+uv pip install -e '.[test]'
+# Run the functional test for sft
+uv run bash tests/functional/sft.sh
+```
+
+At the end of each functional test, the metric checks will be printed as well as
+whether they pass or fail. Here is an example:
+
+```text
+                              Metric Checks
+┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┓
+┃ Status ┃ Check                          ┃ Value             ┃ Message ┃
+┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━┩
+│ PASS   │ data["train/loss"]["9"] < 1500 │ 817.4517822265625 │         │
+└────────┴────────────────────────────────┴───────────────────┴─────────┘
+```
+
+### Running Functional Tests in a Hermetic Environment
+
+For environments lacking necessary dependencies (e.g., `gcc`, `nvcc`)
+or where environmental configuration may be problematic, tests can be run
+in docker with this script:
+
+```sh
+CONTAINER=... bash run_functional_in_docker.sh functional/sft.sh
+```
diff --git a/examples/run_grpo_math.py b/examples/run_grpo_math.py
index 4e7732def1..1c30f8ea94 100644
--- a/examples/run_grpo_math.py
+++ b/examples/run_grpo_math.py
@@ -170,7 +170,9 @@ def main():
     args, overrides = parse_args()
 
     if not args.config:
-        args.config = os.path.join(os.path.dirname(__file__), "configs", "grpo_math_1B.yaml")
+        args.config = os.path.join(
+            os.path.dirname(__file__), "configs", "grpo_math_1B.yaml"
+        )
 
     config = load_config(args.config)
     print(f"Loaded configuration from: {args.config}")
diff --git a/nemo_reinforcer/algorithms/grpo.py b/nemo_reinforcer/algorithms/grpo.py
index 94b9f29a2b..bc4f96099a 100644
--- a/nemo_reinforcer/algorithms/grpo.py
+++ b/nemo_reinforcer/algorithms/grpo.py
@@ -445,7 +445,9 @@ def grpo_train(
 
     # Run grpo training (single-turn)
     for batch in dataloader:
-        print(f"\n{'=' * 25} Step {step + 1}/{min(len(dataloader), master_config['grpo']['max_num_steps'])} {'=' * 25}")
+        print(
+            f"\n{'=' * 25} Step {step + 1}/{min(len(dataloader), master_config['grpo']['max_num_steps'])} {'=' * 25}"
+        )
 
         with timer.time("total_step_time"):
             # Prepare batch
diff --git a/nemo_reinforcer/utils/logger.py b/nemo_reinforcer/utils/logger.py
index 2286075cc9..ea98451e96 100644
--- a/nemo_reinforcer/utils/logger.py
+++ b/nemo_reinforcer/utils/logger.py
@@ -189,7 +189,30 @@ def log_hyperparams(self, params: Dict[str, Any]) -> None:
 
 
 def flatten_dict(d: Dict[str, Any], sep: str = ".") -> Dict[str, Any]:
-    """Flatten a nested dictionary."""
+    """Flatten a nested dictionary.
+
+    Handles nested dictionaries and lists by creating keys with separators.
+    For lists, the index is used as part of the key.
+
+    Args:
+        d: Dictionary to flatten
+        sep: Separator to use between nested keys
+
+    Returns:
+        Flattened dictionary with compound keys
+
+    Examples:
+        ```{doctest}
+        >>> flatten_dict({"a": 1, "b": {"c": 2}})
+        {'a': 1, 'b.c': 2}
+
+        >>> flatten_dict({"a": [1, 2], "b": {"c": [3, 4]}})
+        {'a.0': 1, 'a.1': 2, 'b.c.0': 3, 'b.c.1': 4}
+
+        >>> flatten_dict({"a": [{"b": 1}, {"c": 2}]})
+        {'a.0.b': 1, 'a.1.c': 2}
+        ```
+    """
     result = {}
 
     def _flatten(d, parent_key=""):
@@ -198,6 +221,13 @@ def _flatten(d, parent_key=""):
 
             if isinstance(value, dict):
                 _flatten(value, new_key)
+            elif isinstance(value, list):
+                for i, item in enumerate(value):
+                    list_key = f"{new_key}{sep}{i}"
+                    if isinstance(item, dict):
+                        _flatten(item, list_key)
+                    else:
+                        result[list_key] = item
             else:
                 result[new_key] = value
 
diff --git a/ray.sub b/ray.sub
index 1d5e00db4b..0fab8a8fd9 100644
--- a/ray.sub
+++ b/ray.sub
@@ -58,12 +58,13 @@ port=41993
 ip_head=$head_node_ip:$port
 
 # First we start the head of the ray cluster on one of the physical nodes
-# In this case we are giving an entire physical node to the ray head node
-# The ray head node is marked by including --head to the ray start command
+# Set GPU/CPU resources to 0 to avoid scheduling on the head node
 head_cmd=$(cat <<EOF
 env
 cat <<EOFINNER | tee /launch-head.sh
 ray start --head \
+--num-cpus=0 \
+--num-gpus=0 \
 --node-ip-address="$head_node_ip" \
 --port=$port \
 --block
@@ -81,10 +82,7 @@ EOF
 )
 srun $COMMON_SRUN_ARGS --container-name=ray-head --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-head.log bash -x -c "$head_cmd" &
 
-# We now need to start the Ray worker nodes
-# We define the following variables to help
-num_ray_worker_nodes=$((SLURM_JOB_NUM_NODES - 1)) # One physical node is for the ray head node, rest are for ray worker nodes
-export NUM_ACTORS=$((gpus_per_node * num_ray_worker_nodes)) # 8 actors per ray worker node (one for each GPU)
+export NUM_ACTORS=$((gpus_per_node * SLURM_JOB_NUM_NODES))
 
 # Start Ray worker nodes
 # We want 1 Ray worker node per physical node
@@ -92,7 +90,7 @@ export NUM_ACTORS=$((gpus_per_node * num_ray_worker_nodes)) # 8 actors per ray w
 # NOTE: Ports start above 20000 since 10001-10257 frequently ran into conflicts
 min_worker_port=20001
 max_worker_port=20257
-for ((i = 1; i <= num_ray_worker_nodes; i++)); do
+for ((i = 0; i < SLURM_JOB_NUM_NODES; i++)); do
   node_i=${nodes_array[$i]}
     
   worker_cmd=$(cat <<EOF
@@ -113,12 +111,15 @@ done
 echo ret_code=\$?
 EOF
 )
-  srun $COMMON_SRUN_ARGS --container-name=ray-worker-$i --exact --nodes=1 --ntasks=1 --cpus-per-task=$((16 * gpus_per_node)) -w "$node_i" -o $LOG_DIR/ray-worker-$i.log bash -x -c "$worker_cmd" &
+  if [[ $i -eq 0 ]]; then
+    OVERLAP_HEAD_AND_WORKER_ARG="--overlap"
+  fi
+  srun $COMMON_SRUN_ARGS ${OVERLAP_HEAD_AND_WORKER_ARG:-} --container-name=ray-worker-$i --exact --nodes=1 --ntasks=1 --cpus-per-task=$((16 * gpus_per_node)) -w "$node_i" -o $LOG_DIR/ray-worker-$i.log bash -x -c "$worker_cmd" &
   sleep 3
 done
 
 echo "[INFO] Querying head node too early can cause pyxis to fail"
-for i in {30..1}; do
+for i in {20..1}; do
   echo "[INFO] Waiting for $i seconds before querying head node..."
   sleep 1
 done
@@ -155,17 +156,18 @@ echo "All workers connected!"
 # This driver process is responsible for launching a job on the Ray cluster
 CONTAINER_CWD=$(scontrol show job $SLURM_JOB_ID --json | jq -r '.jobs[].current_working_directory')
 if [[ -n "$COMMAND" ]]; then
-  srun --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-driver.log $COMMAND
+  srun --no-container-mount-home --gpus=0 --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-driver.log $COMMAND
 else
   echo "[INFO]: Ray Cluster is idled, run this on the slurm head node to get a shell to the head node:"
   cat <<EOF >$SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh
 # No args launches on the head node
-WORKER_NUM=\${1:-0}
-if [[ \$WORKER_NUM -eq 0 ]]; then
-  srun --gres=gpu:8 -A $SLURM_JOB_ACCOUNT -p $SLURM_JOB_PARTITION --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" --jobid $SLURM_JOB_ID --pty bash
+WORKER_NUM=\${1:-}
+if [[ -z "\$WORKER_NUM" ]]; then
+  # Empty means we are on the head node
+  srun --no-container-mount-home --gpus=0 -A $SLURM_JOB_ACCOUNT -p $SLURM_JOB_PARTITION --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" --jobid $SLURM_JOB_ID --pty bash
 else
   nodes_array=($nodes)
-  srun --gres=gpu:8 -A $SLURM_JOB_ACCOUNT -p $SLURM_JOB_PARTITION --overlap --container-name=ray-worker-\$WORKER_NUM --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "\${nodes_array[\$WORKER_NUM]}" --jobid $SLURM_JOB_ID --pty bash
+  srun --no-container-mount-home --gres=gpu:8 -A $SLURM_JOB_ACCOUNT -p $SLURM_JOB_PARTITION --overlap --container-name=ray-worker-\$WORKER_NUM --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "\${nodes_array[\$WORKER_NUM]}" --jobid $SLURM_JOB_ID --pty bash
 fi
 EOF
   chmod +x $SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh
diff --git a/tests/functional/grpo.sh b/tests/functional/grpo.sh
index cfbc1bc712..16f7e7530d 100755
--- a/tests/functional/grpo.sh
+++ b/tests/functional/grpo.sh
@@ -2,6 +2,8 @@
 
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetchs metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
 
 set -eou pipefail
 
@@ -10,11 +12,12 @@ JSON_METRICS=$LOG_DIR/$(basename $0 .sh).json
 RUN_LOG=$LOG_DIR/$(basename $0 .sh).log
 export RAY_DEDUP_LOGS=0
 export UV_CACHE_DIR=$PROJECT_ROOT/uv_cache
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
 
 mkdir -p $LOG_DIR
 
 cd $PROJECT_ROOT
-uv run $PROJECT_ROOT/examples/run_grpo_math.py \
+python -u $PROJECT_ROOT/examples/run_grpo_math.py \
     cluster.gpus_per_node=2 \
     grpo.max_num_steps=10 \
     logger.tensorboard_enabled=true \
@@ -25,12 +28,8 @@ uv run $PROJECT_ROOT/examples/run_grpo_math.py \
     2>&1 | tee $RUN_LOG
 
 cd $SCRIPT_DIR
-uv run json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
-
-uv run check_metrics.py $JSON_METRICS \
-    'data["timing/train/policy_refit"]["10"] < 3.0' \
-    'data["timing/train/total_step_time"]["10"] < 20.0' \
-    'data["timing/validation/generation"]["10"] < 3.0' \
-    'max(data["train/token_mult_prob_error"]) < 1.05' \
-    'data["validation/avg_length"]["10"] < 1024' \
+python json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+python check_metrics.py $JSON_METRICS \
+    'max(data["train/token_mult_prob_error"]) < 1.1' \
 
diff --git a/tests/functional/hello_world_fsdp_llama/__init__.py b/tests/functional/hello_world_fsdp_llama/__init__.py
deleted file mode 100644
index 341a77c5bc..0000000000
--- a/tests/functional/hello_world_fsdp_llama/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
diff --git a/tests/functional/hello_world_fsdp_llama/example.ipynb b/tests/functional/hello_world_fsdp_llama/example.ipynb
deleted file mode 100644
index 4529efa34f..0000000000
--- a/tests/functional/hello_world_fsdp_llama/example.ipynb
+++ /dev/null
@@ -1,350 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
-      "Requirement already satisfied: ipywidgets in /usr/local/lib/python3.10/dist-packages (8.1.5)\n",
-      "Requirement already satisfied: comm>=0.1.3 in /usr/local/lib/python3.10/dist-packages (from ipywidgets) (0.2.2)\n",
-      "Requirement already satisfied: ipython>=6.1.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets) (8.28.0)\n",
-      "Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.10/dist-packages (from ipywidgets) (5.14.3)\n",
-      "Requirement already satisfied: widgetsnbextension~=4.0.12 in /usr/local/lib/python3.10/dist-packages (from ipywidgets) (4.0.13)\n",
-      "Requirement already satisfied: jupyterlab-widgets~=3.0.12 in /usr/local/lib/python3.10/dist-packages (from ipywidgets) (3.0.13)\n",
-      "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (5.1.1)\n",
-      "Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (0.19.1)\n",
-      "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (0.1.7)\n",
-      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (3.0.48)\n",
-      "Requirement already satisfied: pygments>=2.4.0 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (2.18.0)\n",
-      "Requirement already satisfied: stack-data in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (0.6.3)\n",
-      "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (1.2.2)\n",
-      "Requirement already satisfied: typing-extensions>=4.6 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (4.12.2)\n",
-      "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets) (4.9.0)\n",
-      "Requirement already satisfied: parso<0.9.0,>=0.8.3 in /usr/local/lib/python3.10/dist-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets) (0.8.4)\n",
-      "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.10/dist-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets) (0.7.0)\n",
-      "Requirement already satisfied: wcwidth in /usr/local/lib/python3.10/dist-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=6.1.0->ipywidgets) (0.2.13)\n",
-      "Requirement already satisfied: executing>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.1.0)\n",
-      "Requirement already satisfied: asttokens>=2.1.0 in /usr/local/lib/python3.10/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.4.1)\n",
-      "Requirement already satisfied: pure-eval in /usr/local/lib/python3.10/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.2.3)\n",
-      "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from asttokens>=2.1.0->stack-data->ipython>=6.1.0->ipywidgets) (1.16.0)\n",
-      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable.It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n",
-      "\u001b[0m\n",
-      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
-      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
-     ]
-    }
-   ],
-   "source": [
-    "!pip install -U ipywidgets"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import ray"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2025-03-02 23:13:34,875\tINFO worker.py:1654 -- Connecting to existing Ray cluster at address: 10.65.26.15:41993...\n",
-      "2025-03-02 23:13:34,883\tINFO worker.py:1832 -- Connected to Ray cluster. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d16239ac6cf44259a74cf9c02d97c0a6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/html": [
-       "<div class=\"lm-Widget p-Widget lm-Panel p-Panel jp-Cell-outputWrapper\">\n",
-       "    <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
-       "        <div class=\"jp-RenderedHTMLCommon\" style=\"display: flex; flex-direction: row;\">\n",
-       "  <svg viewBox=\"0 0 567 224\" fill=\"none\" xmlns=\"http://www.w3.org/2000/svg\" style=\"height: 3em;\">\n",
-       "    <g clip-path=\"url(#clip0_4338_178347)\">\n",
-       "        <path d=\"M341.29 165.561H355.29L330.13 129.051C345.63 123.991 354.21 112.051 354.21 94.2307C354.21 71.3707 338.72 58.1807 311.88 58.1807H271V165.561H283.27V131.661H311.8C314.25 131.661 316.71 131.501 319.01 131.351L341.25 165.561H341.29ZM283.29 119.851V70.0007H311.82C331.3 70.0007 342.34 78.2907 342.34 94.5507C342.34 111.271 331.34 119.861 311.82 119.861L283.29 119.851ZM451.4 138.411L463.4 165.561H476.74L428.74 58.1807H416L367.83 165.561H380.83L392.83 138.411H451.4ZM446.19 126.601H398L422 72.1407L446.24 126.601H446.19ZM526.11 128.741L566.91 58.1807H554.35L519.99 114.181L485.17 58.1807H472.44L514.01 129.181V165.541H526.13V128.741H526.11Z\" fill=\"var(--jp-ui-font-color0)\"/>\n",
-       "        <path d=\"M82.35 104.44C84.0187 97.8827 87.8248 92.0678 93.1671 87.9146C98.5094 83.7614 105.083 81.5067 111.85 81.5067C118.617 81.5067 125.191 83.7614 130.533 87.9146C135.875 92.0678 139.681 97.8827 141.35 104.44H163.75C164.476 101.562 165.622 98.8057 167.15 96.2605L127.45 56.5605C121.071 60.3522 113.526 61.6823 106.235 60.3005C98.9443 58.9187 92.4094 54.9203 87.8602 49.0574C83.3109 43.1946 81.0609 35.8714 81.5332 28.4656C82.0056 21.0599 85.1679 14.0819 90.4252 8.8446C95.6824 3.60726 102.672 0.471508 110.08 0.0272655C117.487 -0.416977 124.802 1.86091 130.647 6.4324C136.493 11.0039 140.467 17.5539 141.821 24.8501C143.175 32.1463 141.816 39.6859 138 46.0505L177.69 85.7505C182.31 82.9877 187.58 81.4995 192.962 81.4375C198.345 81.3755 203.648 82.742 208.33 85.3976C213.012 88.0532 216.907 91.9029 219.616 96.5544C222.326 101.206 223.753 106.492 223.753 111.875C223.753 117.258 222.326 122.545 219.616 127.197C216.907 131.848 213.012 135.698 208.33 138.353C203.648 141.009 198.345 142.375 192.962 142.313C187.58 142.251 182.31 140.763 177.69 138L138 177.7C141.808 184.071 143.155 191.614 141.79 198.91C140.424 206.205 136.44 212.75 130.585 217.313C124.731 221.875 117.412 224.141 110.004 223.683C102.596 223.226 95.6103 220.077 90.3621 214.828C85.1139 209.58 81.9647 202.595 81.5072 195.187C81.0497 187.779 83.3154 180.459 87.878 174.605C92.4405 168.751 98.9853 164.766 106.281 163.401C113.576 162.035 121.119 163.383 127.49 167.19L167.19 127.49C165.664 124.941 164.518 122.182 163.79 119.3H141.39C139.721 125.858 135.915 131.673 130.573 135.826C125.231 139.98 118.657 142.234 111.89 142.234C105.123 142.234 98.5494 139.98 93.2071 135.826C87.8648 131.673 84.0587 125.858 82.39 119.3H60C58.1878 126.495 53.8086 132.78 47.6863 136.971C41.5641 141.163 34.1211 142.972 26.7579 142.059C19.3947 141.146 12.6191 137.574 7.70605 132.014C2.79302 126.454 0.0813599 119.29 0.0813599 111.87C0.0813599 104.451 2.79302 97.2871 7.70605 91.7272C12.6191 86.1673 19.3947 82.5947 26.7579 81.6817C34.1211 80.7686 41.5641 82.5781 47.6863 86.7696C53.8086 90.9611 58.1878 97.2456 60 104.44H82.35ZM100.86 204.32C103.407 206.868 106.759 208.453 110.345 208.806C113.93 209.159 117.527 208.258 120.522 206.256C123.517 204.254 125.725 201.276 126.771 197.828C127.816 194.38 127.633 190.677 126.253 187.349C124.874 184.021 122.383 181.274 119.205 179.577C116.027 177.88 112.359 177.337 108.826 178.042C105.293 178.746 102.113 180.654 99.8291 183.44C97.5451 186.226 96.2979 189.718 96.3 193.32C96.2985 195.364 96.7006 197.388 97.4831 199.275C98.2656 201.163 99.4132 202.877 100.86 204.32ZM204.32 122.88C206.868 120.333 208.453 116.981 208.806 113.396C209.159 109.811 208.258 106.214 206.256 103.219C204.254 100.223 201.275 98.0151 197.827 96.97C194.38 95.9249 190.676 96.1077 187.348 97.4873C184.02 98.8669 181.274 101.358 179.577 104.536C177.879 107.714 177.337 111.382 178.041 114.915C178.746 118.448 180.653 121.627 183.439 123.911C186.226 126.195 189.717 127.443 193.32 127.44C195.364 127.443 197.388 127.042 199.275 126.259C201.163 125.476 202.878 124.328 204.32 122.88ZM122.88 19.4205C120.333 16.8729 116.981 15.2876 113.395 14.9347C109.81 14.5817 106.213 15.483 103.218 17.4849C100.223 19.4868 98.0146 22.4654 96.9696 25.9131C95.9245 29.3608 96.1073 33.0642 97.4869 36.3922C98.8665 39.7202 101.358 42.4668 104.535 44.1639C107.713 45.861 111.381 46.4036 114.914 45.6992C118.447 44.9949 121.627 43.0871 123.911 40.301C126.195 37.515 127.442 34.0231 127.44 30.4205C127.44 28.3772 127.038 26.3539 126.255 24.4664C125.473 22.5788 124.326 20.8642 122.88 19.4205ZM19.42 100.86C16.8725 103.408 15.2872 106.76 14.9342 110.345C14.5813 113.93 15.4826 117.527 17.4844 120.522C19.4863 123.518 22.4649 125.726 25.9127 126.771C29.3604 127.816 33.0638 127.633 36.3918 126.254C39.7198 124.874 42.4664 122.383 44.1635 119.205C45.8606 116.027 46.4032 112.359 45.6988 108.826C44.9944 105.293 43.0866 102.114 40.3006 99.8296C37.5145 97.5455 34.0227 96.2983 30.42 96.3005C26.2938 96.3018 22.337 97.9421 19.42 100.86ZM100.86 100.86C98.3125 103.408 96.7272 106.76 96.3742 110.345C96.0213 113.93 96.9226 117.527 98.9244 120.522C100.926 123.518 103.905 125.726 107.353 126.771C110.8 127.816 114.504 127.633 117.832 126.254C121.16 124.874 123.906 122.383 125.604 119.205C127.301 116.027 127.843 112.359 127.139 108.826C126.434 105.293 124.527 102.114 121.741 99.8296C118.955 97.5455 115.463 96.2983 111.86 96.3005C109.817 96.299 107.793 96.701 105.905 97.4835C104.018 98.2661 102.303 99.4136 100.86 100.86Z\" fill=\"#00AEEF\"/>\n",
-       "    </g>\n",
-       "    <defs>\n",
-       "        <clipPath id=\"clip0_4338_178347\">\n",
-       "            <rect width=\"566.93\" height=\"223.75\" fill=\"white\"/>\n",
-       "        </clipPath>\n",
-       "    </defs>\n",
-       "  </svg>\n",
-       "</div>\n",
-       "\n",
-       "        <table class=\"jp-RenderedHTMLCommon\" style=\"border-collapse: collapse;color: var(--jp-ui-font-color1);font-size: var(--jp-ui-font-size1);\">\n",
-       "    <tr>\n",
-       "        <td style=\"text-align: left\"><b>Python version:</b></td>\n",
-       "        <td style=\"text-align: left\"><b>3.10.12</b></td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "        <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
-       "        <td style=\"text-align: left\"><b>2.43.0</b></td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "    <td style=\"text-align: left\"><b>Dashboard:</b></td>\n",
-       "    <td style=\"text-align: left\"><b><a href=\"http://127.0.0.1:8265\" target=\"_blank\">http://127.0.0.1:8265</a></b></td>\n",
-       "</tr>\n",
-       "\n",
-       "</table>\n",
-       "\n",
-       "    </div>\n",
-       "</div>\n"
-      ],
-      "text/plain": [
-       "RayContext(dashboard_url='127.0.0.1:8265', python_version='3.10.12', ray_version='2.43.0', ray_commit='ecdcdc6a6e63dc4bcd6ea16aae256ce4d32a7e2c')"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "ray.init()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "16.0"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "avail_workers = ray.cluster_resources()[\"worker_units\"]\n",
-    "avail_workers"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from examples.hello_world_fsdp_llama import train"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['AutoModelForCausalLM',\n",
-       " 'AutoTokenizer',\n",
-       " 'LlamaDecoderLayer',\n",
-       " 'ModelTrainer',\n",
-       " 'NASS',\n",
-       " 'NodeInfo',\n",
-       " 'RayClusterCoordinator',\n",
-       " 'Worker',\n",
-       " 'WorkerGroupResources',\n",
-       " '__builtins__',\n",
-       " '__cached__',\n",
-       " '__doc__',\n",
-       " '__file__',\n",
-       " '__loader__',\n",
-       " '__name__',\n",
-       " '__package__',\n",
-       " '__spec__',\n",
-       " 'dataclass',\n",
-       " 'dist',\n",
-       " 'fully_shard',\n",
-       " 'init_device_mesh',\n",
-       " 'json',\n",
-       " 'os',\n",
-       " 'random',\n",
-       " 'ray',\n",
-       " 'register_fsdp_forward_method',\n",
-       " 'time',\n",
-       " 'torch']"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "dir(train)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Worker node info: worker_node_info=[NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23'), NodeInfo(node_id='9faf44e0e79cf43aef2eec7b5003ba0ec197a5bf7c07c7e4706bde7f', node_rank=0, node_ip='10.65.26.23')]\n",
-      "Num physical nodes: self.num_physical_nodes=1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[36m(ModelTrainer pid=52042, ip=10.65.26.23)\u001b[0m DEBUG: self.logical_gpu_id=2 (process_id, world_size, physical_node_id, physical_node_ip, master_addr, num_workers_per_node)=(1, 8, 0, '10.65.26.23', '10.65.26.23', 8)\n",
-      "\u001b[36m(ModelTrainer pid=52043, ip=10.65.26.23)\u001b[0m DEBUG: self.logical_gpu_id=4 (process_id, world_size, physical_node_id, physical_node_ip, master_addr, num_workers_per_node)=(6, 8, 0, '10.65.26.23', '10.65.26.23', 8)\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m [Rank 0] Loading model meta-llama/Llama-3.2-1B on CPU...\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m [Rank 0] Starting synthetic training...\n"
-     ]
-    }
-   ],
-   "source": [
-    "worker_resources = train.WorkerGroupResources(\n",
-    "    num_nodes=1, num_gpus_per_node=8, num_cpus_per_worker=16\n",
-    ")\n",
-    "coordinator = train.RayClusterCoordinator(train.ModelTrainer, worker_resources)\n",
-    "coordinator.initialize_workers()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "16.0"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# Workers after consuming 1 node should be 16 -> 8\n",
-    "ray.cluster_resources()[\"worker_units\"]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m /usr/local/lib/python3.10/dist-packages/torch/autograd/graph.py:818: UserWarning: cuDNN SDPA backward got grad_output.strides() != output.strides(), attempting to materialize a grad_output with matching strides... (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/cudnn/MHA.cpp:670.)\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m   return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 0, Loss: 30.59472656\n",
-      "\u001b[36m(ModelTrainer pid=52041, ip=10.65.26.23)\u001b[0m [Rank 2] Loading model meta-llama/Llama-3.2-1B on CPU...\u001b[32m [repeated 7x across cluster]\u001b[0m\n",
-      "\u001b[36m(ModelTrainer pid=52041, ip=10.65.26.23)\u001b[0m [Rank 2] Starting synthetic training...\u001b[32m [repeated 7x across cluster]\u001b[0m\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 1, Loss: 0.55084229\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 2, Loss: 3.00805664\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 3, Loss: 2.95410156\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 4, Loss: 2.15722656\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 5, Loss: 1.41015625\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 6, Loss: 0.75201416\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 7, Loss: 0.16345596\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 8, Loss: 0.07264709\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 0, Step: 9, Loss: 0.07691956\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 0, Loss: 0.03055668\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 1, Loss: 0.00250489\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 2, Loss: 0.00558114\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 3, Loss: 0.03055668\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 4, Loss: 0.05177402\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 5, Loss: 0.06747818\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 6, Loss: 0.06152725\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 7, Loss: 0.04326725\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 8, Loss: 0.02659702\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Epoch: 1, Step: 9, Loss: 0.01250291\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m [Rank 0] Testing loss is close to expected loss: 0.012502908706665039\n",
-      "\u001b[36m(ModelTrainer pid=52039, ip=10.65.26.23)\u001b[0m Yay! Loss was close :)\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "[None, None, None, None, None, None, None, None]"
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "coordinator.run(hf_model_name=\"meta-llama/Llama-3.2-1B\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ray.shutdown()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.12"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/tests/functional/hello_world_fsdp_llama/main.py b/tests/functional/hello_world_fsdp_llama/main.py
deleted file mode 100644
index df5b9bf1c9..0000000000
--- a/tests/functional/hello_world_fsdp_llama/main.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import time
-from ray.job_submission import JobSubmissionClient, JobStatus
-
-
-def main() -> None:
-    client = JobSubmissionClient("http://127.0.0.1:8265")
-    print("Connected to head!", flush=True)
-
-    # HACK: for now just
-    def num_ray_nodes_available() -> int:
-        import ray
-
-        ray.init()
-        num_gpus_per_node = 8  # hard coded
-        num_nodes_avail = (
-            int(ray.cluster_resources()["worker_units"]) // num_gpus_per_node
-        )
-        ray.shutdown()
-        return num_nodes_avail
-
-    job_id = client.submit_job(
-        entrypoint="RAY_DEDUP_LOGS=0 python3 tests/functional/hello_world_fsdp_llama/train.py",
-        runtime_env={
-            # TODO: disabling for now since it causes issues if my hf_home is in my working dir and ray
-            # wants to upload it to all workers. you get an error like this:
-            #   2025-03-02 11:16:48,187 WARNING packaging.py:417 -- File /workspace/hf_home/hub/models--meta-llama--Meta-Llama-3-8b/snapshots/8cde5ca8380496c9a6cc7ef3a8b46a0372a1d920/model-00001-of-00004.safetensors is very large (4746.15MiB). Consider adding this file to the 'excludes' list to skip uploading it: `ray.init(..., runtime_env={'excludes': ['/workspace/hf_home/hub/models--meta-llama--Meta-Llama-3-8b/snapshots/8cde5ca8380496c9a6cc7ef3a8b46a0372a1d920/model-00001-of-00004.safetensors']})`
-            # "working_dir": "./",
-            "driver_args": {
-                # Scope each "workergroup"
-                "trainer": {
-                    "resources": {
-                        # TODO: read this in from cli args eventually, but for now just use all available
-                        "num_nodes": num_ray_nodes_available(),
-                        "num_gpus_per_node": 8,
-                        "num_cpus_per_worker": 16,
-                    },
-                    "hf_model_name": "meta-llama/Llama-3.2-1B",
-                }
-            },
-            "env_vars": {
-                # TODO: hardcoded, parametrize
-                "HF_HOME": "/workspace/hf_home",
-            },
-        },
-    )
-
-    print(f"Launched job: {job_id}", flush=True)
-    prev_logs = ""
-    while True:
-        status = client.get_job_status(job_id)
-        if status in {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}:
-            if status in {JobStatus.STOPPED, JobStatus.FAILED}:
-                logs = client.get_job_logs(job_id)
-                print(logs, flush=True)
-            break
-        time.sleep(5)
-        if status == JobStatus.RUNNING:
-            logs = client.get_job_logs(job_id)
-            print(logs[len(prev_logs) :], flush=True)
-            prev_logs = logs
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/functional/hello_world_fsdp_llama/train.py b/tests/functional/hello_world_fsdp_llama/train.py
deleted file mode 100644
index 61c823f37f..0000000000
--- a/tests/functional/hello_world_fsdp_llama/train.py
+++ /dev/null
@@ -1,329 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import json
-import ray
-import os
-import random
-import torch
-import torch.distributed as dist
-from torch.distributed.device_mesh import init_device_mesh
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy as NASS
-from dataclasses import dataclass
-from torch.distributed.fsdp import FullyShardedDataParallel, MixedPrecision
-from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
-from torch.distributed.fsdp import register_fsdp_forward_method
-
-
-@dataclass
-class WorkerGroupResources:
-    num_nodes: int
-    num_cpus_per_worker: int = 16  # 128 hyperthread / 8 gpu = 16 cpu/gpu
-    num_gpus_per_node: int = 8  # will always be true on slurm
-
-
-@dataclass
-class NodeInfo:
-    node_id: str
-    node_rank: int
-    node_ip: str
-
-
-# Define the coordinator and worker
-class RayClusterCoordinator:
-    def __init__(
-        self, worker_cls: type["ModelTrainer"], worker_resources: WorkerGroupResources
-    ) -> None:
-        self.worker_cls = worker_cls
-        self.worker_resources = worker_resources
-        self.num_workers = (
-            worker_resources.num_nodes * worker_resources.num_gpus_per_node
-        )
-        self.num_workers_per_node = worker_resources.num_gpus_per_node
-
-        ray_available_workers = int(ray.cluster_resources()["worker_units"])
-        assert self.num_workers // self.num_workers_per_node <= ray_available_workers, (
-            f"Only {ray_available_workers} workers available, which is not enough to schedule {self.num_workers} workers with {self.num_workers_per_node} workers per node"
-        )
-
-        self.workers_initialized = False
-
-        worker_node_info, self.num_physical_nodes = self._get_schedulable_worker_info()
-        print(f"Worker node info: {worker_node_info=}")
-        print(f"Num physical nodes: {self.num_physical_nodes=}")
-        # Assume there's one worker per GPU
-        self.workers = [
-            worker_cls.options(
-                num_gpus=1,
-                num_cpus=worker_resources.num_cpus_per_worker,
-                resources={"worker_units": 1},
-                # Use NodeAffinitySchedulingStrategy to ensure each worker is placed on a specific node
-                # node_id: Unique ID of the target node for this worker
-                # soft=False: Strictly enforce placement on the specified node (no fallback to other nodes)
-                scheduling_strategy=NASS(
-                    node_id=worker_node_info[i].node_id, soft=False
-                ),
-            ).remote(
-                i,
-                self.num_workers,
-                worker_node_info[i].node_rank,
-                worker_node_info[i].node_ip,  # TODO: probably can delete this
-                worker_node_info[
-                    0
-                ].node_ip,  # Arbitrarily make the first worker's hots the master
-                self.num_workers_per_node,
-            )
-            for i in range(self.num_workers)
-        ]
-
-    def _get_schedulable_worker_info(self) -> tuple[list[NodeInfo], int]:
-        """Collects information about available worker nodes in the Ray cluster and prepares
-        scheduling information for worker actors.
-
-        This method:
-        1. Identifies all alive worker nodes with 'worker_units' resources
-        2. Sorts them by NodeID for consistent allocation
-        3. Calculates how many physical nodes are needed based on workers per node
-        4. Verifies that enough nodes are available
-        5. Creates a list of NodeInfo objects for each worker
-
-        Returns:
-            tuple: (worker_node_info, num_nodes_required)
-                - worker_node_info: List of NodeInfo objects containing node_id, node_rank, and node_ip for each worker
-                - num_nodes_required: Number of physical nodes needed for all workers
-
-        Raises:
-            AssertionError: If there aren't enough nodes available to schedule all workers
-        """
-        # Get list of alive worker nodes sorted by NodeID for deterministic allocation
-        worker_node_info = []
-        worker_nodes = sorted(
-            [
-                node
-                for node in ray.nodes()
-                if (node["Alive"] and "worker_units" in node["Resources"])
-            ],
-            key=lambda x: x["NodeID"],
-        )
-
-        # Calculate required nodes and verify availability
-        num_nodes_required = self.num_workers // self.num_workers_per_node
-        num_nodes_available = len(worker_nodes)
-        assert num_nodes_required <= num_nodes_available
-
-        # Create worker info entries - one per GPU across all needed nodes
-        worker_nodes = worker_nodes[:num_nodes_required]
-        for worker_node_id, worker_node in enumerate(worker_nodes):
-            for _ in range(self.num_workers_per_node):
-                worker_node_info.append(
-                    NodeInfo(
-                        worker_node["NodeID"], worker_node_id, worker_node["NodeName"]
-                    )
-                )
-
-        return worker_node_info, num_nodes_required
-
-    def initialize_workers(self, **kwargs):
-        self.worker_init_kwargs = kwargs
-        ray.get([w.initialize.remote(**kwargs) for i, w in enumerate(self.workers)])
-        self.workers_initialized = True
-
-    def run(self, *args, **kwargs):
-        if not self.workers_initialized:
-            raise ValueError("""Cannot run workers without initializing them first.
-                              Please call the initialize_workers method of your cluster coordinator first.""")
-
-        worker_results = ray.get([w.run.remote(*args, **kwargs) for w in self.workers])
-        return worker_results
-
-
-class Worker:
-    def __init__(
-        self,
-        process_id,
-        world_size,
-        physical_node_id,
-        physical_node_ip,
-        master_addr: str,
-        num_workers_per_node: int,
-    ):
-        self.process_id = process_id
-        self.world_size = world_size
-        self.physical_node_id = physical_node_id
-        self.host_ip = physical_node_ip
-        self.master_addr = master_addr
-        self.logical_gpu_id = int(os.environ["CUDA_VISIBLE_DEVICES"])
-        print(
-            f"DEBUG: {self.logical_gpu_id=} {(process_id, world_size, physical_node_id, physical_node_ip, master_addr, num_workers_per_node)=}"
-        )
-        self.num_workers_per_node = num_workers_per_node
-
-    def get_process_id(self):
-        return self.process_id
-
-    def get_host_ip(self):
-        return self.host_ip
-
-    def get_logical_gpu_id(self):
-        return self.logical_gpu_id
-
-    def get_physical_node_id(self):
-        return self.physical_node_id
-
-    def initialize(self):
-        # Set distributed training environment variables
-        os.environ["RANK"] = str(self.process_id)
-        os.environ["WORLD_SIZE"] = str(self.world_size)
-        os.environ["LOCAL_RANK"] = str(self.logical_gpu_id)
-        os.environ["LOCAL_WORLD_SIZE"] = str(self.num_workers_per_node)
-        os.environ["MASTER_ADDR"] = self.master_addr
-        os.environ["MASTER_PORT"] = "29500"
-
-        dist.init_process_group("nccl")
-
-    def run(self, *args, **kwargs):
-        raise NotImplementedError
-
-
-@ray.remote
-class ModelTrainer(Worker):
-    def __init__(
-        self,
-        process_id,
-        world_size,
-        physical_node_id,
-        physical_node_ip,
-        master_addr,
-        num_workers_per_node,
-    ):
-        super().__init__(
-            process_id,
-            world_size,
-            physical_node_id,
-            physical_node_ip,
-            master_addr,
-            num_workers_per_node,
-        )
-
-    def run(self, hf_model_name):
-        rank = dist.get_rank()
-        world_size = dist.get_world_size()
-        ####local_device = torch.device(f"cuda:{rank}")
-        ####torch.cuda.set_device(local_device)
-
-        print(f"[Rank {rank}] Loading model {hf_model_name} on CPU...")
-        model = AutoModelForCausalLM.from_pretrained(
-            hf_model_name,
-            device_map="cpu",
-            torch_dtype=torch.bfloat16,
-        )
-
-        tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        # TODO: could oom?
-        # ------------------------------------------------
-        # 3) Move to GPU + Composable FSDP
-        #    (Initialize device mesh, shard submodules, then shard entire model)
-        # ------------------------------------------------
-        model.cuda()
-
-        # Create a device mesh with 'world_size' GPUs in a 1D arrangement.
-        mesh = init_device_mesh("cuda", (world_size,))
-
-        param_dtype = torch.bfloat16
-        reduce_dtype = torch.float32
-        buffer_dtype = torch.float32
-
-        mp = MixedPrecision(
-            param_dtype=param_dtype,
-            reduce_dtype=reduce_dtype,
-            buffer_dtype=buffer_dtype,
-        )
-
-        model = FullyShardedDataParallel(
-            model,
-            device_mesh=mesh,
-            auto_wrap_policy=size_based_auto_wrap_policy,
-            mixed_precision=mp,
-        )
-
-        # Optionally register "generate" as the forward method so FSDP can handle it properly.
-        register_fsdp_forward_method(model, "generate")
-
-        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
-        model.train()
-
-        num_epochs = 2
-        num_batches = 10
-        batch_size = 2
-        seq_length = 16
-        vocab_size = tokenizer.vocab_size or 32000
-
-        print(f"[Rank {rank}] Starting synthetic training...")
-
-        torch.manual_seed(42)
-        torch.cuda.manual_seed_all(42)
-        random.seed(42)
-        for epoch in range(num_epochs):
-            for step in range(num_batches):
-                input_ids = torch.ones(
-                    (batch_size, seq_length), device="cuda", dtype=torch.long
-                ) * (vocab_size - 1)
-                attention_mask = torch.ones_like(input_ids)
-                labels = input_ids.clone()
-
-                optimizer.zero_grad()
-                outputs = model(
-                    input_ids=input_ids, attention_mask=attention_mask, labels=labels
-                )
-                loss = torch.square(outputs.logits.view(-1)[0])
-                loss.backward()
-                optimizer.step()
-
-                if rank == 0:
-                    print(
-                        f"Epoch: {epoch}, Step: {step}, Loss: {loss.item():.8f}",
-                        flush=True,
-                    )
-
-        if rank == 0:
-            expected_loss = 0.012502908706665039
-            print(
-                f"[Rank {rank}] Testing loss is close to expected loss: {expected_loss}"
-            )
-            torch.testing.assert_close(loss.item(), expected_loss)
-            print("Yay! Loss was close :)")
-
-
-if __name__ == "__main__":
-    ray.init(address="auto", logging_level=0)
-    print(json.dumps(json.loads(os.environ["RAY_JOB_CONFIG_JSON_ENV_VAR"]), indent=4))
-    driver_args = json.loads(os.environ["RAY_JOB_CONFIG_JSON_ENV_VAR"])["runtime_env"][
-        "driver_args"
-    ]
-
-    # TODO: very simple, need to think thru CLI
-    trainer_args = driver_args["trainer"]
-    trainer_resources = trainer_args.pop("resources")
-    worker_resources = WorkerGroupResources(**trainer_resources)
-
-    coordinator = RayClusterCoordinator(ModelTrainer, worker_resources)
-    coordinator.initialize_workers()
-    print("Initialized workers")
-    # Get the job configuration set during launch.
-    # This is automatically set by Ray
-    coordinator.run(**trainer_args)
-    print("Finished")
diff --git a/tests/functional/sft.sh b/tests/functional/sft.sh
index 0e2298c983..b9836886d6 100755
--- a/tests/functional/sft.sh
+++ b/tests/functional/sft.sh
@@ -2,6 +2,8 @@
 
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetchs metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
 
 set -eou pipefail
 
@@ -10,23 +12,25 @@ JSON_METRICS=$LOG_DIR/$(basename $0 .sh).json
 RUN_LOG=$LOG_DIR/$(basename $0 .sh).log
 export RAY_DEDUP_LOGS=0
 export UV_CACHE_DIR=$PROJECT_ROOT/uv_cache
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
 
 mkdir -p $LOG_DIR
 
 cd $PROJECT_ROOT
-uv run $PROJECT_ROOT/examples/run_sft.py \
+python -u $PROJECT_ROOT/examples/run_sft.py \
     cluster.gpus_per_node=2 \
     sft.max_num_steps=10 \
     logger.tensorboard_enabled=true \
     logger.log_dir=$LOG_DIR \
     logger.wandb_enabled=false \
+    checkpointing.enabled=false \
     $@ \
     2>&1 | tee $RUN_LOG
 
 cd $SCRIPT_DIR
-uv run json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+python json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
 
-uv run check_metrics.py $JSON_METRICS \
-  'data["train/loss"]["9"] < 600' \
-  'data["timing/train/sft_train_step"]["9"] < 0.25'
+# TODO: loss is very noisy, this check is mainly for sanity of immediate divergence
+python check_metrics.py $JSON_METRICS \
+  'data["train/loss"]["9"] < 1500' \
 
diff --git a/tests/run_functional_in_docker.sh b/tests/run_functional_in_docker.sh
new file mode 100755
index 0000000000..e3fd403ba1
--- /dev/null
+++ b/tests/run_functional_in_docker.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/..)
+
+set -eou pipefail
+
+# Ensure Docker is installed
+if ! command -v docker &> /dev/null; then
+    echo "Error: Docker is not installed or not in PATH."
+    exit 1
+fi
+
+# CONTAINER is expected to be set as an environment variable
+if [[ -z "${CONTAINER:-}" ]]; then
+    echo "Error: CONTAINER environment variable is not set."
+    echo "Usage: CONTAINER=<docker-container> $0 <script to run, e.g., functional/grpo.sh>"
+    exit 1
+fi
+
+if [[ $# -ne 1 ]]; then
+    echo "Error: Did not provide functional test script to run."
+    echo "Usage: CONTAINER=<docker-container> $0 <script to run, e.g., functional/grpo.sh>"
+    exit 1
+fi
+
+TEST_SCRIPT=$(realpath $1)
+CONTAINER=${CONTAINER}
+
+export HF_HOME=${HF_HOME:-$(realpath $SCRIPT_DIR/../hf_home)}
+mkdir -p $HF_HOME
+
+# Check if running in GitLab CI
+INTERACTIVE_FLAG=""
+if [[ "${CI:-false}" != "true" ]]; then
+    # Setting this interactively lets us issue a keyboard interrupt.
+    INTERACTIVE_FLAG="-it"
+fi
+
+# Note: we run as root because:
+#  1. running as ray prevents us from writing into the current working directory
+#  2. running as ourselves (-u $(id -u):$(id -g)) causes torch compile to fail
+#
+# The workaround is we launch the job but set umask 000 so all files created as root are rwxrwxrwx.
+# We have found that 111 does not always work and can leave the filesystem permissions in a bad state.
+
+# Run the script inside the Docker container with GPU support
+docker run -u root $INTERACTIVE_FLAG --ulimit memlock=-1 --ulimit stack=67108864 --rm --gpus '"device=0,1"' -v "$PROJECT_ROOT:$PROJECT_ROOT" -v $HF_HOME:/hf_home -e WANDB_API_KEY -e HF_TOKEN -e HF_HOME=/hf_home -e HOME=/tmp/ -w $SCRIPT_DIR "$CONTAINER" -- bash -x -c "umask 000 && uv run bash -x $TEST_SCRIPT"
diff --git a/tests/run_unit.sh b/tests/run_unit.sh
old mode 100644
new mode 100755
index efc614868b..e826a18a60
--- a/tests/run_unit.sh
+++ b/tests/run_unit.sh
@@ -6,7 +6,6 @@ PROJECT_ROOT=$(realpath ${SCRIPT_DIR}/..)
 set -eou pipefail
 
 cd $SCRIPT_DIR
-GPUS_PER_NODE=$(nvidia-smi -L | grep -c '^GPU')
 
 if ! command -v pytest >/dev/null 2>&1; then
     echo "[ERROR] pytest not found. Make sure it's installed."
@@ -14,14 +13,21 @@ if ! command -v pytest >/dev/null 2>&1; then
 elif ! command -v ray >/dev/null 2>&1; then
     echo "[ERROR] ray binary not installed, which suggests this package is not installed."
     exit 1
-elif [[ $GPUS_PER_NODE -lt 2 ]]; then
-    echo "[ERROR]: Unit tests need at least 2 GPUs, but found $GPUS_PER_NODE"
-    exit 1
 fi
 
-export CUDA_DEVICE_ORDER=PCI_BUS_ID
-nvidia-smi
-export CUDA_VISIBLE_DEVICES=0,1
+# First try to connect to a ray cluster
+if ! ray status &>/dev/null; then
+    # If we cannot, then check if the local machine has at least two gpus to run the tests
+    GPUS_PER_NODE=$(nvidia-smi -L | grep -c '^GPU')
+    if [[ $GPUS_PER_NODE -lt 2 ]]; then
+        echo "[ERROR]: Unit tests need at least 2 GPUs, but found $GPUS_PER_NODE"
+        exit 1
+    fi
+    export CUDA_DEVICE_ORDER=PCI_BUS_ID
+    nvidia-smi
+    export CUDA_VISIBLE_DEVICES=0,1
+fi
+
 export PYTHONPATH=$(realpath ${SCRIPT_DIR}/..):${PYTHONPATH:-}
 export RAY_DEDUP_LOGS=0
 
diff --git a/tests/run_unit_in_docker.sh b/tests/run_unit_in_docker.sh
old mode 100644
new mode 100755
diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py
index 02e7a072b6..c129ee11fb 100644
--- a/tests/unit/algorithms/test_loss_functions.py
+++ b/tests/unit/algorithms/test_loss_functions.py
@@ -11,11 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import torch
 from nemo_reinforcer.algorithms.loss_functions import NLLLoss
 
 
 def test_nll_loss():
+    if not torch.cuda.is_available():
+        pytest.skip("No GPU available")
+
     loss_fn = NLLLoss()
 
     vocab_size = 8