Align DeepSeek V3 test config with TorchTitan shape

[sanketpurandare]

Stacked PRs:

• Add TorchTitan DSv3 equivalence coverage #444
• Move FSDP recompute tagging to the placement compile path #443
• ->Align DeepSeek V3 test config with TorchTitan shape #441

---

Align DeepSeek V3 test config with TorchTitan shape

This refactors AutoParallel's DeepSeek V3 test model around a hierarchical config shape that mirrors TorchTitan's DeepSeek V3 configuration while keeping make_dsv3_config() as the lightweight constructor for tests and examples. The model now reads layer, attention, RoPE, norm, FFN, and MoE settings through config objects, uses lm_head naming, accepts an optional mesh and compute_dtype, and preserves graph-trainer annotations for module FQNs and expert-parallel regions.

The DS3 implementation now performs linear, RMSNorm, attention, dense FFN, shared expert, and output computations through explicit compute-dtype helpers so the debug model can run in the same bfloat16 style expected by the TorchTitan comparison path. The old FORCE_BALANCED_ROUTING and CPU fill-index path are removed, expert execution unwraps DTensor weights locally where needed, and the local_map example uses MixedPrecisionPolicy plus make_dsv3_config() instead of maintaining a separate flat DeepSeekV3ModelArgs/MoEArgs construction path.

The local_map example also now binds each rank to its CUDA device before mesh/DTensor work, seeds DTensor RNG state explicitly, initializes NCCL with device_id, runs backward with autograd multithreading disabled, initializes weights on the rank device, and uses a device-specific final barrier. Those changes clean up the DTensor RNG sync, CUDA context/cuBLAS, and NCCL barrier warnings while keeping the example aligned with the real 2D dp/ep sharding constraints.

Two supporting graph/module fixes are included because they are required by the updated DS3 path. AutoParallel functionalizes index_put_ mutations when the mutation target is a fresh non-input tensor before AOT compilation, with tests that ensure input mutations are left alone. Parallel module construction now preserves non-persistent buffer registration when rebuilding sharded modules, so aliased RoPE buffers such as freqs_cis and rope.cache do not reappear in state_dict().

Authored with Claude.

[sanketpurandare]

does this now match titan's new dsv3 definition?

Yes it does and we can directly pass the config now from TorchTitan, no patching needed, same for annotations. Also by adding a compute type we don't have to force the entire model to be bfloat16, it interfaces nicely with optimizer in torchtitan as well.