Allow sharding config to wrap declared plain inputs

Author: Matrix-Z97

torchtitan/protocols/module.py

@@ -590,20 +590,19 @@ def _redistribute_inputs(
             if mesh is None:
                 continue
 
-            if (
-                not isinstance(value, DTensor)
-                and parallel_dims.spmd_backend == "full_dtensor"
-            ):
-                raise ValueError("Got a plain Tensor under the full_dtensor mode.")
-
-            if not isinstance(value, DTensor) and src_spmd_layout is not None:
-                layout = resolve_placements(src_spmd_layout, mesh)
-                value = DTensor.from_local(
-                    value,
-                    mesh,
-                    layout,
-                    run_check=False,
-                )
+            if not isinstance(value, DTensor):
+                if src_spmd_layout is not None:
+                    layout = resolve_placements(src_spmd_layout, mesh)
+                    value = DTensor.from_local(
+                        value,
+                        mesh,
+                        layout,
+                        run_check=False,
+                    )
+                elif parallel_dims.spmd_backend == "full_dtensor":
+                    raise ValueError(
+                        "Got a plain Tensor under the full_dtensor mode."
+                    )
 
             if isinstance(value, DTensor) and src_spmd_layout is not None:
                 expected = resolve_placements(src_spmd_layout, mesh)