Update on "Making batching rule for F.embedding DTensor-aware"

zou3519 · zou3519 · commit 0bb02bb32c11 · 2025-09-04T14:45:59.000-07:00
`vmap(F.embedding)(DTensor, DTensor)` was failing because F.embedding's
batching rule generates a new tensor via at::arange, at::arange
generates a regular tensor, and DTensor rightfully errors on mixed
DTensor-regular Tensor operations.

This PR fixes the problem by activating DTensor implicit replication on
just the at::arange and the subsequent add operation.

In order to accomplish this I move the DTensor implicit replication flag
to C++ (most batching rules are in C++).

Test Plan:
- new test

[ghstack-poisoned]
diff --git a/aten/src/ATen/functorch/BatchRulesModules.cpp b/aten/src/ATen/functorch/BatchRulesModules.cpp
@@ -7,6 +7,7 @@
 #include <ATen/functorch/BatchRulesHelper.h>
 #include <ATen/functorch/PlumbingHelper.h>
 #include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/DTensorState.h>
 
 #include <utility>
 
@@ -44,8 +45,13 @@ static std::tuple<Tensor, std::optional<int64_t>> embedding_batch_rule(
   const auto weight_ = reshape_dim_into(*weight_bdim, 0, weight);
   auto indices_ = moveBatchDimToFront(indices, indices_bdim);
 
-  const auto range = getStepTensor(indices, batch_size, num_embeddings);
-  indices_ = indices_ + range;
+  {
+    // getStepTensor returns a regular Tensor. If indices_ is a DTensor
+    // we want to allow this mixed DTensor-Tensor operation.
+    at::DTensorAllowImplicitReplication guard;
+    const auto range = getStepTensor(indices, batch_size, num_embeddings);
+    indices_ = indices_ + range;
+  }
   auto result = at::embedding_symint(weight_, indices_, std::move(padding_idx), scale_grad_by_freq, sparse);
   return std::make_tuple(std::move(result), 0);
 }
diff --git a/torchgen/gen_vmap_plumbing.py b/torchgen/gen_vmap_plumbing.py
@@ -166,12 +166,10 @@ def gen_vmap_inplace_plumbing(native_function: NativeFunction) -> str | None:
     unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)
     bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
 
-    # NOTE: [vmap and DTensor Implicit Replication]
     return f"""\
 template <typename batch_rule_t, batch_rule_t batch_rule>
 {sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
   c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-  at::DTensorAllowImplicitReplication guard2;
   auto maybe_layer = maybeCurrentDynamicLayer();
   vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
   int64_t {cur_level_var} = maybe_layer->layerId();
@@ -190,12 +188,10 @@ def gen_vmap_plumbing_no_returns(native_function: NativeFunction) -> str:
     unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)
     bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
 
-    # See NOTE: [vmap and DTensor Implicit Replication]
     return f"""\
 template <typename batch_rule_t, batch_rule_t batch_rule>
 {sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
   c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-  at::DTensorAllowImplicitReplication guard2;
   auto maybe_layer = maybeCurrentDynamicLayer();
   vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
   int64_t {cur_level_var} = maybe_layer->layerId();
@@ -242,28 +238,10 @@ def gen_vmap_plumbing(native_function: NativeFunction) -> str | None:
     bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
 
     wrapped_returns = gen_returns(returns, cur_level_var, results_var)
-    # NOTE: [vmap and DTensor Implicit Replication]
-    # For all batching rules, we assume that new tensors generated by factory
-    # functions are meant to be replicated by DTensor. This saves us
-    # the trouble of rewriting all of them into e.g. using tensor.new_*.
-    # This is fine because batching rules cannot produce tensors with different
-    # shapes or values across ranks (vmap is pretty constrained),
-    # with the exception of random operations (those have their own separate
-    # handling that is not in this file, see BatchRulesRandomness.cpp).
-    #
-    # Why do we prefer the mode to e.g. rewriting calls to factory functions
-    # into their new_* variants?
-    #
-    # Consider batched_tensor + arange(n). This would work via broadcasting.
-    # However, if we changed this into batched_tensor + batched_tensor.new_arange(n),
-    # then we would actually allocate a repeated tensor for the arange (it
-    # would be a batched tensor), and we'd be using more memory than necessary,
-    # unless we implement some sort of lazy batchedtensor materialization thing.
     return f"""\
 template <typename batch_rule_t, batch_rule_t batch_rule>
 {sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
   c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
-  at::DTensorAllowImplicitReplication guard2;
   auto maybe_layer = maybeCurrentDynamicLayer();
   vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
   int64_t {cur_level_var} = maybe_layer->layerId();
@@ -288,7 +266,6 @@ def gen_all_vmap_plumbing(native_functions: Sequence[NativeFunction]) -> str:
 #pragma once
 #include <ATen/Operators.h>
 #include <ATen/functorch/PlumbingHelper.h>
-#include <ATen/DTensorState.h>
 
 namespace at {{ namespace functorch {{
 
