V0.2 Release Plan

[cp5555]

Release Manager

@cp5555

Endgame

• Code freeze: July 5th, 2023
• Bug Bash date: July 8th, 2023
• Release date: July 19th, 2023

Main Features

MS-AMP O3 Optimization

1. • Support pipelining parallelism and tensor parallelism (Support bf16+pipeline+ZeRO1 #81, Related to Support pipeline parallelism and tensor parallelism #40)
2. • Support Zero optimization (Related to Support Zero #34) (feat(deepspeed): Deepspeed Optimizer with MS-AMP support  #71, feat(deepspeed): Deepspeed ZeRO Optimizer with MS-AMP support  #73, Support msamp section in deepspeed config #77, and fix(zero): fix the bug of gradient accumulation in FP8 deepspeed ZeRO #78)

MS-AMP New Features

1. • Support ScalingTensor in functional.linear ([Feature] ScalingTensor support for F.linear #65)
2. • Support customized attributes in FP8Linear ([Feature] Copy custom attributes from Linear to FP8Linear #64)

MS-AMP Bug Fix

1. • Fix bug in all_reduce_grads() when training multiple models (Related to [Bug] LBOptimizer.all_reduce_grads reduces gradients of only a model, even if training several models #62) ([Bug Fixed] Fix all reduce grads when training more than one models #63)
2. • Fix bug in the cast_to_fp16 ([Bug Fixed] update the cast of ScalingTensor-FP16 #76)

MS-AMP Improvement

1. • Performance optimization (Related to Optimize performance to close the gap between TE and MS-AMP #30) (Wrap tex fp8 transpose in linear backward #23 and Update adamw optimizer with vectorized apis #39)
2. • Add DockerFile for CUDA12.1 support (Add Dockerfile for msamp #72)
3. • Support PyTorch 2 (Support pytorch 2.1 #74)
4. • Support FP8 progress group in PyTorch (Related to Support FP8 ProcessGroup in pytorch #50) (Support torch.distributed.all_reduce for fp8 tensor #85)
5. • Add DeepSpeed cifar10 as UT (add cifar10 example using deepspeed and MS-AMP #79)
6. • Cache TE build in pipeline (Cache TE build to save time in pipelines #82)
7. • Remove env PYTHONOPTIMIZE in docker image (Remove env PYTHONOPTIMIZE in docker image #83)

MS-AMP Demo & Documentation

1. • Add RoBERTa FP8 demo
2. • Add GPT-345M and 13B FP8 demo (Add gpt 345M and 13B example MS-AMP-Examples#8)
3. • Support Transformer Engine FP8 in DeiT and Swin-Transformer (Support Transformer Engine FP8 in DeiT and Swin-Transformer MS-AMP-Examples#5)
4. • Improve the homepage (Refine homepage #66, and Add TE results on homepage #69)

Backlog

1. Move extension installation from post install to setup.py (Related to Moving extension installation from post install to setup.py under project root folder #43)
2. Solve NCCL installation issue (Related to nccl buildig failed without specifying NVCC_GENCODE  #44)
3. Support latest Megatron-LM
4. Support auto scaling factor tuning (ASFT) for FP8 collective communication (Related to Auto scaling factor tuning for FP8 collective communication #41)
5. Setup MS-AMP Website (Related to MS-AMP needs a website #29)
6. Remove the dependency on Nvidia Transformer Engine (TE) (Related to Remove the dependency of Transformer Engine #33)
7. Support DistributedDataParallel with FP8-support ([WIP] DistributedDataParallel with FP8-support #75)