Tensor parallel documentation

[apbose]

Tensor parallel Llama3 tutorial illustrating use of torch.distributed and nccl ops

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_llama3.py	2025-01-17 00:49:54.427641+00:00
+++ /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_llama3.py	2025-01-17 00:50:14.910448+00:00
@@ -64,11 +64,11 @@
    device="cuda",
)

with torch.no_grad():
    # The plan is
-    #plan = {
+    # plan = {
    # "attention": PrepareModuleInput(
    #     input_layouts=(Shard(1), None),
    #     desired_input_layouts=(Replicate(), None),
    # ),
    # "attention.wq": ColwiseParallel(),
@@ -82,22 +82,22 @@
    # ),
    # "feed_forward.w1": ColwiseParallel(),
    # "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
    # "feed_forward.w3": ColwiseParallel(),
    # "ffn_norm": SequenceParallel(),
-    #}
+    # }

    model = ParallelTransformer(model_args, device_mesh)

-# %%
-# Model inference with Torch-TensorRT backend
-# -------------------------------------------
-# When we compile the distributed model using Torch-TensorRT backend, pytorch distributed libraries create the sharded model
-# on multiple GPUs and the communicator operations are used for proper communication. In the above, 
-# `ColwiseParallel` and `RowwiseParallel` shard the attention layers in the column or row fashion. 
-# `SequenceParallel` performs sharded computations of the normalization layer
-# `PrepareModuleInput` configures the model input with proper communication operations
+    # %%
+    # Model inference with Torch-TensorRT backend
+    # -------------------------------------------
+    # When we compile the distributed model using Torch-TensorRT backend, pytorch distributed libraries create the sharded model
+    # on multiple GPUs and the communicator operations are used for proper communication. In the above,
+    # `ColwiseParallel` and `RowwiseParallel` shard the attention layers in the column or row fashion.
+    # `SequenceParallel` performs sharded computations of the normalization layer
+    # `PrepareModuleInput` configures the model input with proper communication operations

    torch.manual_seed(0)
    inp = torch.randint(32000, (8, 256), device="cuda")
    python_result = model(inp)
    torch_tensorrt.runtime.set_multi_device_safe_mode(True)

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_llama3.py	2025-01-17 00:57:39.378946+00:00
+++ /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_llama3.py	2025-01-17 00:57:59.494626+00:00
@@ -64,11 +64,11 @@
    device="cuda",
)

with torch.no_grad():
    # The plan is
-    #plan = {
+    # plan = {
    # "attention": PrepareModuleInput(
    #     input_layouts=(Shard(1), None),
    #     desired_input_layouts=(Replicate(), None),
    # ),
    # "attention.wq": ColwiseParallel(),
@@ -82,23 +82,23 @@
    # ),
    # "feed_forward.w1": ColwiseParallel(),
    # "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
    # "feed_forward.w3": ColwiseParallel(),
    # "ffn_norm": SequenceParallel(),
-    #}
+    # }

    model = ParallelTransformer(model_args, device_mesh)

-# %%
-# Model inference with Torch-TensorRT backend
-# -------------------------------------------
-#  When we compile the distributed model using Torch-TensorRT backend, pytorch distributed libraries create the sharded model
-#  on multiple GPUs and the communicator operations are used for proper communication. In the above, 
-# `ColwiseParallel` and `RowwiseParallel` shard the attention layers in the column or row fashion. 
-# `SequenceParallel` performs sharded computations of the normalization layer
-# `PrepareModuleInput` configures the model input with proper communication operations
-# The NCCL operations used in the distributed backend is handled by the TensorRT-LLM NCCL plugins, which causes no graph breaks now
+    # %%
+    # Model inference with Torch-TensorRT backend
+    # -------------------------------------------
+    #  When we compile the distributed model using Torch-TensorRT backend, pytorch distributed libraries create the sharded model
+    #  on multiple GPUs and the communicator operations are used for proper communication. In the above,
+    # `ColwiseParallel` and `RowwiseParallel` shard the attention layers in the column or row fashion.
+    # `SequenceParallel` performs sharded computations of the normalization layer
+    # `PrepareModuleInput` configures the model input with proper communication operations
+    # The NCCL operations used in the distributed backend is handled by the TensorRT-LLM NCCL plugins, which causes no graph breaks now

    torch.manual_seed(0)
    inp = torch.randint(32000, (8, 256), device="cuda")
    python_result = model(inp)
    torch_tensorrt.runtime.set_multi_device_safe_mode(True)