diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index e9dc3b6fa..70bfcdcba 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -41,6 +41,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_MM_BUFFER_SIZE_MB` | Size of preallocated GPU buffer (in MB) for multi-modal feature hashing optimization. When set to a positive value, temporarily moves features to GPU for faster hash computation, then moves them back to CPU to save GPU memory. Larger features benefit more from GPU hashing. Set to `0` to disable. | `0` | | `SGLANG_MM_PRECOMPUTE_HASH` | Enable precomputing of hash values for MultimodalDataItem | `false` | | `SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH` | Enable NCCL for gathering when preparing mlp sync batch under overlap scheduler (without this flag gloo is used for gathering) | `false` | +| `SGLANG_SYMM_MEM_PREALLOC_GB_SIZE` | Size of preallocated GPU buffer (in GB) for NCCL symmetric memory pool to limit memory fragmentation. Only have an effect when server arg `--enable-symm-mem` is set. | `4` | ## DeepGEMM Configuration (Advanced Optimization) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4bc71df40..5dbd90899 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -453,6 +453,9 @@ class Envs: # TokenizerManager SGLANG_REQUEST_STATE_WAIT_TIMEOUT = EnvInt(4) + # Symmetric Memory + SGLANG_SYMM_MEM_PREALLOC_GB_SIZE = EnvInt(-1) + # Aiter SGLANG_USE_AITER_FP8_PER_TOKEN = EnvBool(False) # fmt: on diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 6ed1decd0..0c68283f4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -551,6 +551,7 @@ class Scheduler( self.max_req_input_len, self.random_seed, self.device, + self.forward_stream, _, _, _, @@ -971,7 +972,6 @@ class Scheduler( if self.device == "cpu": self.default_stream.synchronize = lambda: None # No-op for CPU - self.forward_stream: CudaStream = self.device_module.Stream() self.forward_stream_ctx: CudaStreamContext = self.device_module.stream( self.forward_stream ) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index af4000729..7e19c9d0d 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -400,6 +400,7 @@ class TpModelWorker(BaseTpWorker): self.max_req_input_len, self.random_seed, self.device, + self.model_runner.forward_stream, self.model_runner.req_to_token_pool.size, self.model_runner.req_to_token_pool.max_context_len, self.model_runner.token_to_kv_pool.size, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 46fe65ff9..c63ef3016 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -58,6 +58,9 @@ from sglang.srt.distributed import ( set_mscclpp_all_reduce, set_torch_symm_mem_all_reduce, ) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager from sglang.srt.environ import envs @@ -365,6 +368,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Get memory before model loading min_per_gpu_memory = self.init_torch_distributed() + # Init forward stream for overlap schedule + self.forward_stream = torch.get_device_module(self.device).Stream() + # CPU offload set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank)) @@ -593,6 +599,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Initialize piecewise CUDA graph self.init_piecewise_cuda_graphs() + self.prealloc_symmetric_memory_pool() + def init_routed_experts_capturer(self): if not self.server_args.disable_shared_experts_fusion and hasattr( self.model, "num_fused_shared_experts" @@ -2483,6 +2491,27 @@ class ModelRunner(ModelRunnerKVCacheMixin): logger.error(f"IPC weight update failed: {e}") return False, str(e) + def prealloc_symmetric_memory_pool(self): + # PyTorch mempools never de-fragment memory in OOM scenarios, so we need to pre-allocate a large chunk of memory to limit fragmentation. + if ( + self.is_draft_worker + or not self.server_args.enable_symm_mem + or envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get() <= 0 + ): + return + + # Memory allocation is tied to a cuda stream, use the forward stream + with torch.get_device_module(self.device).stream(self.forward_stream): + logger.info( + f"Pre-allocating symmetric memory pool with {envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get()} GiB" + ) + with use_symmetric_memory(get_tp_group()): + torch.empty( + (envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get() * 1024 * 1024 * 1024,), + dtype=torch.uint8, + device=self.device, + ) + def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]): params_dict = dict(model.named_parameters()) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index aea0c65bf..1679b5691 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1036,6 +1036,17 @@ class ServerArgs: if model_config.is_multimodal: self.adjust_mem_fraction_for_vlm(model_config) + # If symm mem is enabled and prealloc size is not set, set it to 4GB + if ( + self.enable_symm_mem + and not envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.is_set() + ): + envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.set(4) + logger.warning( + "Symmetric memory is enabled, setting symmetric memory prealloc size to 4GB as default." + "Use environment variable SGLANG_SYMM_MEM_PREALLOC_GB_SIZE to change the prealloc size." + ) + def _generate_cuda_graph_batch_sizes(self): """ Generate the list of batch sizes for CUDA graph capture based on cuda_graph_max_bs. diff --git a/test/srt/test_deepseek_v3_fp4_4gpu.py b/test/srt/test_deepseek_v3_fp4_4gpu.py index faf291fa5..5a9d86419 100644 --- a/test/srt/test_deepseek_v3_fp4_4gpu.py +++ b/test/srt/test_deepseek_v3_fp4_4gpu.py @@ -207,5 +207,59 @@ class TestDeepseekV3FP4CutlassMoE(CustomTestCase): self.assertGreater(metrics["accuracy"], 0.935) +class TestDeepseekV3FP4SymmetricMemory(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--tp", + "4", + "--attention-backend", + "trtllm_mla", + "--moe-runner-backend", + "flashinfer_trtllm", + "--quantization", + "modelopt_fp4", + "--kv-cache-dtype", + "fp8_e4m3", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', + "--enable-symm-mem", + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k( + self, + ): # Append an "a" to make this test run first (alphabetically) to warm up the server + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=1319, + parallel=1319, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + if is_in_ci(): + write_github_step_summary( + f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["accuracy"]=:.3f}\n' + ) + + self.assertGreater(metrics["accuracy"], 0.935) + + if __name__ == "__main__": unittest.main()