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.

Sign up for GitHub

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

Jump to bottom

Complete ESM2 pretraining #112

Merged
merged 93 commits into from
Aug 29, 2024
Merged

Complete ESM2 pretraining #112

merged 93 commits into from
Aug 29, 2024

Conversation

sichu2023
Copy link
Collaborator

@sichu2023 sichu2023 commented Aug 20, 2024 •
edited
Loading

Based on @farhadrgh 's PR on ESM2 pretraining and loss test.

Difference from the original feature branch

  • Fix bug in pretraining teardown
  • Migrate dataset shuffling from megatron random sampler to PRNGDatasetShuffler
  • Enable pytest on pretraining
  • Enable global batch size calculation
  • Enable model hparam through argparse
  • Support limit_val_batches as a fraction of validation dataset size and add unittest
  • Add megatron_parallel_state_utils to main unittest
  • Add test to ensure samples from validation dataset are consistent with potential shuffling
  • Remove eval_iters from num_val_samples
  • (reverted) Remove train and validation datasets's dependency on total_samples
  • Extract limit_val_samples inference
  • Sync all updates to geneformer

@sichu2023 sichu2023 requested review from pstjohn and farhadrgh August 20, 2024 20:32
@sichu2023
Copy link
Collaborator Author

/build-ci

6 similar comments
@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023
Copy link
Collaborator Author

/build-ci

@ohadmo
Copy link
Member

ohadmo commented Aug 21, 2024

/build-ci

@ohadmo
Copy link
Member

ohadmo commented Aug 21, 2024

/build-ci

@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023 sichu2023 changed the base branch from farhadr/pretrain to v2-main August 22, 2024 03:15
@sichu2023 sichu2023 marked this pull request as draft August 22, 2024 03:25
@sichu2023 sichu2023 marked this pull request as ready for review August 22, 2024 04:26
@sichu2023
Copy link
Collaborator Author

/build-ci

1 similar comment
@ohadmo
Copy link
Member

ohadmo commented Aug 22, 2024

/build-ci

@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023 sichu2023 requested a review from jstjohn August 23, 2024 03:40
pstjohn
pstjohn reviewed Aug 23, 2024
View reviewed changes
Copy link
Collaborator

@pstjohn pstjohn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will cycle back and finish this review later -- but wanted to post some of these comments in the meantime

scripts/protein/esm2/test_esm2_pretrain.py Show resolved Hide resolved
scripts/protein/esm2/test_esm2_pretrain.py Show resolved Hide resolved
scripts/protein/esm2/test_esm2_pretrain.py Show resolved Hide resolved
scripts/protein/esm2/test_esm2_pretrain.py Show resolved Hide resolved
scripts/protein/esm2/test_esm2_pretrain.py Show resolved Hide resolved
sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py Outdated Show resolved Hide resolved
sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py Outdated Show resolved Hide resolved
sub-packages/bionemo-esm2/src/bionemo/esm2/data/tokenizer/__init__.py Show resolved Hide resolved
sub-packages/bionemo-esm2/src/bionemo/esm2/data/tokenizer/__init__.py Show resolved Hide resolved
sub-packages/bionemo-esm2/src/bionemo/esm2/data/tokenizer/__init__.py Outdated Show resolved Hide resolved
@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023
Copy link
Collaborator Author

I will update README.md once this PR is settled.

sichu2023 added 26 commits August 29, 2024 20:57
@sichu2023
sync testing changes on geneformer
be0f114
@sichu2023
support limit_val_batches = None
18dcb67
@sichu2023
add unittest to ensure consistant samples per epoch
2a2ce08
@sichu2023
add unittest on limit_val_batches = None
1867704
@sichu2023
add unittest on float_or_int_or_none
8503b78
@sichu2023
add infer_global_batch_size unittest
09c42e7
@sichu2023
ruff
6974d28
@sichu2023
ruff
e675f94
@sichu2023
ruff
468697a
@sichu2023
fix dict order in tensor_dict_hash
1cdf113
@sichu2023
remove eval_iters from num_val_samples
4923056
@sichu2023
remove comment
72cdab8
@sichu2023
expand test_limit_val_batches_is_none
394582d
@sichu2023
fix unittest and add limit_val_samples = 0 edge case
b70b76b
@sichu2023
Fix unittest and add limit_val_batches = 0 in esm
bb6ef40
@sichu2023
rename PRNGDatasetShuffler
2653033
@sichu2023
drop total_samples in create_valid_dataset
ab679a9
@sichu2023
drop total_samples from create_train_dataset
bfc4769
@sichu2023
ruff and extract infer_num_val_samples
cd37c0f
@sichu2023
transfer infer_num_samples to geneformer
ea7240f
@sichu2023
ruff
ceac497
@sichu2023
use hosted tokenizer files
aaf3405
@sichu2023
add back num_samples in esm dataset
8669a7d
@sichu2023
ruff
740f2b6
@sichu2023
fix valid shuffling
8c0daad
@sichu2023
num_train_samples in dataset shuffling
aaabf34
@sichu2023
Copy link
Collaborator Author

/build-ci

@sichu2023 sichu2023 merged commit 76ec468 into v2-main Aug 29, 2024
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@gwarmstrong gwarmstrong gwarmstrong left review comments

@jstjohn jstjohn jstjohn approved these changes

@farhadrgh farhadrgh farhadrgh approved these changes

@pstjohn pstjohn pstjohn approved these changes

@jomitchellnv jomitchellnv Awaiting requested review from jomitchellnv

@malcolmgreaves malcolmgreaves Awaiting requested review from malcolmgreaves

@skothenhill-nv skothenhill-nv Awaiting requested review from skothenhill-nv

@DejunL DejunL Awaiting requested review from DejunL

@guoqing-zhou guoqing-zhou Awaiting requested review from guoqing-zhou

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@sichu2023 @ohadmo @jstjohn @pstjohn @farhadrgh @gwarmstrong