From e0949867f51c78abab9007affb873bde7ecbb38b Mon Sep 17 00:00:00 2001 From: xrsrke Date: Tue, 5 Dec 2023 14:07:43 +0700 Subject: [PATCH] [Refactor] Refactor MoE with data parallelism's tests --- pipegoose/testing/utils.py | 15 +++++++++++- tests/nn/data_parallel/test_data_parallel.py | 11 +++------ .../test_hybrid_expert_parallel.py | 23 ++++++++++--------- 3 files changed, 29 insertions(+), 20 deletions(-) diff --git a/pipegoose/testing/utils.py b/pipegoose/testing/utils.py index 7542497..097e777 100644 --- a/pipegoose/testing/utils.py +++ b/pipegoose/testing/utils.py @@ -2,7 +2,7 @@ import random import socket from functools import partial -from typing import Callable +from typing import Callable, Tuple import pytest import torch @@ -118,3 +118,16 @@ def calculate_parameter_similarity(module1: nn.Module, module2: nn.Module, rtol: def count_model_parameters(model): return sum(p.numel() for p in model.parameters()) + + +def get_microbatch( + inputs, labels, parallel_context: ParallelContext, parallel_mode: ParallelMode +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + local_rank = parallel_context.get_local_rank(parallel_mode) + world_size = parallel_context.get_world_size(parallel_mode) + + input_chunks = torch.chunk(inputs["input_ids"], chunks=world_size, dim=0) + attention_chunks = torch.chunk(inputs["attention_mask"], chunks=world_size, dim=0) + label_chunks = torch.chunk(labels, chunks=world_size, dim=0) + + return input_chunks[local_rank], attention_chunks[local_rank], label_chunks[local_rank] diff --git a/tests/nn/data_parallel/test_data_parallel.py b/tests/nn/data_parallel/test_data_parallel.py index 11bf58c..e296379 100644 --- a/tests/nn/data_parallel/test_data_parallel.py +++ b/tests/nn/data_parallel/test_data_parallel.py @@ -9,6 +9,7 @@ from pipegoose.nn import DataParallel from pipegoose.testing.utils import ( calculate_parameter_similarity, + get_microbatch, init_parallel_context, skip_if_no_cuda, spawn, @@ -89,13 +90,6 @@ def test_parallelize_a_transformer_and_inference(model, tokenizer, data_parallel def run_backward_a_parallelized_transformers( rank, world_size, port, tensor_parallel_size, pipeline_parallel_size, data_parallel_size, kwargs ): - def get_microbatch(inputs, labels): - local_rank = parallel_context.get_local_rank(ParallelMode.DATA) - input_chunks = torch.chunk(inputs["input_ids"], chunks=world_size, dim=0) - attention_chunks = torch.chunk(inputs["attention_mask"], chunks=world_size, dim=0) - label_chunks = torch.chunk(labels, chunks=world_size, dim=0) - return input_chunks[local_rank], attention_chunks[local_rank], label_chunks[local_rank] - model = deepcopy(kwargs["model"]) UPDATED_MODEL = deepcopy(kwargs["updated_model"]) LR = kwargs["lr"] @@ -106,7 +100,8 @@ def get_microbatch(inputs, labels): rank, world_size, port, tensor_parallel_size, pipeline_parallel_size, data_parallel_size ) - input_ids, attention_mask, labels = get_microbatch(inputs, labels) + # NOTE: each model replicas only train on a subset of data + input_ids, attention_mask, labels = get_microbatch(inputs, labels, parallel_context, ParallelMode.DATA) parallelized_model = DataParallel(model, parallel_context).parallelize() optim = SGD(parallelized_model.parameters(), lr=LR) diff --git a/tests/nn/expert_parallel/test_hybrid_expert_parallel.py b/tests/nn/expert_parallel/test_hybrid_expert_parallel.py index 045dce2..facc6bd 100644 --- a/tests/nn/expert_parallel/test_hybrid_expert_parallel.py +++ b/tests/nn/expert_parallel/test_hybrid_expert_parallel.py @@ -13,7 +13,7 @@ from pipegoose.nn.data_parallel.data_parallel import DataParallel from pipegoose.nn.expert_parallel.loss import ExpertLoss from pipegoose.nn.expert_parallel.routers import SwitchNoisePolicy, Top1Router -from pipegoose.testing.utils import init_parallel_context, spawn +from pipegoose.testing.utils import get_microbatch, init_parallel_context, spawn MODEL_NAME = "bigscience/bloom-560m" @@ -57,15 +57,20 @@ def run_expert_parallel_with_data_parallel( pipeline_parallel_size, data_parallel_size, ) + # NOTE: each model replicas only train on a subset of data + input_ids, attention_mask, labels = get_microbatch( + kwargs["input"], kwargs["labels"], parallel_context, ParallelMode.EXPERT + ) + loss_func = ExpertLoss(nn.CrossEntropyLoss()) + model = ExpertParallel(model, NUM_EXPERTS, mapping=mapping, router=router, parallel_context=parallel_context).parallelize() model = DataParallel(model, parallel_context).parallelize() - loss_func = ExpertLoss(nn.CrossEntropyLoss()) optim = Adam(model.parameters(), lr=1e-3) - outputs = model(**kwargs["input"]) + outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits[..., :-1, :].view(-1, outputs.logits.shape[-1]) - labels = kwargs["labels"][..., 1:].view(-1).to(logits.device) + labels = labels[..., 1:].view(-1).to(logits.device) loss = loss_func(logits, labels) optim.zero_grad() @@ -76,11 +81,7 @@ def run_expert_parallel_with_data_parallel( expert_grads = torch.chunk(expert_grads, chunks=data_parallel_size, dim=0) # NOTE: check if expert grads are the same across data parallel dimension - assert all( - torch.all(torch.eq(expert_grads[i], expert_grads[j])) - for i in range(len(expert_grads)) - for j in range(i + 1, len(expert_grads)) - ) + assert torch.allclose(*expert_grads) optim.step() @@ -100,8 +101,8 @@ def test_expert_parallel_with_data_parallel(model, tokenizer): noise_policy = SwitchNoisePolicy() router = Top1Router(noise_policy, NUM_EXPERTS, D_MODEL) - text = "Persistence is all you need." - input = tokenizer(text, return_tensors="pt") + text = ["Persistence is all you need.", "Attention is all you need."] + input = tokenizer(text, return_tensors="pt", padding=True) kwargs = { "input": input,