megatron-core 0.16.0rc0.dev101206__tar.gz → 0.16.0rc0.dev101273__tar.gz

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101206
+Version: 0.16.0rc0.dev101273
 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>
@@ -30,7 +30,7 @@ Classifier: Topic :: Utilities
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: torch
-Requires-Dist: numpy<2.0.0
+Requires-Dist: numpy
 Requires-Dist: packaging>=24.2
 Provides-Extra: mlm
 Requires-Dist: flask-restful; extra == "mlm"
@@ -40,7 +40,7 @@ Requires-Dist: wandb; extra == "mlm"
 Requires-Dist: transformers; extra == "mlm"
 Provides-Extra: dev
 Requires-Dist: nvidia-modelopt[torch]; sys_platform != "darwin" and extra == "dev"
-Requires-Dist: transformer-engine[pytorch]<2.10.0,>=2.9.0a0; extra == "dev"
+Requires-Dist: transformer-engine[core_cu13,pytorch]<2.10.0,>=2.9.0a0; extra == "dev"
 Requires-Dist: nvidia-resiliency-ext; extra == "dev"
 Requires-Dist: tqdm; extra == "dev"
 Requires-Dist: einops~=0.8; extra == "dev"
@@ -48,15 +48,15 @@ Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "dev"
 Requires-Dist: nvtx~=0.2; extra == "dev"
 Requires-Dist: multi-storage-client~=0.27; extra == "dev"
 Requires-Dist: opentelemetry-api~=1.33.1; extra == "dev"
-Requires-Dist: setuptools<80.0.0; extra == "dev"
 Requires-Dist: mamba-ssm~=2.2; extra == "dev"
 Requires-Dist: causal-conv1d~=1.5; extra == "dev"
 Requires-Dist: nv-grouped-gemm~=1.1; extra == "dev"
 Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "dev"
-Requires-Dist: av<16.0.0; extra == "dev"
+Requires-Dist: av; extra == "dev"
 Requires-Dist: flashinfer-python; extra == "dev"
 Requires-Dist: wget; extra == "dev"
 Requires-Dist: onnxscript; extra == "dev"
+Requires-Dist: fastapi~=0.50; extra == "dev"
 Provides-Extra: lts
 Requires-Dist: tqdm; extra == "lts"
 Requires-Dist: einops~=0.8; extra == "lts"
@@ -64,15 +64,15 @@ Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "lts"
 Requires-Dist: nvtx~=0.2; extra == "lts"
 Requires-Dist: multi-storage-client~=0.27; extra == "lts"
 Requires-Dist: opentelemetry-api~=1.33.1; extra == "lts"
-Requires-Dist: setuptools<80.0.0; extra == "lts"
 Requires-Dist: mamba-ssm~=2.2; extra == "lts"
 Requires-Dist: causal-conv1d~=1.5; extra == "lts"
 Requires-Dist: nv-grouped-gemm~=1.1; extra == "lts"
 Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "lts"
-Requires-Dist: av<16.0.0; extra == "lts"
+Requires-Dist: av; extra == "lts"
 Requires-Dist: flashinfer-python; extra == "lts"
 Requires-Dist: wget; extra == "lts"
 Requires-Dist: onnxscript; extra == "lts"
+Requires-Dist: fastapi~=0.50; extra == "lts"
 
 <div align="center">
 

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron/core/dist_checkpointing/exchange_utils.py

@@ -63,7 +63,7 @@ class ShardDistribution(NamedTuple):
 def _shard_size(sh_ten: ShardedTensor):
     """Returns size in bytes of a given sharded tensor."""
     if sh_ten.flattened_range is None:
-        numel = np.product(sh_ten.local_shape)
+        numel = np.prod(sh_ten.local_shape)
     else:
         numel = sh_ten.flattened_range.stop - sh_ten.flattened_range.start
     return numel * torch._utils._element_size(sh_ten.dtype)

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron/core/dist_checkpointing/mapping.py

@@ -216,7 +216,7 @@ class ShardedTensor(ShardedBase):
             )
 
         # TODO: np.unravel_index?
-        mask = np.zeros(np.product(self.local_shape), dtype=bool)
+        mask = np.zeros(np.prod(self.local_shape), dtype=bool)
         mask[self.flattened_range] = True
         return np.nonzero(mask.reshape(self.local_shape))
 

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron/core/dist_checkpointing/validation.py

@@ -519,7 +519,7 @@ def _validate_sharding_for_key_flattened(tensors_by_shard):
         all_slices.append((sharding.flattened_range.start, sharding.flattened_range.stop))
 
     starts, stops = map(np.asarray, zip(*sorted(all_slices)))
-    expected_size = np.product(local_shape)
+    expected_size = np.prod(local_shape)
     if starts[0] != 0 or stops[-1] != expected_size or not np.all(starts[1:] == stops[:-1]):
         raise CheckpointingException(
             f"Flattened ranges dont cover the whole shard {tensors_by_shard[0]} of size {expected_size}. Ranges: {(starts, stops)}"

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -115,5 +115,9 @@ class BlockAllocator:
         This resets the available block count to the entire memory pool
         (except for the dummy block).
         """
-
+        # Reset block bag to so we start consuming
+        # from the beginning of the pool for UVM performance
+        self.block_bag = torch.arange(
+            self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
+        )
         self.total_avail = self.total_count - 1

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron/core/package_info.py

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

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101206
+Version: 0.16.0rc0.dev101273
 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>
@@ -30,7 +30,7 @@ Classifier: Topic :: Utilities
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: torch
-Requires-Dist: numpy<2.0.0
+Requires-Dist: numpy
 Requires-Dist: packaging>=24.2
 Provides-Extra: mlm
 Requires-Dist: flask-restful; extra == "mlm"
@@ -40,7 +40,7 @@ Requires-Dist: wandb; extra == "mlm"
 Requires-Dist: transformers; extra == "mlm"
 Provides-Extra: dev
 Requires-Dist: nvidia-modelopt[torch]; sys_platform != "darwin" and extra == "dev"
-Requires-Dist: transformer-engine[pytorch]<2.10.0,>=2.9.0a0; extra == "dev"
+Requires-Dist: transformer-engine[core_cu13,pytorch]<2.10.0,>=2.9.0a0; extra == "dev"
 Requires-Dist: nvidia-resiliency-ext; extra == "dev"
 Requires-Dist: tqdm; extra == "dev"
 Requires-Dist: einops~=0.8; extra == "dev"
@@ -48,15 +48,15 @@ Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "dev"
 Requires-Dist: nvtx~=0.2; extra == "dev"
 Requires-Dist: multi-storage-client~=0.27; extra == "dev"
 Requires-Dist: opentelemetry-api~=1.33.1; extra == "dev"
-Requires-Dist: setuptools<80.0.0; extra == "dev"
 Requires-Dist: mamba-ssm~=2.2; extra == "dev"
 Requires-Dist: causal-conv1d~=1.5; extra == "dev"
 Requires-Dist: nv-grouped-gemm~=1.1; extra == "dev"
 Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "dev"
-Requires-Dist: av<16.0.0; extra == "dev"
+Requires-Dist: av; extra == "dev"
 Requires-Dist: flashinfer-python; extra == "dev"
 Requires-Dist: wget; extra == "dev"
 Requires-Dist: onnxscript; extra == "dev"
+Requires-Dist: fastapi~=0.50; extra == "dev"
 Provides-Extra: lts
 Requires-Dist: tqdm; extra == "lts"
 Requires-Dist: einops~=0.8; extra == "lts"
@@ -64,15 +64,15 @@ Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "lts"
 Requires-Dist: nvtx~=0.2; extra == "lts"
 Requires-Dist: multi-storage-client~=0.27; extra == "lts"
 Requires-Dist: opentelemetry-api~=1.33.1; extra == "lts"
-Requires-Dist: setuptools<80.0.0; extra == "lts"
 Requires-Dist: mamba-ssm~=2.2; extra == "lts"
 Requires-Dist: causal-conv1d~=1.5; extra == "lts"
 Requires-Dist: nv-grouped-gemm~=1.1; extra == "lts"
 Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "lts"
-Requires-Dist: av<16.0.0; extra == "lts"
+Requires-Dist: av; extra == "lts"
 Requires-Dist: flashinfer-python; extra == "lts"
 Requires-Dist: wget; extra == "lts"
 Requires-Dist: onnxscript; extra == "lts"
+Requires-Dist: fastapi~=0.50; extra == "lts"
 
 <div align="center">
 

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/megatron_core.egg-info/requires.txt

@@ -1,9 +1,9 @@
 torch
-numpy<2.0.0
+numpy
 packaging>=24.2
 
 [dev]
-transformer-engine[pytorch]<2.10.0,>=2.9.0a0
+transformer-engine[core_cu13,pytorch]<2.10.0,>=2.9.0a0
 nvidia-resiliency-ext
 tqdm
 einops~=0.8
@@ -11,15 +11,15 @@ tensorstore!=0.1.46,!=0.1.72,~=0.1
 nvtx~=0.2
 multi-storage-client~=0.27
 opentelemetry-api~=1.33.1
-setuptools<80.0.0
 mamba-ssm~=2.2
 causal-conv1d~=1.5
 nv-grouped-gemm~=1.1
 megatron-energon[av_decode]~=6.0
-av<16.0.0
+av
 flashinfer-python
 wget
 onnxscript
+fastapi~=0.50
 
 [dev:sys_platform != "darwin"]
 nvidia-modelopt[torch]
@@ -31,15 +31,15 @@ tensorstore!=0.1.46,!=0.1.72,~=0.1
 nvtx~=0.2
 multi-storage-client~=0.27
 opentelemetry-api~=1.33.1
-setuptools<80.0.0
 mamba-ssm~=2.2
 causal-conv1d~=1.5
 nv-grouped-gemm~=1.1
 megatron-energon[av_decode]~=6.0
-av<16.0.0
+av
 flashinfer-python
 wget
 onnxscript
+fastapi~=0.50
 
 [mlm]
 flask-restful

{megatron_core-0.16.0rc0.dev101206 → megatron_core-0.16.0rc0.dev101273}/pyproject.toml

@@ -20,7 +20,7 @@ dynamic = ["version", "readme"]
 description = "Megatron Core - a library for efficient and scalable training of transformer based models"
 requires-python = ">=3.10"
 license = { text = "Apache 2.0" }
-dependencies = ["torch", "numpy<2.0.0", "packaging>=24.2"]
+dependencies = ["torch", "numpy", "packaging>=24.2"]
 authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
 maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
 keywords = [
@@ -68,7 +68,7 @@ mlm = ["flask-restful", "sentencepiece", "tiktoken", "wandb", "transformers"]
 
 dev = [
     "nvidia-modelopt[torch]; sys_platform != 'darwin'",
-    "transformer-engine[pytorch]>=2.9.0a0,<2.10.0",
+    "transformer-engine[pytorch,core_cu13]>=2.9.0a0,<2.10.0",
     "nvidia-resiliency-ext",
     "tqdm",
     "einops~=0.8",
@@ -76,15 +76,15 @@ dev = [
     "nvtx~=0.2",
     "multi-storage-client~=0.27",
     "opentelemetry-api~=1.33.1",
-    "setuptools<80.0.0",
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
     "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
-    "av<16.0.0",                                        # At the time, av 16.0.0 is not compatible with Python 3.12
+    "av",
     "flashinfer-python",
     "wget",
     "onnxscript",
+    "fastapi~=0.50",                                          # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0
 ]
 
 lts = [
@@ -94,15 +94,15 @@ lts = [
     "nvtx~=0.2",
     "multi-storage-client~=0.27",
     "opentelemetry-api~=1.33.1",
-    "setuptools<80.0.0",
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
     "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
-    "av<16.0.0",                          # At the time, av 16.0.0 is not compatible with Python 3.12
+    "av",
     "flashinfer-python",
     "wget",
     "onnxscript",
+    "fastapi~=0.50",                      # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0
 ]
 
 [dependency-groups]
@@ -146,7 +146,7 @@ linting = [
     "pylint==3.2.6",
 ]
 ci = ["python-gitlab", "slack-sdk", "pandas"]
-flash_mla = ["flash_mla"]
+no_pypi_wheels = ["flash_mla", "emerging_optimizers"]
 
 [tool.uv]
 default-groups = ["linting", "build", "test"]
@@ -173,7 +173,7 @@ override-dependencies = [
 flash_mla = [
     { git = "https://github.com/deepseek-ai/FlashMLA", rev = "9edee0c022cd0938148a18e334203b0aab43aa19" },
 ]
-transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "release_v2.9" } # on `release_v2.9`
+# transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "release_v2.9" } # on `release_v2.9`
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "01a9a8ba360f7b2908728ad0516e0ad9d936966d" }
 emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "fb1add873e7851ec34b48581ea1b15761b73d189" }