Repository: pytorch
Issue: Batching rule not implemented for aten::narrow.Tensor
Labels: triaged,module: functorch

Hi,

I was trying out and benchmarking different single-device Mixture of Experts layer implementations, and I naively assumed vmap could be of help in this task. In some cases a simple vmap-based implementation was faster than my other implementations. However, I quickly came to the conclusion that this implementation is impractical due to its memory requirements. I am not sure that this is the only thing making my implementation memory inefficient, but I presume it is due to the copies. It cannot avoid copies because `.select()` uses Python ints, which are naturally unsupported by vmap, and indexing with a one-dimensional LongTensor or using `.index_select()` creates copies. Some [have proposed](https://github.com/pytorch/functorch/issues/337#issuecomment-994965422) to allow passing tensors to `.select()`, but I think no work is being done towards making that happen - correct me if I am wrong.

I tried to work around this by using `torch.narrow()`, but it seems to be unsupported by vmap:
```
RuntimeError: Batching rule not implemented for aten::narrow.Tensor; the fallback path doesn't work on out= or view ops.
```
Note that dynamic models may become more popular with time, and MoE layers are a very hot research topic recently. Judging by the other linked issue I speculate that JAX already supports this case, so functorch is behind in this aspect. I therefore kindly request `torch.narrow()` support in vmap.

I paste the code for my implementation that reproduces this below. I tested this on pytorch 2.0.0:
```python
import math

import functorch
import torch
import torch.nn as nn

class NarrowVmapExperts(nn.Module):
    def __init__(self,
                 dim,
                 num_experts,
                 hidden_dim=None,
                 activation=nn.GELU):
        super().__init__()
        self.hidden_dim = dim * 4 if hidden_dim is None else hidden_dim
        self.num_experts = num_experts
        # assumes homogeneous experts without biases
        w1 = torch.zeros(num_experts, dim, hidden_dim)
        w2 = torch.zeros(num_experts, hidden_dim, dim)
        self.init_(w1)
        self.init_(w2)
        self.w1 = nn.Parameter(w1)
        self.w2 = nn.Parameter(w2)
        self.act = activation()
        self.batched_expert_forward = functorch.vmap(self.single_expert_forward)
        self.batched_forward = functorch.vmap(self.one_sample_forward)

    @staticmethod
    def init_(t):
        dim = t.size(-1)
        std = 1 / math.sqrt(dim)
        t.uniform_(-std, std)

    def single_expert_forward(self, x, expert_index):
        # x is of size (dim)
        # expert_index is of size () - a scalar
        x = self.act(x @ self.w1.narrow(0, expert_index, 1).squeeze())
        x = x @ self.w2.narrow(0, expert_index, 1).squeeze()
        return x

    def one_sample_forward(self, x, expert_indices, expert_scores):
        # x is of size (dim)
        # expert indices is of size (k)
        # expert scores is of size (k)
        x = x.expand(expert_indices.size(0), -1)
        x = self.batched_expert_forward(x, expert_indices)
        x = torch.sum(x * expert_scores.view(-1, 1), dim=0)
        return x

    def forward(self, x, expert_indices, expert_scores):
        return self.batched_forward(x, expert_indices, expert_scores)


class TopKMoE(nn.Module):
    def __init__(self, num_experts, in_dim, out_dim, hidden_dim, k=2, noisy_gating=True, experts_class=NarrowVmapExperts):
        assert k <= num_experts
        super().__init__()
        self.num_experts = num_experts
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.hidden_dim = hidden_dim
        self.k = k
        self.noisy_gating = noisy_gating
        # instantiate experts
        self.experts = experts_class(dim=in_dim, num_experts=num_experts, hidden_dim=self.hidden_dim)
        self.w_g = nn.Parameter(torch.zeros(in_dim, num_experts), requires_grad=True)
        self.w_noise = nn.Parameter(torch.zeros(in_dim, num_experts), requires_grad=True)

    def gate(self, x):
        x_w_g = x @ self.w_g
        if self.training and self.noisy_gating:
            x_w_noise = x @ self.w_noise
            h_x = x_w_g + torch.randn_like(x_w_noise) * nn.functional.softplus(x_w_noise)
            top_logits, top_indices = h_x.topk(self.k, dim=-1, sorted=False)
            g_x = torch.softmax(top_logits, dim=-1)
        else:
            top_logits, top_indices = x_w_g.topk(self.k, dim=-1, sorted=False)
            g_x = torch.softmax(top_logits, dim=-1)
        return g_x, top_indices

    def forward(self, x, return_g=False):
        # x is of size (batch_size, sequence_length, dim)
        orig_size = x.size()
        x = x.view(-1, x.size(-1))
        g_x, g_indices = self.gate(x)
        out = self.experts(x, g_indices, g_x)
        out = out.view(orig_size)
        if return_g:
            return out, g_x.view(*orig_size[:-1], -1)
        else:
            return out

batch_size = 128
sequence_len = 64
num_experts = 64
dim = 256
hidden_dim = 384
layers = 9
k = 4

# instantiate
moe_layer = TopKMoE(num_experts, dim, dim, hidden_dim, k)
experts = moe_layer.experts

# check output size for a batch of inputs
x = torch.randn(batch_size, sequence_len, dim)
moe_out, g_x = moe_layer(x, return_g=True)
print(f'moe_out.size(): {moe_out.size()}')
print(f'g_x.size(): {g_x.size()}')
```

I also accept suggestions how may I implement MoE layers in an alternative, efficient manner - forgive me if, because of my limited knowledge, I missed some other obvious way to implement this efficiently.

cc @zou3519 @Chillee @samdow @soumith @kshitij12345 @janeyx99