ZeRO-Offload CPU, memory utilization

[taehyunzzz]

I am fine-tuning BERT-large for text classification using transformers library.
Using Intel's pcm tool, I plotted memory bandwidth utilization and CPU utilization for 1000 iterations of weight update only (no forward / backward). Code is something like for i in range(1000) : model_engine.step()
The below plots memory bandwidths, memory read/writes (cum : cumulative), and average CPU utilization of 64 virtual cores (32 physical cores in 2 sockets with hyper-threading).
Green plot is the total memory bandwidth. Blue and orange are for socket 0 and 1 respectively.

I am not certain why CPU utilization is bounded to 50%?
Also, why is memory from the 2 sockets not used uniformly? The bandwidth utilization from the 2 sockets are skewed.
Weight update is known to be memory-bound, but it seems that it is actually not using full bandwidth nor compute.

Does this have something to do with the number of sub-processes when running deepspeed?
I noticed that multiple layers of subprocesses were spawned when training on deepspeed while debugging on vscode.

I also worked with a micro-kernel that resembles C++ CPUAdam code. My code runs CPU weight update on a single random 1GB parameter. Again, code looks something like for i in range(1000) : model_engine.step() in C++.
Results showed that all virtual CPU cores are busy, and memory bandwidth for 8 memory channels were used almost evenly.
Still, Memory utilization still shows about 17 GBps per memory channel, whereas maximum bandwidth is reported to be 22 GBps for my system.

I am using Deepspeed 0.7.7 with AVX512.
My system is equipped with dual-socket Xeon Gold 6226R w/ 8 memory channels (4 per socket) and 3090 Ti GPU.

[taehyunzzz]

Found the issue. Seems like it is due to Pytorch default num_threads. When I set the number of threads with torch.set_num_threads(64), it started to work as expected 💯 Check this.