From 69b9368c9359b8941d3ffca1e0db788b5b0a38ca Mon Sep 17 00:00:00 2001 From: baijumeswani Date: Wed, 25 Nov 2020 09:40:26 -0800 Subject: [PATCH] Add unit tests to identify configuration migration scenarios for checkpointing (#5678) --- cmake/onnxruntime_python.cmake | 7 + .../test/python/checkpoint/__init__.py | 0 .../test/python/checkpoint/_test_helpers.py | 175 ++ .../orttraining_test_load_checkpoint.py | 1472 +++++++++++++++++ .../orttraining_test_save_checkpoint.py | 116 ++ .../python/orttraining_distributed_tests.py | 9 +- .../python/orttraining_test_checkpoint.py | 117 ++ 7 files changed, 1895 insertions(+), 1 deletion(-) create mode 100644 orttraining/orttraining/test/python/checkpoint/__init__.py create mode 100644 orttraining/orttraining/test/python/checkpoint/_test_helpers.py create mode 100644 orttraining/orttraining/test/python/checkpoint/orttraining_test_load_checkpoint.py create mode 100644 orttraining/orttraining/test/python/checkpoint/orttraining_test_save_checkpoint.py create mode 100644 orttraining/orttraining/test/python/orttraining_test_checkpoint.py diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index b253f83d71..f6b2e2f2f3 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -203,6 +203,9 @@ file(GLOB onnxruntime_python_test_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/test/python/*.py" "${ORTTRAINING_SOURCE_DIR}/test/python/*.py" ) +file(GLOB onnxruntime_python_checkpoint_test_srcs CONFIGURE_DEPENDS + "${ORTTRAINING_SOURCE_DIR}/test/python/checkpoint/*.py" +) file(GLOB onnxruntime_python_tools_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/python/tools/*.py" ) @@ -241,6 +244,7 @@ add_custom_command( COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/transformers COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/quantization COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/quantization/operators + COMMAND ${CMAKE_COMMAND} -E make_directory $/checkpoint COMMAND ${CMAKE_COMMAND} -E copy ${ONNXRUNTIME_ROOT}/__init__.py $/onnxruntime/ @@ -256,6 +260,9 @@ add_custom_command( COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_test_srcs} $ + COMMAND ${CMAKE_COMMAND} -E copy + ${onnxruntime_python_checkpoint_test_srcs} + $/checkpoint/ COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_backend_srcs} $/onnxruntime/backend/ diff --git a/orttraining/orttraining/test/python/checkpoint/__init__.py b/orttraining/orttraining/test/python/checkpoint/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/orttraining/orttraining/test/python/checkpoint/_test_helpers.py b/orttraining/orttraining/test/python/checkpoint/_test_helpers.py new file mode 100644 index 0000000000..f0c83c2112 --- /dev/null +++ b/orttraining/orttraining/test/python/checkpoint/_test_helpers.py @@ -0,0 +1,175 @@ +import os +import pickle +from itertools import islice + +import torch +import torch.distributed as dist + +from onnxruntime import set_seed +from onnxruntime.training import amp, checkpoint, optim, orttrainer +from orttraining_test_orttrainer_frontend import _load_pytorch_transformer_model +from onnxruntime.capi._pybind_state import set_cuda_device_id, get_mpi_context_world_rank, get_mpi_context_world_size + +def _train(trainer, train_data, batcher_fn, total_batch_steps = 5, seed = 1): + """Runs train_step total_batch_steps number of times on the given trainer""" + for i in range(total_batch_steps): + torch.manual_seed(seed) + set_seed(seed) + data, targets = batcher_fn(train_data, i*35) + trainer.train_step(data, targets) + +def makedir(checkpoint_dir): + """Creates a directory if checkpoint_dir does not exist""" + if not os.path.exists(checkpoint_dir): + os.makedirs(checkpoint_dir, exist_ok = True) + +def _save(trainer, checkpoint_dir, state_dict_key_name): + """Saves the ORTTrainer checkpoint and the complete state dictionary to the given checkpoint_dir directory""" + + # save current model parameters as a checkpoint + makedir(checkpoint_dir) + checkpoint.experimental_save_checkpoint(trainer, checkpoint_dir) + state_dict = checkpoint.experimental_state_dict(trainer) + pickle.dump({state_dict_key_name : state_dict}, open(os.path.join(checkpoint_dir, state_dict_key_name+'.pkl'), "wb")) + +def _chunkify(sequence, num_chunks): + """Breaks down a given sequence into num_chunks chunks""" + quo, rem = divmod(len(sequence), num_chunks) + return (sequence[i * quo + min(i, rem):(i + 1) * quo + min(i + 1, rem)] for i in range(num_chunks)) + +def _setup_test_infra(world_rank, world_size): + """distributed setup just for testing purposes""" + os.environ['RANK'] = str(world_rank) + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['MASTER_ADDR'] = '127.0.0.1' + os.environ['MASTER_PORT'] = '29500' + + set_cuda_device_id(world_rank) + + dist.init_process_group(backend='nccl', world_size=world_size, rank=world_rank) + +def distributed_setup(func): + """Decorator function for distributed tests. + + Sets up distributed environment by extracting the following variables from MPI context + - world_rank + - world_size + - device + + Also sets up the infrastructure required for the distributed tests such as setting up the torch distributed initialization + """ + def setup(checkpoint_dir): + world_rank = get_mpi_context_world_rank() + world_size = get_mpi_context_world_size() + device = 'cuda:' + str(world_rank) + + _setup_test_infra(world_rank, world_size) + + func(world_rank, world_size, device, checkpoint_dir=checkpoint_dir) + + return setup + +def create_orttrainer_and_load_checkpoint(device, trainer_opts, checkpoint_dir): + """Instantiate and load checkpoint into trainer + + - Instantiates the ORTTrainer with given input trainer_opts configuration for a simple transformer model + - Loads the checkpoint from directory checkpoint_dir into the trainer + - Runs eval_step on the trainer so the trainer onnx graph is initialized + - Returns the trainer state_dict and the pytorch model + """ + seed = 1 + torch.manual_seed(seed) + set_seed(seed) + + # PyTorch transformer model setup + learning_rate = 0.1 + optim_config = optim.LambConfig(lr=learning_rate) + model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device) + trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts)) + + # load checkpoint into trainer + checkpoint.experimental_load_checkpoint(trainer, checkpoint_dir) + + # run an eval step to innitialize the graph + torch.manual_seed(seed) + set_seed(seed) + data, targets = batcher_fn(train_data, 0) + trainer.eval_step(data, targets) + + return checkpoint.experimental_state_dict(trainer), model + +def split_state_dict(state_dict): + """Given a flat state dictionary, split it into optimizer, fp32_param, fp16_param hierarchical dictionary and return""" + + optimizer_keys = ['Moment_1_', 'Moment_2_', 'Update_Count_', 'Step_'] + split_sd = {'optimizer': {}, 'fp32_param': {}, 'fp16_param': {}} + for k, v in state_dict.items(): + mode = 'fp32_param' + for optim_key in optimizer_keys: + if k.startswith(optim_key): + mode = 'optimizer' + break + if k.endswith('_fp16'): + mode = 'fp16_param' + split_sd[mode][k] = v + return split_sd + +def _split_name(name): + """Splits given state name (model or optimizer state name) into the param_name, optimizer_key, view_num and the fp16_key""" + name_split = name.split('_view_') + view_num = None + if(len(name_split) > 1): + view_num = int(name_split[1]) + optimizer_key = '' + fp16_key = '' + if name_split[0].startswith('Moment_1'): + optimizer_key = 'Moment_1_' + elif name_split[0].startswith('Moment_2'): + optimizer_key = 'Moment_2_' + elif name_split[0].startswith('Update_Count'): + optimizer_key = 'Update_Count_' + elif name_split[0].endswith('_fp16'): + fp16_key = '_fp16' + param_name = name_split[0] + if optimizer_key != '': + param_name = param_name.split(optimizer_key)[1] + param_name = param_name.split('_fp16')[0] + return param_name, optimizer_key, view_num, fp16_key + +def aggregate_states(aggregated_states, state_dict): + """Concatenate existing aggregated state dict values with given state_dict values""" + + for key, value in state_dict.items(): + weight_name, optimizer_key, view_num, fp16_key = _split_name(key) + if view_num is not None: + # parameter is sharded + param_name = optimizer_key + weight_name + fp16_key + + if param_name in aggregated_states and optimizer_key not in ['Update_Count_']: + # found a previous shard of the param, concatenate shards ordered by ranks + aggregated_states[param_name] = torch.cat((aggregated_states[param_name], value)) + else: + aggregated_states[param_name] = value + else: + aggregated_states[key] = value + +def create_orttrainer_and_save_checkpoint(device, trainer_opts, checkpoint_dir, state_dict_key_name='state_dict'): + learning_rate = 0.1 + seed = 1 + + torch.manual_seed(seed) + set_seed(seed) + + optim_config = optim.LambConfig(lr=learning_rate) + model, model_desc, loss_fn, batcher_fn, train_data, _, _ = _load_pytorch_transformer_model(device) + trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=loss_fn, options=orttrainer.ORTTrainerOptions(trainer_opts)) + + if 'distributed' in trainer_opts: + train_data = next(islice(_chunkify(train_data, trainer_opts['distributed']['world_size']), trainer_opts['distributed']['world_rank'], None)) + + # run train steps + _train(trainer, train_data, batcher_fn) + + # save current model parameters as a checkpoint + if checkpoint_dir: + _save(trainer, checkpoint_dir, state_dict_key_name) diff --git a/orttraining/orttraining/test/python/checkpoint/orttraining_test_load_checkpoint.py b/orttraining/orttraining/test/python/checkpoint/orttraining_test_load_checkpoint.py new file mode 100644 index 0000000000..eb8dd62ab5 --- /dev/null +++ b/orttraining/orttraining/test/python/checkpoint/orttraining_test_load_checkpoint.py @@ -0,0 +1,1472 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +################################################################################ +# Refer to orttraining_test_checkpoint.py for an overview about Checkpoint tests +################################################################################ + +import os +import pickle +from numpy.testing import assert_allclose +import numpy as np +import argparse +import glob + +import torch +import torch.distributed as dist + +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import onnxruntime +from onnxruntime.training import checkpoint +from _test_helpers import distributed_setup, create_orttrainer_and_load_checkpoint, split_state_dict, aggregate_states + +global_fp16_fp32_atol = 1e-3 + +def test_load_from_single_node_full_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_single_node_mixed_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_post_checkpoint + assert_allclose(value, state_dict_post_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_single_node_mixed_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + +def test_load_from_single_node_full_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_post_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_pre_checkpoint + assert_allclose(value, state_dict_pre_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_data_parallelism_full_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_data_parallelism_mixed_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_post_checkpoint + assert_allclose(value, state_dict_post_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_data_parallelism_mixed_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + +def test_load_from_data_parallelism_full_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_post_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_pre_checkpoint + assert_allclose(value, state_dict_pre_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_pre_checkpoint[key]) + +def test_load_from_distributed_zero_full_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp32_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key]) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + + # load state into pytorch and compare + checkpoint_files = checkpoint._list_checkpoint_files(checkpoint_dir, "ORT_checkpoint") + agg_checkpoint = checkpoint._CombineZeroCheckpoint(checkpoint_files) + agg_state_dict = agg_checkpoint.aggregate_checkpoints() + model.load_state_dict(agg_state_dict, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, agg_state_dict[key]) + +def test_load_from_distributed_zero_mixed_precision_into_single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp16 states are not sharded + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + + # load state into pytorch and compare + checkpoint_files = checkpoint._list_checkpoint_files(checkpoint_dir, "ORT_checkpoint") + agg_checkpoint = checkpoint._CombineZeroCheckpoint(checkpoint_files) + agg_state_dict = agg_checkpoint.aggregate_checkpoints() + model.load_state_dict(agg_state_dict, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, agg_state_dict[key]) + +def test_load_from_distributed_zero_mixed_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp16 states are not sharded + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp16_param'][key]) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + +def test_load_from_distributed_zero_full_precision_into_single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp32 states are not sharded + for key, value in state_dict_post_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_pre_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_single_node_full_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_single_node_mixed_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_post_checkpoint + assert_allclose(value, state_dict_post_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_single_node_mixed_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + +@distributed_setup +def test_load_from_single_node_full_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_post_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_pre_checkpoint + assert_allclose(value, state_dict_pre_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_data_parallelism_full_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_data_parallelism_mixed_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_post_checkpoint + assert_allclose(value, state_dict_post_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_post_checkpoint[key]) + + # load state into pytorch and compare + model.load_state_dict(state_dict_pre_checkpoint, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_data_parallelism_mixed_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_pre_checkpoint.items(): + assert_allclose(value, state_dict_post_checkpoint[key]) + +@distributed_setup +def test_load_from_data_parallelism_full_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = state['state_dict'] + + # compare all states + for key, value in state_dict_post_checkpoint.items(): + if key.endswith('_fp16'): + full_precision_key = key[:-5] + assert full_precision_key in state_dict_pre_checkpoint + assert_allclose(value, state_dict_pre_checkpoint[full_precision_key], atol=global_fp16_fp32_atol) + else: + assert_allclose(value, state_dict_pre_checkpoint[key]) + +@distributed_setup +def test_load_from_distributed_zero_full_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp32_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key]) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + + # load state into pytorch and compare + checkpoint_files = checkpoint._list_checkpoint_files(checkpoint_dir, "ORT_checkpoint") + agg_checkpoint = checkpoint._CombineZeroCheckpoint(checkpoint_files) + agg_state_dict = agg_checkpoint.aggregate_checkpoints() + model.load_state_dict(agg_state_dict, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, agg_state_dict[key]) + +@distributed_setup +def test_load_from_distributed_zero_mixed_precision_into_data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp16 states are not sharded + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + + # load state into pytorch and compare + checkpoint_files = checkpoint._list_checkpoint_files(checkpoint_dir, "ORT_checkpoint") + agg_checkpoint = checkpoint._CombineZeroCheckpoint(checkpoint_files) + agg_state_dict = agg_checkpoint.aggregate_checkpoints() + model.load_state_dict(agg_state_dict, strict=False) + state_dict_pytorch = model.state_dict() + for key, value in state_dict_pytorch.items(): + assert_allclose(value, agg_state_dict[key]) + +@distributed_setup +def test_load_from_distributed_zero_mixed_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp16 states are not sharded + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp16_param'][key]) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_distributed_zero_full_precision_into_data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + num_states = len(glob.glob1(checkpoint_dir,"state_dict*")) + optimizer_states = dict() + for rank in range(num_states): + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(rank)]) + + # compare all states + # fp32 states are not sharded + for key, value in state_dict_post_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_pre_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, state_dict_pre_checkpoint['optimizer']) + + # compare optimizer states + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_post_checkpoint['optimizer'][key].size()), state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_single_node_full_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp32_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key]) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_single_node_mixed_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_single_node_mixed_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp16_param'][key]) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_single_node_full_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_post_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_pre_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_data_parallelism_full_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp32_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key]) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_data_parallelism_mixed_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_data_parallelism_mixed_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp16_param'][key]) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_data_parallelism_full_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict']) + + # compare all states + # model states + for key, value in state_dict_post_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_pre_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # round about way of checking optimizer states. Save state dicts into temporary folder, read them and aggregate them. + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl'), "wb") as f: + pickle.dump(state_dict_post_checkpoint, f) + dist.barrier() + + if world_rank == 0: + num_states = len(glob.glob1(checkpoint_dir,"distributed_state*")) + optimizer_states = dict() + for rank in range(num_states): + rank_state_dict = None + with open(os.path.join(checkpoint_dir, 'distributed_state_'+str(rank)+'.pkl'), 'rb') as f: + rank_state_dict = pickle.load(f) + + # collect optimizer states for later comparison since they are sharded + aggregate_states(optimizer_states, rank_state_dict['optimizer']) + + # compare optimizer states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + for key, value in optimizer_states.items(): + assert_allclose(value.reshape(state_dict_pre_checkpoint['optimizer'][key].size()), state_dict_pre_checkpoint['optimizer'][key]) + """ + + dist.barrier() + os.remove(os.path.join(checkpoint_dir, 'distributed_state_'+str(world_rank)+'.pkl')) + +@distributed_setup +def test_load_from_distributed_zero_full_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(world_rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(world_rank)]) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp32_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key]) + + # compare optimizer states + for key, value in state_dict_pre_checkpoint['optimizer'].items(): + assert_allclose(value, state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_distributed_zero_mixed_precision_into_distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(world_rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(world_rank)]) + + # compare all states + # model states + """ + TODO: Uncomment this after Checkpoint redesign. Current implementation does not support it + fp32 weights pre checkpoint are sharded. But since this is a one to one mapping (from distributed zero to distributed zero albeit + mixed to full precision), the fp32 weights are not aggregated before copying into the new trainer + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + """ + + # compare optimizer states + for key, value in state_dict_pre_checkpoint['optimizer'].items(): + assert_allclose(value, state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_distributed_zero_mixed_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(world_rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(world_rank)]) + + # compare all states + # model states + for key, value in state_dict_pre_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_post_checkpoint['fp16_param'][key]) + + # compare optimizer states + for key, value in state_dict_pre_checkpoint['optimizer'].items(): + assert_allclose(value, state_dict_post_checkpoint['optimizer'][key]) + +@distributed_setup +def test_load_from_distributed_zero_full_precision_into_distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + + # extract state dictionaries to compare + state_dict_post_checkpoint, model = create_orttrainer_and_load_checkpoint(device, opts, checkpoint_dir) + state_dict_post_checkpoint = split_state_dict(state_dict_post_checkpoint) + + state = None + with open(os.path.join(checkpoint_dir, 'state_dict_'+str(world_rank)+'.pkl'), 'rb') as f: + state = pickle.load(f) + state_dict_pre_checkpoint = split_state_dict(state['state_dict_'+str(world_rank)]) + + # compare all states + # model states + for key, value in state_dict_post_checkpoint['fp16_param'].items(): + assert_allclose(value, state_dict_pre_checkpoint['fp32_param'][key[:-5]], atol=global_fp16_fp32_atol) + + # compare optimizer states + for key, value in state_dict_pre_checkpoint['optimizer'].items(): + assert_allclose(value, state_dict_post_checkpoint['optimizer'][key]) + +function_map = { + # all config to single node config + 'test_load_from_single_node_full_precision_into_single_node_full_precision': test_load_from_single_node_full_precision_into_single_node_full_precision, + 'test_load_from_single_node_mixed_precision_into_single_node_mixed_precision': test_load_from_single_node_mixed_precision_into_single_node_mixed_precision, + 'test_load_from_single_node_mixed_precision_into_single_node_full_precision': test_load_from_single_node_mixed_precision_into_single_node_full_precision, + 'test_load_from_single_node_full_precision_into_single_node_mixed_precision': test_load_from_single_node_full_precision_into_single_node_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_single_node_full_precision': test_load_from_data_parallelism_full_precision_into_single_node_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_single_node_full_precision': test_load_from_data_parallelism_mixed_precision_into_single_node_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_single_node_mixed_precision': test_load_from_data_parallelism_mixed_precision_into_single_node_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_single_node_mixed_precision': test_load_from_data_parallelism_full_precision_into_single_node_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_single_node_full_precision': test_load_from_distributed_zero_full_precision_into_single_node_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_single_node_full_precision': test_load_from_distributed_zero_mixed_precision_into_single_node_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_single_node_mixed_precision': test_load_from_distributed_zero_mixed_precision_into_single_node_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_single_node_mixed_precision': test_load_from_distributed_zero_full_precision_into_single_node_mixed_precision, + + # all config to data parallel node config + 'test_load_from_single_node_full_precision_into_data_parallelism_full_precision': test_load_from_single_node_full_precision_into_data_parallelism_full_precision, + 'test_load_from_single_node_mixed_precision_into_data_parallelism_full_precision': test_load_from_single_node_mixed_precision_into_data_parallelism_full_precision, + 'test_load_from_single_node_mixed_precision_into_data_parallelism_mixed_precision': test_load_from_single_node_mixed_precision_into_data_parallelism_mixed_precision, + 'test_load_from_single_node_full_precision_into_data_parallelism_mixed_precision': test_load_from_single_node_full_precision_into_data_parallelism_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_data_parallelism_full_precision': test_load_from_data_parallelism_full_precision_into_data_parallelism_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_data_parallelism_full_precision': test_load_from_data_parallelism_mixed_precision_into_data_parallelism_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_data_parallelism_mixed_precision': test_load_from_data_parallelism_mixed_precision_into_data_parallelism_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_data_parallelism_mixed_precision': test_load_from_data_parallelism_full_precision_into_data_parallelism_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_data_parallelism_full_precision': test_load_from_distributed_zero_full_precision_into_data_parallelism_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_data_parallelism_full_precision': test_load_from_distributed_zero_mixed_precision_into_data_parallelism_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_data_parallelism_mixed_precision': test_load_from_distributed_zero_mixed_precision_into_data_parallelism_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_data_parallelism_mixed_precision': test_load_from_distributed_zero_full_precision_into_data_parallelism_mixed_precision, + + # all config to distributed zero node config + 'test_load_from_single_node_full_precision_into_distributed_zero_full_precision': test_load_from_single_node_full_precision_into_distributed_zero_full_precision, + 'test_load_from_single_node_mixed_precision_into_distributed_zero_full_precision': test_load_from_single_node_mixed_precision_into_distributed_zero_full_precision, + 'test_load_from_single_node_mixed_precision_into_distributed_zero_mixed_precision': test_load_from_single_node_mixed_precision_into_distributed_zero_mixed_precision, + 'test_load_from_single_node_full_precision_into_distributed_zero_mixed_precision': test_load_from_single_node_full_precision_into_distributed_zero_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_distributed_zero_full_precision': test_load_from_data_parallelism_full_precision_into_distributed_zero_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_distributed_zero_full_precision': test_load_from_data_parallelism_mixed_precision_into_distributed_zero_full_precision, + 'test_load_from_data_parallelism_mixed_precision_into_distributed_zero_mixed_precision': test_load_from_data_parallelism_mixed_precision_into_distributed_zero_mixed_precision, + 'test_load_from_data_parallelism_full_precision_into_distributed_zero_mixed_precision': test_load_from_data_parallelism_full_precision_into_distributed_zero_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_distributed_zero_full_precision': test_load_from_distributed_zero_full_precision_into_distributed_zero_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_distributed_zero_full_precision': test_load_from_distributed_zero_mixed_precision_into_distributed_zero_full_precision, + 'test_load_from_distributed_zero_mixed_precision_into_distributed_zero_mixed_precision': test_load_from_distributed_zero_mixed_precision_into_distributed_zero_mixed_precision, + 'test_load_from_distributed_zero_full_precision_into_distributed_zero_mixed_precision': test_load_from_distributed_zero_full_precision_into_distributed_zero_mixed_precision +} +parser = argparse.ArgumentParser(description='Test saved states of trainers to loaded states') +parser.add_argument('--scenario', choices=function_map.keys(), help='training scenario to test saved and loaded states', required=True) +parser.add_argument('--checkpoint_dir', help='path to the saved states directory', required=True) +args = parser.parse_args() +function_map[args.scenario](checkpoint_dir=args.checkpoint_dir) diff --git a/orttraining/orttraining/test/python/checkpoint/orttraining_test_save_checkpoint.py b/orttraining/orttraining/test/python/checkpoint/orttraining_test_save_checkpoint.py new file mode 100644 index 0000000000..0ce24ebe34 --- /dev/null +++ b/orttraining/orttraining/test/python/checkpoint/orttraining_test_save_checkpoint.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +################################################################################ +# Refer to orttraining_test_checkpoint.py for an overview about Checkpoint tests +################################################################################ + +import argparse +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from _test_helpers import distributed_setup, create_orttrainer_and_save_checkpoint + +def single_node_full_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/full_precision/'): + opts = {'device' : {'id' : device}, + 'debug' : {'deterministic_compute': True}} + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir) + +def single_node_mixed_precision(device = 'cuda', checkpoint_dir = 'checkpoint_dir/single_node/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'debug' : {'deterministic_compute': True} + } + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir) + +@distributed_setup +def data_parallelism_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir if world_rank == 0 else None) + +@distributed_setup +def data_parallelism_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/data_parallelism/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True + }, + 'debug' : {'deterministic_compute': True} + } + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir if world_rank == 0 else None) + +@distributed_setup +def distributed_zero_full_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/full_precision/'): + opts = { + 'device' : {'id' : device}, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir, state_dict_key_name='state_dict_'+str(world_rank)) + +@distributed_setup +def distributed_zero_mixed_precision(world_rank, world_size, device, checkpoint_dir = 'checkpoint_dir/distributed_zero/mixed_precision/'): + opts = { + 'device' : {'id' : device}, + 'mixed_precision': + { + 'enabled': True + }, + 'distributed' : + { + 'world_rank' : world_rank, + 'world_size' : world_size, + 'allreduce_post_accumulation' : True, + 'deepspeed_zero_optimization': + { + 'stage': 1 + } + }, + 'debug' : {'deterministic_compute': True} + } + create_orttrainer_and_save_checkpoint(device, opts, checkpoint_dir, state_dict_key_name='state_dict_'+str(world_rank)) + +function_map = { + 'single_node_full_precision': single_node_full_precision, + 'single_node_mixed_precision': single_node_mixed_precision, + 'data_parallelism_full_precision': data_parallelism_full_precision, + 'data_parallelism_mixed_precision': data_parallelism_mixed_precision, + 'distributed_zero_full_precision': distributed_zero_full_precision, + 'distributed_zero_mixed_precision': distributed_zero_mixed_precision +} +parser = argparse.ArgumentParser(description='Save states of trainers') +parser.add_argument('--scenario', choices=function_map.keys(), help='training scenario to save states', required=True) +parser.add_argument('--checkpoint_dir', help='path to the directory where checkpoints can be saved', required=True) +args = parser.parse_args() +function_map[args.scenario](checkpoint_dir = args.checkpoint_dir) diff --git a/orttraining/orttraining/test/python/orttraining_distributed_tests.py b/orttraining/orttraining/test/python/orttraining_distributed_tests.py index a4553130e4..159aef275c 100644 --- a/orttraining/orttraining/test/python/orttraining_distributed_tests.py +++ b/orttraining/orttraining/test/python/orttraining_distributed_tests.py @@ -19,6 +19,13 @@ def parse_arguments(): parser.add_argument("--cwd", help="Path to the current working directory") return parser.parse_args() +def run_checkpoint_tests(cwd, log): + log.debug('Running: Checkpoint tests') + + command = [sys.executable, 'orttraining_test_checkpoint.py'] + + run_subprocess(command, cwd=cwd, log=log).check_returncode() + def main(): import torch ngpus = torch.cuda.device_count() @@ -31,7 +38,7 @@ def main(): log.info("Running distributed tests pipeline") - # TODO: Add distributed test suite here. + run_checkpoint_tests(cwd, log) return 0 diff --git a/orttraining/orttraining/test/python/orttraining_test_checkpoint.py b/orttraining/orttraining/test/python/orttraining_test_checkpoint.py new file mode 100644 index 0000000000..f637f7875f --- /dev/null +++ b/orttraining/orttraining/test/python/orttraining_test_checkpoint.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import subprocess +import os +import shutil +import sys +import torch +from checkpoint._test_helpers import makedir + +def _single_run(execution_file, scenario, checkopint_dir): + assert subprocess.call([sys.executable, execution_file, '--scenario', scenario, '--checkpoint_dir', checkopint_dir]) == 0 + +def _distributed_run(execution_file, scenario, checkopint_dir): + assert subprocess.call(['mpirun', '-n', str(ngpus), '-x', 'NCCL_DEBUG=INFO', sys.executable, execution_file, '--scenario', scenario, '--checkpoint_dir', checkopint_dir]) == 0 + +checkpoint_dir = os.path.abspath('checkpoint/checkpoint_dir/') +makedir(checkpoint_dir) + +ngpus = torch.cuda.device_count() + +# test workflow: +# - there are a total of three files that are used for checkpointing tests: +# - orttraining_test_checkpoint.py: co-ordinating all the checkpoint tests +# - orttraining_test_save_checkpoint.py: responsible for saving all checkpoint files and trained states +# - orttraining_test_load_checkpoint.py: loading the saved checkpoints and the saved states and asserting whether +# the saved states match the loaded states. +# - and a total of 36 tests encompassing checkpointing tests: +# - from [onnxruntime orttrainer][full_precision, mixed_precision][single node training, data parallel training, distributed zero training] to +# [onnxruntime orttrainer, pytorch][full_precision, mixed_precision][single node training, data parallel training, distributed zero training] +# - all tests cannot be written in the same process because: +# - some of them require to be run in a distributed environment (using mpirun) while others can be run using a single process. +# - there is a known limitation where the distributed training run context is implemented as a singleton, so in the same process, no more than one +# orttrainer can be instantiated. Hence the need to run these tests in different processes one at a time. +# - workflow: +# - orttraining_test_checkpoint.py calls orttraining_test_save_checkpoint.py to save following files to disk +# - ORTTrainer checkpoint files through the experimental_save_checkpoint method +# - ORTTrainer states through pickle after extracting all the states of the ORTTrainer through the experimental_state_dict method +# - for each configuration across [onnxruntime orttrainer][full_precision, mixed_precision][single node training, data parallel training, distributed zero training] +# - orttraining_test_checkpoint.py calls orttraining_test_load_checkpoint.py to load each checkpoint into each orttrainer configuration +# - Saved ORTTrainer checkpoint files are loaded into an ORTTrainer using the experimental_load_checkpoint method for each ORTTrainer configuration. +# - Saved states are loaded into a python dictionary (called the state dictionary) through pickle +# - state dictionary is extracted from the ORTTrainer after it has loaded the checkpoint file and the onnx graph has been initialized (by calling eval_step) +# through the experimental_state_dict method. +# - the loaded state dictionary (through pickle) is compared against the extracted state dictionary for: +# - equality (or new equality) of model states +# - equality (or new equality) of optimizer states +# - In some cases the comparison is not directly possible; for example single node trainer to a distributed zero trainer because the extracted state +# dictionary is a distributed one and cannot be compared against a single node trainer directly. +# - First these states are saved using pickle for each rank to a file on disk +# - Wait for all ranks to complete writing the file to disk using barrier() +# - Load all states and aggregate them into 1 state dictionary +# - Compare this aggregated state dictionary against the original one loaded from disk. + +save_checkpoint_file = os.path.join('checkpoint', 'orttraining_test_save_checkpoint.py') +load_checkpoint_file = os.path.join('checkpoint', 'orttraining_test_load_checkpoint.py') + +single_node_full_precision_path = os.path.join(checkpoint_dir, 'single_node', 'full_precision') +single_node_mixed_precision_path = os.path.join(checkpoint_dir, 'single_node', 'mixed_precision') +data_parallelism_full_precision_path = os.path.join(checkpoint_dir, 'data_parallelism', 'full_precision') +data_parallelism_mixed_precision_path = os.path.join(checkpoint_dir, 'data_parallelism', 'mixed_precision') +distributed_zero_full_precision_path = os.path.join(checkpoint_dir, 'distributed_zero', 'full_precision') +distributed_zero_mixed_precision_path = os.path.join(checkpoint_dir, 'distributed_zero', 'mixed_precision') + +# save all checkpoint files (pre-checkpoint) +_single_run(save_checkpoint_file, 'single_node_full_precision', single_node_full_precision_path) +_single_run(save_checkpoint_file, 'single_node_mixed_precision', single_node_mixed_precision_path) +_distributed_run(save_checkpoint_file, 'data_parallelism_full_precision', data_parallelism_full_precision_path) +_distributed_run(save_checkpoint_file, 'data_parallelism_mixed_precision', data_parallelism_mixed_precision_path) +_distributed_run(save_checkpoint_file, 'distributed_zero_full_precision', distributed_zero_full_precision_path) +_distributed_run(save_checkpoint_file, 'distributed_zero_mixed_precision', distributed_zero_mixed_precision_path) + +# load checkpoint files (post-checkpoint) +# going to single node trainer +_single_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_single_node_full_precision', single_node_full_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_single_node_full_precision', single_node_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_single_node_mixed_precision', single_node_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_single_node_mixed_precision', single_node_full_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_single_node_full_precision', data_parallelism_full_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_single_node_full_precision', data_parallelism_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_single_node_mixed_precision', data_parallelism_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_single_node_mixed_precision', data_parallelism_full_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_single_node_full_precision', distributed_zero_full_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_single_node_full_precision', distributed_zero_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_single_node_mixed_precision', distributed_zero_mixed_precision_path) +_single_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_single_node_mixed_precision', distributed_zero_full_precision_path) + +# going to data parallel trainer +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_data_parallelism_full_precision', single_node_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_data_parallelism_full_precision', single_node_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_data_parallelism_mixed_precision', single_node_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_data_parallelism_mixed_precision', single_node_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_data_parallelism_full_precision', data_parallelism_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_data_parallelism_full_precision', data_parallelism_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_data_parallelism_mixed_precision', data_parallelism_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_data_parallelism_mixed_precision', data_parallelism_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_data_parallelism_full_precision', distributed_zero_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_data_parallelism_full_precision', distributed_zero_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_data_parallelism_mixed_precision', distributed_zero_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_data_parallelism_mixed_precision', distributed_zero_full_precision_path) + +# going to distributed zero trainer +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_distributed_zero_full_precision', single_node_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_distributed_zero_full_precision', single_node_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_mixed_precision_into_distributed_zero_mixed_precision', single_node_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_single_node_full_precision_into_distributed_zero_mixed_precision', single_node_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_distributed_zero_full_precision', data_parallelism_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_distributed_zero_full_precision', data_parallelism_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_mixed_precision_into_distributed_zero_mixed_precision', data_parallelism_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_data_parallelism_full_precision_into_distributed_zero_mixed_precision', data_parallelism_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_distributed_zero_full_precision', distributed_zero_full_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_distributed_zero_full_precision', distributed_zero_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_mixed_precision_into_distributed_zero_mixed_precision', distributed_zero_mixed_precision_path) +_distributed_run(load_checkpoint_file, 'test_load_from_distributed_zero_full_precision_into_distributed_zero_mixed_precision', distributed_zero_full_precision_path) + +shutil.rmtree(checkpoint_dir)