diff --git a/tests/unit_tests/test_tensor_parallel.py b/tests/unit_tests/test_tensor_parallel.py
new file mode 100644
index 0000000000..204125d237
--- /dev/null
+++ b/tests/unit_tests/test_tensor_parallel.py
@@ -0,0 +1,77 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from unittest.mock import patch
+
+from torch.distributed.tensor import Replicate, Shard
+
+from torchtitan.distributed import tensor_parallel
+from torchtitan.distributed.tensor_parallel import NoParallel
+
+
+class _FakeDTensor:
+    def __init__(self, placements, local_value=None):
+        self.placements = placements
+        self.local_value = local_value
+        self.redistribute_calls = []
+
+    def redistribute(self, *, placements, async_op):
+        self.redistribute_calls.append((placements, async_op))
+        self.placements = placements
+        return self
+
+    def to_local(self):
+        return self.local_value
+
+
+class TestNoParallel(unittest.TestCase):
+    def test_prepare_output_handles_nested_outputs(self):
+        output_layout = Replicate()
+        first = _FakeDTensor((Shard(0),))
+        second = _FakeDTensor((Replicate(),))
+        metadata = object()
+        outputs = (first, {"nested": [second, metadata]})
+
+        with patch.object(tensor_parallel, "DTensor", _FakeDTensor):
+            result = NoParallel._prepare_output_fn(
+                output_layout,
+                False,
+                None,
+                outputs,
+                None,
+            )
+
+        self.assertIs(result[0], first)
+        self.assertEqual(first.placements, (output_layout,))
+        self.assertEqual(first.redistribute_calls, [((output_layout,), True)])
+        self.assertIs(result[1]["nested"][0], second)
+        self.assertEqual(second.redistribute_calls, [])
+        self.assertIs(result[1]["nested"][1], metadata)
+
+    def test_prepare_output_converts_each_dtensor_leaf_to_local(self):
+        output_layout = Replicate()
+        first_local = object()
+        second_local = object()
+        first = _FakeDTensor((Replicate(),), first_local)
+        second = _FakeDTensor((Shard(0),), second_local)
+
+        with patch.object(tensor_parallel, "DTensor", _FakeDTensor):
+            result = NoParallel._prepare_output_fn(
+                output_layout,
+                True,
+                None,
+                [first, (second,)],
+                None,
+            )
+
+        self.assertIs(result[0], first_local)
+        self.assertIs(result[1][0], second_local)
+        self.assertEqual(second.redistribute_calls, [((output_layout,), True)])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/torchtitan/distributed/tensor_parallel.py b/torchtitan/distributed/tensor_parallel.py
index b608da4712..892adc9e54 100644
--- a/torchtitan/distributed/tensor_parallel.py
+++ b/torchtitan/distributed/tensor_parallel.py
@@ -14,6 +14,7 @@
 from torch.distributed.tensor import distribute_module, DTensor, Replicate
 from torch.distributed.tensor.parallel import ParallelStyle
 from torch.distributed.tensor.placement_types import Placement
+from torch.utils._pytree import tree_map
 
 
 class NoParallel(ParallelStyle):
@@ -70,12 +71,19 @@ def _prepare_output_fn(
         output_layout: Placement,
         use_local_output: bool,
         mod: nn.Module,
-        outputs: DTensor,
+        outputs: Any,
         device_mesh: DeviceMesh,
-    ) -> torch.Tensor | DTensor:
-        if outputs.placements != (output_layout,):
-            outputs = outputs.redistribute(placements=(output_layout,), async_op=True)
-        return outputs.to_local() if use_local_output else outputs
+    ) -> Any:
+        def prepare_output(output: Any) -> Any:
+            if not isinstance(output, DTensor):
+                return output
+            if output.placements != (output_layout,):
+                output = output.redistribute(
+                    placements=(output_layout,), async_op=True
+                )
+            return output.to_local() if use_local_output else output
+
+        return tree_map(prepare_output, outputs)
 
     def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
         return distribute_module(
@@ -92,4 +100,4 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
                 self.output_layout,
                 self.use_local_output,
             ),
-        )
+        )
\ No newline at end of file