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

Optimize the Falcon block for inference #500

Merged
merged 18 commits into from
Sep 4, 2023
Merged

Optimize the Falcon block for inference #500

Changes from 1 commit
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
ca4d091
Optimize Falcon block for inference
mryab Sep 3, 2023
1f006c5
Fix class names
mryab Sep 3, 2023
1fc22bd
Post-rebase changes
mryab Sep 3, 2023
67764fe
Fix formatting, reduce diff
mryab Sep 3, 2023
111bf7e
Fix the test
mryab Sep 3, 2023
ce401f1
Make cos_cached/sin_cached buffers
mryab Sep 3, 2023
2c1452d
Fix buffer registration
mryab Sep 3, 2023
ae30427
Make the block compatible with other architectures
mryab Sep 3, 2023
841a0d5
Improve test and compatibility
mryab Sep 3, 2023
d56f57a
Fix rotary embeddings
mryab Sep 3, 2023
cfaf6c1
Fix rotary embeddings
mryab Sep 3, 2023
1f2ef79
WIP disable graphs
mryab Sep 4, 2023
b941df5
Fix formatting
mryab Sep 4, 2023
177669e
Rollback CUDA graphs
mryab Sep 4, 2023
91f6248
Run tests on CUDA and CPU,
mryab Sep 4, 2023
ea6c037
Enable CUDA graphs only on CUDA
mryab Sep 4, 2023
2c27c19
Do not fuse split_heads with qkv
mryab Sep 4, 2023
52baffb
Update test_optimized_layers
mryab Sep 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix rotary embeddings
  • Loading branch information
@mryab
mryab committed Sep 4, 2023
commit d56f57acd2255b41a9d02eef4675c0122a85ab40
13 changes: 10 additions & 3 deletions src/petals/models/falcon/block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,13 @@ def __init__(self, head_dim: int, base=10000):
self.input_surface = None
self.static_outputs = None

self.cos_sin(
seq_len=INFERENCE_MAX_LENGTH,
past_key_values_length=0,
device=self.inv_freq.device,
dtype=torch.get_default_dtype(),
)

def _optimized_apply_rotary(self, query, key, cos, sin):
if self.cuda_graph is None:
self.cuda_graph = torch.cuda.CUDAGraph()
Expand Down Expand Up @@ -80,11 +87,11 @@ def cos_sin(self, seq_len: int, past_key_values_length: int, device="cpu", dtype
emb = emb.float()

self.register_buffer("cos_cached", emb.cos()[None, :, :].type(dtype))
self.register_buffer("sin_cached", emb.cos()[None, :, :].type(dtype))
self.register_buffer("sin_cached", emb.sin()[None, :, :].type(dtype))

return (
self.cos_cached[:, past_key_values_length : seq_len + past_key_values_length],
self.sin_cached[:, past_key_values_length : seq_len + past_key_values_length],
self.cos_cached[:, past_key_values_length : seq_len + past_key_values_length].type(dtype),
self.sin_cached[:, past_key_values_length : seq_len + past_key_values_length].type(dtype),
)

def forward(self, query, key, past_key_values_length=0):
Expand Down