fusion-bench 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

fusion_bench/optim/muon.py

@@ -0,0 +1,339 @@
+# Copied from https://github.com/KellerJordan/Muon
+import torch
+import torch.distributed as dist
+
+
+def zeropower_via_newtonschulz5(G, steps: int):
+    """
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert (
+        G.ndim >= 2
+    )  # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.mT
+        B = (
+            b * A + c * A @ A
+        )  # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X
+
+
+def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True):
+    momentum.lerp_(grad, 1 - beta)
+    update = grad.lerp_(momentum, beta) if nesterov else momentum
+    if update.ndim == 4:  # for the case of conv filters
+        update = update.view(len(update), -1)
+    update = zeropower_via_newtonschulz5(update, steps=ns_steps)
+    update *= max(1, grad.size(-2) / grad.size(-1)) ** 0.5
+    return update
+
+
+class Muon(torch.optim.Optimizer):
+    """
+    Muon - MomentUm Orthogonalized by Newton-schulz
+
+    https://kellerjordan.github.io/posts/muon/
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. For efficient orthogonalization we use a Newton-Schulz iteration, which has the
+    advantage that it can be stably run in bfloat16 on the GPU.
+
+    Muon should only be used for hidden weight layers. The input embedding, final output layer,
+    and any internal gains or biases should be optimized using a standard method such as AdamW.
+    Hidden convolutional weights can be trained using Muon by viewing them as 2D and then
+    collapsing their last 3 dimensions.
+
+    Arguments:
+        lr: The learning rate, in units of spectral norm per update.
+        weight_decay: The AdamW-style weight decay.
+        momentum: The momentum. A value of 0.95 here is usually fine.
+    """
+
+    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum)
+        assert (
+            isinstance(params, list)
+            and len(params) >= 1
+            and isinstance(params[0], torch.nn.Parameter)
+        )
+        params = sorted(params, key=lambda x: x.size(), reverse=True)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params = group["params"]
+            params_pad = params + [torch.empty_like(params[-1])] * (
+                dist.get_world_size() - len(params) % dist.get_world_size()
+            )
+            for base_i in range(len(params))[:: dist.get_world_size()]:
+                if base_i + dist.get_rank() < len(params):
+                    p = params[base_i + dist.get_rank()]
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                    update = muon_update(
+                        p.grad, state["momentum_buffer"], beta=group["momentum"]
+                    )
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update.reshape(p.shape), alpha=-group["lr"])
+                dist.all_gather(
+                    params_pad[base_i : base_i + dist.get_world_size()],
+                    params_pad[base_i + dist.get_rank()],
+                )
+
+        return loss
+
+
+class SingleDeviceMuon(torch.optim.Optimizer):
+    """
+    Muon variant for usage in non-distributed settings.
+    """
+
+    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    # continue
+                    p.grad = torch.zeros_like(p)  # Force synchronization
+                state = self.state[p]
+                if len(state) == 0:
+                    state["momentum_buffer"] = torch.zeros_like(p)
+                update = muon_update(
+                    p.grad, state["momentum_buffer"], beta=group["momentum"]
+                )
+                p.mul_(1 - group["lr"] * group["weight_decay"])
+                p.add_(update.reshape(p.shape), alpha=-group["lr"])
+
+        return loss
+
+
+def adam_update(grad, buf1, buf2, step, betas, eps):
+    buf1.lerp_(grad, 1 - betas[0])
+    buf2.lerp_(grad.square(), 1 - betas[1])
+    buf1c = buf1 / (1 - betas[0] ** step)
+    buf2c = buf2 / (1 - betas[1] ** step)
+    return buf1c / (buf2c.sqrt() + eps)
+
+
+class MuonWithAuxAdam(torch.optim.Optimizer):
+    """
+    Distributed Muon variant that can be used for all parameters in the network, since it runs an
+    internal AdamW for the parameters that are not compatible with Muon. The user must manually
+    specify which parameters shall be optimized with Muon and which with Adam by passing in a
+    list of param_groups with the `use_muon` flag set.
+
+    The point of this class is to allow the user to have a single optimizer in their code, rather
+    than having both a Muon and an Adam which each need to be stepped.
+
+    You can see an example usage below:
+
+    https://github.com/KellerJordan/modded-nanogpt/blob/master/records/052525_MuonWithAuxAdamExample/b01550f9-03d8-4a9c-86fe-4ab434f1c5e0.txt#L470
+    ```
+    hidden_matrix_params = [p for n, p in model.blocks.named_parameters() if p.ndim >= 2 and "embed" not in n]
+    embed_params = [p for n, p in model.named_parameters() if "embed" in n]
+    scalar_params = [p for p in model.parameters() if p.ndim < 2]
+    head_params = [model.lm_head.weight]
+
+    from muon import MuonWithAuxAdam
+    adam_groups = [dict(params=head_params, lr=0.22), dict(params=embed_params, lr=0.6), dict(params=scalar_params, lr=0.04)]
+    adam_groups = [dict(**g, betas=(0.8, 0.95), eps=1e-10, use_muon=False) for g in adam_groups]
+    muon_group = dict(params=hidden_matrix_params, lr=0.05, momentum=0.95, use_muon=True)
+    param_groups = [*adam_groups, muon_group]
+    optimizer = MuonWithAuxAdam(param_groups)
+    ```
+    """
+
+    def __init__(self, param_groups):
+        for group in param_groups:
+            assert "use_muon" in group
+            if group["use_muon"]:
+                group["params"] = sorted(
+                    group["params"], key=lambda x: x.size(), reverse=True
+                )
+                # defaults
+                group["lr"] = group.get("lr", 0.02)
+                group["momentum"] = group.get("momentum", 0.95)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(group.keys()) == set(
+                    ["params", "lr", "momentum", "weight_decay", "use_muon"]
+                )
+            else:
+                # defaults
+                group["lr"] = group.get("lr", 3e-4)
+                group["betas"] = group.get("betas", (0.9, 0.95))
+                group["eps"] = group.get("eps", 1e-10)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(group.keys()) == set(
+                    ["params", "lr", "betas", "eps", "weight_decay", "use_muon"]
+                )
+        super().__init__(param_groups, dict())
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if group["use_muon"]:
+                params = group["params"]
+                params_pad = params + [torch.empty_like(params[-1])] * (
+                    dist.get_world_size() - len(params) % dist.get_world_size()
+                )
+                for base_i in range(len(params))[:: dist.get_world_size()]:
+                    if base_i + dist.get_rank() < len(params):
+                        p = params[base_i + dist.get_rank()]
+                        if p.grad is None:
+                            # continue
+                            p.grad = torch.zeros_like(p)  # Force synchronization
+                        state = self.state[p]
+                        if len(state) == 0:
+                            state["momentum_buffer"] = torch.zeros_like(p)
+                        update = muon_update(
+                            p.grad, state["momentum_buffer"], beta=group["momentum"]
+                        )
+                        p.mul_(1 - group["lr"] * group["weight_decay"])
+                        p.add_(update.reshape(p.shape), alpha=-group["lr"])
+                    dist.all_gather(
+                        params_pad[base_i : base_i + dist.get_world_size()],
+                        params_pad[base_i + dist.get_rank()],
+                    )
+            else:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+                    state["step"] += 1
+                    update = adam_update(
+                        p.grad,
+                        state["exp_avg"],
+                        state["exp_avg_sq"],
+                        state["step"],
+                        group["betas"],
+                        group["eps"],
+                    )
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss
+
+
+class SingleDeviceMuonWithAuxAdam(torch.optim.Optimizer):
+    """
+    Non-distributed variant of MuonWithAuxAdam.
+    """
+
+    def __init__(self, param_groups):
+        for group in param_groups:
+            assert "use_muon" in group
+            if group["use_muon"]:
+                # defaults
+                group["lr"] = group.get("lr", 0.02)
+                group["momentum"] = group.get("momentum", 0.95)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(group.keys()) == set(
+                    ["params", "lr", "momentum", "weight_decay", "use_muon"]
+                )
+            else:
+                # defaults
+                group["lr"] = group.get("lr", 3e-4)
+                group["betas"] = group.get("betas", (0.9, 0.95))
+                group["eps"] = group.get("eps", 1e-10)
+                group["weight_decay"] = group.get("weight_decay", 0)
+                assert set(group.keys()) == set(
+                    ["params", "lr", "betas", "eps", "weight_decay", "use_muon"]
+                )
+        super().__init__(param_groups, dict())
+
+    @torch.no_grad()
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if group["use_muon"]:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+                    update = muon_update(
+                        p.grad, state["momentum_buffer"], beta=group["momentum"]
+                    )
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update.reshape(p.shape), alpha=-group["lr"])
+            else:
+                for p in group["params"]:
+                    if p.grad is None:
+                        # continue
+                        p.grad = torch.zeros_like(p)  # Force synchronization
+                    state = self.state[p]
+                    if len(state) == 0:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+                    state["step"] += 1
+                    update = adam_update(
+                        p.grad,
+                        state["exp_avg"],
+                        state["exp_avg_sq"],
+                        state["step"],
+                        group["betas"],
+                        group["eps"],
+                    )
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss

fusion_bench/programs/__init__.py

@@ -7,11 +7,13 @@ from fusion_bench.utils.lazy_imports import LazyImporter
 _import_structure = {
     "base_program": ["BaseHydraProgram"],
     "fabric_fusion_program": ["FabricModelFusionProgram"],
+    "fusion_program": ["ModelFusionProgram"],
 }
 
 if TYPE_CHECKING:
     from .base_program import BaseHydraProgram
     from .fabric_fusion_program import FabricModelFusionProgram
+    from .fusion_program import ModelFusionProgram
 else:
     sys.modules[__name__] = LazyImporter(
         __name__,

fusion_bench/programs/fabric_fusion_program.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, List, Optional, Union  # noqa: F401
 
 import lightning as L
-from lightning.fabric.utilities.rank_zero import rank_zero_only
+from lightning_utilities.core.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 from torch import nn
 from tqdm.auto import tqdm
@@ -236,7 +236,7 @@ class FabricModelFusionProgram(
 
         # create symbol link to hydra output directory
         if (
-            self.fabric.is_global_zero
+            rank_zero_only.rank == 0
             and self.log_dir is not None
             and os.path.abspath(self.log_dir) != os.path.abspath(get_hydra_output_dir())
         ):

fusion_bench/programs/fusion_program.py

@@ -0,0 +1,271 @@
+import json
+import os
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+import lightning as L
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from tqdm.auto import tqdm
+
+from fusion_bench import (
+    BaseAlgorithm,
+    BaseHydraProgram,
+    BaseModelPool,
+    BaseTaskPool,
+    RuntimeConstants,
+    auto_register_config,
+    get_rankzero_logger,
+    import_object,
+    instantiate,
+    timeit_context,
+)
+from fusion_bench.utils.json import print_json
+from fusion_bench.utils.rich_utils import print_bordered, print_config_tree
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class ModelFusionProgram(BaseHydraProgram):
+    method: BaseAlgorithm
+    modelpool: BaseModelPool
+    taskpool: Optional[BaseTaskPool] = None
+
+    _config_mapping = BaseHydraProgram._config_mapping | {
+        "_method": "method",
+        "_modelpool": "modelpool",
+        "_taskpool": "taskpool",
+        "fast_dev_run": "fast_dev_run",
+        "seed": "seed",
+        "path": "path",
+    }
+
+    def __init__(
+        self,
+        method: DictConfig,
+        modelpool: DictConfig,
+        taskpool: Optional[DictConfig] = None,
+        *,
+        print_config: bool = True,
+        dry_run: bool = False,
+        report_save_path: Optional[str] = None,
+        merged_model_save_path: Optional[str] = None,
+        merged_model_save_kwargs: Optional[DictConfig] = None,
+        fast_dev_run: bool = False,
+        seed: Optional[int] = None,
+        print_function_call: bool = True,
+        path: DictConfig = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._method = method
+        self._modelpool = modelpool
+        self._taskpool = taskpool
+        self.report_save_path = report_save_path
+        self.merged_model_save_path = merged_model_save_path
+        self.merged_model_save_kwargs = merged_model_save_kwargs
+        self.fast_dev_run = fast_dev_run
+        self.seed = seed
+        self.path = path
+        RuntimeConstants.debug = fast_dev_run
+        RuntimeConstants.print_function_call = print_function_call
+        if path is not None:
+            RuntimeConstants.cache_dir = path.get("cache_dir", None)
+
+        if print_config:
+            print_config_tree(
+                self.config,
+                print_order=["method", "modelpool", "taskpool"],
+            )
+        if dry_run:
+            log.info("The program is running in dry-run mode. Exiting.")
+            exit(0)
+
+    def _instantiate_and_setup(
+        self, config: DictConfig, compat_load_fn: Optional[str] = None
+    ):
+        R"""
+        Instantiates and sets up an object based on the provided configuration.
+
+        This method performs the following steps:
+        1. Checks if the configuration dictionary contains the key "_target_".
+        2. If "_target_" is not found (for v0.1.x), attempts to instantiate the object using a compatible load function if provided.
+           - Logs a warning if "_target_" is missing.
+           - If `compat_load_fn` is provided, imports the function and uses it to instantiate the object.
+           - If `compat_load_fn` is not provided, raises a ValueError.
+        3. If "_target_" is found (for v.0.2.0 and above), attempts to import and instantiate the object using the `instantiate` function.
+           - Ensures the target can be imported.
+           - Uses the `instantiate` function with `_recursive_` set based on the configuration.
+        4. Sets the `_program` attribute of the instantiated object to `self` if the object has this attribute.
+        5. Returns the instantiated and set up object.
+        """
+        if "_target_" not in config:
+            log.warning(
+                "No '_target_' key found in config. Attempting to instantiate the object in a compatible way."
+            )
+            if compat_load_fn is not None:
+                compat_load_fn = import_object(compat_load_fn)
+                if rank_zero_only.rank == 0:
+                    print_bordered(
+                        OmegaConf.to_yaml(config),
+                        title="instantiate compat object",
+                        style="magenta",
+                        code_style="yaml",
+                    )
+                obj = compat_load_fn(config)
+            else:
+                raise ValueError(
+                    "No load function provided. Please provide a load function to instantiate the object."
+                )
+        else:
+            # try to import the object from the target
+            # this checks if the target is valid and can be imported
+            import_object(config._target_)
+            obj = instantiate(
+                config,
+                _recursive_=config.get("_recursive_", False),
+            )
+        if hasattr(obj, "_program"):
+            obj._program = self
+        return obj
+
+    def save_merged_model(self, merged_model):
+        """
+        Saves the merged model to the specified path.
+        """
+        if self.merged_model_save_path is not None:
+            # path to save the merged model, use "{log_dir}" to refer to the logger directory
+            save_path: str = self.merged_model_save_path
+            if "{log_dir}" in save_path and self.log_dir is not None:
+                save_path = save_path.format(log_dir=self.log_dir)
+
+            if os.path.dirname(save_path):
+                os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+            # save the merged model
+            if self.merged_model_save_kwargs is not None:
+                merged_model_save_kwargs = self.merged_model_save_kwargs
+            else:
+                merged_model_save_kwargs = {}
+            with timeit_context(f"Saving the merged model to {save_path}"):
+                self.modelpool.save_model(
+                    merged_model,
+                    save_path,
+                    **merged_model_save_kwargs,
+                )
+        else:
+            print("No save path specified for the merged model. Skipping saving.")
+
+    def evaluate_merged_model(
+        self,
+        taskpool: BaseTaskPool,
+        merged_model: Union[nn.Module, Dict, Iterable],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Union[Dict, List, Any]:
+        """
+        Evaluates the merged model using the provided task pool.
+
+        Depending on the type of the merged model, this function handles the evaluation differently:
+        - If the merged model is an instance of `nn.Module`, it directly evaluates the model.
+        - If the merged model is a dictionary, it extracts the model from the dictionary and evaluates it.
+          The evaluation report is then updated with the remaining dictionary items.
+        - If the merged model is an iterable, it recursively evaluates each model in the iterable.
+        - Raises a `ValueError` if the merged model is of an invalid type.
+
+        Args:
+            taskpool: The task pool used for evaluating the merged model.
+            merged_model: The merged model to be evaluated. It can be an instance of `nn.Module`, a dictionary, or an iterable.
+            *args: Additional positional arguments to be passed to the `evaluate` method of the taskpool.
+            **kwargs: Additional keyword arguments to be passed to the `evaluate` method of the taskpool.
+
+        Returns:
+            The evaluation report. The type of the report depends on the type of the merged model:
+            - If the merged model is an instance of `nn.Module`, the report is a dictionary.
+            - If the merged model is a dictionary, the report is a dictionary updated with the remaining dictionary items.
+            - If the merged model is an iterable, the report is a list of evaluation reports.
+        """
+        if isinstance(merged_model, nn.Module):
+            report = taskpool.evaluate(merged_model, *args, **kwargs)
+            return report
+        elif isinstance(merged_model, Dict):
+            report = {}
+            for key, item in merged_model.items():
+                if isinstance(item, nn.Module):
+                    report[key] = taskpool.evaluate(item, *args, **kwargs)
+                elif key == "models":
+                    # for multi-model evaluation
+                    report[key] = self.evaluate_merged_model(
+                        taskpool, item, *args, **kwargs
+                    )
+                else:
+                    # metadata
+                    report[key] = item
+            return report
+        elif isinstance(merged_model, Iterable):
+            return [
+                self.evaluate_merged_model(taskpool, m, *args, **kwargs)
+                for m in tqdm(merged_model, desc="Evaluating models")
+            ]
+        else:
+            raise ValueError(f"Invalid type for merged model: {type(merged_model)}")
+
+    def run(self):
+        """
+        Executes the model fusion program.
+        """
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+
+        log.info("Running the model fusion program.")
+        # setup the modelpool, method, and taskpool
+        log.info("loading model pool")
+        self.modelpool = self._instantiate_and_setup(
+            self._modelpool,
+            compat_load_fn="fusion_bench.compat.modelpool.load_modelpool_from_config",
+        )
+        log.info("loading method")
+        self.method = self._instantiate_and_setup(
+            self._method,
+            compat_load_fn="fusion_bench.compat.method.load_algorithm_from_config",
+        )
+        if self._taskpool is not None:
+            log.info("loading task pool")
+            self.taskpool = self._instantiate_and_setup(
+                self._taskpool,
+                compat_load_fn="fusion_bench.compat.taskpool.load_taskpool_from_config",
+            )
+
+        self.method.on_run_start()
+        merged_model = self.method.run(self.modelpool)
+        self.method.on_run_end()
+
+        if merged_model is None:
+            log.info(
+                "No merged model returned by the method. Skipping saving and evaluation."
+            )
+        else:
+            self.save_merged_model(merged_model)
+            if self.taskpool is not None:
+                report = self.evaluate_merged_model(self.taskpool, merged_model)
+                try:
+                    if rank_zero_only.rank == 0:
+                        print_json(report, print_type=False)
+                except Exception as e:
+                    log.warning(f"Failed to pretty print the report: {e}")
+                    log.info(report)
+                if self.report_save_path is not None:
+                    # save report (Dict) to a file
+                    # if the directory of `save_report` does not exists, create it
+                    if (
+                        "{log_dir}" in self.report_save_path
+                        and self.path.log_dir is not None
+                    ):
+                        self.report_save_path = self.report_save_path.format(
+                            log_dir=self.path.log_dir
+                        )
+                    os.makedirs(os.path.dirname(self.report_save_path), exist_ok=True)
+                    json.dump(report, open(self.report_save_path, "w"))
+            else:
+                log.info("No task pool specified. Skipping evaluation.")

fusion_bench/scripts/cli.py

@@ -20,8 +20,8 @@ log = logging.getLogger(__name__)
 
 
 def _get_default_config_path():
-    for config_dir in ["fusion_bench_config", "config"]:
-        for config_path_root in [os.getcwd(), PROJECT_ROOT_PATH]:
+    for config_path_root in [os.getcwd(), PROJECT_ROOT_PATH]:
+        for config_dir in ["config", "fusion_bench_config"]:
             config_path = os.path.join(config_path_root, config_dir)
             if os.path.exists(config_path) and os.path.isdir(config_path):
                 return os.path.abspath(config_path)