[Model] Add K2 Horizon FP8 checkpoint support (#38033)
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
co-authored by
Xiaoyu Zhang
parent
e980c1a2f1
commit
756d0e0a85
@@ -661,10 +661,10 @@ def _validate_mova_config(
|
|||||||
"the released checkpoints persist float32 dtype metadata but "
|
"the released checkpoints persist float32 dtype metadata but "
|
||||||
"their weights and validated runtime contract are BF16."
|
"their weights and validated runtime contract are BF16."
|
||||||
)
|
)
|
||||||
if quant_config is not None:
|
if quant_config is not None and quant_config.get_name() != "compressed_tensors":
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"Native xLLM/K2 Horizon serving does not support quantized "
|
"Native xLLM/K2 Horizon serving supports only "
|
||||||
"model weights"
|
"compressed-tensors quantized model weights"
|
||||||
)
|
)
|
||||||
|
|
||||||
runtime = get_exec()
|
runtime = get_exec()
|
||||||
@@ -1686,6 +1686,14 @@ class XllmModel(nn.Module):
|
|||||||
class XllmForCausalLM(nn.Module):
|
class XllmForCausalLM(nn.Module):
|
||||||
fall_back_to_pt_during_load = False
|
fall_back_to_pt_during_load = False
|
||||||
|
|
||||||
|
# Quantized checkpoints store these projections separately. This mapping
|
||||||
|
# lets quantization configs resolve fused runtime modules and their ignore
|
||||||
|
# lists consistently.
|
||||||
|
packed_modules_mapping = {
|
||||||
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||||
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||||
|
}
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: PretrainedConfig,
|
config: PretrainedConfig,
|
||||||
|
|||||||
@@ -266,15 +266,35 @@ def test_native_xllm_requires_bfloat16(monkeypatch):
|
|||||||
_validate_mova_config(config, quant_config=None)
|
_validate_mova_config(config, quant_config=None)
|
||||||
|
|
||||||
|
|
||||||
def test_native_xllm_rejects_quantized_weights(monkeypatch):
|
def test_native_xllm_accepts_compressed_tensors_quantization(monkeypatch):
|
||||||
config = XllmConfig(num_values=0, num_experts=0)
|
config = XllmConfig(num_values=0, num_experts=0)
|
||||||
|
quant_config = SimpleNamespace(get_name=lambda: "compressed_tensors")
|
||||||
|
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16)
|
||||||
|
|
||||||
|
with get_context().override_server_args(**_native_runtime_config()):
|
||||||
|
_validate_mova_config(config, quant_config=quant_config)
|
||||||
|
|
||||||
|
|
||||||
|
def test_native_xllm_rejects_other_quantization(monkeypatch):
|
||||||
|
config = XllmConfig(num_values=0, num_experts=0)
|
||||||
|
quant_config = SimpleNamespace(get_name=lambda: "awq")
|
||||||
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16)
|
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16)
|
||||||
|
|
||||||
with (
|
with (
|
||||||
get_context().override_server_args(**_native_runtime_config()),
|
get_context().override_server_args(**_native_runtime_config()),
|
||||||
pytest.raises(ValueError, match="does not support quantized"),
|
pytest.raises(ValueError, match="supports only compressed-tensors"),
|
||||||
):
|
):
|
||||||
_validate_mova_config(config, quant_config=object())
|
_validate_mova_config(config, quant_config=quant_config)
|
||||||
|
|
||||||
|
|
||||||
|
def test_native_xllm_declares_quantized_fused_module_mapping():
|
||||||
|
assert XllmForCausalLM.packed_modules_mapping == {
|
||||||
|
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||||
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||||
|
}
|
||||||
|
assert K2HorizonForCausalLM.packed_modules_mapping == (
|
||||||
|
XllmForCausalLM.packed_modules_mapping
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_native_xllm_accepts_bfloat16_without_expert_remapping(monkeypatch):
|
def test_native_xllm_accepts_bfloat16_without_expert_remapping(monkeypatch):
|
||||||
|
|||||||
Reference in New Issue
Block a user