fix(fsdp): preserve already-sharded embedding units

[yuhezhang-ai]

What does this PR do?

Fix LLaDA initialization failing with a duplicate fully_shard assertion on FSDPEmbedding.

Our LLaDA recipe uses GSAI-ML/LLaDA-8B-Base, whose checkpoint configuration sets weight_tying: false. Its separate input embedding and output projection are stored in a ModuleDict. Recursive FSDP sharding wraps them before the separate embedding pass introduced in #3708. Skip modules that are already FSDP units so the embedding pass does not wrap them again.

Reject distinct tied input/output modules when either is inside a ModuleList or ModuleDict, including descendants. This check runs before recursive wrapping can replace their shared parameter and silently break the tie. Ordinary tied tables outside those containers remain allowed, as does a single module registered once and returned by both embedding getters.

If LLaDA's optional weight_tying=True mode is explicitly enabled, its tie_weights() registers the same module under both wte and ff_out. That duplicate-registration layout still fails inside recursive sharding on both main and this PR and remains outside this fix. Our shipped LLaDA recipe does not enable that mode.

Changelog

• Preserve already-sharded embedding units and their existing reshard policy.
• Keep passing model to the embedding helper. Leave both recursive sharding implementations unchanged; no additional ownership parameter is needed.
• Raise a clear error for the unsupported tied container layout before any FSDP wrapping.
• Document that an explicit reshard_after_forward=True can cause an extra head all-gather for a container-hosted head; the fused loss handles the sharded weight correctly.

Validation

• Three focused distributed test files: 207 passed, 2 skipped on the simplified implementation at 4202c825a999c43466e3cfeb018594b7da9d4d1b.
• Real FSDP tests on one and two CPU/Gloo ranks, using native mesh dimensions and both the default and Qwen3.5 dtype-based strategies. Supported layouts compare logits, per-parameter gradients, global gradient norm, and one SGD update against an unsharded fp32 reference.
• Cover all four untied table placements, ordinary tied tables outside containers, and single-module tying inside/outside containers with the shared module registered once. Unsupported tied placements, including a nested head, must raise while parameter identities and the shared weight remain intact and no module has been FSDP-wrapped.
• Ruff formatting/lint and git diff --check passed.
• Scoped LLaDA CI passed on the earlier revision 2b748d032cdeccab8a95fc723a8942eef158679d: 10 training steps on 8 H100 GPUs, checkpoint save/consolidation, and single-GPU reload/generation. Scoped CI has not been rerun for this simplification; normal PR CI is being rerun.

Before your PR is "Ready for review"

• Read and followed the contributor guidelines; commits include DCO sign-off.
• Added necessary regression coverage.
• Documentation assessed; no public API or configuration additions require documentation updates.

Additional Information

Fixes AMINT-329.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuhezhang-ai]

/ok to test 1ef407a

[yuhezhang-ai]

/ok to test 2b748d0

[yuhezhang-ai]

/ok to test 4202c82