save_final_model calls trainer.model.save_pretrained directly instead of HF Trainer.save_model, skipping ZeRO-3/FSDP state consolidation and rank gating (CR-A02-012) #401

Wave 1 · theme D1-5: Training, merging and export correctness · highest severity High · 1 finding(s) · effort small

Summary

A ZeRO-3 LoRA run (documented: 'LoRA/DoRA adapters work well with both DeepSpeed and FSDP') completes. adapter_model.safetensors holds empty or sharded tensors, written concurrently by every rank, and the gates evaluate a broken or base-equivalent adapter that is then shipped.

Primary finding

CR-A02-012 — save_final_model calls trainer.model.save_pretrained directly instead of HF Trainer.save_model, skipping ZeRO-3/FSDP state consolidation and rank gating

severity High · reachability documented-opt-in · intent contract-violation · confidence confirmed · independently confirmed · not reproduced (code trace)

What happens. For adapter saves, save_final_model calls self.trainer.model.save_pretrained(final_path) and falls back to trainer.save_model only if that raises (trainer.py:1716-1730). With merge_adapters it calls merge_and_unload().save_pretrained (1736-1740). PEFT 0.19.1 get_peft_model_state_dict builds the adapter from model.state_dict() with no DeepSpeed ZeRO-3 gather (its docstring says the unwrapped model is expected), and save_pretrained defaults to is_main_process=True. HF Trainer.save_model gathers explicitly: _zero3_consolidated_16bit_state_dict(exclude_frozen_parameters=True) for PEFT+ZeRO-3, accelerator.get_state_dict for FSDP, and writes only when args.should_save. Under ZeRO-3 (a shipped preset, 'zero3'/'zero3_offload') or FSDP full_shard, the direct path is expected to write partitioned or empty adapter tensors on every rank without raising. All post-training gates then evaluate that artefact, and it is what gets promoted.

Where.

• forgelm/trainer.py:1735-1740 (merge_and_unload + merged.save_pretrained)
• forgelm/trainer.py:1720 (self.trainer.model.save_pretrained)
• forgelm/trainer.py:1729
• forgelm/trainer.py:1737
• forgelm/trainer.py:633-637 (zero3 presets)
• forgelm/trainer.py:1612,1620 (callers, run on every rank, no rank gating)
• configs/deepspeed/zero3.json
• docs/reference/distributed_training.md:163
• forgelm/trainer.py:591-597 (_apply_distributed_config deepspeed branch)
• forgelm/config.py:347-348

Evidence.

• Code reading plus installed-library introspection: PeftModel.save_pretrained has 'is_main_process' (default True); get_peft_model_state_dict source contains no zero3 handling; transformers Trainer.save_model source contains the zero3 consolidated-state-dict branch and should_save gating. Not executed: needs multi-GPU DeepSpeed.

Recommended fix. Under any distributed strategy, always use self.trainer.save_model(final_path) (rank-gated and consolidating) and save the tokenizer on rank 0 only. For merge_adapters under ZeRO-3, gather first or refuse with a clear error. Needs multi-GPU hardware (DeepSpeed ZeRO-3) to confirm.

Acceptance criteria

• A regression test reproduces the failure scenario and asserts the corrected behaviour.
• Documentation listed above is corrected in English and in its Turkish mirror in the same PR.
• Every finding in this issue is resolved or explicitly closed with a reason.

_{Found by the October 2026 ForgeLM code and documentation review against commit cab2947. Code links point at that commit; line numbers may have moved since. Finding IDs (CR-…, DOC-…, TECH-…) are stable references for this backlog.}