[BUG] Is mcore-0.12.0 checkpoint resume correct?

Open

#1,570 opened on May 9, 2025 10 comments 0 reactions 1 assignee Assigned to @yanring View on GitHub

bug community-request waiting-on-customer

Dominant language: Python
Stars: 17.9k
Forks: 4.5k
Avg merge: 4d 3h
Merged PRs (30d): 272

Description

**Describe the bug**
While using chckpoint recume in mcore-0.12, it cannot produce same loss value in the first forward, and the subsequent loss deviation is relatively large. But mcore-0.11 can generate the same loss in the first step, and the subsequent deviation is very small.

**To Reproduce**
test scripts:
```sh
#!/bin/bash
TOKENIZER_MODEL="/workspace/Models/DeepSeek-V3"

export CUDA_DEVICE_MAX_CONNECTIONS=1
MASTER_ADDR=${MASTER_ADDR:-"localhost"}
MASTER_PORT=${MASTER_PORT:-"6000"}
NNODES=${NNODES:-"1"}
NODE_RANK=${RANK:-"0"}
GPUS_PER_NODE=${GPUS_PER_NODE:-"8"}
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

TP=${TP:-"2"}
PP=${PP:-"2"}
EP=${EP:-"2"}

## train params
SEQ_LEN=${SEQ_LEN:-"4096"}
GB=${GB:-"32"} # global batch size
MB=${MB:-"1"} # micro batch size
TRAIN_ITERS=${TRAIN_ITERS:-"10"}
LR_DECAY_ITERS=${LR_DECAY_ITERS:-"7"}
LR_WARMUP_ITERS=${LR_WARMUP_ITERS:-"3"}

EXTRA_ARGS=${EXTRA_ARGS:-""}

# deterministic-mode, cannot use flash-attn, H800 Nan Bug
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export NCCL_ALGO=Ring
EXTRA_ARGS="$EXTRA_ARGS --deterministic-mode"

## model params
NUM_LAYERS=8
NUM_DENS_LAYERS=1
HIDDEN_SIZE=2048
NUM_ATTN_HEADS=16
INTERMEDIATE_SIZE=10944

MLA_ARGS=(
# --use-precision-aware-optimizer
# --exp-avg-dtype fp16
# --exp-avg-sq-dtype fp16

--no-rope-fusion
--qk-layernorm
# --q-lora-rank 1536
--kv-lora-rank 512
--multi-latent-attention
--qk-head-dim 128
--qk-pos-emb-head-dim 64
--v-head-dim 128
--rotary-scaling-factor 40
)

MOE_ARGS=(
--num-experts 64
--moe-router-topk 6
--moe-ffn-hidden-size 1408
--moe-shared-expert-intermediate-size 2816 # n_shared_experts * ffn_size, 2*1408
--moe-layer-freq "([0]*$NUM_DENS_LAYERS+[1]*($NUM_LAYERS-$NUM_DENS_LAYERS))"

--expert-model-parallel-size $EP
--moe-token-dispatcher-type alltoall
--moe-grouped-gemm
--moe-router-enable-expert-bias
--moe-router-score-function "sigmoid"
# --moe-router-load-balancing-type seq_aux_loss
--moe-router-load-balancing-type aux_loss
--moe-aux-loss-coeff 0.0001
--moe-router-pre-softmax
--moe-router-topk-scaling-factor 6
--moe-per-layer-logging

--moe-router-group-topk 3
--moe-router-num-groups 8

#--attention-backend unfused
# --moe-permute-fusion # must install TE>=V2.1
)

FP8_ARGS=(
# --fp8-format hybrid
# --fp8-margin 0
# --fp8-interval 1
# --fp8-amax-history-len 1024
# --fp8-amax-compute-algo max
# --fp8-param-gather
)

DISTRIBUTED_ARGS=(
--nproc_per_node $GPUS_PER_NODE
--nnodes $NNODES
--node_rank $NODE_RANK
--master_addr $MASTER_ADDR
--master_port $MASTER_PORT
)

MODEL_ARGS=(
--seq-length ${SEQ_LEN}
--num-layers ${NUM_LAYERS}
--hidden-size ${HIDDEN_SIZE}
--ffn-hidden-size ${INTERMEDIATE_SIZE}
--num-attention-heads ${NUM_ATTN_HEADS}
--max-position-embeddings ${SEQ_LEN}

# --add-qkv-bias # qwen2 use qkv bias
--disable-bias-linear
--init-method-std 0.01
--attention-dropout 0.0
--hidden-dropout 0.0
--normalization RMSNorm
--position-embedding-type rope
--swiglu
# FIXME: masked_softmax_fusion?
--no-masked-softmax-fusion
--no-position-embedding
## qwen2 specific
--rotary-base 1000000
--norm-epsilon 1e-6
--untie-embeddings-and-output-weights
)

DATA_ARGS=(
--tokenizer-type HuggingFaceTokenizer
--tokenizer-model ${TOKENIZER_MODEL}
--split 97,2,1
)

if [ -n "${DATA_PATH}" ]; then
DATA_ARGS+=(
--data-path $DATA_PATH
)
else
DATA_ARGS+=(
--mock-data
)
fi

TRAINING_ARGS=(
--micro-batch-size ${MB}
--global-batch-size ${GB}
--lr 1e-4
--train-iters ${TRAIN_ITERS}
--lr-decay-iters ${LR_DECAY_ITERS}
--lr-warmup-iters ${LR_WARMUP_ITERS}
--lr-decay-style cosine
--min-lr 1.0e-5
--weight-decay 0.1
--clip-grad 1.0
--bf16
)

MODEL_PARALLEL_ARGS=(
--tensor-model-parallel-size ${TP}
--pipeline-model-parallel-size ${PP}
--use-distributed-optimizer
# --use-flash-attn
--use-mcore-models
# --overlap-grad-reduce
# --overlap-param-gather
--sequence-parallel
--no-async-tensor-model-parallel-allreduce
# --tp-comm-overlap-rs-dgrad
)

LOGGING_ARGS=(
# --load ./ckpts/1f1b_dsv2_lite_ep2_pp2_dp2 \
--async-save \
--save-interval 5 \
--save ./ckpts/1f1b_dsv2_lite_ep2_pp2_dp2 \
# --no-load-optim \
# --no-load-rng
--eval-interval 1000 \
--eval-iters 0 \
--log-interval 1 \
--log-throughput \
--log-timers-to-tensorboard \
--log-validation-ppl-to-tensorboard \
--log-world-size-to-tensorboard \
--tensorboard-dir "Tensorboards/dual_pipeV_dsv2_lite_ep2_pp2_dp2" \
)

if [ -n "${WANDB_API_KEY}" ]; then
LOGGING_ARGS+=(
--wandb-project ${WANDB_PROJECT:-"test-pretrain"}
--wandb-exp-name ${WANDB_NAME:-"test"}_${MODEL_SIZE}
)
fi

echo "
torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
${MODEL_ARGS[@]} \
${DATA_ARGS[@]} \
${TRAINING_ARGS[@]} \
${MODEL_PARALLEL_ARGS[@]} \
${LOGGING_ARGS[@]} \
${MLA_ARGS[@]} \
${MOE_ARGS[@]} \
${FP8_ARGS[@]} \
${EXTRA_ARGS}
"

NCCL_IB_GID_INDEX=3 NCCL_NET_GDR_LEVEL=SYS torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
${MODEL_ARGS[@]} \
${DATA_ARGS[@]} \
${TRAINING_ARGS[@]} \
${MODEL_PARALLEL_ARGS[@]} \
${LOGGING_ARGS[@]} \
${MLA_ARGS[@]} \
${MOE_ARGS[@]} \
${FP8_ARGS[@]} \
${EXTRA_ARGS}

```

**Expected behavior**
Checkpoint resume can produce same loss.

**Environment (please complete the following information):**
- Megatron-LM tag core-0.12 and core-0.11
- TransformerEngine-v1.13
- Image: nvcr.io/nvidia/pytorch:24.07-py3

**Proposed fix**

**Additional context**

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

[BUG] Is mcore-0.12.0 checkpoint resume correct?

Description

Contributor guide

Assessment

Get new issues in your inbox