megatron-core 0.16.0rc0.dev101638__tar.gz → 0.16.0rc0.dev102658__tar.gz

{megatron_core-0.16.0rc0.dev101638 → megatron_core-0.16.0rc0.dev102658}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101638
+Version: 0.16.0rc0.dev102658
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_core-0.16.0rc0.dev101638 → megatron_core-0.16.0rc0.dev102658}/megatron/core/inference/contexts/dynamic_context.py

@@ -16,7 +16,10 @@ from megatron.core.inference.inference_request import DynamicInferenceRequest
 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 @@ class DynamicInferenceContext(BaseInferenceContext):
         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-0.16.0rc0.dev102658/megatron/core/inference/unified_memory.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+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
+
+from megatron.core.utils import is_torch_min_version
+
+try:
+    if is_torch_min_version("2.8.0"):
+        from torch.cuda.memory import MemPool
+    else:
+        from torch.cuda import MemPool
+    _has_mem_pool = True
+except ImportError:
+    _has_mem_pool = False
+
+
+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)

{megatron_core-0.16.0rc0.dev101638 → megatron_core-0.16.0rc0.dev102658}/megatron/core/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 16
 PATCH = 0
-PRE_RELEASE = 'rc0.dev101638'
+PRE_RELEASE = 'rc0.dev102658'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)

{megatron_core-0.16.0rc0.dev101638 → megatron_core-0.16.0rc0.dev102658}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101638
+Version: 0.16.0rc0.dev102658
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

megatron_core-0.16.0rc0.dev101638/megatron/core/inference/unified_memory.py

@@ -1,89 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-import os
-import warnings
-from pathlib import Path
-
-from torch.cuda.memory import CUDAPluggableAllocator
-from torch.utils.cpp_extension import CUDA_HOME, load_inline
-
-from megatron.core.utils import is_torch_min_version
-
-try:
-    if is_torch_min_version("2.8.0"):
-        from torch.cuda.memory import MemPool
-    else:
-        from torch.cuda import MemPool
-    _has_mem_pool = True
-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)