Very interesting, thanks for sharing! I’m not using the HuggingFace trainer, but I went ahead and ran with deepspeed with my current code.
I normally run:
ENV=var python -m train --arg1 val1 --arg2 val2

I ran:
ENV=var deepspeed ./train.py --arg1 val1 --arg2 val2

I don’t see any difference in runtime/memory usage, though. Is there anything else I need to do? My ds_report looks like this:
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires one of the following commands '['llvm-config', 'llvm-config-9']', but it does not exist!
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------

Did you pre-build the extension ops? I’ve got a CUDA Version mismatch that’s preventing me from pre-building them with DS_BUILD_OPS=1 pip install deepspeed.