Repository: pytorch
Issue: Extended functionality for torch.quantization.fuse_modules
Labels: oncall: quantization,triaged

### 🚀 The feature, motivation and pitch

The method `torch.quantization.fuse_modules` supports many of the common fusion strategies, i.e., conv+bn, conv+bn+relu, etc. However, there are additional fusion operations that are useful in practice that could be interesting. Specifically, cascades of bn+linear layers can actually be fused trivially using the following. The documentation contains the algebraic derivation of the fusion.

```python
@torch.no_grad()
def fuse_batch_norm_1d_into_linear(norm: nn.BatchNorm1d, linear: nn.Linear, epsilon: float=1e-12) -> None:
    """
    Fuse a batch norm module into the linear layer that follows it.

    Args:
        norm: The batch norm layer that occurs before the convolution layer.
        linear: The linear layer to fuse the batch norm into.
        epsilon: A small value for numerical stability.

    Returns:
        None

    Details:
        This function de-composes the fusion into four simple steps. Assume that the
        cascade of a 1d batch normalization into a linear layer is formulated
        as follows where \f$x\f$ is the input vector, \f$\mu, \sigma\f$ are the
        moving statistics of the batch norm, \f$\gamma, \beta\f$ are the learned
        affine parameters of the batch norm, and \f$W, b\f$ are the weights and
        biases of the linear layer.

        \f$y = \Big[ \frac{x - \mu}{\sigma} \odot \gamma + \beta \Big] \cdot W + b\f$

        1. Apply the distributive property to group \f$\beta\f$ with the bias \f$b\f$.
           This allows \f$\beta\f$ to be absorbed by the bias of the linear layer:

           \f$y = \Big[ \frac{x - \mu}{\sigma} \odot \gamma \Big] \cdot W + \beta \cdot W + b\f$

           Update: \f$b \gets \beta \cdot W + b\f$

        2. Apply the associative law for scalar and dot product to group \f$\gamma\f$
           with the weight \f$W\f$. This allows \f$\gamma\f$ to be absorbed by the weight:

           \f$y = \Big[ \frac{x - \mu}{\sigma} \Big] \cdot \big[ W \odot \gamma \big] + b\f$

           Update: \f$W \gets W \odot \gamma\f$

        3. Apply the associative law for scalar and dot product to group \f$\sigma\f$
           with the weight \f$W\f$. This allows \f$\sigma\f$ to be absorbed by the weight:

           \f$y = \big[ x - \mu \big] \cdot \Big[ W \odot \frac{1}{\sigma} \Big] + b\f$

           Update: \f$W \gets W \odot \frac{1}{\sigma}\f$

        4. Apply the distributive property to group \f$\mu\f$ with the bias \f$b\f$.
           This allows \f$\mu\f$ to be absorbed by the bias:

           \f$y = x \cdot W - \mu \cdot W + b\f$

           Update: \f$b \gets b - \mu \cdot W\f$

        This leaves the final simplified linear form with the batch norm analytically
        integrated into the calculation. The batch norm can now be replaced by the
        fused linear layer:

        \f$y = x \cdot W + b\f$

    """
    # 1. Apply distributive property to group β with the bias.
    offset = norm.bias @ linear.weight.T
    if linear.bias is None:
        linear.bias = nn.Parameter(offset)
    else:
        linear.bias[:] = linear.bias + offset
    norm.bias.fill_(0.0)  # Reset β to identity.
    # 2. Apply associative law for scalar and dot product to group γ with weight.
    linear.weight[:] = linear.weight * norm.weight
    norm.weight.fill_(1.0)  # Reset γ to identity.
    # 3. Apply associative law for scalar and dot product to group Var[x] with weight.
    linear.weight[:] = linear.weight / norm.running_var.add(epsilon).sqrt()
    norm.running_var[:] = 1.0  # reset Var[x] to identity.
    # 4. Apply distributive property to group E[x] with bias.
    offset = norm.running_mean @ linear.weight.T
    linear.bias[:] = linear.bias - offset
    norm.running_mean[:] = 0.0  # reset E[x] to identity.
```

This same concept can be applied to bn+conv, though the derivation is less straight forward when supporting strided convolution, group convolution, etc. Happy to provide the derivation and code for that if these are features the PyTorch community would be interested in adding to the library directly. I certainly find them useful in practice!

### Alternatives

I'm aware that `torch.quantization.fuse_modules` can be augmented using `fuse_custom_config_dict`, but perhaps directly integrating these fusion policies into PyTorch could be helpful. I certainly find them useful in practice!


### Additional context

_No response_

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim