[fx] Add concrete info prop

[Cypher30]

What's New?

In this PR, I provide ConcreteInfoProp to facilitate the meta info estimation, we also provide the time profiler that track the running time of each node. You could check the docstring of function _profile_concrete for more information.

For example, now you could run

import torch
from torch.fx import symbolic_trace
from colossalai.fx.passes import ConcreteInfoProp

BATCH_SIZE = 2
DIM_IN = 4
DIM_HIDDEN = 16
DIM_OUT = 16
model = torch.nn.Sequential(
    torch.nn.Linear(DIM_IN, DIM_HIDDEN), 
    torch.nn.Linear(DIM_HIDDEN, DIM_OUT),
    ).cuda()
input_sample = torch.rand(BATCH_SIZE, DIM_IN, device="cuda")
gm = symbolic_trace(model)
interp = ConcreteInfoProp(gm)
interp.run(input_sample)
print(interp.summary(unit='kb')) 

and get the following results

 Op type       Op             Forward time             Backward time    SAVE_FWD_IN    FWD_OUT    FWD_TMP    BWD_OUT    BWD_TMP
-----------  -------  -----------------------  ------------------------  -------------  ---------  ---------  ---------  ---------
placeholder  input_1                    0.0 s                     0.0 s          False    0.00 KB    0.00 KB    0.00 KB    0.00 KB
call_module       _0  0.0003993511199951172 s     0.00706791877746582 s          False    0.50 KB    0.00 KB    0.03 KB    0.66 KB
call_module       _1   6.29425048828125e-05 s  0.00018286705017089844 s          False    0.50 KB    0.00 KB    0.12 KB    0.81 KB
     output   output                    0.0 s                     0.0 s           True    0.00 KB    0.00 KB    0.00 KB    0.00 KB

[Cypher30]

Here is the shortcut for you guys to check what I'm doing!!

[Cypher30]

this part is for retrieving param size in args and kwargs