Repository: SqueezeBits/Torch-TRTLLM
Pull request #62: [FOC-945] refactor: reorder parallelization pass
Author: login: junstar92
Created: 2025-03-04T05:58:14Z | Merged: 2025-03-10T06:07:26Z

### Summary

- Refactored `transform` function to run parallelization passes at the first
  - `trtllm_export` exported only a Torch model to an inlined graph module
  -  Breaking down the custom pass and then adding pre and post pass
  - Passes for pre-custom transformation: `StashLoraSubgraphs`, `ConstantFolding`, `AddTRTLLMInputs`, `ResolveDynamicReshape`, `PropagateTensorParallelism`, `ParallelizeLinear`
  - Passes for post-custom transformation: `ParallelizePipeline`
  - Other passes are the same as before
- Updated `compare.sh` to remove used artifacts such as checkpoints and built engines


Diff:
diff --git a/scripts/compare.sh b/scripts/compare.sh
index fe5e5902..c6a4d8d1 100755
--- a/scripts/compare.sh
+++ b/scripts/compare.sh
@@ -21,6 +21,7 @@ declare REBUILD_NATIVE=false
 declare RERUN_NATIVE=false
 declare TP_SIZE=1
 declare PP_SIZE=1
+declare CLEANUP=false
 
 # Global variables for engine and artifact directories
 declare DITTO_ENGINE_DIR=""
@@ -103,6 +104,8 @@ print_help() {
     echo "  --rebuild-native       Rebuild native TensorRT-LLM engine even if it already exists"
     echo "  --rerun-native         Run native TensorRT-LLM engine even if output files exist"
     echo
+    echo "  --cleanup              Cleanup all engines and artifacts after comparison"
+    echo
     echo "[Example: Simply run meta-llama/Llama-2-7b-chat-hf]"
     echo "  $0 meta-llama/Llama-2-7b-chat-hf"
     echo
@@ -271,6 +274,10 @@ parse_args() {
                 fi
                 shift 2
                 ;;
+            --cleanup)
+                CLEANUP=true
+                shift
+                ;;
             -*)
                 echo "Error: Unrecognized option: $1"
                 echo "Try '$0 --help' for more information"
@@ -584,6 +591,9 @@ native_build() {
     local build_cmd="$TRTLLM_BUILD_SCRIPT $TRTLLM_BUILD_ARGS"
 
     DEBUG_ARTIFACTS_DIR=$TRTLLM_ARTIFACTS_DIR rich_execute "$build_cmd" "$TRTLLM_ENGINE_DIR/build.log" "build native TensorRT-LLM engine"
+    if [ "$CLEANUP" = true ]; then
+        rm -r $TRTLLM_CKPT_DIR
+    fi
 }
 
 # Compare outputs between two files and print result
@@ -655,6 +665,10 @@ compare_outputs() {
         fi
     done
 
+    if [ "$CLEANUP" = true ]; then
+        rm -r $TRTLLM_ENGINE_DIR -r $DITTO_ENGINE_DIR
+    fi
+
     return $all_succeeded
 }
 
diff --git a/scripts/test.sh b/scripts/test.sh
index 24401cbf..e3450f74 100755
--- a/scripts/test.sh
+++ b/scripts/test.sh
@@ -48,7 +48,7 @@ main() {
 
     for MODEL_SPECIFIC_ARG in "${MODEL_SPECIFIC_ARGS[@]}"; do
         echo "Testing $(echo ${MODEL_SPECIFIC_ARG} | sed 's/[[:space:]]\{2,\}/\n    /g')"
-        ${SCRIPT_DIR}/compare.sh ${MODEL_SPECIFIC_ARG} $@
+        ${SCRIPT_DIR}/compare.sh ${MODEL_SPECIFIC_ARG} --cleanup $@
     done
 }
 
diff --git a/src/ditto/api.py b/src/ditto/api.py
index 3f8517de..838160f1 100644
--- a/src/ditto/api.py
+++ b/src/ditto/api.py
@@ -42,8 +42,8 @@
 from .fx import generate_trtllm_engine_config
 from .fx.utils import find_output_node
 from .inline import inline
-from .literals import DTypeLiteral, PassName
-from .transform import parallelize, transform
+from .literals import DTypeLiteral
+from .transform import transform
 from .types import BuiltInConstant, verify
 
 
@@ -121,80 +121,64 @@ def trtllm_build(
         mapping=mapping,
     )
 
-    graph_module = trtllm_export(
-        model,
-        argument_hint,
-        model_config,
-        model.config.torch_dtype,
+    graph_module = trtllm_export(model, argument_hint)
+
+    for rank, transformed_graph_module in transform(
+        graph_module,
+        argument_hint=argument_hint,
+        model_config=model_config,
+        dtype=model.config.torch_dtype,
         run_matmuls_in_fp32=run_matmuls_in_fp32,
         run_activations_in_model_dtype=run_activations_in_model_dtype,
         run_routers_in_model_dtype=run_routers_in_model_dtype,
         extra_passes=[add_outputs(debug_node_names)] if debug_node_names else None,
-    )
-
-    for filename, component in build_trtllm_engine_components(
-        graph_module,
-        argument_hint,
-        build_config=TRTLLMBuildConfig.merge(profile_config, model_config),
-        trt_config=trt_config,
-        engine_cache=engine_cache,
-        network_name=get_network_name(model),
-        debug_node_names=debug_node_names,
     ):
-        save_component(output_dir, filename, component)
+        for filename, component in build_trtllm_engine_components(
+            rank,
+            transformed_graph_module,
+            argument_hint,
+            build_config=TRTLLMBuildConfig.merge(profile_config, model_config),
+            trt_config=trt_config,
+            engine_cache=engine_cache,
+            network_name=get_network_name(model),
+            debug_node_names=debug_node_names,
+        ):
+            save_component(output_dir, filename, component)
 
 
 def trtllm_export(
     model: PreTrainedModel | PeftModel,
     argument_hint: TRTLLMArgumentHint,
-    model_config: TRTLLMModelConfig,
-    dtype: torch.dtype,
     *,
-    run_matmuls_in_fp32: bool = False,
-    run_activations_in_model_dtype: bool = True,
-    skipped_optimizers: list[PassName] | None = None,
-    extra_passes: list[Callable[[GraphModule], GraphModule]] | None = None,
     enable_experimental_decompositions: bool = False,
-    run_routers_in_model_dtype: bool = False,
 ) -> GraphModule:
-    """Export a PyTorch model to a graph module optimized for TensorRT-LLM.
+    """Export a PyTorch model to an inlined graph module.
 
     This function performs several steps:
     1. Exports the PyTorch model to an exported program using torch.export
     2. Inlines the exported program into a graph module
-    3. Optimizes the graph module for TensorRT-LLM compatibility
 
     Args:
         model (PreTrainedModel | PeftModel): The PyTorch model to export
         argument_hint (TRTLLMArgumentHint): Configuration for input arguments
-        model_config (TRTLLMModelConfig): Model configurations
-        dtype (torch.dtype): Data type for the model
-        run_matmuls_in_fp32 (bool, optional): Whether to run matrix multiplications in FP32. Defaults to False.
-        run_activations_in_model_dtype (bool, optional): Whether to run activations in model dtype. Defaults to True.
-        skipped_optimizers (list[PassName] | None, optional): List of optimization passes to skip. Defaults to None.
-        extra_passes (list[Callable[[GraphModule], GraphModule]] | None, optional): Additional transformation passes to
-            apply. Defaults to None.
         enable_experimental_decompositions (bool, optional): Whether to enable experimental decompositions.
             Defaults to False.
-        run_routers_in_model_dtype (bool, optional): Whether to run linear layers for routers in MoE models in model
-            dtype instead of FP32. Defaults to False.
 
     Returns:
-        GraphModule: The optimized graph module for TensorRT-LLM
+        GraphModule: The inlined graph module
     """
     logger.info("Running torch.export")
     hints: dict[str, TensorTypeHint | BuiltInConstant] = {
         INPUT_IDS: argument_hint.batched_input_ids,
         "use_cache": False,
     }
-    device = model.device
-    arguments = TorchExportArguments.from_hints(device=device, **hints)
+    arguments = TorchExportArguments.from_hints(device=model.device, **hints)
     logger.opt(lazy=True).debug("{x}", x=lambda: arguments)
     exported_program = export(model, arguments)
     save_for_debug("exported_program", exported_program)
 
     logger.info("Inlining the exported program into a graph module")
-    logger.opt(lazy=True).debug("Memory Footprint: {m}", m=lambda: get_memory_footprint(device))
+    logger.opt(lazy=True).debug("Memory Footprint: {m}", m=lambda: get_memory_footprint(model.device))
     graph_module = inline(
         exported_program,
         class_name=get_network_name(model),
@@ -204,25 +188,14 @@ def trtllm_export(
     del exported_program
     torch.cuda.empty_cache()
 
-    logger.opt(lazy=True).debug("Memory Footprint: {m}", m=lambda: get_memory_footprint(device))
-    logger.info("Optimizing the graph module")
-    return transform(
-        graph_module,
-        argument_hint=argument_hint,
-        model_config=model_config,
-        dtype=dtype,
-        skipped_optimizers=skipped_optimizers,
-        run_matmuls_in_fp32=run_matmuls_in_fp32,
-        run_activations_in_model_dtype=run_activations_in_model_dtype,
-        run_routers_in_model_dtype=run_routers_in_model_dtype,
-        extra_passes=extra_passes,
-    )
+    return graph_module
 
 
 EngineComponent: TypeAlias = TRTLLMEngineConfig | bytes | dict[str, torch.Tensor] | LoraConfig
 
 
 def build_trtllm_engine_components(
+    rank: int,
     graph_module: GraphModule,
     argument_hint: TRTLLMArgumentHint,
     *,
@@ -235,6 +208,7 @@ def build_trtllm_engine_components(
     """Build TensorRT-LLM engine components from a graph module.
 
     Args:
+        rank (int): The rank of the engine component
         graph_module (GraphModule): The graph module to convert to TensorRT engines
         argument_hint (TRTLLMArgumentHint): Configuration for input arguments
         build_config (TRTLLMBuildConfig): Configuration for building TensorRT-LLM engines
@@ -249,49 +223,48 @@ def build_trtllm_engine_components(
             - TRTLLMEngineConfig | bytes | dict[str, torch.Tensor]: Either the engine configuration, serialized engine
                 bytes, or Lora state dicts.
     """
-    for rank, graph_module_per_rank in parallelize(graph_module, argument_hint):
-        if rank == 0:
-            yield (
-                "config.json",
-                generate_trtllm_engine_config(
-                    graph_module_per_rank,
-                    build_config,
-                    argument_hint.mapping,
-                    architecture=network_name,
-                ),
-            )
-            if (
-                lora_state_dicts := verify(
-                    graph_module_per_rank.meta.pop("lora_state_dicts", {}),
-                    as_type=dict[int, dict[str, torch.Tensor]],
-                )
-            ) and (
-                peft_configs := verify(
-                    graph_module_per_rank.meta.pop("peft_configs", {}),
-                    as_type=dict[int, LoraConfig],
-                )
-            ):
-                for lora_task_uid, state_dict in lora_state_dicts.items():
-                    yield f"lora/{lora_task_uid}/adapter_model.safetensors", state_dict
-                for lora_task_uid, lora_config in peft_configs.items():
-                    yield f"lora/{lora_task_uid}/adapter_config.json", lora_config
-        logger.opt(lazy=True).debug("Memory Footprint: {m}", m=get_memory_footprint)
-        save_for_debug(f"graph_module_rank{rank}", graph_module_per_rank)
-        logger.info(
-            "Building TensorRT engine{for_rank}",
-            for_rank=f" for rank {rank}" if argument_hint.mapping.world_size > 1 else "",
-        )
+    if rank == 0:
         yield (
-            f"rank{rank}.engine",
-            convert(
-                graph_module_per_rank,
-                argument_hint,
-                trt_config or TensorRTConfig(),
-                engine_cache=engine_cache,
-                network_name=network_name,
-                debug_node_names=debug_node_names,
+            "config.json",
+            generate_trtllm_engine_config(
+                graph_module,
+                build_config,
+                argument_hint.mapping,
+                architecture=network_name,
             ),
         )
+        if (
+            lora_state_dicts := verify(
+                graph_module.meta.pop("lora_state_dicts", {}),
+                as_type=dict[int, dict[str, torch.Tensor]],
+            )
+        ) and (
+            peft_configs := verify(
+                graph_module.meta.pop("peft_configs", {}),
+                as_type=dict[int, LoraConfig],
+            )
+        ):
+            for lora_task_uid, state_dict in lora_state_dicts.items():
+                yield f"lora/{lora_task_uid}/adapter_model.safetensors", state_dict
+            for lora_task_uid, lora_config in peft_configs.items():
+                yield f"lora/{lora_task_uid}/adapter_config.json", lora_config
+    logger.opt(lazy=True).debug("Memory Footprint: {m}", m=get_memory_footprint)
+    save_for_debug(f"graph_module_rank{rank}", graph_module)
+    logger.info(
+        "Building TensorRT engine{for_rank}",
+        for_rank=f" for rank {rank}" if argument_hint.mapping.world_size > 1 else "",
+    )
+    yield (
+        f"rank{rank}.engine",
+        convert(
+            graph_module,
+            argument_hint,
+            trt_config or TensorRTConfig(),
+            engine_cache=engine_cache,
+            network_name=network_name,
+            debug_node_names=debug_node_names,
+        ),
+    )
 
 
 def save_component(
diff --git a/src/ditto/arguments/trtllm_argument_hint.py b/src/ditto/arguments/trtllm_argument_hint.py
index c71a5528..85fad989 100644
--- a/src/ditto/arguments/trtllm_argument_hint.py
+++ b/src/ditto/arguments/trtllm_argument_hint.py
@@ -312,7 +312,17 @@ def all_reduce_workspace(self) -> TensorTypeHint | None:
         pointers_per_rank = 7
         pointers_of_counter = 2
         workspace_size = pointers_per_rank * self.mapping.tp_size + pointers_of_counter
-        return TensorTypeHint(shape=(workspace_size,), dtype=torch.int64)
+        return TensorTypeHint(
+            shape=(
+                DynamicDimension(
+                    name="workspace_size",
+                    min=workspace_size,
+                    opt=workspace_size,
+                    max=workspace_size,
+                ),
+            ),
+            dtype=torch.int64,
+        )
 
     @model_serializer(mode="wrap")
     def serialize_model(
diff --git a/src/ditto/fx/optimize.py b/src/ditto/fx/optimize.py
index 0f8aa3e5..6b1279db 100644
--- a/src/ditto/fx/optimize.py
+++ b/src/ditto/fx/optimize.py
@@ -23,7 +23,6 @@
 from ..constants import FX_TRANSFORM_MAXIMUM_ITERATION
 from ..literals import PassName
 from .passes import (
-    AddTRTLLMInputs,
     BindUnmatchedLoraProtos,
     CanonicalizeCopy,
     CastMMToFP32,
@@ -64,7 +63,6 @@
     RewritePowAsMul,
     RewriteReshapeAsUnsqueeze,
     RewriteSplitAsSlices,
-    StashLoraSubgraphs,
     WrapRoPESubgraphs,
     WrapSDPASubgraphs,
 )
@@ -212,16 +210,18 @@ def get_trtllm_conversion_transform(
         Callable[[GraphModule], GraphModule]: A function that applies TRT-LLM conversion passes to a graph module
     """
     passes: list[type[GraphOptimizationPass] | GraphOptimizationPass] = [
-        AddTRTLLMInputs(argument_hint=argument_hint),
         *get_trtllm_output_adaptation_passes(model_config.gather_context_logits),
-        StashLoraSubgraphs,
         CastRouterToFP32,
         ReplaceMoEByMoEPlugin(dtype=dtype),
         FuseQKVProjections,
         FuseGatedMLPProjections,
         WrapRoPESubgraphs,
         RewriteIndexAsSingleSlice,
-        ReplaceSDPAByGPTAttentionPlugin(dtype=dtype),
+        ReplaceSDPAByGPTAttentionPlugin(
+            dtype=dtype,
+            tp_size=argument_hint.mapping.tp_size,
+            tp_rank=argument_hint.mapping.tp_rank,
+        ),
         IndexLayers,
         BindUnmatchedLoraProtos,
         PopLoraPlugins(argument_hint=argument_hint),
diff --git a/src/ditto/fx/passes/parallelize_linear.py b/src/ditto/fx/passes/parallelize_linear.py
index 42cf390a..f3d05889 100644
--- a/src/ditto/fx/passes/parallelize_linear.py
+++ b/src/ditto/fx/passes/parallelize_linear.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from itertools import accumulate
-
 import tensorrt as trt
 import torch
 from pydantic import Field
@@ -22,9 +20,9 @@
 
 from ...configs import TRTLLMMapping
 from ...types import DataType
-from ..nodes import GetAttr, Permute, Reshape, Slice
+from ..nodes import Embedding, Expand, GetAttr, Permute, Reshape, Slice
 from ..subgraphs import FusedLinear, Linear
-from ..targets import AllGatherPlugin, AllReducePlugin, AllReducePluginInputs, GPTAttentionPlugin
+from ..targets import AllGatherPlugin, AllReducePlugin, AllReducePluginInputs
 from ..utils import forget_all_descendant_fake_tensors, get_val
 from .infra import (
     GraphOptimizationPass,
@@ -48,7 +46,9 @@ class ParallelizeLinear(GraphOptimizationPass):
     mapping: TRTLLMMapping = Field(frozen=True)
 
     def call(self, graph_module: GraphModule) -> PassResult:
-        first_node_rewritten: Node | None = None
+        if self.mapping.tp_size == 1:
+            return PassResult(graph_module=graph_module, modified=False)
+
         overall_modified = False
         graph = graph_module.graph
         lm_head = Linear.find_last(graph)
@@ -59,36 +59,46 @@ def call(self, graph_module: GraphModule) -> PassResult:
             linear = fused_linear.linear if fused_linear else linear
             tp_type = node.meta.get("tp_type", TensorParallelType.NONE)
             if tp_type == TensorParallelType.COLUMN:
-                slice_dim_sizes: list[int] | None = None
-                if gpt_attn_plugin := get_gpt_plugin(linear):
-                    q_dim = self.mapping.tp_size * gpt_attn_plugin.num_heads * gpt_attn_plugin.head_size
-                    kv_dim = self.mapping.tp_size * gpt_attn_plugin.num_kv_heads * gpt_attn_plugin.head_size
-                    slice_dim_sizes = [q_dim, kv_dim, kv_dim]
-                elif fused_linear:
-                    slice_dim_sizes = []
+                slice_dim: list[int] | None = None
+                if fused_linear:
+                    slice_dim = [0]
                     for slice_node in fused_linear.slices:
-                        assert isinstance(slice_node.start, int) and isinstance(
-                            slice_node.end, int
-                        ), "slice must have int start and end"
-                        slice_dim_sizes.append(slice_node.end - slice_node.start)
+                        slice_dim.append(slice_node.end)
                         self.parallelize_slice(slice_node)
-                self.parallelize_column_linear(
-                    linear,
-                    gather_output=linear == lm_head,
-                    slice_dim_sizes=slice_dim_sizes,
-                )
+
+                self.parallelize_column_linear(linear, slice_dim=slice_dim, gather_output=linear == lm_head)
             elif tp_type == TensorParallelType.ROW:
                 self.parallelize_row_linear(linear)
 
-            if first_node_rewritten is None:
-                first_node_rewritten = node
-
             overall_modified = True
 
-        if first_node_rewritten is not None:
-            forget_all_descendant_fake_tensors(first_node_rewritten)
+        if overall_modified:
+            # resolve reformatting nodes
+            embedding: Embedding | None = None
+            for node in graph.nodes:
+                if node.meta.get("tp_type", TensorParallelType.NONE) == TensorParallelType.NONE:
+                    embedding = embedding or Embedding.specialize_from(node)
+                    continue
+
+                if (reshape := Reshape.specialize_from(node) or Expand.specialize_from(node)) and -1 in reshape.shape:
+                    assert (
+                        reshape.shape.count(-1) == 1
+                    ), "reshape and expand node must have exactly one dynamic dimension"
 
-        return PassResult(graph_module=graph_module, modified=overall_modified, require_fake_tensor_prop=True)
+                    shard_dim_idx = get_shard_dim_idx(reshape)
+
+                    new_shape = list(reshape.shape)
+                    new_shape[shard_dim_idx] = new_shape[shard_dim_idx] // self.mapping.tp_size
+                    args, kwargs = reshape.args_kwargs(shape=new_shape)
+                    reshape.node.args = args
+                    reshape.node.kwargs = kwargs
+
+            assert embedding is not None, "embedding node not found"
+            forget_all_descendant_fake_tensors(embedding.node)
+
+        return PassResult(
+            graph_module=graph_module, modified=overall_modified, require_fake_tensor_prop=overall_modified
+        )
 
     def get_name_of_attr(self, from_name: str) -> str:
         """Create the name of the attribute with the tensor parallel rank.
@@ -105,53 +115,51 @@ def parallelize_column_linear(
         self,
         linear: Linear,
         *,
+        slice_dim: list[int] | None = None,
         gather_output: bool = True,
-        slice_dim_sizes: list[int] | None = None,
     ) -> None:
         """Parallelize the linear subgraph in column direction.
 
         Args:
             linear (Linear): The linear subgraph to be parallelized
+            slice_dim (list[int], optional): The dimension to slice. Defaults to None.
             gather_output (bool, optional): Whether to gather the output. Defaults to True.
-            slice_dim_sizes (list[int] | None, optional): The dimension sizes to slice the weight. Defaults to None.
         """
         if not (weight := GetAttr.specialize_from(linear.weight_node)):
             return
+        weight_tensor = weight.tensor
         local_out_features = linear.out_features // self.mapping.tp_size
-        if slice_dim_sizes:
-            assert all(
-                dim % self.mapping.tp_size == 0 for dim in slice_dim_sizes
-            ), "dimension to be sliced must be divisible by tp_size"
-            parallelized_weight_tensor = torch.cat(
-                list(zip(*self.slice_weight(weight.tensor, slice_dim_sizes, linear.has_transposed_weight)))[
-                    self.mapping.tp_rank
+        if linear.has_transposed_weight:
+            weight_tensor = weight_tensor.T
+
+        if slice_dim:
+            parallelized_weight_tensor = torch.concat(
+                [
+                    weight_tensor[
+                        :,
+                        slice_dim[i]
+                        + (slice_dim[i + 1] - slice_dim[i])
+                        // self.mapping.tp_size
+                        * self.mapping.tp_rank : slice_dim[i]
+                        + (slice_dim[i + 1] - slice_dim[i]) // self.mapping.tp_size * (self.mapping.tp_rank + 1),
+                    ]
+                    for i in range(len(slice_dim) - 1)
                 ],
-                dim=0 if linear.has_transposed_weight else 1,
+                dim=1,
             )
         else:
-            slice_size = local_out_features
-            if linear.has_transposed_weight:
-                parallelized_weight_tensor = weight.tensor[
-                    slice_size * self.mapping.tp_rank : slice_size * (self.mapping.tp_rank + 1), :
-                ]
-            else:
-                parallelized_weight_tensor = weight.tensor[
-                    :, slice_size * self.mapping.tp_rank : slice_size * (self.mapping.tp_rank + 1)
-                ]
+            parallelized_weight_tensor = weight_tensor[
+                :, local_out_features * self.mapping.tp_rank : local_out_features * (self.mapping.tp_rank + 1)
+            ]
 
-        assert (
-            parallelized_weight_tensor.shape
-            == (
-                local_out_features,
-                linear.in_features,
-            )
-            if linear.has_transposed_weight
-            else (
-                linear.in_features,
-                local_out_features,
-            )
+        assert parallelized_weight_tensor.shape == (
+            linear.in_features,
+            local_out_features,
         ), "unexpected shape of parallelized weight"
 
+        if linear.has_transposed_weight:
+            parallelized_weight_tensor = parallelized_weight_tensor.T
+
         graph = weight.node.graph
         with graph.inserting_before(weight.node):
             parallelized_weight = GetAttr.create(
@@ -166,14 +174,23 @@ def parallelize_column_linear(
             and linear.add is not None
             and (bias := GetAttr.specialize_from(linear.bias_node))
         ):
-            if slice_dim_sizes:
-                parallelized_bias_tensor = torch.cat(
-                    list(zip(*self.slice_bias(bias.tensor, slice_dim_sizes)))[self.mapping.tp_rank]
+            if slice_dim:
+                parallelized_bias_tensor = torch.concat(
+                    [
+                        bias.tensor[
+                            slice_dim[i]
+                            + (slice_dim[i + 1] - slice_dim[i])
+                            // self.mapping.tp_size
+                            * self.mapping.tp_rank : slice_dim[i]
+                            + (slice_dim[i + 1] - slice_dim[i]) // self.mapping.tp_size * (self.mapping.tp_rank + 1),
+                        ]
+                        for i in range(len(slice_dim) - 1)
+                    ],
+                    dim=0,
                 )
             else:
-                slice_size = local_out_features
                 parallelized_bias_tensor = bias.tensor[
-                    slice_size * self.mapping.tp_rank : slice_size * (self.mapping.tp_rank + 1)
+                    local_out_features * self.mapping.tp_rank : local_out_features * (self.mapping.tp_rank + 1)
                 ]
 
             with graph.inserting_before(bias.node):
@@ -252,8 +269,6 @@ def parallelize_row_linear(
     def parallelize_slice(self, slice_node: Slice) -> None:
         """Parallelize the slice node.
 
-        This function just updates the slice node to be parallelized in place.
-
         Args:
             slice_node (Slice): The slice node to be parallelized
         """
@@ -263,55 +278,6 @@ def parallelize_slice(self, slice_node: Slice) -> None:
         slice_node.node.args = args
         slice_node.node.kwargs = kwargs
 
-    def slice_weight(
-        self, weight: torch.Tensor, slice_dim_sizes: list[int], is_transposed: bool
-    ) -> list[list[torch.Tensor]]:
-        """Slice the weight tensor into multiple tensors.
-
-        Args:
-            weight (torch.Tensor): The weight tensor to be sliced
-            slice_dim_sizes (list[int]): The dimension sizes to slice the weight
-            is_transposed (bool): Whether the weight is transposed
-
-        Returns:
-            list[list[torch.Tensor]]: The sliced weight tensors for each tensor parallel rank
-        """
-        cumulative_slice_offsets = list(accumulate([0, *slice_dim_sizes]))
-        sliced_weights: list[list[torch.Tensor]] = []
-        for idx, slice_dim_size in enumerate(slice_dim_sizes):
-            slice_size = slice_dim_size // self.mapping.tp_size
-            if is_transposed:
-                weight_to_be_sliced = weight[cumulative_slice_offsets[idx] : cumulative_slice_offsets[idx + 1], :]
-                slices = [
-                    weight_to_be_sliced[slice_size * i : slice_size * (i + 1), :] for i in range(self.mapping.tp_size)
-                ]
-            else:
-                weight_to_be_sliced = weight[:, cumulative_slice_offsets[idx] : cumulative_slice_offsets[idx + 1]]
-                slices = [
-                    weight_to_be_sliced[:, slice_size * i : slice_size * (i + 1)] for i in range(self.mapping.tp_size)
-                ]
-            sliced_weights.append(slices)
-        return sliced_weights
-
-    def slice_bias(self, bias: torch.Tensor, slice_dim_sizes: list[int]) -> list[list[torch.Tensor]]:
-        """Slice the bias tensor into multiple tensors.
-
-        Args:
-            bias (torch.Tensor): The bias tensor to be sliced
-            slice_dim_sizes (list[int]): The dimension sizes to slice the bias
-
-        Returns:
-            list[list[torch.Tensor]]: The sliced bias tensors for each tensor parallel rank
-        """
-        cumulative_slice_offsets = list(accumulate([0, *slice_dim_sizes]))
-        sliced_biases: list[list[torch.Tensor]] = []
-        for idx, slice_dim_size in enumerate(slice_dim_sizes):
-            slice_size = slice_dim_size // self.mapping.tp_size
-            bias_to_be_sliced = bias[cumulative_slice_offsets[idx] : cumulative_slice_offsets[idx + 1]]
-            slices = [bias_to_be_sliced[slice_size * i : slice_size * (i + 1)] for i in range(self.mapping.tp_size)]
-            sliced_biases.append(slices)
-        return sliced_biases
-
 
 # [TODO] This function is only supported for 2D tensor, it should be extended to support arbitrary dimensions
 # pylint: disable-next=unused-argument
@@ -385,18 +351,25 @@ def insert_allreduce_plugin(
     to.replace_all_uses_with(allreduce, delete_user_cb=lambda user: user is not allreduce)
 
 
-def get_gpt_plugin(linear: Linear) -> GPTAttentionPlugin | None:
-    """Get the GPTAttentionPlugin from the linear node.
+def get_shard_dim_idx(reshape: Reshape | Expand) -> int:
+    """Get the shard dimension index of the reshape or expand node.
 
     Args:
-        linear (Linear): The linear node to get the GPTAttentionPlugin from
+        reshape (Reshape | Expand): The reshape or expand node
 
     Returns:
-        GPTAttentionPlugin | None: The GPTAttentionPlugin from the linear node
+        int: The shard dimension index
     """
-    return (
-        target
-        if len(users := list(linear.output_node.users)) == 1
-        and isinstance(target := users[0].target, GPTAttentionPlugin)
-        else None
-    )
+    dynamic_dim_idx = reshape.shape.index(-1)
+    if dynamic_dim_idx == 0 and len(reshape.shape) == 2:
+        shard_dim_idx = 1
+    elif dynamic_dim_idx == 1 and len(reshape.shape) in [3, 4]:
+        shard_dim_idx = 2
+    elif dynamic_dim_idx == 2 and len(reshape.shape) in [3, 4]:
+        shard_dim_idx = 1 if len(reshape.shape) == 4 else 0
+    elif dynamic_dim_idx == 3 and len(reshape.shape) == 5:
+        shard_dim_idx = 1
+    else:
+        raise ValueError(f"unexpected shape of reformatting node: {reshape.shape}")
+
+    return shard_dim_idx
diff --git a/src/ditto/fx/passes/parallelize_pipeline.py b/src/ditto/fx/passes/parallelize_pipeline.py
index a5686dd0..5c3d6a6c 100644
--- a/src/ditto/fx/passes/parallelize_pipeline.py
+++ b/src/ditto/fx/passes/parallelize_pipeline.py
@@ -20,7 +20,7 @@
 from ...arguments import TRTLLMArgumentHint
 from ...constants import INPUT_IDS
 from ...types import DataType
-from ..nodes import AddTensorTensor, BinaryElementwise, Placeholder, SymSizeInt, Unsqueeze
+from ..nodes import AddTensorTensor, BinaryElementwise, SymSizeInt, Unsqueeze
 from ..targets import GPTAttentionPlugin, RecvPlugin, SendPlugin
 from ..utils import find_closest_common_descendant, get_val
 from .infra import (
@@ -31,7 +31,7 @@
 )
 
 
-# pylint: disable=too-many-locals
+# pylint: disable=too-many-locals, too-many-boolean-expressions
 class ParallelizePipeline(GraphOptimizationPass):
     """Parallelize the pipeline in the graph (Pipeline Parallelism).
 
@@ -42,6 +42,9 @@ class ParallelizePipeline(GraphOptimizationPass):
     argument_hint: TRTLLMArgumentHint
 
     def call(self, graph_module: GraphModule) -> PassResult:
+        if self.argument_hint.mapping.pp_size == 1:
+            return PassResult(graph_module=graph_module, modified=False)
+
         overall_modified = False
         graph = graph_module.graph
         mapping = self.argument_hint.mapping
@@ -62,17 +65,15 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 and (pipeline_input_node := find_input_node_of_pipeline(node))
                 and (input_ids_node := find_input_ids_node(graph))
                 and self.argument_hint.hidden_states_input is not None
+                and (placeholders := {p.name: p for p in graph.find_nodes(op="placeholder")})
+                and (hidden_states_input := placeholders.get("hidden_states_input")) is not None
             ):
-                with graph.inserting_after(input_ids_node):
-                    hidden_states_input = Placeholder.create(
-                        graph, "hidden_states_input", self.argument_hint.hidden_states_input
-                    )
                 input_ids_node.replace_all_uses_with(
-                    hidden_states_input.node,
+                    hidden_states_input,
                     delete_user_cb=lambda user: SymSizeInt.specialize_from(user) is None,
                 )
 
-                recv_node = insert_recv_plugin(graph, hidden_states_input.node, mapping.prev_pp_rank)
+                recv_node = insert_recv_plugin(graph, hidden_states_input, mapping.prev_pp_rank)
                 pipeline_input_node.replace_all_uses_with(recv_node)
                 if (
                     len(users := list(recv_node.users)) == 2
diff --git a/src/ditto/fx/passes/propagate_tensor_parallelism.py b/src/ditto/fx/passes/propagate_tensor_parallelism.py
index 5b329a20..4ddf40c0 100644
--- a/src/ditto/fx/passes/propagate_tensor_parallelism.py
+++ b/src/ditto/fx/passes/propagate_tensor_parallelism.py
@@ -17,8 +17,7 @@
 from torch.fx import Node
 
 from ...configs.trtllm.pretrained import TRTLLMMapping
-from ..subgraphs.linear import Linear
-from ..targets import GPTAttentionPlugin
+from ..nodes import MM
 from .infra import NodewiseOptimizationPass, NodewisePassResult
 
 
@@ -59,16 +58,11 @@ class PropagateTensorParallelism(NodewiseOptimizationPass):
     mapping: TRTLLMMapping
 
     def rewrite(self, node: Node) -> dict[Node, NodewisePassResult]:
+        if self.mapping.tp_size == 1:
+            return {}
         prev_tp_type = get_previous_tp_type(node)
 
-        if isinstance(node.target, GPTAttentionPlugin):
-            node.target.num_heads = node.target.num_heads // self.mapping.tp_size
-            node.target.num_kv_heads = (node.target.num_kv_heads + self.mapping.tp_size - 1) // self.mapping.tp_size
-            node.target.tp_size = self.mapping.tp_size
-            node.target.tp_rank = self.mapping.tp_rank
-            node.meta["tp_type"] = prev_tp_type
-            return {}
-        if not Linear.configure_from(node):
+        if not MM.specialize_from(node):
             node.meta["tp_type"] = prev_tp_type
             return {}
 
diff --git a/src/ditto/fx/passes/replace_sdpa_by_gpt_attention_plugin.py b/src/ditto/fx/passes/replace_sdpa_by_gpt_attention_plugin.py
index 9fdbfc86..f026cf3b 100644
--- a/src/ditto/fx/passes/replace_sdpa_by_gpt_attention_plugin.py
+++ b/src/ditto/fx/passes/replace_sdpa_by_gpt_attention_plugin.py
@@ -35,6 +35,8 @@ class ReplaceSDPAByGPTAttentionPlugin(GraphOptimizationPass):
     """
 
     dtype: torch.dtype
+    tp_size: int
+    tp_rank: int
 
     # pylint: disable-next=too-many-locals
     def call(self, graph_module: GraphModule) -> PassResult:
@@ -92,6 +94,8 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 num_kv_heads=mha.num_kv_heads_per_group,
                 layer_idx_in_cache_pool=layer_idx,
                 head_size=mha.embed_dim,
+                tp_size=self.tp_size,
+                tp_rank=self.tp_rank,
                 type_id=DataType(self.dtype).to(trt.DataType),
                 **mha.rope_config.model_dump(),
             )
@@ -109,7 +113,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 # See https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
                 # pylint: disable-next=invalid-name
                 N, *others, Hq, _, Ev = mha.output_shape  # noqa: N806
-                out_reshape = Reshape.create(graph, plugin_node, [N, *others, -1, Hq, Ev])
+                out_reshape = Reshape.create(graph, plugin_node, [N, *others, -1, Hq // self.tp_size, Ev])
                 dims = [*range(4 + len(others))]
                 dims[-2], dims[-3] = dims[-3], dims[-2]
                 out_permute = Permute.create(graph, out_reshape, dims)
diff --git a/src/ditto/fx/passes/resolve_dynamic_reshape.py b/src/ditto/fx/passes/resolve_dynamic_reshape.py
index 600c2d79..163045d2 100644
--- a/src/ditto/fx/passes/resolve_dynamic_reshape.py
+++ b/src/ditto/fx/passes/resolve_dynamic_reshape.py
@@ -14,7 +14,7 @@
 
 from torch.fx import Node
 
-from ..nodes import Reshape, SymSizeInt
+from ..nodes import Expand, Reshape, SymSizeInt
 from .infra import ModifiedInsideThePass, NodewiseOptimizationPass, NodewisePassResult
 
 
@@ -31,7 +31,7 @@ class ResolveDynamicReshape(NodewiseOptimizationPass):
 
     def rewrite(self, node: Node) -> dict[Node, NodewisePassResult]:
         if (
-            (reshape := Reshape.specialize_from(node))
+            (reshape := Reshape.specialize_from(node) or Expand.specialize_from(node))
             and -1 not in reshape.shape
             and (len([dim for dim in reshape.shape if isinstance(dim, Node) and SymSizeInt.specialize_from(dim)]) == 1)
         ):
diff --git a/src/ditto/fx/subgraphs/moe.py b/src/ditto/fx/subgraphs/moe.py
index 77b6cb91..1722abb8 100644
--- a/src/ditto/fx/subgraphs/moe.py
+++ b/src/ditto/fx/subgraphs/moe.py
@@ -20,6 +20,7 @@
 
 from ...types import SymbolicInteger
 from ..nodes import MM, Gemm, IndexPut, MulTensorTensor, SelectInt, Softmax, ToCopy
+from ..utils import get_tensor_metadata
 from .gated_mlp import GatedMLP
 from .linear import Linear
 from .one_hot import OneHot
@@ -137,7 +138,10 @@ def expert_hidden_size(self) -> int:
         Returns:
             int: Hidden dimension size for expert networks
         """
-        return self.expert_weights[0][0].meta["tensor_meta"].shape[0]
+        assert (
+            expert_weight := get_tensor_metadata(self.expert_weights[0][0])
+        ) is not None, "expert weight's metadata is not found"
+        return expert_weight.shape[0]
 
     @property
     def expert_inter_size(self) -> int:
@@ -146,7 +150,10 @@ def expert_inter_size(self) -> int:
         Returns:
             int: Intermediate dimension size for expert networks
         """
-        return self.expert_weights[0][0].meta["tensor_meta"].shape[1]
+        assert (
+            expert_weight := get_tensor_metadata(self.expert_weights[0][0])
+        ) is not None, "expert weight's metadata is not found"
+        return expert_weight.shape[1]
 
     @property
     def shared_expert_intermediate_size(self) -> int:
@@ -155,7 +162,10 @@ def shared_expert_intermediate_size(self) -> int:
         Returns:
             int: Intermediate dimension size for the shared expert network
         """
-        return self.shared_expert_weights[0][0].meta["tensor_meta"].shape[1]
+        assert (
+            shared_expert_weight := get_tensor_metadata(self.shared_expert_weights[0][0])
+        ) is not None, "shared expert weight's metadata is not found"
+        return shared_expert_weight.shape[1]
 
     @classmethod
     def configure_from(cls, node: Node) -> Self | None:
diff --git a/src/ditto/fx/subgraphs/one_hot.py b/src/ditto/fx/subgraphs/one_hot.py
index 26a831a8..302f96a5 100644
--- a/src/ditto/fx/subgraphs/one_hot.py
+++ b/src/ditto/fx/subgraphs/one_hot.py
@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
 from torch.fx import Node
 from typing_extensions import Self
 
-from ..nodes import ArangeStartStep, EqTensor
+from ..nodes import ArangeStartStep, EqTensor, GetAttr
 from .subgraph import Subgraph
 
 
@@ -32,19 +33,28 @@ class OneHot(Subgraph):
     """
 
     eq: EqTensor
-    arange: ArangeStartStep
+    arange: ArangeStartStep | GetAttr
     this: Node
 
     @classmethod
     def configure_from(cls, node: Node) -> Self | None:
         if not (
             (eq := EqTensor.specialize_from(node))
-            and (arange_and_other := eq.specialize_either_side_as(ArangeStartStep))
+            and (
+                arange_and_other := eq.specialize_either_side_as(ArangeStartStep)
+                or eq.specialize_either_side_as(GetAttr)
+            )
         ):
             return None
 
         arange, other = arange_and_other
-        if not isinstance(other, Node):
+        if not (
+            isinstance(other, Node)
+            and isinstance(arange, GetAttr)
+            and torch.equal(arange.tensor, torch.arange(0, arange.tensor.numel()).to(arange.tensor.device))
+        ):
+            # Note: We assume that the arange is a tensor that would be generated by `torch.arange` operator
+            # if it is a GetAttr node.
             return None
 
         return cls(
diff --git a/src/ditto/transform.py b/src/ditto/transform.py
index f00f327c..4787b449 100644
--- a/src/ditto/transform.py
+++ b/src/ditto/transform.py
@@ -26,14 +26,18 @@
 from .arguments import TRTLLMArgumentHint
 from .configs import TRTLLMModelConfig
 from .contexts import ignore_symbolic_shapes_warning
-from .debug import save_for_debug
+from .debug import get_memory_footprint, save_for_debug
 from .fx import (
+    AddTRTLLMInputs,
+    ConstantFolding,
     ForgetSubmodules,
     ParallelizeLinear,
     ParallelizePipeline,
     Plugin,
     PropagateTensorParallelism,
     ResetCodeGen,
+    ResolveDynamicReshape,
+    StashLoraSubgraphs,
     fake_tensor_prop_on_node_creation,
     get_level1_transform,
     get_optimization_transform,
@@ -54,9 +58,15 @@ def transform(
     run_activations_in_model_dtype: bool = True,
     run_routers_in_model_dtype: bool = False,
     extra_passes: list[Callable[[GraphModule], GraphModule]] | None = None,
-) -> GraphModule:
+) -> Generator[tuple[int, GraphModule], None, None]:
     """Transform a PyTorch GraphModule by applying a series of optimization passes.
 
+    This function performs the following steps:
+    1. Runs the post-inlining passes
+    2. Runs the pre-custom passes for each rank
+    3. Runs the custom passes for each rank
+    4. Runs the post-custom passes for each rank
+
     Args:
         graph_module (GraphModule): The input PyTorch GraphModule to transform
         argument_hint (TRTLLMArgumentHint): Hints about the arguments for optimization
@@ -71,8 +81,11 @@ def transform(
             dtype instead of FP32. Defaults to False.
 
     Returns:
-        GraphModule: The transformed PyTorch GraphModule after applying optimization passes
+        Generator[tuple[int, GraphModule], None, None]: A generator of tuples containing the rank and the transformed
+            GraphModule after applying optimization passes
     """
+    logger.opt(lazy=True).debug("Memory Footprint: {m}", m=get_memory_footprint)
+    logger.info("Optimizing the graph module")
     post_inline_pass_manager = DynamoPassManager.build_from_passlist(
         [
             ResetCodeGen().as_transform(),
@@ -93,61 +106,47 @@ def transform(
 
     save_for_debug("initial_graph_module", graph_module)
 
-    custom_pass_manager = DynamoPassManager.build_from_passlist(
-        [
-            get_optimization_transform(
-                argument_hint,
-                model_config,
-                dtype,
-                skipped_optimizers=skipped_optimizers,
-                run_matmuls_in_fp32=run_matmuls_in_fp32,
-                run_activations_in_model_dtype=run_activations_in_model_dtype,
-                run_routers_in_model_dtype=run_routers_in_model_dtype,
-            ),
-            *(extra_passes or []),
-        ]
-    )
-    logger.debug("Running custom passes")
-    with fake_tensor_prop_on_node_creation(graph_module), ignore_symbolic_shapes_warning():
-        return custom_pass_manager(graph_module)
-
-
-def parallelize(
-    graph_module: GraphModule,
-    argument_hint: TRTLLMArgumentHint,
-) -> Generator[tuple[int, GraphModule], None, None]:
-    """Parallelize the graph module.
-
-    This function is used to parallelize the graph module for each rank, and update the mapping's rank in the argument
-    hint.
-
-    Args:
-        graph_module (GraphModule): The input graph module to parallelize
-        argument_hint (TRTLLMArgumentHint): The argument hint of the parallelized graph module
-
-    Returns:
-        Generator[tuple[int, GraphModule], None, None]:
-            A generator that yields the parallelized graph module for each rank
-    """
-    if argument_hint.mapping.world_size == 1:
-        argument_hint.mapping.rank = 0
-        yield 0, graph_module
-        return
-
-    logger.info("Parallelizing the graph module")
-    save_for_debug("graph_module_before_parallelization", graph_module)
     for rank in range(argument_hint.mapping.world_size):
-        logger.debug(f"Running parallelize passes for rank {rank}")
-        copied_graph_module = copy_graph_module(graph_module)
         argument_hint.mapping.rank = rank
-        parallelize_pass_manager = DynamoPassManager()
-        if argument_hint.mapping.tp_size > 1:
-            parallelize_pass_manager.add_pass(PropagateTensorParallelism(mapping=argument_hint.mapping).as_transform())
-            parallelize_pass_manager.add_pass(ParallelizeLinear(mapping=argument_hint.mapping).as_transform())
-        if argument_hint.mapping.pp_size > 1:
-            parallelize_pass_manager.add_pass(ParallelizePipeline(argument_hint=argument_hint).as_transform())
+        copied_graph_module = copy_graph_module(graph_module)
+        pre_custom_pass_manager = DynamoPassManager.build_from_passlist(
+            [
+                StashLoraSubgraphs().as_transform(),
+                ConstantFolding().as_transform(),
+                AddTRTLLMInputs(argument_hint=argument_hint).as_transform(),
+                ResolveDynamicReshape().as_transform(),
+                PropagateTensorParallelism(mapping=argument_hint.mapping).as_transform(),
+                ParallelizeLinear(mapping=argument_hint.mapping).as_transform(),
+            ]
+        )
+        logger.debug(f"Running pre-custom passes for rank {rank}")
+        with fake_tensor_prop_on_node_creation(copied_graph_module), ignore_symbolic_shapes_warning():
+            copied_graph_module = pre_custom_pass_manager(copied_graph_module)
+
+        custom_pass_manager = DynamoPassManager.build_from_passlist(
+            [
+                get_optimization_transform(
+                    argument_hint,
+                    model_config,
+                    dtype,
+                    skipped_optimizers=skipped_optimizers,
+                    run_matmuls_in_fp32=run_matmuls_in_fp32,
+                    run_activations_in_model_dtype=run_activations_in_model_dtype,
+                    run_routers_in_model_dtype=run_routers_in_model_dtype,
+                ),
+                *(extra_passes or []),
+            ]
+        )
+        logger.debug(f"Running custom passes for rank {rank}")
+        with fake_tensor_prop_on_node_creation(copied_graph_module), ignore_symbolic_shapes_warning():
+            copied_graph_module = custom_pass_manager(copied_graph_module)
+
+        post_pass_manager = DynamoPassManager.build_from_passlist(
+            [ParallelizePipeline(argument_hint=argument_hint).as_transform()]
+        )
+        logger.debug(f"Running post-custom passes for rank {rank}")
         with fake_tensor_prop_on_node_creation(copied_graph_module), ignore_symbolic_shapes_warning():
-            yield rank, parallelize_pass_manager(copied_graph_module)
+            yield rank, post_pass_manager(copied_graph_module)
 
 
 def copy_graph_module(graph_module: GraphModule) -> GraphModule:


Reviews:
user: login: jiwoong-choi
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-10T04:35:41Z
user: login: junstar92
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-10T04:56:58Z
user: login: jiwoong-choi
state: APPROVED
author_association: NONE
submitted_at: 2025-03-10T05:18:01Z

Review comments:
diff_hunk: @@ -653,8 +654,11 @@ compare_outputs() {
         if [ $? -ne 0 ]; then
             all_succeeded=1
         fi
+        rm -r $DITTO_ENGINE_DIR/lora
     done
 
+    rm $TRTLLM_ENGINE_DIR/rank*.engine $DITTO_ENGINE_DIR/rank*.engine
path: scripts/compare.sh
user: login: jiwoong-choi
body: Can you add a flag like `--cleanup` to run the new `rm ...` commands optionally?
created_at: 2025-03-10T04:05:55Z
author_association: NONE
diff_hunk: @@ -312,7 +312,17 @@ def all_reduce_workspace(self) -> TensorTypeHint | None:
         pointers_per_rank = 7
         pointers_of_counter = 2
         workspace_size = pointers_per_rank * self.mapping.tp_size + pointers_of_counter
-        return TensorTypeHint(shape=(workspace_size,), dtype=torch.int64)
+        return TensorTypeHint(
+            shape=(
+                DynamicDimension(
+                    name="workspace_size",
+                    min=workspace_size,
+                    opt=workspace_size,
+                    max=workspace_size,
+                ),
path: src/ditto/arguments/trtllm_argument_hint.py
user: login: jiwoong-choi
body: Why do we need this non-dynamic `DynamicDimension`?
created_at: 2025-03-10T04:11:40Z
author_association: NONE
diff_hunk: @@ -59,36 +59,46 @@ def call(self, graph_module: GraphModule) -> PassResult:
             linear = fused_linear.linear if fused_linear else linear
             tp_type = node.meta.get("tp_type", TensorParallelType.NONE)
             if tp_type == TensorParallelType.COLUMN:
-                slice_dim_sizes: list[int] | None = None
-                if gpt_attn_plugin := get_gpt_plugin(linear):
-                    q_dim = self.mapping.tp_size * gpt_attn_plugin.num_heads * gpt_attn_plugin.head_size
-                    kv_dim = self.mapping.tp_size * gpt_attn_plugin.num_kv_heads * gpt_attn_plugin.head_size
-                    slice_dim_sizes = [q_dim, kv_dim, kv_dim]
-                elif fused_linear:
-                    slice_dim_sizes = []
+                slice_dim: list[int] | None = None
+                if fused_linear:
+                    slice_dim = [0]
                     for slice_node in fused_linear.slices:
-                        assert isinstance(slice_node.start, int) and isinstance(
-                            slice_node.end, int
-                        ), "slice must have int start and end"
-                        slice_dim_sizes.append(slice_node.end - slice_node.start)
+                        slice_dim.append(slice_node.end)
                         self.parallelize_slice(slice_node)
-                self.parallelize_column_linear(
-                    linear,
-                    gather_output=linear == lm_head,
-                    slice_dim_sizes=slice_dim_sizes,
-                )
+
+                self.parallelize_column_linear(linear, slice_dim=slice_dim, gather_output=linear == lm_head)
             elif tp_type == TensorParallelType.ROW:
                 self.parallelize_row_linear(linear)
 
-            if first_node_rewritten is None:
-                first_node_rewritten = node
-
             overall_modified = True
 
-        if first_node_rewritten is not None:
-            forget_all_descendant_fake_tensors(first_node_rewritten)
+        if overall_modified:
+            # resolve reformatting nodes
+            embedding: Embedding | None = None
+            for node in graph.nodes:
+                if node.meta.get("tp_type", TensorParallelType.NONE) == TensorParallelType.NONE:
+                    embedding = embedding or Embedding.specialize_from(node)
+                    continue
+
+                if reshape := Reshape.specialize_from(node) or Expand.specialize_from(node):
path: src/ditto/fx/passes/parallelize_linear.py
user: login: jiwoong-choi
body: ```suggestion
                if (reshape := Reshape.specialize_from(node) or Expand.specialize_from(node)) and -1 in reshape.shape:
```
created_at: 2025-03-10T04:15:02Z
author_association: NONE
diff_hunk: @@ -312,7 +312,17 @@ def all_reduce_workspace(self) -> TensorTypeHint | None:
         pointers_per_rank = 7
         pointers_of_counter = 2
         workspace_size = pointers_per_rank * self.mapping.tp_size + pointers_of_counter
-        return TensorTypeHint(shape=(workspace_size,), dtype=torch.int64)
+        return TensorTypeHint(
+            shape=(
+                DynamicDimension(
+                    name="workspace_size",
+                    min=workspace_size,
+                    opt=workspace_size,
+                    max=workspace_size,
+                ),
path: src/ditto/arguments/trtllm_argument_hint.py
user: login: junstar92
body: It is not related to this PR. This change is just to remove the error message `'IOptimizationProfile::getDimensions: Error Code 4: API Usage Error (...)`. It is the same as `host_runtime_perf_knobs`, `host_max_attention_window_sizes`, ... 
created_at: 2025-03-10T04:56:58Z
author_association: MEMBER

PR comments:
user: login: junstar92
created_at: 2025-03-04T05:59:07Z
body: Tests are running. I will attach the results when the tests are done.
author_association: MEMBER
user: login: junstar92
created_at: 2025-03-05T08:31:46Z
body: There's something wrong with phi-4 & tp-2 and Qwen1.5-MoE. I will fix and update it soon.
author_association: MEMBER
user: login: junstar92
created_at: 2025-03-07T03:58:20Z
body: @kkimmk @huijjj @jiwoong-choi I fixed the errors with phi-4 with tp 2 and MoE. Please review this PR.
author_association: MEMBER