Skip to content

Fix GaLore 8-bit updates with ZeRO CPU offload - #6432

Open
sdjasj wants to merge 1 commit into
hpcaitech:mainfrom
sdjasj:agent/fix-galore-cpu-offload-6399
Open

Fix GaLore 8-bit updates with ZeRO CPU offload#6432
sdjasj wants to merge 1 commit into
hpcaitech:mainfrom
sdjasj:agent/fix-galore-cpu-offload-6399

Conversation

@sdjasj

@sdjasj sdjasj commented Jul 29, 2026

Copy link
Copy Markdown

What changed

  • stage each CPU-offloaded distributed GaLore parameter shard and gradient on the accelerator immediately before its bitsandbytes update
  • move non-paged 8-bit optimizer state and GaLore projector state with the parameter
  • return the updated parameter, gradient, and optimizer state to CPU after every parameter update
  • add a two-step distributed regression that verifies the offloaded state remains on CPU between steps

Why

bitsandbytes only provides GPU update kernels. LowLevelZeroPlugin(cpu_offload=True) placed the master parameter, gradient, quantized states, maps, and block maxima on CPU, so the kernel rejected every input as non-GPU.

Staging one parameter at a time preserves CPU offload between updates and avoids materializing all offloaded optimizer state on the GPU at once.

Impact

GaLoreAdamW8bit can now train with FP32 LowLevelZeroPlugin CPU offload.

Fixes #6399.

Validation

  • python -m pytest tests/test_optimizer/test_dist_galore.py -q (4 GPUs; 1 passed)
  • the regression covers two consecutive CPU-offload optimizer steps and verifies master parameters and non-paged optimizer state return to CPU
  • compileall and git diff --check: passed

@sdjasj
sdjasj force-pushed the agent/fix-galore-cpu-offload-6399 branch from 82b4cfa to 4a5ea63 Compare July 29, 2026 13:09
@sdjasj
sdjasj marked this pull request as ready for review August 1, 2026 12:03
@sdjasj
sdjasj requested a review from a team as a code owner August 1, 2026 12:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG]: RuntimeError When Using GaLoreAdamW8bit with LowLevelZeroPlugin (fp32 + CPU Offload)

1 participant