fusion-bench 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

fusion_bench/modelpool/resnet_for_image_classification.py

@@ -0,0 +1,208 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    override,
+)
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+
+from fusion_bench import BaseModelPool, auto_register_config, get_rankzero_logger
+from fusion_bench.tasks.clip_classification import get_classnames, get_num_classes
+
+if TYPE_CHECKING:
+    from torchvision.models import ResNet as TorchVisionResNet
+
+log = get_rankzero_logger(__name__)
+
+
+def load_torchvision_resnet(
+    model_name: str, weights: Optional[str], num_classes: Optional[int]
+) -> "TorchVisionResNet":
+    import torchvision.models
+
+    model_fn = getattr(torchvision.models, model_name)
+    model: "TorchVisionResNet" = model_fn(weights=weights)
+
+    if num_classes is not None:
+        model.fc = nn.Linear(model.fc.in_features, num_classes)
+
+    return model
+
+
+def load_transformers_resnet(
+    config_path: str, pretrained: bool, dataset_name: Optional[str]
+):
+    from transformers import AutoConfig, ResNetForImageClassification
+
+    if pretrained:
+        model = ResNetForImageClassification.from_pretrained(config_path)
+    else:
+        config = AutoConfig.from_pretrained(config_path)
+        model = ResNetForImageClassification(config)
+
+    if dataset_name is None:
+        return model
+
+    classnames = get_classnames(dataset_name)
+    id2label = {i: c for i, c in enumerate(classnames)}
+    label2id = {c: i for i, c in enumerate(classnames)}
+    model.config.id2label = id2label
+    model.config.label2id = label2id
+
+    model.classifier[1] = (
+        nn.Linear(
+            model.classifier[1].in_features,
+            len(classnames),
+        )
+        if model.config.num_labels > 0
+        else nn.Identity()
+    )
+    return model
+
+
+@auto_register_config
+class ResNetForImageClassificationPool(BaseModelPool):
+    def __init__(self, type: str, **kwargs):
+        super().__init__(**kwargs)
+        assert type in ["torchvision", "transformers"]
+
+    def load_processor(
+        self, stage: Literal["train", "val", "test"] = "test", *args, **kwargs
+    ):
+        if self.type == "torchvision":
+            from torchvision import transforms
+
+            to_tensor = transforms.ToTensor()
+            normalize = transforms.Normalize(
+                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+            )
+            if stage == "train":
+                train_transform = transforms.Compose(
+                    [
+                        transforms.RandomResizedCrop(224),
+                        transforms.RandomHorizontalFlip(),
+                        to_tensor,
+                        normalize,
+                    ]
+                )
+                return train_transform
+            else:
+                val_transform = transforms.Compose(
+                    [
+                        transforms.Resize(256),
+                        transforms.CenterCrop(224),
+                        to_tensor,
+                        normalize,
+                    ]
+                )
+                return val_transform
+
+        elif self.type == "transformers":
+            from transformers import AutoImageProcessor
+
+            if self.has_pretrained:
+                config_path = self._models["_pretrained_"].config_path
+            else:
+                for model_cfg in self._models.values():
+                    if isinstance(model_cfg, str):
+                        config_path = model_cfg
+                        break
+                    if "config_path" in model_cfg:
+                        config_path = model_cfg["config_path"]
+                        break
+            return AutoImageProcessor.from_pretrained(config_path)
+
+    @override
+    def load_model(self, model_name_or_config: Union[str, DictConfig], *args, **kwargs):
+        log.debug(f"Loading model: {model_name_or_config}", stacklevel=2)
+        if (
+            isinstance(model_name_or_config, str)
+            and model_name_or_config in self._models
+        ):
+            model_name_or_config = self._models[model_name_or_config]
+
+        if self.type == "torchvision":
+            from torchvision.models import (
+                resnet18,
+                resnet34,
+                resnet50,
+                resnet101,
+                resnet152,
+            )
+
+            match model_name_or_config:
+                case "resnet18":
+                    model = resnet18()
+                case "resnet34":
+                    model = resnet34()
+                case "resnet50":
+                    model = resnet50()
+                case "resnet101":
+                    model = resnet101()
+                case "resnet152":
+                    model = resnet152()
+                case dict() | DictConfig() as model_config:
+                    if "dataset_name" in model_config:
+                        num_classes = get_num_classes(model_config["dataset_name"])
+                        if "num_classes" in model_config:
+                            assert (
+                                num_classes == model_config["num_classes"]
+                            ), f"num_classes mismatch: {num_classes} vs {model_config['num_classes']}"
+                    elif "num_classes" in model_config:
+                        num_classes = model_config["num_classes"]
+                    else:
+                        num_classes = None
+                    model = load_torchvision_resnet(
+                        model_name=model_config["model_name"],
+                        weights=model_config.get("weights", None),
+                        num_classes=num_classes,
+                    )
+                case _:
+                    raise ValueError(
+                        f"Invalid model_name_or_config type: {type(model_name_or_config)}"
+                    )
+        elif self.type == "transformers":
+            match model_name_or_config:
+                case str() as model_path:
+                    from transformers import AutoModelForImageClassification
+
+                    model = AutoModelForImageClassification.from_pretrained(model_path)
+                case dict() | DictConfig() as model_config:
+
+                    model = load_transformers_resnet(
+                        config_path=model_config["config_path"],
+                        pretrained=model_config.get("pretrained", False),
+                        dataset_name=model_config.get("dataset_name", None),
+                    )
+                case _:
+                    raise ValueError(
+                        f"Invalid model_name_or_config type: {type(model_name_or_config)}"
+                    )
+
+            # override forward to return logits only
+            original_forward = model.forward
+            model.forward = lambda pixel_values, **kwargs: original_forward(
+                pixel_values=pixel_values, **kwargs
+            ).logits
+            model.original_forward = original_forward
+        else:
+            raise ValueError(f"Unknown model type: {self.type}")
+        return model
+
+    @override
+    def save_model(self, model, path, *args, **kwargs):
+        if self.type == "torchvision":
+            torch.save(model.state_dict(), path)
+        elif self.type == "transformers":
+            model.save_pretrained(path)
+            self.load_processor().save_pretrained(path)
+        else:
+            raise ValueError(f"Unknown model type: {self.type}")

fusion_bench/models/__init__.py

@@ -1,10 +1,36 @@
 # flake8: noqa F401
-from fusion_bench.utils import LazyStateDict
-
-from . import separate_io, utils
-from .hf_utils import (
-    create_default_model_card,
-    load_model_card_template,
-    save_pretrained_with_remote_code,
-)
-from .parameter_dict import ParameterDictModel
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+from . import utils
+
+_extra_objects = {
+    "utils": utils,
+}
+_import_structure = {
+    "hf_utils": [
+        "create_default_model_card",
+        "load_model_card_template",
+        "save_pretrained_with_remote_code",
+    ],
+    "parameter_dict": ["ParameterDictModel"],
+    "separate_io": ["separate_load", "separate_save"],
+}
+
+if TYPE_CHECKING:
+    from .hf_utils import (
+        create_default_model_card,
+        load_model_card_template,
+        save_pretrained_with_remote_code,
+    )
+    from .parameter_dict import ParameterDictModel
+    from .separate_io import separate_load, separate_save
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/optim/__init__.py

@@ -1,2 +1,40 @@
-from . import exception, lr_scheduler
-from .mezo import MeZO
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+from . import lr_scheduler
+
+_extra_objects = {
+    "lr_scheduler": lr_scheduler,
+}
+_import_structure = {
+    "exception": [
+        "NoClosureError",
+        "NoSparseGradientError",
+        "NegativeLRError",
+        "NegativeStepError",
+        "ZeroParameterSizeError",
+    ],
+    "mezo": ["MeZO"],
+    "muon": ["Muon"],
+}
+
+if TYPE_CHECKING:
+    from .exception import (
+        NegativeLRError,
+        NegativeStepError,
+        NoClosureError,
+        NoSparseGradientError,
+        ZeroParameterSizeError,
+    )
+    from .mezo import MeZO
+    from .muon import Muon
+
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/optim/lr_scheduler/__init__.py

@@ -1 +1,27 @@
-from .linear_warmup import *
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+_import_structure = {
+    "linear_warmup": [
+        "BaseLinearWarmupScheduler",
+        "LinearWarmupScheduler",
+        "CosineDecayWithWarmup",
+        "PolySchedulerWithWarmup",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .linear_warmup import (
+        BaseLinearWarmupScheduler,
+        CosineDecayWithWarmup,
+        LinearWarmupScheduler,
+        PolySchedulerWithWarmup,
+    )
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+    )

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__,