[Cleanup] Miscellaneous Refactors #1607
Conversation
wesleytruong
requested review from
tianyu-l,
fegin,
wwwjn and
wconstab
as code owners
meta-cla
bot
added
the
CLA Signed
torchtitan/components/validate.py
Outdated
| @@ -82,8 +82,9 @@ def validate( | |||
| step: int, | |||
| ) -> None: | |||
| # Set model to eval mode | |||
| for model in model_parts: | |||
| model.eval() | |||
| model = model_parts[0] | |||
There was a problem hiding this comment.
why keeping this? I think we only need this in non-PP case.
torchtitan/components/validate.py
Outdated
| @@ -174,7 +175,8 @@ def validate( | |||
| module.reshard() | |||
There was a problem hiding this comment.
After you switch the order of checkpoint and validate, I think this reshard part is not necessary and creates extra overhead -- we are doing validate (reshard) -> next train (unshard), which we could've avoided?
Could you check if it works well when checkpoint and validate happens on the same step? If so we can remove this code.
There was a problem hiding this comment.
Sure, in my test it looks like removing this reshard doesn't affect memory usage or loss so I think it should be fine to remove.
| Reshard | No Reshard |
|---|---|
 |
 |
torchtitan/train.py
Outdated
| model_args, job_config.model.hf_assets_path | ||
| model_args, | ||
| job_config.model.hf_assets_path | ||
| if job_config.checkpoint.enable_checkpoint |
There was a problem hiding this comment.
Why we need this change here? hf_assets_path can be built during loading HF weights as well, even if it's not used, and we could leave the checkpointing related logic in checkpoint.py
There was a problem hiding this comment.
Sorry you're right, I was trying to do some clean up in order to suppress this warning if the user is not intending to save in HF since in that case we wouldn't need a model.safetensor.index.json, but this would interfere with if we want to load from HF from the hf_assets_path. One other idea to suppress this would be to move this error to the checkpointer, but then it would make this warning difficult to overload such as we want to may want to do in Flux. https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/flux/model/state_dict_adapter.py#L53-L57
There was a problem hiding this comment.
I think it's ok to show the warning to user. When do we want to suppress the warning? I could only think if a model has only one safetensor file, which out the model.safetensor.index.json.
If a model checkpoint has multiple safetensor file, then both saving and loading needs to check the model.safetensor.index.json file exists.
There was a problem hiding this comment.
When do we want to suppress the warning?
E.g. when checkpointing is not enabled at all.
If a model checkpoint has multiple safetensor file, then both saving and loading needs to check the
model.safetensor.index.jsonfile exists.
I believe only saving requires it. For loading it should be optional, but I'm not sure if it helps load faster. cc @ankitageorge to confirm.
There was a problem hiding this comment.
Ya only saving requires it
tianyu-l
added
the
release blocking
wesleytruong
force-pushed
the
validation_pp_fix
branch
from
a4962f3 to
19bd2df
Compare
torchtitan/train.py
Outdated
| job_config.model.hf_assets_path | ||
| if job_config.checkpoint.enable | ||
| and job_config.checkpoint.last_save_in_hf | ||
| else None, |
There was a problem hiding this comment.
hmmm even for load you need the path? https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpoint.py#L543
I think there are two ways:
- always leave the warning there as is
- only pass the
hf_assets_pathin whencheckpoint.enable=True
I feel since 2 is not clean enough (as in we don't differentiate save vs. load), so I'm OK with 1.
cc @wwwjn for your opinion.
There was a problem hiding this comment.
I'm OK with 1 as well. We can leave the warning and let user aware of this
wesleytruong
changed the title
wesleytruong
force-pushed
the
validation_pp_fix
branch
from
a4dca2d to
ff95c80
Compare
| @@ -398,13 +398,13 @@ class Parallelism: | |||
|
|
|||
| @dataclass | |||
| class Checkpoint: | |||
| enable_checkpoint: bool = False | |||
| enable: bool = False | |||
There was a problem hiding this comment.
nice, I was thinking to do this as well.
torchtitan/train.py
Outdated
| model_args, | ||
| ( | ||
| job_config.model.hf_assets_path | ||
| if job_config.checkpoint.enable |
There was a problem hiding this comment.
Could you please move using checkpoint information out of the trainer, if possible? We deliberately hide job_config.checkpoint inside Checkpointer and let the trainer always call checkpointer API. This usage looks okay but would be good to hide if possible.
This reverts commit cd337db.
This PR makes several miscellaneous refactors to clean up
torchtitanbefore release.Changes:
model_partsto eval mode inValidatorclass to support PP (Bug fix)checkpoint.enable_checkpoint -> checkpoint.enable(Refactor)validation.enabled -> validation.enable(Refactor)