[wip] OLMo3 anneals

README.md

@@ -315,3 +315,10 @@ All PMR CLI commands support the following options:
 | `--script` | `-s` | None | Path to script file or directory to execute |
 
 Note that you can provide either `--command` or `--script`, but not both. When using `--script` with a directory path, all executable files in that directory will be distributed across the instances.
+
+
+# Midtraining utilities
+I (mj) built out some utilities to minimize manual labor for common tasks for midtraining.
+- [`scripts/gs2weka.py`](scripts/gs2weka.py): This script finds the latest checkpoint for a given model configuration and copies it from Google Cloud Storage to Weka storage using olmo-cookbook. Run with `python scripts/gs2weka.py <yaml_file>` to automatically detect your Beaker account and process the latest checkpoint, or use --beaker-name to specify a different account name.
+- [`scripts/convert_from_config.py`](scripts/convert_from_config.py): This script finds the latest checkpoint in Weka storage for a given model configuration and converts it to HuggingFace format using olmo-cookbook-eval. Run with `python scripts/convert_from_config.py <yaml_file>` to automatically detect your Beaker account and convert the latest checkpoint, with optional --overwrite flag to reconvert existing checkpoints.
+- [`scripts/olmo3_midtrain_eval.sh`](scripts/olmo3_midtrain_eval.sh): This script runs OLMo3 midtraining evaluations on a given checkpoint path using two different task suites (midtrain and main). Run with `bash scripts/olmo3_midtrain_eval.sh <checkpoint_path>` where the checkpoint path should point to a converted HuggingFace format checkpoint (e.g., ending in -hf).

pyproject.toml

@@ -51,7 +51,7 @@ checkpoints = [
     "boto3"
 ]
 all = [
-    "ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@c779ca546cc3194e73e7491aaefcdffbed042c65",
+    "ai2-olmo-core @ git+https://github.com/allenai/OLMo-core.git@tylerr/olmo3-scripts-swafix-foreachopt",
     "beaker-py>=1,<2",
     "GitPython>=3.0,<4.0",
     "wandb",

scripts/convert_from_config.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python3
+"""
+Script to process YAML file and run olmo-cookbook command with latest checkpoint
+"""
+import argparse
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import yaml
+
+
+def run_command(cmd, shell=False, errs_okay=False):
+    """Run a shell command and return stdout"""
+    try:
+        if shell:
+            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=True)
+        else:
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        return result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        print(f"Error running command: {' '.join(cmd) if isinstance(cmd, list) else cmd}")
+        print(f"Error: {e.stderr}")
+        if not errs_okay:
+            sys.exit(1)
+        raise e
+
+
+def get_yaml_name(yaml_file):
+    """Extract the 'name' attribute from YAML file"""
+    try:
+        with open(yaml_file, "r") as f:
+            data = yaml.safe_load(f)
+
+        if "name" not in data:
+            print(f"Error: 'name' attribute not found in {yaml_file}")
+            sys.exit(1)
+
+        return data["name"]
+    except Exception as e:
+        print(f"Error reading YAML file {yaml_file}: {e}")
+        sys.exit(1)
+
+
+def get_beaker_name():
+    """Get the NAME from 'beaker account whoami' output"""
+    output = run_command(["beaker", "account", "whoami"])
+
+    # Parse the table output to extract NAME
+    lines = output.strip().split("\n")
+    if len(lines) < 2:
+        print("Error: Unexpected output from 'beaker account whoami'")
+        sys.exit(1)
+
+    # Look for the data row (skip header)
+    for line in lines[1:]:
+        parts = line.split()
+        if len(parts) >= 2:
+            return parts[1]  # NAME is the second column
+
+    print("Error: Could not extract NAME from beaker account whoami output")
+    sys.exit(1)
+
+
+def find_latest_checkpoint(beaker_name, yaml_name):
+    """Find the latest checkpoint directory in weka"""
+
+    weka_path = f"weka://oe-training-default/ai2-llm/checkpoints/{beaker_name}/{yaml_name}-*"
+
+    # Convert weka:// path to s3:// path for s5cmd
+    s3_path = weka_path.replace("weka://oe-training-default/", "s3://oe-training-default/")
+
+    # Add wildcard to check for any files in the directory
+    s3_path_wildcard = f"{s3_path}/*"
+
+    print(f"Checking if weka path exists: {weka_path}")
+    print(f"Using s5cmd to check: {s3_path_wildcard}")
+
+    cmd = [
+        "s5cmd",
+        "--profile",
+        "WEKA",
+        "--endpoint-url",
+        "https://weka-aus.beaker.org:9000",
+        "ls",
+        s3_path_wildcard,
+    ]
+
+    try:
+        output = run_command(cmd, errs_okay=True)
+        if not output:
+            print(f"No checkpoints found with prefix: {prefix}")
+            sys.exit(1)
+
+        # Get all matching paths
+        paths = output.strip().split("\n")
+
+        # Sort paths to get the latest one (lexicographically)
+        paths = [_.split(" ")[-1].strip() for _ in paths]
+        ckpts = set()
+        for p in paths:
+            re_string = yaml_name + r"-[0-9a-f]{8}/step\d+/"
+            if re.match(re_string, p):
+                ckpts.add(re.match(re_string, p).group())
+        assert (
+            len(ckpts) > 0
+        ), "No valid checkpoints found??? [this should assert should never fail if we got here to begin with]"
+        max_ckpt = max(ckpts)
+        print(max_ckpt)
+        return "weka://oe-training-default/ai2-llm/checkpoints/%s/%s" % (beaker_name, max_ckpt)
+
+    except subprocess.CalledProcessError as e:
+        print("No weka paths found!")
+        print(
+            f"Make sure you have access to weka://oe-training-deafult/ai2-llm/checkpoints/{beaker_name}/{yaml_name}-* directories"
+        )
+        raise e
+        # sys.exit(1)
+    except Exception as e:
+        print("ERR CODE ", e)
+        raise e
+
+
+def check_hf_path_exists(latest_ckpt):
+    """Check if the corresponding weka path already exists"""
+    # Convert gs:// path to weka:// path
+    hf_path = latest_ckpt.rstrip("/") + "-hf/*"
+
+    print(f"Checking if weka path exists: {hf_path}")
+    cmd = [
+        "s5cmd",
+        "--profile",
+        "WEKA",
+        "--endpoint-url",
+        "https://weka-aus.beaker.org:9000",
+        "ls",
+        hf_path,
+    ]
+
+    # Convert weka:// path to s3:// path for s5cmd
+    hf_path = hf_path.replace("weka://oe-training-default/", "s3://oe-training-default/")
+
+    print(f"Checking if weka path exists: {hf_path}")
+    cmd = [
+        "s5cmd",
+        "--profile",
+        "WEKA",
+        "--endpoint-url",
+        "https://weka-aus.beaker.org:9000",
+        "ls",
+        hf_path,
+    ]
+
+    try:
+        # Run the command - if it succeeds, the path exists
+        output = run_command(cmd, errs_okay=True)
+        print(f"✅ Weka path exists - found %s files:" % len(output.split("\n")))
+        return True
+    except subprocess.CalledProcessError as e:
+        # If the command fails, the path doesn't exist
+        print(f"❌ Weka path does not exist (s5cmd failed as expected)")
+        return False
+
+
+def run_olmo_cookbook(weka_path):
+    """Run the olmo-cookbook command with the GCS path"""
+    print("Converting %s" % weka_path)
+    weka_path = weka_path.replace("weka://", "/").rstrip("/")
+    cmd = [
+        "olmo-cookbook-eval",
+        "convert",
+        weka_path,
+        "-t",
+        "olmo-core-v2",
+        "--use-beaker",
+        "--huggingface-transformers-git-url",
+        "https://github.com/2015aroras/transformers.git",
+        "--huggingface-transformers-commit-hash",
+        "ae3889ced6ed7362e5883671fc6dc4cb4fece5fa",
+        "--olmo-core-v2-commit-hash",
+        "57a04d0b69047d797c96eede056a211e75b5914a",
+    ]
+    print(f"Running: {' '.join(cmd)}")
+
+    try:
+        # Run the command and stream output in real-time
+        process = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, universal_newlines=True
+        )
+
+        beaker_url = None
+        beaker_url_pattern = re.compile(r"https://beaker\.org/ex/[A-Z0-9]+")
+
+        for line in process.stdout:
+            print(line, end="")
+
+            # Look for the beaker URL in the output
+            match = beaker_url_pattern.search(line)
+            if match:
+                beaker_url = match.group(0)
+
+        process.wait()
+
+        if process.returncode != 0:
+            print(f"Error: olmo-cookbook command failed with return code {process.returncode}")
+            sys.exit(1)
+
+        # Print the extracted Beaker URL
+        if beaker_url:
+            print(f"\n" + "=" * 60)
+            print(f"🔗 Beaker Experiment URL: {beaker_url}")
+            print(f"=" * 60)
+            return beaker_url
+        else:
+            print("\nWarning: Could not extract Beaker experiment URL from output")
+
+    except Exception as e:
+        print(f"Error running olmo-cookbook: {e}")
+        sys.exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Process YAML file and run olmo-cookbook with latest checkpoint")
+    parser.add_argument("yaml_file", help="Path to the YAML file")
+    parser.add_argument("--beaker-name", required=False, default=None)
+    parser.add_argument("--overwrite", required=False, type=bool, default=False)
+    args = parser.parse_args()
+
+    # Validate input file exists
+    if not Path(args.yaml_file).exists():
+        print(f"Error: YAML file {args.yaml_file} does not exist")
+        sys.exit(1)
+
+    print(f"Processing YAML file: {args.yaml_file}")
+
+    # Step 1: Get name from YAML
+    yaml_name = get_yaml_name(args.yaml_file)
+    print(f"YAML name: {yaml_name}")
+
+    # Step 2: Get beaker name
+    if args.beaker_name == None:
+        beaker_name = get_beaker_name()
+    else:
+        beaker_name = args.beaker_name
+    print(f"Beaker name: {beaker_name}")
+
+    # Step 3: Find latest checkpoint
+    print(
+        f"Searching for checkpoints with prefix: weka://oe-training-default/ai2-llm/checkpoints/{beaker_name}/{yaml_name}-"
+    )
+    latest_checkpoint = find_latest_checkpoint(beaker_name, yaml_name)
+    print(f"Latest checkpoint: {latest_checkpoint}")
+
+    # Step 4: Check if weka path already exists
+    if check_hf_path_exists(latest_checkpoint) and not args.overwrite:
+        print(f"\n🚫 Converted checkpoint already exists in weka storage. Skipping cookbook command.")
+        print(f"The checkpoint has already been copied to weka://oe-training-default/")
+        return
+
+    # Step 5: Run olmo-cookbook command
+    run_olmo_cookbook(latest_checkpoint)
+
+
+if __name__ == "__main__":
+    main()