0.23.1 Release - fix for torch 2.6

Author: kevalmorabia97

.vscode/settings.json

@@ -3,7 +3,7 @@
     "editor.rulers": [
         100,
         120
-    ], // 100 for black auto-formatter, 120 for hard limit in ruff
+    ], // 100 for ruff auto-formatter, 120 for hard limit in ruff
     "[python]": {
         "editor.defaultFormatter": "charliermarsh.ruff",
         "editor.formatOnSave": true,
@@ -40,4 +40,7 @@
         "--no-cov",
     ],
     "evenBetterToml.schema.enabled": false, // disable toml/json schema since we have custom fields
+    "python.analysis.extraPaths": [
+        "./tests/" // add tests to python path just like pytest does in pyproject.toml
+    ]
 }

examples/llm_distill/README.md

@@ -6,7 +6,7 @@ In this tutorial, we demonstrate how to use Model Optimizer to perform teacher-s
 
 Checkout the stand-along distillation example in the [NVIDIA NeMo repository](https://github.com/NVIDIA/NeMo/tree/main/examples/nlp/language_modeling/megatron_gpt_distillation.py).
 
-You can also look at the tutorial notebooks [here](https://github.com/NVIDIA/NeMo/tree/main/tutorials/llm/llama-3/pruning-distillation) which showcase the usage of Minitron pruning followed by distillation for Llama 3.1 8B step-by-step in NeMo framework.
+You can also look at the tutorial notebooks [here](https://github.com/NVIDIA/NeMo/tree/main/tutorials/llm/llama/pruning-distillation) which showcase the usage of Minitron pruning followed by distillation for Llama 3.1 8B step-by-step in NeMo framework.
 
 ## Knowledge Distillation (KD) for HuggingFace Models
 

examples/llm_distill/main.py

@@ -197,7 +197,7 @@ def train():
             if not os.path.isfile(modelopt_state_path):
                 raise FileNotFoundError("`modelopt_state.pt` not found with checkpoint.")
             logger.info(f"Loading modelopt state from {modelopt_state_path}")
-            modelopt_state = torch.load(modelopt_state_path)
+            modelopt_state = torch.load(modelopt_state_path, weights_only=False)
             mto.restore_from_modelopt_state(model, modelopt_state)
 
         logger.info("Beginning training...")

examples/llm_qat/README.md

@@ -61,7 +61,7 @@ torch.save(mto.modelopt_state(model), "modelopt_quantizer_states.pt")
 
 # To resume training from a checkpoint or load the final QAT model for evaluation,
 # load the quantizer states before loading the model weights
-# mto.restore_from_modelopt_state(model, torch.load("modelopt_quantizer_states.pt"))
+# mto.restore_from_modelopt_state(model, torch.load("modelopt_quantizer_states.pt", weights_only=False))
 # After loading the quantizer states, load the model weights
 # model.load_state_dict(state_dict_from_last_checkpoint)
 

examples/llm_sparsity/requirements.txt

@@ -1,4 +1,3 @@
 flash-attn
 sentencepiece>=0.2.0
 tensorboardX
-torch>=2.1.0

examples/pruning/README.md

@@ -10,7 +10,7 @@ Model Optimizer can be used in one of the following complementary pruning modes
 
 ## Documentation
 
-Checkout the [Quick Start: Pruning](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/4_pruning.html) and the detailed [Optimization Guide](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/2_pruning.html) in the Model Optimizer documentation for more information on how to use the above pruning algorithms in Model Optimizer.
+Checkout the [Quick Start: Pruning](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/5_pruning.html) and the detailed [Optimization Guide](https://nvidia.github.io/TensorRT-Model-Optimizer/guides/2_pruning.html) in the Model Optimizer documentation for more information on how to use the above pruning algorithms in Model Optimizer.
 
 ## Algorithms
 

examples/speculative_decoding/requirements.txt

@@ -3,4 +3,3 @@ openai
 py7zr
 sentencepiece>=0.2.0
 tensorboardX
-torch>=2.1.0

modelopt/torch/export/distribute.py

@@ -91,7 +91,7 @@ def read_configs_and_weights_from_rank(
             raise ValueError("NFSWorkspace is not initialized!")
         state_path = self._get_state_path(target_rank)
         if state_path.exists():
-            state = torch.load(state_path, map_location="cpu")
+            state = torch.load(state_path, map_location="cpu", weights_only=False)
             return state["config"], state["weight"]
         else:
             return None, None

modelopt/torch/export/layer_utils.py

@@ -822,8 +822,8 @@ def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer):
                 # for Int8 SQ case, we split the weight scaling factor into two parts.
                 weight_scaling_factors = torch.chunk(weight_scaling_factor, 2, dim=0)
 
-            config.fc = build_linear_config(fc_linear)
-            config.gate = build_linear_config(fc_linear)
+            config.fc = build_linear_config(fc_linear, LINEAR_COLUMN)
+            config.gate = build_linear_config(fc_linear, LINEAR_COLUMN)
             config.fc.weight = weights[0]
             config.gate.weight = weights[1]
             if weight_scaling_factors is not None:

modelopt/torch/opt/conversion.py

@@ -488,7 +488,7 @@ def restore_from_modelopt_state(model: ModelLike, modelopt_state: dict[str, Any]
 
         # Restore the previously saved modelopt state followed by model weights
         mto.restore_from_modelopt_state(
-            model, torch.load("modelopt_state.pt")
+            model, torch.load("modelopt_state.pt", weights_only=False)
         )  # Restore modelopt state
         model.load_state_dict(torch.load("model_weights.pt"), ...)  # Load the model weights
 
@@ -561,6 +561,7 @@ def restore(model: ModelLike, f: Union[str, os.PathLike, BinaryIO], **kwargs) ->
 
     # load checkpoint
     kwargs.setdefault("map_location", "cpu")
+    kwargs.setdefault("weights_only", False)
     objs = torch.load(f, **kwargs)
 
     # restore model architecture