Merge branch 'main' into chuck/te_hf_ckpt

.github/workflows/pr-cpu.yaml

@@ -23,6 +23,10 @@ jobs:
           container: mosaicml/pytorch:2.3.0_cpu-python3.11-ubuntu20.04
           markers: "not gpu"
           pytest_command: "coverage run -m pytest"
+        - name: "cpu-2.3.1"
+          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+          markers: "not gpu"
+          pytest_command: "coverage run -m pytest"
     name: ${{ matrix.name }}
     if: github.repository_owner == 'mosaicml'
     with:

.github/workflows/pr-gpu.yaml

@@ -24,6 +24,11 @@ jobs:
           markers: "gpu"
           pytest_command: "coverage run -m pytest"
           pip_deps: "[all]"
+        - name: "gpu-2.3.1"
+          container: mosaicml/llm-foundry:2.3.1_cu121-latest
+          markers: "gpu"
+          pytest_command: "coverage run -m pytest"
+          pip_deps: "[all]"
     name: ${{ matrix.name }}
     if: github.repository_owner == 'mosaicml'
     with:

README.md

@@ -230,7 +230,7 @@ python data_prep/convert_dataset_hf.py \
 # Train an MPT-125m model for 10 batches
 composer train/train.py \
   train/yamls/pretrain/mpt-125m.yaml \
-  data_local=my-copy-c4 \
+  variables.data_local=my-copy-c4 \
   train_loader.dataset.split=train_small \
   eval_loader.dataset.split=val_small \
   max_duration=10ba \

llmfoundry/eval/datasets/__init__.py

@@ -22,6 +22,18 @@
     tokenizer_needs_prefix_space,
     trim_context,
 )
+from llmfoundry.registry import icl_datasets
+
+icl_datasets.register(
+    'multiple_choice',
+    func=InContextLearningMultipleChoiceTaskDataset,
+)
+icl_datasets.register('schema', func=InContextLearningSchemaTaskDataset)
+icl_datasets.register('language_modeling', func=InContextLearningLMTaskDataset)
+icl_datasets.register(
+    'generation_task_with_answers',
+    func=InContextLearningGenerationTaskWithAnswersDataset,
+)
 
 __all__ = [
     'InContextLearningDataset',