Add Support for DeepSpeed ZeRO Stage 1

[mjdenkowski]

This PR adds training support for DeepSpeed with ZeRO stage 1:

• Install DeepSpeed: pip install deepspeed
• Launch training with DeepSpeed: deepspeed --no_python ... sockeye-train ...
• Optionally specify FP16 or BF16 mode with --deepspeed-fp16 or --deepspeed-bf16
• DeepSpeed automatically passes the --local_rank option to Sockeye, which activates DeepSpeed mode.

Sockeye's DeepSpeed mode does the following:

• Initialize DeepSpeed with ZeRO stage 1 (optimizer state sharding).
• Skip model tracing because DeepSpeed makes runtime modifications to the model that may not be JIT-compatible.
• Guarantee model outputs are in FP32 before computing losses.
• Have DeepSpeed automatically handle several aspects of training: moving the model to the correct device, accumulating gradients across multiple updates, synchronizing states for model updates, and stepping the learning rate scheduler and optimizer.
• Use DeepSpeed to save/load model parameters, optimizer states, and learning rate scheduler states.

DeepSpeed saves checkpoints using its own format. The new convert_deepspeed module recovers standard Sockeye parameter files from these checkpoints. At the end of training, checkpoints are automatically converted. Checkpoints can also be converted manually using the CLI.

DeepSpeed is an optional dependency (like Apex) and these changes are backward compatible. Training results with DeepSpeed ZeRO stage 1 are identical for this PR and the deepspeed branch. Training results for standard FP32 mode, PyTorch AMP, and Apex AMP are identical for this PR and the main branch.

Pull Request Checklist

• [x] Changes are complete (if posting work-in-progress code, prefix your pull request title with '[WIP]' until you can check this box.
• [x] Unit tests pass (pytest)
• [x] Were system tests modified? If so did you run these at least 5 times to account for the variation across runs?
• [x] System tests pass (pytest test/system)
• [x] Passed code style checking (./style-check.sh)
• [x] You have considered writing a test
• [x] Updated major/minor version in sockeye/__init__.py. Major version bump if this is a backwards incompatible change.
• [x] Updated CHANGELOG.md

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[mjdenkowski]

Great addition, thanks!

I am wondering if there is a good way to test some of the added functionality in Github actions. I guess its difficult to install deepspeed in the workflows. Would it be possible to test the checkpoint conversion module somehow?

That's a good idea. It's now possible to pip install deepspeed on CPU-only machines. Training won't run (GPUs required), but some utility functions will.

I've added a test that runs checkpoint conversion with existing model files if the deepspeed package is installed. I've also updated our workflows to install deepspeed, which should be compatible with all of our environments. We can see if anything breaks.

[mjdenkowski]

Our automatic tests now cover installing deepspeed and converting a DeepSpeed checkpoint.