Remove dtensor dependency in Tensor Parallel

[3outeille]

Motivation

• API stays the same but implementation details are different. It will be easier for us to debug whihc distributed calls are called we differentiated them by ourselves (example of VocabParallelEmbedding) and thus can put a breakpoint.

• Benchmark shows speedup with torch.compile on/off

• An example with ColwiseParallel

class _AllReduceBackward(torch.autograd.Function):
    """Identity forward, all-reduce backward. Used before colwise layers (f in Megatron)."""

    @staticmethod
    def forward(ctx, x, device_mesh):
        ctx.device_mesh = device_mesh
        return x

    @staticmethod
    def backward(ctx, grad_output):
        device_mesh = ctx.device_mesh
        if device_mesh.size() == 1:
            return grad_output, None
        dist.all_reduce(grad_output, op=dist.ReduceOp.SUM, group=device_mesh.get_group())
        return grad_output, None


def all_reduce_backward(x, device_mesh):
    """Identity forward, all-reduce backward. Use before colwise layers."""
    return _AllReduceBackward.apply(x, device_mesh)
    
class ColwiseParallel(TensorParallelLayer):
    """
    Column-wise parallel: weight is sharded on dim -2 (output features).
    Forward: input replicated -> output sharded on last dim.
    If gather_output=True, output is all-gathered to produce full tensor.
    """

    def __init__(self, gather_output: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.gather_output = gather_output

    def _prepare_input_fn(self, mod, inputs, device_mesh):
        input_tensor = inputs[0] if inputs else inputs
        return all_reduce_backward(input_tensor, device_mesh)

    def _prepare_output_fn(self, mod, outputs, device_mesh):
        if self.gather_output:
            return all_gather(outputs, device_mesh)
        return outputs

    def shard_tensor(
        self,
        param,
        param_type=None,
        param_casting_dtype=None,
        to_contiguous=None,
        rank=None,
        device_mesh=None,
        tensor_idx=None,
    ):
        device_mesh = self.device_mesh
        empty_param = self.empty_param
        rank = self.rank
        if param_type == "bias":
            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1, tensor_idx)
        else:
            parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -2, tensor_idx)
        parameter = parameter.to(param_casting_dtype)
        return parameter, None

    def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
        distribute_module(
            module,
            device_mesh,
            self._prepare_input_fn,
            self._prepare_output_fn,
        )

Benchmark

• Scripts

#!/usr/bin/env python
"""
Profiling script for Tensor Parallel using torch.profiler.

Run on BOTH branches:

    # WITH DTensor (base branch)
    git checkout v5-test_tensor_parallel_moe
    torchrun --nproc_per_node=2 profile_tp_trace.py

    # WITHOUT DTensor (PR branch)
    git checkout v5-tensor-parallel-dtensor
    torchrun --nproc_per_node=2 profile_tp_trace.py

    # With torch.compile enabled
    torchrun --nproc_per_node=2 profile_tp_trace.py --compile

Open traces with Perfetto UI: https://ui.perfetto.dev/
"""

import argparse
import os
import subprocess

import torch
import torch.distributed as dist
from torch.profiler import ProfilerActivity, profile
from transformers import AutoModelForCausalLM


def get_rank():
    if dist.is_initialized():
        return dist.get_rank()
    return 0


def log_rank0(msg):
    if get_rank() == 0:
        print(msg)


def get_current_branch():
    """Get current git branch name."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
            capture_output=True,
            text=True,
            check=True,
        )
        return result.stdout.strip()
    except Exception:
        return "unknown"


def parse_args():
    parser = argparse.ArgumentParser(description="Profile Tensor Parallel with torch.profiler")
    parser.add_argument(
        "--compile",
        action="store_true",
        help="Enable torch.compile for the model",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    
    model_id = "meta-llama/Llama-3.2-3B"
    batch_size = 4
    seq_length = 512
    output_dir = "./profiler_traces"
    
    torch.manual_seed(42)
    
    branch = get_current_branch()
    
    log_rank0("=" * 60)
    log_rank0("Tensor Parallel Profiling with torch.profiler")
    log_rank0("=" * 60)
    log_rank0(f"Branch: {branch}")
    log_rank0(f"Model: {model_id}")
    log_rank0(f"torch.compile: {args.compile}")
    
    # Load model with tensor parallel
    log_rank0("\nLoading model with tensor parallel...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.float16,
        tp_plan="auto",
    )
    
    # Apply torch.compile if requested
    if args.compile:
        log_rank0("Compiling model with torch.compile...")
        model = torch.compile(model)
    
    dist.barrier()
    device = model.device
    rank = get_rank()
    log_rank0(f"Model loaded on device: {device}")
    
    # Create input tensors
    vocab_size = model.config.vocab_size
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_length), device=device)
    labels = torch.randint(0, vocab_size, (batch_size, seq_length), device=device)
    
    model.train()
    
    # Warmup
    log_rank0("\nRunning warmup...")
    for _ in range(3):
        model.zero_grad()
        outputs = model(input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
    torch.cuda.synchronize()
    dist.barrier()
    
    # Profile one forward + backward pass
    log_rank0("\nProfiling forward + backward pass...")
    
    os.makedirs(output_dir, exist_ok=True)
    
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
        with_flops=True,
    ) as prof:
        model.zero_grad()
        
        # Forward pass
        outputs = model(input_ids, labels=labels)
        loss = outputs.loss
        
        # Backward pass
        loss.backward()
        
        torch.cuda.synchronize()
    
    dist.barrier()
    
    # Export trace with branch name and compile status
    compile_suffix = "_compiled" if args.compile else ""
    trace_file = os.path.join(output_dir, f"trace_{branch}{compile_suffix}_rank{rank}.json")
    prof.export_chrome_trace(trace_file)
    log_rank0(f"\nTrace saved to: {trace_file}")
    
    # Print summary table
    if rank == 0:
        compile_status = " (compiled)" if args.compile else ""
        print("\n" + "=" * 60)
        print(f"PROFILER SUMMARY - {branch}{compile_status}")
        print("=" * 60)
        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
    
    log_rank0(f"\nOpen trace with Perfetto UI: https://ui.perfetto.dev/")


if __name__ == "__main__":
    main()

• Baseline

• No Dtensor

• Baseline + torch.compile

• No Dtensor + torch.compile

[3outeille]

run-slow: afmoe, apertus, arcee, aria, bamba, cohere, cohere2, cwm, dbrx, deepseek_v2, deepseek_v3, diffllama, doge, dots1, emu3, ernie4_5

[github-actions]

This comment contains run-slow, running the specified jobs:

models: ["models/afmoe", "models/apertus", "models/arcee", "models/aria", "models/bamba", "models/cohere", "models/cohere2", "models/cwm", "models/dbrx", "models/deepseek_v2", "models/deepseek_v3", "models/diffllama", "models/doge", "models/dots1", "models/emu3", "models/ernie4_5"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

✅ No failing test specific to this PR 🎉 !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

Will review in details later on, make sure you write explicitly the motivation behind removing dtensor please (In terms of API etc) !

[ArthurZucker]

this is not really optimal, we should build / copy the utils or just create a dummy Dtensor with meta device just to let it handle the shape! But we cannot ship without shape checking IMO!

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: afmoe, apertus, arcee, aria, bamba, cohere, cohere2, cwm, dbrx, deepseek_v2, deepseek_v3, diffllama, doge, dots1, emu3, ernie4_5

[3outeille]

run-slow: afmoe, apertus, arcee, aria, bamba, cohere, cohere2, cwm, dbrx, deepseek_v2, deepseek_v3, diffllama, doge, dots1, emu3, ernie4_5

[github-actions]

This comment contains run-slow, running the specified jobs:

models: ["models/afmoe", "models/apertus", "models/arcee", "models/aria", "models/bamba", "models/cohere", "models/cohere2", "models/cwm", "models/dbrx", "models/deepseek_v2", "models/deepseek_v3", "models/diffllama", "models/doge", "models/dots1", "models/emu3", "models/ernie4_5"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

✅ No failing test specific to this PR 🎉 !

[3outeille]

merging this PR into #42809 to unblock me. It will provide more more thorough testing

[ArthurZucker]

Ok!