-
Notifications
You must be signed in to change notification settings - Fork 87
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge remote-tracking branch 'origin/main' into dclm
- Loading branch information
Showing
133 changed files
with
10,270 additions
and
5,177 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| name: Launch Llama 2 Small Fast | ||
|
|
||
| on: | ||
| workflow_run: | ||
| workflows: ["Build and Push Docker TPU Images"] | ||
| types: | ||
| - completed | ||
| branches: [main, "experiment/*"] | ||
| # pull_request: | ||
| workflow_dispatch: | ||
|
|
||
| jobs: | ||
| test: | ||
| if: (github.event.pull_request.head.repo.full_name == github.repository) | ||
| runs-on: ubuntu-latest | ||
| env: | ||
| TPU_ZONE: "us-central2-b" | ||
| TPU_TYPE: "v4-32" | ||
|
|
||
| steps: | ||
| - name: Checkout code | ||
| uses: actions/checkout@v2 | ||
|
|
||
| - name: Set up Google Cloud SDK | ||
| uses: google-github-actions/setup-gcloud@v1 | ||
| with: | ||
| project_id: ${{ secrets.GCP_PROJECT_ID }} | ||
|
|
||
| - name: Authenticate to Google Cloud | ||
| uses: google-github-actions/auth@v1 | ||
| with: | ||
| credentials_json: ${{ secrets.GCP_SA_KEY }} | ||
|
|
||
| - name: Configure Google Cloud | ||
| run: | | ||
| gcloud config set project ${{ secrets.GCP_PROJECT_ID }} | ||
| REGION=${TPU_ZONE%-*} | ||
| echo "$REGION" | ||
| gcloud auth configure-docker $REGION-docker.pkg.dev | ||
| - name: Install locally | ||
| run: | | ||
| python -m pip install --upgrade pip | ||
| pip install -e .[test] "jax[cpu]==0.4.30" | ||
| - name: Launch Small Fast TPU Train LM job | ||
| run: | | ||
| export TPU_NAME=small-fast-${{ github.run_id }} | ||
| export WANDB_API_KEY=${{ secrets.WANDB_API_KEY }} | ||
| export RUN_ID=small_fast_${{ github.run_id }} | ||
| export HF_TOKEN=${{ secrets.HF_TOKEN }} | ||
| cat > .config <<EOF | ||
| env: | ||
| WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} | ||
| WANDB_ENTITY: stanford-mercury | ||
| WANDB_PROJECT: levanter | ||
| HF_TOKEN: ${{ secrets.HF_TOKEN }} | ||
| EOF | ||
| python infra/launch.py -e CI 1 --foreground --tpu_name ${TPU_NAME} --run_id $RUN_ID --zone ${TPU_ZONE} --tpu_type ${TPU_TYPE} --preemptible -- \ | ||
| python -m levanter.main.train_lm \ | ||
| --config_path config/llama_small_fast.yaml \ | ||
| --trainer.checkpointer.base_path gs://levanter-checkpoints/llama-itest/ \ | ||
| --trainer.checkpointer.save_interval 10m | ||
| --trainer.num_train_steps 10000 | ||
| - name: Cleanup | ||
| if: ${{ always() }} | ||
| run: | | ||
| export TPU_NAME=small-fast-${{ github.run_id }} | ||
| gcloud compute tpus queued-resources delete $TPU_NAME --zone ${TPU_ZONE} --quiet --force |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -150,5 +150,9 @@ ledger.json | |
| /checkpoints | ||
| *.jaxpr | ||
|
|
||
| # local execution commands | ||
| local_*.sh | ||
|
|
||
| # aider | ||
| .aider* | ||
|
|
||
| .benchmarks | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| data: | ||
| configs: | ||
| owt: | ||
| train_urls: | ||
| - "gs://pubmed-mosaic/openwebtext-sharded/openwebtext_train.{1..128}-of-128.jsonl.gz" | ||
| validation_urls: | ||
| - "gs://pubmed-mosaic/openwebtext-sharded/openwebtext_val.{1..8}-of-8.jsonl.gz" | ||
| wikitext: | ||
| id: dlwh/wikitext_103_detokenized | ||
| train_weights: | ||
| owt: 0.6 | ||
| wikitext: 0.4 | ||
| tokenizer: gpt2 | ||
| cache_dir: "gs://levanter-data/tokenized/data_mix" | ||
| supervised_data: | ||
| validation_urls: | ||
| - "gs://marin-us-central2/benchmarks/mmlu/mmlu-*-dev-evaluation.jsonl.gz" | ||
| cache_dir: "gs://marin-us-central2/benchmarks/tokenized-gpt2/mmlu/" | ||
| model: | ||
| type: gpt2 | ||
| hidden_dim: 768 | ||
| num_heads: 12 | ||
| num_layers: 12 | ||
| seq_len: 1024 | ||
| gradient_checkpointing: true | ||
| scale_attn_by_inverse_layer_idx: true | ||
| trainer: | ||
| tracker: | ||
| project: "levanter" | ||
| tags: [ "openwebtext+wiki", "gpt2", "itest"] | ||
|
|
||
| mp: p=f32,c=bfloat16 | ||
| model_axis_size: 1 | ||
|
|
||
| train_batch_size: 256 | ||
| num_train_steps: 20000 | ||
| optimizer: | ||
| learning_rate: 1E-3 | ||
| weight_decay: 0.1 | ||
| warmup: 0.01 |
Oops, something went wrong.