mirror of https://github.com/hpcaitech/ColossalAI
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
import torch.nn as nn
|
|
import torch
|
|
from colossalai.auto_parallel.solver.graph_analysis import GraphAnalyser
|
|
from colossalai.fx import ColoTracer, ColoGraphModule
|
|
|
|
|
|
class LinearModel(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = nn.Linear(4, 4)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.linear2 = nn.Linear(4, 4)
|
|
|
|
def forward(self, x1, x2):
|
|
x1 = x1 * 2
|
|
x1 = self.linear1(x1)
|
|
x1 = self.relu(x1)
|
|
x1 = self.linear2(x1)
|
|
out = x1 + x2
|
|
return out
|
|
|
|
|
|
def test_liveness_analysis():
|
|
model = LinearModel()
|
|
tracer = ColoTracer()
|
|
graph = tracer.trace(model,
|
|
meta_args={
|
|
'x1': torch.rand(4, 4, device='meta'),
|
|
'x2': torch.rand(4, 4, device='meta')
|
|
})
|
|
gm = ColoGraphModule(root=model, graph=graph, class_name=model.__class__.__name__)
|
|
|
|
graph_analyser = GraphAnalyser(gm)
|
|
liveness_list = graph_analyser.liveness_analysis()
|
|
stage_count = len(liveness_list)
|
|
|
|
# if a LiveStage is covered by another LiveStage, we just keep the larger one.
|
|
assert stage_count == 1
|
|
|
|
# a variable named `relu` must exist
|
|
# and this live var must have inplace = True
|
|
assert liveness_list[0].all_live_vars.exists('relu')
|
|
relu_var = liveness_list[0].all_live_vars.get('relu')
|
|
assert relu_var.is_inplace
|
|
|
|
# the unique vars must be fewer than the all vars since in-place ops exist
|
|
all_live_vars = liveness_list[0].all_live_vars
|
|
unique_live_vars = liveness_list[0].unique_live_vars
|
|
assert len(unique_live_vars) + 1 == len(all_live_vars)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_liveness_analysis()
|