Code

Poc elastic training#1310

Code

Draft

lukebaumann wants to merge 27 commits intomainAI-Hypercomputer/maxtext:mainfrom poc-elastic-trainingAI-Hypercomputer/maxtext:poc-elastic-trainingCopy head branch name to clipboard

+3,232-1,046

Commits

Commits on Feb 28, 2025

• Add Pathways Benchmarking Recipes for Scale Testing

  SujeethJinesh

  committed
  a410df6

  Copy full SHA for a410df6

Commits on Mar 1, 2025

• Updated input pipeline to use values from elastic utils.

  Show description for ccd8c26
  lukebaumann

  committed
  ccd8c26

  Copy full SHA for ccd8c26

• Checking in benchmark changes

  lukebaumann

  committed
  9fdab26

  Copy full SHA for 9fdab26

• checking in. async checkpointing does not work. no checkpointing or synchronized checkpointing might

  lukebaumann

  committed
  62b6eb0

  Copy full SHA for 62b6eb0

Commits on Mar 3, 2025

• Wait until all slices are available before starting. Fixed a bug with the number of values returned from reshard_fn

  lukebaumann

  committed
  33516b9

  Copy full SHA for 33516b9

Commits on Mar 4, 2025

• Updated to three maybe_* functions

  lukebaumann

  committed
  0298e13

  Copy full SHA for 0298e13

• Added an initial save at the beginning

  lukebaumann

  committed
  ef3d7c5

  Copy full SHA for ef3d7c5

• Removed restoring from checkpoint during reshard_handler

  lukebaumann

  committed
  7d1cf08

  Copy full SHA for 7d1cf08

Commits on Mar 5, 2025

• Updates for the refactor

  lukebaumann

  committed
  d5eb1b4

  Copy full SHA for d5eb1b4

• Using shaurya's images, 20 checkpoint period, and sync checkpointing

  lukebaumann

  committed
  f05b93c

  Copy full SHA for f05b93c

• Testing on smaller cluster and turning off profiler

  lukebaumann

  committed
  4f3c822

  Copy full SHA for 4f3c822

• Updated simulator init args

  lukebaumann

  committed
  6cb73e9

  Copy full SHA for 6cb73e9

• Using simulator and simplified maybe_reshard_up/down return values

  lukebaumann

  committed
  2511d6e

  Copy full SHA for 2511d6e

• Fixed the bug. Indented to the wrong level so jit was within the context manager of the old mesh when resharding up

  lukebaumann

  committed
  8b6c8ff

  Copy full SHA for 8b6c8ff

• Removed the simulator

  lukebaumann

  committed
  f4110ed

  Copy full SHA for f4110ed

• Removed an unecessary block

  lukebaumann

  committed
  60c3583

  Copy full SHA for 60c3583

• Using the big cluster

  lukebaumann

  committed
  7a779ab

  Copy full SHA for 7a779ab

• Added back one block

  lukebaumann

  committed
  e5ccfe9

  Copy full SHA for e5ccfe9

• Turning off profiler

  lukebaumann

  committed
  c7d3492

  Copy full SHA for c7d3492

Commits on Mar 7, 2025

• Added NOT_FOUND to exception types that we try to handle with resharding. This attempts to address part of the issue with and elastic event at the same time as a profiler

  lukebaumann

  committed
  220c70f

  Copy full SHA for 220c70f

• Fix to maybe_reshard_up for handling the retry logic

  lukebaumann

  committed
  ba07bf4

  Copy full SHA for ba07bf4

Commits on Mar 8, 2025

• Fixed the retry logic for reshard up (down was fixed prior). Fixed several linting issues

  lukebaumann

  committed
  c0ad90c

  Copy full SHA for c0ad90c

• Fixed another bug for reshard down retry logic

  lukebaumann

  committed
  8fd5628

  Copy full SHA for 8fd5628

• Checking in. Things are not working

  lukebaumann

  committed
  b575d37

  Copy full SHA for b575d37

• Cleaned up some random logs. I believe there is a race condition for the snapshot

  lukebaumann

  committed
  6b34aac

  Copy full SHA for 6b34aac

• Bumping the reshard check period to 15

  lukebaumann

  committed
  ce94037

  Copy full SHA for ce94037

Commits on Mar 9, 2025

• Changed name from save to snapshot. Added a snapshot deque. Refactored some code from elastic handler into elastic utils

  lukebaumann

  committed
  1fedd26

  Copy full SHA for 1fedd26