Slow checkpoint saving time (6 mins to save an 8B model checkpoint in sync mode)

[vwxyzjn]

It takes ~6 minutes to save a checkpoint using non async mode. Is this expected?

Sync mode

[rank0]:[titan] 2025-06-15 21:31:48,968 - root - INFO - TensorBoard logging enabled. Logs will be saved at ./outputs/tb/20250615-2131 
[rank0]:[titan] 2025-06-15 21:31:48,969 - root - INFO - CUDA capacity: NVIDIA H100 80GB HBM3 with 79.10GiB memory                     
[rank0]:[titan] 2025-06-15 21:31:49,083 - root - INFO - Model llama3 8B size: 8,030,261,248 total parameters                          
[rank0]:[titan] 2025-06-15 21:31:49,084 - root - INFO - Applied full activation checkpointing to the model                            
[rank0]:[titan] 2025-06-15 21:31:49,164 - root - INFO - Applied FSDP to the model                                                     
[rank0]:[titan] 2025-06-15 21:31:49,505 - root - INFO - Peak FLOPS used for computing MFU: 9.890e+14                                  
[rank0]:[titan] 2025-06-15 21:31:49,505 - root - INFO - CUDA memory usage for model: 3.95GiB(4.99%)                                   
[rank0]:[titan] 2025-06-15 21:31:49,535 - root - INFO - Checkpointing active. Checkpoints will be loaded from and saved to ./outputs/c
heckpoint                                                                                                                             
[rank0]:[titan] 2025-06-15 21:31:49,535 - root - INFO - Trainer is initialized with local batch size 1, global batch size 64, gradient
 accumulation steps 8, sequence length 8192, total steps 1000 (warmup 40).                                                            
[rank0]:[titan] 2025-06-15 21:31:49,535 - root - INFO - Loading the checkpoint from assets/models/dcp/llama3.1-8B.                    
[rank0]:[titan] 2025-06-15 21:32:02,935 - root - INFO - [GC] GC collection for checkpoint loading. 0.01 seconds.                      
[rank0]:[titan] 2025-06-15 21:32:02,935 - root - INFO - Finished loading the checkpoint in 13.40 seconds.                             
[rank0]:[titan] 2025-06-15 21:32:02,935 - root - INFO - Training starts at step 1.                                                    
[rank0]:[titan] 2025-06-15 21:32:15,816 - root - INFO - step:  1  loss:  2.4292  memory: 29.18GiB(36.90%)  tps: 2,452  tflops: 141.98 
 mfu: 14.36%                                                                                                                          
[rank0]:[titan] 2025-06-15 21:32:15,816 - root - INFO - Saving the checkpoint (or staging if async is enabled).                       
[rank0]:[titan] 2025-06-15 21:38:31,430 - root - INFO - [GC] GC collection invoked by checkpointer. 0.04 seconds.                     
[rank0]:[titan] 2025-06-15 21:38:31,431 - root - INFO - Finished saving the checkpoint (or staging if async is enabled)in 375.61 secon
ds.                                                                                                                                   
[rank0]:[titan] 2025-06-15 21:38:31,431 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40          
[rank0]:[titan] 2025-06-15 21:40:09,439 - root - INFO - step: 10  loss:  2.3602  memory: 36.65GiB(46.33%)  tps: 1,245  tflops: 72.12  
mfu: 7.29%  

Async mode:

rank0]:[titan] 2025-06-15 21:44:35,889 - root - INFO - step:  1  loss:  2.4292  memory: 29.18GiB(36.90%)  tps: 2,327  tflops: 134.74  mfu: 13.62%
[rank0]:[titan] 2025-06-15 21:44:35,890 - root - INFO - Saving the checkpoint (or staging if async is enabled).
[rank0]:[titan] 2025-06-15 21:44:35,898 - root - INFO - [GC] GC collection invoked by checkpointer. 0.01 seconds.
[rank0]:[titan] 2025-06-15 21:44:47,661 - root - INFO - [GC] GC collection invoked by checkpointer. 0.00 seconds.
[rank0]:[titan] 2025-06-15 21:44:47,672 - root - INFO - Finished saving the checkpoint (or staging if async is enabled)in 11.78 seconds.
[rank0]:[titan] 2025-06-15 21:44:47,672 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40
[rank0]:/home/ubuntu/code/thirdparty/torchtitan/.venv/lib/python3.13/site-packages/torch/distributed/checkpoint/filesystem.py:111: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
[rank0]:  if tensor.storage().size() != tensor.numel():
[rank0]:[titan] 2025-06-15 21:46:26,319 - root - INFO - step: 10  loss:  2.3601  memory: 36.64GiB(46.33%)  tps: 5,341  tflops: 309.34  mfu: 31.28%

Reproduction: check out #1300 and run

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" uv run ./run_train.sh \
  --model.tokenizer_path assets/tokenizer/Meta-Llama-3.1-8B-tokenizer.model \
  --training.max_seq_len 131072 \
  --checkpoint.initial_load_path "assets/models/dcp/llama3.1-8B" \
  --profiling.no_enable_profiling \
  --checkpoint.enable_checkpoint \
  --checkpoint.async_mode async \
  --activation_checkpoint.mode full \
  --training.global_batch_size 64 \
  --lr_scheduler.warmup_steps 40 \
  --optimizer.lr 1e-5

[tianyu-l]

8B parameters translates into 8B x 3 (model parameters + optimizer states 2x) x 4 (float32 dtype) = 96 GB data, plus the states for dataloader and lr scheduler.

Assuming the (PCIe) bandwidth from H100 GPU to CPU is ~50GB/s, it should take no more than a few seconds?

Have you checked how large the checkpoint is?

cc @fegin

[vwxyzjn]

Yeah the checkpoint is about 90G.

ubuntu@costa-tt-lk7k7:~/code/thirdparty/torchtitan/outputs/checkpoint/step-1000$ ls
__0_0.distcp  __1_0.distcp  __2_0.distcp  __3_0.distcp  __4_0.distcp  __5_0.distcp  __6_0.distcp  __7_0.distcp
ubuntu@costa-tt-lk7k7:~/code/thirdparty/torchtitan/outputs/checkpoint/step-1000$ du -sh .
90G     .

Could you see if you all can reproduce it? Below is my dependencies.
https://gist.github.com/vwxyzjn/ad38da9253496f39efa999efc7c7c8cf

[fegin]

Staging should be pretty fast, even with a very bad PCIe (v.3), it won't take longer than 30 seconds for staging a 90G checkpoint.

@Saiteja64, for async checkpoint, TorchTitan just calls DCP.async_save().

cc., @wwwjn for ramping up checkpointing related issues.

[vwxyzjn]

FYI claude tells me to just test cp

ubuntu@costa-tt-lk7k7:~/code/outputs/checkpoint/step-1$ du -sh .
80G     .
ubuntu@costa-tt-lk7k7:~/code/outputs/checkpoint/step-1$ cd ..
ubuntu@costa-tt-lk7k7:~/code/outputs/checkpoint$ time cp step-1/ step-2
cp: -r not specified; omitting directory 'step-1/'

real    0m0.002s
user    0m0.000s
sys     0m0.002s
ubuntu@costa-tt-lk7k7:~/code/outputs/checkpoint$ time cp -r step-1/ step-2

real    0m31.127s
user    0m0.000s
sys     0m0.007s

It also feels like machine-dependent perhaps? In a new run (potentially a new machine) with async checkpoint I got:

[rank0]:[titan] 2025-06-16 21:35:58,929 - root - INFO - step:  1  loss:  1.4989  memory: 29.16GiB(36.87%)  tps: 2,230  tflops: 129.16  mfu: 13.06%
[rank0]:[titan] 2025-06-16 21:35:58,930 - root - INFO - Saving the checkpoint (or staging if async is enabled).
[rank0]:[titan] 2025-06-16 21:35:58,946 - root - INFO - [GC] GC collection invoked by checkpointer. 0.02 seconds.
[rank0]:[titan] 2025-06-16 21:36:05,206 - root - INFO - [GC] GC collection invoked by checkpointer. 0.00 seconds.
[rank0]:[titan] 2025-06-16 21:36:05,220 - root - INFO - Finished saving the checkpoint (or staging if async is enabled)in 6.29 seconds.
[rank0]:[titan] 2025-06-16 21:36:05,220 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40
[rank0]:/home/ubuntu/code/.venv/lib/python3.12/site-packages/torch/distributed/checkpoint/filesystem.py:116: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
[rank0]:  if tensor.storage().size() != tensor.numel():
[rank0]:[titan] 2025-06-16 21:37:45,161 - root - INFO - step: 10  loss:  1.4308  memory: 36.64GiB(46.33%)  tps: 5,552  tflops: 321.56  mfu: 32.51%

[fegin]

@vwxyzjn You can try this command:

-> % dd if=/dev/zero of=/home/fegin/local/test_disk bs=1G count=10
10+0 records in
10+0 records out
10737418240 bytes (11 GB, 10 GiB) copied, 1.85036 s, 5.8 GB/s

Also, I tried the same configuration and the staging time is ~5 seconds. So I guess I won't be able to reproduce on my local machine. What's the GPU and PCIe of the "old machine"?

[vwxyzjn]

I got this:

ubuntu@costa-ww-wrnk9:~/code/outputs$ touch test.txt
ubuntu@costa-ww-wrnk9:~/code/outputs$ dd if=/dev/zero of=test.txt bs=1G count=10
10+0 records in
10+0 records out
10737418240 bytes (11 GB, 10 GiB) copied, 4.86356 s, 2.2 GB/s
ubuntu@costa-ww-wrnk9:~/code/outputs$ 

I also tried using nvme drive:

ubuntu@costa-ww-wrnk9:/mnt/nvme$ sudo touch test.txt
ubuntu@costa-ww-wrnk9:/mnt/nvme$ sudo dd if=/dev/zero of=test.txt bs=1G count=10
10+0 records in
10+0 records out
10737418240 bytes (11 GB, 10 GiB) copied, 4.78167 s, 2.2 GB/s
ubuntu@costa-ww-wrnk9:/mnt/nvme$ 

The slow checkpoint time without async model can reproduce it again.

[rank0]:  return collate([torch.as_tensor(b) for b in batch], collate_fn_map=collate_fn_map)
[rank0]:[titan] 2025-06-16 23:40:42,708 - root - INFO - step:  1  loss:  1.4989  memory: 29.16GiB(36.87%)  tps: 2,456  tflops: 142.22  mfu: 14.38%
[rank0]:[titan] 2025-06-16 23:40:42,709 - root - INFO - Saving the checkpoint (or staging if async is enabled).
[rank0]:[titan] 2025-06-16 23:47:16,568 - root - INFO - [GC] GC collection invoked by checkpointer. 0.01 seconds.
[rank0]:[titan] 2025-06-16 23:47:16,568 - root - INFO - Finished saving the checkpoint (or staging if async is enabled)in 393.86 seconds.
[rank0]:[titan] 2025-06-16 23:47:16,568 - root - INFO - Synchronizing and adjusting timeout for all ProcessGroups to 0:01:40
[rank0]:[titan] 2025-06-16 23:48:54,634 - root - INFO - step: 10  loss:  1.4306  memory: 36.65GiB(46.33%)  tps: 1,199  tflops: 69.44  mfu: 7.02%
[rank0]:[titan] 2025-06-16 23:50:43,643 - root - INFO - step: 20  loss:  1.4760  memory: 36.65GiB(46.33%)  tps: 6,012  tflops: 348.18  mfu: 35.21%

It's AWS's 8xNVIDIA H100 80GB HBM3 box. I believe it's this one: ml.p5.48xlarge, which has the following configurations according to here.

Intel Sapphire Rapids CPU and PCIe Gen5 between the CPU and GPU in P5en instances; 3rd Gen AMD EPYC processors (AMD EPYC 7R13) and PCIe Gen4 between the CPU and GPU in P5 and P5e instances.

[vwxyzjn]

So an update on this: it turns out we just needed to save the model on an nvme SSD mounted at /mnt/nvme instead of the default elastic block store.

Once we do that, the checkpoint speed is much much faster (400 seconds -> 30 seconds)

There is a massive difference in speed:

## 500GB image drive:

dd if=/dev/zero \
    of="/opt/sagemaker/test_10G_on_sagemaker.img" \
    bs=1G count=10 oflag=direct status=progress
10737418240 bytes (11 GB) copied, 40.276593 s, 267 MB/s
10+0 records in
10+0 records out
10737418240 bytes (11 GB) copied, 40.2767 s, 267 MB/s
27TB NVMe drive:

## /mnt/nvme 27 TB:

dd if=/dev/zero \
    of="/opt/dlami/nvme/test_10G_on_nvme.img" \
    bs=1G count=10 oflag=direct status=progress
6442450944 bytes (6.4 GB) copied, 1.027345 s, 6.3 GB/s
10+0 records in
10+0 records out
10737418240 bytes (11 GB) copied, 1.50334 s, 7.1 GB/s

# lsblk -d -o NAME,TRAN,SIZE,MODEL
NAME    TRAN    SIZE MODEL
nvme0n1 nvme    100G Amazon Elastic Block Store
nvme1n1 nvme    500G Amazon Elastic Block Store
nvme4n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme2n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme7n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme9n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme3n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme5n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme6n1 nvme    3.5T Amazon EC2 NVMe Instance Storage
nvme8n1 nvme    3.5T Amazon EC2 NVMe Instance Storage

Happy to close this issue now, but feel free to let me know if you all have further questions or things you'd like me to try.