Add unit tests to identify configuration migration scenarios for checkpointing (#5678)

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 $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/transformers
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/quantization
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/quantization/operators
+  COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/checkpoint
   COMMAND ${CMAKE_COMMAND} -E copy
       ${ONNXRUNTIME_ROOT}/__init__.py
       $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/
@@ -256,6 +260,9 @@ add_custom_command(
   COMMAND ${CMAKE_COMMAND} -E copy
       ${onnxruntime_python_test_srcs}
       $<TARGET_FILE_DIR:${test_data_target}>
+  COMMAND ${CMAKE_COMMAND} -E copy
+      ${onnxruntime_python_checkpoint_test_srcs}
+      $<TARGET_FILE_DIR:${test_data_target}>/checkpoint/
   COMMAND ${CMAKE_COMMAND} -E copy
       ${onnxruntime_backend_srcs}
       $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/backend/

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)

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)

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
 

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)