[QUESTION] Checkpoint storage format

[syx11237744]

Your question
Ask a clear and concise question about Megatron-LM.
Could you let me know which version I should revert to if I want to use the previous checkpoint storage format, which is stored as .pt? Or are there any other methods to save it as a .pt file?Thank you!

DATASET_PATH=/share/root/out_file/sum.jsonl
SAVE_PATH=/share/sunyuanxu/out_file/sum
VOCAB_FILE=gpt2_/vocab.json
MERGE_FILE=gpt2_/merges.txt

export CUDA_DEVICE_MAX_CONNECTIONS=1

GPUS_PER_NODE=1

Change for multinode config

MASTER_ADDR=localhost
MASTER_PORT=6000
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

CHECKPOINT_PATH=/share/root/checkpoint/cp
DATA_PATH=/share/root/out_file/sum_text_document
DISTRIBUTED_ARGS="
--nproc_per_node $GPUS_PER_NODE
--nnodes $NNODES
--node_rank $NODE_RANK
--master_addr $MASTER_ADDR
--master_port $MASTER_PORT
"

GPT_ARGS="
--num-layers 24
--hidden-size 1024
--num-attention-heads 16
--seq-length 1024
--max-position-embeddings 1024
--micro-batch-size 32
--global-batch-size 256
--lr 0.00015
--train-iters 1000
--lr-decay-iters 320000
--lr-decay-style cosine
--min-lr 1.0e-5
--weight-decay 1e-2
--lr-warmup-fraction .01
--clip-grad 1.0
--fp16
--attention-softmax-in-fp32
"

DATA_ARGS="
--data-path $DATA_PATH
--vocab-file $VOCAB_FILE
--merge-file $MERGE_FILE
--split 949,50,1
"

OUTPUT_ARGS="
--log-interval 100
--save-interval 10000
--eval-interval 1000
--eval-iters 10
"

torchrun $DISTRIBUTED_ARGS Megatron-LM/pretrain_gpt.py
$GPT_ARGS
$DATA_ARGS
$OUTPUT_ARGS
--distributed-backend nccl
--save $CHECKPOINT_PATH
--load $CHECKPOINT_PATH

[Advait251206]

Short answer: you don’t need to downgrade Megatron-LM to get .pt checkpoints — but in current versions the default checkpoint format has changed to a sharded / distributed format, and .pt is no longer the primary format.

However, you can still force or recover .pt-style checkpoints depending on how you save them.

---

1. Why .pt format disappeared

Older Megatron-LM versions saved checkpoints as:

model_optim_rng.pt

This contained:

• model weights
• optimizer state
• RNG state

Newer versions changed this because:

• models became too large for single-file checkpoints
• tensor parallelism requires sharding
• ZeRO-like optimizations split optimizer state
• multi-node training needs distributed saves

So now Megatron uses:

iter_XXXX/
  mp_rank_00/
  mp_rank_01/
  ...
  optimizer_states/
  rng_states/

👉 This is the sharded checkpoint format

---

2. Can you still get .pt checkpoints?

✔ Yes, but only in limited scenarios

You can get .pt style checkpoints if:

Option A — disable tensor parallelism

Run with:

• --tensor-model-parallel-size 1
• --pipeline-model-parallel-size 1

Then Megatron can still produce legacy-style full checkpoints.

---

Option B — use older Megatron-LM commit

If you really need legacy .pt format, you must revert to:

Megatron-LM version before the checkpoint refactor (~2022–early 2023 era)

But this is NOT recommended because:

• lacks modern optimizations
• incompatible with newer configs
• not compatible with modern NeMo / TRT-LLM pipelines

---

Option C — manually consolidate checkpoints

Modern approach:

tools/merge_checkpoints.py

or:

tools/checkpoint_util.py --load ... --save-converted ...

You can convert:

sharded checkpoint → single .pt file

---

3. Recommended modern workflow (best practice)

Instead of .pt, use:

✔ Distributed checkpoint format

iter_XXXX/

Then use:

• load_checkpoint() APIs
• Megatron checkpoint loader
• NeMo converter (if needed)
• TensorRT-LLM conversion scripts

---

4. Why .pt is discouraged now

Because .pt:

• does not scale beyond single-GPU assumptions
• breaks with tensor/pipeline parallelism
• is memory inefficient for large models
• cannot represent optimizer sharding cleanly

---

5. If your goal is portability (important)

If you want checkpoints for:

✔ inference (TensorRT-LLM)

Use:

• HF export OR
• Megatron → HF conversion scripts

✔ training resume

Use:

• native Megatron sharded checkpoints (recommended)

---

6. Direct answer to your question

❓ “Which version should I revert to for .pt format?”

👉 Roughly:

• Megatron-LM pre tensor-parallel checkpoint refactor (≈ 2022 era)

BUT:

⚠ not recommended for modern training

---

❓ “Can I still save as .pt in current version?”

👉 Only if:

• TP = 1
• PP = 1
• and legacy save utilities are used

Otherwise:

❌ not supported by default

---

7. Bottom line

• ❌ .pt is legacy format and no longer default
• ✔ modern Megatron uses sharded checkpoints
• ✔ you can convert back to .pt if needed
• ⚠ downgrading is possible but strongly discouraged
• ✔ best practice is to adapt to sharded format

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206