megatron-core 0.16.0rc0.dev129362__tar.gz → 0.16.0rc0.dev129924__tar.gz

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev129362
+Version: 0.16.0rc0.dev129924
 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.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/inference/contexts/attention_context/mha_metadata.py

@@ -43,6 +43,7 @@ class MHAMetadata(MetadataBase):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
+        decode_only: bool = False,
     ):
         """
         Args:
@@ -52,6 +53,7 @@ class MHAMetadata(MetadataBase):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
+            decode_only: bool
         """
         if padded_active_request_count is None:
             padded_active_request_count = real_batch_size
@@ -98,9 +100,11 @@ class MHAMetadata(MetadataBase):
             padded_active_request_count,
             is_cumulative_tensor=True,
         )
-        self._max_seqlen_q = padded_active_token_count
-        if torch.all(self._query_lengths_buf[:padded_active_request_count] <= 1):
+
+        if decode_only:
             self._max_seqlen_q = 1
+        else:
+            self._max_seqlen_q = max(2, padded_active_token_count)
         self._max_seqlen_k = self.max_seqlen
 
         self.state_data = {
@@ -148,6 +152,7 @@ class GraphedMHAMetadata(MHAMetadata):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
+        decode_only: bool = False,
     ):
         """
         Args:
@@ -157,6 +162,7 @@ class GraphedMHAMetadata(MHAMetadata):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
+            decode_only: bool
         """
         super().update(
             request_query_lengths,
@@ -165,6 +171,7 @@ class GraphedMHAMetadata(MHAMetadata):
             padded_active_token_count,
             real_batch_size,
             padded_active_request_count,
+            decode_only,
         )
 
     def reset(self):
@@ -184,6 +191,7 @@ class NonGraphedMHAMetadata(MHAMetadata):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
+        decode_only: bool = False,
     ):
         """
         Args:
@@ -193,6 +201,7 @@ class NonGraphedMHAMetadata(MHAMetadata):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
+            decode_only: bool
         """
         super().update(
             request_query_lengths,
@@ -201,6 +210,7 @@ class NonGraphedMHAMetadata(MHAMetadata):
             padded_active_token_count,
             real_batch_size,
             padded_active_request_count,
+            decode_only,
         )
         if len(self.state_data["query_lengths"]) > 0:
             self.state_data["max_seqlen_q"] = torch.max(self.state_data["query_lengths"]).item()

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/inference/contexts/dynamic_context.py

@@ -4,7 +4,7 @@ import math
 import warnings
 from contextlib import nullcontext
 from enum import Enum
-from typing import List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -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
@@ -46,6 +49,17 @@ try:
 except ImportError:
     HAVE_FLASHINFER = False
 
+try:
+    import wandb  # pylint: disable=unused-import
+
+    HAVE_WANDB = True
+except ImportError:
+    HAVE_WANDB = False
+    wandb = None
+
+if TYPE_CHECKING:
+    import wandb as WandbModule
+
 
 class ContextOverflowError(Exception):
     """Base exception for when a new request does not fit.
@@ -223,6 +237,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             levels will be included to control other tensors within the context.
         use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation.
         If None, defaults to using flash-infer if available.
+        metrics_writer (Optional['WandbModule']): Wandb module for writing metrics.
     """
 
     def __init__(
@@ -248,6 +263,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         use_cuda_graphs_for_non_decode_steps: bool = True,
         use_flashinfer_fused_rope: bool = False,
         unified_memory_level: Optional[int] = 0,
+        metrics_writer: Optional['WandbModule'] = None,
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
@@ -257,6 +273,8 @@ class DynamicInferenceContext(BaseInferenceContext):
                 block_size_tokens == 64
             ), "Flash MLA requires a block size of 64. Set --inference-dynamic-batching-block-size 64 to fix this assert"
 
+        self.metrics_writer = metrics_writer
+
         # Per partition num heads and hidden size.
         projection_size = kv_channels * num_attention_heads
         if tensor_model_parallel_size is None:
@@ -323,16 +341,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
@@ -951,6 +973,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             padded_active_token_count=self.padded_active_token_count,
             real_batch_size=real_req_batch_size,
             padded_active_request_count=self.padded_active_request_count,
+            decode_only=self.is_decode_only(),
         )
         # All attention metadata calculations are now handled by MHAMetadata.update()
 
@@ -1561,3 +1584,67 @@ class DynamicInferenceContext(BaseInferenceContext):
 
         # Convert each log prob tensor into a list
         return [lp.tolist() for lp in selected_log_probs_list]
+
+    def get_kvcache_utilization_stats(self) -> dict:
+        """Compute KV cache buffer utilization stats for the current step.
+
+        Returns a dictionary with counts and percentages for both allocated block
+        usage (overall buffer occupancy) and active usage (blocks referenced by
+        currently active requests this step).
+
+        Return:
+            {
+            'total_blocks': int,
+            'allocated_blocks': int,
+            'active_unique_blocks': int,
+            'allocated_utilization': float,
+            'active_utilization': float,
+            'active_request_count': int,
+            'paused_request_count': int,
+            'gtd_block_count': int,
+            }
+        """
+        # Total usable blocks exclude the reserved dummy block.
+        total_blocks = max(self.block_allocator.block_count_total - 1, 1)
+        block_count_avail = int(self.block_allocator.block_count_avail)
+
+        # Overall allocated blocks in the buffer right now.
+        allocated_blocks = (self.block_allocator.block_count_total - 1) - block_count_avail
+        allocated_blocks = int(max(0, allocated_blocks))
+
+        # Active unique blocks referenced by current active requests only.
+        active_start = self.paused_request_count
+        active_end = self.total_request_count
+        if active_end > active_start:
+            active_rows = self.request_to_kv_block_ids[active_start:active_end]
+            # Filter valid block ids (>= 0) and count unique ids.
+            valid_ids = active_rows[active_rows >= 0]
+            if valid_ids.numel() > 0:
+                unique_ids = torch.unique(valid_ids)
+                active_unique_blocks = int(unique_ids.numel())
+            else:
+                active_unique_blocks = 0
+        else:
+            active_unique_blocks = 0
+
+        allocated_utilization = float(allocated_blocks) / float(total_blocks)
+        active_utilization = float(active_unique_blocks) / float(total_blocks)
+
+        # Diagnostic helpers
+        num_non_gtd_blocks = max(0, block_count_avail - int(self.gtd_block_count))
+        total_request_count = int(self.total_request_count)
+        return {
+            'total_blocks': int(total_blocks),
+            'allocated_blocks': int(allocated_blocks),
+            'active_unique_blocks': int(active_unique_blocks),
+            'allocated_utilization': allocated_utilization,
+            'active_utilization': active_utilization,
+            'active_request_count': int(self.get_active_request_count()),
+            'paused_request_count': int(self.paused_request_count),
+            'gtd_block_count': int(self.gtd_block_count),
+            'block_count_avail': int(block_count_avail),
+            'num_non_gtd_blocks': int(num_non_gtd_blocks),
+            'active_token_count': int(self.active_token_count),
+            'total_request_count': int(total_request_count),
+            'max_requests': int(self.max_requests),
+        }

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/inference/engines/dynamic_engine.py

@@ -57,6 +57,14 @@ try:
 except:
     HAVE_MSGPACK = False
 
+try:
+    import wandb
+
+    HAVE_WANDB = True
+except ImportError:
+    HAVE_WANDB = False
+    wandb = None
+
 
 def format_mem_bytes(mem_bytes):
     """Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes."""
@@ -89,6 +97,8 @@ class DynamicInferenceEngine(AbstractEngine):
         static_sampling (bool): If True, all requests are assumed to have the same
             sampling parameters. This avoids needing to loop through all requests and
             their sampling parameters every generation step, improving latency.
+        inference_logging_step_interval (int): The step interval at which to log
+        inference metrics to wandb. Defaults to 0, which means no logging.
     """
 
     def __init__(
@@ -101,6 +111,7 @@ class DynamicInferenceEngine(AbstractEngine):
         track_paused_request_events: bool = False,
         enable_chunked_prefill: bool = True,
         static_sampling: bool = False,
+        inference_logging_step_interval: int = 0,
     ):
 
         if enable_cuda_graph is not None:
@@ -137,6 +148,32 @@ class DynamicInferenceEngine(AbstractEngine):
         self.enable_chunked_prefill = enable_chunked_prefill
         self.static_sampling = static_sampling
 
+        self.inference_logging_step_interval = inference_logging_step_interval
+        # Configure wandb to use separate step counter for inference metrics (only once)
+        if self.inference_logging_step_interval > 0 and self.context.metrics_writer is not None:
+            logging.info(
+                f"\033[1;93m[INFERENCE]\033[0m "
+                f"\033[1;95mLogging inference metrics to wandb (rank {torch.distributed.get_rank()})\033[0m"
+            )
+            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
+                # Make all inference/* metrics use inference_step as their x-axis
+                # This allows inference and training to have independent step counters
+                context.metrics_writer.define_metric(
+                    "inference/*", step_metric="inference/inference_step"
+                )
+                # Initialize inference step offset by querying existing run history
+                self.inference_step_offset = 0
+                if wandb.run is not None:
+                    api_run = wandb.Api().run(
+                        f"{wandb.run.entity}/{wandb.run.project}/{wandb.run.id}"
+                    )
+                    max_step = 0
+                    for row in api_run.scan_history(keys=["inference/inference_step"]):
+                        val = row.get("inference/inference_step")
+                        if isinstance(val, (int, float)) and int(val) > max_step:
+                            max_step = int(val)
+                    self.inference_step_offset = int(max_step)
+
         # Initialize the asyncio loop if it has not already been initialized.
         # TODO: Start the engine loop here.
         self._loop = get_asyncio_loop()
@@ -780,6 +817,41 @@ class DynamicInferenceEngine(AbstractEngine):
             self.request_completion_futures[failed_request_id].set_result(failed_request)
         self.failed_request_ids.clear()
 
+        # Log KV cache utilization stats to W&B
+        if (
+            self.inference_logging_step_interval > 0
+            and self.step_count > 0
+            and self.step_count % self.inference_logging_step_interval == 0
+            and self.context.metrics_writer is not None
+        ):
+
+            # Get KV cache utilization stats from dynamic context
+            kv_stats = self.context.get_kvcache_utilization_stats()
+
+            # Prepare metrics dictionary with all stats
+            # Use 'inference/' prefix for all metrics to separate from training metrics
+            metrics = {
+                'inference/inference_step': int(self.inference_step_offset + int(self.step_count)),
+                'inference/step_time_s': float(step_time),
+                'inference/waiting_queue_len': int(len(self.waiting_request_ids)),
+                'inference/total_requests_dict_size': int(len(self.requests)),
+            }
+            # Add KV stats with inference/ prefix
+            # Convert utilization metrics from 0-1 range to 0-100 percentage range for better visualization
+            for key, value in kv_stats.items():
+                if 'utilization' in key:
+                    # Convert to percentage (0-100) and group under kvcache_utilization
+                    metrics[f'inference/{key}'] = float(value * 100.0)
+                else:
+                    metrics[f'inference/{key}'] = value
+
+            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
+                self.context.metrics_writer.log(metrics, commit=True)
+            else:
+                raise ValueError(
+                    f"Unsupported metrics writer type: {type(self.context.metrics_writer)}"
+                )
+
         # Print context state.
         if verbose:
             context = self.context

megatron_core-0.16.0rc0.dev129924/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.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/models/gpt/gpt_model.py

@@ -375,14 +375,13 @@ class GPTModel(LanguageModule):
                 )
                 or self.config.flash_decode
             )
-            and rotary_pos_cos is not None
             and inference_context.is_static_batching()
         ):
             current_batch_size = input_ids.shape[0]
             sequence_len_offset = torch.tensor(
                 [inference_context.sequence_len_offset] * current_batch_size,
                 dtype=torch.int32,
-                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
+                device=torch.cuda.current_device(),
             )
         else:
             sequence_len_offset = None

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/package_info.py

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

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/transformer/attention.py

@@ -48,15 +48,26 @@ except ImportError:
     rearrange = None
 
 try:
-    from flashattn_hopper.flash_attn_interface import _flash_attn_forward
-    from flashattn_hopper.flash_attn_interface import (
+    from flash_attn_3.flash_attn_interface import _flash_attn_forward
+    from flash_attn_3.flash_attn_interface import (
         flash_attn_with_kvcache as flash_attn3_with_kvcache,
     )
 
     HAVE_FA3 = True
-except:
+except ImportError as e:
     HAVE_FA3 = False
 
+if not HAVE_FA3:
+    try:
+        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
+        from flashattn_hopper.flash_attn_interface import (
+            flash_attn_with_kvcache as flash_attn3_with_kvcache,
+        )
+
+        HAVE_FA3 = True
+    except ImportError as e:
+        pass
+
 try:
     from flash_mla import flash_mla_with_kvcache, get_mla_metadata
 

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/transformer/cuda_graphs.py

@@ -1182,7 +1182,11 @@ class CudaGraphManager(torch.nn.Module):
 
             if runner is None:
                 if _CudagraphGlobalRecord.cudagraph_created:
-                    assert False
+                    assert False, (
+                        f"`cudagraph_created` is set to True but no matching cudagraph "
+                        f"runners were found. This module has {len(self.cudagraph_runners)} "
+                        f"existing runners. Use `get_mismatch_errors` to debug mismatches."
+                    )
                 else:
                     runner = _CudaGraphRunner(
                         megatron_module,

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/transformer/dot_product_attention.py

@@ -126,6 +126,8 @@ class DotProductAttention(MegatronModule):
                     )
                 ),
             )
+            if config.perform_initialization:
+                self.softmax_offset = config.init_method(self.softmax_offset)
         else:
             raise ValueError("Softmax type not supported")
 

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/transformer/moe/router.py

@@ -66,6 +66,8 @@ class Router(ABC, MegatronModule):
         """Reset the router parameters."""
         if self.config.perform_initialization:
             self.config.init_method(self.weight)
+            if self.bias is not None:
+                self.config.init_method(self.bias)
         self.weight.data = self.weight.data.to(dtype=self.config.params_dtype)
         setattr(self.weight, 'sequence_parallel', self.config.sequence_parallel)
         if self.bias is not None:

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron/core/transformer/pipeline_parallel_layer_layout.py

@@ -15,8 +15,11 @@ logger = logging.getLogger(__name__)
 class PipelineParallelLayerLayout:
     """Configuration of custom pipeline parallel layer partitioning."""
 
-    def __repr__(self):
-        return self.input_data
+    def __repr__(self) -> str:
+        if isinstance(self.input_data, str):
+            return self.input_data
+        else:
+            return str(self.input_data)
 
     def __init__(self, layout: str | list, pipeline_model_parallel_size: int):
         """Initialize PipelineParallelLayerLayout from a list or a str.

{megatron_core-0.16.0rc0.dev129362 → megatron_core-0.16.0rc0.dev129924}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev129362
+Version: 0.16.0rc0.dev129924
 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>