Checkpointing with DLIO

[A-Tarraf]

Hello,

The part about the checkpointing is broken in the documentation (https://dlio-benchmark.readthedocs.io/en/latest/config.html#checkpoint). Also, I get some inconsistency when I try to run a modified version of resnet50 with checkpointing (the file is provided at the end).

When I execute the benchmark with the following:

mpirun -np 100 dlio_benchmark workload=resnet50_my_a100 ++workload.workflow.generate_data=True ++workload.workflow.train=True

It prints out the following:

[OUTPUT] 2025-03-19T00:33:57.908099 Starting data generation
[OUTPUT] 2025-03-19T00:34:03.282143 Generation done
[OUTPUT] 2025-03-19T00:34:03.345361 Total number of parameters in the model: 3172149248
[OUTPUT] 2025-03-19T00:34:26.549998 Model size: 0.0591 GB
[OUTPUT] 2025-03-19T00:34:26.552479 Optimizer state size: 35.4516 GB
[OUTPUT] 2025-03-19T00:34:26.552583 Total checkpoint size: 35.5107 GB
[OUTPUT] 2025-03-19T00:34:26.555036 Max steps per epoch: 12 = 500 * 1024 / 400 / 100 (samples per file * num files / batch size / comm size)
[OUTPUT] 2025-03-19T00:34:27.549667 Starting epoch 1: 12 steps expected
[OUTPUT] 2025-03-19T00:34:27.549969 Starting block 1
...
[OUTPUT] 2025-03-19T00:34:53.602374 Epoch 1 - Block 1 [Training] Accelerator Utilization [AU] (%): 23.9001
[OUTPUT] 2025-03-19T00:34:53.602470 Epoch 1 - Block 1 [Training] Throughput (samples/second): 19953.6376
[OUTPUT] 2025-03-19T00:34:53.602537 Epoch 1 - Block 1 [Training] Computation time per step (second): 0.4356+/-0.0005 (set value: {'mean': 0.435})
[OUTPUT] 2025-03-19T00:34:53.614024 Starting saving checkpoint 1 after total step 12 for epoch 1
[OUTPUT] 2025-03-19T00:34:53.821952 Saved model checkpoint in 0.0302 seconds
[OUTPUT] 2025-03-19T00:34:53.837066 Saved optimizer checkpoint in 0.0150 seconds
[OUTPUT] 2025-03-19T00:34:53.837464 Finished saving checkpoint 1 for epoch 1 in 0.2234 s; Throughput: 158.9274 GB/s
[OUTPUT] 2025-03-19T00:34:53.839193 Ending epoch 1 - 12 steps completed in 26.29 s
...
[METRIC] ==========================================================
[METRIC] Number of Simulated Accelerators: 100 
[METRIC] Training Accelerator Utilization [AU] (%): 6.7870 (8.8023)
[METRIC] Training Throughput (samples/second): 5663.3138 (7344.8845)
[METRIC] Training I/O Throughput (MB/second): 619.2741 (803.1511)
[METRIC] train_au_meet_expectation: fail
[METRIC] Checkpoint save duration (seconds): 0.1669 (0.0337)
[METRIC] Checkpoint save I/O Throughput (GB/second): 220.3871 (38.1355)
[METRIC] ==========================================================

However, when I look into the repository for the checkpoints, I see that the size of the checkpoints is not even close to what the benchmark states:

drwxr-sr-x 2 xxx xxx 25K Mar 19 00:34 global_epoch1_step13
drwxr-sr-x 2 xxx xxx 25K Mar 19 00:36 global_epoch4_step13

I can also see the checkpoints, but their size is not correct:

xxx/checkpoints/resnet50 $ ls
global_epoch10_step13  global_epoch1_step13  global_epoch4_step13  global_epoch7_step13

xxx/checkpoints/resnet50 $ du -sh .
237K	.

xxx/checkpoints/resnet50 $ ls global_epoch10_step13 
checkpoint                               zero_pp_rank_0_mp_rank_0_model_states.pb-1.data-00000-of-00001  zero_pp_rank_0_mp_rank_0_optim_states.pb-1.index
model_states-0.pb-1.data-00000-of-00001  zero_pp_rank_0_mp_rank_0_model_states.pb-1.index
model_states-0.pb-1.index                zero_pp_rank_0_mp_rank_0_optim_states.pb-1.data-00000-of-00001

This is what my resnet50_my_a100.yaml file contains:

model:
  name: resnet50
  type: cnn
  model_size: 499153191
  num_layers: 50

framework: tensorflow

workflow:
  generate_data: True
  train: True
  checkpoint: True

dataset:
  num_files_train: 1024
  num_samples_per_file: 500 #1251
  record_length_bytes: 114660.07
  record_length_bytes_resize: 150528
  data_folder: data/resnet50
  format: tfrecord

train:
  computation_time: 0.435
  epochs: 10

reader:
  data_loader: tensorflow
  read_threads: 8
  computation_threads: 8
  batch_size: 400
  dont_use_mmap: True

checkpoint:
  checkpoint_folder: checkpoints/resnet50
  checkpoint_after_epoch: 1
  epochs_between_checkpoints: 3
  # model_size: 4034713312
  # type: all_ranks

metric:
  au: 0.90

Am I doing something wrong?

[A-Tarraf]

I looked deeper into this, and it is unrelated to how I execute the benchmark. I tried the standard bert_v100 with 10 instead of 500 files, but I had the same issue.

However, it worked with the llama 7b benchmarks generating checkpoints of a total of 900Gb:

 mpiexec -np 900 --oversubscribe --hostfile /lustre/project/nhr-admire/tarraf/tmp/hostfile_mpi --map-by node   dlio_benchmark  workload=llama_7b_zero3