From ca37aa13048ffe5c5c85e96b7358825533622a1f Mon Sep 17 00:00:00 2001 From: Aditya Ranjan Date: Wed, 6 Mar 2024 10:57:32 -0800 Subject: [PATCH 1/2] initialize grad_input to None --- axonn/intra_layer/fully_connected.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/axonn/intra_layer/fully_connected.py b/axonn/intra_layer/fully_connected.py index 6d2cfae..718f747 100644 --- a/axonn/intra_layer/fully_connected.py +++ b/axonn/intra_layer/fully_connected.py @@ -106,6 +106,8 @@ def backward(ctx, grad_output): if dist.get_world_size(ctx.backward_all_reduce_group) > 1 or ( not overlap_reduce_scatter ): + grad_input=None + if ctx.needs_input_grad[0]: grad_input = grad_output.matmul(weight) handle = dist.all_reduce( @@ -135,6 +137,8 @@ def backward(ctx, grad_output): grad_weight = None # weight gradients are not ready yet return grad_input, grad_weight, None, None, None, None, None, None, None else: + grad_input=None + if ctx.needs_input_grad[1]: grad_weight = ( grad_output.reshape(-1, grad_output.shape[-1]) From a4840086680d51c2317db3b2914b4924cb8b6725 Mon Sep 17 00:00:00 2001 From: Aditya Ranjan Date: Wed, 6 Mar 2024 11:10:49 -0800 Subject: [PATCH 2/2] minor --- axonn/intra_layer/fully_connected.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/axonn/intra_layer/fully_connected.py b/axonn/intra_layer/fully_connected.py index 718f747..5b75346 100644 --- a/axonn/intra_layer/fully_connected.py +++ b/axonn/intra_layer/fully_connected.py @@ -106,7 +106,7 @@ def backward(ctx, grad_output): if dist.get_world_size(ctx.backward_all_reduce_group) > 1 or ( not overlap_reduce_scatter ): - grad_input=None + grad_input, grad_weight = None, None if ctx.needs_input_grad[0]: grad_input = grad_output.matmul(weight) @@ -137,7 +137,7 @@ def backward(ctx, grad_output): grad_weight = None # weight gradients are not ready yet return grad_input, grad_weight, None, None, None, None, None, None, None else: - grad_input=None + grad_input, grad_weight = None, None if ctx.needs_input_grad[1]: grad_weight = (