mirror of https://github.com/hpcaitech/ColossalAI
1089 lines
33 KiB
Python
1089 lines
33 KiB
Python
import gc
|
|
from copy import deepcopy
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.testing import assert_close
|
|
|
|
import colossalai
|
|
from colossalai.cluster import ProcessGroupMesh
|
|
from colossalai.pipeline.p2p import PipelineP2PCommunication
|
|
from colossalai.pipeline.stage_manager import PipelineStageManager
|
|
from colossalai.testing import rerun_if_address_is_in_use, spawn
|
|
|
|
# info of model
|
|
IN_DIM = 8192
|
|
OUT_DIM = 8192
|
|
NUM_LAYER = 3
|
|
|
|
|
|
# A simple MLP
|
|
class MlpModel(nn.Module):
|
|
def __init__(self, in_dim=IN_DIM, out_dim=OUT_DIM, num_layers=NUM_LAYER):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList([nn.Linear(in_dim, out_dim, bias=None) for _ in range(num_layers)])
|
|
|
|
def forward(self, x):
|
|
for layer in self.layers:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
# Step1: dx = w*dy
|
|
def backward_b(loss, x, model):
|
|
print(f"Before bwd b: {torch.cuda.memory_allocated()/1024**3 :.3f} GB")
|
|
torch.autograd.backward(loss, inputs=x, retain_graph=True)
|
|
print(f"After bwd b: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
# Step1: dx = w*dy; for layer not last
|
|
def backward_b_not_last(tensors, grad, x, model):
|
|
print(f"Before bwd b: {torch.cuda.memory_allocated()/1024**3 :.3f} GB")
|
|
torch.autograd.backward(tensors=tensors, grad_tensors=grad, inputs=x, retain_graph=True)
|
|
print(f"After bwd b: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
def backward_w(loss, model):
|
|
print(f"Before bwd w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
torch.autograd.backward(loss, inputs=list(model.parameters()))
|
|
print(f"After bwd w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
# Step2: dummy dw = x*dy
|
|
def backward_w_not_last(tensors, grad, model):
|
|
print(f"Before bwd w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
torch.autograd.backward(tensors=tensors, grad_tensors=grad, inputs=list(model.parameters()))
|
|
print(f"After bwd w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
# In this poc, we check feasibility of spliting dx and dw in bwd propagation
|
|
def test_dx_dw_split():
|
|
device = "cuda:0"
|
|
model = nn.Linear(8, 8, bias=None).to(device=device)
|
|
print(f"model numel {get_model_numel(model)}") # 4GB
|
|
x = torch.rand(8, 8).to(device=device)
|
|
ref_model = deepcopy(model)
|
|
ref_x = x.clone()
|
|
|
|
# first step
|
|
x.requires_grad_()
|
|
loss = model(x).sum()
|
|
backward_b(loss, x, model)
|
|
for p in model.parameters():
|
|
assert p.grad is None
|
|
assert x.grad is not None
|
|
backward_w(loss, model)
|
|
for p in model.parameters():
|
|
assert p.grad is not None
|
|
|
|
# # second step
|
|
# loss = model(x).sum()
|
|
# backward_b(loss, x, model)
|
|
# backward_w(loss, model)
|
|
|
|
ref_x.requires_grad_()
|
|
ref_loss = ref_model(ref_x).sum()
|
|
ref_loss.backward()
|
|
|
|
assert torch.equal(x.grad, ref_x.grad)
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
assert torch.equal(p1.grad, p2.grad)
|
|
|
|
|
|
# In this poc, we check nsync of spliting dx and dw in bwd propagation in following order:
|
|
# fwd1 --> fwd2 --> dx1 --> dx2 --> dw1 --> dw2
|
|
def test_double_dx_dw_split_nsync():
|
|
device = "cuda:0"
|
|
model = nn.Linear(8, 8, bias=None).to(device=device)
|
|
# print(f"model numel {get_model_numel(model)}") # 4GB
|
|
x1 = torch.rand(8, 8).to(device=device)
|
|
x2 = torch.rand(8, 8).to(device=device)
|
|
ref_model = deepcopy(model)
|
|
ref_x1 = x1.clone()
|
|
ref_x2 = x2.clone()
|
|
|
|
# first step
|
|
x1.requires_grad_()
|
|
x2.requires_grad_()
|
|
ref_x1.requires_grad_()
|
|
ref_x2.requires_grad_()
|
|
|
|
# loss for dx_dw bwd
|
|
loss1 = model(x1).sum()
|
|
loss2 = model(x2).sum()
|
|
|
|
# loss for common bwd
|
|
ref_loss1 = ref_model(ref_x1).sum()
|
|
ref_loss2 = ref_model(ref_x2).sum()
|
|
|
|
# dx1
|
|
backward_b(loss1, x1, model)
|
|
for p in model.parameters():
|
|
assert p.grad is None
|
|
assert x1.grad is not None
|
|
|
|
# dx2
|
|
backward_b(loss2, x2, model)
|
|
|
|
# dw1
|
|
backward_w(loss1, model)
|
|
for p in model.parameters():
|
|
assert p.grad is not None
|
|
|
|
# common bwd 1
|
|
ref_loss1.backward()
|
|
|
|
# assert dx1 & dw1 == bwd 1
|
|
assert_close(x1.grad, ref_x1.grad)
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
# dw2
|
|
backward_w(loss2, model)
|
|
|
|
# common bwd 2
|
|
ref_loss2.backward()
|
|
|
|
# assert dx2 & dw2 == bwd 2
|
|
assert_close(x2.grad, ref_x2.grad)
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
print(f"bwd2:\n p1 {p1.grad},\n p2 {p2.grad}\n")
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
|
|
# In this poc, we check sync of spliting dx and dw in bwd propagation in following order:
|
|
# fwd1 --> fwd2 --> dx1 --> dw1 --> dx2 --> dw2
|
|
def test_double_dx_dw_split_sync():
|
|
device = "cuda:0"
|
|
model = nn.Linear(8, 8, bias=None).to(device=device)
|
|
x1 = torch.rand(8, 8).to(device=device)
|
|
x2 = torch.rand(8, 8).to(device=device)
|
|
|
|
ref_model = deepcopy(model)
|
|
ref_x1 = x1.clone()
|
|
ref_x2 = x2.clone()
|
|
|
|
x1.requires_grad_()
|
|
x2.requires_grad_()
|
|
ref_x1.requires_grad_()
|
|
ref_x2.requires_grad_()
|
|
|
|
############
|
|
# step1:
|
|
############
|
|
print(f"Step1\n")
|
|
|
|
# loss1
|
|
loss1 = model(x1).sum()
|
|
|
|
# ref_loss1
|
|
ref_loss1 = ref_model(ref_x1).sum()
|
|
|
|
# dx1
|
|
backward_b(loss1, x1, model)
|
|
for p in model.parameters():
|
|
assert p.grad is None
|
|
assert x1.grad is not None
|
|
|
|
# dw1
|
|
backward_w(loss1, model)
|
|
for p in model.parameters():
|
|
assert p.grad is not None
|
|
|
|
# common bwd 1
|
|
ref_loss1.backward()
|
|
|
|
# assert dx1 & dw1 == bwd 1
|
|
assert_close(x1.grad, ref_x1.grad)
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
############
|
|
# step2:
|
|
############
|
|
print(f"Step2\n")
|
|
|
|
# loss2
|
|
loss2 = model(x2).sum()
|
|
|
|
# ref_loss2
|
|
ref_loss2 = ref_model(ref_x2).sum()
|
|
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
# dx2
|
|
backward_b(loss2, x2, model)
|
|
|
|
# dw2
|
|
backward_w(loss2, model)
|
|
|
|
# common bwd 2
|
|
ref_loss2.backward()
|
|
|
|
# assert dx2 & dw2 == bwd 2
|
|
assert_close(x2.grad, ref_x2.grad)
|
|
for p1, p2 in zip(model.parameters(), ref_model.parameters()):
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
|
|
# In this poc, we check if a memory leak has occurred after del input & loss(with graph)
|
|
def mem_dx_dw():
|
|
device = "cuda:0"
|
|
print(f"Before init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
model = MlpModel().to(device=device)
|
|
print(f"model numel {get_model_numel(model)}") # 4GB
|
|
print(f"After init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
print(f"Before init x1&2&3: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
x1 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
x2 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
x3 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
|
|
x1.requires_grad_()
|
|
x2.requires_grad_()
|
|
x3.requires_grad_()
|
|
print(f"After init x1&2&3: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step1:
|
|
############
|
|
print(f"\nStep1")
|
|
|
|
# loss1
|
|
print(f"Before Fwd x1: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
loss1 = model(x1).sum()
|
|
print(f"After Fwd x1: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
print(f"Before loss1: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
print(f"After loss1: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
# dx1
|
|
backward_b(loss1, x1, model)
|
|
|
|
# dw1
|
|
backward_w(loss1, model)
|
|
|
|
del loss1, x1
|
|
# del x1
|
|
# del y1
|
|
print(f"After del x1&y1: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step2:
|
|
############
|
|
print(f"\nStep2")
|
|
|
|
# loss2
|
|
loss2 = model(x2).sum()
|
|
|
|
# dx2
|
|
backward_b(loss2, x2, model)
|
|
|
|
# dw2
|
|
backward_w(loss2, model)
|
|
|
|
del x2, loss2
|
|
# del x2
|
|
# del y2
|
|
print(f"After del x2&y2: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step3:
|
|
############
|
|
print(f"\nStep3")
|
|
|
|
# loss3
|
|
loss3 = model(x3).sum()
|
|
|
|
# dx2
|
|
backward_b(loss3, x3, model)
|
|
|
|
# dw2
|
|
backward_w(loss3, model)
|
|
|
|
# del x3
|
|
# del y3
|
|
del x3, loss3
|
|
|
|
print(f"After del x3&y3: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
param_ids = [id(p) for p in model.parameters()]
|
|
for obj in gc.get_objects():
|
|
if torch.is_tensor(obj) and id(obj) not in param_ids:
|
|
print(obj)
|
|
|
|
|
|
# In this poc, we check if a memory leak has occurred after del input & loss(with graph) & activation
|
|
def activation_dx_dw():
|
|
device = "cuda:0"
|
|
# model = nn.Linear(IN_DIM, OUT_DIM, bias=None).to(device=device)
|
|
print(f"Before init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
model = MlpModel().to(device=device)
|
|
x1 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
x2 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
x3 = torch.rand(IN_DIM, OUT_DIM).to(device=device)
|
|
|
|
x1.requires_grad_()
|
|
x2.requires_grad_()
|
|
x3.requires_grad_()
|
|
print(f"After init Model, x1,x2,x3: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step1:
|
|
############
|
|
print(f"\nStep1")
|
|
|
|
# loss1
|
|
output1 = model(x1)
|
|
loss1 = output1.sum()
|
|
|
|
# dx1
|
|
backward_b(loss1, x1, model)
|
|
|
|
# dw1
|
|
backward_w(loss1, model)
|
|
|
|
# del loss1, x1
|
|
del loss1, x1, output1
|
|
print(f"After del : {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step2:
|
|
############
|
|
print(f"\nStep2")
|
|
|
|
# loss2
|
|
output2 = model(x2)
|
|
loss2 = output2.sum()
|
|
|
|
# dx2
|
|
backward_b(loss2, x2, model)
|
|
|
|
# dw2
|
|
backward_w(loss2, model)
|
|
|
|
# del x2, loss2
|
|
del x2, loss2, output2
|
|
print(f"After del : {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
############
|
|
# step3:
|
|
############
|
|
print(f"\nStep3")
|
|
|
|
# loss3
|
|
output3 = model(x3)
|
|
loss3 = output3.sum()
|
|
|
|
# dx2
|
|
backward_b(loss3, x3, model)
|
|
|
|
# dw2
|
|
backward_w(loss3, model)
|
|
|
|
# del x3, loss3
|
|
del x3, loss3, output3
|
|
|
|
print(f"After del : {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
# In this poc, we apply model chunk instead of layer
|
|
def model_chunk_dx_dw():
|
|
device = "cuda:0"
|
|
num_layers = 4
|
|
print(f"Before init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
model = MlpModel(in_dim=4096, out_dim=4096, num_layers=num_layers).to(device=device)
|
|
input = torch.rand(4096, 4096, requires_grad=True).to(device=device)
|
|
|
|
input_base = input.clone()
|
|
|
|
model_base = deepcopy(model)
|
|
|
|
##########################
|
|
# Fwd bwd for dx dw
|
|
##########################
|
|
|
|
model_chunk_0 = torch.nn.Sequential() # for layer 1 & 2
|
|
model_chunk_1 = torch.nn.Sequential() # for layer 3 & 4
|
|
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx < 2:
|
|
model_chunk_0.append(sub_model)
|
|
else:
|
|
model_chunk_1.append(sub_model)
|
|
|
|
print(f"After init Model & input: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Step1:chunk 0 fwd
|
|
##########################
|
|
output1 = model_chunk_0(input)
|
|
|
|
# detach output1; then output1 for chunk 0, output1_dt for chunk 1;
|
|
output1_dt = output1.detach()
|
|
output1_dt.requires_grad_()
|
|
print(f"After chunk0 fwd (include detach output1): {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Step2:chunk 1 fwd
|
|
##########################
|
|
output2 = model_chunk_1(output1_dt)
|
|
|
|
print(f"After chunk1 fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Step3:chunk 1 bwd b: dx=w*dy & bwd w:dw=x*dy
|
|
##########################
|
|
loss = output2.mean()
|
|
backward_b(loss, output1_dt, model_chunk_1)
|
|
backward_w(loss, model_chunk_1)
|
|
|
|
print(f"After chunk1 bwd b & w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Step4:chunk 0 bwd b: dx=w*dy & bwd w:dw=x*dy
|
|
##########################
|
|
# dx = w*dy
|
|
backward_b_not_last(tensors=output1, grad=output1_dt.grad, x=input, model=model_chunk_0)
|
|
backward_w_not_last(tensors=output1, grad=output1_dt.grad, model=model_chunk_0)
|
|
|
|
print(f"After chunk0 bwd b & w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Fwd bwd for base
|
|
##########################
|
|
|
|
# fwd & bwd
|
|
output_base = model_base(input_base)
|
|
|
|
loss_base = output_base.mean()
|
|
|
|
loss_base.backward()
|
|
print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Assert param
|
|
##########################
|
|
|
|
assert_close(output2, output_base)
|
|
assert_close(output2.grad, output_base.grad)
|
|
|
|
for p1, p2 in zip(model.parameters(), model_base.parameters()):
|
|
assert_close(p1, p2)
|
|
assert_close(p1.grad, p2.grad)
|
|
|
|
del output1, output1_dt, output2, loss, loss_base, output_base
|
|
print(f"After del: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
|
|
# In this poc, we apply model chunk and a pp group for communication
|
|
def model_chunk_dx_dw_communication(
|
|
rank: int,
|
|
world_size: int,
|
|
port: int,
|
|
):
|
|
# init dist
|
|
colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost")
|
|
pg_mesh = ProcessGroupMesh(world_size)
|
|
stage_manager = PipelineStageManager(pg_mesh, pipeline_axis=0, enable_interleave=True, num_model_chunks=2)
|
|
rank = dist.get_rank()
|
|
comm = PipelineP2PCommunication(stage_manager, overlap_p2p=False)
|
|
|
|
print(f"{stage_manager.get_rank()}")
|
|
|
|
# init model and input
|
|
num_layers = 4
|
|
print(f"Before init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};")
|
|
model = MlpModel(in_dim=4096, out_dim=4096, num_layers=num_layers).to(rank)
|
|
input = torch.rand(4096, 4096, requires_grad=True).to(rank)
|
|
|
|
input_base = input.clone()
|
|
model_base = deepcopy(model)
|
|
|
|
if rank == 0:
|
|
model_chunk_0 = torch.nn.Sequential().to(rank) # for layer 1 & 2 on rank0
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx < 2:
|
|
model_chunk_0.append(sub_model)
|
|
else:
|
|
model_chunk_1 = torch.nn.Sequential().to(rank) # for layer 3 & 4 on rank1
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx >= 2:
|
|
model_chunk_1.append(sub_model)
|
|
|
|
print(
|
|
f"After init Model & input: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
##########################
|
|
# Step1:chunk 0 fwd
|
|
##########################
|
|
if rank == 0:
|
|
output1 = model_chunk_0(input)
|
|
print(
|
|
f"After chunk0 fwd (include detach output1): {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
# send y(output1_dt) to next stage
|
|
comm.send_forward(output1, stage_manager.get_next_rank())
|
|
|
|
##########################
|
|
# Step2:chunk 1 fwd
|
|
##########################
|
|
if rank == 1:
|
|
# recv y(output1_dt) from prev stage
|
|
output1_dt_rank1, wait_handles = comm.recv_forward(stage_manager.get_prev_rank())
|
|
output1_dt_rank1.requires_grad_()
|
|
output2 = model_chunk_1(output1_dt_rank1)
|
|
|
|
print(
|
|
f"After chunk1 fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
##########################
|
|
# Step3:chunk 1 on device_1 bwd b: dx=w*dy & bwd w:dw=x*dy
|
|
##########################
|
|
if rank == 1:
|
|
loss = output2.mean()
|
|
backward_b(loss, output1_dt_rank1, model_chunk_1)
|
|
backward_w(loss, model_chunk_1)
|
|
|
|
print(f"After chunk1 bwd b & w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
# send bwd output1_dt_rank1 from rank1 to rank 0
|
|
comm.send_backward(output1_dt_rank1.grad, stage_manager.get_prev_rank())
|
|
##########################
|
|
# Step4:chunk 0 on device_0 bwd b: dx=w*dy & bwd w:dw=x*dy
|
|
##########################
|
|
|
|
if rank == 0:
|
|
# recv bwd output1_dt_rank1 from rank1 to rank 0
|
|
output1_dt_rank0_grad, _ = comm.recv_backward(stage_manager.get_next_rank())
|
|
|
|
backward_b_not_last(tensors=output1, grad=output1_dt_rank0_grad, x=input, model=model_chunk_0)
|
|
backward_w_not_last(tensors=output1, grad=output1_dt_rank0_grad, model=model_chunk_0)
|
|
|
|
print(f"After chunk0 bwd b & w: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Fwd bwd for base
|
|
##########################
|
|
# fwd & bwd
|
|
output_base = model_base(input_base)
|
|
loss_base = output_base.mean()
|
|
loss_base.backward()
|
|
print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Assert param
|
|
##########################
|
|
# assert output
|
|
if rank == 1:
|
|
assert_close(output2, output_base)
|
|
assert_close(output2.grad, output_base.grad)
|
|
|
|
# assert model param & grad
|
|
if rank == 0:
|
|
count = 0
|
|
for (chunk_name, chunk_param), (base_name, base_param) in zip(
|
|
model_chunk_0.named_parameters(), model_base.named_parameters()
|
|
):
|
|
if count < 2:
|
|
assert_close(chunk_param, base_param)
|
|
assert_close(chunk_param.grad, base_param.grad)
|
|
count += 1
|
|
if rank == 1:
|
|
count = 0
|
|
for (chunk_name, chunk_param), (base_name, base_param) in zip(
|
|
model_chunk_1.named_parameters(), model_base.named_parameters()
|
|
):
|
|
if count >= 2:
|
|
assert_close(chunk_param, base_param)
|
|
assert_close(chunk_param.grad, base_param.grad)
|
|
count += 1
|
|
# clean memory
|
|
if rank == 0:
|
|
del output1, output1_dt_rank0_grad
|
|
if rank == 1:
|
|
del output2, loss, output1_dt_rank1
|
|
del loss_base, output_base
|
|
print(f"After del: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};")
|
|
|
|
|
|
# fwd schedule
|
|
def schedule_f(
|
|
stage_manager: PipelineStageManager,
|
|
comm: PipelineP2PCommunication,
|
|
input: torch.Tensor,
|
|
model_chunk: torch.nn.ModuleList,
|
|
model_chunk_id: int,
|
|
):
|
|
# chunk_id == 0
|
|
if model_chunk_id == 0:
|
|
# recv fwd from prev
|
|
if stage_manager.is_first_stage(ignore_chunk=True):
|
|
input = input # get local input
|
|
else:
|
|
prev_rank = stage_manager.get_prev_rank()
|
|
input, wait_handles = comm.recv_forward(prev_rank)
|
|
|
|
# fwd step
|
|
output = model_chunk[model_chunk_id](input)
|
|
|
|
# send fwd to next
|
|
if stage_manager.is_last_stage(ignore_chunk=True):
|
|
return input, output, None # return local output
|
|
else:
|
|
next_rank = stage_manager.get_next_rank()
|
|
comm.send_forward(output, next_rank)
|
|
|
|
# chunk_id == 1
|
|
if model_chunk_id == 1:
|
|
# recv fwd from next
|
|
if stage_manager.is_last_stage(ignore_chunk=True):
|
|
input = input # get local input
|
|
else:
|
|
next_rank = stage_manager.get_next_rank()
|
|
input, wait_handles = comm.recv_forward(next_rank)
|
|
|
|
# fwd step
|
|
output = model_chunk[model_chunk_id](input)
|
|
|
|
# send fwd to prev
|
|
if stage_manager.is_first_stage(ignore_chunk=True):
|
|
loss = output.mean()
|
|
return input, output, loss # return local output
|
|
else:
|
|
prev_rank = stage_manager.get_prev_rank()
|
|
comm.send_forward(output, prev_rank)
|
|
return input, output, None
|
|
|
|
|
|
# bwd b schedule
|
|
def schedule_b(
|
|
stage_manager: PipelineStageManager,
|
|
comm: PipelineP2PCommunication,
|
|
input: torch.Tensor, # x
|
|
output: torch.Tensor, # y
|
|
output_grad: torch.Tensor, # dy
|
|
model_chunk: torch.nn.ModuleList,
|
|
model_chunk_id: int,
|
|
):
|
|
# chunk_id == 0
|
|
if model_chunk_id == 0:
|
|
|
|
# recv bwd from next
|
|
if stage_manager.is_last_stage(ignore_chunk=True):
|
|
output_grad = output_grad # get dy from local
|
|
else:
|
|
next_rank = stage_manager.get_next_rank()
|
|
output_grad, _ = comm.recv_backward(next_rank)
|
|
|
|
# bwd step
|
|
backward_b_not_last(tensors=output, grad=output_grad, x=input, model=model_chunk[model_chunk_id])
|
|
backward_w_not_last(tensors=output, grad=output_grad, model=model_chunk[model_chunk_id])
|
|
|
|
# send bwd to prev
|
|
if stage_manager.is_first_stage(ignore_chunk=True):
|
|
return input.grad
|
|
else:
|
|
prev_rank = stage_manager.get_prev_rank()
|
|
comm.send_backward(input.grad, prev_rank)
|
|
|
|
# chunk_id == 1
|
|
if model_chunk_id == 1:
|
|
# recv bwd from prev
|
|
if stage_manager.is_first_stage(ignore_chunk=True):
|
|
output_grad = output_grad
|
|
else:
|
|
prev_rank = stage_manager.get_prev_rank()
|
|
output_grad, _ = comm.recv_backward(next_rank=prev_rank)
|
|
|
|
# bwd step
|
|
if stage_manager.is_first_stage(ignore_chunk=True):
|
|
backward_b(loss=output_grad, x=input, model=model_chunk[model_chunk_id])
|
|
backward_w(loss=output_grad, model=model_chunk[model_chunk_id])
|
|
else:
|
|
# commom bwd step
|
|
backward_b_not_last(tensors=output, grad=output_grad, x=input, model=model_chunk[model_chunk_id])
|
|
backward_w_not_last(tensors=output, grad=output_grad, model=model_chunk[model_chunk_id])
|
|
|
|
# send bwd to next
|
|
if stage_manager.is_last_stage(ignore_chunk=True):
|
|
return input.grad
|
|
else:
|
|
next_rank = stage_manager.get_next_rank()
|
|
comm.send_backward(input.grad, next_rank)
|
|
|
|
return input.grad
|
|
|
|
|
|
# bwd w schedule (dw already splite in schedule b)
|
|
def schedule_w():
|
|
pass
|
|
|
|
|
|
# In this poc, we apply a scheduling method for each rank: schedule_f --> schedule_b --> schedule_w
|
|
def model_chunk_dx_dw_comm_interleaved(
|
|
rank: int,
|
|
world_size: int,
|
|
port: int,
|
|
):
|
|
# init dist
|
|
colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost")
|
|
pg_mesh = ProcessGroupMesh(world_size)
|
|
stage_manager = PipelineStageManager(pg_mesh, pipeline_axis=0, enable_interleave=True, num_model_chunks=world_size)
|
|
rank = dist.get_rank()
|
|
comm = PipelineP2PCommunication(stage_manager, overlap_p2p=False)
|
|
|
|
# init model and input
|
|
num_layers = 8
|
|
in_dim = out_dim = 2048
|
|
print(f"Before init Model: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};")
|
|
model = MlpModel(in_dim=in_dim, out_dim=out_dim, num_layers=num_layers).to(rank)
|
|
input0 = torch.rand(in_dim, out_dim, requires_grad=True).to(rank)
|
|
|
|
input_base = input0.clone()
|
|
model_base = deepcopy(model)
|
|
|
|
if rank == 0:
|
|
# layer 0 & 7 to chunk 0 on rank0
|
|
chunk_0 = torch.nn.ModuleList().to(rank)
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx == 0 or idx == 7:
|
|
chunk_0.append(sub_model)
|
|
elif rank == 1:
|
|
# layer 1 & 6 to chunk 1 on rank1
|
|
chunk_1 = torch.nn.ModuleList().to(rank)
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx == 1 or idx == 6:
|
|
chunk_1.append(sub_model)
|
|
elif rank == 2:
|
|
# layer 2 & 5 to chunk 2 on rank2
|
|
chunk_2 = torch.nn.ModuleList().to(rank)
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx == 2 or idx == 5:
|
|
chunk_2.append(sub_model)
|
|
else:
|
|
# layer 3 & 4 to chunk 3 on rank3
|
|
chunk_3 = torch.nn.Sequential().to(rank)
|
|
for idx, sub_model in enumerate(model.layers):
|
|
if idx == 3 or idx == 4:
|
|
chunk_3.append(sub_model)
|
|
|
|
print(
|
|
f"After init Model & input: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
# buffer use to save input and output
|
|
|
|
##########################
|
|
# Step1: fwd
|
|
##########################
|
|
######
|
|
# fwd 1->4
|
|
######
|
|
# chunk 0 id 0 (layer 0) fwd
|
|
if rank == 0:
|
|
chunk_id = 0
|
|
input0, output0, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input0,
|
|
model_chunk=chunk_0,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 0 id 0 (layer 0)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
# chunk 1 id 0 (layer 1) fwd
|
|
if rank == 1:
|
|
chunk_id = 0
|
|
input1, output1, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_1,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 1 id 0 (layer 1)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
# chunk 2 id 0 (layer 2) fwd
|
|
if rank == 2:
|
|
chunk_id = 0
|
|
input2, output2, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_2,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 2 id 0 (layer 2)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
# chunk 3 id 0 (layer 3) fwd
|
|
if rank == 3:
|
|
chunk_id = 0
|
|
input3, output3, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_3,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 3 id 0 (layer 3)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
######
|
|
# fwd 4->1
|
|
######
|
|
|
|
if rank == 3:
|
|
chunk_id = 1
|
|
input4, output4, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=output3,
|
|
model_chunk=chunk_3,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 3 id 1 (layer 4)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
if rank == 2:
|
|
chunk_id = 1
|
|
input5, output5, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_2,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 2 id 1 (layer 5)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
if rank == 1:
|
|
chunk_id = 1
|
|
input6, output6, _ = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_1,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
print(
|
|
f"chunk 1 id 1 (layer 6)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
if rank == 0:
|
|
chunk_id = 1
|
|
input7, output7, loss = schedule_f(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=None,
|
|
model_chunk=chunk_0,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
# print(f"fwd output {output7}")
|
|
print(
|
|
f"chunk 0 id 1 (layer 7)fwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};"
|
|
)
|
|
|
|
##########################
|
|
# Step2: bwd
|
|
##########################
|
|
######
|
|
# bwd rank 4->1
|
|
######
|
|
# chunk 0 id 1 (layer 7) bwd
|
|
if rank == 0:
|
|
chunk_id = 1
|
|
input_grad7 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input7, # x
|
|
output=output7, # y
|
|
output_grad=loss, # dy
|
|
model_chunk=chunk_0,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# # chunk 1 id 1 (layer 6) bwd
|
|
if rank == 1:
|
|
chunk_id = 1
|
|
input_grad6 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input6, # x
|
|
output=output6, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_1,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# chunk 2 id 1 (layer 5) bwd
|
|
if rank == 2:
|
|
chunk_id = 1
|
|
input_grad5 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input5, # x
|
|
output=output5, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_2,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# chunk 3 id 1 (layer 4) bwd
|
|
if rank == 3:
|
|
chunk_id = 1
|
|
input_grad4 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input4, # x
|
|
output=output4, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_3,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
######
|
|
# bwd rank 1->4
|
|
######
|
|
|
|
# chunk 3 id 0 (layer 3) bwd
|
|
if rank == 3:
|
|
chunk_id = 0
|
|
input_grad3 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input3, # x
|
|
output=output3, # y
|
|
output_grad=input_grad4, # dy
|
|
model_chunk=chunk_3,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# chunk 2 id 0 (layer 2) bwd
|
|
if rank == 2:
|
|
chunk_id = 0
|
|
input_grad2 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input2, # x
|
|
output=output2, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_2,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# chunk 1 id 0 (layer 1) bwd
|
|
if rank == 1:
|
|
chunk_id = 0
|
|
input_grad1 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input1, # x
|
|
output=output1, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_1,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
# chunk 0 id 0 (layer 0) bwd
|
|
if rank == 0:
|
|
chunk_id = 0
|
|
input_grad0 = schedule_b(
|
|
stage_manager=stage_manager,
|
|
comm=comm,
|
|
input=input0, # x
|
|
output=output0, # y
|
|
output_grad=None, # dy
|
|
model_chunk=chunk_0,
|
|
model_chunk_id=chunk_id,
|
|
)
|
|
|
|
##########################
|
|
# Fwd bwd for base
|
|
##########################
|
|
# fwd & bwd
|
|
output_base = model_base(input_base)
|
|
loss_base = output_base.mean()
|
|
loss_base.backward()
|
|
print(f"After base fwd & bwd: {torch.cuda.memory_allocated()/1024**3 :.3f} GB;")
|
|
|
|
##########################
|
|
# Assert close
|
|
##########################
|
|
# assert output
|
|
if rank == 0:
|
|
assert_close(output7, output_base)
|
|
|
|
# assert weight
|
|
if rank == 0:
|
|
# layer 0
|
|
assert_close(chunk_0[0].weight, model_base.layers[0].weight)
|
|
assert_close(chunk_0[0].weight.grad, model_base.layers[0].weight.grad)
|
|
# layer 7
|
|
assert_close(chunk_0[1].weight, model_base.layers[7].weight)
|
|
assert_close(chunk_0[1].weight.grad, model_base.layers[7].weight.grad)
|
|
if rank == 1:
|
|
# layer 1
|
|
assert_close(chunk_1[0].weight, model_base.layers[1].weight)
|
|
assert_close(chunk_1[0].weight.grad, model_base.layers[1].weight.grad)
|
|
# layer 6
|
|
assert_close(chunk_1[1].weight, model_base.layers[6].weight)
|
|
assert_close(chunk_1[1].weight.grad, model_base.layers[6].weight.grad)
|
|
|
|
if rank == 2:
|
|
# layer 2
|
|
assert_close(chunk_2[0].weight, model_base.layers[2].weight)
|
|
assert_close(chunk_2[0].weight.grad, model_base.layers[2].weight.grad)
|
|
# layer 5
|
|
assert_close(chunk_2[1].weight, model_base.layers[5].weight)
|
|
assert_close(chunk_2[1].weight.grad, model_base.layers[5].weight.grad)
|
|
|
|
if rank == 3:
|
|
# layer 3
|
|
assert_close(chunk_3[0].weight, model_base.layers[3].weight)
|
|
assert_close(chunk_3[0].weight.grad, model_base.layers[3].weight.grad)
|
|
# layer 4
|
|
assert_close(chunk_3[1].weight, model_base.layers[4].weight)
|
|
assert_close(chunk_3[1].weight.grad, model_base.layers[4].weight.grad)
|
|
|
|
# clean memory
|
|
if rank == 0:
|
|
del input0, output0, input_grad0, input7, output7, input_grad7, loss
|
|
if rank == 1:
|
|
del input1, output1, input_grad1, input6, output6, input_grad6
|
|
if rank == 2:
|
|
del input2, output2, input_grad2, input5, output5, input_grad5
|
|
if rank == 3:
|
|
del input3, output3, input_grad3, input4, output4, input_grad4
|
|
del loss_base, output_base
|
|
|
|
print(f"After del: {torch.cuda.memory_allocated()/1024**3 :.3f} GB on device {stage_manager.get_rank()};")
|
|
|
|
|
|
@rerun_if_address_is_in_use()
|
|
def test_dx_dw_dist():
|
|
spawn(
|
|
model_chunk_dx_dw_comm_interleaved,
|
|
nprocs=4,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_dx_dw_dist()
|