tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

fix: apply fix for only CP mode

Browse Source
This commit is contained in:
NanoCode012
2026-02-25 14:49:46 +07:00
parent 0d0122cabe
commit eb13054672
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/axolotl/core/trainers/base.py
View File

@@ -720,10 +720,14 @@ class AxolotlTrainer(
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
LOG.info(f"Saving model checkpoint to {output_dir}") LOG.info(f"Saving model checkpoint to {output_dir}")
# fix for Context Parallel save # fix for Context Parallel save: CP eval invalidates tensor storage
if state_dict is None: # pointers, so clone to CPU to get fresh valid storage for safetensors
state_dict = self.accelerator.get_state_dict(self.model) if (
if state_dict is not None: state_dict is not None
and self.axolotl_cfg
and self.axolotl_cfg.context_parallel_size
and self.axolotl_cfg.context_parallel_size > 1
):
state_dict = { state_dict = {
k: v.detach().cpu() if isinstance(v, torch.Tensor) else v k: v.detach().cpu() if isinstance(v, torch.Tensor) else v
for k, v in state_dict.items() for k, v in state_dict.items()