Revert "Add simple CPU offloading support" (#2252)

We'll re-add the commit to correctly ack Kaichao's authorship
This commit is contained in:
Ying Sheng
2024-11-28 23:36:55 -08:00
committed by GitHub
parent 4f2ee48ed1
commit 4057ea82c9
9 changed files with 29 additions and 173 deletions
+5 -8
View File
@@ -48,7 +48,6 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils import make_layers
class OlmoeMoE(nn.Module):
@@ -262,13 +261,11 @@ class OlmoeModel(nn.Module):
config.vocab_size,
config.hidden_size,
)
self.layers = make_layers(
config.num_hidden_layers,
lambda idx, prefix: OlmoeDecoderLayer(
config=config,
quant_config=quant_config,
layer_id=idx,
),
self.layers = nn.ModuleList(
[
OlmoeDecoderLayer(config, layer_id, quant_config=quant_config)
for layer_id in range(config.num_hidden_layers)
]
)
self.norm = RMSNorm(config.hidden_size, eps=1e-5)