github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[fx] fix wrong variable name in solver rotor (#1502) 

* [fx] fix wrong variable name in solver rotor

* [fx] fix wrong variable name in solver rotor

* code modification
pull/1506/head^2
Boyuan Yao 2022-08-26 15:47:08 +08:00 committed by GitHub
parent 3b6a5e2593
commit 31fffd3fc5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/fx/passes/algorithms/ckpt_solver_rotor.py
View File

@ -73,7 +73,7 @@ def _compute_table(chain: Chain, mmax) -> Tuple:
return (opt, what)
def _rec(chain, lmin, lmax, cmem, opt_table):
def _rec(chain: Chain, lmin, lmax, cmem, opt_table):
""" chain : the class describing the AC graph
lmin : index of the first forward to execute
lmax : upper bound index of the last forward to execute (not included)
@ -97,14 +97,14 @@ def _rec(chain, lmin, lmax, cmem, opt_table):
if what[cmem][lmin][lmax][0]:
sequence.insert(ForwardEnable(lmin))
sequence.insert_sequence(_rec(chain, lmin + 1, lmax, cmem - chain.cbweigth[lmin + 1], opt_table))
sequence.insert_sequence(_rec(chain, lmin + 1, lmax, cmem - chain.cbweight[lmin + 1], opt_table))
sequence.insert(Backward(lmin))
else:
j = what[cmem][lmin][lmax][1]
sequence.insert(ForwardCheck(lmin))
for k in range(lmin + 1, j):
sequence.insert(ForwardNograd(k))
sequence.insert_sequence(_rec(chain, j, lmax, cmem - chain.cweigth[j], opt_table))
sequence.insert_sequence(_rec(chain, j, lmax, cmem - chain.cweight[j], opt_table))
sequence.insert_sequence(_rec(chain, lmin, j - 1, cmem, opt_table))
return sequence

14
colossalai/fx/passes/algorithms/utils.py
View File

@ -44,9 +44,9 @@ class Forward(Operation):
def __repr__(self):
return "{n}_{i}".format(n=self.name, i=self.index)
def cost(self, chain):
def cost(self, chain: Chain):
if chain is not None:
return chain.fweigth[self.index]
return chain.fweight[self.index]
else:
return 1
@ -80,9 +80,9 @@ class Forwards(Operation):
def __repr__(self):
return "F_{i}->{j}".format(i=self.index[0], j=self.index[1])
def cost(self, chain):
def cost(self, chain: Chain):
if chain is not None:
return sum(chain.fweigth[self.index[0]:self.index[1] + 1])
return sum(chain.fweight[self.index[0]:self.index[1] + 1])
else:
return (self.index[1] - self.index[0] + 1)
@ -99,9 +99,9 @@ class Backward(Operation):
def __repr__(self):
return "B_{i}".format(i=self.index)
def cost(self, chain):
def cost(self, chain: Chain):
if chain is not None:
return chain.bweigth[self.index]
return chain.bweight[self.index]
else:
return 1
@ -126,7 +126,7 @@ class MemoryAccess(Operation):
def __repr__(self):
return "{n}_{i}".format(n=self.name, i=self.index)
def cost(self, chain):
def cost(self, chain: Chain):
return 0