rename DOLMA -> OLMo (allenai#86)

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -17,7 +17,7 @@ body:
 
       ```python
       # All necessary imports at the beginning
-      import dolma
+      import olmo
 
       # A succinct reproducing example trimmed down to the essential parts:
       assert False is True, "Oh no!"

.github/workflows/main.yml

@@ -101,15 +101,15 @@ jobs:
         if: always()
         run: |
           . .venv/bin/activate
-          pip uninstall -y dolma
+          pip uninstall -y olmo
 
   gpu_tests:
     name: GPU Tests
     runs-on: ubuntu-latest
     timeout-minutes: 15
     env:
       BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
-      BEAKER_IMAGE: dolma-torch2-test
+      BEAKER_IMAGE: olmo-torch2-test
       BEAKER_WORKSPACE: ai2/llm-testing
     steps:
       - name: Determine current commit SHA (pull request)

.github/workflows/pr_checks.yml

@@ -9,7 +9,7 @@ on:
     branches:
       - main
     paths:
-      - 'dolma/**'
+      - 'olmo/**'
 
 jobs:
   changelog:

CONTRIBUTING.md

@@ -77,8 +77,8 @@ When you're ready to contribute code to address an open issue, please follow the
 
     Then you can create and activate a new Python environment by running:
 
-        conda create -n dolma python=3.9
-        conda activate dolma
+        conda create -n olmo python=3.9
+        conda activate olmo
 
     Once your virtual environment is activated, you can install your local clone in "editable mode" with
 
@@ -139,13 +139,13 @@ When you're ready to contribute code to address an open issue, please follow the
 
     We also strive to maintain high test coverage, so most contributions should include additions to [the unit tests](https://github.com/allenai/LLM/tree/main/tests). These tests are run with [`pytest`](https://docs.pytest.org/en/latest/), which you can use to locally run any test modules that you've added or changed.
 
-    For example, if you've fixed a bug in `dolma/a/b.py`, you can run the tests specific to that module with
+    For example, if you've fixed a bug in `olmo/a/b.py`, you can run the tests specific to that module with
 
         pytest -v tests/a/b_test.py
 
     To check the code coverage locally in this example, you could run
 
-        pytest -v --cov dolma.a.b tests/a/b_test.py
+        pytest -v --cov olmo.a.b tests/a/b_test.py
 
     If your contribution involves additions to any public part of the API, we require that you write docstrings
     for each function, method, class, or module that you add.

Makefile

@@ -1,5 +1,5 @@
 # If you update this, also update BEAKER_IMAGE in .github/workflows/main.yml
-IMAGE_NAME_BASE = dolma-torch2
+IMAGE_NAME_BASE = olmo-torch2
 # If you update this, also update BEAKER_WORKSPACE in .github/workflows/main.yml
 BEAKER_WORKSPACE = ai2/llm-testing
 
@@ -24,7 +24,7 @@ beaker-info :
 .PHONY : images
 images : gantry-image test-image
 
-PHONY : base-image
+.PHONY : base-image
 base-image :
 	docker build -f docker/Dockerfile.base -t $(IMAGE_NAME_BASE)-base .
 
@@ -91,7 +91,7 @@ gantry-run-ib :
 		--env NCCL_DEBUG=INFO \
 		--env SCRATCH_DIR=/tmp/scratch \
 		--env FLASH_DIR=/tmp/flash \
-		--env WANDB_PROJECT=dolma-beaker-ib \
+		--env WANDB_PROJECT=olmo-beaker-ib \
 		--env-secret WANDB_API_KEY=WANDB_API_KEY \
 		--replicas 4 \
 		--leader-selection \
@@ -103,8 +103,8 @@ gantry-run-ib :
 
 .PHONY : check-cpu-install
 check-cpu-install :
-	@python -c 'from dolma import check_install; check_install(cuda=False)'
+	@python -c 'from olmo import check_install; check_install(cuda=False)'
 
 .PHONY : check-cuda-install
 check-cuda-install :
-	@python -c 'from dolma import check_install; check_install(cuda=True)'
+	@python -c 'from olmo import check_install; check_install(cuda=True)'

README.md

@@ -1,4 +1,4 @@
-# DOLMA: Delightful Open Language Model from AI2
+# OLMo: Delightful Open Language Model from AI2
 
 ## Setup
 
@@ -24,7 +24,7 @@ gantry run \
   --nfs \
   --priority preemptible \
   --gpus 8 \
-  --beaker-image dolma-torch2-gantry \
+  --beaker-image olmo-torch2-gantry \
   --cluster 'ai2/*-cirrascale' \
   --allow-dirty \
   -- composer scripts/train.py configs/1.2b-c4.yaml
@@ -36,7 +36,7 @@ Train the 70B model on c4 with gantry across multiple nodes:
 gantry run \
   --workspace ai2/llm-testing \
   --priority "high" \
-  --beaker-image dolma-torch2-gantry \
+  --beaker-image olmo-torch2-gantry \
   --cluster ai2/general-cirrascale-a100-80g-ib \
   --gpus 8 \
   --nfs \
@@ -45,7 +45,7 @@ gantry run \
   --env NCCL_DEBUG=INFO \
   --env SCRATCH_DIR=/tmp/scratch \
   --env FLASH_DIR=/tmp/flash \
-  --env WANDB_PROJECT=dolma-beaker-ib \
+  --env WANDB_PROJECT=olmo-beaker-ib \
   --env-secret WANDB_API_KEY=WANDB_API_KEY \
   --replicas 4 \
   --leader-selection \

RELEASE_PROCESS.md

@@ -2,7 +2,7 @@
 
 ## Steps
 
-1. Update the version in `dolma/version.py`.
+1. Update the version in `olmo/version.py`.
 
 3. Run the release script:
 

conftest.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from dolma.config import (
+from olmo.config import (
     DataConfig,
     ModelConfig,
     OptimizerConfig,
@@ -11,7 +11,7 @@
     TokenizerConfig,
     TrainConfig,
 )
-from dolma.tokenizer import Tokenizer
+from olmo.tokenizer import Tokenizer
 
 TEST_MODEL = "gpt2"
 

docker/Dockerfile.gantry

@@ -4,11 +4,11 @@
 # To build and push the image to Beaker, run 'make gantry-image'.
 # To test the image after pushing to Beaker, run 'make gantry-test'.
 
-FROM dolma-torch2-base
+FROM olmo-torch2-base
 
 WORKDIR /stage
 
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-WORKDIR /app/dolma
+WORKDIR /app/olmo

docker/Dockerfile.test

@@ -4,7 +4,7 @@
 #
 # To build and push the image to Beaker, run 'make test-image'.
 
-FROM dolma-torch2-base
+FROM olmo-torch2-base
 
 COPY scripts/test_entrypoint.sh /entrypoint.sh
 RUN chmod +x /entrypoint.sh

dolma/__init__.py → olmo/__init__.py

@@ -12,4 +12,4 @@ def check_install(cuda: bool = False):
         assert torch.cuda.is_available(), "CUDA is not available!"
         print("CUDA available")
 
-    print(f"DOLMA v{VERSION} installed")
+    print(f"OLMo v{VERSION} installed")

dolma/composer.py → olmo/composer.py

@@ -23,16 +23,16 @@
     TrainConfig,
 )
 from .data import DataCollator, MemMapDataset
-from .exceptions import DolmaConfigurationError
-from .model import Dolma, LayerNormBase
+from .exceptions import OlmoConfigurationError
+from .model import LayerNormBase, Olmo
 from .optim import DecoupledLionW
 
 log = logging.getLogger(__name__)
 
 __all__ = [
     "TrainBatchPerplexity",
-    "ComposerDolmaLM",
-    "DolmaConsoleLogger",
+    "ComposerOlmoLM",
+    "OlmoConsoleLogger",
     "build_dataloader",
     "build_optimizer",
     "build_scheduler",
@@ -76,10 +76,10 @@ def compute(self) -> torch.Tensor:
         return torch.exp(self.loss)
 
 
-class ComposerDolmaLM(ComposerModel):
-    def __init__(self, model_or_config: Union[Dolma, ModelConfig]):
+class ComposerOlmoLM(ComposerModel):
+    def __init__(self, model_or_config: Union[Olmo, ModelConfig]):
         super().__init__()
-        self.model = Dolma(model_or_config) if isinstance(model_or_config, ModelConfig) else model_or_config
+        self.model = Olmo(model_or_config) if isinstance(model_or_config, ModelConfig) else model_or_config
         self.config = self.model.config
         self.num_fwd_flops = self.model.num_fwd_flops
 
@@ -131,7 +131,7 @@ def flops_per_batch(self, batch: BatchDict):
         return self.num_fwd_flops * 3 * batch["input_ids"].shape[0]
 
 
-class DolmaConsoleLogger(ConsoleLogger):
+class OlmoConsoleLogger(ConsoleLogger):
     metrics_to_log: Set[str] = {"loss/train/total", "trainer/global_step", "metrics/*"}
 
     def log_metrics(self, metrics: dict[str, float], step: Optional[int] = None) -> None:
@@ -285,7 +285,7 @@ def calculate_batch_size_info(
     global_batch_size: int, device_microbatch_size: Union[int, str]
 ) -> Tuple[int, Union[str, int], Union[str, int]]:
     if global_batch_size % dist.get_world_size() != 0:
-        raise DolmaConfigurationError(
+        raise OlmoConfigurationError(
             f"Global batch size {global_batch_size} is not divisible by {dist.get_world_size()} "
             "as a result, the batch size would be truncated, please adjust `global_batch_size` "
             f"to be divisible by world size, {dist.get_world_size()}."
@@ -303,7 +303,7 @@ def calculate_batch_size_info(
             device_microbatch_size = device_batch_size
         device_grad_accum = math.ceil(device_batch_size / device_microbatch_size)
     else:
-        raise DolmaConfigurationError(f"Not sure how to parse {device_microbatch_size=}")
+        raise OlmoConfigurationError(f"Not sure how to parse {device_microbatch_size=}")
 
     return device_batch_size, device_microbatch_size, device_grad_accum
 
@@ -324,7 +324,7 @@ def update_batch_size_info(cfg: TrainConfig):
         elif isinstance(cfg.device_train_microbatch_size, int):
             cfg.device_eval_batch_size = cfg.device_train_microbatch_size
         else:
-            raise DolmaConfigurationError(
+            raise OlmoConfigurationError(
                 f"Not sure how to parse device_train_microbatch_size={cfg.device_train_microbatch_size}"
             )
     return cfg

dolma/config.py → olmo/config.py

@@ -22,7 +22,7 @@
 from omegaconf.errors import OmegaConfBaseException
 
 from .aliases import PathOrStr
-from .exceptions import DolmaConfigurationError
+from .exceptions import OlmoConfigurationError
 
 __all__ = [
     "ActivationType",
@@ -90,7 +90,7 @@ def new(cls: Type[C], **kwargs) -> C:
                 conf = om.merge(conf, kwargs)
             return cast(C, om.to_object(conf))
         except OmegaConfBaseException as e:
-            raise DolmaConfigurationError(str(e))
+            raise OlmoConfigurationError(str(e))
 
     @classmethod
     def load(cls: Type[C], path: PathOrStr, overrides: Optional[List[str]] = None) -> C:
@@ -103,7 +103,7 @@ def load(cls: Type[C], path: PathOrStr, overrides: Optional[List[str]] = None) -
                 conf = om.merge(conf, om.from_dotlist(overrides))
             return cast(C, om.to_object(conf))
         except OmegaConfBaseException as e:
-            raise DolmaConfigurationError(str(e))
+            raise OlmoConfigurationError(str(e))
 
     def save(self, path: PathOrStr) -> None:
         """Save to a YAML file."""
@@ -155,7 +155,7 @@ class BlockType(StrEnum):
 @dataclass
 class ModelConfig(BaseConfig):
     """
-    DOLMA (model) configuration.
+    OLMo (model) configuration.
     """
 
     # Note that the defaults for these attributes are equivalent to the base GPT2 model.
@@ -401,7 +401,7 @@ class CompilerConfig(BaseConfig):
 @dataclass
 class TrainConfig(BaseConfig):
     """
-    DOLMA training configuration.
+    OLMo training configuration.
     """
 
     run_name: Optional[str] = None

olmo/exceptions.py

@@ -0,0 +1,19 @@
+__all__ = ["OlmoError", "OlmoConfigurationError", "OlmoCliError"]
+
+
+class OlmoError(Exception):
+    """
+    Base class for all custom OLMo exceptions.
+    """
+
+
+class OlmoConfigurationError(OlmoError):
+    """
+    An error with a configuration file.
+    """
+
+
+class OlmoCliError(OlmoError):
+    """
+    An error from incorrect CLI usage.
+    """