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

Enable torch.compile for loss_fn #597

Merged
merged 16 commits into from
Dec 8, 2023
Merged

Enable torch.compile for loss_fn #597

Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
dc7e966
compile loss fn for all workloads
BoyuanFeng Dec 5, 2023
fa411fd
reduce eval_batch_size by half for criteo1tb_pytorch
BoyuanFeng Dec 6, 2023
dfb8ceb
further reduce eval_batch_size by half for criteo1tb_pytorch
BoyuanFeng Dec 6, 2023
30062f8
reduce to 1/4
BoyuanFeng Dec 6, 2023
fa3dddd
reduce eval_batch_size for criteo1tb
BoyuanFeng Dec 6, 2023
1ddfa09
recover eval_batch_size for criteo1tb
BoyuanFeng Dec 6, 2023
bfe74ce
temporary commenting out workload.eval_model for criteo1tb
BoyuanFeng Dec 6, 2023
7a112f5
disable torch.compile for criteo1tb
BoyuanFeng Dec 7, 2023
4b53433
recover evaluation
BoyuanFeng Dec 7, 2023
c4029ac
comment out eval to avoid tokenizer for librispeech
BoyuanFeng Dec 7, 2023
ea72827
recover eval_model
BoyuanFeng Dec 7, 2023
a57759b
comment out eval
BoyuanFeng Dec 7, 2023
7132021
specify workloads that benefits from compiling loss_fn
BoyuanFeng Dec 7, 2023
bcf1ca0
enable evaluation
BoyuanFeng Dec 7, 2023
33ecf8e
nit
BoyuanFeng Dec 7, 2023
a134d08
Merge branch 'dev' into feat/pt2-compile-loss-fn
BoyuanFeng Dec 8, 2023
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
5 changes: 5 additions & 0 deletions submission_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,9 @@ def train_once(
]
eager_backend_workloads = ['librispeech_deepspeech']
aot_eager_backend_workloads = []
loss_compilation_workloads = [
'fastmri', 'librispeech_deepspeech', 'ogbg', 'wmt'
]
base_workload = workloads.get_base_workload_name(workload_name)
if base_workload in compile_error_workloads:
logging.warning(
Expand All @@ -247,6 +250,8 @@ def train_once(
else:
logging.info('Performing `torch.compile`.')
model_params = torch.compile(model_params)
if base_workload in loss_compilation_workloads:
workload.loss_fn = torch.compile(workload.loss_fn)
logging.info('Initializing optimizer.')
with profiler.profile('Initializing optimizer'):
optimizer_state = init_optimizer_state(workload,
Expand Down
Loading