Fix usage of return_bias argument

[ksivaman]

No description provided.

[timmoon10]

I'm not sure about this API. I can imagine use-cases with use_bias=False and return_bias=True:

defdo_linear(x, use_bias):
layer=Linear(bias=use_bias, return_bias=True)
y, b=layer(x)
ifbisnotNone:
y+=breturny

This seems more Pythonic than the proposed API:

defdo_linear(x, use_bias):
layer=Linear(bias=use_bias, return_bias=True)
outs=layer(x)
ifuse_bias:
returnouts[0] +outs[1]
else:
returnouts

Pinging @sudhakarsingh27.

[ksivaman]

Re: I can imagine use-cases with use_bias=False and return_bias=True

For this use case, the user can set both of those arguments to True and then the model will return the additive bias from forward without doing the actual addition. The point of this PR is that for this use case, we don't want return_bias to override the bias argument.

[sudhakarsingh27]

Re: I can imagine use-cases with use_bias=False and return_bias=True

I agree, that the first way seems more pythonic and we can only do that if we make sure the return tuple is fixed in length from the concerned API. In addition to out and (out, bias), the concerned API (LayerNormLinear) can also emit (out, bias, ln_out) and so we'd probably need to fix 3-tuple as the output. If we fix 3 tuple, then a codebase that uses this API (e.g. megatron) should change too since it could expect 2-tuple output from original ColumnParallelLinear layer. So, although approach 1 seems more pythonic, it still wouldn't make things more general (to be able to swap 2 almost equivalent layers with varying length n-tuples)

For this use case, the user can set both of those arguments to True and then the model will return the additive bias from forward without doing the actual addition. The point of this PR is that for this use case, we don't want return_bias to override the bias argument.

Right, this PR probably isn't the right place to discuss this tangent, let's probably create a new PR.

[timmoon10]

My preferred behavior:

bias	return_bias	Creates bias tensor	Applies bias	Returns
True	True	Yes	No	out, bias
True	False	Yes	Yes	out
False	True	No	No	out, None
False	False	No	No	out

It seems our contention is about the bias=False, return_bias=True case. The current implementation is problematic since it creates the bias tensor, but this PR changes the return behavior.

[timmoon10]

LGTM

[ksivaman]

/te-ci