Moving the checkpointing util funtions to a spearate file to avoid circular dependency #1623
Conversation
wang2yn84
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
vipannalla,
RissyRan,
richjames0,
rni418,
gagika,
shralex,
yangyuwei,
SurbhiJainUSC,
hengtaoguo,
A9isha,
wyzhang,
mitalisi,
gpolovets1 and
mailvijayasingh
as code owners
There was a problem hiding this comment.
Thanks @wang2yn84 for the nice refactor. Was the circular dependency causing some type of error? Or is the change for cleaner code?
| @@ -52,6 +52,7 @@ | |||
| import jax | |||
| import gc | |||
| from MaxText import max_logging | |||
| from MaxText import maxtext_checkpointing_utils | |||
There was a problem hiding this comment.
Looks like from MaxText import max_utils is no longer needed. Can you remove it from this file and other ones where it isn't needed either?
| ) | ||
| state = maxtext_utils.init_decode_state(None, params) | ||
|
|
||
| state = max_utils.unbox_logicallypartioned(state) |
There was a problem hiding this comment.
Is this change trying to get rid of the max_utils dependency?
There was a problem hiding this comment.
The change is to move all the util functions that depends on MaxText.checkpointing to this file so that other files can depends on maxtext_utils without circular dependency issue. maxtext_checkpointing_utils can depend on max_utils and maxtext_utils without any issue.
Hi Branden, in the current implementation, we have to carefully avoid depend on maxtext_utils from multihost_dataloading. There are couple of util functions in multihost_dataloading we should place in maxtext_utils but we can't. |
| "Skipping saving quantized checkpoint as save_quantized_params_path is null." | ||
|
|
||
|
|
||
| def setup_training_state(model, data_iterator, tx, config, rng, mesh, checkpoint_manager): |
There was a problem hiding this comment.
I don't think this method fits in this file - this define the training state, its needed even without checkpointing
There was a problem hiding this comment.
Chatted offline, agreed to rename this new util file to maxtext_state_initializion_utils.py.
| return state, state_mesh_annotations, state_mesh_shardings, data_iterator | ||
|
|
||
|
|
||
| def setup_decode_state(model, config, rng, mesh, checkpoint_manager): |
There was a problem hiding this comment.
This is not a checkpointing method, its needed even when checkpointing is turned off
| import orbax.checkpoint.experimental.emergency.replicator_checkpoint_manager as emergency_replicator_checkpoint_manager | ||
|
|
||
|
|
||
| def setup_initial_state( |
There was a problem hiding this comment.
this is not a checkpointing function, its needed even when checkpointing is turned off
wang2yn84
force-pushed
the
lance-refactor
branch
3 times, most recently
from
fa152d3 to
f3a0a49
Compare
|
|
||
|
|
||
|
|
||
| def save_quantized_checkpoint_if_configured(config, params): |
There was a problem hiding this comment.
Can you have some inference folks look at this change @vipannalla or @mitalisi
wang2yn84
force-pushed
the
lance-refactor
branch
2 times, most recently
from
ba55551 to
83950d0
Compare
…rcular dependency.
wang2yn84
force-pushed
the
lance-refactor
branch
from
83950d0 to
d34944c
Compare
Description
The circular dependency is multihost_dataloading -> maxtext_utils -> checkpointing -> multihost_dataloading. That's why in the previous development, we removed the deps of maxtext_utils from multihost_dataloading to avoid the issue, but that's just temporarily. Moving all the maxtext_util functions that depends on checkpointing to a separate util function solves this issue.
Tests
Presubmit tests.
Checklist
Before submitting this PR, please make sure (put X in square brackets):