Need a concise example of Tensor Parallelism (TP) training using Trainer/SFTTrainer.
Author: meet-minimalist
Pull request: no
Labels: name: Documentation

name: Feature request

name: Tensor Parallel

### Feature request

I have checked the code and there are few places which talk about TP. I saw from_pretrained method for model contains tp_plan and device_mesh. I also checked that the TrainingArgument can take parallelism_config which defines the TP/CP plan along with FSDP. However, I am not able to successfully stitch things together to make the only TP based training work. Please help.

Ref: 
- https://github.com/huggingface/transformers/blob/main/examples/3D_parallel.py

### Motivation

Need to enable only TP based training, but no tutorial or example is available.

### Your contribution

Given proper understanding and proper guidance, I can come up with clean example and documentation for the same.

Comments:
You can take a look at [Torch Native Parallelism guide](https://github.com/huggingface/accelerate/tree/main/examples/torch_native_parallelism) in [accelerate](https://huggingface.co/docs/accelerate) and the [tp_hsdp.yaml](https://github.com/huggingface/accelerate/blob/main/examples/torch_native_parallelism/configs/tp_hsdp.yaml) example to get some intuition.

Hope it helps 😄 


-----

## Conceptual Understanding: The Missing Glue

Think of your GPUs as a team of chefs and the model as a giant, complex recipe.

  * **Data Parallelism (DP):** Each chef (GPU) gets a full copy of the entire recipe (model). They each cook a different small dish (batch of data) simultaneously.
  * **Fully Sharded Data Parallelism (FSDP):** The recipe book (model weights) is torn into pieces. Each chef gets only one piece. When they need a part of the recipe they don't have, they quickly ask another chef for it.
  * **Tensor Parallelism (TP):** This is what you want. The recipe itself is too big for one chef's station. So, you split a single, very large cooking step (like a huge matrix multiplication in a transformer layer) across all the chefs. Chef 1 handles the first half of the ingredients, Chef 2 handles the second half. They work *simultaneously on the same step* for a single dish (batch).

The "glue" that makes this work in Hugging Face is the **`device_mesh`**. For TP-only, you need to tell `accelerate` to arrange all your GPUs in a single line (a 1D mesh) and assign them to the 'tensor\_parallelism' axis.

-----

## Step-by-Step Guide to TP-Only Training

### Step 1: Configure `accelerate`

First, ensure `accelerate` is configured for multi-GPU on a single machine. If you haven't done this, run `accelerate config` in your terminal and answer the questions like this (assuming you have 4 GPUs):

```bash
$ accelerate config
...
Which type of machine are you using? ([0] This machine, [1] AWS, ...) -> 0
Which type of multi-GPU training are you using? ([0] No, [1] DeepSpeed, [2] FSDP, [3] Megatron-LM) -> 0  # Choose 'No' for basic torch.distributed
...
How many processes in total will you use? [1]: -> 4 # Your number of GPUs
...
```

This creates a `default_config.yaml` file that `accelerate launch` will use.

### Step 2: The Training Script

This is where you'll define the device mesh and the `parallelism_config`. The key is to create a 1D mesh with a size equal to your number of GPUs.

1.  **Define the `parallelism_config`**: This dictionary is the modern way to control 3D parallelism. For TP-only, you set `tensor_parallelism` to `True` and specify the size. You also need to tell it which parts of the model to parallelize. For Llama-style models, this is typically the attention and MLP layers.
2.  **Create the `device_mesh`**: You'll get the world size (number of GPUs) from the environment and create a 1D mesh, e.g., `DeviceMesh(range(world_size), (world_size,))`.
3.  **Load the model within the mesh context**: Using the `init_empty_weights()` and the device mesh as a context manager, you load your model. `accelerate` will see the active device mesh and the `parallelism_config` in your `TrainingArguments` and automatically shard the model weights during initialization.

### Step 3: Launch the Script

You must use `accelerate launch` to run your script. This command reads your `accelerate` config and sets up the distributed environment variables that your script needs to know the world size and process rank.

```bash
accelerate launch --num_processes=4 your_tp_script.py
```

Notice that `--num_processes` must match `tensor_parallelism_size`.

-----

## Complete, Runnable Example

Here is a full script that puts everything together. This example uses a small Llama model and a dummy dataset so you can run it directly and see that it works without needing real data.

**File: `run_tp_only.py`**

```python
import os
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
)
from accelerate import init_empty_weights, infer_auto_device_map
from accelerate.utils import get_current_device
from accelerate.device_map import DeviceMap

# Get the number of GPUs from the environment
# accelerate launch sets this automatically
world_size = int(os.environ.get("WORLD_SIZE", 1))
device = get_current_device()

# --- 1. Define the Parallelism Configuration ---
# This is the new, unified way to configure TP/DP/FSDP
# For TP-only, we enable `tensor_parallelism` and set its size.
# The `tensor_parallelism_config` specifies which modules to shard.
# The keys ('self_attn', 'mlp', etc.) must match attribute names in the model class.
parallelism_config = {
    "tensor_parallelism": True,
    "tensor_parallelism_size": world_size,
    "tensor_parallelism_config": {
        # A list of patterns to match module names for sharding
        "self_attn": "col", # Shard Q,K,V and Output projection
        "mlp": "col",       # Shard the first MLP layer
    },
}

# --- 2. Set up Training Arguments ---
training_args = TrainingArguments(
    output_dir="./tp_llama_test",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    num_train_epochs=1,
    logging_steps=1,
    # Pass the parallelism_config to the trainer
    parallelism_config=parallelism_config,
    # Important for TP: Standard DataParallelism must be disabled
    ddp_find_unused_parameters=False,
)

# --- 3. Load Tokenizer and Model ---
# Model to be trained
model_name = "meta-llama/Llama-2-7b-hf" # Using a real model for correct architecture
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# The magic happens here!
# We use init_empty_weights() and the device_mesh to load the model.
# accelerate will automatically shard the weights based on our parallelism_config.
with init_empty_weights():
    # Inside this context, the model is created on the 'meta' device (no memory used)
    # The device_map="auto" along with the active parallelism config will shard it correctly.
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto", # Critical for accelerate to handle sharding
    )

print(f"Model loaded on device: {model.device}")
# You can inspect the model to see the sharded layers. They will be 'Shard(টেনসরপ্যারালাল)' objects.
# For example, on rank 0 you might see:
# model.layers[0].self_attn.q_proj = Shard(টেনসরপ্যারালাল(in_features=4096, out_features=1024, bias=False))
# And on rank 1:
# model.layers[0].self_attn.q_proj = Shard(টেনসরপ্যারালাল(in_features=4096, out_features=1024, bias=False))
# The out_features are divided by world_size (4096 / 4 = 1024 in this example if using 4 GPUs).
print(model)


# --- 4. Create a Dummy Dataset ---
# A simple dataset with a few examples to make the script runnable
dummy_data = {
    "text": [
        "This is a sample sentence for training.",
        "Tensor parallelism is a cool technique.",
        "We need to test the training loop.",
        "This example should run without errors."
    ]
}
dataset = Dataset.from_dict(dummy_data)

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)

tokenized_dataset = dataset.map(tokenize_function, batched=True)


# --- 5. Initialize Trainer and Start Training ---
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

print("Starting Tensor Parallel training...")
trainer.train()
print("Training complete!")

```

**To Run This Example:**

Assuming you have 2 GPUs:

```bash
accelerate launch --num_processes=2 run_tp_only.py
```

Assuming you have 4 GPUs:

```bash
accelerate launch --num_processes=4 run_tp_only.py
```

## Key Pitfalls and Things to Check

  * **`parallelism_config` is Key**: The old `tp_plan` argument is being phased out. Using `parallelism_config` in `TrainingArguments` is the correct and modern approach.
  * **Module Names**: The keys in `tensor_parallelism_config` (`"self_attn"`, `"mlp"`) must **exactly match** the attribute names of the `nn.Module`s in your model architecture. Use `print(model)` to inspect the model's structure and find the correct names.
  * **`accelerate launch` is a Must**: Running `python your_script.py` will not work. It won't set up the distributed environment, and your script will fail.
  * **Data Parallelism**: For a pure TP setup, ensure you are not accidentally enabling another form of parallelism. In this example, since each GPU is part of the TP group, the `per_device_train_batch_size` applies to the *entire model*. The global batch size will be `per_device_train_batch_size * gradient_accumulation_steps`.

hi @AkibDa , throwing the below error, while running the above script.

   ` from accelerate.utils import get_current_device
ImportError: cannot import name 'get_current_device' from 'accelerate.utils'`
Is there any specific version?


`ValueError: You can't train a model that has been loaded with `device_map='auto'` in any distributed mode. Please rerun your script specifying `--num_processes=1` or by launching with `python {{myscript.py}}`.`


Hi @quic-akuruvil 

I think for the import error  get_current_device:

That function isn’t present in older accelerate releases. 

Either:upgrade: `pip install -U accelerate (works with ≥0.21+)`, or

just don’t use it—there’s no need for it in the script provided by @AkibDa 


The script is still throwing a ValueError, as pointed out by @quic-akuruvil even after I upgrade the 'accelerate'.

 `device_map = auto` tells HF to slice the models across the devices for inference and when I lauch multiple processes with Trainer. It wants to controls how the model is split but since it already see "auto-sharded" model so it throws the error
`ValueError: You can't train a model that has been loaded with device_map='auto'in any distributed mode`

TP-Only training example(works with Trainer)

```
!pip install -U transformers accelerate datasets
import os, torch
from datasets import Dataset
from accelerate import ParallelismConfig
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)

# TP size: number of GPUs you launch with accelerate 
tp_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count() or 1))

pc = ParallelismConfig(
    tp_size=tp_size,   # TP only
    cp_size=1,
    dp_shard_size=1,
    dp_replicate_size=1,
)

model_id = "meta-llama/Llama-2-7b-hf"   # any TP-capable causal LM

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

# NOTE: no device_map="auto", no init_empty_weights() for training
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto",
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
)

# dummy text
texts = [
    "This is a sample sentence for training.",
    "Tensor parallelism is a cool technique.",
    "We need to test the training loop.",
    "This example should run without errors.",
]
ds = Dataset.from_dict({"text": texts})

def tokenize_fn(ex):
    out = tok(ex["text"], truncation=True, padding="max_length", max_length=128)
    out["labels"] = out["input_ids"]
    return out

train_ds = ds.map(tokenize_fn, batched=True, remove_columns=["text"])
collator = DataCollatorForLanguageModeling(tok, mlm=False)

args = TrainingArguments(
    output_dir="./tp_llama_test",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    num_train_epochs=1,
    logging_steps=1,
    save_steps=0,
    report_to=[],
    remove_unused_columns=False,
    bf16=torch.cuda.is_available(),
    ddp_find_unused_parameters=False,
    accelerator_config={"parallelism_config": pc.to_dict()},
)

trainer = Trainer(model=model, args=args, train_dataset=train_ds, data_collator=collator)
print(f"TP size: {tp_size}")
trainer.train()
```


Run the script with 2GPUs in this example:
`
accelerate launch --num_processes 2 tp_only_trainer.py`




> # The magic happens here!
> # We use init_empty_weights() and the device_mesh to load the model.
> # accelerate will automatically shard the weights based on our parallelism_config.
> with init_empty_weights():

Why this kind of empty weights initialization is needed here. Could you explain more the gist of it? It will go into error state with this kind of model loading. 

Hi @sergiopaniego ,

I am trying to run the example script available at this location of accelerate: [https://github.com/huggingface/accelerate/blob/61bcdaa45d35a77478decc78ce204b54dc108812/examples/torch_native_parallelism/nd_parallel.py](url)

But getting below error:
`  self._try_replicate_spec_for_scalar_tensor(
[rank0]:   File "/home/akuruvil/FT/ft_env/lib/python3.10/site-packages/torch/distributed/tensor/_dispatch.py", line 468, in _try_replicate_spec_for_scalar_tensor
[rank0]:     raise RuntimeError(
[rank0]: RuntimeError: aten._foreach_add_.List: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!`

Could you take a look at this. TP + DDP setup can be enabled without CP right. I have enabled only TP+DDP using this.

Hi @sergiopaniego  

Thanks for the detailed discussion in this thread — it helped clarify the correct setup.

From what I understand, the TP-only configuration that works reliably with
'Trainer' is the one using 'ParallelismConfig' via 'accelerator_config',
without using  device_map="auto" or 'init_empty_weights()'.

I’d be happy to turn this working approach into a clean, minimal
documentation example and add a short note explaining the common pitfalls
(e.g. why device_map="auto" doesn’t work with Trainer).

Before I start, could you confirm where you’d prefer this example to live
(Trainer docs vs examples/)?


I’ve been looking into this and ran into the same confusion around how tp_plan, device_mesh, and parallelism_config are supposed to work together when using Tensor Parallelism with Trainer.

I’d like to contribute a short documentation section (and possibly a minimal example) that focuses on TP-only training, without FSDP or DeepSpeed, and explains how these pieces fit together. I’ll use examples/3D_parallel.py as reference but keep things as simple and focused as possible.

Please let me know if this sounds like a reasonable direction, or if you’d prefer it structured differently.