Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Mitchish65 #388

Merged
merged 72 commits into from
Jan 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
72 commits
Select commit Hold shift + click to select a range
f661bfd
Adds a 65B config
dirkgr Nov 15, 2023
28cf940
Make the config actually 65B
dirkgr Nov 16, 2023
0c122e9
More parameters
dirkgr Nov 16, 2023
5401904
Set to the most memory efficient setting
dirkgr Nov 20, 2023
a4af8af
Change some names
dirkgr Nov 20, 2023
cc72acd
Default to local checkpointing
dirkgr Nov 20, 2023
c609fb4
Script for the 65b config
dirkgr Nov 20, 2023
addc631
Put some learnings into the script
dirkgr Nov 21, 2023
135cf11
Decay everything
dirkgr Nov 21, 2023
ce380b4
New image version for logging into nodes
dirkgr Nov 21, 2023
518b2e1
Merge remote-tracking branch 'origin/epwalsh/threaded-data-loading' i…
dirkgr Nov 21, 2023
2b59d48
Read the longrunfix data
dirkgr Nov 21, 2023
12fa63c
Mokey patch FSDP stream management
dirkgr Nov 21, 2023
74a7670
3 streams
dirkgr Nov 21, 2023
ddd69c5
Revert "3 streams"
dirkgr Nov 21, 2023
36f2b17
Revert "Mokey patch FSDP stream management"
dirkgr Nov 21, 2023
f52fa18
Smaller global batch for batch size warmup
dirkgr Nov 21, 2023
8952655
Grad clipping warmup
dirkgr Nov 22, 2023
573c0c0
Adjust the run script for batch size warmup
dirkgr Nov 22, 2023
14ea8fa
Typo
dirkgr Nov 22, 2023
750202f
Mitchich 35
dirkgr Nov 22, 2023
471945a
Improved 50
dirkgr Nov 22, 2023
addac71
35
dirkgr Nov 22, 2023
6bae22b
Load the model in turns
dirkgr Nov 28, 2023
4ad3982
Merge remote-tracking branch 'origin/main' into mitchish65
dirkgr Nov 28, 2023
caafbec
Fix missing import
dirkgr Nov 30, 2023
5ee7fb6
Double the batch size
dirkgr Nov 30, 2023
2d5f5d5
More nodes
dirkgr Nov 30, 2023
b4626ab
Load twice as fast
dirkgr Nov 30, 2023
966ca47
2T tokens
dirkgr Nov 30, 2023
67592d8
Load models in a separate thread, so we can keep making NCCL calls th…
dirkgr Nov 30, 2023
0b8b3dc
Merge remote-tracking branch 'origin/main' into mitchish65
dirkgr Nov 30, 2023
999a323
Need a device
dirkgr Nov 30, 2023
21dd204
That's not how you call reduce
dirkgr Nov 30, 2023
41aa520
Not sure why positional args don't work
dirkgr Nov 30, 2023
beab360
Forgot `load_path`
dirkgr Nov 30, 2023
71ed346
Got the wrong timeout error
dirkgr Nov 30, 2023
6ab70cc
Safetensors
dirkgr Nov 30, 2023
aa2301d
All ranks load at the same time
dirkgr Dec 2, 2023
bcd41c2
Formatting
dirkgr Dec 2, 2023
ac7a80b
Revert "All ranks load at the same time"
dirkgr Dec 2, 2023
600437d
Now something else is slow
dirkgr Dec 2, 2023
72f5666
I feel like this should have been `False` all along.
dirkgr Dec 2, 2023
893d897
Try restoring checkpoints with `summon_full_params()`
dirkgr Dec 2, 2023
834374c
Fixes how we load from an unsharded checkpoint
dirkgr Dec 4, 2023
2f69e80
Merge remote-tracking branch 'origin/mitchish65-summon' into mitchish65
dirkgr Dec 4, 2023
1d86c8b
Call GC a bunch while loading
dirkgr Dec 5, 2023
8cb5e7c
Run longer, and make time limit work
dirkgr Dec 6, 2023
6e4cbfd
Merge remote-tracking branch 'origin/mitchish65-timelimit' into mitch…
dirkgr Dec 6, 2023
e86e9ba
Try closing and reopening the safetensors file every once in a while
dirkgr Dec 6, 2023
e312b04
Fix counts
dirkgr Dec 6, 2023
5452658
We need even more time it seems
dirkgr Dec 8, 2023
72bdc96
Checkpoint even more often, because we clearly can't keep a job going…
dirkgr Dec 10, 2023
00305eb
Add support for ephemeral checkpoints
epwalsh Dec 12, 2023
a3e1a56
fix typo
epwalsh Dec 12, 2023
1ac7666
fix
epwalsh Dec 13, 2023
5ca9a5d
Merge remote-tracking branch 'origin/mitchish65-timelimit' into mitch…
dirkgr Dec 14, 2023
498afaa
Merge remote-tracking branch 'origin/epwalsh/ephemeral-checkpoints' i…
dirkgr Dec 14, 2023
45fc06d
Use the new ephemeral checkpoint feature
dirkgr Dec 14, 2023
d6dcf90
Restores checkpoints directly into flat params from unsharded checkpo…
dirkgr Dec 14, 2023
028d978
Merge branch 'mitchish65' into mitchish65-flatparamrestore
dirkgr Dec 14, 2023
288ccc5
Clean up the full checkpoint restorer and make it useful
dirkgr Dec 14, 2023
4935c0f
Makes checkpoint loading finally work
dirkgr Dec 15, 2023
933cce4
Backwards compatibility
dirkgr Dec 16, 2023
a0d480b
We are not yet ready to double the batch size.
dirkgr Dec 16, 2023
52bcd3f
Need to clean up the key some more
dirkgr Dec 17, 2023
8e7b0fd
Keep Pete's fixes
dirkgr Dec 18, 2023
6425924
Merge remote-tracking branch 'origin/main' into mitchish65
dirkgr Jan 4, 2024
fd95100
Remove unused import
dirkgr Jan 4, 2024
472ea1e
Productivity through sorted imports
dirkgr Jan 4, 2024
f581951
Adds documentation
dirkgr Jan 5, 2024
8bd7452
Merge branch 'main' into mitchish65
dirkgr Jan 5, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 183 additions & 0 deletions configs/mitchish35.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
run_name: mitchish35-001
seed: 6198
dry_run: false

wandb:
name: ${run_name}
project: olmo-large
group: mitchish35

model:
d_model: 7168
n_heads: 56
n_layers: 56
# mlp_ratio: 6
mlp_hidden_size: 37888
weight_tying: false
alibi: false
rope: true
flash_attention: true
attention_dropout: 0.0
attention_layer_norm: false
multi_query_attention: false
include_bias: false
block_type: sequential
layer_norm_type: default
layer_norm_with_affine: false
bias_for_layer_norm: false
attention_layer_norm_with_affine: false
activation_type: swiglu
residual_dropout: 0.0
embedding_dropout: 0.0
max_sequence_length: 2048
vocab_size: 50280
embedding_size: 50304
eos_token_id: 0
pad_token_id: 1
init_device: meta
init_fn: mitchell

compile: null

optimizer:
name: adamw
learning_rate: 1.5e-4
weight_decay: 0.1
decay_norm_and_bias: true
decay_embeddings: true
betas:
- 0.9
- 0.95
metrics_log_interval: 10

scheduler:
name: linear_with_warmup
t_warmup: 5000
alpha_f: 0.1
grad_clip_warmup_steps: 1000
grad_clip_warmup_factor: 5

tokenizer:
identifier: tokenizers/allenai_eleuther-ai-gpt-neox-20b-pii-special.json
truncate_direction: right

save_folder: ${oc.env:CHECKPOINTS_PATH}/${oc.env:SLURM_JOB_ID,${run_name}}
save_overwrite: false

# Sharded checkpoints (best for restarts)
save_interval: 500
save_num_checkpoints_to_keep: -1
sharded_checkpointer: local

# Unsharded checkpoints (for final storage)
save_interval_unsharded: null
save_num_unsharded_checkpoints_to_keep: -1

load_path: null

max_duration: 953674 # 2T tokens
global_train_batch_size: 1024
device_train_microbatch_size: 1

precision: amp_bf16

fsdp:
wrapping_strategy: by_block
precision: mixed

activation_checkpointing: whole_layer

max_grad_norm: 1.0
max_grad_norm_ratio: null

speed_monitor:
window_size: 20

eval_interval: ${save_interval}
eval_subset_num_batches: -1
device_eval_batch_size: ${device_train_microbatch_size}
evaluators:
- label: all-small-ppl-validation
data:
num_workers: 0
drop_last: true
# pin_memory: true
# prefetch_factor: 1
# persistent_workers: false
# timeout: 0
datasets:
4chan-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/4chan/val.npy
c4_100_domains-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/c4_100_domains/val.npy
c4_en-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/c4_en/val.npy
gab-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/gab/val.npy
ice-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/ice/val.npy
m2d2_s2orc-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/m2d2_s2orc/val.npy
m2d2_wiki-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/m2d2_wiki/val.npy
manosphere-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/manosphere/val.npy
mc4_en-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/mc4_en/val.npy
pile-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/pile/val.npy
ptb-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/ptb/val.npy
twitterAEE-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/twitterAEE/val.npy
wikitext_103-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/wikitext_103/val.npy

##########################
# Downstream evaluations #
##########################
- label: piqa
type: downstream

- label: hellaswag
type: downstream

- label: winogrande
type: downstream

- label: openbook_qa
type: downstream

# - label: boolq # requires implemention of the pmi_dc matrix
# type: downstream

- label: sciq
type: downstream

- label: arc_easy
type: downstream

# - label: arc_challenge # requires implemention of the pmi_dc matrix
# type: downstream

- label: copa
type: downstream

- label: rte
type: downstream

- label: commitment_bank
type: downstream

- label: sst2
type: downstream

data:
paths: ${path.glob:${oc.env:DATA_PATH}/v1_5-sample/gpt-neox-20b-pii-special-longrunfix/*.npy}
pad_direction: right
num_workers: 0
drop_last: true
pin_memory: true
prefetch_factor: 16
persistent_workers: true
timeout: 0
183 changes: 183 additions & 0 deletions configs/mitchish50.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
run_name: mitchish50-001
seed: 6198
dry_run: false

wandb:
name: ${run_name}
project: olmo-large
group: mitchish50

model:
d_model: 8192
n_heads: 64
n_layers: 64
# mlp_ratio: 6
mlp_hidden_size: 40960
weight_tying: false
alibi: false
rope: true
flash_attention: true
attention_dropout: 0.0
attention_layer_norm: false
multi_query_attention: false
include_bias: false
block_type: sequential
layer_norm_type: default
layer_norm_with_affine: false
bias_for_layer_norm: false
attention_layer_norm_with_affine: false
activation_type: swiglu
residual_dropout: 0.0
embedding_dropout: 0.0
max_sequence_length: 2048
vocab_size: 50280
embedding_size: 50304
eos_token_id: 0
pad_token_id: 1
init_device: meta
init_fn: mitchell

compile: null

optimizer:
name: adamw
learning_rate: 1.5e-4
weight_decay: 0.1
decay_norm_and_bias: true
decay_embeddings: true
betas:
- 0.9
- 0.95
metrics_log_interval: 10

scheduler:
name: linear_with_warmup
t_warmup: 5000
alpha_f: 0.1
grad_clip_warmup_steps: 1000
grad_clip_warmup_factor: 5

tokenizer:
identifier: tokenizers/allenai_eleuther-ai-gpt-neox-20b-pii-special.json
truncate_direction: right

save_folder: ${oc.env:CHECKPOINTS_PATH}/${oc.env:SLURM_JOB_ID,${run_name}}
save_overwrite: false

# Sharded checkpoints (best for restarts)
save_interval: 500
save_num_checkpoints_to_keep: -1
sharded_checkpointer: local

# Unsharded checkpoints (for final storage)
save_interval_unsharded: null
save_num_unsharded_checkpoints_to_keep: -1

load_path: null

max_duration: 953674 # 2T tokens
global_train_batch_size: 1024
device_train_microbatch_size: 1

precision: amp_bf16

fsdp:
wrapping_strategy: by_block
precision: mixed

activation_checkpointing: whole_layer

max_grad_norm: 1.0
max_grad_norm_ratio: null

speed_monitor:
window_size: 20

eval_interval: ${save_interval}
eval_subset_num_batches: -1
device_eval_batch_size: ${device_train_microbatch_size}
evaluators:
- label: all-small-ppl-validation
data:
num_workers: 0
drop_last: true
# pin_memory: true
# prefetch_factor: 1
# persistent_workers: false
# timeout: 0
datasets:
4chan-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/4chan/val.npy
c4_100_domains-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/c4_100_domains/val.npy
c4_en-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/c4_en/val.npy
gab-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/gab/val.npy
ice-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/ice/val.npy
m2d2_s2orc-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/m2d2_s2orc/val.npy
m2d2_wiki-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/m2d2_wiki/val.npy
manosphere-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/manosphere/val.npy
mc4_en-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/mc4_en/val.npy
pile-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/pile/val.npy
ptb-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/ptb/val.npy
twitterAEE-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/twitterAEE/val.npy
wikitext_103-validation:
- ${oc.env:EVAL_DATA_PATH}/perplexity/v2_small_gptneox20b/wikitext_103/val.npy

##########################
# Downstream evaluations #
##########################
- label: piqa
type: downstream

- label: hellaswag
type: downstream

- label: winogrande
type: downstream

- label: openbook_qa
type: downstream

# - label: boolq # requires implemention of the pmi_dc matrix
# type: downstream

- label: sciq
type: downstream

- label: arc_easy
type: downstream

# - label: arc_challenge # requires implemention of the pmi_dc matrix
# type: downstream

- label: copa
type: downstream

- label: rte
type: downstream

- label: commitment_bank
type: downstream

- label: sst2
type: downstream

data:
paths: ${path.glob:${oc.env:DATA_PATH}/v1_5-sample/gpt-neox-20b-pii-special-longrunfix/*.npy}
pad_direction: right
num_workers: 0
drop_last: true
pin_memory: true
prefetch_factor: 16
persistent_workers: true
timeout: 0
Loading