Method compute_eigenvalue

deepspeed/runtime/eigenvalue.py:71–147

(self, module, device=None, scale=1.0)

        return m

    def compute_eigenvalue(self, module, device=None, scale=1.0):
        block_eigenvalue = []
        param_keys = []
        layers = self.get_layers(module)

        for block in range(self.layer_num):
            model_block = layers[block]

            # We found this randn() has obvious accuracy impact in some cases, save/recover random state here.
            rng_state = torch.random.get_rng_state()
            if device is None:
                v = [
                    torch.randn(p.size()) for p in model_block.parameters()
                    if p.grad is not None and p.grad.grad_fn is not None
                ]
            else:
                v = [
                    torch.randn(p.size(), device=device) for p in model_block.parameters()
                    if p.grad is not None and p.grad.grad_fn is not None
                ]
            torch.random.set_rng_state(rng_state)

            grads = [
                param.grad for param in model_block.parameters()
                if param.grad is not None and param.grad.grad_fn is not None
            ]
            params = [
                param for param in model_block.parameters()
                if param.grad is not None and param.grad.grad_fn is not None
            ]

            layer_keys = [id(p) for p in model_block.parameters()]
            param_keys.append(layer_keys)

            v = self.normalize(v)

            # Disable eigenvalue if the model doesn't support second order gradients computation,
            # e.g. when enabling DS transformer kernel.
            if len(grads) == 0 or len(params) == 0:
                log_dist('The model does NOT support eigenvalue computation.', ranks=[0], level=logging.WARNING)
                return []

            i = 0
            eigenvalue_current, eigenvalue_previous = 1., 0.

            while (i < self.max_iter) and abs(eigenvalue_current) > 0 and (abs(
                (eigenvalue_current - eigenvalue_previous) / eigenvalue_current)
                                                                           >= self.tol):  # test convergence criteria
                eigenvalue_previous = eigenvalue_current

                Hv = torch.autograd.grad(grads, params, grad_outputs=v, only_inputs=True, retain_graph=True)
                #Hv = [hv.float() for hv in Hv]
                Hv = [self.nan_to_num(hv).float() for hv in Hv]

                eigenvalue_current = self.inner_product(Hv, v).item()

                v = self.normalize(Hv)
                v = [x / scale for x in v]