NMFW-50: Colocated correctness tests

[yashaswikarnati]

Summary

Correctness tests for colocated MIMO training — compares heterogeneous TP/DP against TP=1 baseline.

• 9 comparisons per iteration (encoder output, post-communicate, LLM output, loss, input grads, encoder param grads, LLM param grads, encoder weights, LLM weights)
• 3 training iterations with SGD lr=1e-6
• 3 configs: fan-in (TP2/DP4→TP4/DP2), fan-out (TP4/DP2→TP2/DP4), equal (TP4/DP2→TP4/DP2)
• Full determinism stack (NVTE flags, CUBLAS, torch.use_deterministic_algorithms)
• FP32

Test command

uv run python -m torch.distributed.run --nproc_per_node=8 \
    -m pytest "tests/unit_tests/models/test_mimo_colocated_correctness.py::TestColocatedCorrectness::test_correctness[fan_in]" -v

Linear: NMFW-50
Targets: PR #1

🤖 Generated with Claude Code