make backward function explicit in a layer which is a combination of some ops

[guoyejun]

🚀 The feature, motivation and pitch

To get a top-down performance analysis, we'd like to know how the time is consumed in different components of the model. Currently, we are able to know the forward time of the layers, but have no separated time of the bwd time of different layers.

For example, I have a layer (MyLayer) which is the combination of some pytorch standard OPs, see blow. I can add something (for example, record_function, or events at the beginning/end) in the forward function, and I can get the time of forward function. But, I'm unable to know how much time is consumed by the bwd of MyLayer. I can get the OP level profiling from pytorch profiler, but it is not easy to collect them back into MyLayer (many reasons, for example, there are other layers fwd+bwd may call the same OPs).

class MyLayer(torch.nn.Module):
    def __init__(self):
        super(MyLayer, self).__init__()
        self.mlp_up_proj = nn.Linear(args.hs, 4 * args.hs)
        self.mlp_act = nn.GELU(approximate='none')
        self.mlp_down_proj = nn.Linear(4*args.hs, args.hs)
        
    def forward(self, x):
        y1 = self.mlp_up_proj(x)
        y2 = self.mlp_act(y1)
        y3 = self.mlp_down_proj(y2)
        return y3

So, I tried to change MyLayer as below to show my idea. I know it does not work, it even crashes in the runtime. If we can modify the original MyLayer with least change to make the backward function explicit, we can do all kinds of profiling on it.

# it does not work
class MyLayerImpl(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, mlp_up_proj, mlp_act, mlp_down_proj):
        with torch.enable_grad():
            y1 = mlp_up_proj(x)
            y2 = mlp_act(y1)
            y3 = mlp_down_proj(y2)
        y4 = y3.detach()
        ctx.save_for_backward(x, y1, y2, y3)
        return y4

    @staticmethod
    def backward(ctx, dout, *args):
        x, y1, y2, y3 = ctx.saved_tensors
        din = y3.backward(dout)
        return din, None, None, None
    

class MyLayer(torch.nn.Module):
    def __init__(self):
        super(MyLayer, self).__init__()
        self.mlp_up_proj = nn.Linear(args.hs, 4 * args.hs)
        self.mlp_act = nn.GELU(approximate='none')
        self.mlp_down_proj = nn.Linear(4*args.hs, args.hs)
        
    def forward(self, x):
        return MyLayerImpl.apply(x, self.mlp_up_proj, self.mlp_act, self.mlp_down_proj)

Alternatives

No response

Additional context

No response

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @robieta @chaekit @aaronenyeshi @nbcsm @guotuofeng @guyang3532 @gaoteng-git @tiffzhaofb @dzhulgakov @davidberard98

[davidberard98]

@guoyejun is profile(with_stack=True, ...) sufficient for this?

[guoyejun]

@davidberard98 thanks for your reply, but I think profile(with_stack=True, ...) is not sufficient.

From the profiler, we can see that all the backwards under one item torch/autograd/_init__.py(149): backward, while the requirement is that we can know clearly/explicitly the backwards of MyLayer.

Take the below code as example, I can see all the forwards of MyLayer (labeled as "xxxxmylayer:forward"), but it is not easy to find the backwards of MyLayer, what we see is all the backwards for the whole model.

import torch
import argparse
from torch import nn
import torch.nn.functional as F
from torch.profiler import profile, record_function, ProfilerActivity

torch.manual_seed(0)

parser = argparse.ArgumentParser(
                    prog='ProgramName',
                    description='What the program does',
                    epilog='Text at the bottom of help')
parser.add_argument('--mbs', type=int, default=2)
parser.add_argument('--seq', type=int, default=8)
parser.add_argument('--hs', type=int, default=4)
parser.add_argument('--device', type=str, default='cuda')
args = parser.parse_args()


for k, v in vars(args).items():
    print(k, '=', v)

    
class MyLayer(torch.nn.Module):
    def __init__(self):
        super(MyLayer, self).__init__()
        self.mlp_up_proj = nn.Linear(args.hs, 4 * args.hs)
        self.mlp_down_proj = nn.Linear(4*args.hs, args.hs)
        
    def forward(self, x):
        with record_function("xxxxmylayer:forward"):
            y = self.mlp_up_proj(x)
            y = self.mlp_down_proj(y)
            return y


class TestNet(torch.nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.Wqkv = nn.Linear(args.hs, 3 * args.hs)
        self.out_proj = nn.Linear(args.hs, args.hs)
        self.my_layer1 = MyLayer()
        self.my_layer2 = MyLayer()
        
    def forward(self, x):
        qkv = self.Wqkv(x)
        query, key, value = qkv.chunk(3, dim=2)
        new_v = query + key + value
        proj = self.out_proj(new_v)
        proj = self.my_layer1(proj)
        proj = self.my_layer2(proj)
        return proj


def main():
    model = TestNet()
    model.bfloat16().to(args.device)

    x = torch.randn([args.mbs, args.seq, args.hs]).bfloat16().to(args.device)
    x.requires_grad_()
    
    # warm up
    for i in range(3):
        y = model(x)
        loss = y.sum()
        loss.backward()
    
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], 
                 with_stack=True,
                 on_trace_ready=torch.profiler.tensorboard_trace_handler('./timeline', worker_name='worker0')) as prof:
        y = model(x)
        loss = y.sum()
        loss.backward()

    
if __name__ == "__main__":
    main()

It would be nice if we can redefine MyLayer as below, the below code does not work, just to show my idea of the clear&explicit backward of MyLayer.

class MyLayer(torch.nn.NewModule):
    def __init__(self):
        super(MyLayer, self).__init__()
        self.mlp_up_proj = nn.Linear(args.hs, 4 * args.hs)
        self.mlp_down_proj = nn.Linear(4*args.hs, args.hs)
        self.mlp_act = nn.GELU(approximate='none')
        
    def forward(self, x):
        with record_function("xxxxmylayer:forward"):
            y1 = self.mlp_up_proj(x)
            y2 = self.mlp_act(y1)
            y3 = self.mlp_down_proj(y2)
            return y3

    def backward(self, dout, ...)
        with record_function("xxxxmylayer:backward"):
            din = super().backward(dout, ...)
            return din

[sraikund16]

@guoyejun In the profiler output, we generate fwdbwd links between each forward op and backward. So you should be able to get all your fwd ops using record_function to wrap around them and then after you can use the fwdbwd links to get the corresponding fwd ops

[sraikund16]

closing as there is a workaround, can re-open if necessary

[Da1sypetals]

This is important when a model is very very big that I cannot even load the trace of a whole model into the profiler ( it crashes) , and I want to start/stop the profiler inside the backward function. There should be a more elegant way to do so.