diff --git a/megatron/arguments.py b/megatron/arguments.py index f71507898..e3229e6e3 100644 --- a/megatron/arguments.py +++ b/megatron/arguments.py @@ -565,6 +565,9 @@ def _add_training_args(parser): help='Iteration ranges to skip. The values are one or more dash-separated ranges. e.g., 101-200 251-300.') group.add_argument('--abort-on-unmet-fused-kernel-constraints', action='store_true', help="If set to True, the program will abort if the constraints for loading a fused kernel aren't met") + group.add_argument('--stagger_checkpoint_save_load_group_size', type=int, default=None, + help="How many processes can load/save checkpoints at once, None means not to stagger and load all at once") + return parser diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py index 48401a9f1..bb2acb46b 100644 --- a/megatron/model/transformer.py +++ b/megatron/model/transformer.py @@ -605,15 +605,15 @@ def get_slopes_power_of_2(n): return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][ :n - closest_power_of_2] - slopes = torch.Tensor(get_slopes(num_attention_heads)) + slopes = torch.tensor(get_slopes(num_attention_heads)) alibi = slopes.unsqueeze(1).unsqueeze(1) * torch.arange(max_seq_len).unsqueeze(0).unsqueeze(0).expand( num_attention_heads, -1, -1) - + #Select the part of the tensor that corresponds to our tensor parallel index. tp_world_size = mpu.get_tensor_model_parallel_world_size() tp_index = mpu.get_tensor_model_parallel_rank() alibi = alibi.reshape((tp_world_size, -1, *alibi.shape[1:]))[tp_index] - + alibi = alibi.repeat(batch_size, 1, 1) return alibi @@ -629,7 +629,7 @@ class ParallelTransformerLayerPipe(ParallelTransformerLayer): to the next stage in the pipeline. This version is useful if masks are dynamic. - + 2) forward(input, **kwargs) -> output When the mask is static over all samples, it is advantageous to cache the mask and avoid communicating it. diff --git a/megatron/training.py b/megatron/training.py index 346fff1cc..d31ab47f5 100644 --- a/megatron/training.py +++ b/megatron/training.py @@ -55,6 +55,7 @@ from megatron.utils import report_memory, flops_calculator from megatron.global_vars import codecarbon_tracker_start, codecarbon_tracker_stop from megatron.data.dataset_utils import analyze_data_prefix +from megatron.utils import tp_staggered_action import deepspeed @@ -403,7 +404,8 @@ def setup_model_and_optimizer(model_provider_func): # max time. torch.distributed.barrier() timers('load-checkpoint').start() - args.iteration = load_checkpoint(model, optimizer, lr_scheduler) + with tp_staggered_action(): + args.iteration = load_checkpoint(model, optimizer, lr_scheduler) torch.distributed.barrier() timers('load-checkpoint').stop() timers.log(['load-checkpoint']) @@ -759,7 +761,8 @@ def save_checkpoint_and_time(iteration, model, optimizer, lr_scheduler): # all ranks report the max time. torch.distributed.barrier() timers('save-checkpoint').start() - save_checkpoint(iteration, model, optimizer, lr_scheduler) + with tp_staggered_action(): + save_checkpoint(iteration, model, optimizer, lr_scheduler) torch.distributed.barrier() timers('save-checkpoint').stop() timers.log(['save-checkpoint']) diff --git a/megatron/utils.py b/megatron/utils.py index 98d2f611c..fcc944778 100644 --- a/megatron/utils.py +++ b/megatron/utils.py @@ -19,6 +19,7 @@ import sys import warnings from random import randint +from typing import Callable import torch from torch import nn @@ -392,3 +393,112 @@ def found_kill_switch(): return True else: return False + + +class AllocateOnGPU(object): + """ + allocate a model directly on GPU. + + Example: + + + with AllocateOnGPU(dtype=torch.float32, enabled=enabled): + model = MyModel(hidden_dim=4*1024, nlayers=32) + + """ + + _orig_torch_empty = torch.empty + _orig_torch_zeros = torch.zeros + _orig_torch_ones = torch.ones + _orig_torch_full = torch.full + + def __init__(self, dtype, enabled=True): + self.dtype = dtype + self.enabled = enabled + + @staticmethod + def fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype) -> Callable: + def wrapped_fn(*args, **kwargs) -> torch.Tensor: + if kwargs.get("device", None) is None: + kwargs['device'] = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"])) + tensor: torch.Tensor = fn(*args, **kwargs) + if tensor.is_floating_point(): + tensor = tensor.to(target_fp_dtype) + return tensor + return wrapped_fn + + @staticmethod + def get_new_tensor_fn_for_dtype(fn: Callable, dtype: torch.dtype) -> Callable: + def new_tensor(cls, *args) -> torch.Tensor: + device = torch.device('cuda:{}'.format(os.environ["LOCAL_RANK"])) + tensor = fn(0, device=device).new_empty(*args) + if tensor.is_floating_point(): + tensor = tensor.to(dtype) + return tensor + return new_tensor + + def __enter__(self): + if not self.enabled: + return + torch.Tensor.__old_new__ = torch.Tensor.__new__ + torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self._orig_torch_empty, self.dtype) + torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype) + torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype) + torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype) + torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype) + + def __exit__(self, exc_type, exc_value, traceback): + if not self.enabled: + return + torch.Tensor.__new__ = torch.Tensor.__old_new__ + torch.empty = self._orig_torch_empty + torch.zeros = self._orig_torch_zeros + torch.ones = self._orig_torch_ones + torch.full = self._orig_torch_full + +from contextlib import contextmanager + +@contextmanager +def tp_staggered_action(): + """ + + This context manager staggers the operation it manages across multiple ranks + + So for example, if there is not enough CPU RAM to load all checkpoints on the node at once, we + can stagger those one by one or in groups to make the processes consume less CPU memory + concurrently and thus not trigger cgroups's kill action. + + The size of the group is defined by args.stagger_checkpoint_save_load + + """ + + args = get_args() + + stagger_size = args.stagger_checkpoint_save_load_group_size + + if stagger_size is None: + yield + + else: + # staggered load_checkpoint over TP ranks + + tp_size = args.tensor_model_parallel_size + + # XXXX: move this check to arguments.py + assert tp_size % stagger_size == 0, f"tp_size ({tp_size}) needs to be divisible by stagger_checkpoint_save_load_group_size ({stagger_size})." + + # split the TP ranks into groups, each group size of stagger_size + # then load each group together and have the other groups wait + tp_rank = mpu.get_tensor_model_parallel_rank() + ranks = list(range(tp_size)) + groups = [ranks[i:i + stagger_size] for i in range(0, tp_size, stagger_size)] + + for group in groups: + if tp_rank in group: + print(f"XRank={tp_rank} is processing") + yield + else: + print(f"XRank={tp_rank} is waiting") + torch.distributed.barrier() + # deadlocking in save_checkpoint + #torch.distributed.barrier(group=mpu.get_tensor_model_parallel_group()) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 64489a89c..b13263455 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -26,7 +26,7 @@ from megatron.model import GPTModel, GPTModelPipe from megatron.training import pretrain from megatron.utils import get_ltor_masks_and_position_ids, get_prefix_indices -from megatron.utils import average_losses_across_data_parallel_group +from megatron.utils import average_losses_across_data_parallel_group, AllocateOnGPU import deepspeed from deepspeed.runtime.utils import see_memory_usage @@ -41,43 +41,47 @@ def model_provider(pre_process=True, post_process=True): args = get_args() + # while this is mostly a no-op for ZeRO<3 it inits torch.distributed and does other things with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(), remote_device=None if args.remote_device == 'none' else args.remote_device, config_dict_or_path=args.deepspeed_config, enabled=args.zero_stage == 3, mpu=mpu): - if args.deepspeed: - # Precompute the attention mask and store it in args. This avoids having to - # pipeline it as an activation during training. The mask is constant, and thus - # we can reuse it. - attention_mask = torch.tril(torch.ones( - (1, args.seq_length, args.seq_length), device=torch.cuda.current_device())).view( - 1, 1, args.seq_length, args.seq_length) - - # Convert attention mask to binary: - attention_mask = (attention_mask < 0.5) - if args.fp16: - attention_mask = attention_mask.half() - elif args.bf16: - attention_mask = attention_mask.bfloat16() - - # must be bool or the training crashes expecting bool, but getting Half - args.attn_mask = attention_mask.to(torch.bool) - - model = GPTModelPipe( - num_tokentypes=0, - parallel_output=True - ) - # This is a hack to give us a reference to get_batch_pipe from within training.py - # We need to call model.set_batch_fn after deepspeed.initialize - model._megatron_batch_fn = get_batch_pipe - else: - model = GPTModel( - num_tokentypes=0, - parallel_output=True, - pre_process=pre_process, - post_process=post_process - ) + # XXX: make `enabled` configurable or always load on GPU? + with AllocateOnGPU(dtype=args.params_dtype, enabled=True): + if args.deepspeed: + # Precompute the attention mask and store it in args. This avoids having to + # pipeline it as an activation during training. The mask is constant, and thus + # we can reuse it. + attention_mask = torch.tril(torch.ones( + (1, args.seq_length, args.seq_length), device=torch.cuda.current_device())).view( + 1, 1, args.seq_length, args.seq_length) + + # Convert attention mask to binary: + attention_mask = (attention_mask < 0.5) + if args.fp16: + attention_mask = attention_mask.half() + elif args.bf16: + attention_mask = attention_mask.bfloat16() + + # must be bool or the training crashes expecting bool, but getting Half + args.attn_mask = attention_mask.to(torch.bool) + + model = GPTModelPipe( + num_tokentypes=0, + parallel_output=True + ) + # This is a hack to give us a reference to get_batch_pipe from within training.py + # We need to call model.set_batch_fn after deepspeed.initialize + model._megatron_batch_fn = get_batch_pipe + else: + model = GPTModel( + num_tokentypes=0, + parallel_output=True, + pre_process=pre_process, + post_process=post_process + ) + see_memory_usage(f"After Building Model", force=True) return model diff --git a/tests/test_basic.py b/tests/test_basic.py deleted file mode 100644 index d2a60f92c..000000000 --- a/tests/test_basic.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_import(): - import megatron diff --git a/tests/test_model.py b/tests/test_model.py index 8cabc1416..34ef0ee20 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -16,17 +16,17 @@ def get_default_args(): return { # GPT_ARGS "--num-layers": "2", - "--hidden-size": "128", - "--num-attention-heads": "4", - "--seq-length": "256", + "--hidden-size": "64", + "--num-attention-heads": "2", + "--seq-length": "128", "--max-position-embeddings": "256", - "--micro-batch-size": "4", + "--micro-batch-size": "2", "--global-batch-size": "8", "--lr-decay-iters": "320000", "--lr-decay-style": "cosine", "--lr": "0.00015", "--min-lr": "1.0e-5", - "--train-iters": "5000", + "--train-iters": "20", "--tokenizer-type": "PretrainedFromHF", "--tokenizer-name-or-path": "gpt2", "--data-impl": "mmap", @@ -42,8 +42,8 @@ def get_default_args(): # OUTPUT_ARGS "--log-interval": "10", - "--save-interval": "500", - "--eval-interval": "100", + "--save-interval": "10", + "--eval-interval": "10", "--eval-iters": "10", "--checkpoint-activations": "", @@ -184,7 +184,7 @@ def test_prefix_lm_reset_attention_mask(self): equal_vectors(output[0, changed_target_index:], output_changed_target[0, changed_target_index:]) ) ) - # Unchanged changed rows should not change either + # Unchanged rows should not change either self.assertTrue( torch.all( equal_vectors(output[1, :], output_changed_target[1, :]) diff --git a/tests/test_tensor_parallel.py b/tests/test_tensor_parallel.py index e821c9ac7..3a575216f 100644 --- a/tests/test_tensor_parallel.py +++ b/tests/test_tensor_parallel.py @@ -28,17 +28,17 @@ def get_default_args(self): return { # GPT_ARGS "--num-layers": "2", - "--hidden-size": "128", - "--num-attention-heads": "4", - "--seq-length": "256", + "--hidden-size": "64", + "--num-attention-heads": "2", + "--seq-length": "128", "--max-position-embeddings": "256", - "--micro-batch-size": "4", + "--micro-batch-size": "2", "--global-batch-size": "8", "--lr-decay-iters": "320000", "--lr-decay-style": "cosine", "--lr": "0.00015", "--min-lr": "1.0e-5", - "--train-iters": "5000", + "--train-iters": "10", "--tokenizer-type": "PretrainedFromHF", "--tokenizer-name-or-path": "gpt2", "--data-impl": "mmap", @@ -51,15 +51,14 @@ def get_default_args(self): "--attention-dropout": "0", "--hidden-dropout": "0", - # OUTPUT_ARGS "--log-interval": "10", - "--save-interval": "500", - "--eval-interval": "100", + "--save-interval": "10", + "--eval-interval": "10", "--eval-iters": "10", "--checkpoint-activations": "", - + #ds args "--deepspeed": "", "--deepspeed_config":f"{self.test_file_dir_str}/ds_config.json", @@ -67,7 +66,7 @@ def get_default_args(self): "--deepspeed-activation-checkpointing": "" # DATA_ARGS } - + def setUp(self) -> None: super().setUp() @@ -85,14 +84,14 @@ def infer_model(args): MASTER_ADDR="localhost", MASTER_PORT="9991", RANK=str(tp_index), LOCAL_RANK=str(tp_index), WORLD_SIZE=str(tp_size) ) logging.getLogger().critical("Process: starting") - + #Hack import megatron.initialize as init init.git_ds_info = lambda: None with patch('sys.argv', flatten_arguments(command_args)): with mockenv_context(**dist_env): - + def create_model_inputs(tokens): args = get_args() @@ -131,7 +130,7 @@ def create_model_inputs(tokens): model._config.zero_enabled = False _, _ = model.load_checkpoint(load, load_optimizer_states=False, load_lr_scheduler_states=False, load_module_only=True) model._config.zero_enabled = zero_enabled - + if token_ids is None: token_ids = torch.randint(args.padded_vocab_size, (args.micro_batch_size, args.seq_length)) @@ -140,8 +139,8 @@ def create_model_inputs(tokens): token_ids[token_ids == tokenizer.eod] %= args.padded_vocab_size else: token_ids = torch.tensor(token_ids) - - + + model.micro_batches = 1 model.set_batch_fn(create_model_inputs) # process batch @@ -155,42 +154,44 @@ def create_model_inputs(tokens): (input_token_ids_changed[:,changed_index] + 1) % args.padded_vocab_size output = model.eval_batch(iter([token_ids]), compute_loss = False, reduce_output = None)[0] - + output = gather_from_tensor_model_parallel_region(output)[..., :tokenizer.vocab_size] if save != None: args.save = save save_checkpoint(0, [model], None, None) - + return (output[0].detach().cpu().numpy(), token_ids.detach().cpu().numpy()) def test_alibi_tp(self): mp.set_start_method('spawn', force=True) cp_dir = self.get_auto_remove_tmp_dir() - + command_args = self.get_default_args() command_args["--position-embedding-type"] = "alibi" command_args["--tensor-model-parallel-size"] = "1" - + pool = Pool(1) result = pool.map(MegDSTestTP.infer_model, [((0, 1, command_args, None, cp_dir, None))]) pool.close() pool.join() - + output, tokens = result[0] logging.getLogger().info("First done!") command_args["--tensor-model-parallel-size"] = "2" pool = Pool(2) - result = pool.map(MegDSTestTP.infer_model, [((0, 2, command_args, tokens, None, cp_dir)), ((1, 2, command_args, tokens, None, cp_dir))]) + result = pool.map(MegDSTestTP.infer_model, [((0, 2, command_args, tokens, None, cp_dir)), + ((1, 2, command_args, tokens, None, cp_dir))]) pool.close() pool.join() - + output2, tokens = result[0] - logging.getLogger().critical(output-output2) - self.assertTrue(np.allclose(output,output2, atol=5e-3, rtol=0), "Different results when running with TP=1 and TP=2") + print(output-output2) + self.assertTrue(np.allclose(output, output2, atol=5e-3, rtol=0), + "Different results when running with TP=1 and TP=2") if __name__ == '__main__': unittest.main() diff --git a/tests/test_training.py b/tests/test_training.py index 144cfdfc4..79ae08c1e 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -169,6 +169,21 @@ def get_variation_config(self, variation, output_dir, n_samples=None): """.split() + elif variation == "stagger_checkpoints": + + # XXX: don't enable yet as it deadlocks on save_checkpoint + # --stagger_checkpoint_save_load_group_size 1 + new_args = f""" + --rampup-batch-size 2 2 {n_samples} + --train-samples {n_samples} + + --lr-decay-samples 6 + """.split() + + new_ds_args = f""" + --deepspeed_config {self.test_file_dir_str}/ds_config.json + """.split() + elif variation == "bnb": # BitsAndBytes - 8-bit optimizer @@ -283,7 +298,7 @@ def test_kill_switch(self): - @parameterized.expand(["base", "cl", "bnb", "glu", "alibi"]) + @parameterized.expand(["base", "stagger_checkpoints", "cl", "bnb", "glu", "alibi"]) def test_training_all(self, variation): # optional runs