Fix the incorrect step log for profiler after resuming from a checkpoint #293
Conversation
Summary: The profiler currently maintains a counter locally and that counter is not synchronized with the checkpointed train step. This PR fixes the issue.
facebook-github-bot
added
the
CLA Signed
torchtitan/profiling.py
Outdated
| @@ -27,8 +27,7 @@ def maybe_enable_profiling(config: JobConfig, *pos_args, **kwargs): | |||
| trace_dir = os.path.join(dump_dir, save_trace_dir) | |||
| profile_freq = config.profiling.profile_freq | |||
|
|
|||
| _global_iter_count = 0 | |||
|
|
|||
| _global_iter_count = global_step | |||
There was a problem hiding this comment.
This does not sync with the internal state step_num of profiler. Let's remove it and replace all appearance of _global_iter_count by prof.step_num in trace_handler.
Also need to set torch_profiler.step_num = global_step on line 71.
There was a problem hiding this comment.
LGTM
Note: if resuming from a global_step such that global_step % profile_freq == profile_freq - 1, there won't be profile trace at step global_step + 1, which would have been done if not resuming from a checkpoint.
Essentially this is caused by profiler doesn't maintain a state dict. This miss may be avoided if we correctly implement save & load functions more carefully, which doesn't seem to be worth it.
|
In some trainers, the profiling is designed to be done only once and the global step is used to prevent profiling from happening after checkpoint resume. |
…int (pytorch#293) Summary: The profiler currently maintains a counter locally and that counter is not synchronized with the checkpointed train step. This PR fixes the issue.
…int (pytorch#293) Summary: The profiler currently maintains a counter locally and that counter is not synchronized with the checkpointed train step. This PR fixes the issue.
Summary:
The profiler currently maintains a counter locally and that counter is not synchronized with the checkpointed train step. This PR fixes the issue.