├── .gitignore ├── README.md ├── configs ├── experiment │ ├── 124M_ado.yaml │ ├── 124M_balanced.yaml │ ├── 124M_natural.yaml │ ├── 1B_ado.yaml │ ├── 1B_balanced.yaml │ └── 1B_natural.yaml ├── loader │ ├── pile_ado.yaml │ ├── pile_balanced.yaml │ ├── pile_doremi.yaml │ ├── pile_natural.yaml │ └── tasks_sampler_cfg │ │ └── pile_tasks.yaml ├── main.yaml └── model │ ├── 124M.yaml │ ├── 70M.yaml │ ├── nano.yaml │ └── xl.yaml ├── launch.py ├── procdata ├── build_tfds_configs.sh ├── process_dataset.py ├── shard_files.py └── viz_jsonl_zst.py ├── requirements.txt ├── scripts ├── manual_build_tfds.py ├── tpu_commands.sh └── tpu_vm_setup.sh └── src ├── __init__.py ├── data_selectors ├── __init__.py └── ado.py ├── dataloader ├── __init__.py ├── abstract.py ├── adaptive_loader.py └── pile_loader.py ├── eval ├── __init__.py ├── llm_harness_jax_utils.py └── lm_harness_eval.py ├── layers.py ├── model.py ├── sharding.py ├── tfds ├── __init__.py ├── conversion_utils.py └── the_pile_grouped.py └── train.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/.gitignore -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/README.md -------------------------------------------------------------------------------- /configs/experiment/124M_ado.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/124M_ado.yaml -------------------------------------------------------------------------------- /configs/experiment/124M_balanced.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/124M_balanced.yaml -------------------------------------------------------------------------------- /configs/experiment/124M_natural.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/124M_natural.yaml -------------------------------------------------------------------------------- /configs/experiment/1B_ado.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/1B_ado.yaml -------------------------------------------------------------------------------- /configs/experiment/1B_balanced.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/1B_balanced.yaml -------------------------------------------------------------------------------- /configs/experiment/1B_natural.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/experiment/1B_natural.yaml -------------------------------------------------------------------------------- /configs/loader/pile_ado.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/loader/pile_ado.yaml -------------------------------------------------------------------------------- /configs/loader/pile_balanced.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/loader/pile_balanced.yaml -------------------------------------------------------------------------------- /configs/loader/pile_doremi.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/loader/pile_doremi.yaml -------------------------------------------------------------------------------- /configs/loader/pile_natural.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/loader/pile_natural.yaml -------------------------------------------------------------------------------- /configs/loader/tasks_sampler_cfg/pile_tasks.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/loader/tasks_sampler_cfg/pile_tasks.yaml -------------------------------------------------------------------------------- /configs/main.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/main.yaml -------------------------------------------------------------------------------- /configs/model/124M.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/model/124M.yaml -------------------------------------------------------------------------------- /configs/model/70M.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/model/70M.yaml -------------------------------------------------------------------------------- /configs/model/nano.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/model/nano.yaml -------------------------------------------------------------------------------- /configs/model/xl.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/configs/model/xl.yaml -------------------------------------------------------------------------------- /launch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/launch.py -------------------------------------------------------------------------------- /procdata/build_tfds_configs.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/procdata/build_tfds_configs.sh -------------------------------------------------------------------------------- /procdata/process_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/procdata/process_dataset.py -------------------------------------------------------------------------------- /procdata/shard_files.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/procdata/shard_files.py -------------------------------------------------------------------------------- /procdata/viz_jsonl_zst.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/procdata/viz_jsonl_zst.py -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/requirements.txt -------------------------------------------------------------------------------- /scripts/manual_build_tfds.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/scripts/manual_build_tfds.py -------------------------------------------------------------------------------- /scripts/tpu_commands.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/scripts/tpu_commands.sh -------------------------------------------------------------------------------- /scripts/tpu_vm_setup.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/scripts/tpu_vm_setup.sh -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/data_selectors/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/data_selectors/ado.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/data_selectors/ado.py -------------------------------------------------------------------------------- /src/dataloader/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/dataloader/abstract.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/dataloader/abstract.py -------------------------------------------------------------------------------- /src/dataloader/adaptive_loader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/dataloader/adaptive_loader.py -------------------------------------------------------------------------------- /src/dataloader/pile_loader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/dataloader/pile_loader.py -------------------------------------------------------------------------------- /src/eval/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/eval/llm_harness_jax_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/eval/llm_harness_jax_utils.py -------------------------------------------------------------------------------- /src/eval/lm_harness_eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/eval/lm_harness_eval.py -------------------------------------------------------------------------------- /src/layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/layers.py -------------------------------------------------------------------------------- /src/model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/model.py -------------------------------------------------------------------------------- /src/sharding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/sharding.py -------------------------------------------------------------------------------- /src/tfds/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/tfds/conversion_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/tfds/conversion_utils.py -------------------------------------------------------------------------------- /src/tfds/the_pile_grouped.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/tfds/the_pile_grouped.py -------------------------------------------------------------------------------- /src/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yidingjiang/ado/HEAD/src/train.py --------------------------------------------------------------------------------