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

Implement megatron-aware perplexity in torchmetrics #525

Open
wants to merge 29 commits into
base: main
Choose a base branch
from
Open

Implement megatron-aware perplexity in torchmetrics #525

wants to merge 29 commits into from

Conversation

sichu2023
Copy link
Collaborator

@sichu2023 sichu2023 commented Dec 12, 2024 •
edited
Loading

Summary

Replace PerplexityLoggingCallback with parallelism-aware MegatronPerplexityMetric. Pipeline parallel is not supported yet.

Usage

Added perplexity logging in argparse.

python sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py --log-train-ppl --log-val-ppl

Testing

Added under sub-packages/bionemo-llm/tests/bionemo/llm/test_lightning.py

Tests for these changes can be run via:

pytest -v sub-packages/bionemo-llm/tests/bionemo/llm/test_lightning.py

@sichu2023 sichu2023 mentioned this pull request Dec 12, 2024
Draft
@sichu2023
Copy link
Collaborator Author

Previous @pstjohn 's effort
#382

@sichu2023 sichu2023 self-assigned this Dec 12, 2024
@sichu2023 sichu2023 changed the base branch from main to sichu/update-ddp-config December 12, 2024 14:00
@sichu2023 sichu2023 force-pushed the sichu/update-ddp-config branch 2 times, most recently from ecd102a to cfa66d3 Compare December 12, 2024 17:28
@sichu2023 sichu2023 force-pushed the sichu/update-ddp-config branch from cfa66d3 to e1e9037 Compare December 16, 2024 16:59
@sichu2023
Copy link
Collaborator Author

Exponential Validation Loss vs Validation Perplexity
Test run on single device without parallelism.

@sichu2023 sichu2023 force-pushed the sichu/torchmetrics-ppl-logging branch from 2fd1fd7 to 4959e80 Compare December 19, 2024 19:48
@sichu2023 sichu2023 changed the base branch from sichu/update-ddp-config to main December 19, 2024 19:54
@sichu2023 sichu2023 marked this pull request as ready for review December 19, 2024 19:55
@sichu2023 sichu2023 force-pushed the sichu/torchmetrics-ppl-logging branch from 1de355d to 32e0568 Compare December 19, 2024 19:56
@sichu2023
Copy link
Collaborator Author

On multi-device DP
Exponential of Validation Loss and Validation PPL over Time

@sichu2023
Copy link
Collaborator Author

And on multi-device TP x DP.
Uploading Exponential of Validation Loss and Validation PPL over Time (Second Set).png…

@sichu2023 sichu2023 requested review from pstjohn and jstjohn December 19, 2024 23:57
@sichu2023
add tp-aware update method
b2e5f1d
@sichu2023
drop comment on tp-aware normalization
edf1da3
@sichu2023
update comment
35b5abf
@sichu2023
fix cp error
564626c
@sichu2023
drop process_group
bf16c01
@sichu2023
add MegatronPerplexityMetric testing
6234c70
@sichu2023
fix metric device
5db215d
@sichu2023
fix MegatronPerplexityMetric.update
0a3ae38
@sichu2023
clean up test_megatron_perplexity_metric_with_single_microbatch_golde…
0603b4d 
…n_value_without_parallelism
@sichu2023
fix get_random_microbatch
c43d9e4
@sichu2023
add variable length microbatch test
71f9dc0
@sichu2023
ruff
561d3c1
@sichu2023
add back self.log_{train,val}_ppl
f3cc5d6
@sichu2023
add back return in {train,validation}_step
d03a94e
@sichu2023
add argparse
eac8b18
@sichu2023
disable async ckpt save
da76178
@sichu2023
drop pp support
bd6ebd3
@sichu2023
Revert "update ddp config"
7886e86 
This reverts commit a698ad4.
@sichu2023
Revert "update MegatronMixedPrecision"
f2ea013 
This reverts commit 6f14982.
@sichu2023
disable training ppl logging by default
80403ae
@sichu2023
remove ppl callback
dcc84ec
@sichu2023
move pp check to train_esm2.py
fb5ad07
@sichu2023
ruff
78aa7a3
@sichu2023 sichu2023 force-pushed the sichu/torchmetrics-ppl-logging branch from 007c6e0 to 78aa7a3 Compare December 20, 2024 00:09
@sichu2023 sichu2023 enabled auto-merge (squash) December 20, 2024 00:09
@sichu2023
Copy link
Collaborator Author

/build-ci

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@farhadrgh farhadrgh Awaiting requested review from farhadrgh farhadrgh is a code owner

@dorotat-nv dorotat-nv Awaiting requested review from dorotat-nv dorotat-nv is a code owner

@malcolmgreaves malcolmgreaves Awaiting requested review from malcolmgreaves malcolmgreaves is a code owner

@skothenhill-nv skothenhill-nv Awaiting requested review from skothenhill-nv skothenhill-nv is a code owner

@pstjohn pstjohn Awaiting requested review from pstjohn pstjohn is a code owner

@jstjohn jstjohn Awaiting requested review from jstjohn jstjohn is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

@sichu2023 sichu2023

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@sichu2023 @jstjohn @pstjohn