mirror of https://github.com/hpcaitech/ColossalAI
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.fx import GraphModule
|
|
|
|
from colossalai.fx import ColoTracer as Tracer
|
|
from colossalai.testing import clear_cache_before_run
|
|
|
|
|
|
class ControlFlowModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = nn.Linear(10, 10)
|
|
self.linear2 = nn.Linear(10, 10)
|
|
|
|
def forward(self, x, y):
|
|
x1 = self.linear1(x)
|
|
y1 = self.linear2(y)
|
|
|
|
if x1.dim() == 2:
|
|
return x1 + y1
|
|
else:
|
|
return x1 - y1
|
|
|
|
|
|
@clear_cache_before_run()
|
|
def test_control_flow():
|
|
model = ControlFlowModel()
|
|
tracer = Tracer()
|
|
graph_branch_true = tracer.trace(
|
|
model, meta_args={"x": torch.rand(4, 10, device="meta"), "y": torch.rand(4, 10, device="meta")}
|
|
)
|
|
graph_branch_false = tracer.trace(
|
|
model, meta_args={"x": torch.rand(10, device="meta"), "y": torch.rand(4, 10, device="meta")}
|
|
)
|
|
|
|
gm_branch_true = GraphModule(model, graph_branch_true, model.__class__.__name__)
|
|
gm_branch_false = GraphModule(model, graph_branch_false, model.__class__.__name__)
|
|
gm_branch_true.recompile()
|
|
gm_branch_false.recompile()
|
|
|
|
# test the true branch
|
|
x = torch.rand(4, 10)
|
|
y = torch.rand(4, 10)
|
|
assert torch.all(model(x, y) == gm_branch_true(x, y))
|
|
assert torch.all(gm_branch_false(x, y) != gm_branch_true(x, y))
|
|
|
|
# test the true branch
|
|
x = torch.rand(10)
|
|
y = torch.rand(4, 10)
|
|
assert torch.all(model(x, y) == gm_branch_false(x, y))
|
|
assert torch.all(gm_branch_false(x, y) != gm_branch_true(x, y))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_control_flow()
|