Package: ai2-olmo
OLMo: Open Language ModelOLMo is a repository for training and using AI2's state-of-the-art open language models.
It is built by scientists, for scientists.InstallationFirst installPyTorchaccording to the instructions specific to your operating system.To install from source (recommended for training/fine-tuning) run:gitclonehttps://github.com/allenai/OLMo.gitcdOLMo
pipinstall-e.[all]Otherwise you can install the model code by itself directly from PyPI with:pipinstallai2-olmoModels overviewThe core models in the OLMo family released so far are (all trained on theDolma dataset):ModelTraining TokensContext LengthOLMo 1B3 Trillion2048OLMo 7B2.5 Trillion2048OLMo 7B Twin 2T2 Trillion2048Fine-tuningTo 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. Seescripts/prepare_tulu_data.pyfor 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 theconfigs/directory that you can use as a starting point. The most important thing is to make sure the model parameters (themodelfield 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:Updateload_pathto point to the checkpoint you want to start from.Setreset_trainer_statetotrue.Updatedata.pathsto point to thetoken_ids.npyfile you generated.Optionally updatedata.label_mask_pathsto point to thelabel_mask.npyfile you generated, unless you don't need special masking for the loss.Updateevaluatorsto add/remove in-loop evaluations.Once you're satisfied with your training config, you can launch the training job viatorchrun. 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_stateNote: passing CLI overrides like--reset_trainer_stateis only necessary if you didn't update those fields in your config.InferenceYou can utilize our HuggingFace integration to run inference on the olmo checkpoints:fromhf_olmoimport*# registers the Auto* classesfromtransformersimportAutoModelForCausalLM,AutoTokenizerolmo=AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B")tokenizer=AutoTokenizer.from_pretrained("allenai/OLMo-7B")message=["Language modeling is "]inputs=tokenizer(message,return_tensors='pt',return_token_type_ids=False)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])Alternatively, with the huggingface pipeline abstraction:fromtransformersimportpipelineolmo_pipe=pipeline("text-generation",model="allenai/OLMo-7B")print(olmo_pipe("Language modeling is"))Inference on finetuned checkpointsIf you finetune the model using the code above, you can use the conversion script to convert a native OLMo checkpoint to a HuggingFace-compatible checkpointpythonhf_olmo/convert_olmo_to_hf.py--checkpoint-dir/path/to/checkpointQuantizationolmo=AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B",torch_dtype=torch.float16,load_in_8bit=True)# requires bitsandbytesThe 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.EvaluationAdditional tools for evaluating OLMo models are available at theOLMo Evalrepo.