diff --git a/python/sglang/srt/constants.py b/python/sglang/srt/constants.py index 76d3b1cd3ad5..2d47c6337f93 100644 --- a/python/sglang/srt/constants.py +++ b/python/sglang/srt/constants.py @@ -10,3 +10,5 @@ ] HEALTH_CHECK_RID_PREFIX = "HEALTH_CHECK" + +GIB_BYTES = 1073741824 # 1024**3 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 291e2c2de3bd..c78400db3024 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1,4 +1,5 @@ # Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/layers/quantization/fp8.py +# Online quantization adapted from https://github.com/vllm-project/vllm/pull/29196 and https://github.com/vllm-project/vllm/pull/31914 from __future__ import annotations @@ -65,6 +66,7 @@ ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin +from sglang.srt.layers.quantization.online_quantization import CopyNumelCounter from sglang.srt.layers.quantization.unquant import ( UnquantizedFusedMoEMethod, UnquantizedLinearMethod, @@ -179,6 +181,9 @@ def __init__( raise ValueError("MXFP8 requires weight_block_size=[1, 32].") self.weight_block_size = weight_block_size + if not is_checkpoint_fp8_serialized: + logger.info("Quantizing model to FP8 on the fly during loading.") + def get_name(self) -> str: return "mxfp8" if self.use_mxfp8 else "fp8" @@ -330,6 +335,11 @@ def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]): self.use_aiter_fp8_per_token = envs.SGLANG_USE_AITER_FP8_PER_TOKEN.get() self.use_per_token_if_dynamic = False + if self.block_quant and not self.is_checkpoint_fp8_serialized: + raise ValueError( + f"block_quant={self.block_quant} is not supported along online quantization (is_checkpoint_fp8_serialized={self.is_checkpoint_fp8_serialized})." + ) + def validate_block_quant_shapes( self, input_size: int, @@ -387,7 +397,7 @@ def create_weights( layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition layer.orig_dtype = params_dtype - weight_loader = extra_weight_attrs.get("weight_loader") + original_weight_loader = extra_weight_attrs.get("weight_loader") if self.block_quant: block_n, block_k = self.quant_config.weight_block_size @@ -400,13 +410,27 @@ def create_weights( skip_block_quant_check, ) - # Create the weight - weight_dtype = ( - torch.float8_e4m3fn if self.is_checkpoint_fp8_serialized else params_dtype - ) + # Wrap weight loader for online quantization if checkpoint is not fp8 serialized + if not self.is_checkpoint_fp8_serialized: + layer.weight_scale = None + layer._loaded_numel = 0 + device = torch.device("meta") + weight_dtype = params_dtype + + weight_loader = self.get_online_weight_loader(layer, original_weight_loader) + else: + weight_loader = original_weight_loader + weight_dtype = torch.float8_e4m3fn + device = torch.get_default_device() + + layer._load_device = torch.get_default_device() + weight = ModelWeightParameter( data=torch.empty( - output_size_per_partition, input_size_per_partition, dtype=weight_dtype + output_size_per_partition, + input_size_per_partition, + dtype=weight_dtype, + device=device, ), input_dim=1, output_dim=0, @@ -414,19 +438,15 @@ def create_weights( ) layer.register_parameter("weight", weight) - # If checkpoint is serialized fp8, load them. - # Otherwise, wait until process_weights_after_loading. - if self.is_checkpoint_fp8_serialized: - # WEIGHT SCALE + if not self.is_checkpoint_fp8_serialized: + layer.input_scale = None + else: if self.block_quant: if hasattr(self.quant_config, "activation_scheme"): assert self.quant_config.activation_scheme == "dynamic" elif hasattr(self.quant_config, "linear_activation_scheme"): assert self.quant_config.linear_activation_scheme == "dynamic" - if self.use_mxfp8 and not self.is_checkpoint_fp8_serialized: - raise ValueError( - "MXFP8 requires fp8-serialized checkpoint for linear layers." - ) + scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32 scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.empty scale = BlockQuantScaleParameter( @@ -444,6 +464,7 @@ def create_weights( scale[:] = torch.finfo(torch.float32).min layer.register_parameter("weight_scale_inv", scale) else: + # If checkpoint is serialized fp8, load the scales. scale = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, @@ -469,6 +490,88 @@ def create_weights( else: layer.register_parameter("input_scale", None) + def get_online_weight_loader(self, layer, original_weight_loader): + def online_fp8_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: int | None = None, + ): + assert param.device.type == "meta" + + if layer._loaded_numel == 0: + layer.weight = ModelWeightParameter( + data=torch.empty_like(param.data, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=online_fp8_weight_loader, + ) + + # Move to device for faster quantization. At this point, loaded weights are already materialized on CPU RAM. + loaded_weight = loaded_weight.to(layer._load_device) + + param = layer.weight + + kwargs = {} + if shard_id is not None: + kwargs["loaded_shard_id"] = shard_id + + # In case TP>1, the weight loader logic uses narrow so we can not directly rely on `param.shape` or `loaded_weight.shape`. + copy_numel_counter = CopyNumelCounter() + with copy_numel_counter: + # Loads the quantized weight. + original_weight_loader(param, loaded_weight, **kwargs) + + layer._loaded_numel += copy_numel_counter.copied_numel + target_loaded_numel = layer.weight.numel() + + assert ( + layer._loaded_numel <= target_loaded_numel + ), f"target_loaded_numel={target_loaded_numel}, layer._loaded_numel={layer._loaded_numel}" + + # Delay online quantization until all tensor shards (e.g. q_proj, k_proj, v_proj) are loaded, to avoid having to re-quantize later on. + if layer._loaded_numel == target_loaded_numel: + full_loaded_weight = layer.weight + + if self.use_mxfp8: + # MXFP8 quantization + qweight, weight_scale = mxfp8_group_quantize(full_loaded_weight) + # Register weight_scale_inv if not already registered + if ( + not hasattr(layer, "weight_scale_inv") + or layer.weight_scale_inv is None + ): + layer.register_parameter( + "weight_scale_inv", + Parameter(weight_scale, requires_grad=False), + ) + else: + layer.weight_scale_inv.data = weight_scale + layer.weight_scale_inv.requires_grad_(False) + layer.weight_scale_inv.format_ue8m0 = True + elif ( + self.cutlass_fp8_supported + or self.use_marlin + or (_use_aiter and self.use_aiter_fp8_per_token) + ): + # apply per-channel quantization default as + # cutlass sgl-kernel and marlin only support per-channel scale + qweight, weight_scale = per_token_group_quant_fp8( + full_loaded_weight, loaded_weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + if _use_aiter and self.use_aiter_fp8_per_token: + self.use_per_token_if_dynamic = True + qweight = shuffle_weight(qweight.contiguous(), (16, 16)) + layer.weight_scale = weight_scale + else: + # per-tensor quantization + qweight, weight_scale = input_to_float8(full_loaded_weight) + layer.weight_scale = weight_scale + + layer.weight = torch.nn.Parameter(qweight, requires_grad=False) + + return online_fp8_weight_loader + def process_weights_after_loading_block_quant(self, layer: Module) -> None: # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: @@ -490,7 +593,10 @@ def process_weights_after_loading_block_quant(self, layer: Module) -> None: return elif self.use_mxfp8: if not self.is_checkpoint_fp8_serialized: - self._quantize_mxfp8_weights(layer) + # Quantization already happened during weight loading via get_online_weight_loader. + assert layer.weight.dtype == torch.float8_e4m3fn + assert hasattr(layer, "weight_scale_inv") + assert layer.weight_scale_inv.format_ue8m0 return # MXFP8 scales are stored as UE8M0 uint8; no requantization here. # Keep parameter object to preserve weight_loader attrs for hot reload. @@ -609,39 +715,16 @@ def process_weights_after_loading(self, layer: Module) -> None: if self.block_quant: self.process_weights_after_loading_block_quant(layer) else: - layer.weight = Parameter(layer.weight.data, requires_grad=False) - - # If checkpoint not serialized fp8, quantize the weights. if not self.is_checkpoint_fp8_serialized: - if ( - self.cutlass_fp8_supported - or self.use_marlin - or (_use_aiter and self.use_aiter_fp8_per_token) - ): - # apply per-channel quantization default as - # cutlass sgl-kernel and marlin only support per-channel scale - qweight, weight_scale = per_token_group_quant_fp8( - layer.weight, layer.weight.shape[-1] - ) - weight_scale = weight_scale.t().contiguous() - if _use_aiter and self.use_aiter_fp8_per_token: - self.use_per_token_if_dynamic = True - qweight = shuffle_weight(qweight.contiguous(), (16, 16)) - else: - # per-tensor quantization - qweight, weight_scale = input_to_float8(layer.weight) + assert layer.weight_scale is not None + layer.weight.data = layer.weight.data.t() - # Update the layer with the new values. - layer.weight = Parameter(qweight.t(), requires_grad=False) - layer.weight_scale = Parameter(weight_scale, requires_grad=False) - layer.input_scale = None + layer.weight = Parameter(layer.weight.data, requires_grad=False) + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) # If checkpoint is fp8, handle that there are N scales for N # shards in a fused module - else: - layer.weight_scale = Parameter( - layer.weight_scale.data, requires_grad=False - ) + if self.is_checkpoint_fp8_serialized: if ( hasattr(self.quant_config, "activation_scheme") and self.quant_config.activation_scheme == "static" @@ -827,6 +910,11 @@ def __init__(self, quant_config: Fp8Config): is_sm100_supported() or is_sm90_supported() or is_sm120_supported() ), "cutlass_fp8 MoE requires SM90, SM100, or SM120 GPUs" + if not self.quant_config.is_checkpoint_fp8_serialized and _use_hip_int4: + raise NotImplementedError( + f"Online MOE FP8 quantization (is_checkpoint_fp8_serialized={self.quant_config.is_checkpoint_fp8_serialized}) along SGLANG_INT4_WEIGHT=1 is not supported at the moment. Please open an issue." + ) + @staticmethod def is_deepgemm_moe_runner_backend_enabled() -> bool: """Check if MoE will actually use DeepGEMM runner for FP8.""" @@ -857,8 +945,21 @@ def create_weights( self.with_bias = with_bias from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + original_weight_loader = extra_weight_attrs.get("weight_loader") + if self.quant_config.is_checkpoint_fp8_serialized: params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn + weight_loader = original_weight_loader + device = torch.get_default_device() + else: + # Online quantization: use original dtype and meta device + weight_loader = self.get_online_weight_loader(layer, original_weight_loader) + device = torch.device("meta") + + layer._load_device = torch.get_default_device() + layer._w13_loaded_numel = 0 + layer._w2_loaded_numel = 0 + tp_size = get_tensor_model_parallel_world_size() w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes( @@ -921,6 +1022,7 @@ def create_weights( 2 * intermediate_size_per_partition, hidden_size // 8, dtype=params_dtype, + device=device, ), requires_grad=False, ) @@ -930,6 +1032,7 @@ def create_weights( hidden_size, intermediate_size_per_partition // 8, dtype=params_dtype, + device=device, ), requires_grad=False, ) @@ -940,6 +1043,7 @@ def create_weights( w13_up_dim, hidden_size, dtype=params_dtype, + device=device, ), requires_grad=False, ) @@ -949,6 +1053,7 @@ def create_weights( hidden_size, w2_up_dim, dtype=params_dtype, + device=device, ), requires_grad=False, ) @@ -956,6 +1061,7 @@ def create_weights( extra_weight_attrs.update( {"weight_padded": weight_padded}, ) + extra_weight_attrs["weight_loader"] = weight_loader layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) @@ -1111,6 +1217,117 @@ def create_weights( layer.w13_input_scale = None layer.w2_input_scale = None + def get_online_weight_loader(self, layer, original_weight_loader): + def online_fp8_moe_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ): + # Determine which weight parameter we're loading (w13 or w2) + is_w13 = "w13" in weight_name + is_w2 = "w2" in weight_name + + # Initialize weight on device if first load + if is_w13 and layer._w13_loaded_numel == 0: + layer.w13_weight = torch.nn.Parameter( + torch.empty_like(param.data, device=layer._load_device), + requires_grad=False, + ) + param = layer.w13_weight + elif is_w2 and layer._w2_loaded_numel == 0: + layer.w2_weight = torch.nn.Parameter( + torch.empty_like(param.data, device=layer._load_device), + requires_grad=False, + ) + param = layer.w2_weight + + # Move to device for faster quantization + loaded_weight = loaded_weight.to(layer._load_device) + + if is_w13: + param = layer.w13_weight + elif is_w2: + param = layer.w2_weight + + # In case TP>1, the weight loader logic uses narrow so we can not directly rely on `param.shape` or `loaded_weight.shape`. + copy_numel_counter = CopyNumelCounter() + with copy_numel_counter: + original_weight_loader( + param, loaded_weight, weight_name, shard_id, expert_id + ) + + if is_w13: + layer._w13_loaded_numel += copy_numel_counter.copied_numel + target_loaded_numel = layer.w13_weight.numel() + current_loaded = layer._w13_loaded_numel + elif is_w2: + layer._w2_loaded_numel += copy_numel_counter.copied_numel + target_loaded_numel = layer.w2_weight.numel() + current_loaded = layer._w2_loaded_numel + else: + raise ValueError("Expected w13 or w2.") + + assert ( + current_loaded <= target_loaded_numel + ), f"target_loaded_numel={target_loaded_numel}, current_loaded={current_loaded}" + + # Delay online quantization until all tensor shards (e.g. w1 and w3) are loaded, to avoid having to re-quantize later on. + if is_w13 and layer._w13_loaded_numel == target_loaded_numel: + if not self.use_mxfp8: + self._quantize_w13_online(layer) + elif is_w2 and layer._w2_loaded_numel == target_loaded_numel: + if not self.use_mxfp8: + self._quantize_w2_online(layer) + + # For MXFP8, call _process_mxfp8_moe_weights once when both w13 and w2 are fully loaded + if self.use_mxfp8: + w13_target_numel = layer.w13_weight.numel() + w2_target_numel = layer.w2_weight.numel() + if ( + layer._w13_loaded_numel == w13_target_numel + and layer._w2_loaded_numel == w2_target_numel + ): + self._process_mxfp8_moe_weights(layer, quantize=True) + + return online_fp8_moe_weight_loader + + def _quantize_w13_online(self, layer): + """Quantize w13_weight after all weights are loaded.""" + # If ROCm, fp8_dtype will be float8_e4m3fnuz (MI300x HW) + w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype) + + # Re-initialize w13_scale because we directly quantize + # merged w13 weights and generate a single scaling factor. + layer.w13_weight_scale = torch.nn.Parameter( + torch.ones( + layer.num_local_experts, + dtype=torch.float32, + device=w13_weight.device, + ), + requires_grad=False, + ) + + for expert in range(layer.num_local_experts): + w13_weight[expert, :, :], layer.w13_weight_scale[expert] = scaled_fp8_quant( + layer.w13_weight.data[expert, :, :] + ) + + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + + def _quantize_w2_online(self, layer): + """Quantize w2_weight after all weights are loaded.""" + # If ROCm, fp8_dtype will be float8_e4m3fnuz (MI300x HW) + w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype) + + for expert in range(layer.num_local_experts): + w2_weight[expert, :, :], layer.w2_weight_scale[expert] = scaled_fp8_quant( + layer.w2_weight.data[expert, :, :] + ) + + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + def process_weights_after_loading_block_quant(self, layer: Module) -> None: # If ROCm, normalize the weights and scales to e4m3fnuz if _is_fp8_fnuz: @@ -1159,9 +1376,17 @@ def process_weights_after_loading_block_quant(self, layer: Module) -> None: ), "Fp8MoEMethod on CPU requires that CPU has AMX support" _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) elif self.use_mxfp8: - self._process_mxfp8_moe_weights( - layer, quantize=not self.quant_config.is_checkpoint_fp8_serialized - ) + if not self.quant_config.is_checkpoint_fp8_serialized: + # Quantization already happened during weight loading via get_online_weight_loader. + assert layer.w13_weight.dtype == torch.float8_e4m3fn + assert layer.w2_weight.dtype == torch.float8_e4m3fn + assert hasattr(layer, "w13_weight_scale_inv") + assert hasattr(layer, "w2_weight_scale_inv") + assert layer.w13_weight_scale_inv.format_ue8m0 + assert layer.w2_weight_scale_inv.format_ue8m0 + else: + # Checkpoint is already FP8 serialized, just need to swizzle the scales + self._process_mxfp8_moe_weights(layer, quantize=False) else: # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0 from sglang.srt.layers import deep_gemm_wrapper @@ -1394,11 +1619,8 @@ def _copy_or_rebind(param: Parameter, new_value: torch.Tensor) -> None: _copy_or_rebind(layer.w2_weight_scale_inv, w2_s) layer.w13_weight.requires_grad_(False) layer.w2_weight.requires_grad_(False) - layer.w13_weight_scale_inv.requires_grad_(False) layer.w2_weight_scale_inv.requires_grad_(False) - layer.w13_weight_scale_inv.format_ue8m0 = True layer.w2_weight_scale_inv.format_ue8m0 = True - layer.w13_input_scale = None layer.w2_input_scale = None if ( @@ -1421,29 +1643,9 @@ def process_weights_after_loading(self, layer: Module) -> None: # If checkpoint is fp16 or bfloat16, quantize in place. elif not self.quant_config.is_checkpoint_fp8_serialized: - # If ROCm, fp8_dtype will be float8_e4m3fnuz (MI300x HW) - w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype) - w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype) - - # Re-initialize w13_scale because we directly quantize - # merged w13 weights and generate a single scaling factor. - layer.w13_weight_scale = torch.nn.Parameter( - torch.ones( - layer.num_local_experts, - dtype=torch.float32, - device=w13_weight.device, - ), - requires_grad=False, - ) - for expert in range(layer.num_local_experts): - w13_weight[expert, :, :], layer.w13_weight_scale[expert] = ( - scaled_fp8_quant(layer.w13_weight.data[expert, :, :]) - ) - w2_weight[expert, :, :], layer.w2_weight_scale[expert] = ( - scaled_fp8_quant(layer.w2_weight.data[expert, :, :]) - ) - layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + # Online quantization already happened during weight loading via get_online_weight_loader. + assert layer.w13_weight.dtype == fp8_dtype + assert layer.w2_weight.dtype == fp8_dtype if _is_hip: self.process_weights_hip_scale_padding(layer) diff --git a/python/sglang/srt/layers/quantization/online_quantization.py b/python/sglang/srt/layers/quantization/online_quantization.py new file mode 100644 index 000000000000..cafde9fb16de --- /dev/null +++ b/python/sglang/srt/layers/quantization/online_quantization.py @@ -0,0 +1,22 @@ +import torch +from torch.utils._python_dispatch import TorchDispatchMode + + +class CopyNumelCounter(TorchDispatchMode): + """ + Tracks total number of elements modified with `copy_`. Useful for keeping + track of weight loading where underlying weights can be arbitrarily + transformed (such as with `narrow`) before calling copy. + """ + + def __init__(self): + super().__init__() + self.copied_numel = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + out = func(*args, **kwargs) + if func == torch.ops.aten.copy_.default: + self.copied_numel += args[0].numel() + return out diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 36fc1a33a936..6f9c5928603a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -35,6 +35,7 @@ import numpy as np import torch +from sglang.srt.constants import GIB_BYTES from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, get_remote_instance_transfer_engine_info_per_rank, @@ -80,6 +81,7 @@ get_model_architecture, set_default_torch_dtype, ) +from sglang.srt.utils.common import is_cuda_alike # Constants for memory management DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = ( @@ -730,6 +732,14 @@ def load_model( def load_weights_and_postprocess(model, weights, target_device): model.load_weights(weights) + # Used in test_online_quantization.py to verify memory savings when using online quantization. + if is_cuda_alike(): + peak_memory = torch.cuda.max_memory_allocated() + logger.debug( + "Peak GPU memory after loading weights: %s GiB", + f"{peak_memory / GIB_BYTES:.3f}", + ) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is not None: diff --git a/test/registered/quant/test_online_quantization.py b/test/registered/quant/test_online_quantization.py new file mode 100644 index 000000000000..b7f413e29965 --- /dev/null +++ b/test/registered/quant/test_online_quantization.py @@ -0,0 +1,132 @@ +import io +import os +import re + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu") +register_amd_ci(est_time=106, suite="stage-b-test-small-1-gpu-amd") +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.srt.utils.common import is_cuda_alike +from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestOnlineQuantizationMemoryLoad(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.SGLANG_USE_AITER = os.environ.get("SGLANG_USE_AITER", None) + + # DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE has a shape not compatible with aiter. + os.environ["SGLANG_USE_AITER"] = "0" + + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = io.StringIO() + cls.stderr = io.StringIO() + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--quantization", + "fp8", + "--tensor-parallel-size", + "1", + "--log-level", + "debug", + "--disable-cuda-graph", # See https://github.com/sgl-project/sglang/issues/18002 + ], + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + + # Extract and display peak GPU memory from logs + combined_output = cls.stdout.getvalue() + cls.stderr.getvalue() + peak_memory = cls._extract_peak_memory(combined_output) + + if is_cuda_alike() and not peak_memory: + raise ValueError("Should have found peak memory") + + cls.peak_memory = float(peak_memory) + + @classmethod + def _extract_peak_memory(cls, log_output): + """Extract peak GPU memory value from log output.""" + # Search for the log message pattern + pattern = r"Peak GPU memory after loading weights:\s+([\d.]+)\s+GiB" + match = re.search(pattern, log_output) + if match: + return match.group(1) + return None + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + cls.stdout.close() + cls.stderr.close() + + if cls.SGLANG_USE_AITER: + os.environ["SGLANG_USE_AITER"] = cls.SGLANG_USE_AITER + + +class TestOnlineQuantizationMemoryLoadDense(TestOnlineQuantizationMemoryLoad): + model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN + + def test_peak_memory(self): + if not is_cuda_alike(): + self.skipTest("not is_cuda_alike") + + # Original BF16 model: 2.887 GiB + assert self.peak_memory < 2 + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=500, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + + # Qwen/Qwen3-8B hits >0.9 as well. + # Original model reference accuracy: ~0.608 + self.assertGreater(metrics["accuracy"], 0.58) + + +class TestOnlineQuantizationMemoryLoadMOE(TestOnlineQuantizationMemoryLoad): + model = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE + + def test_peak_memory(self): + if not is_cuda_alike(): + self.skipTest("not is_cuda_alike") + + # Original BF16 model: 26.695 GiB + assert self.peak_memory < 15.5 + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=500, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + + # Original model reference accuracy: ~0.626 + self.assertGreater(metrics["accuracy"], 0.58)