cleaner code

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Author: shengliangxu

modelopt/recipe/loader.py

@@ -88,7 +88,10 @@ def _load_recipe_from_file(recipe_file: Path | Traversable) -> ModelOptRecipeBas
     plus a ``quant_cfg`` mapping and an optional ``algorithm`` for PTQ recipes.
     """
     raw = _load_raw_config(recipe_file)
-    assert isinstance(raw, dict), f"Recipe file {recipe_file} must be a YAML mapping."
+    if not isinstance(raw, dict):
+        raise ValueError(
+            f"Recipe file {recipe_file} must be a YAML mapping, got {type(raw).__name__}."
+        )
     data = _resolve_imports(raw)
 
     metadata = data.get("metadata", {})
@@ -121,7 +124,10 @@ def _load_recipe_from_dir(recipe_dir: Path | Traversable) -> ModelOptRecipeBase:
         )
 
     recipe_data = _load_raw_config(recipe_file)
-    assert isinstance(recipe_data, dict), f"Recipe file {recipe_file} must be a YAML mapping."
+    if not isinstance(recipe_data, dict):
+        raise ValueError(
+            f"Recipe file {recipe_file} must be a YAML mapping, got {type(recipe_data).__name__}."
+        )
     metadata = recipe_data.get("metadata", {})
     recipe_type = metadata.get("recipe_type")
     if recipe_type is None:
@@ -138,14 +144,21 @@ def _load_recipe_from_dir(recipe_dir: Path | Traversable) -> ModelOptRecipeBase:
             raise ValueError(
                 f"Cannot find quantize in {recipe_dir}. Looked for: quantize.yml, quantize.yaml"
             )
-        # Resolve imports: imports are in recipe.yml, quantize data is separate
+        # Resolve imports from both recipe.yaml and quantize.yaml
         quantize_data = _load_raw_config(quantize_file)
-        assert isinstance(quantize_data, dict), f"{quantize_file} must be a YAML mapping."
+        if not isinstance(quantize_data, dict):
+            raise ValueError(
+                f"{quantize_file} must be a YAML mapping, got {type(quantize_data).__name__}."
+            )
+        # Resolve quantize.yaml's own imports first (if any)
+        if "imports" in quantize_data:
+            quantize_data = _resolve_imports(quantize_data)
+        # Then resolve recipe.yaml's imports applied to the quantize data
         combined: dict[str, Any] = {"quantize": quantize_data}
         imports = recipe_data.get("imports")
         if imports:
             combined["imports"] = imports
-        combined = _resolve_imports(combined)
+            combined = _resolve_imports(combined)
         return ModelOptPTQRecipe(
             recipe_type=RecipeType.PTQ,
             description=metadata.get("description", "PTQ recipe."),

modelopt/torch/opt/config_loader.py

@@ -198,6 +198,28 @@ def _lookup(ref_name: str, context: str) -> Any:
             )
         return import_map[ref_name]
 
+    def _resolve_dict_value(d: dict[str, Any], key: str) -> None:
+        """Resolve ``$import`` in a dict value: ``key: {$import: name, ...inline}``."""
+        val = d[key]
+        if not isinstance(val, dict) or _IMPORT_KEY not in val:
+            return
+        ref = val.pop(_IMPORT_KEY)
+        inline_keys = dict(val)
+        ref_names = ref if isinstance(ref, list) else [ref]
+
+        merged: dict[str, Any] = {}
+        for rname in ref_names:
+            snippet = _lookup(rname, f"{key} of {d}")
+            if not isinstance(snippet, dict):
+                raise ValueError(
+                    f"$import {rname!r} in {key} must resolve to a dict, "
+                    f"got {type(snippet).__name__}."
+                )
+            merged.update(snippet)
+
+        merged.update(inline_keys)
+        d[key] = merged
+
     def _resolve_list(entries: list[Any]) -> list[Any]:
         """Resolve $import markers in a list of entries."""
         resolved: list[Any] = []
@@ -215,27 +237,11 @@ def _resolve_list(entries: list[Any]) -> list[Any]:
                         f"list, got {type(imported).__name__}."
                     )
                 resolved.extend(imported)
-            elif (
-                isinstance(entry, dict)
-                and isinstance(entry.get("cfg"), dict)
-                and _IMPORT_KEY in entry["cfg"]
-            ):
-                ref = entry["cfg"].pop(_IMPORT_KEY)
-                inline_keys = dict(entry["cfg"])
-                ref_names = ref if isinstance(ref, list) else [ref]
-
-                merged: dict[str, Any] = {}
-                for rname in ref_names:
-                    snippet = _lookup(rname, f"cfg of {entry}")
-                    if not isinstance(snippet, dict):
-                        raise ValueError(
-                            f"$import {rname!r} in cfg must resolve to a dict, "
-                            f"got {type(snippet).__name__}."
-                        )
-                    merged.update(snippet)
-
-                merged.update(inline_keys)
-                entry["cfg"] = merged
+            elif isinstance(entry, dict):
+                # Resolve $import in any dict value within the entry
+                for key in list(entry):
+                    if isinstance(entry.get(key), dict) and _IMPORT_KEY in entry[key]:
+                        _resolve_dict_value(entry, key)
                 resolved.append(entry)
             else:
                 resolved.append(entry)