Skip to content

v1.15.0: FSDP2 activation memory, dtensor improvements

Latest
Latest

Choose a tag to compare

View all tags
@SunMarc SunMarc released this 09 Sep 13:30
· 1 commit to main since this release
v1.15.0
6afc1e5

v1.15.0: FSDP2 activation memory, dtensor improvements

FSDP2

A large batch of FSDP2 work this release: two fixes that cut activation memory at long sequence lengths, tied-embedding support on torch >= 2.13, and a round of checkpointing correctness and scale fixes.

Activation checkpointing was wrapping each child of the matched layer (self_attn, mlp, the norms) instead of the layer itself, so every inter-child activation stayed saved for backward. It now wraps the layer.

There's also a new FSDP2-only activation_checkpointing_offload, which moves the remaining per-layer checkpoint inputs to pinned CPU memory. Gradients are exactly those of plain activation checkpointing:

# fsdp2.yaml
fsdp_config:
  fsdp_version: 2
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_activation_checkpointing: true
  fsdp_activation_checkpointing_offload: true
accelerate launch --config_file fsdp2.yaml train.py
  • FSDP2 activation checkpointing: wrap the matched transformer layer itself, not each of its children by @qgallouedec in #4172
  • Add FSDP2 activation_checkpointing_offload: offload checkpointed layer inputs to pinned CPU memory by @qgallouedec in #4175
  • Fix FSDP2 tied-embedding models on torch >= 2.13: put the output embedding in the same fully_shard group by @qgallouedec in #4171
  • Fix FSDP2/PEFT/FULL_STATE_DICT dropping every rank's adapter shard except rank 0 by @AmineDiro in #4206
  • FSDP2: per-rank torch.save/load for SHARDED_STATE_DICT to fix 2800+ NPU checkpoint timeout by @gygdh-001 in #4105
  • Fix FSDP sharded checkpoint path resolution by @HaomingSong in #4119
  • Raise a clear error when FSDP is enabled on a mesh with no shard dimension by @qgallouedec in #4180

DTensor

Two fixes for DTensor-sharded models, which you hit with FSDP2, tensor parallelism, or any N-D parallelism setup: gradient clipping no longer fails on the foreach op when plain tensors and DTensors are mixed, and prepare_model leaves an already-sharded model where it is:

Offloading & Quantization

An entire model can now be dispatched to disk, including tied weights — useful for tools like llm-compressor that compress large models on machines that can't hold them:

Trackers

Custom trackers can be registered by name and then selected from log_with= like any built-in one:

from accelerate import Accelerator
from accelerate.tracking import register_tracker_class

register_tracker_class(MyTracker)  # MyTracker.name == "my_tracker"
accelerator = Accelerator(log_with="my_tracker")

Device support

Neuron gains a torch dynamo backend (so --torch-compile works with the Transformers Trainer) and MPS is now reported and handled properly by accelerate env and find_executable_batch_size.

CLI

  • Handle missing accelerate executable in env command by @tandede in #4168
  • Fix accelerate launch --cpu never setting the KMP variables by @yupengtang in #4222
  • Fix estimate-memory for timm>=1.0.29 by adding the hf-hub: prefix by @iamsharduld in #4213

Minor fixes

  • Refuse context parallelism for models with sliding-window or chunked attention layers by @qgallouedec in #4177
  • Fix convert_model_to_fp8_ao converting the first and last linear layers by @vineethsaivs in #4147
  • Fix get_non_persistent_buffers mutating module._non_persistent_buffers_set by @sohumt123 in #4116
  • fix: preserve dataloader iteration state when resuming from checkpoint by @TomQunChao in #4071
  • Fix load_accelerator_state only restoring one RNG backend by @Rakshit-gen in #4217
  • Reach through the DeepSpeed optimizer wrapper in AcceleratedOptimizer.eval() by @vineethsaivs in #4130
  • Fix compile_regions running the uncompiled module by @hjinnkim in #4188
  • Forward reduce_batch_size_fn through the find_executable_batch_size decorator form by @vineethsaivs in #4160
  • Fix distributedType -> distributed_type typo in downcast_bf16 guard by @uttam12331 in #4149
  • Fix duplicated words in user-facing messages by @Sreekant13 in #4091
  • Add missing type hints in checkpointing.py by @RudrenduPaul in #4153
  • Add missing type hints in utils/memory.py by @RudrenduPaul in #4152
  • docs: fix parameter names and a ghost entry in docstrings by @DaoyuanLi2816 in #4076
  • docs: fix documented parameter names that do not match the signature by @massimiliano1991 in #4156
  • docs: fix garbled sentence in CONTRIBUTING.md by @simpleqt in #4231
  • docs: fix three dead links left by the docs restructure by @simpleqt in #4230
  • Fix mixed precision argument typo in examples README by @noobyalan in #4192
  • Make doc builds faster by @mishig25 in #4103
  • Add security policy by @SunMarc in #4134
  • Bump pre-commit ruff hook to v0.13.1 to match the quality extra by @devangpratap in #4089
  • Fix flaky behavior in gated-model CLI test by making the test deterministic by @sywangyi in #4073
  • [test] Update min version torch by @SunMarc in #4170
  • Fix CI by @SunMarc in #4229
  • chore(deps): bump the actions group across 1 directory with 7 updates by @dependabot[bot] in #4167
  • chore(deps): bump the actions group across 1 directory with 4 updates by @dependabot[bot] in #4182
  • chore: update build_and_run_tests.yml by @hf-security-analysis[bot] in #4169

New Contributors

Full Changelog: v1.14.0...v1.15.0

Contributors

shoemoney, iamsharduld, and 30 other contributors
Assets 2
Loading