Inference | Lazy compile UVM allocator.

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
@@ -276,16 +279,30 @@ 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:
+            # >>>
+            # 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:
+            #     self.unified_memory_mempool = create_unified_mempool()
+            # +++
+            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,8 +2,8 @@
 
 import os
 import warnings
+from enum import Enum, auto
 from pathlib import Path
-
 from torch.cuda.memory import CUDAPluggableAllocator
 from torch.utils.cpp_extension import CUDA_HOME, load_inline
 
@@ -18,72 +18,128 @@
 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
+# >>>
+# _has_unified_memory = False
+# _compilation_attempted = False
+# +++
+class CompilationState(Enum):
+    UNATTEMPTED = auto()
+    FAILURE = auto()
+    SUCCESS = auto()
+
+
+_compilation_state = CompilationState.UNATTEMPTED
+# <<<
+
+
+class UnifiedMemoryUnsupportedError(Exception):
+    pass
+
+
+def compile_allocator():
+    """Attempt to compile UVM allocator."""
+
+    # >>>
+    # global _compilation_attempted, _has_unified_memory
+
+    # if _compilation_attempted:
+    #     return
+    # _compilation_attempted = True
+    # +++
+    global _compilation_state
+
+    if _compilation_state != CompilationState.UNATTEMPTED:
+        return
+    # <<<
+
+    _mempool_c_src = r"""
+    #include <cuda_runtime_api.h>
+    #include <cstddef>
+
+    #define EXPORT extern "C"
 
-# 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)
+    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()
+            # >>>
+            # _has_unified_memory = True
+            _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.
+    # global _has_unified_memory
+    # >>>
+    from lutil import pax
+    # pax("_has_unified_memory")
+    pax("_compilation_state")
+    # <<<
+    if not _has_unified_memory:
+        raise UnifiedMemoryUnsupportedError()
+    else:
+        return MemPool(allocator=_alloc)