numpy: The latest release of numpy is installed alongside all torch builds. Previously, only the python3-numpy package from the Ubuntu distribution in use was included in the final torch images.
Additionally, the torch:nccl build with Ubuntu 20.04 × CUDA 12.2.2 now uses an updated base image featuring NCCL v2.21.5-1
Stability
Steps that were liable to intermittent failure due to network errors during the torch build are now retried after a random delay instead of cancelling the build run.
Approximately 1/3 of the available CPUs are used for concurrent compilation jobs during torch builds, instead of all CPUs.
Available disk space is reported prior to compiling torch, to better catch storage issues during CI.
Fixes
A version of triton compiled from source is now included with torch-nightly builds, since custom triton source commits are used during PyTorch development that are not necessarily compatible with the distributed versions available through package managers.
The newest update of deepspeed included in torch-extras has had ahead-of-time compilation for its CCL_COMM op disabled to be able to compile without oneCCL installed.
Performance
The main ml-containers build workflow now accepts an optional cache-key argument for more granular caching.
The cache-key and image-name parameters are used together to form a complete cache key for a given image.
Different torch/torch-extras/nightly-torch/nightly-torch-extras image flavours (corresponding to their base images) now use different cache keys.
This leads to a greatly improved cache hit rate, as previously, dozens of completely incompatible image builds were competing for the same cache spots.
Improved cache hit rates at build time should correspond to better image pull times when newer images can more often share perfectly identical layers with older images.
Updates for PyTorch, Apex, DeepSpeed, Flash Attention, and
nccl-testsBase Images, Fix for Triton intorch-nightly, and Caching ImprovementsWhat a mouthful of a title!
Updates
This change contains the following library version updates:
torch: 2.2.2 → 2.3.0vision: 0.17.2 → 0.18.0audio: 2.2.2 → 2.3.0apex:2386a91→a7de60edeepspeed: 0.12.6 → 0.14.2flash_attn: 2.4.2 → 2.5.8numpy: The latest release ofnumpyis installed alongside alltorchbuilds. Previously, only thepython3-numpypackage from the Ubuntu distribution in use was included in the finaltorchimages.Additionally, the
torch:ncclbuild with Ubuntu 20.04 × CUDA 12.2.2 now uses an updated base image featuring NCCL v2.21.5-1Stability
torchbuild are now retried after a random delay instead of cancelling the build run.torchbuilds, instead of all CPUs.torch, to better catch storage issues during CI.Fixes
tritoncompiled from source is now included withtorch-nightlybuilds, since customtritonsource commits are used during PyTorch development that are not necessarily compatible with the distributed versions available through package managers.deepspeedincluded intorch-extrashas had ahead-of-time compilation for itsCCL_COMMop disabled to be able to compile without oneCCL installed.Performance
ml-containersbuild workflow now accepts an optionalcache-keyargument for more granular caching.cache-keyandimage-nameparameters are used together to form a complete cache key for a given image.torch/torch-extras/nightly-torch/nightly-torch-extrasimage flavours (corresponding to their base images) now use different cache keys.ccachecache size was increased to 5 GiB.