MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / Eigenvalue

Class Eigenvalue

deepspeed/runtime/eigenvalue.py:13–153  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

11
12
13class Eigenvalue(object):
14
15 def __init__(self,
16 verbose=False,
17 max_iter=100,
18 tol=1e-2,
19 stability=0,
20 gas_boundary_resolution=1,
21 layer_name='',
22 layer_num=0):
23 super().__init__()
24
25 self.verbose = verbose
26 self.max_iter = max_iter
27 self.tol = tol
28 self.stability = stability
29 self.gas_boundary_resolution = gas_boundary_resolution
30 self.layer_name = layer_name
31 self.layer_num = layer_num
32
33 assert len(self.layer_name) > 0 and layer_num > 0
34
35 log_dist(
36 f'enabled eigenvalue with verbose={verbose}, max_iter={max_iter}, tol={tol}, stability={stability}, gas_boundary_resolution={gas_boundary_resolution}, layer_name={layer_name}, layer_num={layer_num}',
37 ranks=[0])
38
39 # Replace all nan/pos-inf/neg-inf to zero
40 def nan_to_num(self, x):
41 if required_torch_version(min_version=1.8):
42 return torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
43 else:
44 # Fallback to numpy based implementation for backwards-compatibility with PyTorch 1.7 or older versions.
45 device = x.device
46 x = x.cpu().numpy()
47 x = np.nan_to_num(x=x, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
48 return torch.from_numpy(x).to(device)
49
50 def normalize(self, v):
51 norm_squared = self.inner_product(v, v)
52 norm = norm_squared**0.5 + self.stability
53 normalized_vectors = [vector / norm for vector in v]
54 normalized_vectors = [self.nan_to_num(vector) for vector in normalized_vectors]
55 return normalized_vectors
56
57 def inner_product(self, xs, ys):
58 return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)])
59
60 def get_layers(self, module):
61 scope_names = self.layer_name.split('.')
62 assert len(scope_names) > 0
63
64 m = module
65 for name in scope_names:
66 assert hasattr(m, name), "layer_name configuration is invalid."
67 m = getattr(m, name)
68
69 return m
70

Callers 1

_configure_eigenvalueMethod · 0.90

Calls

no outgoing calls

Tested by

no test coverage detected