Remove dtensor dependency in Tensor Parallel
Author: 3outeille
Pull request: yes

# 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`
```python
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

```python
#!/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
<img width="1504" height="594" alt="image" src="https://github.com/user-attachments/assets/8128c819-0d4f-4b7c-80b3-48c9ebf8e253" />

- No Dtensor 
 <img width="1622" height="627" alt="image" src="https://github.com/user-attachments/assets/c555e816-0db0-4205-b134-f14470c636a4" />
 
- Baseline + `torch.compile`
<img width="1875" height="634" alt="image" src="https://github.com/user-attachments/assets/ee433c15-5183-415b-b9cf-40319605f272" />

- No Dtensor + `torch.compile`
<img width="1620" height="576" alt="image" src="https://github.com/user-attachments/assets/c9e39c32-ace1-4254-a00c-ee86f06faa88" />



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

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: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/20849505202)

✅ No failing test specific to this PR 🎉 !


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_43157). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[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

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

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: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/20918795329)

✅ No failing test specific to this PR 🎉 !


merging this PR into https://github.com/huggingface/transformers/pull/42809 to unblock me. It will provide more more thorough testing 

Ok!