scr: initial integration for Microsoft Megatron-Deepspeed

[adammoody]

Integrates SCR into Megatron-DeepSpeed.

This requires the work here to integrate SCR into DeepSpeed: adammoody/DeepSpeed#1

This adds new Megatron arguments for enabling and configuring SCR within a training run:

--scr - enable SCR for checkpointing
--scr-interval - number of steps between defensive checkpoints
--scr-seconds - number of seconds between defensive checkpoints
--scr-overhead - maximum percent runtime overhead to allow for defensive checkpoints

With SCR, one can write checkpoints to node-local storage, which is often faster than the parallel file system at large scale. By default, SCR writes to /dev/shm, and it can be configured to use a node-local SSD if available. Those checkpoints can be transferred to the parallel file system upon a failure, or they can be discarded if there is no failure. This allows one to increase the defensive checkpoint frequency of a training run without adding significant run time or requiring more disk space.

For example, one might configure a run to save a checkpoint every 100 steps in /dev/shm but only persist a checkpoint on the parallel file system every 2000 steps:

--scr
--scr-interval=100
--save-interval=2000

For best performance, one should also enable parallel writes of the layer checkpoint files in the DeepSpeed configuration file:

  "checkpoint": {
    "parallel_write": {
      "pipeline_stage": true
    }
  }