phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-03 17:29:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

rename quant types to use _mixed naming

Browse source
This commit is contained in:
Ashwin Bharambe 2025-04-07 12:57:58 -07:00
parent b239c57c54
commit 76004eacb4
3 changed files with 11 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/apis/inference/inference.py
View file

@ -97,18 +97,18 @@ class QuantizationType(Enum):
"""Type of model quantization to run inference with. """Type of model quantization to run inference with.
:cvar bf16: BFloat16 typically this means _no_ quantization :cvar bf16: BFloat16 typically this means _no_ quantization
:cvar fp8: 8-bit floating point quantization :cvar fp8_mixed: 8-bit floating point quantization with mixed precision
:cvar int4: 4-bit integer quantization :cvar int4_mixed: 4-bit integer quantization with mixed precision
""" """
bf16 = "bf16" bf16 = "bf16"
fp8 = "fp8" fp8_mixed = "fp8_mixed"
int4 = "int4" int4_mixed = "int4_mixed"
@json_schema_type @json_schema_type
class Fp8QuantizationConfig(BaseModel): class Fp8QuantizationConfig(BaseModel):
type: Literal["fp8"] = "fp8" type: Literal["fp8_mixed"] = "fp8_mixed"
@json_schema_type @json_schema_type
@ -124,7 +124,7 @@ class Int4QuantizationConfig(BaseModel):
:param scheme: Quantization scheme to use. Defaults to "int4_weight_int8_dynamic_activation" :param scheme: Quantization scheme to use. Defaults to "int4_weight_int8_dynamic_activation"
""" """
type: Literal["int4"] = "int4" type: Literal["int4_mixed"] = "int4_mixed"
scheme: Optional[str] = "int4_weight_int8_dynamic_activation" scheme: Optional[str] = "int4_weight_int8_dynamic_activation"

2
llama_stack/models/llama/llama4/quantization/loader.py
View file

@ -91,7 +91,7 @@ def convert_to_quantized_model(
log_status(f"Rank {rank}: Quantizing int4 weights from bf16") log_status(f"Rank {rank}: Quantizing int4 weights from bf16")
def apply_quantization(_, weight): def apply_quantization(_, weight):
return quantize_int4(weight, output_device=torch.device("cuda")) return quantize_int4(weight, fp8_activation_scale_ub, output_device=torch.device("cuda"))
else: else:
fp8_scales_path = os.path.join(checkpoint_dir, f"fp8_scales_{rank}.pt") fp8_scales_path = os.path.join(checkpoint_dir, f"fp8_scales_{rank}.pt")

8
llama_stack/providers/inline/inference/meta_reference/generators.py
View file

@ -133,9 +133,9 @@ class Llama4Generator:
ckpt_dir = model_checkpoint_dir(resolved_model.descriptor()) ckpt_dir = model_checkpoint_dir(resolved_model.descriptor())
if config.quantization: if config.quantization:
if config.quantization.type == "fp8": if config.quantization.type == "fp8_mixed":
quantization_mode = QuantizationMode.fp8_mixed quantization_mode = QuantizationMode.fp8_mixed
elif config.quantization.type == "int4": elif config.quantization.type == "int4_mixed":
quantization_mode = QuantizationMode.int4_mixed quantization_mode = QuantizationMode.int4_mixed
elif config.quantization.type == "bf16": elif config.quantization.type == "bf16":
quantization_mode = None quantization_mode = None
@ -226,9 +226,9 @@ class Llama3Generator:
ckpt_dir = model_checkpoint_dir(resolved_model.descriptor()) ckpt_dir = model_checkpoint_dir(resolved_model.descriptor())
if config.quantization: if config.quantization:
if config.quantization.type == "fp8": if config.quantization.type == "fp8_mixed":
quantization_mode = QuantizationMode.fp8_mixed quantization_mode = QuantizationMode.fp8_mixed
elif config.quantization.type == "int4": elif config.quantization.type == "int4_mixed":
quantization_mode = QuantizationMode.int4_mixed quantization_mode = QuantizationMode.int4_mixed
elif config.quantization.type == "bf16": elif config.quantization.type == "bf16":
quantization_mode = None quantization_mode = None