refactor(distributed): establish shared native TP replica ownership

Summary

Follow-up to #3822: move TP-replica initialization and gradient-reduction ownership into shared model parallelization contracts and the distributed backends, so we can eventually remove the optimizer-boundary synchronization helper.

#3822 keeps the general correctness fix. Its native FSDP/DDP experiment was withdrawn from the merge scope in 6a946a85a: it validated useful building blocks, but the Nano-specific hybrid still needed the general helper elsewhere and did not demonstrate an end-to-end Nano speedup.

Why this needs separate work

Parameter replication does not specify what the local gradient means:

• Redundant full computation produces complete gradients that should be averaged across TP.
• Head-/sequence-sharded computation can produce partial contributions to a replicated parameter that must be summed.
• TP/EP/model-owned local shards must not be synchronized as replicas.
• An existing FSDP/DDP/DTensor reduction must not be applied again by another owner.

Initialization is a separate lifecycle problem: adapters may be initialized after parallelization, and rank-specific training RNG must not silently become a common model-initialization RNG.

Backend support also matters. In the inspected Megatron-FSDP stack, copying a DTensor gradient's local storage into its reduction buffer does not first resolve its TP Partial placement. Globally changing Q/K norm plans to Partial gradients would therefore lose required sums on that path. Re-check this behavior against the supported backend versions when implementing the migration.

TorchTitan is a useful reference for explicit, shared layout contracts with model-specific declarations, rather than a universal catch-all reduction:
shared decoder layouts,
Qwen3 Q/K norm layouts,
sharding protocol.

Preserved prototype and evidence

Prototype commit: e58de84bf81fb544754c8b47cfa17772485ab407. Its parent commits and tests remain in Git history; reverting them does not discard the experiment.

• Native Nano FSDP placement: validated dense HF Nano TP2, with DP-only ownership retained for TP-sharded projections. It deliberately falls back for unvalidated compositions; it is not a universal model plan.
• Production-strategy reference regression: 40 local GPU cases passed, including DP4 x TP2, accumulation, clipping, activation checkpointing and optimizer updates.
• Native draft DDP regression: combined DP x TP ownership with an asynchronous FP32 reduction hook. Separate four-rank CPU/NCCL reference checks covered initialization, BF16/FP32, clipping and partial accumulation windows. This can be evaluated as an independent, narrower migration.
• Exact-head scoped CI 67196187 passed Nano full SFT and Nano PEFT, with all six checkpoint phases passing and exact restored model/buffer/optimizer state on all eight ranks.
• A four-GPU DCP probe restored checkpoints from legacy DP-only ownership into native ownership exactly, then matched the next Adam update.
• Prototype GitHub CI passed lint/import/type checks, all unit suites and all non-GB200 L2 suites. The overall run failed in six GB200 suites; logs show NCCL network-plugin initialization failures. Do not treat it as an all-platform green result.

Performance is not established generally: a matched Nano PEFT probe changed by roughly +/-2% depending on warmup cutoff. A separate synthetic 67M-parameter BF16 draft benchmark measured 3.255 ms versus 2.624 ms per update (about 19% lower); this is not an end-to-end model-training speedup claim.

Proposed scope and acceptance criteria

• Define explicit replica/partial/shard semantics in the model plans and a shared ownership mechanism, without model-name exceptions in generic training code.
• Assign exactly one gradient-reduction owner per parameter for each supported backend; document or reject unsupported compositions.
• Cover initialization after model surgery/PEFT, without changing dropout/data RNG semantics.
• Resolve FSDP2 and Megatron-FSDP behavior for Partial gradients, communication dtype and gradient accumulation before changing shared norm styles.
• Validate representative dense, hybrid, MoE, PEFT, speculative and diffusion paths; include TP1/TP2 with DP composition and supported CP/PP combinations, not unsupported configurations.
• Compare gradients, global norms and optimizer updates to explicit references; include accumulation 1/>1/trailing windows, clipping on/off, BF16/FP32, frozen/unused/tied parameters and EP/model-owned exclusions.
• Verify checkpoint compatibility, exact restored state and continued optimizer updates.
• Measure end-to-end step time and memory on representative workloads; quantify any overlap benefit and extra FSDP wrapping cost.
• Remove helper calls only for fully owned paths, with a clear migration plan for remaining paths and no double reduction.

Potential benefits: fewer explicit recipe synchronization calls, better communication bucketing/overlap during backward, and clearer distributed ownership. None eliminates the mathematically required communication.

AMINT-314 (independent-process restarted-forward numerical drift after exact restoration) remains separate.

Start by reading nemo_automodel/components/models/nemotron_v3/parallelization.py and the referenced functional tests test_tp_replica_fsdp.py and test_tp_replica_ddp.py, then inspect the shared model parallelization contracts and distributed backends. Compare their ownership and initialization behavior with the stated replica, partial, and shard semantics. Done means supported paths have one reduction owner, preserve initialization and checkpoint behavior, pass the listed reference validations, and remove helper calls only where ownership is complete.