"docs/vscode:/vscode.git/clone" did not exist on "b2f7745774fc3d5105c0e0babbfb656bfcacfcac"
Unverified Commit fecbb7c7 authored by Wei's avatar Wei Committed by GitHub
Browse files

[Bugfix][gpt-oss] passing the cache config in gpt-oss (#23613)


Signed-off-by: default avatarWei Wei <wwei6@meta.com>
parent 6dab89b8
...@@ -174,12 +174,15 @@ class TransformerBlock(torch.nn.Module): ...@@ -174,12 +174,15 @@ class TransformerBlock(torch.nn.Module):
def __init__( def __init__(
self, self,
config: GptOssConfig, config: GptOssConfig,
cache_config: CacheConfig,
quant_config: QuantizationConfig, quant_config: QuantizationConfig,
prefix: str = "", prefix: str = "",
): ):
super().__init__() super().__init__()
self.layer_idx = extract_layer_index(prefix) self.layer_idx = extract_layer_index(prefix)
self.attn = OAIAttention(config, prefix=f"{prefix}.attn") self.attn = OAIAttention(config,
prefix=f"{prefix}.attn",
cache_config=cache_config)
self.mlp = MLPBlock(config, self.mlp = MLPBlock(config,
self.layer_idx, self.layer_idx,
quant_config=quant_config, quant_config=quant_config,
...@@ -203,6 +206,7 @@ class GptOssModel(nn.Module): ...@@ -203,6 +206,7 @@ class GptOssModel(nn.Module):
): ):
super().__init__() super().__init__()
self.config = vllm_config.model_config.hf_config self.config = vllm_config.model_config.hf_config
self.cache_config = vllm_config.cache_config
self.quant_config = vllm_config.quant_config self.quant_config = vllm_config.quant_config
self.parallel_config = vllm_config.parallel_config self.parallel_config = vllm_config.parallel_config
self.config.hidden_size = self.config.hidden_size self.config.hidden_size = self.config.hidden_size
...@@ -213,6 +217,7 @@ class GptOssModel(nn.Module): ...@@ -213,6 +217,7 @@ class GptOssModel(nn.Module):
self.layers = torch.nn.ModuleList([ self.layers = torch.nn.ModuleList([
TransformerBlock( TransformerBlock(
self.config, self.config,
cache_config=self.cache_config,
quant_config=self.quant_config, quant_config=self.quant_config,
prefix=maybe_prefix(prefix, f"block.{layer_idx}"), prefix=maybe_prefix(prefix, f"block.{layer_idx}"),
) for layer_idx in range(self.config.num_hidden_layers) ) for layer_idx in range(self.config.num_hidden_layers)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment