Release notes for the SageMaker model parallelism library

Fixed an ImportError issue when using moe_load_balancing=balanced in torch.sagemaker.moe.moe_config.MoEConfig for expert parallelism.

Fixed a fine-tuning issue where the torch.sagemaker.transform call raised KeyError when load_state_dict_from_rank0 is enabled.

Fixed an out-of-memory (OOM) error raised when loading large Mixture of Experts (MoE) models, such as Mixtral 8x22B, for fine-tuning.

SMP Docker container for PyTorch v2.2.0 with CUDA v12.1

658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.2.0-gpu-py310-cu121

Added a new core feature, expert parallelism, to support Mixture of Experts transformer models. To learn more, see Expert parallelism.

SMP Docker container for PyTorch v2.2.0 with CUDA v12.1

658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.2.0-gpu-py310-cu121

Added support for FP8 training of the following Hugging Face transformer models on P5 instances with Transformer Engine integration:

Fixed a bug where tensors were not guaranteed to be contiguous before the AllGather collective call during tensor parallelism training.

Added support for PyTorch v2.2.0.

Upgraded the SMDDP library to v2.2.0.

Upgraded the FlashAttention library to v2.3.3.

Upgraded the NCCL library to v2.19.4.

Discontinued support for Transformer Engine versions before v1.2.0.

The SMP Activation offloading feature currently does not work. Use the native PyTorch activation offloading instead.

Included a patch to fix the performance regression discussed in the issue thread at https://github.com/pytorch/pytorch/issues/117748 in the PyTorch GitHub repository.

SMP Docker container for PyTorch v2.2.0 with CUDA v12.1

658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.2.0-gpu-py310-cu121

Added support for PyTorch v2.1.2.

Discontinued support for Hugging Face Transformers v4.31.0.

An issue is discovered that fine-tuning of the Hugging Face Llama 2 model with attn_implementation=flash_attention_2 and FSDP causes the model to diverge. For reference, see the issue ticket in the Hugging Face Transformers GitHub repository. To avoid the divergence issue, use attn_implementation=sdpa. Alternatively, use the SMP transformer model implementation by setting up use_smp_implementation=True.

SMP Docker container for PyTorch v2.1.2 with CUDA v12.1

658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.1.2-gpu-py310-cu121

https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/smp-v2/

A new torch.sagemaker package, entirely revamped from the previous smdistributed.modelparallel.torch package in SMP v1.x.

Support for PyTorch 2.0.1.

Support for PyTorch FSDP.

Tensor parallelism implementation by integrating with the Transformer Engine library.

Support for both SageMaker Training and SageMaker HyperPod.

SMP v2 revamped the APIs entirely and provides the torch.sagemaker package. Mostly, you only need to initialize with the torch.sagemaker.init() module and pass model parallel configuration parameters. With this new package, you can significantly simplify code modifications in your training script. To learn more about adapting your training script to use SMP v2, see Get started with the SageMaker model parallelism library v2.

If you've used SMP v1 for training Hugging Face Transformer models and want to reuse the models in SMP v2, see Upgrade from SMP v1 to SMP v2.

For PyTorch FSDP training, you should use SMP v2.

Activation checkpointing currently only works with the following wrapping policies with FSDP.

To use Activation offloading, FSDP activation checkpointing type must be REENTRANT.

When running with tensor parallel enabled with the sharded data parallel degree set to 1, you must use backend = nccl. The smddp backend option is not supported in this scenario.

Transformer Engine is required to use PyTorch with the SMP library even when not using tensor parallelism.

Starting from this release, the documentation for the SageMaker model parallelism library is fully available in this Amazon SageMaker Developer Guide. In favor of this complete developer guide for SMP v2 in the Amazon SageMaker Developer Guide, the additional reference for SMP v1.x in the SageMaker Python SDK documentation is deprecated. If you still need the documentation for SMP v1.x, the developer guide for SMP v1.x is available at (Archived) SageMaker model parallelism library v1.x, and the SMP Python library v1.x reference is available in the SageMaker Python SDK v2.199.0 documentation.

Discontinued support for TensorFlow.

There is no pipeline parallelism support in SMP v2.

There is no support for the DeepSpeed library in favor of native PyTorch FSDP.

SMP Docker container for PyTorch v2.0.1 with CUDA v12.1

658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.0.1-gpu-py310-cu121