DeepSpeed

54fd5317 - Merge with parent branch; Activation checkpointing support non-tensor args/return values