fix: warn on using single device sharding (#1274)

* fix: warn on using single device sharding

* fix: warn if users are creating single device mesh

* Update src/Sharding.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -832,8 +832,14 @@ GenerateCompileOptions(int64_t device_id, const int64_t *mesh_ids,
   options.executable_build_options.set_num_replicas(num_replicas);
   options.executable_build_options.set_num_partitions(num_partitions);
 
-  if (num_replicas > 1 || num_partitions > 1) {
-    assert(device_id < 0);
+  if (device_id < 0) {
+    if (num_replicas == 1 && num_partitions == 1) {
+      llvm::errs()
+          << "[libReactantExtra] num_replicas & num_partitions are both 1, but "
+             "device_id is negative. This can happen if you are sharding with "
+             "a single device.\n";
+    }
+
     assert(num_replicas * num_partitions == num_mesh_ids);
 
     options.executable_build_options.set_use_spmd_partitioning(

src/Sharding.jl

@@ -54,6 +54,15 @@ struct Mesh{D,ID<:AbstractVector{Int}}
         axis_names::NTuple{D,Union{String,Symbol}},
         axis_sizes::Dims{D},
     ) where {D}
+        @assert length(logical_device_ids) ≥ 1
+        if length(logical_device_ids) == 1
+            @warn "Constructing a single device mesh is not well supported and is \
+                   equivalent to not specifying any sharding. If you want to mock \
+                   multi-device setup on a single cpu host, set the environment variable \
+                   XLA_FLAGS=\"--xla_force_host_platform_device_count=12\" before loading \
+                   Reactant.jl and force reactant to use `cpu` devices using \
+                   `Reactant.set_default_backend(\"cpu\")`." maxlog = 1
+        end
         return new{D,typeof(logical_device_ids)}(
             sorted_device_ids, logical_device_ids, Symbol.(axis_names), axis_sizes
         )