ml-explore / ml-explore/mlx

FSDP saving state improvement

Open
#3,401 2 comments 0 reactions 1 assignee View on GitHub

@nastya236 is already working on this.

Since Apr 12, 2026.

cuda distributed enhancement
Dominant language
C++
Stars
28.5k
Forks
2.3k
Avg merge
3d 8h
Merged PRs (30d)
62

Description

Currently in FSDP skips intermediate reconstruction of tree. This results in the following issues:

  • grouping for the parameters might be non-deterministic -> loaded state might be incorrect
  • state tight to the number of groups -- so to the size of the communication, which is wrong design

While current strategy works deterministically if everything is fixed from run to run, however improvement is needed.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.