Fix NoParallel output hook for multi-output modules (#3374)

[GodlyDonuts]

Summary

Fixes #3374. NoParallel's output hook crashes on modules that return more than one value.

NoParallel._prepare_output_fn assumed the wrapped module returns a single DTensor and accessed outputs.placements directly. Modules that return a tuple — e.g. the MoE router gate, which returns (logits, ...) — hit:

AttributeError: 'tuple' object has no attribute 'placements'

(full traceback in #3374; reproducible via the eager path of the script in #3345).

Fix

Map the existing redistribute / to_local logic over every output with torch.utils._pytree.tree_map, applying the configured output_layout / use_local_output to each DTensor output and leaving non-DTensor outputs (ints, None, etc.) untouched:

def _prepare(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, outputs)

The single-output path is unchanged: a lone DTensor is a pytree leaf, so tree_map applies _prepare to it directly and returns the same result as before. This matches the resolution suggested in the issue ("apply the same output_layout and local_output_grad_placements to every output via a tree_map").

Behavioral impact

Loss-neutral. For single-output modules the result is identical; for multi-output modules it replaces a hard crash with the intended per-output handling. No numerical change.

Testing

Adds tests/unit_tests/test_tensor_parallel.py (CPU / gloo via DTensorTestBase), covering:

• test_single_output_unchanged — regression: single-output modules behave exactly as before.
• test_multi_output — the bug: a module returning (tensor, tensor, int) now works; both DTensor outputs are redistributed/localized and the non-tensor output passes through.
• test_multi_output_keeps_dtensor_when_not_local — use_local_output=False keeps every output a DTensor.

Verified the new test_multi_output fails on main with the original AttributeError and passes with this change. black clean on all touched files.

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[tianyu-l]

This NoParallel class is being deprecated, as mentioned in #3374 (comment)

The last remaining usage is in https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/transformers_modeling_backend/parallelize.py

cc @mreso @HosseinKaviani-H

[GodlyDonuts]

Thanks @tianyu-l — that makes sense, no point hardening a class that's on its way out.

Since the last consumer is experiments/transformers_modeling_backend/parallelize.py, I'd be glad to migrate it off NoParallel so the class can be deleted outright — that closes #3374 more cleanly than this patch does. A couple of questions for @mreso / @HosseinKaviani-H before I touch your experiment, since the replacement isn't a straight swap:

• The Identity-placeholder cases (tok_embeddings / norm / lm_head when the stage doesn't own them under PP) are pure boundary conversion — PrepareModuleInput / PrepareModuleOutput look sufficient there.
• The per-layer cases (q_a_proj, q_a_layernorm, kv_a_proj_with_mqa, kv_a_layernorm) lean on NoParallel also replicating their params as DTensors on the TP mesh (for grad-norm clipping / fused optimizer). Is the intended target there the config-based full-DTensor path from [Full DTensor] Config-based Full DTensor for Llama3 #3159, or just a replicate plan?

If this migration is already on your plate, I'll happily close this PR. Otherwise point me at the preferred replacement and I'll put it up.

[HosseinKaviani-H]

@GodlyDonuts @tianyu-l

Thanks for your note. NoParallel isn't part of my SFT PR (my only change to this file is the Context Parallel / FlexAttention wrapping, the NoParallel usages predate my branch).

Regarding your questions

Identity placeholders (tok_embeddings/norm/lm_head under PP): no params, pure boundary conversion PrepareModuleInput/PrepareModuleOutput is the right swap.

DeepSeek MLA latents (q_a_proj, kv_a_proj_with_mqa, …): I think these need the replicate-params-as-DTensor behavior (for grad-norm clipping + fused optimizer), not just boundary conversion - so they need a real replicate plan, not PrepareModule. I didn't write these tho. @acisseJZhong @mreso can comment on this.

One ask since this file is shared: please keep the dense and TP/PP paths working and land the import removal in the same PR so the experiment keeps importing. Either way, the potential changes don't seem invasive for my PR.

[GodlyDonuts]

Thanks @HosseinKaviani-H, this is really helpful — and good to know the NoParallel usages predate your branch.

To capture where this leaves the migration:

• Identity placeholders (embeddings / norm / head under PP): agreed, these are pure boundary conversion and map cleanly onto PrepareModuleInput / PrepareModuleOutput.
• DeepSeek MLA latents (the q/kv down-projections and their norms): also agreed — these need the replicate-params-as-DTensor behavior, not just boundary conversion.

The MLA case is the one genuinely open question. There's no off-the-shelf replicate ParallelStyle in the pinned torch (only the PrepareModule helpers), so it's an architectural call rather than a drop-in swap. @acisseJZhong @mreso — for those param-bearing modules, would you prefer a small replicate plan that preserves today's behavior, or routing them through the config-based full-DTensor path (#3159)?

Once you point me at the preferred direction, I'm glad to do the migration and land it as a single PR that keeps both the dense and TP/PP paths working and removes the NoParallel import in the same change. Happy to use this PR as the vehicle or open a fresh one — whichever you'd rather.