github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] Add pofo sequence annotation (#1637)

Browse Source 
* [autoparallel] annotate pofo sequence

* [autoparallel] remove unused print

* [autoparallel] fix some code
This commit is contained in:
Boyuan Yao
2022-09-24 01:52:57 +08:00
committed by GitHub
parent 04bbabeea8
commit f921733621
3 changed files with 133 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -54,7 +54,8 @@ class Offload(Operation):
super().__init__()
self.index = index
self.name = "Off"
if has_bar:
self.has_bar = has_bar
if self.has_bar:
self.name += "wBar"
def __repr__(self):
@@ -67,7 +68,8 @@ class Prefetch(Operation):
super().__init__()
self.index = index
self.name = "Pre"
if has_bar:
self.has_bar = has_bar
if self.has_bar:
self.name += "wBar"
def __repr__(self):