●Change8

v1.14.0

📦 accelerateView on GitHub →
✨ 7 features🐛 24 fixes🔧 9 symbols

Summary

This release focuses heavily on FSDP2 improvements, including better dtype handling, sharding of embeddings/norms, and QLoRA stability. It also introduces end-to-end support for AMD ROCm devices.

Migration Steps

  1. If you rely on the package name `pytorch-triton-xpu`, update your dependencies to use `triton-xpu` instead.

✨ New Features

  • FSDP2 now fully shards embedding and norm layers.
  • Added support for dynamic batch size in BatchSamplerShard with the `even_batches` option.
  • Added support for AMD ROCm devices, enabling end-to-end acceleration.
  • Added padded allgather and broadcast operations for Neuron devices to reduce recompilation.
  • Improved offloading support for quantized models, including Torchao, int8, and tied-weight handling.
  • Added `max` reduction type utility.
  • The package name `pytorch-triton-xpu` was renamed to `triton-xpu`.

🐛 Bug Fixes

  • Fixed dtype mismatch when loading full state dict in FSDP2.
  • Fixed region compilation issues in FSDPv2.
  • Fixed mixed-dtype AssertionError by casting the model to a uniform dtype before calling `fully_shard` in FSDP2.
  • Fixed QLoRA crashes by auto-excluding non-floating frozen Params4bit from `fully_shard` in FSDP2.
  • Fixed FSDP2 auto-wrap policy ignoring the `_no_split_modules` fallback.
  • Fixed key-based matching in `fsdp2_load_full_state_dict`.
  • Added missing `model_has_params4bit` guard to `fsdp2_load_full_state_dict` call.
  • Fixed issue where FSDP1 keys (REMOVED / NOT_YET_IMPLEMENTED) were leaked when transitioning to FSDP2.
  • Prevented double-wrapping models when calling `prepare_model()`.
  • Fixed int8 offload hook detachment statistics restoration.
  • Fixed `keep_in_fp32_modules` not working for tied weights in `load_and_quantize_model`.
  • Fixed `dtype_byte_size` calculation for FP8 fnuz / e8m0fnu dtypes.
  • Fixed iterable dataset sharding condition when `n_shards == num_processes`.
  • Fixed implicit padding in `split_between_processes` when `apply_padding=False` and `num_samples < num_processes`.
  • Allowed Flash Attention kernels in DeepSpeed Sharded Prediction (SP).
  • Conditionally imported `torch.distributed.algorithms.join` in `accelerator.py`.
  • Fixed `is_hf_initialized` attribute state.
  • Fixed MLU backend not being part of the `_prepare_backend` elif chain.
  • Fixed notebook launcher CUDA initialization.
  • Relaxed numerical tolerance for XPU in `test_big_modeling`.
  • Fixed Gloo backend error during `test_load_checkpoint_and_dispatch_with_broadcast` on XPU.
  • Fixed `TrackioTracker.log()` ignoring the `step` parameter.
  • Fixed MLflowTracker.store_init_configuration mutating the caller's config dictionary.
  • Added missing Neuron device case handling.

Affected Symbols

fsdp2_load_full_state_dictfully_shardprepare_model()BatchSamplerShardload_and_quantize_modelTrackioTracker.log()MLflowTracker.store_init_configurationaccelerator.pyParallelismConfig.get_device_mesh