try fix model loading quen3.6-35b

This commit is contained in:
Dobromir Popov
2026-07-07 18:36:29 +02:00
parent f1eea5b6d4
commit cdd2699e63
5 changed files with 368 additions and 28 deletions

View File

@@ -134,8 +134,9 @@ class TorchModelShard:
self.model.to(self.device)
except Exception as exc:
if _looks_like_oom(exc):
memory_kind = "VRAM" if self.device.type == "cuda" else "RAM"
raise InsufficientVRAMError(
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} "
f"with {quantization} quantization; choose a smaller shard or lower quantization"
) from exc
raise
@@ -411,7 +412,7 @@ def _should_partial_materialize_shard(
return False
if total_layers_hint is None:
return False
return not (shard_start == 0 and shard_end >= total_layers_hint - 1)
return True
def _load_partial_model_from_snapshot(
@@ -476,7 +477,7 @@ def _load_partial_model_from_snapshot(
)
with init_empty_weights_fn():
model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype)
model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype)
tie_weights = getattr(model, "tie_weights", None)
if callable(tie_weights):
tie_weights()
@@ -498,7 +499,7 @@ def _load_partial_model_from_snapshot(
for tensor_name in names:
set_tensor_fn(
model,
tensor_name,
_checkpoint_tensor_name_for_model(model, tensor_name),
device,
value=handle.get_tensor(tensor_name),
dtype=dtype,
@@ -569,38 +570,74 @@ def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
return torch.bfloat16
def _causal_lm_config(cfg: Any) -> Any:
"""Use the text decoder config for composite VLM/MoE presets."""
get_text_config = getattr(cfg, "get_text_config", None)
if callable(get_text_config):
try:
return get_text_config()
except Exception:
pass
text_config = getattr(cfg, "text_config", None)
if text_config is not None:
return text_config
return cfg
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
inner = getattr(model, "model", None)
if inner is not None and hasattr(inner, "language_model"):
return tensor_name
if ".language_model." in tensor_name:
return tensor_name.replace(".language_model.", ".")
return tensor_name
def _transformer_backbone(model: Any) -> Any:
if hasattr(model, "model"):
inner = model.model
language_model = getattr(inner, "language_model", None)
if language_model is not None:
return language_model
return inner
if hasattr(model, "transformer"):
return model.transformer
raise ModelBackendError(
"unsupported HuggingFace model architecture: no transformer backbone found"
)
def _model_layers(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "layers"):
return model.model.layers
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
return model.transformer.h
backbone = _transformer_backbone(model)
for attr in ("layers", "h", "blocks"):
layers = getattr(backbone, attr, None)
if layers is not None:
return layers
raise ModelBackendError(
"unsupported HuggingFace model architecture: no transformer layers found"
)
def _embed_tokens(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
return model.model.embed_tokens
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
return model.transformer.wte
backbone = _transformer_backbone(model)
for attr in ("embed_tokens", "wte"):
embed = getattr(backbone, attr, None)
if embed is not None:
return embed
raise ModelBackendError(
"unsupported HuggingFace model architecture: no token embeddings found"
)
def _position_embeddings(model: Any) -> Any | None:
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
return model.transformer.wpe
return None
backbone = _transformer_backbone(model)
return getattr(backbone, "wpe", None)
def _rotary_embedding_module(model: Any) -> Any | None:
if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
return model.model.rotary_emb
if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
return model.transformer.rotary_emb
return None
backbone = _transformer_backbone(model)
return getattr(backbone, "rotary_emb", None)
def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
@@ -627,10 +664,11 @@ def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> l
def _final_norm(model: Any) -> Any | None:
if hasattr(model, "model") and hasattr(model.model, "norm"):
return model.model.norm
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
return model.transformer.ln_f
backbone = _transformer_backbone(model)
for attr in ("norm", "ln_f", "final_layer_norm"):
norm = getattr(backbone, attr, None)
if norm is not None:
return norm
return None
@@ -743,7 +781,12 @@ def _looks_like_oom(exc: BaseException) -> bool:
current: BaseException | None = exc
while current is not None:
text = str(current).lower()
if "out of memory" in text or "cuda error: out of memory" in text:
if (
"out of memory" in text
or "cuda error: out of memory" in text
or "paging file is too small" in text
or "os error 1455" in text
):
return True
current = current.__cause__ or current.__context__
return False