Tensor-parallelize the DeepSeek V3 transformer layer

[wujingyue]

cc @syed-ahmed

[github-actions]

Review updated until commit 98cadce

Description

• Added multidevice test for DeepSeek V3 transformer layer

• Parallelized transformer layer using Rowwise and Colwise parallelism

• Moved test from test_deepseek_v3.py to multidevice/test_deepseek_v3.py

---

Changes walkthrough 📝

	Relevant files
Enhancement	test_deepseek_v3.py Add multidevice test for DeepSeek V3 transformer layer      tests/python/multidevice/test_deepseek_v3.py Added new test file for multidevice testing of DeepSeek V3 transformer layer Implemented setup_process_group fixture for initializing process group Added default_tensor_type context manager for setting default tensor type and device Implemented test_transformer_layer to test parallelized transformer layer +143/-0
Other	test_deepseek_v3.py Remove old test_transformer_layer                                                tests/python/test_deepseek_v3.py Removed old test_transformer_layer function +0/-60

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Timeout Risk The test timed out once when downloading the model configuration. This could be a transient issue, but it's worth investigating to ensure it doesn't happen consistently. # This test timed out once when downloading # "/deepseek-ai/DeepSeek-V3/resolve/main/configuration_deepseek.py" (cf. # http://nv/eCm). I consider this a one-off, but please let me know if this # error becomes consistent. Hardcoded Port The default port for the process group initialization is hardcoded. This could lead to conflicts if multiple tests are run simultaneously. Consider using a dynamic port assignment. backend="nccl", init_method="tcp://localhost:29500", world_size=communicator.size(), Device Mesh Initialization The device mesh is initialized with a hardcoded device type ("cuda"). This could be problematic if the test is run in an environment without CUDA support. Consider making the device type configurable. mesh = dist.device_mesh.init_device_mesh("cuda", [d])

[wujingyue]

!test

[wujingyue]

!test

[kevinstephano]

LGTM.

[wujingyue]

!test

[wujingyue]

!test

[wujingyue]

!test