Foundation piece for the ModelParallel layout-map work in #2807. Before the per-model layout maps can be tested consistently, they all need one place to assert sharding. This tracks adding TestCase.run_distribution_test, which checks the sharding spec (and fails loudly if a regex matched nothing), requires every rank>=2 weight to be either mapped or explicitly declared replicated, and runs a real forward/backward/optimizer step plus a parity check against an undistributed twin at two mesh shapes. It also registers the multi_device marker and adds a distribution-test CI job on 8 simulated devices.
No model files change here; the per-model PRs build on top. Handled by #2809.
Foundation piece for the ModelParallel layout-map work in #2807. Before the per-model layout maps can be tested consistently, they all need one place to assert sharding. This tracks adding
TestCase.run_distribution_test, which checks the sharding spec (and fails loudly if a regex matched nothing), requires every rank>=2 weight to be either mapped or explicitly declared replicated, and runs a real forward/backward/optimizer step plus a parity check against an undistributed twin at two mesh shapes. It also registers themulti_devicemarker and adds a distribution-test CI job on 8 simulated devices.No model files change here; the per-model PRs build on top. Handled by #2809.