nonuniform tensor parallelism#2149

Draft

daiyaanarfeen wants to merge 9 commits intoNVIDIA:mainfrom

daiyaanarfeen:nonuniform-tp

Commits on Mar 4, 2026

feat: non-intrusive nonuniform tensor parallelism implementation
daiyaanarfeen
committed
test: add comprehensive unit tests for nonuniform TP
daiyaanarfeen
committed
fix: remove read-only shape attribute assignment in test
daiyaanarfeen
committed
fix: remove another shape assignment in test_ntp_backward_hook_core_gpu
daiyaanarfeen
committed
test: add end-to-end NTP test for 8 GPUs without mocking
daiyaanarfeen
committed
fix: skip spare ranks gracefully in end-to-end test
daiyaanarfeen
committed
refactor: move NTP config to separate NonuniformTPConfig class
daiyaanarfeen
committed
chore: remove unused typing imports from distributed_data_parallel_config.py
daiyaanarfeen
committed
test: fix integration tests to actually run without distributed mock
daiyaanarfeen
committed