Added tensor_parallelism examples (#3047)

Co-authored-by: cehongwang <[email protected]>
pytorch · Aug 23, 2024 · 846fdd2 · 846fdd2
1 parent 2f7dbe8
commit 846fdd2
Show file tree

Hide file tree

Showing 4 changed files with 712 additions and 1 deletion.
diff --git a/examples/distributed_inference/README.md b/examples/distributed_inference/README.md
@@ -11,4 +11,6 @@ See the examples started with `data_parallel` for more details.
 
 2. Tensor parallel distributed inference
 
-In development.
+Here we use torch.distributed as an example, but compilation with tensor parallelism is agnostic to the implementation framework as long as the module is properly sharded.
+
+torchrun --nproc_per_node=2 tensor_parallel_llama2.py