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

Calculate step duration and MFU after a training run is done #67

Open
tengyifei opened this issue Jan 31, 2025 · 0 comments
Open

Calculate step duration and MFU after a training run is done #67

tengyifei opened this issue Jan 31, 2025 · 0 comments
Assignees
@tengyifei

Comments

@tengyifei
Copy link
Collaborator

No description provided.

@tengyifei tengyifei self-assigned this Jan 31, 2025
tengyifei added a commit that referenced this issue Jan 31, 2025
@tengyifei
Script to compute step duration
f622b16 
Precursor to #67.

Imports https://github.com/AI-Hypercomputer/tpu-recipes/blob/main/utils/profile_convert.py
and improves it.

Specifically, I noticed sometimes there is an empty gap between two step
markers in the profile. So if we averaged event durations, that would
overestimate the MFU. Instead, this now averages the delta between the
starting time offset of neighboring events.

Now that we can print step time from the profile, I removed the step
time from the training loop. That added a bunch of delays and is
actually pretty inaccurate (1.7s vs 1.85s in local testing).

Tested:

  XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 python3 torchprime/torch_xla_models/train.py mesh.fsdp=8 profile_step=4 model=llama-3-8b
  XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 python3 torchprime/torch_xla_models/train.py mesh.fsdp=8 profile_step=4 model=mixtral-8x7b
@tengyifei tengyifei mentioned this issue Jan 31, 2025
Merged
tengyifei added a commit that referenced this issue Feb 3, 2025
@tengyifei
Script to compute step duration (#70)
a6b6630 
Precursor to #67.

Imports https://github.com/AI-Hypercomputer/tpu-recipes/blob/main/utils/profile_convert.py
and improves it.

Specifically, I noticed sometimes there is an empty gap between two step
markers in the profile. So if we averaged event durations, that would
overestimate the MFU. Instead, this now averages the delta between the
starting time offset of neighboring events.

Now that we can print step time from the profile, I removed the step
time from the training loop. That added a bunch of delays and is
actually pretty inaccurate (1.7s vs 1.85s in local testing).

Tested:

  XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 python3 torchprime/torch_xla_models/train.py mesh.fsdp=8 profile_step=4 model=llama-3-8b
  XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 python3 torchprime/torch_xla_models/train.py mesh.fsdp=8 profile_step=4 model=mixtral-8x7b
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@tengyifei tengyifei

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@tengyifei