DeepSpeed

cf4d815f - Unit tests for activation checkpointing arbitrary inputs/outputs