From 59f337fd8eead35f67f82409901aa3daaed42190 Mon Sep 17 00:00:00 2001 From: zhangjianshe <1603088851@qq.com> Date: Fri, 3 Jul 2026 15:59:49 +0800 Subject: [PATCH] Allow sharding config to wrap declared plain inputs --- torchtitan/protocols/module.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/torchtitan/protocols/module.py b/torchtitan/protocols/module.py index 08d1a366c3..9cc4331d3e 100644 --- a/torchtitan/protocols/module.py +++ b/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)