[Distributed checkpoint] save_state_dict and load_state_dict fail on multi node with torch2.0

[nkflash]

🐛 Describe the bug

I am working on module checkpoint by distributed checkpoint.
There are two pods(one is master with rank 0, other is worker with rank 1)
I found only master save checkpoint file, worker save nothing(but create the directory)
Code looks like:

            if config.checkpoint_iters is not None and self.iter_num % config.checkpoint_iters == 0:
                model_state_dict = model.state_dict()
                sd = {'module': model_state_dict, 'iteration': self.iter_num}
                fs_storage_writer = FileSystemWriter("/tmp/checkpoint")
                save_state_dict(
                    state_dict=sd,
                    storage_writer=fs_storage_writer,
                )

    if args.start_from_checkpoint:
        sd = {'module': model.state_dict(), 'iteration': 0}
        fs_storage_reader = torch.distributed.checkpoint.FileSystemReader("/tmp/checkpoint")
        load_state_dict(
            state_dict=sd,
            storage_reader=fs_storage_reader,
        )

When I load the cp I got such error:

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @WZ @wz337

Versions

Collecting environment information...
PyTorch version: 2.0.0a0+1767026
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.24.1
Libc version: glibc-2.31

Python version: 3.8.10 (default, Mar 13 2023, 10:26:41) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-113-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 12.1.66
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 470.129.06
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.8.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.8.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Byte Order: Little Endian
Address sizes: 43 bits physical, 48 bits virtual
CPU(s): 256
On-line CPU(s) list: 0-255
Thread(s) per core: 2
Core(s) per socket: 64
Socket(s): 2
NUMA node(s): 8
Vendor ID: AuthenticAMD
CPU family: 23
Model: 49
Model name: AMD EPYC 7742 64-Core Processor
Stepping: 0
Frequency boost: enabled
CPU MHz: 3261.352
CPU max MHz: 2250.0000
CPU min MHz: 1500.0000
BogoMIPS: 4491.95
Virtualization: AMD-V
L1d cache: 4 MiB
L1i cache: 4 MiB
L2 cache: 64 MiB
L3 cache: 512 MiB
NUMA node0 CPU(s): 0-15,128-143
NUMA node1 CPU(s): 16-31,144-159
NUMA node2 CPU(s): 32-47,160-175
NUMA node3 CPU(s): 48-63,176-191
NUMA node4 CPU(s): 64-79,192-207
NUMA node5 CPU(s): 80-95,208-223
NUMA node6 CPU(s): 96-111,224-239
NUMA node7 CPU(s): 112-127,240-255
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate sme ssbd mba sev ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.22.2
[pip3] pytorch-quantization==2.1.2
[pip3] torch==2.0.0a0+1767026
[pip3] torch-tensorrt==1.4.0.dev0
[pip3] torchtext==0.13.0a0+fae8e8c
[pip3] torchvision==0.15.0a0
[pip3] triton==2.0.0
[conda] Could not collect

[nkflash]

checkpoint file on master pod looks like:

[wz337]

@nkflash Do you have a repo or a toy model that we can use to repro this?

[nkflash]

@nkflash Do you have a repo or a toy model that we can use to repro this?

Try this git@github.com:nkflash/minGPT.git
use pre-train branch.
I just use minGPT to test distributed checkpoint.

Use train.sh to launch the test

@wz337

[nkflash]

@nkflash Do you have a repo or a toy model that we can use to repro this?

Try this git@github.com:nkflash/minGPT.git use pre-train branch. I just use minGPT to test distributed checkpoint.

Use train.sh to launch the test

@wz337

@wz337 any progress?

[kumpera]

Are you using a network file system? Distributed Checkpoint current implementation depends on writing to a place that is shared by all ranks.

Looks like this is not the case.

[nkflash]

Are you using a network file system? Distributed Checkpoint current implementation depends on writing to a place that is shared by all ranks.

Looks like this is not the case.

@kumpera At first I used shared FS to test, got same error. So I used local FS to test.

I retest it with shared FS, same error as bellow

[kumpera]

I managed to reproduce your bug and it's due to inconsistent usage of models during loading and saving.

When saving, you pass self.model to save_state_dict which is an instance of DDP.
When loading you pass model to load_state_dict which is an instance of GPT.

Why does this matter? Because DDP inserts a module. prefix to all keys in the state_dictand that causes checkpointing to fail.

The following diff made it work for me:

index aa76515..24af86c 100644
--- a/mingpt/trainer.py
+++ b/mingpt/trainer.py
@@ -68,7 +68,8 @@ class Trainer:
         self.iter_num = 0
         self.dir_locate = 'checkpoint'
         if config.start_from_checkpoint:
-            sd = {'module': model.state_dict(), 'iteration': 0}
+            sd = {'module': self.model.state_dict(), 'iteration': 0}
+
             fs_storage_reader = FileSystemReader(self.dir_locate)
             load_state_dict(
                 state_dict=sd,
@@ -76,7 +77,7 @@ class Trainer:
             )
             print(f'start train from checkpoint, iterator:{sd["iteration"]}')
             self.iter_num = sd["iteration"]
-            model.load_state_dict(sd["module"])
+            self.model.load_state_dict(sd["module"])
         # variables that will be assigned to trainer class later for logging and etc
         self.iter_time = 0.0
         self.iter_dt = 0.0

We should improve pytorch to emit warnings both during loading and saving if this situation is detected.

[nkflash]

I managed to reproduce your bug and it's due to inconsistent usage of models during loading and saving.

When saving, you pass self.model to save_state_dict which is an instance of DDP. When loading you pass model to load_state_dict which is an instance of GPT.

Why does this matter? Because DDP inserts a module. prefix to all keys in the state_dictand that causes checkpointing to fail.

The following diff made it work for me:

index aa76515..24af86c 100644
--- a/mingpt/trainer.py
+++ b/mingpt/trainer.py
@@ -68,7 +68,8 @@ class Trainer:
         self.iter_num = 0
         self.dir_locate = 'checkpoint'
         if config.start_from_checkpoint:
-            sd = {'module': model.state_dict(), 'iteration': 0}
+            sd = {'module': self.model.state_dict(), 'iteration': 0}
+
             fs_storage_reader = FileSystemReader(self.dir_locate)
             load_state_dict(
                 state_dict=sd,
@@ -76,7 +77,7 @@ class Trainer:
             )
             print(f'start train from checkpoint, iterator:{sd["iteration"]}')
             self.iter_num = sd["iteration"]
-            model.load_state_dict(sd["module"])
+            self.model.load_state_dict(sd["module"])
         # variables that will be assigned to trainer class later for logging and etc
         self.iter_time = 0.0
         self.iter_dt = 0.0

We should improve pytorch to emit warnings both during loading and saving if this situation is detected.

Great! Thank you for pointing out my mistakes.