Training script

.github/workflows/checks.yml

@@ -17,6 +17,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
       - name: setup python
         uses: actions/setup-python@v5
         with:
@@ -31,11 +33,11 @@ jobs:
       - name: dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt
+          pip install -r requirements-nocuda.txt
           pip install -e .
       - name: black
         run: black --check .
       - name: isort
-        run: isort --profile black --check .
+        run: isort --check .
       - name: pytest
         run: pytest

.gitignore

@@ -6,6 +6,11 @@ __pycache__/
 # C extensions
 *.so
 
+bin
+include
+lib64
+pyvenv.cfg
+
 # Distribution / packaging
 .Python
 build/
@@ -158,3 +163,15 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# ignore wandb files
+**/wandb/*
+**/*.wandb
+**/wandb-summary.json
+**/wandb-metadata.json
+
+# scratch notebook
+notebooks/scratch.ipynb
+
+# dsstore
+.DS_Store

.pre-commit-config.yaml

@@ -8,5 +8,3 @@ repos:
     rev: 5.13.2
     hooks:
       - id: isort
-        name: isort (python)
-        args: ["--profile", "black"]

.vscode/launch.json

@@ -0,0 +1,34 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "run_training 256",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "scripts/run_training.py",
+            "console": "integratedTerminal",
+            "args": "--debug --train_sample_limit=256"
+            //"args": "${command:pickArgs}"
+        },
+        {
+            "name": "run_training --help",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "scripts/run_training.py",
+            "console": "integratedTerminal",
+            "args": "--help"
+            //"args": "${command:pickArgs}"
+        },
+        {
+            "name": "run training with debug plus custom args",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "scripts/run_training.py",
+            "console": "integratedTerminal",
+            "args": "--debug ${command:pickArgs}"
+        }
+    ]
+}

.vscode/settings.json

@@ -7,9 +7,5 @@
         "source.organizeImports": "explicit"
     },
     "python.analysis.typeCheckingMode": "basic",
-    "isort.args": [
-        "--profile",
-        "black"
-    ],
     "black-formatter.importStrategy": "fromEnvironment",
 }

README.md

@@ -49,3 +49,6 @@ When you save a file vscode should automatically format it. Otherwise, pre-commi
      - comment important sections of the code in _Files changed_ tab
    - when it's ready, add the relevant stakeholders as reviewers
 4. after the comments are resolved and PR is approved, merge it using _Squash and merge_
+
+# Incrementing Versions
+When making a new release, increment the version in `delphi/__init__.py`

notebooks/training_demo.ipynb

@@ -0,0 +1,45 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from delphi.train.config.utils import get_presets_by_name\n",
+    "from delphi.train.training import run_training\n",
+    "from delphi.train.utils import ModelTrainingState\n",
+    "from delphi.train.run_context import RunContext\n",
+    "\n",
+    "\n",
+    "def train() -> tuple[ModelTrainingState, RunContext]:\n",
+    "    config = get_presets_by_name()[\"v0-llama2-100k\"]\n",
+    "    config.wandb_config.entity = \"jaiwithani\"\n",
+    "    return run_training(config)\n",
+    "\n",
+    "model_train_result = train()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "tinyevals",
+   "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.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "delphi"
+dynamic = ["version"]
+
+[tool.setuptools.dynamic]
+version = {attr = "delphi.__version__"}
+
+
+[tool.black]
+extend-exclude = 'src/llama2c'
+
+[tool.isort]
+profile = 'black'
+known_third_party = ['llama2c', 'wandb']
+extend_skip = ['src/llama2c']
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

requirements-nocuda.txt

@@ -0,0 +1,30 @@
+# this is a separate requirements.txt file for use in github actions
+# this omits packages that cannot be installed in github actions due
+# to hardware limitations (e.g. no GPU). All packages here are automatically
+# included when installing from requirements.txt
+torch==2.1.2
+datasets==2.16.1
+tqdm==4.66.1
+ipywidgets==8.1.1
+nbformat==5.9.2
+pytest==7.4.4
+black==23.12.1
+jaxtyping==0.2.25
+beartype==0.16.4
+pre-commit==3.6.0
+isort==5.13.2
+chardet==5.2.0
+sentencepiece==0.1.99
+protobuf==4.25.2
+plotly==5.18.0
+wandb==0.16.3
+spacy==3.7.2
+pandas==1.3.4
+dacite==1.8.1
+
+# temporarily installing transformers from main until 4.39.0 comes out (for mamba support)
+transformers @ git+https://github.com/huggingface/transformers@main
+# transformers==4.39.0  TODO: use this once 4.39.0 releases
+
+# spacy-transformers requires transformers <= 4.37.0, temporarily disabling
+# spacy-transformers>=1.3.4

requirements.txt

@@ -1,19 +1,8 @@
-torch==2.1.2
-datasets==2.16.1
-transformers==4.36.2
-tqdm==4.66.1
-ipywidgets==8.1.1
-nbformat==5.9.2
-pytest==7.4.4
-black==23.12.1
-jaxtyping==0.2.25
-beartype==0.16.4
-pre-commit==3.6.0
-isort==5.13.2
-spacy==3.7.2
-chardet==5.2.0
-sentencepiece==0.1.99
-protobuf==4.25.2
-plotly==5.18.0
-spacy-transformers==1.3.4
-pandas==1.3.4
+# most packages are specified in requirements-gh.txt, and new packages should be placed
+# there UNLESS they cannot be installed without CUDA support, in which case they should go here.
+-r requirements-nocuda.txt
+
+# these libs support better mamba implementations in transformers, 
+# but require CUDA/nvcc, so they won't work on MacOS.
+mamba_ssm==1.2.0.post1; sys_platform != 'darwin'
+causal-conv1d==1.2.0.post2; sys_platform != 'darwin'

scripts/run_training.py

@@ -0,0 +1,151 @@
+import argparse
+import logging
+import os
+from dataclasses import fields, is_dataclass
+from itertools import chain
+from pathlib import Path
+from typing import Any, Optional, _GenericAlias  # type: ignore
+
+from delphi.constants import CONFIG_PRESETS_DIR
+from delphi.train.config import (
+    GigaConfig,
+    build_config_from_files_and_overrides,
+    get_preset_paths,
+    get_user_config_path,
+)
+from delphi.train.training import run_training
+from delphi.train.utils import save_results
+
+
+def _unoptionalize(t: type) -> type:
+    if isinstance(t, _GenericAlias):
+        return t.__args__[0]
+    else:
+        return t
+
+
+def get_preset_args(args: argparse.Namespace) -> list[Path]:
+    cands = []
+    for preset in get_preset_paths():
+        if hasattr(args, preset.stem) and getattr(args, preset.stem):
+            cands.append(preset)
+    return cands
+
+
+def get_config_files(args: argparse.Namespace) -> list[Path]:
+    user_config_path = get_user_config_path()
+    cands = [user_config_path] if user_config_path.exists() else []
+    cands += get_preset_args(args)
+    config_files = list(chain(*args.config_file)) if args.config_file else []
+    cands += map(Path, config_files)
+    configs = []
+    for candpath in cands:
+        if candpath.exists():
+            configs.append(candpath)
+            logging.info(f"Found config file {candpath}...")
+        else:
+            raise FileNotFoundError(candpath, f"Config file {candpath} does not exist.")
+    return configs
+
+
+def add_dataclass_args_recursively(
+    parser: argparse.ArgumentParser,
+    dc: type[object],
+    default_group: argparse._ArgumentGroup,
+    group: Optional[argparse._ArgumentGroup] = None,
+    prefix: str = "",
+):
+    for field in fields(dc):  # type: ignore
+        # if field is an Optional type, strip it to the actual underlying type
+        _type = _unoptionalize(field.type)
+        if is_dataclass(_type):
+            _group = group or parser.add_argument_group(f"{field.name}")
+            add_dataclass_args_recursively(
+                parser,
+                _type,
+                default_group,
+                _group,
+                prefix=f"{prefix}{field.name}.",
+            )
+        else:
+            _group = group or default_group
+            _group.add_argument(
+                f"--{prefix}{field.name}",
+                type=_type,
+                required=False,
+                help=f"Default: {field.default}"
+                if field.default != field.default_factory
+                else f"Must be specified as part of {_group.title}",
+            )
+
+
+def setup_parser() -> argparse.ArgumentParser:
+    # Setup argparse
+    parser = argparse.ArgumentParser(description="Train a delphi model")
+    parser.add_argument(
+        "--config_file",
+        help=(
+            "Path to json file(s) containing config values. Specific values can be overridden with --arguments. "
+            "e.g. `--config_file primary_config.json secondary_config.json --log_interval 42`. "
+            'If passing multiple configs with overlapping args, use "priority" key to specify precedence, e.g. {"priority": 100} '
+            f'overrides {{"priority": 99}} See preset configs in {CONFIG_PRESETS_DIR}'
+        ),
+        action="append",
+        nargs="*",
+        required=False,
+        type=str,
+    )
+    config_arg_group = parser.add_argument_group("Config arguments")
+    add_dataclass_args_recursively(parser, GigaConfig, config_arg_group)
+    preset_arg_group = parser.add_argument_group("Preset configs")
+    for preset in sorted(get_preset_paths()):
+        preset_arg_group.add_argument(
+            f"--{preset.stem}",
+            help=f"Use {preset.stem} preset config",
+            action="store_true",
+        )
+    return parser
+
+
+def var_args_to_dict(config_vars: dict[str, Any]) -> dict[str, Any]:
+    # {"a.b.c" = 4} to {"a": {"b": {"c": 4}}}
+    d = {}
+    for k, v in config_vars.items():
+        if v is None:
+            continue
+        cur = d
+        subkeys = k.split(".")
+        for subkey in subkeys[:-1]:
+            if subkey not in cur:
+                cur[subkey] = {}
+            cur = cur[subkey]
+        cur[subkeys[-1]] = v
+    return d
+
+
+def args_to_dict(args: argparse.Namespace) -> dict[str, Any]:
+    # at the toplevel, filter for args corresponding to field names in GigaConfig
+    field_names = set(field.name for field in fields(GigaConfig))
+    config_vars = {
+        k: v for k, v in vars(args).items() if k.split(".")[0] in field_names
+    }
+    return var_args_to_dict(config_vars)
+
+
+def main():
+    parser = setup_parser()
+    args = parser.parse_args()
+
+    config_files = get_config_files(args)
+    args_dict = args_to_dict(args)
+    config = build_config_from_files_and_overrides(config_files, args_dict)
+
+    # run training
+    results, run_context = run_training(config)
+    final_out_dir = os.path.join(config.output_dir, "final")
+    save_results(config, results, run_context, final_out_dir)
+    print(f"Saved results to {final_out_dir}")
+
+
+if __name__ == "__main__":
+    main()