Skip to content

Commit

Permalink
[Misc] Medusa supports custom bias
Browse files Browse the repository at this point in the history
  • Loading branch information
skylee-01 committed Nov 15, 2024
1 parent b311efd commit 2250501
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions vllm/model_executor/models/medusa.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@

class ResidualBlock(nn.Module):

def __init__(self, hidden_size: int, num_layers: int) -> None:
def __init__(self, config: VllmConfig, hidden_size: int, num_layers: int) -> None:

Check failure on line 17 in vllm/model_executor/models/medusa.py

View workflow job for this annotation

GitHub Actions / ruff (3.12)

Ruff (E501)

vllm/model_executor/models/medusa.py:17:81: E501 Line too long (86 > 80)
super().__init__()

self.layers = nn.ModuleList([
nn.Linear(hidden_size, hidden_size, bias=False)
nn.Linear(hidden_size, hidden_size, bias=getattr(config, "medusa_fc_bias", False))

Check failure on line 21 in vllm/model_executor/models/medusa.py

View workflow job for this annotation

GitHub Actions / ruff (3.12)

Ruff (E501)

vllm/model_executor/models/medusa.py:21:81: E501 Line too long (94 > 80)
for _ in range(num_layers)
])
self.act = nn.SiLU()
Expand Down Expand Up @@ -49,7 +49,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.config = config
self.blocks = nn.ModuleList([
ResidualBlock(hidden_size=self.config.hidden_size,
ResidualBlock(config=config, hidden_size=self.config.hidden_size,
num_layers=self.config.num_hidden_layers)
for _ in range(self.config.num_heads)
])
Expand Down

0 comments on commit 2250501

Please sign in to comment.