Mattj/augusta

src/cookbook/model/builder.py

@@ -2,13 +2,18 @@
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
-from olmo_core.data import (
-    DataMix,
-    NumpyDataLoaderConfig,
-    NumpyDatasetConfig,
-    NumpyDatasetType,
-    TokenizerConfig,
+from cookbook.aliases import SourceInstance, WandbConfig
+from cookbook.data.dataset import MixtureBuilder
+from cookbook.model.config import (
+    MODEL_TO_LR_MAP,
+    DefaultOptimizerProperties,
+    ModelTrainConfig,
+    SupportedTokenizers,
+    WrappedTransformerConfig,
 )
+from cookbook.model.evaluators import DownstreamEvaluators
+from cookbook.model.schedulers import WSD
+from olmo_core.data import DataMix, NumpyDataLoaderConfig, NumpyDatasetConfig, NumpyDatasetType, TokenizerConfig
 from olmo_core.data.types import NumpyDatasetDType
 from olmo_core.nn.transformer import TransformerConfig
 from olmo_core.optim import AdamWConfig, CosWithWarmup, OptimGroupOverride, Scheduler
@@ -27,18 +32,6 @@
     WandBCallback,
 )
 
-from cookbook.aliases import SourceInstance, WandbConfig
-from cookbook.data.dataset import MixtureBuilder
-from cookbook.model.config import (
-    MODEL_TO_LR_MAP,
-    DefaultOptimizerProperties,
-    ModelTrainConfig,
-    SupportedTokenizers,
-    WrappedTransformerConfig,
-)
-from cookbook.model.evaluators import DownstreamEvaluators
-from cookbook.model.schedulers import WSD
-
 logger = logging.getLogger(__name__)
 
 
@@ -189,8 +182,15 @@ def __init__(
         if any(substring in cluster for substring in ["jupiter", "saturn"]) and weka:
             self.root_dir = f"/weka/oe-training-default/ai2-llm"
             logger.info(f"Using Weka bucket as root dir: {self.root_dir}")
-            self.checkpoint_dir = f"{self.root_dir}/checkpoints/{self.beaker_user.lower()}/{self.run_name}"
-
+        elif "augusta" in cluster:
+            try:
+                assert not weka
+            except AssertionError as e:
+                logger.info("Can't be on Augusta and weka!")
+                raise e
+            self.data_dir = self.root_dir = "gs://ai2-llm"
+
+        self.checkpoint_dir = f"{self.root_dir}/checkpoints/{self.beaker_user.lower()}/{self.run_name}"
         self.dataset_cache = f"{self.root_dir}/{self.beaker_user.lower()}/{self.run_name}/dataset-cache"
 
     def get_tokenizer_config(self, tokenizer) -> TokenizerConfig:

src/cookbook/recipes/love2code/train-SAMPLE-augusta.yaml

@@ -0,0 +1,21 @@
+name: "train-SAMPLE-augusta"
+description: "Just trying out cookbook + augusta"
+budget: "ai2/oe-training"
+workspace: "ai2/learn2code"
+nodes: 1
+gpus: 8
+preemptible: true
+max_tokens: 113_184_153_600 # 5xC multiplier
+sequence_length: 2048
+seed: 1337
+model: "olmo2_1B"
+tokenizer: "dolma2"
+priority: high
+cluster: ai2/augusta-google-1
+weka: false
+dataset:
+  sources:
+  - name: sample
+    target_ratio: 1.0
+    paths:
+      - gs://ai2-llm/preprocessed/love2code/python_only/python_tokens/part-000-00000.npy

src/cookbook/utils/config.py

@@ -3,23 +3,13 @@
 from typing import List, Tuple, cast
 
 import yaml
-from olmo_core.launch.beaker import (
-    BeakerEnvSecret,
-    BeakerLaunchConfig,
-    BeakerWekaBucket,
-)
-from olmo_core.train.callbacks import ConfigSaverCallback, WandBCallback
-from olmo_core.utils import get_default_device, seed_all
 
-from cookbook.aliases import (
-    ExperimentConfig,
-    ExperimentGroup,
-    ExperimentInstance,
-    SourceConfig,
-    SourceInstance,
-)
+from cookbook.aliases import ExperimentConfig, ExperimentGroup, ExperimentInstance, SourceConfig, SourceInstance
 from cookbook.model.builder import TransformerConfigBuilder
 from cookbook.utils.data import normalize_source_paths
+from olmo_core.launch.beaker import BeakerEnvSecret, BeakerLaunchConfig, BeakerWekaBucket
+from olmo_core.train.callbacks import ConfigSaverCallback, WandBCallback
+from olmo_core.utils import get_default_device, seed_all
 
 logger = logging.getLogger(__name__)
 
@@ -190,6 +180,8 @@ def mk_launch_configs(group: ExperimentGroup, beaker_user: str) -> list[BeakerLa
                 BeakerEnvSecret(name="AWS_CREDENTIALS", secret=f"{beaker_user}_AWS_CREDENTIALS"),
                 BeakerEnvSecret(name="R2_ENDPOINT_URL", secret="R2_ENDPOINT_URL"),
                 BeakerEnvSecret(name="WEKA_ENDPOINT_URL", secret="WEKA_ENDPOINT_URL"),
+                BeakerEnvSecret(name="GS_INTEROP_KEY", secret="GS_INTEROP_KEY"),
+                BeakerEnvSecret(name="GS_INTEROP_SECRET", secret="GS_INTEROP_SECRET"),
             ],
             setup_steps=[
                 'git clone "$REPO_URL"',

src/cookbook/utils/data.py

@@ -1,4 +1,6 @@
 import concurrent.futures
+import hashlib
+import json
 import logging
 import os
 import pathlib
@@ -7,22 +9,18 @@
 from urllib.parse import urlparse
 
 import s3fs
+from tqdm import tqdm
+
+from cookbook.aliases import SourceConfig
 from olmo_core.aliases import PathOrStr
 from olmo_core.data.types import NumpyDatasetDType
 from olmo_core.io import get_file_size, is_url, normalize_path
 from olmo_core.utils import OLMoEnvironmentError
-from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 logging.getLogger("botocore").setLevel(logging.WARNING)
 
 
-import hashlib
-import json
-
-from cookbook.aliases import SourceConfig
-
-
 def _bytes_to_tokens(num_bytes: int, dtype: NumpyDatasetDType) -> int:
     """
     Convert bytes to tokens based on the dtype.
@@ -65,9 +63,12 @@ def get_token_counts_and_ratios(
             parsed = urlparse(path)
             if parsed.scheme == "s3":
                 continue
-            if parsed.scheme == "weka":
+            elif parsed.scheme == "weka":
                 client_kwargs["endpoint_url"] = os.environ.get("WEKA_ENDPOINT_URL")
-
+            elif parsed.scheme == "gs":
+                client_kwargs["endpoint_url"] = "https://storage.googleapis.com"
+                client_kwargs["key"] = os.environ.get("GS_INTEROP_KEY")
+                client_kwargs["secret"] = os.environ.get("GS_INTEROP_SECRET")
     fs = s3fs.S3FileSystem(client_kwargs={**client_kwargs})
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=64) as executor: