Skip to content

Commit ecd951d

Browse files
wanchaolwanchaol
authored
Update TP examples to align with tutorials (#1243)
as titled
1 parent 7df10c2 commit ecd951d

File tree

6 files changed

+685
-71
lines changed

6 files changed

+685
-71
lines changed

‎.github/workflows/main_distributed.yaml‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
name: Run Distributed Examples
2+
3+
on:
4+
push:
5+
branches: [ main ]
6+
pull_request:
7+
branches: [ main ]
8+
schedule:
9+
# Every day at 3:00am
10+
- cron: '0 3 * * *'
11+
12+
13+
jobs:
14+
test:
15+
16+
runs-on: 4-core-ubuntu-gpu-t4
17+
18+
steps:
19+
- uses: actions/checkout@v2
20+
- name: Set up Python 3.8
21+
uses: actions/setup-python@v2
22+
with:
23+
python-version: 3.8
24+
- name: Install PyTorch
25+
run: |
26+
python -m pip install --upgrade pip
27+
pip install --pre torch -f https://download.pytorch.org/whl/nightly/cu118/torch_nightly.html
28+
- name: Run Tests
29+
run: |
30+
./run_distributed_examples.sh "run_all,clean"
31+
- name: Open issue on failure
32+
if: ${{ failure() && github.event_name == 'schedule' }}
33+
uses: rishabhgupta/git-action-issue@v2
34+
with:
35+
token: ${{ secrets.GITHUB_TOKEN }}
36+
title: Daily CI failed
37+
body: Commit ${{ github.sha }} daily scheduled [CI run](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}) failed, please check why
38+
assignees: ''

‎distributed/tensor_parallelism/README.md‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
1-
# PyTorch Tensor Parallelism for distributed training
1+
# PyTorch native Tensor Parallel for distributed training
22

3-
This example demonstrates SPMD Megatron-LM style tensor parallel by using
4-
PyTorch native Tensor Parallelism APIs, which include:
3+
This example demonstrates SPMD Megatron-LM style Tensor Parallel by using
4+
PyTorch native Tensor Parallel APIs, which include:
55

6-
1. High-level APIs for module-level parallelism with a dummy MLP model.
7-
2. Model agnostic ops for `DistributedTensor`, such as `Linear` and `RELU`.
8-
3. A E2E demo of tensor parallel for a given toy model (Forward/backward + optimization).
6+
1. Simple module-level Tensor Parallelism on a dummy MLP model.
7+
2. Simple module-level Tensor Parallelism with Sequence Parallel inputs/outputs on a dummy MLP model.
8+
3. A E2E demo of Fully Sharded Data Parallel + Tensor Parallel (with Sequence Parallel) on a example Llama2 model.
99

10-
More details about the design can be found:
11-
https://github.com/pytorch/pytorch/issues/89884
10+
More details about the PyTorch native Tensor Parallel APIs, please see PyTorch docs:
11+
https://pytorch.org/docs/stable/distributed.tensor.parallel.html
1212

1313
```
1414
pip install -r requirements.txt

‎distributed/tensor_parallelism/fsdp_tp_example.py‎

Lines changed: 80 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,12 @@
11
import sys
2+
import os
23
import torch
34
import torch.distributed as dist
45
import torch.nn as nn
56
import torch.nn.functional as F
67

7-
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
8-
from torch.distributed.tensor.parallel import (
9-
parallelize_module,
10-
ColwiseParallel,
11-
RowwiseParallel,
12-
)
13-
14-
import os
158
from log_utils import rank_log, get_logger, verify_min_gpu_count
169

17-
1810
# ---- GPU check ------------
1911
_min_gpu_count = 4
2012

@@ -23,13 +15,24 @@
2315
sys.exit()
2416
# ---------------------------
2517

26-
from torch.distributed._tensor.device_mesh import init_device_mesh
18+
from llama2_model import Transformer, ModelArgs
19+
20+
from torch.distributed.device_mesh import init_device_mesh
21+
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
22+
from torch.distributed._tensor import Shard, Replicate
23+
from torch.distributed.tensor.parallel import (
24+
parallelize_module,
25+
ColwiseParallel,
26+
RowwiseParallel,
27+
PrepareModuleInput,
28+
SequenceParallel
29+
)
2730

2831

2932
"""
3033
This is the script to test 2D Parallel which combines Tensor/Sequence
31-
parallel with Fully Sharded Data Parallel (TP/SP + FSDP) on a toy model
32-
in the SPMD style. We show an E2E working flow from forward, backward
34+
parallel with Fully Sharded Data Parallel (TP/SP + FSDP) on a example
35+
Llama2 model. We show an E2E working flow from forward, backward
3336
and optimization.
3437
3538
We enabled Fully Sharded Data Parallel + Tensor Parallel in
@@ -53,41 +56,10 @@
5356
[0, 8, ..., 8N-8], [1, 9, ..., 8N-7], ..., [7, 15, ..., 8N-1]
5457
======================================================================
5558
56-
More details can be seen in the slide:
57-
https://docs.google.com/presentation/d/17g6WqrO00rP3MsxbRENsPpjrlSkwiA_QB4r93_eB5is/
59+
More details can be seen in the PyTorch tutorials:
60+
https://pytorch.org/tutorials/intermediate/TP_tutorial.html
5861
"""
5962

60-
61-
def find_multiple(n: int, k: int) -> int:
62-
"""function to find resizing multiple for SwiGLU MLP"""
63-
if n % k == 0:
64-
return n
65-
return n + k - (n % k)
66-
67-
68-
class MLP_swiglu(nn.Module):
69-
"""SwiGLU to showcase a Llama style MLP model"""
70-
71-
def __init__(self, mlp_dim: int = 1024) -> None:
72-
super().__init__()
73-
hidden_dim = 4 * mlp_dim
74-
scaled_hidden = int(2 * hidden_dim / 3)
75-
rounded_hidden = find_multiple(scaled_hidden, 256)
76-
77-
self.in_proj = nn.Linear(mlp_dim, rounded_hidden, bias=False)
78-
self.gate_proj = nn.Linear(mlp_dim, rounded_hidden, bias=False)
79-
self.out_proj = nn.Linear(rounded_hidden, mlp_dim, bias=False)
80-
81-
def forward(self, x: torch.Tensor) -> torch.Tensor:
82-
x = F.silu(self.in_proj(x)) * self.gate_proj(x)
83-
x = self.out_proj(x)
84-
return x
85-
86-
87-
"""
88-
Main body of the demo of a basic version of tensor parallel by using
89-
PyTorch native APIs.
90-
"""
9163
tp_size = 2
9264
logger = get_logger()
9365

@@ -120,26 +92,72 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
12092
# to mimic the behavior of the dataloader.
12193
dp_rank = dp_mesh.get_local_rank()
12294

123-
# create model and move it to GPU with id rank
124-
_mlp_dim = 1024
125-
base_model_swiglu = MLP_swiglu(mlp_dim=_mlp_dim).to("cuda")
126-
127-
128-
# Custom parallelization plan for the swiglu MLP model
129-
custom_tp_model = parallelize_module(
130-
module=base_model_swiglu,
131-
device_mesh=tp_mesh,
132-
parallelize_plan={
133-
"in_proj": ColwiseParallel(),
134-
"gate_proj": ColwiseParallel(),
135-
"out_proj": RowwiseParallel(),
136-
},
95+
# create model and move it to GPU - init"cuda"_mesh has already mapped GPU ids.
96+
simple_llama2_config = ModelArgs(dim=256, n_layers=2, n_heads=16, vocab_size=32000)
97+
98+
model = Transformer.from_model_args(simple_llama2_config).to("cuda")
99+
100+
# init model weights
101+
model.init_weights()
102+
103+
# parallelize the first embedding and the last linear out projection
104+
model = parallelize_module(
105+
model,
106+
tp_mesh,
107+
{
108+
"tok_embeddings": RowwiseParallel(
109+
input_layouts=Replicate(),
110+
),
111+
"output": ColwiseParallel(
112+
input_layouts=Shard(1),
113+
output_layouts=Replicate()
114+
),
115+
"norm": SequenceParallel(),
116+
"layers.0": PrepareModuleInput(
117+
input_layouts=(Replicate(), None),
118+
desired_input_layouts=(Shard(1), None),
119+
use_local_output=True,
120+
),
121+
}
137122
)
138123

139-
rank_log(_rank, logger, f"Model after parallelization {custom_tp_model=}\n")
124+
for layer_id, transformer_block in enumerate(model.layers):
125+
layer_tp_plan = {
126+
"attention": PrepareModuleInput(
127+
input_layouts=(Shard(1), None),
128+
desired_input_layouts=(Replicate(), None),
129+
),
130+
"attention.wq": ColwiseParallel(),
131+
"attention.wk": ColwiseParallel(),
132+
"attention.wv": ColwiseParallel(),
133+
"attention.wo": RowwiseParallel(output_layouts=Shard(1)),
134+
"attention_norm": SequenceParallel(),
135+
"feed_forward": PrepareModuleInput(
136+
input_layouts=(Shard(1),),
137+
desired_input_layouts=(Replicate(),),
138+
),
139+
"feed_forward.w1": ColwiseParallel(),
140+
"feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
141+
"feed_forward.w3": ColwiseParallel(),
142+
"ffn_norm": SequenceParallel(),
143+
}
144+
145+
# Adjust attention module to use the local number of heads
146+
attn_layer = transformer_block.attention
147+
attn_layer.n_heads = attn_layer.n_heads // tp_mesh.size()
148+
attn_layer.n_kv_heads = attn_layer.n_kv_heads // tp_mesh.size()
149+
150+
# Custom parallelization plan for the model
151+
parallelize_module(
152+
module=transformer_block,
153+
device_mesh=tp_mesh,
154+
parallelize_plan=layer_tp_plan
155+
)
140156

141157
# Init FSDP using the dp device mesh
142-
sharded_model = FSDP(custom_tp_model, device_mesh=dp_mesh, use_orig_params=True)
158+
sharded_model = FSDP(model, device_mesh=dp_mesh, use_orig_params=True)
159+
160+
rank_log(_rank, logger, f"Model after parallelization {sharded_model=}\n")
143161

144162
# Create an optimizer for the parallelized and sharded model.
145163
lr = 3e-3
@@ -156,7 +174,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
156174
for i in range(num_iterations):
157175
# seeding with dp_rank to ensure identical inputs for TP groups
158176
torch.manual_seed(i + dp_rank)
159-
inp = torch.rand(batch_size, _mlp_dim, device="cuda")
177+
inp = torch.randint(32000, (8, 256), device="cuda")
160178

161179
output = sharded_model(inp)
162180
output.sum().backward()

0 commit comments

Comments
 (0)