biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Use runtime profiling to replace manual memory analyzers (#81)

Browse Source
This commit is contained in:
Zhuohan Li
2023-05-19 11:35:44 -06:00
committed by GitHub
parent 825d8892b5
commit f756799b84
14 changed files with 211 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
cacheflow/model_executor/models/gpt2.py
View File

@@ -58,7 +58,8 @@ class GPT2Attention(nn.Module):
self.c_proj = RowParallelLinear(self.hidden_size, self.hidden_size, bias=True,
input_is_parallel=True,
perform_initialization=False)
self.attn = GPTCacheFlowAttention(scale=self.scale)
self.attn = GPTCacheFlowAttention(self.num_heads, self.head_dim,
scale=self.scale)
def forward(
self,

3
cacheflow/model_executor/models/gpt_neox.py
View File

@@ -62,7 +62,8 @@ class GPTNeoXAttention(nn.Module):
scaling = self.head_size ** -0.5
rotary_dim = int(self.head_size * config.rotary_pct)
assert rotary_dim % 2 == 0
self.attn = GPTNeoXCacheFlowAttention(scaling, rotary_dim)
self.attn = GPTNeoXCacheFlowAttention(self.num_heads, self.head_size,
scaling, rotary_dim)
def forward(
self,

3
cacheflow/model_executor/models/llama.py
View File

@@ -104,7 +104,8 @@ class LlamaAttention(nn.Module):
input_is_parallel=True,
perform_initialization=False,
)
self.attn = GPTNeoXCacheFlowAttention(self.scaling, self.head_dim)
self.attn = GPTNeoXCacheFlowAttention(self.num_heads, self.head_dim,
self.scaling, rotary_dim=self.head_dim)
def forward(
self,

3
cacheflow/model_executor/models/opt.py
View File

@@ -74,7 +74,8 @@ class OPTAttention(nn.Module):
self.out_proj = RowParallelLinear(embed_dim, embed_dim, bias=bias,
input_is_parallel=True,
perform_initialization=False)
self.attn = GPTCacheFlowAttention(scale=self.scaling)
self.attn = GPTCacheFlowAttention(self.num_heads, self.head_dim,
scale=self.scaling)
def forward(
self,