Function normalize_model_name

verl/verl/utils/model.py:209–237

Transform the model name in each model_chunk in each pp stage into the name in inference engine

(name, pp_rank, vpp_rank, pp_size, vpp_size, num_layers)

    """

    def normalize_model_name(name, pp_rank, vpp_rank, pp_size, vpp_size, num_layers):
        """
        Transform the model name in each model_chunk in each pp stage into the name in inference engine
        """
        if vpp_size > 1:
            # print(f'try to bind vpp params to inference engine...')
            layers_per_pp = num_layers // pp_size
            layers_per_vpp = layers_per_pp // vpp_size
            pp_offset = layers_per_vpp * pp_rank
            vpp_offset = (layers_per_vpp * pp_size) * vpp_rank
            layer_offset = pp_offset + vpp_offset
        else:
            layers_per_pp = num_layers // pp_size
            layer_offset = layers_per_pp * pp_rank

        if layer_name in name:  # belong to an intermediate layer
            split_name = name.split('.')
            # find the num next to split_name
            for i, name in enumerate(split_name):
                if name == layer_name:
                    break
            layer_num_idx = i + 1
            # check the name
            assert len(split_name) >= layer_num_idx + 1, f'split_name = {split_name}'
            assert split_name[layer_num_idx].isdigit(), f'split_name = {split_name}'
            # increment layer_num_idx by layer_offset
            split_name[layer_num_idx] = str(int(split_name[layer_num_idx]) + layer_offset)
            name = '.'.join(split_name)  # weight name in inference_tp_model
        return name

    pp_size = len(params)
    normalized_name_to_param = {}