-
Notifications
You must be signed in to change notification settings - Fork 1.1k
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.
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
Pretrain openwebtext script with speed monitoring #147
Conversation
if not torch.cuda.is_available(): | ||
return 0 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We could make this function get the device type so that it can work for XLA devices too. @gkroiz Do you know where we can get the FLOPs for the different TPU archs?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi, https://cloud.google.com/tpu/docs/system-architecture-tpu-vm seems to have some numbers for v2, v3, and v4, but only for bfloat16. Let me see if I can find a document that covers FLOPs for other levels of precision
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Leaving this for a follow-up
Part of #123
openwebtext with FSDP requires Lightning-AI/pytorch-lightning#17832
Follow-ups: