Skip to content

Conversation

@qihqi
Copy link
Collaborator

@qihqi qihqi commented Jun 13, 2024

No description provided.

@qihqi qihqi requested review from FanhaiLu1 and wang2yn84 June 13, 2024 17:44
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason why we don't run int8 on v5e-1?

----| ------- | ------ |---------- | -------------|-----------------|------------------|----------------------
2024-05-14 | TPU v5e-8 | bfloat16 | 512 | 2048 | 1024 | 1024 | 8700
2024-05-14 | TPU v5e-8 | int8 | 1024 | 2048 | 1024 | 1024 | 8746
2024-06-13 | TPU v5e-1 | bfloat16 | 1024 | 2048 | 1024 | 1024 | 4249
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The v5e-1 number looks great! The gap is huge, I'm wondering the v5-8 data parallel didn't compute each token locally.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1. Since we are sharding on batch for Gemma 2B, even at naive case (duplicate weights) v5e-8 should have 8 * 4249 = 33k toks/sec

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The prefills are replicated (8 chips computing the same thing) because we are sharding on batch and prefill is of batch=1. There might be other issues. Overall we should recommand v5e-1 for gemma-2b instead.

@qihqi qihqi merged commit 7526a90 into main Jun 17, 2024
@qihqi qihqi deleted the qihqi-patch-1 branch June 17, 2024 17:21
@qihqi
Copy link
Collaborator Author

qihqi commented Jun 17, 2024

Any reason why we don't run int8 on v5e-1?

Yeah, we should, we havnet gotten to it yet.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants