Fix metrics (#15998)
This commit is contained in:
@@ -555,7 +555,7 @@ def fused_experts_impl(
|
|||||||
gemm1_alpha,
|
gemm1_alpha,
|
||||||
gemm1_limit,
|
gemm1_limit,
|
||||||
)
|
)
|
||||||
elif _is_hip or _is_cuda:
|
elif _is_cuda or _is_hip:
|
||||||
if not filter_expert:
|
if not filter_expert:
|
||||||
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
||||||
else:
|
else:
|
||||||
@@ -575,7 +575,7 @@ def fused_experts_impl(
|
|||||||
elif activation == "gelu" and is_gated:
|
elif activation == "gelu" and is_gated:
|
||||||
assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
|
assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
|
||||||
assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
|
assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
|
||||||
if _is_hip or _is_cuda:
|
if _is_cuda or _is_hip:
|
||||||
if not filter_expert:
|
if not filter_expert:
|
||||||
gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -810,9 +810,11 @@ def _apply_activation(x, ACTIVATION_TYPE: tl.constexpr):
|
|||||||
x = x.to(tl.float32)
|
x = x.to(tl.float32)
|
||||||
if ACTIVATION_TYPE == "silu":
|
if ACTIVATION_TYPE == "silu":
|
||||||
return x * tl.sigmoid(x)
|
return x * tl.sigmoid(x)
|
||||||
else:
|
elif ACTIVATION_TYPE == "gelu":
|
||||||
kAlpha = 0.7978845608028654
|
kAlpha = 0.7978845608028654
|
||||||
return 0.5 * x * (1 + tanh(kAlpha * (x + 0.044715 * x * x * x)))
|
return 0.5 * x * (1 + tanh(kAlpha * (x + 0.044715 * x * x * x)))
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported activation: {ACTIVATION_TYPE}")
|
||||||
|
|
||||||
|
|
||||||
@triton.jit
|
@triton.jit
|
||||||
|
|||||||
@@ -200,7 +200,7 @@ class DataParallelController:
|
|||||||
|
|
||||||
self.init_dispatcher()
|
self.init_dispatcher()
|
||||||
|
|
||||||
self.watchdog = Watchdog.create(
|
self.soft_watchdog = Watchdog.create(
|
||||||
debug_name="DataParallelController",
|
debug_name="DataParallelController",
|
||||||
watchdog_timeout=server_args.soft_watchdog_timeout,
|
watchdog_timeout=server_args.soft_watchdog_timeout,
|
||||||
soft=True,
|
soft=True,
|
||||||
@@ -564,7 +564,7 @@ class DataParallelController:
|
|||||||
def event_loop(self):
|
def event_loop(self):
|
||||||
while True:
|
while True:
|
||||||
while True:
|
while True:
|
||||||
self.watchdog.feed()
|
self.soft_watchdog.feed()
|
||||||
try:
|
try:
|
||||||
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
|
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
|
||||||
except zmq.ZMQError:
|
except zmq.ZMQError:
|
||||||
|
|||||||
@@ -116,7 +116,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
|
|||||||
self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss"
|
self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss"
|
||||||
self.disable_tokenizer_batch_decode = server_args.disable_tokenizer_batch_decode
|
self.disable_tokenizer_batch_decode = server_args.disable_tokenizer_batch_decode
|
||||||
|
|
||||||
self.watchdog = Watchdog.create(
|
self.soft_watchdog = Watchdog.create(
|
||||||
debug_name="DetokenizerManager",
|
debug_name="DetokenizerManager",
|
||||||
watchdog_timeout=server_args.soft_watchdog_timeout,
|
watchdog_timeout=server_args.soft_watchdog_timeout,
|
||||||
soft=True,
|
soft=True,
|
||||||
@@ -136,12 +136,12 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
|
|||||||
def event_loop(self):
|
def event_loop(self):
|
||||||
"""The event loop that handles requests"""
|
"""The event loop that handles requests"""
|
||||||
while True:
|
while True:
|
||||||
with self.watchdog.disable():
|
with self.soft_watchdog.disable():
|
||||||
recv_obj = self.recv_from_scheduler.recv_pyobj()
|
recv_obj = self.recv_from_scheduler.recv_pyobj()
|
||||||
output = self._request_dispatcher(recv_obj)
|
output = self._request_dispatcher(recv_obj)
|
||||||
if output is not None:
|
if output is not None:
|
||||||
self.send_to_tokenizer.send_pyobj(output)
|
self.send_to_tokenizer.send_pyobj(output)
|
||||||
self.watchdog.feed()
|
self.soft_watchdog.feed()
|
||||||
|
|
||||||
def trim_matched_stop(
|
def trim_matched_stop(
|
||||||
self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
|
self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
|
||||||
|
|||||||
@@ -78,7 +78,13 @@ class SchedulerMetricsMixin:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if self.enable_metrics:
|
if self.enable_metrics:
|
||||||
|
if self.server_args.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||||
|
engine_type = "prefill"
|
||||||
|
elif self.server_args.disaggregation_mode == DisaggregationMode.DECODE:
|
||||||
|
engine_type = "decode"
|
||||||
|
else:
|
||||||
engine_type = "unified"
|
engine_type = "unified"
|
||||||
|
|
||||||
labels = {
|
labels = {
|
||||||
"model_name": self.server_args.served_model_name,
|
"model_name": self.server_args.served_model_name,
|
||||||
"engine_type": engine_type,
|
"engine_type": engine_type,
|
||||||
|
|||||||
@@ -237,6 +237,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
if speculative_algorithm.is_none()
|
if speculative_algorithm.is_none()
|
||||||
else server_args.speculative_num_draft_tokens
|
else server_args.speculative_num_draft_tokens
|
||||||
)
|
)
|
||||||
|
self.validate_total_tokens = True
|
||||||
|
|
||||||
def init_tokenizer_and_processor(self):
|
def init_tokenizer_and_processor(self):
|
||||||
server_args = self.server_args
|
server_args = self.server_args
|
||||||
@@ -424,7 +425,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
|
|
||||||
if self.server_args.gc_warning_threshold_secs > 0.0:
|
if self.server_args.gc_warning_threshold_secs > 0.0:
|
||||||
configure_gc_warning(self.server_args.gc_warning_threshold_secs)
|
configure_gc_warning(self.server_args.gc_warning_threshold_secs)
|
||||||
self.watchdog = Watchdog.create(
|
self.soft_watchdog = Watchdog.create(
|
||||||
debug_name="TokenizerManager",
|
debug_name="TokenizerManager",
|
||||||
watchdog_timeout=self.server_args.soft_watchdog_timeout,
|
watchdog_timeout=self.server_args.soft_watchdog_timeout,
|
||||||
soft=True,
|
soft=True,
|
||||||
@@ -723,9 +724,10 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
"""Validates that the input token count and the requested token count doesn't exceed the model's context length."""
|
"""Validates that the input token count and the requested token count doesn't exceed the model's context length."""
|
||||||
# FIXME: unify the length validation logic with the one in the scheduler.
|
# FIXME: unify the length validation logic with the one in the scheduler.
|
||||||
_max_req_len = self.context_len
|
_max_req_len = self.context_len
|
||||||
|
|
||||||
input_token_num = len(input_ids) if input_ids is not None else 0
|
input_token_num = len(input_ids) if input_ids is not None else 0
|
||||||
input_token_num += self.reserve_input_token_num
|
input_token_num += self.reserve_input_token_num
|
||||||
|
|
||||||
|
# Validate input length
|
||||||
if input_token_num >= self.context_len:
|
if input_token_num >= self.context_len:
|
||||||
if self.server_args.allow_auto_truncate:
|
if self.server_args.allow_auto_truncate:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -741,16 +743,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
f"model's context length ({self.context_len} tokens)."
|
f"model's context length ({self.context_len} tokens)."
|
||||||
)
|
)
|
||||||
|
|
||||||
if isinstance(obj, EmbeddingReqInput) and self.is_generation:
|
# Validate total tokens (input + max_new_tokens)
|
||||||
raise ValueError(
|
|
||||||
"This model does not appear to be an embedding model by default. "
|
|
||||||
"Please add `--is-embedding` when launching the server or try another model."
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check total tokens (input + max_new_tokens)
|
|
||||||
max_new_tokens = obj.sampling_params.get("max_new_tokens")
|
max_new_tokens = obj.sampling_params.get("max_new_tokens")
|
||||||
if (
|
if (
|
||||||
max_new_tokens is not None
|
self.validate_total_tokens
|
||||||
|
and max_new_tokens is not None
|
||||||
and (max_new_tokens + input_token_num) >= _max_req_len
|
and (max_new_tokens + input_token_num) >= _max_req_len
|
||||||
):
|
):
|
||||||
if self.server_args.allow_auto_truncate:
|
if self.server_args.allow_auto_truncate:
|
||||||
@@ -773,10 +770,18 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
)
|
)
|
||||||
raise ValueError(error_msg)
|
raise ValueError(error_msg)
|
||||||
|
|
||||||
# Matryoshka embeddings validations
|
# Validate embedding requests
|
||||||
|
if isinstance(obj, EmbeddingReqInput) and self.is_generation:
|
||||||
|
raise ValueError(
|
||||||
|
"This model does not appear to be an embedding model by default. "
|
||||||
|
"Please add `--is-embedding` when launching the server or try another model."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate Matryoshka embeddings
|
||||||
if isinstance(obj, EmbeddingReqInput):
|
if isinstance(obj, EmbeddingReqInput):
|
||||||
self._validate_for_matryoshka_dim(obj)
|
self._validate_for_matryoshka_dim(obj)
|
||||||
|
|
||||||
|
# Validate custom logit processor
|
||||||
if isinstance(obj, GenerateReqInput):
|
if isinstance(obj, GenerateReqInput):
|
||||||
if (
|
if (
|
||||||
obj.return_hidden_states
|
obj.return_hidden_states
|
||||||
@@ -839,8 +844,18 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _validate_input_ids_in_vocab(
|
def _validate_input_ids_in_vocab(
|
||||||
self, input_ids: List[int], vocab_size: int
|
self, input_ids: Union[List[int], List[List[int]]], vocab_size: int
|
||||||
) -> None:
|
) -> None:
|
||||||
|
# Handle both single sequence and batch of sequences
|
||||||
|
if isinstance(input_ids[0], list):
|
||||||
|
# Batch of sequences
|
||||||
|
for seq in input_ids:
|
||||||
|
if any(id >= vocab_size for id in seq):
|
||||||
|
raise ValueError(
|
||||||
|
f"The input_ids {seq} contains values greater than the vocab size ({vocab_size})."
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Single sequence
|
||||||
if any(id >= vocab_size for id in input_ids):
|
if any(id >= vocab_size for id in input_ids):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"The input_ids {input_ids} contains values greater than the vocab size ({vocab_size})."
|
f"The input_ids {input_ids} contains values greater than the vocab size ({vocab_size})."
|
||||||
@@ -1420,11 +1435,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
async def handle_loop(self):
|
async def handle_loop(self):
|
||||||
"""The event loop that handles requests"""
|
"""The event loop that handles requests"""
|
||||||
while True:
|
while True:
|
||||||
with self.watchdog.disable():
|
with self.soft_watchdog.disable():
|
||||||
recv_obj = await self.recv_from_detokenizer.recv_pyobj()
|
recv_obj = await self.recv_from_detokenizer.recv_pyobj()
|
||||||
self._result_dispatcher(recv_obj)
|
self._result_dispatcher(recv_obj)
|
||||||
self.last_receive_tstamp = time.time()
|
self.last_receive_tstamp = time.time()
|
||||||
self.watchdog.feed()
|
self.soft_watchdog.feed()
|
||||||
|
|
||||||
def _handle_batch_output(
|
def _handle_batch_output(
|
||||||
self,
|
self,
|
||||||
@@ -1819,6 +1834,14 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
):
|
):
|
||||||
meta_info[attr_name] = getattr(recv_obj, attr_name)[index]
|
meta_info[attr_name] = getattr(recv_obj, attr_name)[index]
|
||||||
|
|
||||||
|
def _request_has_grammar(self, obj: GenerateReqInput) -> bool:
|
||||||
|
return (
|
||||||
|
obj.sampling_params.get("json_schema", None)
|
||||||
|
or obj.sampling_params.get("regex", None)
|
||||||
|
or obj.sampling_params.get("ebnf", None)
|
||||||
|
or obj.sampling_params.get("structural_tag", None)
|
||||||
|
)
|
||||||
|
|
||||||
def collect_metrics(self, state: ReqState, recv_obj: BatchStrOutput, i: int):
|
def collect_metrics(self, state: ReqState, recv_obj: BatchStrOutput, i: int):
|
||||||
completion_tokens = (
|
completion_tokens = (
|
||||||
recv_obj.completion_tokens[i]
|
recv_obj.completion_tokens[i]
|
||||||
@@ -1856,13 +1879,6 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
state.last_completion_tokens = completion_tokens
|
state.last_completion_tokens = completion_tokens
|
||||||
|
|
||||||
if state.finished:
|
if state.finished:
|
||||||
has_grammar = (
|
|
||||||
state.obj.sampling_params.get("json_schema", None)
|
|
||||||
or state.obj.sampling_params.get("regex", None)
|
|
||||||
or state.obj.sampling_params.get("ebnf", None)
|
|
||||||
or state.obj.sampling_params.get("structural_tag", None)
|
|
||||||
)
|
|
||||||
|
|
||||||
retraction_count = (
|
retraction_count = (
|
||||||
recv_obj.retraction_counts[i]
|
recv_obj.retraction_counts[i]
|
||||||
if getattr(recv_obj, "retraction_counts", None)
|
if getattr(recv_obj, "retraction_counts", None)
|
||||||
@@ -1876,7 +1892,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
completion_tokens,
|
completion_tokens,
|
||||||
recv_obj.cached_tokens[i],
|
recv_obj.cached_tokens[i],
|
||||||
state.finished_time - state.created_time,
|
state.finished_time - state.created_time,
|
||||||
has_grammar,
|
self._request_has_grammar(state.obj),
|
||||||
retraction_count,
|
retraction_count,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -511,8 +511,6 @@ def ci_download_with_validation_and_retry(
|
|||||||
kwargs["disable"] = True
|
kwargs["disable"] = True
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
|
||||||
|
|
||||||
# Retry loop for handling corrupted downloads
|
# Retry loop for handling corrupted downloads
|
||||||
for attempt in range(max_retries):
|
for attempt in range(max_retries):
|
||||||
hf_folder = snapshot_download(
|
hf_folder = snapshot_download(
|
||||||
|
|||||||
@@ -40,32 +40,20 @@ from sglang.srt.layers.quantization.modelopt_quant import (
|
|||||||
ModelOptFp4Config,
|
ModelOptFp4Config,
|
||||||
ModelOptFp8Config,
|
ModelOptFp8Config,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.model_loader.ci_weight_validation import (
|
||||||
|
ci_download_with_validation_and_retry,
|
||||||
|
ci_validate_and_cleanup_local_snapshot,
|
||||||
|
)
|
||||||
from sglang.srt.utils import find_local_repo_dir, log_info_on_rank0, print_warning_once
|
from sglang.srt.utils import find_local_repo_dir, log_info_on_rank0, print_warning_once
|
||||||
from sglang.utils import is_in_ci
|
from sglang.utils import is_in_ci
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from fastsafetensors import SafeTensorsFileLoader, SingleGroup
|
from fastsafetensors import SafeTensorsFileLoader, SingleGroup
|
||||||
except ImportError:
|
except ImportError as e:
|
||||||
|
SafeTensorsFileLoader = SingleGroup = None
|
||||||
class PlaceholderModule:
|
|
||||||
def __init__(self, name):
|
|
||||||
self.name = name
|
|
||||||
|
|
||||||
def __getattr__(self, name):
|
|
||||||
raise ImportError(f"Please install {self.name}")
|
|
||||||
|
|
||||||
fastsafetensors = PlaceholderModule("fastsafetensors")
|
|
||||||
SafeTensorsFileLoader = None
|
|
||||||
SingleGroup = None
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# use system-level temp directory for file locks, so that multiple users
|
|
||||||
# can share the same lock without error.
|
|
||||||
# lock files in the temp directory will be automatically deleted when the
|
|
||||||
# system reboots, so users will not complain about annoying lock files
|
|
||||||
temp_dir = tempfile.gettempdir()
|
|
||||||
|
|
||||||
|
|
||||||
def enable_hf_transfer():
|
def enable_hf_transfer():
|
||||||
"""automatically activates hf_transfer"""
|
"""automatically activates hf_transfer"""
|
||||||
@@ -82,10 +70,11 @@ def enable_hf_transfer():
|
|||||||
enable_hf_transfer()
|
enable_hf_transfer()
|
||||||
|
|
||||||
|
|
||||||
class DisabledTqdm(tqdm):
|
# use system-level temp directory for file locks, so that multiple users
|
||||||
def __init__(self, *args, **kwargs):
|
# can share the same lock without error.
|
||||||
kwargs["disable"] = True
|
# lock files in the temp directory will be automatically deleted when the
|
||||||
super().__init__(*args, **kwargs)
|
# system reboots, so users will not complain about annoying lock files
|
||||||
|
temp_dir = tempfile.gettempdir()
|
||||||
|
|
||||||
|
|
||||||
def get_lock(
|
def get_lock(
|
||||||
@@ -169,6 +158,12 @@ def replace_substrings(key: str, substring_mapping: dict[str, str]) -> str:
|
|||||||
return key
|
return key
|
||||||
|
|
||||||
|
|
||||||
|
class DisabledTqdm(tqdm):
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
kwargs["disable"] = True
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
# TODO(woosuk): Move this to other place.
|
# TODO(woosuk): Move this to other place.
|
||||||
def get_quant_config(
|
def get_quant_config(
|
||||||
model_config: ModelConfig,
|
model_config: ModelConfig,
|
||||||
@@ -194,6 +189,7 @@ def get_quant_config(
|
|||||||
if hf_quant_config is not None:
|
if hf_quant_config is not None:
|
||||||
hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
|
hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
|
||||||
return quant_cls.from_config(hf_quant_config)
|
return quant_cls.from_config(hf_quant_config)
|
||||||
|
|
||||||
# In case of bitsandbytes/QLoRA, get quant config from the adapter model.
|
# In case of bitsandbytes/QLoRA, get quant config from the adapter model.
|
||||||
if model_config.quantization == "bitsandbytes":
|
if model_config.quantization == "bitsandbytes":
|
||||||
if (
|
if (
|
||||||
@@ -204,9 +200,9 @@ def get_quant_config(
|
|||||||
model_name_or_path = load_config.model_loader_extra_config[
|
model_name_or_path = load_config.model_loader_extra_config[
|
||||||
"qlora_adapter_name_or_path"
|
"qlora_adapter_name_or_path"
|
||||||
]
|
]
|
||||||
|
|
||||||
else:
|
else:
|
||||||
model_name_or_path = model_config.model_path
|
model_name_or_path = model_config.model_path
|
||||||
|
|
||||||
is_local = os.path.isdir(model_name_or_path)
|
is_local = os.path.isdir(model_name_or_path)
|
||||||
if not is_local:
|
if not is_local:
|
||||||
# Download the config files.
|
# Download the config files.
|
||||||
@@ -357,10 +353,6 @@ def _find_local_hf_snapshot_dir_unlocked(
|
|||||||
# Only perform cache validation and cleanup in CI to avoid
|
# Only perform cache validation and cleanup in CI to avoid
|
||||||
# unnecessary overhead for regular users
|
# unnecessary overhead for regular users
|
||||||
if is_in_ci() and local_weight_files:
|
if is_in_ci() and local_weight_files:
|
||||||
from sglang.srt.model_loader.ci_weight_validation import (
|
|
||||||
ci_validate_and_cleanup_local_snapshot,
|
|
||||||
)
|
|
||||||
|
|
||||||
is_valid = ci_validate_and_cleanup_local_snapshot(
|
is_valid = ci_validate_and_cleanup_local_snapshot(
|
||||||
model_name_or_path, found_local_snapshot_dir, local_weight_files
|
model_name_or_path, found_local_snapshot_dir, local_weight_files
|
||||||
)
|
)
|
||||||
@@ -443,23 +435,10 @@ def download_weights_from_hf(
|
|||||||
allow_patterns = [pattern]
|
allow_patterns = [pattern]
|
||||||
break
|
break
|
||||||
|
|
||||||
# Only perform validation and retry in CI to avoid overhead for regular users
|
|
||||||
if is_in_ci():
|
|
||||||
from sglang.srt.model_loader.ci_weight_validation import (
|
|
||||||
ci_download_with_validation_and_retry,
|
|
||||||
)
|
|
||||||
|
|
||||||
return ci_download_with_validation_and_retry(
|
|
||||||
model_name_or_path=model_name_or_path,
|
|
||||||
allow_patterns=allow_patterns,
|
|
||||||
ignore_patterns=ignore_patterns,
|
|
||||||
cache_dir=cache_dir,
|
|
||||||
revision=revision,
|
|
||||||
max_retries=max_retries,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
# Simple download without validation for non-CI environments
|
|
||||||
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
log_info_on_rank0(logger, f"Using model weights format {allow_patterns}")
|
||||||
|
|
||||||
|
if not is_in_ci():
|
||||||
|
# Simple download without validation for non-CI environments
|
||||||
hf_folder = snapshot_download(
|
hf_folder = snapshot_download(
|
||||||
model_name_or_path,
|
model_name_or_path,
|
||||||
allow_patterns=allow_patterns,
|
allow_patterns=allow_patterns,
|
||||||
@@ -470,6 +449,16 @@ def download_weights_from_hf(
|
|||||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||||
)
|
)
|
||||||
return hf_folder
|
return hf_folder
|
||||||
|
else:
|
||||||
|
# Only perform validation and retry in CI to avoid overhead for regular users
|
||||||
|
return ci_download_with_validation_and_retry(
|
||||||
|
model_name_or_path=model_name_or_path,
|
||||||
|
allow_patterns=allow_patterns,
|
||||||
|
ignore_patterns=ignore_patterns,
|
||||||
|
cache_dir=cache_dir,
|
||||||
|
revision=revision,
|
||||||
|
max_retries=max_retries,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def download_safetensors_index_file_from_hf(
|
def download_safetensors_index_file_from_hf(
|
||||||
|
|||||||
Reference in New Issue
Block a user