def native_batch_norm(
input: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
running_mean: Optional[Tensor],
running_var: Optional[Tensor],
training: bool,
momentum: float,
eps: float,
) -> Tuple[Tensor, Tensor, Tensor]:
reduction_dims = [0] + list(range(2, input.dim()))
computation_dtype = utils.get_computation_dtype(input.dtype)
if training:
output, mean, rstd = normalize(input, reduction_dims, eps)
save_mean = _squeeze_multiple(mean, reduction_dims)
save_rstd = _squeeze_multiple(rstd, reduction_dims)
if running_mean is not None:
running_mean.copy_(momentum * save_mean +
(1 - momentum) * running_mean)
if running_var is not None:
n = input.numel() / input.shape[1]
# This doesn't strictly match eager's numerics, which accumulates var sum and then directly applies the correction
# But... that would require re-implementing var here, for negligible numerics gain on a tensor whose
# numerics probably don't matter.
unbiased_var = torch.var(input, reduction_dims,
unbiased=False) * (n / (n - 1))
running_var.copy_(momentum * unbiased_var +
(1 - momentum) * running_var)
else:
assert running_mean is not None and running_var is not None
running_mean = running_mean.to(dtype=computation_dtype)
running_var = running_var.to(dtype=computation_dtype)
mean = running_mean
invstd = 1 / (torch.sqrt(running_var + eps))
# Very annoying inconsistency where CPU and CUDA give different shapes
if input.device.type != "cpu":
save_mean = running_mean
save_rstd = invstd
else:
save_mean = input.new_zeros((0, ))
save_rstd = input.new_zeros((0, ))
mean = _unsqueeze_to_dim(mean, input.dim() - 1)
invstd = _unsqueeze_to_dim(invstd, input.dim() - 1)
output = (input - mean) * invstd
if weight is None:
weight = input.new_ones(())
if bias is None:
bias = input.new_zeros(())
weight = _unsqueeze_to_dim(weight, input.dim() - 1)
bias = _unsqueeze_to_dim(bias, input.dim() - 1)
output = output * weight + bias
if input.device.type == "cpu":
save_mean = save_mean.to(dtype=input.dtype)
save_rstd = save_rstd.to(dtype=input.dtype)
return output.to(dtype=input.dtype), save_mean, save_rstd
def normalize(input, norm_dims, eps):
computation_dtype = utils.get_computation_dtype(input.dtype)
input_acc = input.to(dtype=computation_dtype)
biased_var = torch.var(input_acc,
dim=norm_dims,
unbiased=False,
keepdim=True)
mean = torch.mean(input_acc, dim=norm_dims, keepdim=True)
rstd = torch.rsqrt(biased_var + eps)
out = (input - mean) * rstd
return out, mean, rstd
def native_layer_norm_backward(
grad_out: Tensor,
input: Tensor,
normalized_shape: List[int],
mean: Tensor,
rstd: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
output_mask: List[bool],
) -> Tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
input_shape = input.shape
input_ndim = input.dim()
computation_dtype = utils.get_computation_dtype(input.dtype)
grad_out_cast, input_cast, weight_cast, bias_cast = [
x.to(computation_dtype) if x is not None else x
for x in (grad_out, input, weight, bias)
]
assert grad_out_cast is not None
axis = input_ndim - len(normalized_shape)
inner_dims = input_shape[axis:]
outer_dims = input_shape[:axis]
inner_dim_indices: List[int] = []
outer_dim_indices: List[int] = []
for i in range(input_ndim):
if i >= axis:
inner_dim_indices.append(i)
else:
outer_dim_indices.append(i)
N = prod(inner_dims) # type: ignore[arg-type]
M = prod(outer_dims) # type: ignore[arg-type]
if M <= 0 or N <= 0:
return (
input.new_zeros(input_shape),
input.new_zeros(input_shape[axis:]),
input.new_zeros(input_shape[axis:]),
)
x_hat = (input_cast - mean) * rstd
if weight_cast is not None:
grad_x_hat = grad_out_cast * weight_cast
else:
grad_x_hat = grad_out_cast
a = grad_x_hat * N
b = torch.sum(grad_x_hat, inner_dim_indices, True)
c1 = torch.mul(grad_x_hat, x_hat)
c2 = torch.sum(c1, inner_dim_indices, True)
c3 = torch.mul(x_hat, c2)
inner = a - b - c3
d_input: Optional[Tensor] = None
d_weight: Optional[Tensor] = None
d_bias: Optional[Tensor] = None
if output_mask[0]:
d_input = (rstd / N) * inner
if output_mask[1] and weight_cast is not None:
if len(outer_dim_indices) > 0:
d_weight = torch.sum(grad_out_cast * x_hat, outer_dim_indices,
False)
else:
d_weight = grad_out_cast * x_hat
if output_mask[2] and bias_cast is not None:
if len(outer_dim_indices) > 0:
d_bias = torch.sum(grad_out_cast, outer_dim_indices, False)
else:
d_bias = grad_out_cast
return (
_maybe_cast(d_input, input.dtype),
_maybe_cast(d_weight, input.dtype),
_maybe_cast(d_bias, input.dtype),
)
def native_batch_norm_backward(
grad_out: Tensor,
input: Tensor,
weight: Optional[Tensor],
running_mean: Optional[Tensor],
running_var: Optional[Tensor],
save_mean: Optional[Tensor],
save_invstd: Optional[Tensor],
train: bool,
eps: float,
output_mask: List[bool],
) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
input_dtype = input.dtype
computation_dtype = utils.get_computation_dtype(input.dtype)
(
grad_out_cast,
input_cast,
weight_cast,
running_mean_cast,
running_var_cast,
save_mean_cast,
save_invstd_cast,
) = [
x.to(computation_dtype) if x is not None else x for x in (
grad_out,
input,
weight,
running_mean,
running_var,
save_mean,
save_invstd,
)
]
input_shape = input.shape
input_rank = input.dim()
assert input_rank >= 2, "rank of the input must be at least 2"
axis = 1
num_features = prod(list(input_shape)) / input_shape[axis]
mean = save_mean_cast
invstd = save_invstd_cast
if train:
assert save_mean_cast is not None and save_invstd_cast is not None
else:
assert running_mean_cast is not None and running_var_cast is not None
mean = running_mean_cast
invstd = torch.rsqrt(running_var_cast + eps)
broadcast_mask: List[int] = [1] * input_rank
broadcast_mask[axis] = input_shape[axis]
reduction_axes: List[int] = []
for i in range(input_rank):
if i != axis:
reduction_axes.append(i)
mean = torch.reshape(mean, broadcast_mask) # type: ignore[arg-type]
norm = 1.0 / num_features
grad_output_sum = torch.sum(grad_out_cast,
reduction_axes) # type: ignore[arg-type]
dot_p = torch.sum(grad_out_cast * (input_cast - mean), reduction_axes)
grad_mean = torch.reshape(grad_output_sum * norm, broadcast_mask)
proj_scale = torch.reshape(torch.mul(dot_p * norm, invstd * invstd),
broadcast_mask) # type: ignore[operator]
if weight_cast is None:
grad_scale = torch.reshape(
invstd, broadcast_mask) * 1.0 # type: ignore[arg-type]
else:
grad_scale = torch.reshape(invstd * weight_cast, broadcast_mask)
if train:
proj = (input_cast - mean) * proj_scale
grad_input = ((grad_out_cast - proj) - grad_mean) * grad_scale
else:
grad_input = grad_out_cast * grad_scale
if output_mask[1]:
grad_weight = dot_p * invstd
else:
grad_weight = None # "None" doesn't work with vjp, should use zeros for vjp
if output_mask[2]:
grad_bias = grad_output_sum
else:
grad_bias = None # "None" doesn't work with vjp, should use zeros for vjp
return (
grad_input.to(input_dtype),
_maybe_cast(grad_weight, input_dtype),
_maybe_cast(grad_bias, input_dtype),
)
