deepspeedai--deepspeed/deepspeed/checkpoint/ds_to_universal.py

#!/usr/bin/env python

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

from functools import partial
from itertools import chain
import argparse
import glob
import itertools
import math
from concurrent.futures import ProcessPoolExecutor
import os
import re
import shutil
import torch
import tqdm
#from pprint import pprint

from deepspeed.checkpoint import DeepSpeedCheckpoint
from deepspeed.checkpoint import (
    OPTIMIZER_STATE_DICT,
    ZERO_STAGE,
    BASE_OPTIMIZER_STATE,
    SINGLE_PARTITION_OF_FP32_GROUPS,
    PARAM_GROUPS,
    PARAM_SLICE_MAPPINGS,
    PARAM_SHAPES,
    PARAM,
    CAT_DIM,
    PARAM_N_SUB_PARAMS,
    SUB_PARAM_SHAPE,
    VOCAB_TENSOR,
    UNIVERSAL_CHECKPOINT_INFO,
    UNIVERSAL_CHECKPOINT_VERSION_KEY,
    UNIVERSAL_CHECKPOINT_VERSION_VALUE,
    VOCABULARY_PARAMETER_PATTERNS,
    PIPELINE_REPLICATED_PARAMETER_PATTERNS,
    TP_REPLICATED_PARAMETER_PATTERNS,
    PARAMETER_TO_AVERAGE_PATTERNS,
    PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
    PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0,
    PARAMETER_WITH_SUB_PARAMS,
    AUTOEP_LAYERS_KEY,
    AUTOEP_LAYERS_KEY_LEGACY,
    EP_IS_EXPERT_PARAM,
    EP_NUM_EXPERTS,
    EXPERT_PARAMETER_PATTERNS,
    SubparamShape,
)
from deepspeed.checkpoint.autoep_zero3_metadata import (
    is_autoep_zero3_partitioned_entry,
    validate_autoep_zero3_partitioned_metadata,
)


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, required=True, help='Input DeepSpeed Checkpoint folder')
    parser.add_argument('--output_folder', type=str, required=True, help='Output DeepSpeed checkpoint folder')
    parser.add_argument('--num_extract_workers',
                        default=4,
                        type=int,
                        help='How many parallel processes to extract zero shards')
    parser.add_argument(
        '--num_merge_workers',
        default=2,
        type=int,
        help=
        'How many parallel processes to merge tp slices (more memory intensive, use much fewer than --num_extract_workers))'
    )
    parser.add_argument('--keep_temp_folder',
                        action='store_true',
                        help='Preserve temporary folder of intermediate checkpoint slice files. Useful for debugging.')
    parser.add_argument('--no_strict',
                        dest='strict',
                        action='store_false',
                        help='Do not perform validity checks on converted checkpoint.')
    parser.add_argument('--inject_missing_state',
                        action='store_true',
                        help='Inject missing checkpoint state into the checkpoint if it is absent.')
    args = parser.parse_args()
    print(f'args = {args}')
    return args


def atoi(text):
    return int(text) if text.isdigit() else text


def natural_keys(text):
    '''
    alist.sort(key=natural_keys) sorts in human order
    http://nedbatchelder.com/blog/200712/human_sorting.html
    (See Toothy's implementation in the comments)
    '''
    return [atoi(c) for c in re.split(r'(\d+)', text)]


def _create_checkpoint_paths(base_folder, iteration, tp_degree, pp_degree):
    path_list = []
    iter_folder = f'iter_{iteration:07d}'
    for i in range(0, tp_degree):
        path_list.append([])
        for j in range(0, pp_degree):
            rank_folder = f'mp_rank_{i:02d}' if pp_degree == 1 else f'mp_rank_{i:02d}_{j:03d}'
            ckpt_path = os.path.join(rank_folder, 'model_optim_rng.pt')
            path_list[i].append(os.path.join(base_folder, iter_folder, ckpt_path))

    return path_list


def _save_checkpoint(file_path, chkpt_sd):
    dir, _ = os.path.split(file_path)
    os.makedirs(dir, exist_ok=True)
    torch.save(chkpt_sd, file_path)


def extract_zero_shards(dir, ds_checkpoint, indices_3D):
    pp_index, tp_index, dp_index = indices_3D
    sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=pp_index, tp_index=tp_index, dp_index=dp_index)

    # pprint(f"Processing {dp_index=} {pp_index=}, {tp_index=}")

    optim_sd = sd[OPTIMIZER_STATE_DICT]
    param_slice_mappings = optim_sd[PARAM_SLICE_MAPPINGS]
    universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO)
    pipeline_replicated_params = universal_checkpoint_info.get(PIPELINE_REPLICATED_PARAMETER_PATTERNS, [])
    # print(f'{pipeline_replicated_params=}')

    # dict
    state_groups = optim_sd[BASE_OPTIMIZER_STATE]["state"]
    # list
    fp32_groups = optim_sd[SINGLE_PARTITION_OF_FP32_GROUPS]
    param_groups_cnt = len(state_groups)

    for param_group_id in range(param_groups_cnt):

        flat_state = dict(
            exp_avg=state_groups[param_group_id]["exp_avg"],
            exp_avg_sq=state_groups[param_group_id]["exp_avg_sq"],
            fp32=fp32_groups[param_group_id],
        )

        if "step" in state_groups[param_group_id]:
            flat_state["step"] = state_groups[param_group_id]["step"]

        for name, fragment_mapping in param_slice_mappings[param_group_id].items():
            if pp_index > 0 and any(re.match(pattern, name) for pattern in pipeline_replicated_params):
                # Skip tied weights that are replicated in first and last pp stages
                continue

            # pprint(f"dpt{dp_index}{pp_index}{tp_index} {param_group_id} {name} => {fragment_mapping.start}:{fragment_mapping.numel}")
            for state_key in flat_state.keys():
                dump_param_fragment(dir, tp_index, dp_index, state_key, flat_state[state_key], name,
                                    fragment_mapping.start, fragment_mapping.numel)


def extract_zero_shards_stage3(optim_files, param_shapes, dp_degree, temp_dir, dp_index, exclude_param_names=None):
    exclude_param_names = exclude_param_names or set()
    state_dict = torch.load(optim_files[dp_index], map_location='cpu', weights_only=False)
    optim_sd = state_dict[OPTIMIZER_STATE_DICT]
    partition_groups = optim_sd.get('ds_zero_partition_groups') or []

    for idx, sub_group_shape in enumerate(param_shapes):
        flat_state = dict(
            exp_avg=optim_sd['optimizer_state_dict']['state'][idx]["exp_avg"],
            exp_avg_sq=optim_sd['optimizer_state_dict']['state'][idx]["exp_avg_sq"],
            fp32=optim_sd['fp32_flat_groups'][idx],
        )
        partition_metadata = partition_groups[idx] if idx < len(partition_groups) else {}
        partition_count = partition_metadata.get('partition_count', dp_degree)
        partition_rank = partition_metadata.get('partition_rank', dp_index)
        offset = 0
        for name, shape in sub_group_shape.items():
            unpartitioned_numel = _shape_numel(shape)
            partitioned_numel, _ = _zero_partitioned_param_info(unpartitioned_numel, partition_count)
            padding_free_numel = max(0, min(partitioned_numel,
                                            unpartitioned_numel - partition_rank * partitioned_numel))
            if name not in exclude_param_names:
                for state_key in flat_state.keys():
                    dump_param_fragment(temp_dir, 0, dp_index, state_key, flat_state[state_key], name, offset,
                                        padding_free_numel)
            offset += partitioned_numel


cnt = 0


def dp_index_to_str(dp_index):
    return f"{dp_index:0>2d}"


def dump_param_fragment(dir, tp_index, dp_index, state_name, state_flat_tensor, param_name, offset, numel):

    global cnt  # temp hack

    param_base_path = os.path.join(dir, param_name, str(tp_index))
    os.makedirs(param_base_path, exist_ok=True)

    cnt += 1

    path = os.path.join(param_base_path, f"{state_name}.{dp_index_to_str(dp_index)}")

    #print(f"{param_name}: {offset}: {numel} => {path}")

    # State might be a python int or a tensor
    if state_name != "step" and torch.is_tensor(state_flat_tensor):
        state_flat_tensor = state_flat_tensor.narrow(0, offset, numel).clone()
    _save_checkpoint(path, state_flat_tensor)


def _merge_zero_shards(param_base_path, state, tp_degree, slice_shape=None):
    slices = []
    for tp_index in range(tp_degree):
        prefix_path = os.path.join(param_base_path, str(tp_index), f"{state}")
        paths = glob.glob(f"{prefix_path}.*")

        if len(paths) == 0:
            continue

        pattern = re.compile(f"{prefix_path}\\.([0-9]+)")
        dp_indices = set()
        for p in paths:
            m = pattern.match(p)
            if m:
                dp_indices.add(int(m.group(1)))
            else:
                raise ValueError(f"Cannot parse dp_rank from {p}")

        paths = [f"{prefix_path}.{dp_index_to_str(dp_index)}" for dp_index in sorted(list(dp_indices))]
        shards = [torch.load(p, weights_only=False) for p in paths]

        if state == "step":
            assert all(v == shards[0] for v in shards), "All shards must have the same step value"
            slice = shards[0]
        else:
            if slice_shape is None:
                slice = torch.cat(shards, dim=0)
            else:
                slice = torch.cat(shards, dim=0).reshape(slice_shape)

        slices.append(slice)
    return slices


def merge_tp_slices(ds_checkpoint, dir, slice_dir, tp_degree, name_and_shape):

    name, shape = name_and_shape
    slice_base_path = os.path.join(slice_dir, name)
    param_base_path = os.path.join(dir, name)

    universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO)
    replicated_parameters = universal_checkpoint_info.get(TP_REPLICATED_PARAMETER_PATTERNS, [])
    parameters_to_average = universal_checkpoint_info.get(PARAMETER_TO_AVERAGE_PATTERNS, [])
    parameters_with_row_parallelism = universal_checkpoint_info.get(PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, [])
    vocabulary_parameters = universal_checkpoint_info.get(VOCABULARY_PARAMETER_PATTERNS, [])
    parameters_with_2_sub_params_cat_dim_0 = universal_checkpoint_info.get(PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, [])
    parameter_with_sub_params = universal_checkpoint_info.get(PARAMETER_WITH_SUB_PARAMS, [])

    unmatched_patterns = set(replicated_parameters + parameters_to_average + parameters_with_row_parallelism +
                             vocabulary_parameters + parameters_with_2_sub_params_cat_dim_0)
    unmatched_patterns.update(chain.from_iterable(SubparamShape(**s).patterns for s in parameter_with_sub_params))

    def get_matched_pattern(patterns_, name_):
        matched_ = [pattern_ for pattern_ in patterns_ if re.match(pattern_, name_)]
        assert len(matched_) <= 1, f'Got more than one matching patterns={matched_} for {name_}'
        if matched_:
            pattern_ = matched_[0]
            unmatched_patterns.discard(pattern_)
            return pattern_
        return None

    def get_matched_sub_params_pattern(name_):
        for subparam_shape_dict in parameter_with_sub_params:
            subparam_shape = SubparamShape(**subparam_shape_dict)
            for pattern_ in subparam_shape.patterns:
                if re.match(pattern_, name_):
                    unmatched_patterns.discard(pattern_)
                    return subparam_shape
        return None

    matched_sub_params_shape = get_matched_sub_params_pattern(name)

    step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, shape)
    if step_merged:
        _save_checkpoint(os.path.join(param_base_path, "step.pt"), step_merged[0])

    for state in ("fp32", "exp_avg", "exp_avg_sq"):
        slices = _merge_zero_shards(slice_base_path, state, tp_degree, shape)
        final_path = os.path.join(param_base_path, f"{state}.pt")

        #print(f"Expected shape: {shape}")
        #print(f"Fragment sizes:", list(frag.shape for frag in slices))
        ckpt_dict = {}
        if get_matched_pattern(replicated_parameters, name):
            if len(slices) > 1:
                assert all([slices[0].equal(other_slice) for other_slice in slices[1:]])
            param = slices[0]
            # print(f'replicate {name} using first slice')
        elif get_matched_pattern(parameters_to_average, name):
            param = sum(slices) / len(slices)
            # print(f'merge {name} using average')
        elif get_matched_pattern(parameters_with_2_sub_params_cat_dim_0, name):
            cat_dim = 0
            chunked_slices = [torch.chunk(s, 2, dim=cat_dim) for s in slices]
            merged_chunks_0 = torch.cat([s[0] for s in chunked_slices], dim=cat_dim)
            merged_chunks_1 = torch.cat([s[1] for s in chunked_slices], dim=cat_dim)
            param = torch.cat([merged_chunks_0, merged_chunks_1], dim=cat_dim)
            ckpt_dict[CAT_DIM] = cat_dim
            ckpt_dict[PARAM_N_SUB_PARAMS] = 2
        elif matched_sub_params_shape:
            merged_chunks = []
            partition_dim = matched_sub_params_shape.partition_dim

            sub_dim_sizes = matched_sub_params_shape.shape[partition_dim]
            if not isinstance(sub_dim_sizes, tuple):
                sub_dim_sizes = (sub_dim_sizes, )

            partition_shape = [sum(d) if isinstance(d, tuple) else d for d in matched_sub_params_shape.shape]
            partition_shape = [d // tp_degree if i == partition_dim else d for i, d in enumerate(partition_shape)]
            slices = [s.view(partition_shape) for s in slices]

            offset = 0
            for sub_dim_size in sub_dim_sizes:
                part_sub_dim_size = sub_dim_size // tp_degree
                merged_chunks.append(
                    torch.cat([s.narrow(partition_dim, offset, part_sub_dim_size) for s in slices], dim=partition_dim))
                offset += part_sub_dim_size
            param = torch.cat(merged_chunks, dim=partition_dim)
            ckpt_dict[SUB_PARAM_SHAPE] = matched_sub_params_shape
        else:
            cat_dim = 1 if get_matched_pattern(parameters_with_row_parallelism, name) else 0
            # print(f"merge {name} with CAT DIM: {cat_dim}")
            param = torch.cat(slices, dim=cat_dim)
            ckpt_dict[CAT_DIM] = cat_dim

        if get_matched_pattern(vocabulary_parameters, name):
            #print(f"Before {param.shape=}")
            # strip padding
            original_vocab_size = universal_checkpoint_info['original_vocab_size']
            param = param[:original_vocab_size, :]
            ckpt_dict[VOCAB_TENSOR] = True
            #print(f"After {param.shape=}")

        #print(f"Final shape: {param.shape}")
        ckpt_dict[PARAM] = param
        _save_checkpoint(final_path, ckpt_dict)

    return unmatched_patterns


def merge_zero3_slices(dp_degree, dir, slice_dir, name):
    slice_base_path = os.path.join(slice_dir, name)
    param_base_path = os.path.join(dir, name)

    for state in ("fp32", "exp_avg", "exp_avg_sq"):
        slices = _merge_zero_shards(slice_base_path, state, 1)
        final_path = os.path.join(param_base_path, f"{state}.pt")
        _save_checkpoint(final_path, slices[0])


def _do_parallel_work(do_work, work_chunks, num_workers):
    results = []
    if num_workers > 1:
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            future_list = [executor.submit(do_work, work) for work in work_chunks]
            for f in tqdm.tqdm(future_list):
                results.append(f.result())
    else:
        # No parallel pass for unit testing
        # We can't create child processes in tests
        for work in tqdm.tqdm(work_chunks):
            results.append(do_work(work))
    return results


def _extract_zero_shard_files(args, ds_checkpoint, temp_dir):
    _3d_range_list = list(
        itertools.product(range(ds_checkpoint.pp_degree), range(ds_checkpoint.tp_degree),
                          range(ds_checkpoint.dp_degree)))
    #pprint(f'{_3d_range_list=}')

    do_work = partial(extract_zero_shards, temp_dir, ds_checkpoint)
    _do_parallel_work(do_work, _3d_range_list, args.num_extract_workers)


def _extract_zero_shard_files_stage3(args, optim_files, param_shapes, dp_degree, temp_dir, exclude_param_names=None):
    do_work = partial(extract_zero_shards_stage3,
                      optim_files,
                      param_shapes,
                      dp_degree,
                      temp_dir,
                      exclude_param_names=exclude_param_names)
    _do_parallel_work(do_work, list(range(dp_degree)), args.num_extract_workers)


def _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir):
    zero_output_folder = os.path.join(args.output_folder, "zero")
    do_work = partial(merge_tp_slices, ds_checkpoint, zero_output_folder, temp_dir, ds_checkpoint.tp_degree)
    unmatched_patterns_lists = _do_parallel_work(do_work, list(slice_shapes.items()), args.num_merge_workers)

    # verify that all patterns were used
    # if a pattern was not used by any of the workers, then it was not used at all -> assert/alert
    sets = [set(lst) for lst in unmatched_patterns_lists]
    unmatched_patterns = list(set.intersection(*sets))
    if args.strict:
        assert not unmatched_patterns, f'Unused patterns={unmatched_patterns} while merging tp slices'
    elif unmatched_patterns:
        print(f'Warning: Unused patterns={unmatched_patterns} while merging tp slices')


def _merge_zero3_slice_files(args, param_keys, dp_degree, temp_dir):
    zero_output_folder = os.path.join(args.output_folder, "zero")
    do_work = partial(merge_zero3_slices, dp_degree, zero_output_folder, temp_dir)
    _do_parallel_work(do_work, param_keys, args.num_merge_workers)


def _zero_partitioned_param_info(unpartitioned_numel, world_size):
    remainder = unpartitioned_numel % world_size
    padding_numel = (world_size - remainder) if remainder else 0
    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
    return partitioned_numel, padding_numel


def _shape_numel(shape):
    if hasattr(shape, "numel"):
        return shape.numel()
    return math.prod(shape)


def _zero3_rank_from_file(path):
    match = re.search(r'(?:bf16_)?zero_pp_rank_([0-9]+)_mp_rank_', os.path.basename(path))
    if match is None:
        raise ValueError(f"Cannot parse ZeRO rank from checkpoint file name: {path}")
    return int(match.group(1))


def _get_autoep_metadata(model_state):
    autoep_metadata = model_state.get(AUTOEP_LAYERS_KEY)
    if autoep_metadata is None:
        autoep_metadata = model_state.get(AUTOEP_LAYERS_KEY_LEGACY)
    return autoep_metadata


def _uses_zero3_partitioned_autoep_metadata(autoep_metadata):
    if not isinstance(autoep_metadata, list):
        return False
    _validate_zero3_partitioned_autoep_metadata(autoep_metadata, require_partitioned=False)
    return any(is_autoep_zero3_partitioned_entry(entry) for entry in autoep_metadata)


def _validate_zero3_partitioned_autoep_metadata(autoep_metadata, require_partitioned=True):
    validate_autoep_zero3_partitioned_metadata(autoep_metadata,
                                               require_partitioned=require_partitioned,
                                               version_context="This converter")


def _autoep_expert_param_info(autoep_metadata):
    info = {}
    if not isinstance(autoep_metadata, list):
        return info
    _validate_zero3_partitioned_autoep_metadata(autoep_metadata)
    for entry in autoep_metadata:
        if not isinstance(entry, dict):
            continue
        if not is_autoep_zero3_partitioned_entry(entry):
            continue
        prefix = entry.get('expert_key_prefix')
        if not prefix:
            continue
        for wname in ('w1', 'w2', 'w3'):
            info[f"{prefix}.{wname}"] = entry
    return info


def _autoep_expert_param_names_by_rank(model_files):
    expert_param_names = set()
    metadata_by_rank = {}
    for model_file in model_files:
        rank = _zero3_rank_from_file(model_file)
        model_state = torch.load(model_file, map_location=torch.device('cpu'), weights_only=False)
        autoep_metadata = _get_autoep_metadata(model_state)
        if autoep_metadata is not None:
            metadata_by_rank[rank] = autoep_metadata
            if _uses_zero3_partitioned_autoep_metadata(autoep_metadata):
                expert_param_names.update(_autoep_expert_param_info(autoep_metadata))
    return expert_param_names, metadata_by_rank


def _rank_map_from_files(files, description):
    rank_map = {}
    for path in files:
        rank = _zero3_rank_from_file(path)
        if rank in rank_map:
            raise RuntimeError(f"Duplicate ZeRO rank {rank} in {description} files: "
                               f"{rank_map[rank]} and {path}")
        rank_map[rank] = path
    return rank_map


def _validate_zero3_model_optim_rank_sets(model_files, optim_files):
    model_rank_map = _rank_map_from_files(model_files, "model-state")
    optim_rank_map = _rank_map_from_files(optim_files, "optimizer-state")
    model_ranks = set(model_rank_map)
    optim_ranks = set(optim_rank_map)
    if model_ranks != optim_ranks:
        raise RuntimeError("ZeRO-3 checkpoint model/optimizer rank sets do not match: "
                           f"model_only={sorted(model_ranks - optim_ranks)}, "
                           f"optim_only={sorted(optim_ranks - model_ranks)}")
    if not model_ranks:
        raise RuntimeError("ZeRO-3 checkpoint has no model/optimizer rank files")
    return model_rank_map, optim_rank_map


def _validate_autoep_expert_shapes(model_states_by_rank, metadata_by_rank):
    for rank, autoep_metadata in metadata_by_rank.items():
        if not _uses_zero3_partitioned_autoep_metadata(autoep_metadata):
            continue
        expert_info = _autoep_expert_param_info(autoep_metadata)
        param_shapes = model_states_by_rank[rank][PARAM_SHAPES]
        zero_shape_names = {name for sub_group_shape in param_shapes for name in sub_group_shape}
        missing = set(expert_info) - zero_shape_names
        if missing:
            raise RuntimeError(f"AutoEP expert parameters are missing from rank {rank} ZeRO param_shapes: "
                               f"{sorted(missing)}")
        frozen_shapes = model_states_by_rank[rank].get('frozen_param_shapes') or {}
        frozen_experts = set(expert_info).intersection(frozen_shapes)
        if frozen_experts:
            raise RuntimeError("AutoEP frozen expert parameters cannot be converted from the ZeRO-3 "
                               f"partition-native format yet: {sorted(frozen_experts)}")


def _save_zero3_autoep_universal_tensor(output_dir, param_name, state_key, tensor, num_experts):
    param_dir = os.path.join(output_dir, "zero", param_name)
    os.makedirs(param_dir, exist_ok=True)
    _save_checkpoint(
        os.path.join(param_dir, f"{state_key}.pt"),
        {
            PARAM: tensor,
            CAT_DIM: 0,
            EP_IS_EXPERT_PARAM: True,
            EP_NUM_EXPERTS: num_experts,
        },
    )


def _consolidate_zero3_autoep_expert_states(output_dir, model_files, optim_files):
    model_rank_map, optim_rank_map = _validate_zero3_model_optim_rank_sets(model_files, optim_files)
    model_states_by_rank = {
        rank: torch.load(model_file, map_location=torch.device('cpu'), weights_only=False)
        for rank, model_file in model_rank_map.items()
    }
    optim_states_by_rank = {
        rank: torch.load(optim_file, map_location=torch.device('cpu'), weights_only=False)
        for rank, optim_file in optim_rank_map.items()
    }
    metadata_by_rank = {
        rank: _get_autoep_metadata(model_state)
        for rank, model_state in model_states_by_rank.items() if _get_autoep_metadata(model_state) is not None
    }
    _validate_autoep_expert_shapes(model_states_by_rank, metadata_by_rank)

    expert_fragments = {}
    num_experts_by_param = {}
    expected_dp_world_by_param_rank = {}
    expected_ep_ranks_by_param = {}

    for rank, model_state in model_states_by_rank.items():
        optim_state = optim_states_by_rank.get(rank)
        if optim_state is None:
            raise FileNotFoundError(f"Missing ZeRO optimizer checkpoint for rank {rank}")

        autoep_metadata = _get_autoep_metadata(model_state)
        if not _uses_zero3_partitioned_autoep_metadata(autoep_metadata):
            continue

        expert_info = _autoep_expert_param_info(autoep_metadata)
        param_shapes = model_state[PARAM_SHAPES]
        zero_optim_state = optim_state[OPTIMIZER_STATE_DICT]
        partition_groups = zero_optim_state.get('ds_zero_partition_groups') or []

        for sub_group_id, sub_group_shape in enumerate(param_shapes):
            optimizer_sub_state = zero_optim_state['optimizer_state_dict']['state'][sub_group_id]
            flat_state = {
                'fp32': zero_optim_state['fp32_flat_groups'][sub_group_id],
                'exp_avg': optimizer_sub_state.get('exp_avg'),
                'exp_avg_sq': optimizer_sub_state.get('exp_avg_sq'),
            }
            partition_metadata = partition_groups[sub_group_id] if sub_group_id < len(partition_groups) else {}
            partition_count = partition_metadata.get('partition_count', len(model_states_by_rank))
            partition_rank = partition_metadata.get('partition_rank', rank)

            offset = 0
            for param_name, shape in sub_group_shape.items():
                unpartitioned_numel = _shape_numel(shape)
                partitioned_numel, _ = _zero_partitioned_param_info(unpartitioned_numel, partition_count)
                padding_free_numel = max(
                    0, min(partitioned_numel, unpartitioned_numel - partition_rank * partitioned_numel))

                layer_info = expert_info.get(param_name)
                if layer_info is not None:
                    ep_rank = layer_info['ep_rank']
                    num_experts_by_param[param_name] = layer_info['num_experts']
                    expected_dp_world_by_param_rank[(param_name,
                                                     ep_rank)] = layer_info['expert_data_parallel_world_size']
                    expected_ep_ranks_by_param[param_name] = set(range(layer_info['ep_size']))
                    for state_key, flat_tensor in flat_state.items():
                        if flat_tensor is None:
                            raise RuntimeError(f"Missing optimizer state '{state_key}' for AutoEP expert "
                                               f"parameter {param_name} on ZeRO rank {rank}")
                        fragment = flat_tensor.narrow(0, offset, padding_free_numel).clone()
                        key = (param_name, state_key, ep_rank)
                        expert_fragments.setdefault(key, []).append((partition_rank, fragment, shape))

                offset += partitioned_numel

    grouped_by_param = {}
    for (param_name, state_key, ep_rank), fragments in expert_fragments.items():
        grouped_by_param.setdefault((param_name, state_key), {})[ep_rank] = fragments

    for (param_name, state_key), ep_rank_fragments in grouped_by_param.items():
        missing_ep_ranks = expected_ep_ranks_by_param[param_name] - set(ep_rank_fragments)
        if missing_ep_ranks:
            raise RuntimeError(f"Missing AutoEP universal fragments for {param_name}/{state_key} EP ranks: "
                               f"{sorted(missing_ep_ranks)}")
        ep_tensors = []
        for ep_rank in sorted(ep_rank_fragments):
            fragments = sorted(ep_rank_fragments[ep_rank], key=lambda item: item[0])
            expected_dp_world = expected_dp_world_by_param_rank[(param_name, ep_rank)]
            partition_ranks = [partition_rank for partition_rank, _, _ in fragments]
            if len(partition_ranks) != len(set(partition_ranks)):
                raise RuntimeError(f"Duplicate AutoEP expert-DP partition ranks for {param_name}/{state_key} "
                                   f"EP rank {ep_rank}: {partition_ranks}")
            if set(partition_ranks) != set(range(expected_dp_world)):
                raise RuntimeError(f"Incomplete AutoEP expert-DP fragments for {param_name}/{state_key} "
                                   f"EP rank {ep_rank}: got {sorted(partition_ranks)}, "
                                   f"expected {list(range(expected_dp_world))}")
            shape = fragments[0][2]
            if any(tuple(fragment_shape) != tuple(shape) for _, _, fragment_shape in fragments):
                raise RuntimeError(f"Inconsistent AutoEP expert fragment shapes for {param_name}/{state_key} "
                                   f"EP rank {ep_rank}")
            full_flat = torch.cat([fragment for _, fragment, _ in fragments], dim=0)[:_shape_numel(shape)]
            ep_tensors.append(full_flat.view(shape))

        if not ep_tensors:
            continue
        full_expert_tensor = torch.cat(ep_tensors, dim=0)
        if full_expert_tensor.shape[0] != num_experts_by_param[param_name]:
            raise RuntimeError(f"AutoEP universal tensor for {param_name}/{state_key} has wrong expert dimension: "
                               f"got {full_expert_tensor.shape[0]}, expected {num_experts_by_param[param_name]}")
        _save_zero3_autoep_universal_tensor(output_dir, param_name, state_key, full_expert_tensor,
                                            num_experts_by_param[param_name])


def _parse_model_states_stage3(files):
    return torch.load(files[0], map_location=torch.device('cpu'), weights_only=False)[PARAM_SHAPES]


def _save_optimizer_state(args, ds_checkpoint):
    sharded_states = [BASE_OPTIMIZER_STATE, PARAM_SLICE_MAPPINGS, SINGLE_PARTITION_OF_FP32_GROUPS]
    sd = ds_checkpoint.get_zero_checkpoint_state(pp_index=0, tp_index=0, dp_index=0)

    optim_sd = sd[OPTIMIZER_STATE_DICT]
    output_sd = {k: v for k, v in optim_sd.items() if k not in sharded_states}
    output_sd[PARAM_GROUPS] = optim_sd[BASE_OPTIMIZER_STATE][PARAM_GROUPS]
    zero_output_folder = os.path.join(args.output_folder, "zero")
    output_file_path = os.path.join(zero_output_folder, "optimizer_state.pt")
    _save_checkpoint(output_file_path, output_sd)


def _save_optimizer_state_stage3(args, optim_files):
    sd = torch.load(optim_files[0], map_location=torch.device('cpu'), weights_only=False)
    output_sd = sd[OPTIMIZER_STATE_DICT]
    output_sd[PARAM_GROUPS] = output_sd[OPTIMIZER_STATE_DICT][PARAM_GROUPS]
    zero_output_folder = os.path.join(args.output_folder, "zero")
    output_file_path = os.path.join(zero_output_folder, "optimizer_state.pt")
    _save_checkpoint(output_file_path, output_sd)


def _get_optim_files(checkpoint_dir):
    return _get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")


def _filter_zero3_optim_files(optim_files):
    return [f for f in optim_files if re.match(r'(?:bf16_)?zero_pp_rank_', os.path.basename(f))]


def _get_model_state_files(checkpoint_dir):
    return _get_checkpoint_files(checkpoint_dir, "*_model_states.pt")


def _is_expert_model_state_file(checkpoint_file):
    basename = os.path.basename(checkpoint_file)
    return basename.startswith('layer_') and '_expert_' in basename


def _get_zero3_model_state_files(checkpoint_dir):
    model_files = [f for f in _get_model_state_files(checkpoint_dir) if not _is_expert_model_state_file(f)]

    if len(model_files) == 0:
        raise FileNotFoundError(f"can't find ZeRO Stage 3 model state files in directory '{checkpoint_dir}'")

    return model_files


def _raise_if_stage3_autoep_universal_conversion(model_files):
    for model_file in model_files:
        model_state = torch.load(model_file, map_location=torch.device('cpu'), weights_only=False)
        autoep_metadata = model_state.get(AUTOEP_LAYERS_KEY)
        if autoep_metadata is None:
            autoep_metadata = model_state.get(AUTOEP_LAYERS_KEY_LEGACY)

        if autoep_metadata is not None:
            raise NotImplementedError("Stage 3 universal checkpoint conversion with AutoEP is not supported. "
                                      "Use regular same-topology ZeRO-3 checkpoint load for AutoEP checkpoints.")


def _get_checkpoint_files(checkpoint_dir, glob_pattern):
    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)

    if len(ckpt_files) == 0:
        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")

    return ckpt_files


def _get_zero_stage(optim_files):
    state_dict = torch.load(optim_files[0], map_location=torch.device('cpu'), weights_only=False)
    optimizer_state = state_dict[OPTIMIZER_STATE_DICT]
    zero_stage = optimizer_state.get(ZERO_STAGE, 1)
    return zero_stage


def _inject_missing_state(ds_checkpoint):
    if UNIVERSAL_CHECKPOINT_INFO not in ds_checkpoint.global_state:
        sd = torch.load(ds_checkpoint.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False)
        if UNIVERSAL_CHECKPOINT_INFO not in sd:
            ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO] = {}
            ds_checkpoint.global_state[UNIVERSAL_CHECKPOINT_INFO][
                UNIVERSAL_CHECKPOINT_VERSION_KEY] = UNIVERSAL_CHECKPOINT_VERSION_VALUE


def _check_for_required_state(ds_checkpoint):
    universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO)
    assert universal_checkpoint_info is not None, f'Required {UNIVERSAL_CHECKPOINT_INFO} state is missing in checkpoint. Verify that client creates this state.'


def _classify_autoep_expert_file_consolidation(autoep_metadata, expert_files):
    if autoep_metadata is not None:
        return 'autoep'
    if expert_files:
        return 'native_moe'
    return 'none'


def main(args):
    print('Convert DeepSpeed Checkpoint to Universal Checkpoint')

    print(f'Converting DeepSpeed checkpoint in {args.input_folder} to Universal checkpoint in {args.output_folder}')

    optim_files = _get_optim_files(args.input_folder)
    zero3_optim_files = _filter_zero3_optim_files(optim_files)
    zero_stage = _get_zero_stage(zero3_optim_files or optim_files)
    if zero_stage > 2 and zero3_optim_files:
        optim_files = zero3_optim_files

    if zero_stage <= 2:
        ds_checkpoint = DeepSpeedCheckpoint(args.input_folder)
        if args.inject_missing_state:
            _inject_missing_state(ds_checkpoint)
        else:
            _check_for_required_state(ds_checkpoint)

        iteration = ds_checkpoint.get_iteration()
        #_create_latest_file(args.output_folder, iteration)
        checkpoint_paths = _create_checkpoint_paths(args.output_folder, iteration, ds_checkpoint.tp_degree,
                                                    ds_checkpoint.pp_degree)

        slice_shapes = []
        for mp_rank_file in ds_checkpoint.mp_rank_files:
            mp_sd = torch.load(mp_rank_file, map_location=torch.device('cpu'), weights_only=False)
            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('*** 2.5. Consolidating AutoEP expert files')
        from deepspeed.checkpoint.autoep_universal import (
            consolidate_autoep_expert_files,
            consolidate_autoep_optimizer_states,
        )

        # Load AutoEP metadata from main checkpoint
        main_sd = torch.load(ds_checkpoint.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False)
        autoep_metadata = main_sd.get(AUTOEP_LAYERS_KEY)
        if autoep_metadata is None:
            autoep_metadata = main_sd.get(AUTOEP_LAYERS_KEY_LEGACY)

        # Check for expert files in checkpoint directory
        expert_files = glob.glob(os.path.join(args.input_folder, 'layer_*_expert_*_model_states.pt'))
        autoep_expert_file_type = _classify_autoep_expert_file_consolidation(autoep_metadata, expert_files)

        if autoep_expert_file_type == 'autoep':
            consolidate_autoep_expert_files(args.input_folder, args.output_folder, autoep_metadata)
            ep_size = autoep_metadata[0]['ep_size'] if autoep_metadata else 1
            consolidate_autoep_optimizer_states(args.input_folder, args.output_folder, autoep_metadata, ep_size)
            print(f'    Consolidated {len(autoep_metadata)} AutoEP layer(s)')
        elif autoep_expert_file_type == 'native_moe':
            print(f'    Found {len(expert_files)} expert checkpoint file(s) but no AutoEP metadata; '
                  'assuming native DeepSpeed MoE and skipping AutoEP consolidation')
        else:
            print('    No AutoEP layers found, skipping')

        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, injecting AutoEP metadata into UNIVERSAL_CHECKPOINT_INFO
        for f in glob.glob(os.path.join(args.input_folder, 'mp*')):
            if autoep_metadata is not None:
                # Load -> update with AutoEP metadata -> save
                mp_sd = torch.load(f, map_location=torch.device('cpu'), weights_only=False)
                if UNIVERSAL_CHECKPOINT_INFO not in mp_sd:
                    mp_sd[UNIVERSAL_CHECKPOINT_INFO] = {}
                mp_sd[UNIVERSAL_CHECKPOINT_INFO][EXPERT_PARAMETER_PATTERNS] = [r'\.experts\.w[123]$']
                mp_sd[UNIVERSAL_CHECKPOINT_INFO][AUTOEP_LAYERS_KEY] = autoep_metadata
                out_path = os.path.join(args.output_folder, os.path.basename(f))
                torch.save(mp_sd, out_path)
            else:
                shutil.copy2(f, args.output_folder)

    else:
        # Stage 3 path
        model_files = _get_zero3_model_state_files(args.input_folder)
        autoep_expert_param_names, autoep_metadata_by_rank = _autoep_expert_param_names_by_rank(model_files)
        has_autoep_metadata = any(metadata is not None for metadata in autoep_metadata_by_rank.values())
        has_zero3_partitioned_autoep = any(
            _uses_zero3_partitioned_autoep_metadata(metadata) for metadata in autoep_metadata_by_rank.values())
        if has_autoep_metadata and not has_zero3_partitioned_autoep:
            raise NotImplementedError("Stage 3 universal checkpoint conversion for AutoEP requires the "
                                      "partition-native AutoEP ZeRO-3 checkpoint format.")
        if not has_zero3_partitioned_autoep:
            autoep_expert_param_names = set()
        else:
            _validate_zero3_model_optim_rank_sets(model_files, optim_files)
        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,
                                         exclude_param_names=autoep_expert_param_names)

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

        if has_zero3_partitioned_autoep:
            print('*** 2.5. Consolidating AutoEP ZeRO-3 expert states')
            _consolidate_zero3_autoep_expert_states(args.output_folder, model_files, optim_files)

        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, filtering out expert files
        for f in glob.glob(os.path.join(args.input_folder, '*model_states.pt')):
            if _is_expert_model_state_file(f):
                continue
            if has_zero3_partitioned_autoep:
                model_state = torch.load(f, map_location=torch.device('cpu'), weights_only=False)
                autoep_metadata = _get_autoep_metadata(model_state)
                if UNIVERSAL_CHECKPOINT_INFO not in model_state:
                    model_state[UNIVERSAL_CHECKPOINT_INFO] = {}
                model_state[UNIVERSAL_CHECKPOINT_INFO][EXPERT_PARAMETER_PATTERNS] = [r'.*\.experts\.w[123]$']
                model_state[UNIVERSAL_CHECKPOINT_INFO][AUTOEP_LAYERS_KEY] = autoep_metadata
                torch.save(model_state, os.path.join(args.output_folder, os.path.basename(f)))
            else:
                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)

    print('*** Done!')


if __name__ == "__main__":
    args = parse_arguments()
    main(args)