def map_key(key, name_pt, num_layers):
    '''Map state dict key from checkpoint to the current model'''
    num_pt = int(name_pt[6:8])
    if num_pt == 0:
        return f'language_model.embedding.{key}'
    if num_pt == 1:
        return f'language_model.projector.{key}'
    num_state = num_pt - 3
    if 0 <= num_state < num_layers:
        return f'language_model.transformer.layers.{num_state}.{key}'
    if num_pt == num_layers + 4:
        return f'language_model.output_layer.{key}'
Map state dict key from checkpoint to the current model
megatron_lm/megatron/checkpointing.py