Batching custom cuda function with vmap

viciopoli: <p>Hi everyone,</p>
<p>I’m trying to use <code>torch.vmap</code> to speed up my project.<br>
I have a custom cuda function, and I’m trying to wrap it into vmap to work with batches.</p>
<p>I’ll write here a toy code to better understand what’s going on:</p>
<pre><code class="lang-auto">class MyModule(nn.module):
    ....
    def forward(self, x):
        data = my_custom_cuda_func(x)
        ....

my_module1 = MyModule()
my_module2 = MyModule()

params, buffers = stack_module_state([my_module1, my_module2])

base_model = copy.deepcopy(my_module1)
base_model = base_model.to('meta')

def fmodel(params, buffers, x):
    return functional_call(base_model, (params, buffers), (x))

one_shot_model = vmap(fmodel)

test_data = ..... # a tensor of the shape [2, B, C]

results = one_shot_model(params, buffers, test_data)
</code></pre>
<p>I’m not sure that I’m properly understanding how the vmap works, but I would expect that when forward is called it receives a tensor of shape [B,C]. However, when I run the code, I get the error:</p>
<blockquote>
<p>RuntimeError: Batching rule not implemented for my_custom_cuda_func. We could not generate a fallback.</p>
</blockquote>
<p>I tried also to do wrap <code>my_custom_cuda_func</code> in vmap as follow:</p>
<pre><code class="lang-auto">data = vmap(my_custom_cuda_func)(x)
</code></pre>
<p>but I get the same error. Do you have any suggestions?<br>
Thanks!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/viciopoli">@viciopoli</a>,</p>
<p>When using <code>torch.func.vmap</code> it does not automatically generate a batching rule for custom functions (as noted at the bottom of the <code>torch.func.vmap</code> docs: <a href="https://pytorch.org/docs/stable/generated/torch.func.vmap.html#:~:text=vmap%20does%20not,Previous" class="inline-onebox" rel="noopener nofollow ugc">torch.func.vmap — PyTorch 2.3 documentation</a>).</p>
<p>In the case of a custom function, you need to manually specify the <code>vmap</code> rule (so <code>torch.func.vmap</code> can figure out how to vectorize over your function). This can be implemented via a <code>torch.autograd.Function</code> object, where you manually specify the <code>vmap</code> attribute. You can read more in the documentation: <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.3 documentation</a></p>
<p>In this case, you must also manually specify the <code>backward</code> formula for the custom <code>torch.autograd.Function</code> object in order to compute gradients.</p>

viciopoli: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> ,</p>
<p>Thanks for the answer, I tried to implement the batched version of my function by simply iterate over the non-btached function, I followed the <a href="https://github.com/pytorch/pytorch/blob/main/functorch/writing_batching_rules.md" rel="noopener nofollow ugc">writing_batching_rules</a> guide:</p>
<pre><code class="lang-auto">std::vector&lt;torch::Tensor&gt; my_custom_func(
        torch::Tensor x) {
  ...
  return some_other_func(x);
}

std::pair&lt;std::vector&lt;std::vector&lt;torch::Tensor&gt;&gt;, std::optional&lt;int64_t&gt;&gt; my_custom_func_batch_rule(
        const std::vector&lt;torch::Tensor&gt;&amp; x
        const std::optional&lt;int64_t&gt;&amp; batch_dim) {
    std::vector&lt;std::vector&lt;torch::Tensor&gt;&gt; results;
    for (size_t i = 0; i &lt; rays_o_batch.size(); ++i) {
        auto sample_result = sample_pts_on_rays(x[i]);
        results.push_back(sample_result);
    }
    return std::make_tuple(std::move(results), 0);
}
</code></pre>
<p>Is this the way of proceeding?<br>
After doing this I tried to register the batched function using VMAP_SUPPORT:</p>
<pre><code class="lang-auto">TORCH_LIBRARY(utils, m) {
  m.def("my_custom_func", &amp;my_custom_func);
  VMAP_SUPPORT(my_custom_func, my_custom_func_batch_rule)
  ....
}
</code></pre>
<p>However, it seems I’m missing something since when compiling I get the error:</p>
<blockquote>
<p>error: ‘VMAP_SUPPORT’ was not declared in this scope</p>
</blockquote>
<p>Do you know if there is a more in-depth guide for implementing these kinds of operations?</p>
<p>Thanks</p>
<p><strong>UPDATE</strong></p>
<pre><code class="lang-auto">#include &lt;ATen/functorch/BatchRulesHelper.h&gt;
</code></pre>
<p>solves the <code>VMAP_SUPPORT</code> error.</p>

AlphaBetaGamma96: <p>I don’t think simply looping over the data within the <code>torch.func.vmap</code> will work, but I’m not 100% sure.</p>
<p>From my understanding of <code>torch.func.vmap</code>, it removes the batch dim of your input Tensor so ‘looping’ over the batch dim no longer makes sense (as it no longer exists).</p>
<p>I think it’d be best to get a dev’s opinion on this! Apologies for the tag, <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

viciopoli: <p>hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>,</p>
<p>Yes, I think that’s right, what I see when debugging the Python code is a <code>BatchedTensor</code>, that has the same shape of the tensor that my function works with.</p>
<p>I also noticed that in the batch_rule declarations in <a href="https://github.dev/pytorch/pytorch/blob/main/aten/src/ATen/functorch/BatchRulesViews.cpp" rel="noopener nofollow ugc">BatchRulesViews.cpp</a> there is a <code>moveBatchDimToFront</code> call on the input tensor almost for in every batch_rule. Probably that’s the way of getting a normal tensor out of the <code>BatchedTensor</code>.</p>
<p>But, yes looking forward to hearing other suggestions!</p>
<p>Thanks.</p>

ptrblck: <p>Your batching rule should allow the usage of a batched input, but as <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> pointed out, using a for loop would still launch multiple kernels (one for each sample if I understand your code correctly) and might thus not yield a significant speedup (moving the loop from Python to C++ could still help). You could check if the actual operator (<code>sample_pts_on_rays</code>) could be rewritten in a way to accept batched inputs.</p>

viciopoli: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Thanks for your answer.<br>
Sorry my bad, <code>sample_pts_on_rays</code> should be there <code>my_custom_func</code>. I was trying to call several times the function that does not work with batches. So, ideally, I should rewrite everything such that batched inputs are supported, but what exactly is a batched input?<br>
When I debug the code in the Python script that calls the function, I see a <code>BatchedTensor</code>, but how do you work with this Tensor type?</p>
<p>Also, there is probably something wrong with the function declaration, since I get <code>my_custom_func_generated_plumbing</code> error when trying to VMAP_SUPPORT(my_custom_func, my_custom_func_batch_rule). Do you have any suggestions on why this could happen?</p>
<p>Thanks for your time!<br>
Vincenzo</p>

AlphaBetaGamma96: <p>I don’t think you should directly interact with the <code>BatchedTensor</code> type, as I think <code>torch.func</code> handles that explicitly under the hood. I think the best approach would be to look for the look at <code>writing_batching_rules</code> doc you mentioned above: <a href="https://github.com/pytorch/pytorch/blob/main/functorch/writing_batching_rules.md" class="inline-onebox" rel="noopener nofollow ugc">pytorch/functorch/writing_batching_rules.md at main · pytorch/pytorch · GitHub</a></p>

viciopoli: <p>I see, ok I’ll try to figure this out and I’ll update this post!</p>
<p>Thanks</p>