Sockeye

Added

- log beam search avg output vocab size

Changed

- common base Search for GreedySearch and BeamSearch
- .pylintrc: suppress warnings about deprecated pylint warning suppressions

Fixed

- Send skip_nvs and nvs_thresh args now to Translator constructor in sockeye-translate instead of ignoring them.

Added

- Added training support for [DeepSpeed](https://www.deepspeed.ai/).
- Installation: `pip install deepspeed`
- Usage: `deepspeed --no_python ... sockeye-train ...`
- DeepSpeed mode uses Zero Redundancy Optimizer (ZeRO) stage 1 ([Rajbhandari et al., 2019](https://arxiv.org/abs/1910.02054v3)).
- Run in FP16 mode with `--deepspeed-fp16` or BF16 mode with `--deepspeed-bf16`.

Added

- Clean up GPU and CPU memory used during training initialization before starting the main training loop.

Changed

- Refactored training code in advance of adding DeepSpeed support:
- Moved logic for flagging interleaved key-value parameters from layers.py to model.py.
- Refactored LearningRateScheduler API to be compatible with PyTorch/DeepSpeed.
- Refactored optimizer and learning rate scheduler creation to be modular.
- Migrated to ModelWithLoss API, which wraps a Sockeye model and its losses in a single module.
- Refactored primary and secondary worker logic to reduce redundant calculations.
- Refactored code for saving/loading training states.
- Added utility code for managing model/training configurations.

Removed

- Removed unused training option `--learning-rate-t-scale`.

Added

- Added `sockeye-train` and `sockeye-translate` option `--clamp-to-dtype` that clamps outputs of transformer attention, feed-forward networks, and process blocks to the min/max finite values for the current dtype. This can prevent inf/nan values from overflow when running large models in float16 mode. See: https://discuss.huggingface.co/t/t5-fp16-issue-is-fixed/3139

Added

- Added support for offline model quantization with `sockeye-quantize`.
- Pre-quantizing a model avoids the load-time memory spike of runtime quantization. For example, a float16 model loads directly as float16 instead of loading as float32 then casting to float16.