"""Train GPT or BLOOM models.""" # pylint: disable=protected-access,too-many-lines import argparse import logging import math import os import sys import time from concurrent.futures import ProcessPoolExecutor from typing import Optional from modelling_RW import RWForCausalLM import model_config as model_config_lib import numpy as np import smdistributed.modelparallel import smdistributed.modelparallel.torch as smp import torch import torch.utils.data import transformers from data_pipeline import create_pretraining_dataloader # pylint: disable=wrong-import-order from learning_rates import AnnealingLR # pylint: disable=wrong-import-order from memory_tracker import memory_status, memory_status_cpu # pylint: disable=wrong-import-order from sdp_utils import build_param_id_to_buffer, build_param_id_to_offset, log_param_norms from smdistributed.modelparallel.torch.nn import FusedLayerNorm # pylint: disable=import-error from torch import optim from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoConfig, set_seed from transformers.trainer_utils import is_main_process # pylint: enable=import-error logging.getLogger("torch.distributed.distributed_c10d").setLevel(logging.ERROR) if not sys.warnoptions: import warnings warnings.simplefilter("ignore") def get_learning_rate_scheduler(optimizer, args): """Get learning rate scheduler.""" # Add linear learning rate scheduler. if args.lr_decay_iters is not None: num_iters = args.lr_decay_iters else: num_iters = args.max_steps num_iters = max(1, num_iters) init_step = 0 warmup_iter = args.warmup * num_iters plateau_iter = warmup_iter + args.plateau * num_iters lr_scheduler = AnnealingLR( optimizer, start_lr=args.lr, warmup_iter=warmup_iter, plateau_iter=plateau_iter, total_iters=num_iters, decay_style=args.lr_decay_style, last_iter=init_step, min_lr=args.min_lr, use_checkpoint_lr_scheduler=args.load_partial or args.load_full, override_lr_scheduler=False, ) return lr_scheduler def get_param_groups_by_weight_decay(module): """Get param groups.""" weight_decay_params = {"params": []} no_weight_decay_params = {"params": [], "weight_decay": 0.0} param_ids = set() for module_ in module.modules(): if isinstance(module_, FusedLayerNorm): for p in list(module_._parameters.values()): # pylint: disable=invalid-name if p is not None and id(p) not in param_ids: no_weight_decay_params["params"].append(p) param_ids.add(id(p)) else: for n, p in list( # pylint: disable=invalid-name module_._parameters.items() # pylint: disable=protected-access ): if p is not None and n != "bias" and id(p) not in param_ids: weight_decay_params["params"].append(p) param_ids.add(id(p)) for n, p in list( # pylint: disable=invalid-name module_._parameters.items() # pylint: disable=protected-access ): if p is not None and n == "bias" and id(p) not in param_ids: no_weight_decay_params["params"].append(p) param_ids.add(id(p)) if not no_weight_decay_params["params"]: return [weight_decay_params] return weight_decay_params, no_weight_decay_params # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, input_ids, attention_mask, args): """Train step.""" if args.logits_output: output = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids) loss = output["loss"] else: loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"] model.backward(loss) if args.logits_output: return output return loss # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def test_step(model, input_ids, attention_mask): """Test step.""" loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"] return loss def eval_model(model, dataloader, num_batches, use_bert_data): """Eval model.""" model = model.eval() n_batches = 0 loss = 0.0 with torch.no_grad(): for batch_idx, input_data in enumerate(dataloader): if use_bert_data: input_ids, _, attention_mask, _, _ = input_data else: input_ids, attention_mask = input_data if batch_idx >= num_batches: break loss += test_step(model, input_ids, attention_mask).reduce_mean() n_batches += 1 if n_batches > 0: torch.distributed.all_reduce(loss, group=smp.get_dp_process_group()) loss /= smp.dp_size() loss /= n_batches loss = loss.item() ppl = math.exp(loss) else: loss = -1.0 ppl = -1.0 return loss, ppl def train( # pylint: disable=too-many-arguments,too-many-branches,too-many-locals,too-many-statements model, optimizer, lr_scheduler, model_config, start_train_path_index, start_batch_index, num_params, total_steps, args, param_id_to_buffer, ): """Eval model.""" if args.enable_memory_profiling > 0: memory_status_cpu(msg="before train step") model.train() if args.parallel_proc_data_processing: pool = ProcessPoolExecutor(1) dp_rank = smp.dp_rank() if not args.prescaled_batch else smp.rdp_rank() dp_size = smp.dp_size() if not args.prescaled_batch else smp.rdp_size() data_type = "BERT" if args.use_bert_data else "GPT" if args.use_bert_data: train_paths = sorted( [ os.path.join(args.training_dir, p) for p in os.listdir(args.training_dir) if os.path.isfile(os.path.join(args.training_dir, p)) and "training" in p ] ) else: if args.zipped_data > 0: file_extension = ".json.gz" else: file_extension = ".json" train_paths = sorted( [ os.path.join(args.training_dir, p) for p in os.listdir(args.training_dir) if p.endswith(file_extension) ] ) train_dataloader = create_pretraining_dataloader( [train_paths[start_train_path_index]], args.train_batch_size, args.max_context_width, seed=args.seed, dp_rank=dp_rank, dp_size=dp_size, shuffle=args.same_seed < 1, zipped=args.zipped_data > 0, use_last_file_only=args.fast_validation > 0, data_type=data_type, ) if args.validation_freq is not None: # load all validation examples if smp.rank() == 0: logging.info("Creating val dataloader") if args.use_bert_data: val_paths = sorted( [ os.path.join(args.test_dir, p) for p in os.listdir(args.test_dir) if os.path.isfile(os.path.join(args.test_dir, p)) and "testing" in p ] ) else: if args.zipped_data > 0: file_extension = ".json.gz" else: file_extension = ".json" val_paths = sorted( [ os.path.join(args.test_dir, p) for p in os.listdir(args.test_dir) if p.endswith(file_extension) ] ) val_dataloader = create_pretraining_dataloader( val_paths, args.val_batch_size, args.max_context_width, seed=args.seed, dp_rank=dp_rank, dp_size=dp_size, shuffle=True, zipped=args.zipped_data > 0, use_last_file_only=args.fast_validation > 0, data_type=data_type, ) if smp.rank() == 0: logging.info("Created val dataloader of size %d.", len(val_dataloader)) start = time.time() throughputs = [] to_save = {"loss": [], "val_loss": []} loss_metric = 0 def grad_accumulation_boundary(batch_idx): return batch_idx % args.gradient_accumulation == args.gradient_accumulation - 1 def should_record(): # only record the ranks that in the tp group that contains global rank 0 if smp.tp_size() > 1: tp_group = smp.get_tp_group() return 0 in tp_group return smp.rank() == 0 # Set the same seed for computation set_seed(args.seed) for index in range(start_train_path_index, args.epochs * len(train_paths)): next_train_path_index = (index + 1) % len(train_paths) curr_train_path_index = index % len(train_paths) if total_steps >= args.max_steps: break if args.parallel_proc_data_processing: dataset_future = pool.submit( create_pretraining_dataloader, [train_paths[next_train_path_index]], args.train_batch_size, args.max_context_width, seed=args.seed, dp_rank=dp_rank, dp_size=dp_size, shuffle=args.same_seed < 1, zipped=args.zipped_data > 0, use_last_file_only=args.fast_validation > 0, data_type=data_type, ) if smp.rank() == 0: if args.use_bert_data: logging.info( "Reading data from training path %s.", train_dataloader.dataset.input_file ) else: logging.info( "Reading data from training path %s.", train_dataloader.dataset.input_paths ) for batch_idx, input_data in enumerate(train_dataloader): if batch_idx < start_batch_index: if smp.rank() == 0: logging.info( "Resuming from saved batch index %d, skipping batch %d ...", start_batch_index, batch_idx, ) if start_batch_index == len(train_dataloader): # If saving at the last batch of the file, read from the next file start_batch_index = 0 break continue start_batch_index = 0 if args.use_bert_data: input_ids, _, attention_mask, _, _ = input_data else: input_ids, attention_mask = input_data if total_steps >= args.max_steps: break torch.cuda.synchronize() step_start = time.time() if grad_accumulation_boundary(batch_idx - 1): optimizer.zero_grad(set_to_none=True) if args.logits_output: train_output = train_step(model, input_ids, attention_mask, args) loss_mb = train_output["loss"] logits_mb = train_output["logits"] if smp.tp_size() > 1: logits = torch.cat(tuple(logits_mb.outputs), dim=1) # pylint: disable=no-member else: logits = torch.cat(tuple(logits_mb.outputs), dim=0) # pylint: disable=no-member else: # Return value, loss_mb is a StepOutput object loss_mb = train_step(model, input_ids, attention_mask, args) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() if not args.validation_freq: loss_metric = loss.item() if args.enable_memory_profiling > 0: memory_status_cpu("After_train_step_cpu") memory_status(msg="After_train_step") if args.clean_cache > 0: # empty the cache to avoid OOM torch.cuda.empty_cache() if grad_accumulation_boundary(batch_idx): if args.sharded_data_parallel_degree < 1: # as SDP does its own clipping through sdp_gradient_clipping arg in init config optimizer.clip_master_grads(args.grad_clip) optimizer.step() if not (args.fp16 and optimizer.overflow): lr_scheduler.step() if args.enable_memory_profiling > 0: memory_status(msg="After_opt_step") torch.cuda.synchronize() if args.log_param_norms and args.sharded_data_parallel_degree > 1: log_param_norms(model, optimizer, param_id_to_buffer) total_steps += 1 time_elapsed = time.time() - start step_time = time.time() - step_start sample_processed = input_ids.shape[0] * dp_size throughput = sample_processed / step_time throughputs.append(throughput) # Based on the formula in # https://developer.nvidia.com/blog/scaling-language-model-training-to-a-trillion-parameters-using-megatron/ tflops_per_gpu = compute_tflops( throughput, num_params, smp.size(), input_ids.shape[1], log = (batch_idx == 0) ) if not total_steps % args.logging_freq and args.log_reduced_training_loss > 0: loss_detached = loss.detach() torch.distributed.all_reduce(loss_detached, group=smp.get_dp_process_group()) loss_scalar = loss_detached.item() / smp.dp_size() else: loss_scalar = loss.item() if smp.rank() == 0 and not total_steps % args.logging_freq: if args.sharded_data_parallel_degree > 1: gradnorm_str = f", Grad norm: {optimizer._global_grad_norm}" else: gradnorm_str = "" logging.info( "(%ds), Batch %d Loss: %s, Speed: %s samples/sec, TFLOPS/GPU: %s %s", int(time_elapsed), total_steps - 1, loss_scalar, throughput, tflops_per_gpu, gradnorm_str, ) # Compute average throughput and tflops after 30 steps to remove # high variance in initial steps if len(throughputs) > 30: avg_throughput = np.average(throughputs[30:]) avg_tflops = compute_tflops( avg_throughput, num_params, smp.size(), input_ids.shape[1] ) logging.info( f"Batch {total_steps - 1}," + f" Running Avg Speed: {avg_throughput} samples/sec," + f" Running Avg TFLOPS/GPU: {avg_tflops}" ) # evaluate on validation if args.validation_freq and not total_steps % args.validation_freq: # In GPT-NeoX runs with SDPTP, validation runs require a clean cache torch.cuda.empty_cache() cur_state = np.random.get_state() model = model.eval() val_loss, val_ppl = eval_model( model, val_dataloader, args.validation_batches, args.use_bert_data ) if is_main_process(smp.rank()): logging.info( "(%ds) Batch %d Validation loss: %s", int(time.time() - start), total_steps - 1, val_loss, ) logging.info( "(%ds) Batch %d Validation perplexity: %s", int(time.time() - start), total_steps - 1, val_ppl, ) loss_metric = val_loss if args.logits_output: to_save["val_loss"].append(val_loss) model = model.train() if args.preserve_np_state > 0: np.random.set_state(cur_state) # checkpoint if not total_steps % args.checkpoint_freq: user_content = { "cli_args": args.__dict__, "num_params": num_params, "total_steps": total_steps, "start_train_path_index": curr_train_path_index, "model_config": model_config, "start_batch_index": batch_idx + 1, } user_content["lr_scheduler"] = lr_scheduler.state_dict() # buffer_names and param_shapes used to reconstruct the full model # are automatically saved by smp.save_checkpoint() in user_content # for partial checkpoints smp.save_checkpoint( args.checkpoint_dir, tag=f"total_steps{total_steps}", partial=True, model=model, optimizer=optimizer, user_content=user_content, num_kept_partial_checkpoints=args.num_kept_checkpoints, ) if args.logits_output: to_save["loss"].append(loss.item()) if total_steps >= args.max_steps: if should_record() and args.logits_output: to_save["logits"] = logits.detach().cpu() output_file = f"rank_{smp.rank()}_" + args.logits_output torch.save(to_save, os.path.join(args.model_dir, output_file)) logging.info( "logits and loss saved at %s", os.path.join(args.model_dir, output_file) ) break del train_dataloader if args.parallel_proc_data_processing: s = time.time() # pylint: disable=invalid-name train_dataloader = dataset_future.result(timeout=None) wait_time = time.time() - s if wait_time > 1: # TODO if this happens, we should try num_workers>1 in dataloader # pylint: disable=fixme logging.info( "[%d] Waited %s for data loader to be ready. " "Please check if dataloader performance can be " "improved to avoid these waits.", smp.rank(), wait_time, ) else: train_dataloader = create_pretraining_dataloader( [train_paths[next_train_path_index]], args.train_batch_size, args.max_context_width, seed=args.seed, dp_rank=dp_rank, dp_size=dp_size, shuffle=args.same_seed < 1, zipped=args.zipped_data > 0, use_last_file_only=args.fast_validation > 0, data_type=data_type, ) # Using median throughput across all steps, could be more robust. return total_steps, np.median(throughputs) if throughputs else 0, loss_metric def parse_args(): # pylint: disable=too-many-statements """Parse args.""" parser = argparse.ArgumentParser() # hyperparameters sent by the client are passed as command-line arguments to the script. opt_grp = parser.add_argument_group( title="optimization", description="arguments for optimization" ) opt_grp.add_argument( "--train_batch_size", type=int, default=4, help="batch size per dp rank, for tensor parallelism degree 8 with pipeline parallel degree 1 this means 8*this batch size per node", # pylint: disable=line-too-long ) opt_grp.add_argument("--val_batch_size", type=int, default=4) opt_grp.add_argument("--max_steps", "--max_training_steps", type=int, default=5000) opt_grp.add_argument("--seed", type=int, default=12345) opt_grp.add_argument("--same_seed", type=int, default=0) opt_grp.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"]) opt_grp.add_argument("--fp16", default=0, type=int, help="automatic mixed precision training") opt_grp.add_argument("--bf16", default=0, type=int, help="automatic mixed precision training") opt_grp.add_argument("--sharded_data_parallel_degree", default=1, type=int) opt_grp.add_argument("--ddp_dist_backend", type=str, default="auto") opt_grp.add_argument("--grad_clip", default=1.0, type=float, help="gradient clipping") opt_grp.add_argument("--weight_decay", default=0.01, type=float, help="weight decay") opt_grp.add_argument( "--beta1", default=0.9, type=float, help="beta1 parameter for Adam optimizer" ) opt_grp.add_argument( "--beta2", default=0.95, type=float, help="beta2 parameter for Adam optimizer" ) opt_grp.add_argument( "--activation_checkpointing", type=int, default=1, help="enable gradient checkpointing to reduce memory consumption", ) parser.add_argument( "--logging_freq", type=int, default=1, help="number of iterations between logging" ) parser.add_argument( "--log_param_norms", type=int, default=0, help="to log param norms with logging_freq frequency, currently works only for sharded data parallel jobs", # pylint: disable=line-too-long ) parser.add_argument( "--log_reduced_training_loss", type=int, default=0, help="to log training loss after reducing across all data parallel ranks with logging_freq frequency", # pylint: disable=line-too-long ) # I/O io_grp = parser.add_argument_group(title="io", description="location for input and output") io_grp.add_argument("--use_bert_data", type=int, default=0, help="use bert data for training") io_grp.add_argument("--zipped_data", type=int, default=1, help="input data is zipped files") io_grp.add_argument( "--epochs", type=int, default=3, help="times of iterating over the training dataset" ) io_grp.add_argument("--output-data-dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"]) io_grp.add_argument( "--checkpoint-dir", type=str, default="/opt/ml/checkpoints", help="Saves partial checkpoints (model, optimizer) to this dir, and loads latest checkpoint from this if load_partial is specified.", # pylint: disable=line-too-long ) io_grp.add_argument( "--model-dir", type=str, default=os.environ["SM_MODEL_DIR"], help="Saves full model for inference to this dir. Also used if load_full is given to load the model. Note the lack of optimizer state here.", # pylint: disable=line-too-long ) io_grp.add_argument("--training-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"]) io_grp.add_argument("--test-dir", type=str, default=os.environ["SM_CHANNEL_TEST"]) io_grp.add_argument( "--parallel_proc_data_processing", type=int, default=0, help="Load data in parallel with a different process. At any point a process can have two files in memory. With tensor parallelism, each of the 8 processes on an instance will then have 2 files in memory. Depending on file sizes this may or may not be feasible. With pipeline parallelism this was not a problem as only 1 rank on an instance loaded data.", # pylint: disable=line-too-long ) io_grp.add_argument( "--save_final_full_model", type=int, default=0, help="Enabling this will save a combined model only at the end", ) io_grp.add_argument("--load_partial", type=int, default=0, help="Load from partial checkpoints") io_grp.add_argument("--load_full", type=int, default=0, help="Load from full checkpoints") io_grp.add_argument( "--logits_output", type=str, default="", help="Path to save logits and loss" ) io_grp.add_argument("--prescaled_batch", type=int, default=1, help="use prescaled batch") # configure model size model_grp = parser.add_argument_group( title="model", description="arguments to describe model configuration" ) model_grp.add_argument( "--fine_tune", type=int, default=0, help="Fine-tune model from checkpoint or pretrained model", ) model_grp.add_argument("--model_name", type=str, default="", help="HF model name") model_grp.add_argument("--max_context_width", type=int, default=1024) model_grp.add_argument("--vocab_size", type=int, default=50264) model_grp.add_argument("--hidden_width", type=int, default=768) model_grp.add_argument("--intermediate_size", type=int, default=2048) model_grp.add_argument("--num_layers", type=int, default=12) model_grp.add_argument("--num_heads", type=int, default=12) model_grp.add_argument("--num_heads_kv", type=int, default=8) model_grp.add_argument("--resid_pdrop", type=float, default=0.1) model_grp.add_argument("--embd_pdrop", type=float, default=0.1) model_grp.add_argument("--attn_pdrop", type=float, default=0.1) model_grp.add_argument("--alibi", type=float, default=0) model_grp.add_argument("--summary_first_pdrop", type=float, default=0.1) model_grp.add_argument("--use_adamw", type=int, default=0, help="Use adamw optimizer") model_grp.add_argument( "--use_distributed_transformer", type=int, default=1, help="Use distributed transformer" ) model_grp.add_argument( "--checkpoint_sublayers", type=int, default=0, help="Apply activation checkpointing to submodules of each transformer layer", ) model_grp.add_argument("--initializer_range", type=float, default=0.02) smp_grp = parser.add_argument_group(title="smp", description="smp") smp_grp.add_argument("--tensor_parallel_degree", type=int, default=1) smp_grp.add_argument("--pipeline_parallel_degree", type=int, default=1) smp_grp.add_argument("--microbatches", type=int, default=1) smp_grp.add_argument("--active_microbatches", type=int, default=None) smp_grp.add_argument("--optimize", type=str, default="speed") smp_grp.add_argument("--activation_strategy", type=str, default="each") smp_grp.add_argument("--shard_optimizer_state", type=int, default=0) smp_grp.add_argument("--offload_activations", type=int, default=0) smp_grp.add_argument("--fast_mode", type=int, default=0) smp_grp.add_argument("--static_mode", type=int, default=0) smp_grp.add_argument("--delayed_param", type=int, default=0) smp_grp.add_argument("--same_partition_load", type=int, default=0) smp_grp.add_argument( "--attention_in_fp32", type=int, default=0, help="When using FP16 and if the activations overflow, doing the attention computation in fp32 may help. But note that this can substantially increase memory usage and reduce performance. We recommend using bf16 instead which is more numerically stable and would not need this.", # pylint: disable=line-too-long ) smp_grp.add_argument( "--residual_addition_in_fp32", type=int, default=0, help="When using FP16 and if the activations overflow, adding residuals in fp32 may help. But note that this can substantially increase memory usage and reduce performance. We recommend using bf16 instead which is more numerically stable and would not need this.", # pylint: disable=line-too-long ) smp_grp.add_argument("--placement_strategy", type=str, default="cluster") smp_grp.add_argument("--activation_loading_horizon", type=int, default=4) smp_grp.add_argument("--skip_tracing", type=int, default=0) smp_grp.add_argument("--query_key_layer_scaling", type=int, default=0) smp_grp.add_argument("--fused_softmax", type=int, default=1) smp_grp.add_argument("--flash_attention", type=int, default=1) smp_grp.add_argument("--fused_dropout", type=int, default=0) smp_grp.add_argument("--fused_bias_gelu", type=int, default=1) smp_grp.add_argument("--gradient_accumulation", type=int, default=1) smp_grp.add_argument("--model_type", type=str, default="gpt2") smp_grp.add_argument("--rotary_pct", type=float, default=0.25) smp_grp.add_argument("--rotary_emb_base", type=int, default=10000) parser.add_argument( "--num_kept_checkpoints", type=int, default=5, help="how many checkpoints to keep before deleting", ) parser.add_argument( "--checkpoint_freq", type=int, default=10000, help="number of iterations between checkpointing", ) parser.add_argument( "--validation_freq", type=int, default=None, help="number of iterations to print validation loss", ) parser.add_argument( "--validation_batches", type=int, default=10, help="number of batches to estimate validation loss", ) parser.add_argument( "--manual_partition", type=int, default=0, help="evenly distribute layers across the partitions", ) parser.add_argument( "--partition_assignment", type=str, default="", help="number of transformer layers assigned to each partition", ) parser.add_argument( "--preserve_np_state", type=int, default=0, help="Perserve the numpy random state between validation", ) parser.add_argument( "--fast_validation", type=int, default=1, help="Running validation only with the last data file for faster speed", ) parser.add_argument( "--gather_if_shard", type=int, default=1, help="When sharding opt states is enabled, gather the opt checkpoint to rdp rank 0 during saving", # pylint: disable=line-too-long ) parser.add_argument( "--clean_cache", type=int, default=0, help="Clean torch reserved memory at he end of every step", ) parser.add_argument("--use_fsx", type=int, default=0, help="Using FSx for checkpointing") parser.add_argument( "--enable_memory_profiling", type=int, default=0, help="Enable memory profile" ) # learning rate lr_grp = parser.add_argument_group( title="lr", description="arguments for learning rate schedule" ) lr_grp.add_argument("--lr", type=float, default=None, help="Initial learning rate.") lr_grp.add_argument( "--lr_decay_style", type=str, default="linear", choices=["constant", "linear", "cosine", "exponential", "plateau"], help="Learning rate decay function.", ) lr_grp.add_argument( "--lr_decay_iters", type=int, default=None, help="number of iterations to decay learning rate over," " If None defaults to train iters", ) lr_grp.add_argument( "--min_lr", type=float, default=0.0, help="Minumum value for learning rate. The scheduler" "clip values below this threshold.", ) lr_grp.add_argument( "--warmup", type=float, default=0.01, help="Percentage of total iterations to warmup on " "(.01 = 1 percent of all training iters).", ) lr_grp.add_argument( "--plateau", type=float, default=0.4, help="Percentage of total iterations to keep at max if using plateau lr", ) ci_grp = parser.add_argument_group(title="ci", description="ci related settings") ci_grp.add_argument("--ci", default=False, action="store_true", help="Whether enable ci") ci_grp.add_argument("--time_to_train", type=int, help="time to train threshold") ci_grp.add_argument("--throughput", type=float, help="throughput threshold") ci_grp.add_argument("--loss", type=float, help="loss threshold") args, _ = parser.parse_known_args() return args def compute_num_params(model): """Get num params.""" num_params = 0 seen = set() for p in model.parameters(): # pylint: disable=invalid-name if p not in seen: seen.add(p) if hasattr(p, "ds_shape"): num_params += np.prod(p.ds_shape) else: num_params += np.prod(p.size()) return num_params def compute_tflops(throughput, num_params, num_gpus, seq_len, log = False): """Compute TFLOPs.""" tflops = 8 * throughput * num_params / num_gpus * seq_len * 1e-12 if log and smp.rank() == 0: logging.info("Compute tflops: (%s, %s, %s, %s) ==> %s.", throughput, num_params, num_gpus, seq_len, tflops) return tflops def _show_env_vars(rank: Optional[int] = 0): env_var = os.environ if rank is None or smp.rank() == rank: logging.info("Env variables (len = %d):", len(env_var)) count = 0 for key, value in sorted(env_var.items()): logging.info(" env [%03d/%03d] %-20s: `%s`", count, len(env_var), key, value) count += 1 def main(): # pylint: disable=too-many-branches,too-many-locals,too-many-statements """Main function to train GPT.""" args = parse_args() if args.partition_assignment != "" and args.manual_partition == 0: logging.warning("Partition_assignment is set, enable manual_partition.") args.manual_partition = 1 # any value here is overriden by the config set in notebook when launching the sagemaker job smp_config = { "ddp": True, "tensor_parallel_degree": args.tensor_parallel_degree, "pipeline_parallel_degree": args.pipeline_parallel_degree, "microbatches": args.microbatches, "shard_optimizer_state": args.shard_optimizer_state > 0, "prescaled_batch": args.prescaled_batch > 0, "fp16": args.fp16 > 0, "bf16": args.bf16 > 0, "offload_activations": args.offload_activations > 0, "delayed_parameter_initialization": args.delayed_param > 0, "optimize": args.optimize, "placement_strategy": args.placement_strategy, "activation_loading_horizon": args.activation_loading_horizon, "skip_tracing": args.skip_tracing > 0, "auto_partition": not args.manual_partition, "default_partition": 0, "static_mode": args.static_mode > 0, "fast_mode": args.fast_mode > 0, "sharded_data_parallel_degree": args.sharded_data_parallel_degree, "ddp_dist_backend": args.ddp_dist_backend, "sdp_hierarchical_allgather": False, "sdp_gradient_clipping": args.grad_clip, } if args.active_microbatches is not None: smp_config["active_microbatches"] = args.active_microbatches if args.log_param_norms and args.use_distributed_transformer == 1: logging.warning( "Script currently doesn't support logging param norms when using distributed transformer, disabling log_param_norms" # pylint: disable=line-too-long ) smp.init(smp_config) _show_env_vars(0) if smp.rank() == 0: logging.info("Arguments: %s", args.__dict__) logging.info("Transformers version: %s", transformers.__version__) logging.info( "smdistributed.modelparallel version: %s", smdistributed.modelparallel.__version__ ) logging.info("smdistributed config: %s", smp_config) if args.save_final_full_model and smp.rank() == 0: logging.warning( "Note that save_final_full_model only saves the final model at the end " "of all steps. It does not save optimizer state. Optimizer state is only " "saved with partial models which are saved at checkpointing_freq during " "training. If you want to restart training you need partial checkpoints." ) if args.partition_assignment != "": partition_assignment = args.partition_assignment.split(",") msg = ( f"partition_assignment must have the same size as pipeline parallel degree, " f"but getting {len(partition_assignment)} vs {smp.pp_size()}" ) logging.fatal("Will fail with: %s.", msg) raise AssertionError(msg) model_config = AutoConfig.from_pretrained("tiiuae/falcon-7b", trust_remote_code=True) model_config.hidden_size = args.hidden_width model_config.n_layer = args.num_layers model_config.n_head = args.num_heads model_config.n_head_kv = args.num_heads_kv model_config.use_cache = False # the following improves start-up time by skipping proper initialization # of weights in the original model. this is not a problem because DistributedModel # will override those weights anyway when we use distributed transformer. if args.use_distributed_transformer > 0: from transformers.modeling_utils import ( # pylint: disable=import-error,import-outside-toplevel PreTrainedModel, ) PreTrainedModel.init_weights = lambda x: None set_seed(args.seed) if args.enable_memory_profiling > 0: memory_status_cpu(msg="before model creation") if args.fp16 and args.bf16: raise ValueError("FP16 and BF16 cannot be simultaneously enabled.") if args.fp16: dtype = torch.float16 # pylint: disable=no-member elif args.bf16: dtype = torch.bfloat16 # pylint: disable=no-member else: dtype = torch.get_default_dtype() # pylint: disable=no-member if args.fine_tune > 0 and smp.rank() == 0: if args.model_type == "flan_t5": pretrained_model = AutoModelForSeq2SeqLM.from_pretrained( args.model_name or args.model_dir ) else: pretrained_model = AutoModelForCausalLM.from_pretrained( args.model_name or args.model_dir ) model_state_dict = pretrained_model.state_dict() path = os.path.join(args.model_dir, "fullmodel.pt") torch.save(model_state_dict, path) smp.barrier() # About zero_init: # we only want to init with zero for actual model for training, # in disttf case it's used in DistModel wrapper. for others we don't need to set zero init # This is needed only to param_id_to_offset with smp.model_creation( tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0, zero_init=args.use_distributed_transformer == 0, dtype=dtype, distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1, use_alibi=args.alibi > 0, attention_in_fp32=args.attention_in_fp32 > 0, fp32_residual_addition=args.residual_addition_in_fp32 > 0, query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1, fused_softmax=args.fused_softmax > 0, fused_dropout=args.fused_dropout > 0, fused_bias_gelu=args.fused_bias_gelu > 0, flash_attention=args.flash_attention > 0, ): if args.model_type == "flan_t5": model = AutoModelForSeq2SeqLM.from_config(model_config) else: model = RWForCausalLM(model_config) if args.enable_memory_profiling > 0: memory_status_cpu(msg="after model creation") # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) if not args.same_seed: # Set seed by tp_rank to prevent weights from being the same on different tp_ranks set_seed(args.seed + smp.tp_rank()) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. if args.enable_memory_profiling > 0: memory_status_cpu(msg="before dist model creation") model = smp.DistributedModel( model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation ) if args.enable_memory_profiling > 0: memory_status_cpu(msg="after dist model creation") m = model.get_module() # pylint: disable=invalid-name num_params = compute_num_params(m) if smp.rank() == 0: logging.info("# total parameters: %s", num_params) if args.use_distributed_transformer > 0: transformer_layers = m.transformer.seq_layers else: if args.model_type in ["gpt2", "bloom", "falcon"]: transformer_layers = m.transformer.h elif args.model_type == "gpt_neox": transformer_layers = m.gpt_neox.layers elif args.model_type == "flan_t5": transformer_layers = m.encoder.block if args.manual_partition: logging.debug("Manual partition enabled") if args.partition_assignment != "": get_num_layers = lambda x: int( # pylint: disable=unnecessary-lambda-assignment partition_assignment[x] ) total_layers = sum(get_num_layers(pp_rank) for pp_rank in range(smp.pp_size())) msg = ( f"partition_assignment must have the same total transformer layers as model, " f"but getting {total_layers} vs {args.num_layers}" ) logging.fatal("Will fail with: %s.", msg) raise AssertionError(msg) # evenly distribute layers across all partitions div, rem = divmod(args.num_layers, smp.pp_size()) get_num_layers = lambda x: ( # pylint: disable=unnecessary-lambda-assignment div + 1 if x >= smp.pp_size() - rem else div ) assignments = [] # (TODO) This is required for 175B otherwise a hang for partition "8,17,17,18,18,18" # Need further investigation # for pp_rank in reversed(range(smp.pp_size())): for pp_rank in range(smp.pp_size()): nl = get_num_layers(pp_rank) # pylint: disable=invalid-name logging.debug("%s layers assigned to partition %d", nl, pp_rank) assignments += [pp_rank for _ in range(nl)] for i, c in enumerate(transformer_layers.children()): # pylint: disable=invalid-name smp.set_partition(c, assignments[i]) param_groups = get_param_groups_by_weight_decay(m) if args.use_adamw > 0: optimizer = optim.AdamW( param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay ) else: optimizer = optim.Adam( param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay ) if args.activation_checkpointing: # pylint: disable=too-many-nested-blocks if args.use_distributed_transformer or smp.tp_size() > 1: if args.checkpoint_sublayers: for c in transformer_layers.children(): # pylint: disable=invalid-name smp.set_activation_checkpointing(c.attention) smp.set_activation_checkpointing(c.output) else: smp.set_activation_checkpointing( transformer_layers, strategy=args.activation_strategy ) else: for c in transformer_layers.children(): # pylint: disable=invalid-name if args.checkpoint_sublayers: if args.model_type == "gpt2": smp.set_activation_checkpointing(c.attn) smp.set_activation_checkpointing(c.mlp) elif args.model_type in ["gpt_neox", "bloom"]: if args.model_type == "gpt_neox": smp.set_activation_checkpointing(c.attention) elif args.model_type == "bloom": smp.set_activation_checkpointing(c.self_attention) smp.set_activation_checkpointing(c.input_layernorm) smp.set_activation_checkpointing(c.post_attention_layernorm) smp.set_activation_checkpointing(c.mlp) elif args.model_type == "falcon": smp.set_activation_checkpointing(c.self_attention) else: smp.set_activation_checkpointing(c) # Checkpoint decoder layers in T5 if args.model_type == "flan_t5": for c in m.decoder.block.children(): # pylint: disable=invalid-name smp.set_activation_checkpointing(c) if args.log_param_norms > 0 and args.sharded_data_parallel_degree > 1 and args.use_distributed_transformer == 0: param_id_to_offset = build_param_id_to_offset(param_groups) optimizer = smp.DistributedOptimizer( optimizer, static_loss_scale=None, dynamic_loss_scale=True, dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2}, ) if args.fine_tune > 0: smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False) if args.log_param_norms > 0 and args.sharded_data_parallel_degree > 1 and args.use_distributed_transformer == 0: param_id_to_buffer = build_param_id_to_buffer(optimizer, param_id_to_offset) else: param_id_to_buffer = None lr_scheduler = get_learning_rate_scheduler(optimizer, args) if args.enable_memory_profiling > 0: model.register_post_partition_hook( lambda model, optimizer: memory_status(msg="After partition") ) # load after wrapping model and optimizer with smp Distributed... if args.load_full or args.load_partial: if args.load_partial and args.load_full: logging.info( "Since both --load_partial and --load_full set, will try to load from full " "checkpoint. If the intention is to load from partial checkpoint, please don't set " "--load_full" ) partial = not args.load_full path = args.checkpoint_dir if partial else args.model_dir tag = None if partial else "fullmodel.pt" user_content = smp.resume_from_checkpoint(path, tag=tag, partial=partial) total_steps = user_content["total_steps"] if partial else 0 start_train_path_index = user_content.get("start_train_path_index", 0) start_batch_index = user_content.get("start_batch_index", 0) if "lr_scheduler" in user_content: lr_scheduler.load_state_dict(user_content["lr_scheduler"]) else: total_steps = 0 start_train_path_index = 0 start_batch_index = 0 # Add emty cache to clear memory when loaded with partial checkpointing # for SDPTP and GPT NeoX torch.cuda.empty_cache() start = time.time() total_steps, throughput, loss = train( model, optimizer, lr_scheduler, model_config, start_train_path_index, start_batch_index, num_params, total_steps, args, param_id_to_buffer, ) time_to_train = time.time() - start if args.ci: logging.info("[SMP_METRIC]__T5__Time_to_train__%s", time_to_train) logging.info("[SMP_METRIC]__T5__samples/second__%s", throughput) logging.info("[SMP_METRIC]__T5__Loss__%s", loss) if not args.load_partial and not args.load_full: if time_to_train >= args.time_to_train: msg = f"Time to train ({time_to_train}) >= threshold ({args.time_to_train})" logging.fatal("Will fail with: %s.", msg) raise AssertionError(msg) if throughput <= args.throughput: msg = f"Throughput ({throughput}) >= threshold ({args.throughput})" logging.fatal("Will fail with: %s.", msg) raise AssertionError(msg) if args.loss and loss >= args.loss: msg = f"Loss ({loss}) >= threshold ({args.loss})" logging.fatal("Will fail with: %s.", msg) raise AssertionError(msg) if args.save_final_full_model: # saves full model at the end user_content = { "cli_args": args.__dict__, "num_params": num_params, "total_steps": total_steps, "model_config": model_config, } # pylint: disable=line-too-long # You can also get the full model from the SDP checkpoint, by using the following API # > from smp.sharded_data_parallel_checkpoint import get_full_state_dict_from_sharded_data_parallel_checkpoint # > full_model = get_full_state_dict_from_sharded_data_parallel_checkpoint(args.model_dir, tag=f"sharded_data_parallel_final_full_{num_params}", dtype=torch.float32) # > if args.use_distributed_transformer > 0: # translate the state_dict to hf format if distributed transformer is used # > full_model = smp.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(full_model, max_seq_len=args.max_context_width) # Note: the shared parameter will not be reflected so during loading you might need to load with strict=False # pylint: enable=line-too-long smp.save_checkpoint( args.model_dir, tag="fullmodel.pt", partial=False, model=model, user_content=user_content, ) smp.barrier() if smp.rank() == 0: logging.info("SMP training finished successfully") if __name__ == "__main__": logging.basicConfig( format="%(asctime)s,%(msecs)03d %(levelname)-8s " "[%(filename)s:%(lineno)d] %(message)s", datefmt="%Y-%m-%d:%H:%M:%S", level=logging.INFO, ) main()