Inference | Lazy compile UVM allocator.

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,

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

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 <cuda_runtime_api.h>
-#include <cstddef>
-
-#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 <cuda_runtime_api.h>
+    #include <cstddef>
+
+    #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)

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.