Documentation Improvements

.gitignore

@@ -32,7 +32,6 @@ doc/_build/
 *.swp
 .DS_Store
 
-
 # python
 
 *.pyc

README.md

@@ -17,23 +17,20 @@
   </a>
 </p>
 
-OLMo is a repository for training and using AI2's state-of-the-art open language models. 
-It is built by scientists, for scientists.
+OLMo is a repository for training and using AI2's state-of-the-art open language models. It is designed by scientists, for scientists.
 
 ## Installation
 
-First install [PyTorch](https://pytorch.org) according to the instructions specific to your operating system.
+First, install [PyTorch](https://pytorch.org) following the instructions specific to your operating system.
 
-To install from source (recommended for training/fine-tuning) run:
+For training and fine-tuning, we recommend installing from source:
 
 ```bash
 git clone https://github.com/allenai/OLMo.git
 cd OLMo
 pip install -e .[all]
 ```
-
-Otherwise you can install the model code by itself directly from PyPI with:
-
+You can also install from PyPI with:
 ```bash
 pip install ai2-olmo
 ```
@@ -42,38 +39,31 @@ pip install ai2-olmo
 
 ### Overview
 
-The core models in the OLMo family released so far are (all trained on the [Dolma dataset](https://huggingface.co/datasets/allenai/dolma)): 
-| Model | Training Tokens | Context Length | Training Config | W&B Logs | Data Order File(s) ☨ |
-|-------|-----------------|:--------------:|-----------------|----------|--------------------|
-| [OLMo 1B](https://huggingface.co/allenai/OLMo-1B) | 3 Trillion | 2048 | [configs/official/OLMo-1B.yaml](https://github.com/allenai/OLMo/blob/main/configs/official/OLMo-1B.yaml) | [wandb.ai/…/OLMo-1B](https://wandb.ai/ai2-llm/OLMo-1B/reports/OLMo-1B--Vmlldzo2NzY1Njk1) | [epoch 1](https://olmo-checkpoints.org/ai2-llm/olmo-small/46zc5fly/train_data/global_indices.npy) |
-| [OLMo 7B](https://huggingface.co/allenai/OLMo-7B) | 2.5 Trillion | 2048 | [configs/official/OLMo-7B.yaml](https://github.com/allenai/OLMo/blob/main/configs/official/OLMo-7B.yaml) | [wandb.ai/…/OLMo-7B](https://wandb.ai/ai2-llm/OLMo-7B/reports/OLMo-7B--Vmlldzo2NzQyMzk5) | [epoch 1](https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy), [epoch 2](https://olmo-checkpoints.org/ai2-llm/olmo-medium/wd2gxrza/train_data/global_indices.npy) |
-| [OLMo 7B Twin 2T](https://huggingface.co/allenai/OLMo-7B-Twin-2T) | 2 Trillion  | 2048 | [configs/official/OLMo-7B.yaml](https://github.com/allenai/OLMo/blob/main/configs/official/OLMo-7B.yaml) | [wandb.ai/…/OLMo-7B-Twin-2T](https://wandb.ai/ai2-llm/OLMo-7B/reports/OLMo-7B-Twin-2T--Vmlldzo2NzU0NTIz) | [epoch 1](https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy) |
-| [OLMo 7B April 2024](https://huggingface.co/allenai/OLMo-7B-0424-hf) | 2.05 Trillion  | 4096 | [configs/official/OLMo-7B-0424.yaml](https://github.com/allenai/OLMo/blob/main/configs/official/OLMo-7B-0424.yaml) | *Coming soon* | *Coming soon* |
-| [OLMo 7B July 2024](https://huggingface.co/allenai/OLMo-7B-0724-hf) | 2.75 Trillion  | 4096 | [configs/official/OLMo-7B-0724.yaml](https://github.com/allenai/OLMo/blob/main/configs/official/OLMo-7B-0724.yaml) | *Coming soon* | *Coming soon* |
-
-> ☨ *See [Inspecting training data](#inspecting-training-data) below for usage.*
+The core models in the OLMo family released are: 
+| Model | Training Tokens | Context Length | Training Config | W&B Logs |
+|-------|-----------------|:--------------:|-----------------|----------|
+| [OLMo2 7B](https://huggingface.co/allenai/OLMo-2-1124-7B) | 4 Trillion | 4096 | [configs/official-1124/OLMo2-7B-stage1.yaml](https://github.com/allenai/OLMo/blob/main/configs/official-1124/OLMo2-7B-stage1.yaml) | wandb.ai/…/OLMo2-7B (link to come)
+| [OLMo2 13B](https://huggingface.co/allenai/OLMo-2-1124-13B) | 5 Trillion | 4096 | [configs/official-1124/OLMo2-12B-stage1.yaml](https://github.com/allenai/OLMo/blob/main/configs/official-1124/OLMo2-13B-stage1.yaml) | wandb.ai/…/OLMo2-13B (link to come)
 
-### Checkpoints
+For instruction tuned variants of these models, go to
+ * [OLMo2 7B Instruct](https://huggingface.co/allenai/OLMo-2-1124-7B-Instruct)
+ * [OLMo2 13B Instruct](https://huggingface.co/allenai/OLMo-2-1124-13B-Instruct)
 
-URLs to checkpoints at intermediate steps of the models' trainings can be found in the csv files under [`checkpoints/official/`](https://github.com/allenai/OLMo/blob/main/checkpoints/official). These 'directory' URLs cannot currently be directly accessed, but files within the directory are publicly accessible. These URLs can also be provided to the training script to resume training from the checkpoint (see [Training](#training)). Each checkpoint directory consists of:
-
-- `config.yaml`: the config at that training step.
-- `model.pt`, `optim.pt`, `train.pt`: model, optimizer and training state at that training step.
-
-Details about the other types of OLMo checkpoints (including OLMo HF Transformers checkpoints) can be found in [Checkpoints.md](https://github.com/allenai/OLMo/blob/main/docs/Checkpoints.md).
+> ☨ *See [Inspecting training data](#inspecting-training-data) below for usage.*
 
 ## Inference
 
-You can utilize our Hugging Face integration to run inference on the OLMo Transformers checkpoints:
+You can use our Hugging Face integration to run inference on the OLMo Transformers checkpoints:
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
-
-olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-0724-hf")
-tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-7B-0724-hf")
-
+olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-1124-7B")
+tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-1124-7B")
 message = ["Language modeling is "]
 inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
+# optional verifying cuda
+# inputs = {k: v.to('cuda') for k,v in inputs.items()}
+# olmo = olmo.to('cuda')
 response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
 print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
 ```
@@ -82,129 +72,21 @@ Alternatively, with the Hugging Face pipeline abstraction:
 
 ```python
 from transformers import pipeline
-olmo_pipe = pipeline("text-generation", model="allenai/OLMo-7B-0724-hf")
+olmo_pipe = pipeline("text-generation", model="allenai/OLMo-2-1124-7B")
 print(olmo_pipe("Language modeling is"))
 ```
 
-### Inference on finetuned checkpoints
-
-If you finetune the model using the code in [Fine-tuning](#fine-tuning), you can use the conversion script to convert a native OLMo checkpoint to a Hugging Face-compatible checkpoint.
-
-```bash
-python scripts/convert_olmo_to_hf_new.py --input_dir /path/to/olmo/checkpoint --output_dir /path/to/hf/checkpoint/ --tokenizer_json_path tokenizers/allenai_gpt-neox-olmo-dolma-v1_5.json
-```
-
 ### Quantization
 
 ```python
-olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-0724-hf", torch_dtype=torch.float16, load_in_8bit=True)  # requires bitsandbytes
-```
-
-The quantized model is more sensitive to typing / cuda, so it is recommended to pass the inputs as inputs.input_ids.to('cuda') to avoid potential issues.
-
-## Reproducibility
-
-### Training
-
-The configs used to train the official OLMo models are provided in the [`configs/official/`](https://github.com/allenai/OLMo/blob/main/configs/official) directory.
-
-Note that while the training and validation data is public and free to download, the paths to the data within those configs are pointed at a CloudFlare R2 bucket, which requires an API key for programmatic access.
-So in order to use any of these configs to reproduce a training run you'll first have to download the corresponding data to a location of your choosing and then update the paths in the config accordingly.
-
-You can derive the public HTTP URL from an R2 URL by replacing `r2://olmo-data` with `https://olmo-data.org`.
-For example, if the R2 data URL is:
-
-`r2://olmo-data/preprocessed/olmo-mix/v1_5/gpt-neox-20b-pii-special/part-000-00000.npy`
-
-then the corresponding public URL is:
-
-`https://olmo-data.org/preprocessed/olmo-mix/v1_5/gpt-neox-20b-pii-special/part-000-00000.npy`
-
-Once you've updated the data paths in the config you can launch a training run via `torchrun`. For example, to launch the 1B model training on a single 8x GPU node, you would run:
-
-```bash
-torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml
-```
-
-You can use the same method to launch multi-node jobs as well. See [the documentation](https://pytorch.org/docs/stable/elastic/run.html) for `torchrun` to understand the additional arguments you'll need to configure the rendezvous backend / endpoint.
-
-To resume training from a checkpoint, you can pass its path (local or URL)
-to `scripts/train.py` with the `--load_path` arguments. For example, to resume training from step 1000 of the OLMo 1B run:
-
-```bash
-torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml --load_path=https://olmo-checkpoints.org/ai2-llm/olmo-small/w1r5xfzt/step1000-unsharded
-```
-
-### Inspecting training data
-
-You may be interested in inspecting the exact tokens that composed a particular batch during the training of one of the OLMo models.
-We provide tools to do this, but first you'll need to download the data as above (unless you have an R2 API key) and update the corresponding config accordingly.
-
-Then take note of the URL of the data order file you want, which can be found in the [Models Overview](#models-overview) table. For example, the data order file for the first epoch of the OLMo-7B model is [https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy](https://olmo-checkpoints.org/ai2-llm/olmo-small/46zc5fly/train_data/global_indices.npy).
-
-Once you have that you can use this snippet to inspect the data within a particular batch:
-
-```python
-import numpy as np
-from cached_path import cached_path
-
-from olmo.config import TrainConfig
-from olmo.data import build_memmap_dataset
-
-# Update these paths to what you want:
-data_order_file_path = cached_path("https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy")
-train_config_path = "configs/official/OLMo-7B.yaml"
-
-
-cfg = TrainConfig.load(train_config_path)
-dataset = build_memmap_dataset(cfg, cfg.data)
-batch_size = cfg.global_train_batch_size
-global_indices = np.memmap(data_order_file_path, mode="r+", dtype=np.uint32)
-
-
-def get_batch_instances(batch_idx: int) -> list[list[int]]:
-    batch_start = batch_idx * batch_size
-    batch_end = (batch_idx + 1) * batch_size
-    batch_indices = global_indices[batch_start:batch_end]
-    batch_instances = []
-    for index in batch_indices:
-        token_ids = dataset[index]["input_ids"].tolist()
-        batch_instances.append(token_ids)
-    return batch_instances
-
-
-# Get all 2048 x 2048 token IDs in the first batch.
-get_batch_instances(0)
-```
-
-
-## Fine-tuning
-
-To fine-tune an OLMo model using our trainer you'll first need to prepare your dataset by tokenizing it and saving the tokens IDs to a flat numpy memory-mapped array. See [`scripts/prepare_tulu_data.py`](./scripts/prepare_tulu_data.py) for an example with the Tulu V2 dataset, which can be easily modified for other datasets.
-
-Next, prepare your training config. There are many examples in the [`configs/`](https://github.com/allenai/OLMo/blob/main/configs) directory that you can use as a starting point. The most important thing is to make sure the model parameters (the `model` field in the config) match up with the checkpoint you're starting from. To be safe you can always start from the config that comes with the model checkpoint. At a minimum you'll need to make the following changes to the config or provide the corresponding overrides from the command line:
-
-- Update `load_path` to point to the checkpoint you want to start from.
-- Set `reset_trainer_state` to `true`.
-- Update `data.paths` to point to the `token_ids.npy` file you generated.
-- Optionally update `data.label_mask_paths` to point to the `label_mask.npy` file you generated, unless you don't need special masking for the loss.
-- Update `evaluators` to add/remove in-loop evaluations.
-
-Once you're satisfied with your training config, you can launch the training job via `torchrun`. For example:
-
-```
-torchrun --nproc_per_node=8 scripts/train.py {path_to_train_config} \
-    --data.paths=[{path_to_data}/input_ids.npy] \
-    --data.label_mask_paths=[{path_to_data}/label_mask.npy] \
-    --load_path={path_to_checkpoint} \
-    --reset_trainer_state
+olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-1124-7B", torch_dtype=torch.float16, load_in_8bit=True)  # requires bitsandbytes
 ```
 
-Note: passing CLI overrides like `--reset_trainer_state` is only necessary if you didn't update those fields in your config.
+The quantized model is sensitive to input types and CUDA handling. To avoid potential issues, we recommend explicitly converting input IDs to CUDA using: `inputs.input_ids.to('cuda')`
 
 ## Evaluation
 
-Additional tools for evaluating OLMo models are available at the [OLMo Eval](https://github.com/allenai/ai2-olmo-eval) repo.
+Additional tools for evaluating OLMo models are available at the [OLMo Eval](https://github.com/allenai/OLMo-eval) repo.
 
 ## Debugging