[spmd_types] fix dsv3 ep + chunk loss typecheck (#3866)

https://github.com/pytorch/torchtitan/actions/runs/28794699622/job/85382003819

- ChunkedLossWrapper(CrossEntropyLoss) wrapper refactor broke typecheck,
since now it always goes through base loss class
- existing bug for EP + no typecheck

Author: pianpwk

torchtitan/components/loss.py

@@ -327,6 +327,10 @@ def __call__(
             # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
             with spmd.no_typecheck():
                 loss = loss / global_valid_tokens
+                if get_spmd_backend() == "spmd_types":
+                    spmd.assert_type(
+                        loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I}
+                    )
         return loss, {}
 
 
@@ -354,6 +358,10 @@ def __call__(
             # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
             with spmd.no_typecheck():
                 loss = loss / global_valid_tokens
+                if get_spmd_backend() == "spmd_types":
+                    spmd.assert_type(
+                        loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I}
+                    )
         return loss, {}
 
 
@@ -737,9 +745,6 @@ def _chunk(t):
                     logits, label_chunk, global_valid_tokens, **chunk_inputs
                 )
                 metrics = self._combine_chunk_metrics(metrics, chunk_metrics)
-                if get_spmd_backend() == "spmd_types":
-                    # V -> P reinterpret, after exiting local region.
-                    spmd.assert_type(chunk_loss, {"dp": spmd.P, "cp": spmd.P})
                 total_loss = total_loss + chunk_loss.detach()
 
                 if requires_grad:

torchtitan/models/common/token_dispatcher.py

@@ -448,7 +448,7 @@ def dispatch(
                 if get_spmd_backend() == "spmd_types"
                 else self.ep_mesh.get_group()
             )
-            if get_spmd_backend() == "spmd_types":
+            if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
                 num_local_tokens_per_expert_E = spmd.reinterpret_mesh(
                     num_local_tokens_per_expert_E, spmd.current_mesh()
                 )