upstream/main merge and dataset stats

.github/workflows/push-image.yml

@@ -60,8 +60,12 @@ jobs:
           # ghcr_user: ${{ secrets.GHCR_USER }}
 
       # big images fail, trying this
+      # reference for big files in runner: https://github.com/actions/runner-images/issues/10386
       - name: Delete huge unnecessary tools folder
-        run: rm -rf /opt/hostedtoolcache /usr/share/dotnet "$AGENT_TOOLSDIRECTORY"
+        run: rm -rf /opt/hostedtoolcache /usr/share/dotnet "$AGENT_TOOLSDIRECTORY" /usr/local/lib/android/sdk/ndk
+
+      - name: Check remaining disk space
+        run: df -h
 
       - name: Build image
         run: |

Dockerfile

@@ -11,7 +11,7 @@ RUN apt-get update --no-install-recommends && apt-get install -y nginx && mkdir
 ENV HF_HUB_ENABLE_HF_TRANSFER=1
 RUN pip install torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128 --no-cache-dir
 RUN pip install packaging --no-cache-dir
-RUN pip install flash-attn==2.8.0.post2 --no-build-isolation --no-cache-dir
+RUN pip install flash-attn==2.8.0.post2 flashinfer-python>=0.2.7.post1 --no-build-isolation --no-cache-dir
 COPY requirements.txt .
 RUN pip install -r requirements.txt --no-cache-dir
 RUN python -m nltk.downloader punkt

Dockerfile.uv

@@ -11,6 +11,9 @@ RUN echo "UV_CACHE_DIR: ${UV_CACHE_DIR}"
 # setup files
 WORKDIR /stage/
 
+# Install nginx and create conf.d directory
+RUN apt-get update --no-install-recommends && apt-get install -y nginx && mkdir -p /etc/nginx/conf.d && rm -rf /var/lib/apt/lists/*
+
 ENV HF_HUB_ENABLE_HF_TRANSFER=1
 ENV UV_COMPILE_BYTECODE=0
 
@@ -26,6 +29,10 @@ RUN --mount=type=cache,target=${UV_CACHE_DIR} \
 # rerun with extras
 RUN uv sync --extra compile --extra liger
 
+# punkt / punkt_tab
+RUN uv run -m nltk.downloader punkt
+RUN uv run -m nltk.downloader punkt_tab
+
 WORKDIR /stage/
 
 # Copy all runtime files directly to final stage

README.md

@@ -53,7 +53,7 @@ pip install --upgrade pip "setuptools<70.0.0" wheel
 # TODO, unpin setuptools when this issue in flash attention is resolved
 pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
 pip install packaging
-pip install flash-attn==2.7.2.post1 --no-build-isolation
+pip install flash-attn==2.7.2.post2 flashinfer-python>=0.2.7.post1 --no-build-isolation
 pip install -r requirements.txt
 pip install -e .
 python -m nltk.downloader punkt
@@ -62,7 +62,7 @@ python -m nltk.downloader punkt
 * **Local installation with uv (preview)**: We are experimenting with using [uv](https://docs.astral.sh/uv/). You can install via
 ```bash
 uv sync
-uv sync --extra compile --extra liger # to install flash attention and liger-kernel
+uv sync --extra compile --extra liger # to install flash attention, flash infer, and liger-kernel
 ```
 
 

configs/beaker_configs/code_api_setup.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+set -e
+
 # Configuration
 TOTAL_CPUS=$(nproc)
 CODE_SERVER_CPUS=128
@@ -8,7 +10,7 @@ NGINX_PORT=8070
 API_BASE_PORT=1234
 
 # Get leader replica IP
-BEAKER_LEADER_REPLICA_IP=$(getent hosts ${BEAKER_LEADER_REPLICA_HOSTNAME} | awk '{print $1}')
+BEAKER_LEADER_REPLICA_IP=$(getent hosts ${BEAKER_LEADER_REPLICA_HOSTNAME} | head -n 1 | awk '{print $1}')
 
 # Set up environment
 export PYTHONPATH=$REPO_PATH

mason.py

@@ -54,39 +54,35 @@ def parse_env_var(env_var_str: str) -> Dict[str, str]:
         raise argparse.ArgumentTypeError("Environment variable name cannot be empty")
     return {"name": name, "value": value}
 
-
-NFS_CLUSTERS = [
-    "ai2/allennlp-cirrascale",
-    "ai2/aristo-cirrascale",
-    "ai2/climate-cirrascale",
-    "ai2/general-cirrascale",
-    "ai2/general-cirrascale-a5000",
-    "ai2/mosaic-cirrascale",
-    "ai2/mosaic-cirrascale-a100",
-    "ai2/pluto-cirrascale",
-    "ai2/prior-cirrascale",
-    "ai2/s2-cirrascale",
-    "ai2/s2-cirrascale-l40",
-]
-
 WEKA_CLUSTERS = [
     "ai2/jupiter-cirrascale-2",
+    "ai2/jupiter",
     "ai2/saturn-cirrascale",
     "ai2/titan-cirrascale",
+    "ai2/titan",
     "ai2/neptune-cirrascale",
+    "ai2/neptune",
     "ai2/ceres-cirrascale",
+    "ai2/ceres",
     "ai2/triton-cirrascale",
+    "ai2/triton",
     "ai2/rhea-cirrascale",
+    "ai2/rhea",
 ]
 GCP_CLUSTERS = [
-    "ai2/augusta-google-1"
+    "ai2/augusta-google-1",
+    "ai2/augusta"
 ]
 
 INTERCONNECT_CLUSTERS = [
     "ai2/jupiter-cirrascale-2",
+    "ai2/jupiter",
     "ai2/ceres-cirrascale",
+    "ai2/ceres",
     "ai2/titan-cirrascale",
+    "ai2/titan",
     "ai2/augusta-google-1",
+    "ai2/augusta",
 ]
 
 
@@ -238,6 +234,9 @@ def get_env_vars(pure_docker_mode: bool, cluster: List[str], beaker_secrets: Lis
                 whoami: str, resumable: bool, num_nodes: int, additional_env_vars: List[Dict[str, str]],
                 additional_secrets: List[Dict[str, str]]):
     env_vars = []
+    if "VLLM_ATTENTION_BACKEND" not in additional_env_vars:
+        env_vars.append(beaker.EnvVar(name="VLLM_ATTENTION_BACKEND",
+                                      value="FLASHINFER"))
     # Add user-specified environment variables first
     for env_var in additional_env_vars:
         env_vars.append(
@@ -487,16 +486,8 @@ def get_env_vars(pure_docker_mode: bool, cluster: List[str], beaker_secrets: Lis
 def get_datasets(beaker_datasets, cluster: List[str]):
     """if pure docker mode we don't mount the NFS; so we can run it on jupiter2"""
     res = []
-    # if none of the cluster is in weka, we mount the NFS
-    if all(c in NFS_CLUSTERS for c in cluster):
-        res = [
-            beaker.DataMount(
-                source=beaker.DataSource(host_path="/net/nfs.cirrascale"),
-                mount_path="/net/nfs.cirrascale",
-            ),
-        ]
     # if all cluster is in weka, we mount the weka
-    elif all(c in WEKA_CLUSTERS for c in cluster):
+    if all(c in WEKA_CLUSTERS for c in cluster):
         res = [
             beaker.DataMount(
                 source=beaker.DataSource(weka="oe-adapt-default"),
@@ -635,6 +626,8 @@ def remove_arg_from_list(lst: List[str], item: str, remove_value: bool = False):
                             dataset_config_hashes.append(dataset_config_hash)
                     stderr = result.stderr
                     return_code = result.returncode
+                    if return_code != 0:
+                        raise Exception(f"Error code {return_code} when creating cached dataset")
                     console.log("✅✅✅ Finished running the caching command")
 
                 if file in OPEN_INSTRUCT_RESUMABLES and idx != -1 and len(args.auto_checkpoint_state_dir) > 0:
@@ -706,8 +699,13 @@ def remove_arg_from_list(lst: List[str], item: str, remove_value: bool = False):
                     break
 
             commit_hash = get_commit_hash(model_name_or_path, model_revision, "config.json", "model")
-            download_from_hf(model_name_or_path, model_revision) # first download the model
-            path = download_from_hf(model_name_or_path, model_revision) # then get the path
+            if os.path.exists(model_name_or_path):
+                path = model_name_or_path
+                model_name_or_path = os.path.basename(model_name_or_path)
+                console.log(f"Local model is already downloaded, using path basename as model name {model_name_or_path}, note that commit hash is {commit_hash}")
+            else:
+                download_from_hf(model_name_or_path, model_revision) # first download the model
+                path = download_from_hf(model_name_or_path, model_revision) # then get the path
             gs_saved_path = f"gs://ai2-llm/post-training/deletable_cache_models/{model_name_or_path}/{commit_hash}"
             gs_folder = gs_folder_exists(gs_saved_path) # race condition exists, but it's fine since we are launching mason sequentially
             if not gs_folder: