Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
combine_lm_vision_checkpoints.sh 3.29 KB
Newer Older
xingjinliang's avatar
升级megatron_v0.10  
xingjinliang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57 
#/bin/bash
MCORE_LM=$1    # <path_to_mcore_lm_model_folder>
MCORE_VISION=$2   # <path_to_mcore_vision_model_folder>
OUTPUT_DIR=$3   # <path_to_output_folder_for_combined_checkpoint>
MODEL_TYPE=$4   # Model type. Default: Mistral CLIP example.

if [[ $MODEL_TYPE == "nvlm" ]]; then
    # NVLM TP=8
    python examples/multimodal/combine_state_dicts.py \
        --input \
        ${MCORE_LM}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_03/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_03/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_04/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_04/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_05/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_05/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_06/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_06/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_07/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_07/model_optim_rng.pt \
        --prefixes language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model \
        --output \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_03/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_04/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_05/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_06/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_07/model_optim_rng.pt
else
    # Mistral CLIP example TP=4.
    python examples/multimodal/combine_state_dicts.py \
        --input \
        ${MCORE_LM}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${MCORE_LM}/iter_0000001/mp_rank_03/model_optim_rng.pt \
        ${MCORE_VISION}/iter_0000001/mp_rank_03/model_optim_rng.pt \
        --prefixes language_model vision_model language_model vision_model language_model vision_model language_model vision_model \
        --output \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_00/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_01/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_02/model_optim_rng.pt \
        ${OUTPUT_DIR}/iter_0000001/mp_rank_03/model_optim_rng.pt
fi

echo 1 > ${OUTPUT_DIR}/latest_checkpointed_iteration.txt