NVIDIA / NVIDIA/Megatron-LM

Enable optimizer offload → Error occurs when loading the saved checkpoint.

Open
#1,842 3 comments 0 reactions 0 assignees View on GitHub
bug community-request module: moe
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

I enabled --use-precision-aware-optimizer and --optimizer-cpu-offload for training. After training 10 steps on a very small dataset, the loss was 0.31. However, when resuming training from this checkpoint, the loss jumped to 10. If I enable the --no-load-optim option, training continues normally.Here are the hyperparameters I used for training. Could you please let me know if this is a bug or if there is an issue with the parameters I used? Thank you in advance for your reply.
`CHECKPOINT_PATH=./ceshi
VOCAB_FILE=gpt2-vocab.json
MERGE_FILE=gpt2-merges.txt
DATA_PATH=/mnt/beegfs/luojiangang/Megatron-LM/data/my-gpt2_text_document

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

MODEL_ARGS=(
--use-mcore-models
--disable-bias-linear
--seq-length 2048
--max-position-embeddings 32768
--num-layers 1
--hidden-size 4096
--ffn-hidden-size 16384
--num-attention-heads 16
--init-method-std 0.01
--attention-dropout 0.0
--hidden-dropout 0.0
--normalization RMSNorm
--position-embedding-type rope
--swiglu
--untie-embeddings-and-output-weights
--no-masked-softmax-fusion
--no-position-embedding
--rotary-base 1000000
--ckpt-format torch
)

MOE_ARGS=(
--moe-grouped-gemm
--moe-token-dispatcher-type alltoall
--moe-router-load-balancing-type seq_aux_loss
--moe-aux-loss-coeff 1e-2
--overlap-param-gather
--overlap-grad-reduce
--num-experts 8
--moe-router-topk 2
)

DATA_ARGS=(
--data-path $DATA_PATH
--data-cache-path /mnt/beegfs/luojiangang/Megatron-LM/datacache
--vocab-file $VOCAB_FILE
--merge-file $MERGE_FILE
--split 10,0,0
)

TRAINING_ARGS=(
--micro-batch-size 1
--global-batch-size $GLOBAL_SIZE
--lr 1e-4
--train-iters 100
--lr-decay-iters 32
--lr-decay-style cosine
--min-lr 1.0e-5
--weight-decay 0.1
--lr-warmup-iters 5
--clip-grad 1.0
--bf16
--recompute-method uniform
--recompute-granularity full
--recompute-num-layers 1
--use-precision-aware-optimizer
--optimizer-cpu-offload
--no-load-optim
)

MODEL_PARALLEL_ARGS=(
--tensor-model-parallel-size 1
--pipeline-model-parallel-size 1
--expert-model-parallel-size 1
--use-distributed-optimizer
--sequence-parallel
)

LOGGING_ARGS=(
--log-interval 1
--save-interval 5
--eval-interval 1000
--eval-iters 10
--save $CHECKPOINT_PATH
--load $CHECKPOINT_PATH
--tensorboard-dir "${CHECKPOINT_PATH}/tensorboard"
--no-load-rng
)
`

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.