Address PR review feedback: type annotation, checkpoint warning, comment

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

modelopt/torch/opt/searcher.py

@@ -254,6 +254,10 @@ def load_search_checkpoint(self, strict=True) -> bool:
             return False
         # Backward compat: fall back to the original single-file path
         if not os.path.exists(checkpoint):
+            warn_rank_0(
+                f"Per-rank checkpoint {checkpoint} not found, falling back to "
+                f"{self.config['checkpoint']}. Ensure world size matches the original run."
+            )
             checkpoint = self.config["checkpoint"]
         if not os.path.exists(checkpoint):
             warn_rank_0(f"Checkpoint {checkpoint} does not exist! Initializing from scratch.")

modelopt/torch/quantization/algorithms.py

@@ -337,7 +337,7 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     candidate_stats: dict[str, dict[str, list[float]]]
     best: dict[str, Any]
     quantizer_states: dict
-    method_name: str = None
+    method_name: str | None = None
 
     quant_grouping_rules = [
         r"^(.*?)\.(q_proj|k_proj|v_proj)$",  # q_proj, k_proj, v_proj for llama like models
@@ -1356,6 +1356,7 @@ def _resolve_best_recipe(search_state, constraints, verbose=False):
 
 
 def _match_quantizer_cfg(quant_cfg, quantizer_attr):
+    # Last-match-wins to mirror set_quantizer_by_cfg behavior
     matched = None
     for pattern, cfg in quant_cfg.items():
         if fnmatch.fnmatch(quantizer_attr, pattern):