diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index e7745537991..baa25787e83 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -355,6 +355,9 @@ def build_dynamic_engine_setup_prefix( else: cg_str = "--" + # Unified memory (UVM). + uvm_str = f"uvm {int(context.unified_memory_level)}" + # Prompt description prompt_src_str = ( "cli" if args.prompts else @@ -390,6 +393,7 @@ def build_dynamic_engine_setup_prefix( get_model_size_str(model), "dynamic", cg_str, + uvm_str, request_str, buffer_limits_str, guaranteed_fraction_str, diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index bdf4e9fe43a..8185eb0bcd4 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -16,7 +16,10 @@ from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( InferenceWrapperConfig, ) -from megatron.core.inference.unified_memory import create_unified_mempool, has_unified_memory +from megatron.core.inference.unified_memory import ( + UnifiedMemoryUnsupportedError, + create_unified_mempool, +) from megatron.core.inference.utils import tensor_swap from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb from megatron.core.package_info import __version__ as mcore_version @@ -322,16 +325,20 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length + + # Unified memory. self.unified_memory_level = unified_memory_level if unified_memory_level > 0: - if not has_unified_memory and torch.distributed.get_rank() == 0: - warnings.warn( - "Unified memory requested but not available; defaulting to GPU memory." - ) - self.unified_memory_level = 0 - else: + try: self.unified_memory_mempool = create_unified_mempool() + except UnifiedMemoryUnsupportedError: + if torch.distributed.get_rank() == 0: + warnings.warn( + "Unified memory requested but not available; defaulting to GPU memory." + ) + self.unified_memory_level = 0 + # Request and token counts. self.total_request_count = 0 self.active_token_count = 0 self.paused_request_count = 0 diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 016a4f80d39..25634e8e3c3 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -2,6 +2,7 @@ import os import warnings +from enum import Enum, auto from pathlib import Path from torch.cuda.memory import CUDAPluggableAllocator @@ -18,72 +19,109 @@ except ImportError: _has_mem_pool = False -_mempool_c_src = r""" -#include -#include - -#define EXPORT extern "C" - -EXPORT void* managed_malloc(size_t size, int device, void* stream) { - (void)stream; - int cur = -1; - cudaGetDevice(&cur); - if (device != cur && device >= 0) cudaSetDevice(device); - - // cudaMallocManaged allows for more memory to be allocated than the device memory size. - // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. - void* ptr = nullptr; - cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); - if (err != cudaSuccess) return nullptr; - - if (device >= 0) { - // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. - // This is a hint that tries to prevent data from being migrated away from the device. - cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); - // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. - // Even if the memory has to be migrated away from the device, it still does not page fault. - // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, - // but there is no harm in adding this flag as well for future-proofing. - cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); - } - return ptr; -} - -EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { - // Memory allocated with cudaMallocManaged should be released with cudaFree. - (void)size; (void)device; (void)stream; - if (ptr) cudaFree(ptr); -} -""" - -# Avoid linting errors. -has_unified_memory = False -_alloc = None - -# Build the .so upon import; this avoids issues. -if _has_mem_pool: - _extra_ldflags = ["-lcudart"] - if CUDA_HOME: - _cuda_lib = os.path.join(CUDA_HOME, "lib64") - if os.path.isdir(_cuda_lib): - _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"] - try: - _mod = load_inline( - name="managed_alloc_runtime", - cpp_sources=[_mempool_c_src], - functions=[], - with_cuda=True, - extra_ldflags=_extra_ldflags, - verbose=False, - ) - _so_path = Path(_mod.__file__).as_posix() - _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() - has_unified_memory = True - except (RuntimeError, ImportError, OSError): - warnings.warn("Failed to create unified memory mempool.") - - -def create_unified_mempool(): - """Create a unified memory mempool using CUDA managed memory.""" - assert has_unified_memory - return MemPool(allocator=_alloc) + +class CompilationState(Enum): + """Enum to distinguish between unified memory (UVM) compilation states.""" + + UNATTEMPTED = auto() # Compilation has not been attempted. + FAILURE = auto() # Compilation attempted, but failed. + SUCCESS = auto() # Compilation attempted, and succeeded. + + +# Compilation vars. +_compilation_state = CompilationState.UNATTEMPTED +_alloc = None # must remain global until process exit. +_mod = None # must remain global until process exit. + + +class UnifiedMemoryUnsupportedError(Exception): + """Unified memory is not supported on this system.""" + + pass + + +def compile_allocator(): + """Attempt to compile UVM allocator.""" + + global _compilation_state, _alloc, _mod + + if _compilation_state != CompilationState.UNATTEMPTED: + return + + _mempool_c_src = r""" + #include + #include + + #define EXPORT extern "C" + + EXPORT void* managed_malloc(size_t size, int device, void* stream) { + (void)stream; + int cur = -1; + cudaGetDevice(&cur); + if (device != cur && device >= 0) cudaSetDevice(device); + + // cudaMallocManaged allows for more memory to be allocated than the device memory size. + // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. + void* ptr = nullptr; + cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); + if (err != cudaSuccess) return nullptr; + + if (device >= 0) { + // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. + // This is a hint that tries to prevent data from being migrated away from the device. + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); + // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. + // Even if the memory has to be migrated away from the device, it still does not page fault. + // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, + // but there is no harm in adding this flag as well for future-proofing. + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); + } + return ptr; + } + + EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { + // Memory allocated with cudaMallocManaged should be released with cudaFree. + (void)size; (void)device; (void)stream; + if (ptr) cudaFree(ptr); + } + """ + + # Build the .so upon import; this avoids issues. + if _has_mem_pool: + _extra_ldflags = ["-lcudart"] + if CUDA_HOME: + _cuda_lib = os.path.join(CUDA_HOME, "lib64") + if os.path.isdir(_cuda_lib): + _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"] + try: + _mod = load_inline( + name="managed_alloc_runtime", + cpp_sources=[_mempool_c_src], + functions=[], + with_cuda=True, + extra_ldflags=_extra_ldflags, + verbose=False, + ) + _so_path = Path(_mod.__file__).as_posix() + _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() + _compilation_state = CompilationState.SUCCESS + except (RuntimeError, ImportError, OSError): + warnings.warn("Failed to create unified memory mempool.") + _compilation_state = CompilationState.FAILURE + + +def create_unified_mempool() -> MemPool: + """Create a unified memory mempool using CUDA managed memory. + + Returns: + (MemPool) Unified memory mempool. + """ + + # Attempt to compile allocator. + compile_allocator() + + # Return mempool. + if _compilation_state != CompilationState.SUCCESS: + raise UnifiedMemoryUnsupportedError() + else: + return MemPool(allocator=_alloc) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 70c2bf6ed10..1cd9d66ece1 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -875,11 +875,19 @@ def test_calculate_and_store_log_probs(self): @pytest.mark.internal def test_unified_memory(self): - from megatron.core.inference.unified_memory import has_unified_memory - if not has_unified_memory: + from megatron.core.inference.unified_memory import ( + UnifiedMemoryUnsupportedError, + create_unified_mempool, + ) + + # Check UVM support. + try: + create_unified_mempool() + except UnifiedMemoryUnsupportedError: pytest.skip("Unified memory not available due to bad environment.") + # Setup. self._setup_model_parallel_group(1, 1) # Compute number of contexts needed to fill GPU memory.