mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-28 21:17:08 +00:00
1e88811c7a
* [autoparallel] first move. * [autoparallel] add solver rotor. * [autoparallel] add ckpt solvers. * [autoparallel] modify codegen. * [fx] fix annotation in test. * [fx] remove check. * [autoparallel] polish docstring. * [fx] refactor MetaTensor.
64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.fx import GraphModule
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
from colossalai.fx import ColoTracer
|
|
|
|
|
|
class MLP(torch.nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear1 = torch.nn.Linear(4, 4)
|
|
self.linear2 = torch.nn.Linear(4, 4)
|
|
|
|
def forward(self, x):
|
|
x = self.linear1(x)
|
|
x = self.linear2(x)
|
|
return x
|
|
|
|
|
|
# Simple module for demonstration
|
|
class MyModule(torch.nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.mlp_1 = MLP()
|
|
self.mlp_2 = MLP()
|
|
self.output = torch.nn.Linear(4, 4)
|
|
|
|
def forward(self, x):
|
|
x = checkpoint(self.mlp_1, x)
|
|
x = checkpoint(self.mlp_2, x)
|
|
x = self.output(x)
|
|
return x
|
|
|
|
|
|
def test_activation_checkpoint_annotation():
|
|
module = MyModule()
|
|
|
|
# test tracing with activation checkpoint
|
|
tracer = ColoTracer(trace_act_ckpt=True)
|
|
graph = tracer.trace(module)
|
|
gm = GraphModule(module, graph)
|
|
|
|
for node in gm.graph.nodes:
|
|
if node.name in ['mlp_1_linear1', 'mlp_1_linear2']:
|
|
assert node.meta.get('activation_checkpoint', -1) == 0
|
|
|
|
for node in gm.graph.nodes:
|
|
if node.name in ['mlp_2_linear1', 'mlp_2_linear2']:
|
|
assert node.meta.get('activation_checkpoint', -1) == 1
|
|
|
|
tracer = ColoTracer(trace_act_ckpt=False)
|
|
graph = tracer.trace(module)
|
|
gm = GraphModule(module, graph)
|
|
|
|
for node in gm.graph.nodes:
|
|
assert not hasattr(node, 'activation_checkpoint')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_activation_checkpoint_annotation()
|