[CI] Fix broken CI (#6599)
Revert4fb3d5e1b2it breaks E2E Test - vLLM version: v0.15.0 - vLLM main:d7e17aaacd
This commit is contained in:
@@ -1,14 +1,11 @@
|
||||
from typing import Any
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from vllm.model_executor.models.utils import extract_layer_index
|
||||
|
||||
from vllm_ascend.distributed.parallel_state import get_shard_weight_group
|
||||
from vllm_ascend.ops.layer_shard_linear import (
|
||||
is_hidden_layer,
|
||||
post_process_after_loading_for_shard_weight_series,
|
||||
reach_layer_for_shard_weight_series,
|
||||
register_layer_to_shard_weight_series,
|
||||
)
|
||||
is_hidden_layer, post_process_after_loading_for_shard_weight_series,
|
||||
reach_layer_for_shard_weight_series, register_layer_to_shard_weight_series)
|
||||
from vllm_ascend.utils import flashcomm2_enable, o_shard_enable
|
||||
|
||||
|
||||
@@ -29,7 +26,7 @@ class Flashcomm2OShardManager:
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._shard_layers: dict[int, Any] = {}
|
||||
self._shard_layers: Dict[int, Any] = {}
|
||||
|
||||
def flashcomm2_oshard_enable(self):
|
||||
return flashcomm2_enable() and o_shard_enable()
|
||||
@@ -55,10 +52,12 @@ class Flashcomm2OShardManager:
|
||||
self._shard_layers[layer_idx] = layer
|
||||
|
||||
register_layer_to_shard_weight_series(
|
||||
series_name="o_proj", group=get_shard_weight_group(), layer=layer, prefetch_step=prefetch_step
|
||||
)
|
||||
series_name="o_proj",
|
||||
group=get_shard_weight_group(),
|
||||
layer=layer,
|
||||
prefetch_step=prefetch_step)
|
||||
|
||||
def get_layer(self, layer_idx: int) -> Any | None:
|
||||
def get_layer(self, layer_idx: int) -> Optional[Any]:
|
||||
"""Safely retrieves a registered layer by its index.
|
||||
|
||||
Args:
|
||||
|
||||
Reference in New Issue
Block a user