[Python] Why doesn't the loss decrease when I use torch.distributed.pipelining to train the...

My complete code is as follows：

import torch
import torch.nn as nn
from dataclasses import dataclass
import os
import torch.distributed as dist
from torch.distributed.pipelining import pipeline, SplitPoint, PipelineStage, ScheduleGPipe

@dataclass
class ModelArgs:
dim: int = 512
n_layers: int = 8
n_heads: int = 8
vocab_size: int = 10000

class Transformer(nn.Module):
def __init__(self, model_args: ModelArgs):
super().__init__()

self.tok_embeddings = nn.Embedding(model_args.vocab_size, model_args.dim)

# Using a ModuleDict lets us delete layers witout affecting names,
# ensuring checkpoints will correctly save and load.
self.layers = torch.nn.ModuleDict()
for layer_id in range(model_args.n_layers):
self.layers[str(layer_id)] = nn.TransformerDecoderLayer(model_args.dim, model_args.n_heads)

self.norm = nn.LayerNorm(model_args.dim)
self.output = nn.Linear(model_args.dim, model_args.vocab_size)

def forward(self, tokens: torch.Tensor):
# Handling layers being 'None' at runtime enables easy pipeline splitting
h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens

for layer in self.layers.values():
h = layer(h, h)

h = self.norm(h) if self.norm else h
output = self.output(h).float() if self.output else h
return output



global rank, device, pp_group, stage_index, num_stages
def init_distributed():
global rank, device, pp_group, stage_index, num_stages
rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
device = torch.device(f"cuda:{rank}") if torch.cuda.is_available() else torch.device("cpu")
dist.init_process_group()

# This group can be a sub-group in the N-D parallel case
pp_group = dist.new_group()
stage_index = rank
num_stages = world_size

def manual_model_split(model, example_input_microbatch, model_args) -> PipelineStage:
if stage_index == 0:
# prepare the first stage model
for i in range(4, 8):
del model.layers[str(i)]
model.norm = None
model.output = None
stage_input_microbatch = example_input_microbatch

elif stage_index == 1:
# prepare the second stage model
for i in range(4):
del model.layers[str(i)]
model.tok_embeddings = None
stage_input_microbatch = torch.randn(example_input_microbatch.shape[0], example_input_microbatch.shape[1], model_args.dim)

stage = PipelineStage(
model,
stage_index,
num_stages,
device,
input_args=stage_input_microbatch,
)
return stage

if __name__ == "__main__":
init_distributed()
num_microbatches = 4
model_args = ModelArgs()
model = Transformer(model_args)

# Dummy data
x = torch.ones(32, 500, dtype=torch.long)
y = torch.randint(0, model_args.vocab_size, (32, 500), dtype=torch.long)
example_input_microbatch = x.chunk(num_microbatches)[0]

# Option 1: Manual model splitting
stage = manual_model_split(model, example_input_microbatch, model_args)

# Option 2: Tracer model splitting
# stage = tracer_model_split(model, example_input_microbatch)

x = x.to(device)
y = y.to(device)

def tokenwise_loss_fn(outputs, targets):
loss_fn = nn.CrossEntropyLoss()
outputs = outputs.view(-1, model_args.vocab_size)
targets = targets.view(-1)
return loss_fn(outputs, targets)

schedule = ScheduleGPipe(stage, n_microbatches=num_microbatches, loss_fn=tokenwise_loss_fn)

for i in range(30):
if rank == 0:
schedule.step(x)
elif rank == 1:
losses = []
output = schedule.step(target=y, losses=losses)
print(losses)
dist.destroy_process_group()


In fact, this is almost identical to the official documentation given by pytorchhttps://pytorch.org/tutorials/intermediate/pipelining_tutorial.html.

However, when I train for multiple iterations, the loss does not decrease, and it seems that the gradients are not being successfully backpropagated.

[tensor(9.3873, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3760, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3691, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3736, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3819, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3733, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3657, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3706, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3854, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3731, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3607, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3721, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3826, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3835, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3653, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3716, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3833, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3808, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3646, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3771, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3846, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3766, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3658, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3703, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3861, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3805, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3666, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3754, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3862, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3818, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3660, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3697, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3884, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3671, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3722, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3714, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3851, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3783, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3629, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3778, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3850, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3752, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3681, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3722, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3819, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3801, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3680, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3772, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3877, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3742, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3681, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3733, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3833, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3723, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3641, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3819, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3840, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3784, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3694, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3721, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3860, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3736, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3761, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3803, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3877, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3780, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3665, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3701, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3832, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3754, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3593, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3785, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3811, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3751, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3701, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3709, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3828, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3734, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3697, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3715, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3834, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3821, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3702, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3759, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3836, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3745, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3711, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3740, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3912, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3769, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3701, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3720, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3833, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3719, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3691, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3710, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3910, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3804, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3656, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3670, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3843, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3781, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3644, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3727, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3898, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3729, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3655, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3719, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3828, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3692, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3732, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3655, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3855, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3791, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3697, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3767, device='cuda:1', grad_fn=<NllLossBackward0>)]
[tensor(9.3871, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3781, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3682, device='cuda:1', grad_fn=<NllLossBackward0>), tensor(9.3789, device='cuda:1', grad_fn=<NllLossBackward0>)]


Does anyone have any ideas on how to solve this issue? Any help would be greatly appreciated.

Continue reading...