mirror of https://github.com/hpcaitech/ColossalAI
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
|
import operator
|
||
|
|
import torch
|
||
|
|
from torch.fx.proxy import Proxy, Attribute
|
||
|
|
|
||
|
|
__all__ = ['ColoProxy']
|
||
|
|
|
||
|
|
|
||
|
|
class ColoProxy(Proxy):
|
||
|
|
"""
|
||
|
|
ColoProxy is a proxy class which uses meta tensor to handle data-dependent control flow. The original torch.fx proxy
|
||
|
|
cannot be used to infer the condition statement, with this proxy, torch.fx can still run even with if statements.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
proxy = tracer.create_proxy(...)
|
||
|
|
proxy.meta_tensor = torch.empty(4, 2, device='meta')
|
||
|
|
print(len(proxy)) # expect output 4
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
self.meta_tensor = None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def meta_tensor(self):
|
||
|
|
return self.meta_tensor
|
||
|
|
|
||
|
|
@meta_tensor.setter
|
||
|
|
def meta_tensor(self, tensor: torch.Tensor):
|
||
|
|
assert tensor.is_meta, 'Expected to receive a meta tensor, but got a non-meta tensor'
|
||
|
|
self.meta_tensor = tensor
|
||
|
|
|
||
|
|
@property
|
||
|
|
def has_meta_tensor(self):
|
||
|
|
return self.meta_tensor is not None
|
||
|
|
|
||
|
|
def _assert_has_meta(self):
|
||
|
|
assert self.has_meta_tensor, f'Meta tensor is not set for {self.node.name}'
|
||
|
|
|
||
|
|
@property
|
||
|
|
def dtype(self):
|
||
|
|
self._assert_has_meta()
|
||
|
|
return self.meta_tensor.dtype
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
self._assert_has_meta()
|
||
|
|
return len(self.meta_tensor)
|
||
|
|
|
||
|
|
def __bool__(self):
|
||
|
|
self._assert_has_meta()
|
||
|
|
return self.meta_tensor
|
||
|
|
|
||
|
|
def __getattr__(self, k):
|
||
|
|
if k == "metadata":
|
||
|
|
return self.meta_tensor
|
||
|
|
# note: not added to the graph yet, if this is a method call
|
||
|
|
# we peephole optimize to the method invocation
|
||
|
|
return Attribute(self, k)
|
||
|
|
|
||
|
|
def __setitem__(self, indices, values):
|
||
|
|
return self.tracer.create_proxy("call_function", operator.setitem, (self, indices, values), {})
|