fusion-bench 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

fusion_bench/method/gossip/layer_wise_gossip.py

@@ -0,0 +1,434 @@
+import copy
+import gc
+import logging
+import os
+from abc import abstractmethod
+from typing import Any, Callable, List, Mapping, Union, cast  # noqa: F401
+
+import torch
+from lightning.fabric.utilities.rank_zero import rank_zero_only
+from omegaconf import DictConfig
+from torch import Tensor
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+from fusion_bench.compat.method import ModelFusionAlgorithm
+from fusion_bench.compat.modelpool import ModelPool
+from fusion_bench.mixins.lightning_fabric import LightningFabricMixin
+from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.modelpool import (
+    CLIPVisionModelPool,
+    GPT2ForSequenceClassificationPool,
+)
+from fusion_bench.models.wrappers.layer_wise_fusion import (
+    LayerWiseMergedModel,
+    get_layer_wise_weights,
+)
+from fusion_bench.utils.data import load_tensor_from_file
+
+from .entropy_loss import entropy_loss
+
+log = logging.getLogger(__name__)
+
+
+# obtain the current GPU memory usage
+def get_memory_usage(desc):
+    allocated = torch.cuda.memory_allocated() / 1024**2  # 转换为 MB
+    cached = torch.cuda.memory_reserved() / 1024**2  # 转换为 MB
+    return (
+        f"{desc}\nAllocated Memory: {allocated:.2f} MB\nCached Memory: {cached:.2f} MB"
+    )
+
+
+class ModelScheduler:
+    """
+    Manage the storage of models, schedule the order in which models are loaded to GPU
+    transfer data between the CPU and GPu
+    """
+
+    def __init__(
+        self,
+        config: DictConfig,
+        modelpool: ModelPool,
+    ):
+        self.pretrained_model = modelpool.load_model("_pretrained_")
+        self.finetuned_models = [
+            modelpool.load_model(name) for name in modelpool.model_names
+        ]
+        self.num_finetuned_models = len(self.finetuned_models)
+        self.new_finetuned_models = copy.deepcopy(self.finetuned_models)
+        self.finetuned_models_name = [name for name in modelpool.model_names]
+
+        self.config = config
+
+    @torch.no_grad()  # not sure whether to use this
+    def __call__(self, model_id):
+        """
+        return models and relevant data in each step
+        """
+        pretrained_model = copy.deepcopy(self.pretrained_model)
+        if self.config.topo == "ring":
+            finetuned_models = [
+                copy.deepcopy(
+                    self.finetuned_models[(model_id + 1) % self.num_finetuned_models]
+                ),
+                copy.deepcopy(self.finetuned_models[model_id]),
+                copy.deepcopy(
+                    self.finetuned_models[(model_id - 1) % self.num_finetuned_models]
+                ),
+            ]
+        elif "rotate" in self.config.topo:
+            number = self.config.topo.split("_")[1]
+            finetuned_models = [copy.deepcopy(self.finetuned_models[model_id])]
+            for i in range(0, int(number)):
+                finetuned_models.append(
+                    copy.deepcopy(
+                        self.finetuned_models[
+                            (model_id + i + 1) % self.num_finetuned_models
+                        ]
+                    )
+                )
+        # initialize layer-wise weights using the provided configuration `init_values` or load from file if `weights` is provided
+        if self.config.weights is None:
+            layer_wise_weight = get_layer_wise_weights(
+                num_models=len(finetuned_models),
+                num_layers=len(
+                    tuple(
+                        filter(lambda p: p.requires_grad, pretrained_model.parameters())
+                    )
+                ),
+                init_values=self.config.init_values,
+            )
+        else:
+            if isinstance(self.config.weights, str):
+                # self.config.weights is a path to a saved tensor
+                layer_wise_weight = load_tensor_from_file(self.config.weights)
+            else:
+                raise ValueError(f"Unsupported weights format: {self.config.weights}")
+
+        module = LayerWiseMergedModel(
+            layer_wise_weight=layer_wise_weight,
+            pretrained_model=pretrained_model,
+            finetuned_models=finetuned_models,
+            clamp_weights=self.config.clamp_weights,
+            tie_weights=self.config.tie_weights,
+            strict=self.config.strict,
+        )
+        print(f"{layer_wise_weight.size()=}, {layer_wise_weight.numel()=}")
+        return module
+
+    def store_model(self, new_finetuned_model_dict, model_id):
+        """
+        store new finetuned model after every turn of adamerging
+        """
+        self.new_finetuned_models[model_id].load_state_dict(new_finetuned_model_dict)
+
+    def update_models(self):
+        self.finetuned_models = copy.deepcopy(self.new_finetuned_models)
+
+    def get_final_models(self, idx=None):
+        # need a check
+        if idx is not None:
+            return copy.deepcopy(self.finetuned_models[idx])
+
+        final_models = [
+            {"name": name, "model": model}
+            for name, model in zip(self.finetuned_models_name, self.finetuned_models)
+        ]
+        num_finetuned_models = len(self.finetuned_models)
+
+        average_model = copy.deepcopy(self.pretrained_model)
+        state_dict = average_model.state_dict(keep_vars=True)
+        for name, _ in self.finetuned_models[0].named_parameters():
+            state_dict[name].data.zero_()
+        for model in self.finetuned_models:
+            for name, param in model.named_parameters():
+                state_dict[name] = state_dict[name] + 1 / num_finetuned_models * param
+
+        average_model.load_state_dict(state_dict)
+        final_models += [{"name": "average model", "model": average_model}]
+
+        return final_models
+
+    def move_to(self, device):
+        self.pretrained_model.to(device=device)
+        for model in self.finetuned_models:
+            model.to(device=device)
+
+
+class LayerWiseGossipAlgorithm(
+    ModelFusionAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
+    """
+    Implements the Layer-Wise AdaMerging Algorithm.
+
+    This class merges the layers of a pretrained model with those of several fine-tuned models.
+    The merging is controlled by layer-wise weights, which can be initialized based on a provided configuration or loaded from a file.
+    """
+
+    def __init__(self, algorithm_config: DictConfig):
+        """
+        Initialize the LayerWiseAdaMergingAlgorithm with the given configuration.
+
+        Args:
+            algorithm_config (DictConfig): The configuration for the algorithm.
+        """
+        super().__init__(algorithm_config)
+        self._program = None
+
+    @rank_zero_only
+    def save_merging_weights(self, file_path: str, merging_weights: torch.Tensor):
+        """
+        Save the merging weights to a file.
+
+        Args:
+            file_path (str): The path to save the merging weights.
+            merging_weights (torch.Tensor): The merging weights to save.
+        """
+        if self.fabric.is_global_zero and self.config.get(
+            "save_merging_weights", False
+        ):
+            if isinstance(file_path, str) and not file_path.startswith(("/", ".")):
+                # if the file path is not absolute or relative to current working directory, save it in the log directory
+                save_path = os.path.join(self.log_dir, file_path)
+            else:
+                save_path = file_path
+            log.info(f"saving merging weights to {save_path}.")
+            if os.path.dirname(save_path):
+                os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            torch.save(merging_weights.detach().cpu(), save_path)
+
+    def free_gpu_memory(self, module: LayerWiseMergedModel):
+        module.pretrained_model.to("cpu")
+        for model in module.task_vectors:
+            model.to("cpu")
+        del module
+        gc.collect()
+        torch.cuda.empty_cache()
+        log.info(get_memory_usage("after freeing memory, the memory usage of GPU is:"))
+
+    def update_datasets(self, datasets):
+        """
+        for evary epoch of local adamerging, we only use the data set corresponding to the model involved in the fusion
+        """
+        num_datasets = len(datasets)
+        datasets_copy = datasets.copy()
+        if self.config.topo == "ring":
+            for i in range(num_datasets):
+                datasets[i] = (
+                    datasets_copy[i]
+                    .union(datasets_copy[(i + 1) % num_datasets])
+                    .union(datasets_copy[(i - 1) % num_datasets])
+                )
+        elif "rotate" in self.config.topo:
+            number = self.config.topo.split("_")[1]
+            for i in range(num_datasets):
+                datasets[i] = datasets_copy[i]
+                for j in range(0, int(number)):
+                    datasets[i] = datasets[i].union(
+                        datasets_copy[(i + j + 1) % num_datasets]
+                    )
+        return datasets
+
+    def run(self, modelpool: ModelPool):
+        """
+        Run the Layer-Wise AdaMerging Algorithm.
+
+        This method constructs the wrapped model and performs test-time adaptation if necessary.
+
+        Args:
+            modelpool (ModelPool): The model pool containing the pretrained and fine-tuned models.
+
+        Returns:
+            LayerWiseMergedModel: The merged model after test-time adaptation.
+        """
+        log.info("Fusing models using layer-wise adaptive merging.")
+        self.modelpool = modelpool
+        self.log_hyperparams(self.config)
+        self.num_finetuned_models = len(modelpool.model_names)
+        datasets = [{dataset} for dataset in modelpool.model_names]
+
+        with self.profile("construct the wrapped model"):
+            model_scheduler = ModelScheduler(
+                modelpool=self.modelpool, config=self.config
+            )
+
+        if self.config.weights is not None:
+            # skip the test-time adaptation
+            return module.merge_and_unload()
+        else:
+            for step_idx in tqdm(
+                range(self.config.gossip_max_steps),
+                "Gossip merging",
+                dynamic_ncols=True,
+            ):
+                datasets = self.update_datasets(datasets)
+                log.info(f"Gossip merging step:, {step_idx}")
+                for model_id in tqdm(
+                    range(self.num_finetuned_models),
+                    "local admerging",
+                    dynamic_ncols=True,
+                ):
+                    if self.config.gossip_skip_adamerging == True:
+                        # skip adamerging, only merge
+                        with self.profile("construct the local wrapped model"):
+                            module = model_scheduler(model_id)
+                        log.info(
+                            f"skip adamerging, only merge ({modelpool.model_names[model_id]})"
+                        )
+                        model_scheduler.store_model(module.merge_weights(), model_id)
+                        self.free_gpu_memory(module)
+                    else:
+                        with self.profile("construct the local wrapped model"):
+                            module = model_scheduler(model_id)
+
+                        if self.config.improve_dataset == True:
+                            log.info(
+                                f"improved datasets, the datasets used in this local merging is {datasets[model_id]}"
+                            )
+                        else:
+                            log.info(
+                                f"unimproved datasets, the datasets used in this local merging is {modelpool.model_names}"
+                            )
+                        with self.profile("test-time adaptation"):
+                            module = self.test_time_adaptation(
+                                module, datasets[model_id]
+                            )
+                        model_scheduler.store_model(module.merge_weights(), model_id)
+                        log.info(
+                            get_memory_usage(
+                                f"after local merging ({modelpool.model_names[model_id]}), the memory usage of GPU is:"
+                            )
+                        )
+                        self.free_gpu_memory(
+                            module
+                        )  # simulate distributed GPU memory usage as much as possible
+
+                model_scheduler.update_models()
+
+                if "rotate" in self.config.topo:
+                    number = self.config.topo.split("_")[1]
+                    if int(number) == 1 and step_idx >= 20:
+                        self._program.evaluate_merged_model(
+                            self._program.taskpool, model_scheduler.get_final_models()
+                        )
+                        model_scheduler.move_to("cpu")
+                else:
+                    if (
+                        self.config.accuracy_test_interval != 0
+                        and (step_idx + 1) % self.config.accuracy_test_interval == 0
+                    ):
+                        self._program.evaluate_merged_model(
+                            self._program.taskpool, model_scheduler.get_final_models()
+                        )
+                        model_scheduler.move_to("cpu")
+        return model_scheduler.get_final_models()
+
+    def on_test_time_adaptation_start(self):
+        """
+        Something to do before the test-time adaptation starts. Such as setting up the task-specific heads.
+        """
+        pass
+
+    @abstractmethod
+    def get_shuffled_test_loader_iter(self, task: str) -> DataLoader:
+        """
+        Loader of test dataset for test-time adaptation. labels are not needed.
+
+        Args:
+            task (str): The name of the task.
+
+        Returns:
+            DataLoader: The data loader for the test dataset.
+        """
+        pass
+
+    @abstractmethod
+    def compute_logits(self, module, images: Tensor, task: str) -> Tensor:
+        """
+        Compute the logits for the given images and task.
+
+        Args:
+            module: The model module.
+            images (Tensor): The input images.
+            task (str): The name of the task.
+
+        Returns:
+            Tensor: The computed logits.
+        """
+        pass
+
+    def test_time_adaptation(self, module: LayerWiseMergedModel, datasets):
+        """
+        Perform test-time adaptation on the merged model.
+
+        This method adapts the merging weights during test-time to improve performance.
+
+        Args:
+            module (LayerWiseMergedModel): The merged model.
+
+        Returns:
+            LayerWiseMergedModel: The adapted merged model.
+        """
+        self.on_test_time_adaptation_start()
+
+        # configure optimizer
+        if self.config.optimizer == "adam":
+            optimizer = torch.optim.Adam([module.merge_weight], lr=self.config.lr)
+            print(f"{optimizer=}")
+            module, optimizer = self.fabric.setup(module, optimizer)
+            log.info(
+                get_memory_usage(
+                    "after loading models and optimizer, the memory usage of GPU is:"
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {self.config.optimizer}")
+
+        module.train()
+        module.merge_weights()
+        for step_idx in (
+            pbar := tqdm(
+                range(self.config.max_steps if not self.is_debug_mode else 1),
+                ("[DEBUG MODE] " if self.is_debug_mode else "")
+                + "AdaMerging Test-time adaptation",
+                dynamic_ncols=True,
+            )
+        ):
+            # default behavior for first-order optimizers
+            for task in self.modelpool.model_names:
+                if self.config.improve_dataset == True and task not in datasets:
+                    continue
+                with self.profile("data loading"):
+                    batch = next(self.get_shuffled_test_loader_iter(task))
+                with self.profile("forward pass"):
+                    if isinstance(self.modelpool, GPT2ForSequenceClassificationPool):
+                        logits = self.compute_logits(module, batch, task)
+                    elif isinstance(self.modelpool, CLIPVisionModelPool):
+                        logits = self.compute_logits(module, batch[0], task)
+                    loss = entropy_loss(logits)
+                with self.profile("backward pass"):
+                    self.fabric.backward(loss, retain_graph=True)
+
+            with self.profile("optimizer step"):
+                optimizer.step()
+                optimizer.zero_grad()
+            with self.profile("merging weights"):
+                module.merge_weights()
+
+            metrics = {
+                "train/loss": loss.item(),
+                "train/weight_max": module.merge_weight.max().item(),
+                "train/weight_min": module.merge_weight.min().item(),
+                "train/weight_mean": module.merge_weight.mean().item(),
+            }
+            self.fabric.log_dict(metrics, step=step_idx)
+            pbar.set_postfix(metrics)
+
+        self.print_profile_summary()
+        del optimizer
+        gc.collect()
+        torch.cuda.empty_cache()
+        return module

fusion_bench/method/gossip/min_norm_solvers.py

@@ -0,0 +1,227 @@
+# This code is from
+# Multi-Task Learning as Multi-Objective Optimization
+# Ozan Sener, Vladlen Koltun
+# Neural Information Processing Systems (NeurIPS) 2018
+# https://github.com/intel-isl/MultiObjectiveOptimization
+from typing import Union
+
+import numpy as np
+import torch
+
+
+def np_sum(x: Union[torch.Tensor, np.ndarray]) -> float:
+    if isinstance(x, torch.Tensor):
+        return x.sum().item()
+    return np.sum(x)
+
+
+def to_numpy(x: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    return x
+
+
+class MinNormSolver:
+    MAX_ITER = 250
+    STOP_CRIT = 1e-5
+
+    def _min_norm_element_from2(v1v1, v1v2, v2v2):
+        """
+        Analytical solution for min_{c} |cx_1 + (1-c)x_2|_2^2
+        d is the distance (objective) optimzed
+        v1v1 = <x1,x1>
+        v1v2 = <x1,x2>
+        v2v2 = <x2,x2>
+        """
+        if v1v2 >= v1v1:
+            # Case: Fig 1, third column
+            gamma = 0.999
+            cost = v1v1
+            return gamma, cost
+        if v1v2 >= v2v2:
+            # Case: Fig 1, first column
+            gamma = 0.001
+            cost = v2v2
+            return gamma, cost
+        # Case: Fig 1, second column
+        gamma = -1.0 * ((v1v2 - v2v2) / (v1v1 + v2v2 - 2 * v1v2))
+        cost = v2v2 + gamma * (v1v2 - v2v2)
+        return gamma, cost
+
+    def _min_norm_2d(vecs, dps):
+        R"""
+        Find the minimum norm solution as combination of two points
+        This is correct only in 2D
+        ie. min_c |\sum c_i x_i|_2^2 st. \sum c_i = 1 , 1 >= c_1 >= 0 for all i, c_i + c_j = 1.0 for some i, j
+        """
+        dmin = 1e8
+        for i in range(len(vecs)):
+            for j in range(i + 1, len(vecs)):
+                if (i, j) not in dps:
+                    dps[(i, j)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(i, j)] += (
+                            torch.mul(vecs[i][k], vecs[j][k]).sum().data.cpu()
+                        )
+                    dps[(j, i)] = dps[(i, j)]
+                if (i, i) not in dps:
+                    dps[(i, i)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(i, i)] += (
+                            torch.mul(vecs[i][k], vecs[i][k]).sum().data.cpu()
+                        )
+                if (j, j) not in dps:
+                    dps[(j, j)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(j, j)] += (
+                            torch.mul(vecs[j][k], vecs[j][k]).sum().data.cpu()
+                        )
+                c, d = MinNormSolver._min_norm_element_from2(
+                    dps[(i, i)], dps[(i, j)], dps[(j, j)]
+                )
+                if d < dmin:
+                    dmin = d
+                    sol = [(i, j), c, d]
+        return sol, dps
+
+    def _projection2simplex(y):
+        R"""
+        Given y, it solves argmin_z |y-z|_2 st \sum z = 1 , 1 >= z_i >= 0 for all i
+        """
+        m = len(y)
+        sorted_y = np.flip(np.sort(y), axis=0)
+        tmpsum = 0.0
+        tmax_f = (np.sum(y) - 1.0) / m
+        for i in range(m - 1):
+            tmpsum += sorted_y[i]
+            tmax = (tmpsum - 1) / (i + 1.0)
+            if tmax > sorted_y[i + 1]:
+                tmax_f = tmax
+                break
+        return np.maximum(y - tmax_f, np.zeros(y.shape))
+
+    def _next_point(cur_val, grad, n):
+        proj_grad = grad - (np.sum(grad) / n)
+        tm1 = -1.0 * cur_val[proj_grad < 0] / proj_grad[proj_grad < 0]
+        tm2 = (1.0 - cur_val[proj_grad > 0]) / (proj_grad[proj_grad > 0])
+
+        skippers = np_sum(tm1 < 1e-7) + np_sum(tm2 < 1e-7)
+        t = 1
+        if len(tm1[tm1 > 1e-7]) > 0:
+            t = np.min(to_numpy(tm1[tm1 > 1e-7]))
+        if len(tm2[tm2 > 1e-7]) > 0:
+            t = min(t, np.min(to_numpy(tm2[tm2 > 1e-7])))
+
+        next_point = proj_grad * t + to_numpy(cur_val)
+        next_point = MinNormSolver._projection2simplex(next_point)
+        return next_point
+
+    def find_min_norm_element(vecs):
+        R"""
+        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull
+        as min |u|_2 st. u = \sum c_i vecs[i] and \sum c_i = 1.
+        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})
+        Hence, we find the best 2-task solution, and then run the projected gradient descent until convergence
+        """
+        # Solution lying at the combination of two points
+        dps = {}
+        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)
+
+        n = len(vecs)
+        sol_vec = np.zeros(n)
+        sol_vec[init_sol[0][0]] = init_sol[1]
+        sol_vec[init_sol[0][1]] = 1 - init_sol[1]
+
+        if n < 3:
+            # This is optimal for n=2, so return the solution
+            return sol_vec, init_sol[2]
+
+        iter_count = 0
+
+        grad_mat = np.zeros((n, n))
+        for i in range(n):
+            for j in range(n):
+                grad_mat[i, j] = dps[(i, j)]
+
+        while iter_count < MinNormSolver.MAX_ITER:
+            grad_dir = -1.0 * np.dot(grad_mat, sol_vec)
+            new_point = MinNormSolver._next_point(sol_vec, grad_dir, n)
+            # Re-compute the inner products for line search
+            v1v1 = 0.0
+            v1v2 = 0.0
+            v2v2 = 0.0
+            for i in range(n):
+                for j in range(n):
+                    v1v1 += sol_vec[i] * sol_vec[j] * dps[(i, j)]
+                    v1v2 += sol_vec[i] * new_point[j] * dps[(i, j)]
+                    v2v2 += new_point[i] * new_point[j] * dps[(i, j)]
+            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)
+            new_sol_vec = nc * sol_vec + (1 - nc) * new_point
+            change = new_sol_vec - sol_vec
+            if np_sum(np.abs(change)) < MinNormSolver.STOP_CRIT:
+                return sol_vec, nd
+            sol_vec = new_sol_vec
+
+    def find_min_norm_element_FW(vecs):
+        R"""
+        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull
+        as min |u|_2 st. u = \sum c_i vecs[i] and \sum c_i = 1.
+        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})
+        Hence, we find the best 2-task solution, and then run the Frank Wolfe until convergence
+        """
+        # Solution lying at the combination of two points
+        dps = {}
+        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)
+
+        n = len(vecs)
+        sol_vec = np.zeros(n)
+        sol_vec[init_sol[0][0]] = init_sol[1]
+        sol_vec[init_sol[0][1]] = 1 - init_sol[1]
+
+        if n < 3:
+            # This is optimal for n=2, so return the solution
+            return sol_vec, init_sol[2]
+
+        iter_count = 0
+
+        grad_mat = np.zeros((n, n))
+        for i in range(n):
+            for j in range(n):
+                grad_mat[i, j] = dps[(i, j)]
+
+        while iter_count < MinNormSolver.MAX_ITER:
+            t_iter = np.argmin(np.dot(grad_mat, sol_vec))
+
+            v1v1 = np.dot(sol_vec, np.dot(grad_mat, sol_vec))
+            v1v2 = np.dot(sol_vec, grad_mat[:, t_iter])
+            v2v2 = grad_mat[t_iter, t_iter]
+
+            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)
+            new_sol_vec = nc * sol_vec
+            new_sol_vec[t_iter] += 1 - nc
+
+            change = new_sol_vec - sol_vec
+            if np_sum(np.abs(change)) < MinNormSolver.STOP_CRIT:
+                return sol_vec, nd
+            sol_vec = new_sol_vec
+
+
+def gradient_normalizers(grads, losses, normalization_type):
+    gn = {}
+    if normalization_type == "l2":
+        for t in grads:
+            gn[t] = np.sqrt(np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]]))
+    elif normalization_type == "loss":
+        for t in grads:
+            gn[t] = losses[t]
+    elif normalization_type == "loss+":
+        for t in grads:
+            gn[t] = losses[t] * np.sqrt(
+                np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]])
+            )
+    elif normalization_type == "none":
+        for t in grads:
+            gn[t] = 1.0
+    else:
+        print("ERROR: Invalid Normalization Type")
+    return gn