mirror of
https://github.com/axolotl-ai-cloud/axolotl.git
synced 2026-10-12 02:11:20 +00:00
Wing LianandNanoCode012
3035553f07
* Fix FSDP2 QLoRA full checkpoint save and resume Preserve TorchAO 8-bit optimizer codes, scales, and expert ownership in full checkpoints. Restore model and optimizer shards across compatible meshes, and shard packed bitsandbytes expert weights with their quantization metadata. Add distributed checkpoint and packed-expert regression coverage, plus documentation of restore compatibility and checkpoint memory costs. * Cover full-parameter EP model save routing Exercise the Trainer save route through the patched Accelerator getter without PEFT. Verify global expert shapes and values, frozen parameters, buffers, and rank-zero-only assembly. * Remove obsolete EP state-dict gather helper Remove the unused gather_ep_experts_into_state_dict function and its dedicated test worker and test. Full checkpoints use the ownership-aware FSDP2 gather path. * Preserve native FSDP2 full model checkpoint state * Separate packed checkpoint CPU and GPU CI tests * Fix model-only FSDP2 trainer saves and final export routing * Isolate distributed checkpoint tests and subprocess coverage from regular CPU CI * Reject incompatible logical BNB shapes during checkpoint restore * Preserve frozen PEFT state and reject generic packed EP exports * increase distributed cpu timeout: * test: shard packed bnb expert in non-EP FSDP2 checkpoint case and verify restored values * ci: keep gpu-marked CPU-runnable monkeypatch tests in the CPU lanes * feat: reject torchao 4-bit and FP8 optimizers with expert parallelism --------- Co-authored-by: NanoCode012 <[email protected]>