diff --git a/tests/test_auto_parallel/test_tensor_shard/test_gpt/gpt_modules.py b/tests/test_auto_parallel/test_tensor_shard/test_gpt/gpt_modules.py index b66ad1949..22a237131 100644 --- a/tests/test_auto_parallel/test_tensor_shard/test_gpt/gpt_modules.py +++ b/tests/test_auto_parallel/test_tensor_shard/test_gpt/gpt_modules.py @@ -113,6 +113,7 @@ class GPT2Attention(nn.Module): attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]: + # query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2) qkv = self.c_attn(hidden_states) @@ -187,7 +188,6 @@ class GPT2Model(GPT2PreTrainedModel): self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, - token_type_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: input_shape = input_ids.size() @@ -196,8 +196,6 @@ class GPT2Model(GPT2PreTrainedModel): device = input_ids.device - token_type_ids = token_type_ids.view(-1, input_shape[-1]) - past_length = 0 past_key_values = tuple([None] * len(self.h)) @@ -223,9 +221,6 @@ class GPT2Model(GPT2PreTrainedModel): # add_2 hidden_states = inputs_embeds + position_embeds - token_type_embeds = self.wte(token_type_ids) - hidden_states = hidden_states + token_type_embeds - # comment to run pipeline # add_3 output_shape = input_shape + (hidden_states.size(-1),) @@ -239,3 +234,46 @@ class GPT2Model(GPT2PreTrainedModel): hidden_states = hidden_states.view(output_shape) return hidden_states + + +class GPT2LMHeadModel(GPT2PreTrainedModel): + _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.transformer = GPT2Model(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + ): + transformer_outputs = self.transformer( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + lm_logits = self.lm_head(transformer_outputs) + + return lm_logits + + +class GPTLMLoss(nn.Module): + + def __init__(self): + super().__init__() + self.loss_fn = nn.CrossEntropyLoss() + + def forward(self, logits, labels): + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) diff --git a/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_gpt2_performance.py b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_gpt2_performance.py new file mode 100644 index 000000000..87155307f --- /dev/null +++ b/tests/test_auto_parallel/test_tensor_shard/test_gpt/test_gpt2_performance.py @@ -0,0 +1,159 @@ +import copy +import random +from functools import partial +from time import time +from typing import Dict, Optional, Tuple, Union + +import numpy as np +import psutil +import pytest +import torch +import torch.multiprocessing as mp +import torch.nn as nn +import transformers +from torch.fx import GraphModule +from torch.profiler import ProfilerActivity, profile, record_function, schedule, tensorboard_trace_handler + +from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass +from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass +from colossalai.auto_parallel.tensor_shard.constants import BATCHNORM_MODULE_OP +from colossalai.auto_parallel.tensor_shard.sharding_strategy import ShardingSpec +from colossalai.auto_parallel.tensor_shard.solver import ( + CostGraph, + GraphAnalyser, + Solver, + SolverOptions, + StrategiesConstructor, +) +from colossalai.device.device_mesh import DeviceMesh +from colossalai.fx.tracer.tracer import ColoTracer +from colossalai.initialize import launch, launch_from_torch +from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.tensor.shape_consistency import ShapeConsistencyManager, to_global +from colossalai.testing import assert_close, assert_close_loose, parameterize, rerun_if_address_is_in_use +from colossalai.testing.pytest_wrapper import run_on_environment_flag +from colossalai.utils import free_port +from tests.test_auto_parallel.test_tensor_shard.test_gpt.gpt_modules import GPT2LMHeadModel, GPTLMLoss + +BATCH_SIZE = 128 +SEQ_LENGTH = 128 +HIDDEN_DIM = 4096 +NUM_HEADS = 32 +NUM_LAYERS = 4 +VOCAB_SIZE = 50257 +NUM_STEPS = 10 + + +def get_cpu_mem(): + return psutil.Process().memory_info().rss / 1024**2 + + +def get_gpu_mem(): + return torch.cuda.memory_allocated() / 1024**2 + + +def get_mem_info(prefix=''): + return f'{prefix}GPU memory usage: {get_gpu_mem():.2f} MB, CPU memory usage: {get_cpu_mem():.2f} MB' + + +def get_tflops(model_numel, batch_size, seq_len, step_time): + return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12) + + +# Randomly Generated Data +def get_data(batch_size, seq_len, vocab_size): + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device()) + attention_mask = torch.ones_like(input_ids) + return input_ids, attention_mask + + +def main(): + disable_existing_loggers() + launch_from_torch(config={}) + logger = get_dist_logger() + config = transformers.GPT2Config(n_position=SEQ_LENGTH, n_layer=NUM_LAYERS, n_head=NUM_HEADS, n_embd=HIDDEN_DIM) + + model = GPT2LMHeadModel(config=config).to('cuda') + + input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) + attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64) + + meta_input_sample = { + 'input_ids': input_ids.to('meta'), + 'attention_mask': attention_mask.to('meta'), + } + + physical_mesh_id = torch.arange(0, 4) + mesh_shape = (2, 2) + # [[0, 1] + # [2, 3]] + device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True) + shape_consistency_manager = ShapeConsistencyManager() + + tracer = ColoTracer() + + graph = tracer.trace(root=model, meta_args=meta_input_sample) + gm = GraphModule(model, graph, model.__class__.__name__) + gm.recompile() + + graph_analyser = GraphAnalyser(gm) + liveness_list = graph_analyser.liveness_analysis() + solver_options = SolverOptions() + strategies_constructor = StrategiesConstructor(graph, device_mesh, solver_options) + strategies_constructor.build_strategies_and_cost() + + cost_graph = CostGraph(strategies_constructor.leaf_strategies) + cost_graph.simplify_graph() + solver = Solver(gm.graph, strategies_constructor, cost_graph, graph_analyser, memory_budget=-1) + ret = solver.call_solver_serialized_args() + + solution = list(ret[0]) + print(solution) + gm, sharding_spec_dict, origin_spec_dict, comm_actions_dict = runtime_preparation_pass( + gm, solution, device_mesh, strategies_constructor) + gm = runtime_apply_pass(gm) + gm.recompile() + # *******************strategy selected******************* + print("*******************strategy selected*******************") + strategies_list = solution + + nodes = [strategies_vector.node for strategies_vector in strategies_constructor.leaf_strategies] + for index, node in enumerate(nodes): + print(node.name, node.strategies_vector[strategies_list[index]].name) + + # build criterion + criterion = GPTLMLoss() + + optimizer = torch.optim.Adam(gm.parameters(), lr=0.01) + numel = sum([p.numel() for p in model.parameters()]) + logger.info(get_mem_info(prefix='After init model, '), ranks=[0]) + get_tflops_func = partial(get_tflops, numel, BATCH_SIZE, SEQ_LENGTH) + torch.cuda.synchronize() + model.train() + # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + # schedule=schedule(wait=1, warmup=2, active=2), + # on_trace_ready=tensorboard_trace_handler(f'log/dummy_data/bs128_seq128_new'), + # record_shapes=True, + # profile_memory=True) as prof: + # with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA]) as prof: + for n in range(10): + # we just use randomly generated data here + input_ids, attn_mask = get_data(BATCH_SIZE, SEQ_LENGTH, VOCAB_SIZE) + optimizer.zero_grad() + start = time() + outputs = gm(input_ids, attn_mask, sharding_spec_dict, origin_spec_dict, comm_actions_dict) + loss = criterion(outputs, input_ids) + loss.backward() + optimizer.step() + # prof.step() + torch.cuda.synchronize() + step_time = time() - start + logger.info( + f'[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}', + ranks=[0]) + # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)) + torch.cuda.synchronize() + + +if __name__ == '__main__': + main()