import argparse import time from functools import partial import torch from torch import nn from colossalai.fx import ColoTracer from colossalai.fx.passes.adding_split_node_pass import gpipe_dp_split_pass, split_with_split_nodes_pass from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.legacy.pipeline.middleware.adaptor import get_fx_topology from colossalai.legacy.pipeline.rpc._pipeline_schedule import FillDrainPipelineEngine from colossalai.legacy.pipeline.rpc.utils import rpc_run from colossalai.logging import disable_existing_loggers, get_dist_logger from model_zoo import model_builder def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--model_type", type=str, default="gpt2_medium") parser.add_argument("--world_size", type=int, default=2) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--dp_degree", type=int, default=1) parser.add_argument("--tp_degree", type=int, default=1) parser.add_argument("--num_microbatches", type=int, default=2) parser.add_argument("--device", type=str, choices=["cpu", "cuda"], default="cuda") parser.add_argument("--master_addr", type=str, default="localhost") parser.add_argument("--master_port", type=str, default="29011") parser.add_argument("--num_worker_threads", type=int, default=128) return parser.parse_args() 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)) # 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 get_tflops(model_numel, batch_size, seq_len, step_time): return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12) # Create annotated model which is noted where to be splitted. def get_annotated_model(model, data_kwargs, num_stages, num_microbatches): tracer = ColoTracer() meta_args = {k: v.to("meta") for k, v in data_kwargs.items()} graph = tracer.trace(root=model, meta_args=meta_args) gm = torch.fx.GraphModule(model, graph, model.__class__.__name__) interp_meta_args = tuple([v.to("meta") for k, v in data_kwargs.items()]) interp = MetaInfoProp(gm) interp.run(*interp_meta_args) # annotated_model = avgnode_split_pass(gm, num_stages) annotated_model = gpipe_dp_split_pass(gm, num_stages, num_microbatches, mode="block", block_limit=0.01) return annotated_model def create_partition_module(pp_rank: int, num_stages: int, model, data_kwargs, num_microbatches): annotated_model = get_annotated_model(model, data_kwargs, num_stages, num_microbatches) top_module, split_submodules = split_with_split_nodes_pass(annotated_model, merge_output=True) topo = get_fx_topology(top_module) for submodule in split_submodules: if isinstance(submodule, torch.fx.GraphModule): setattr(submodule, "_topo", topo) return split_submodules[pp_rank + 1] def partition(model, data_kwargs, num_microbatches, pp_rank: int, chunk: int, stage_num: int): module = create_partition_module(pp_rank, stage_num, model, data_kwargs, num_microbatches) return module def run_master(args): batch_size = args.batch_size device = args.device world_size = args.world_size stage_num = world_size num_microbatches = args.num_microbatches model_type = args.model_type # batch size per DP degree SEQ_LEN = 1024 VOCAB_SIZE = 50257 NUM_STEPS = 10 WARMUP_STEPS = 1 disable_existing_loggers() logger = get_dist_logger() logger.info( f"{args.model_type}, batch size {batch_size}, num stage {stage_num}, num microbatch {num_microbatches}", ranks=[0], ) torch.manual_seed(123) # build criterion criterion = GPTLMLoss() # warm up pipeline fx partition input_ids, attn_mask = get_data(batch_size, SEQ_LEN, VOCAB_SIZE) warmup_data_kwargs = {"input_ids": input_ids, "attention_mask": attn_mask} # create model logger.info(f"start model_builder") model = model_builder(model_type)(checkpoint=False) logger.info(f"end model_builder") # set 1f1b pipeline engine pp_engine = FillDrainPipelineEngine( partition_fn=partial(partition, model, warmup_data_kwargs, num_microbatches), stage_num=stage_num, num_microbatches=num_microbatches, device=device, chunk=1, criterion=criterion, metric=None, checkpoint=False, ) partition_numels = pp_engine.remote_numels() for rank, numel in partition_numels.items(): logger.info(f"{rank=} numel in the partition:{numel}") # build optim pp_engine.initialize_optimizer(torch.optim.Adam, lr=1e-3) ranks_tflops = {} for n in range(NUM_STEPS): # we just use randomly generated data here input_ids, attn_mask = get_data(batch_size, SEQ_LEN, VOCAB_SIZE) batch = {"input_ids": input_ids, "attention_mask": attn_mask} start = time.time() outputs = pp_engine.forward_backward(batch=batch, labels=input_ids, forward_only=False) step_time = time.time() - start for rank, numel in partition_numels.items(): if rank not in ranks_tflops: ranks_tflops[rank] = [] step_tflops = get_tflops(numel, batch_size, SEQ_LEN, step_time) logger.info( f"Rank{rank} , [{n + 1}/{NUM_STEPS}] , Step time: {step_time:.3f}s, TFLOPS: {get_tflops(numel, batch_size, SEQ_LEN, step_time):.3f}", ranks=[0], ) if n >= WARMUP_STEPS: ranks_tflops[rank].append(step_tflops) median_index = ((NUM_STEPS - WARMUP_STEPS) >> 1) + WARMUP_STEPS gpu_tflops = [] for rank, tflops_list in ranks_tflops.items(): tflops_list.sort() gpu_tflops.append(tflops_list[median_index]) logger.info(f"GPU{rank} Median TFLOPS is {tflops_list[median_index]:.3f}") logger.info(f"Total TFLOPS is {sum(gpu_tflops):.3f}") logger.info(f"Avg TFLOPS per GPU is {sum(gpu_tflops) / world_size:.3f}") if __name__ == "__main__": args = parse_args() rpc_run(args, run_master)