Repository: microsoft/DeepSpeed
File: deepspeed/checkpoint/ds_to_universal.py (Python)

Code before:
            slice_shapes += mp_sd[PARAM_SHAPES]

        # fix back to normal flat dict, merge duplicates for tp>1
        slice_shapes = dict((k, v) for d in slice_shapes for k, v in d.items())
        temp_dir = os.path.join(args.output_folder, 'tmp')

        print('*** 1. Extracting ZeRO fragments')
        _extract_zero_shard_files(args, ds_checkpoint, temp_dir)

        print('*** 2. Merging slices .....')
        _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir)

        print('*** 3. Saving common optimizer states')
        _save_optimizer_state(args, ds_checkpoint)

        if not args.keep_temp_folder:
            shutil.rmtree(temp_dir, ignore_errors=True)

        # Copy mp* files into output folder
        for f in glob.glob(os.path.join(args.input_folder, 'mp*')):
            shutil.copy2(f, args.output_folder)

    else:
        model_files = _get_model_state_files(args.input_folder)
        param_shapes = _parse_model_states_stage3(model_files)
        dp_degree = len(optim_files)

        temp_dir = os.path.join(args.output_folder, 'tmp')

        print('*** 1. Extracting ZeRO fragments')
        _extract_zero_shard_files_stage3(args, optim_files, param_shapes, dp_degree, temp_dir)

        print('*** 2. Merging slices .....')
        param_keys = {key for sub_group_shapes in param_shapes for key in sub_group_shapes.keys()}
        _merge_zero3_slice_files(args, param_keys, dp_degree, temp_dir)

        print('*** 3. Saving common optimizer states')
        _save_optimizer_state_stage3(args, optim_files)

        if not args.keep_temp_folder:
            shutil.rmtree(temp_dir, ignore_errors=True)

        # Copy *model_states files into output folder
        for f in glob.glob(os.path.join(args.input_folder, '*model_states.pt')):
            shutil.copy2(f, args.output_folder)

    # Update latest to output folder
    checkpoint_root_folder, step_folder = os.path.split(args.output_folder)
    latest_file = os.path.join(checkpoint_root_folder, 'latest_universal')
    with open(latest_file, "w") as f:
        f.write(step_folder)

Review comment (performance):
How about let's still use `dp_degree = len(model_files)`? By design, the number of `optim_files` should match the number of `model_files`, since each DP worker saves both a model file and an optimizer file. However, if a developer chooses to resume training using only the weights and deletes the optimizer state files, then `len(optim_files) = 0`, which would cause issues.

Code after:
            slice_shapes += mp_sd[PARAM_SHAPES]

        # fix back to normal flat dict, merge duplicates for tp>1
        slice_shapes = dict((k, v) for d in slice_shapes for k, v in d.items())
        temp_dir = os.path.join(args.output_folder, 'tmp')

        print('*** 1. Extracting ZeRO fragments')
        _extract_zero_shard_files(args, ds_checkpoint, temp_dir)

        print('*** 2. Merging slices .....')
        _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir)

        print('*** 3. Saving common optimizer states')
        _save_optimizer_state(args, ds_checkpoint)

        if not args.keep_temp_folder:
            shutil.rmtree(temp_dir, ignore_errors=True)

        # Copy mp* files into output folder
        for f in glob.glob(os.path.join(args.input_folder, 'mp*')):
            shutil.copy2(f, args.output_folder)

    else:
        model_files = _get_model_state_files(args.input_folder)
        param_shapes = _parse_model_states_stage3(model_files)
        dp_degree = len(model_files)

        temp_dir = os.path.join(args.output_folder, 'tmp')

        print('*** 1. Extracting ZeRO fragments')
        _extract_zero_shard_files_stage3(args, optim_files, param_shapes, dp_degree, temp_dir)

        print('*** 2. Merging slices .....')
        param_keys = {key for sub_group_shapes in param_shapes for key in sub_group_shapes.keys()}
        _merge_zero3_slice_files(args, param_keys, dp_degree, temp_dir)

        print('*** 3. Saving common optimizer states')
        _save_optimizer_state_stage3(args, optim_files)

        if not args.keep_temp_folder:
            shutil.rmtree(temp_dir, ignore_errors=True)

        # Copy *model_states files into output folder
        for f in glob.glob(os.path.join(args.input_folder, '*model_states.pt')):
            shutil.copy2(f, args.output_folder)

    # Update latest to output folder
    checkpoint_root_folder, step_folder = os.path.split(args.output_folder)
    latest_file = os.path.join(checkpoint_root_folder, 'latest_universal')
    with open(latest_file, "w") as f:
        f.write(step_folder)