Skip to content

[ckpt, hardware] fix: make checkpoint CPU offload device-agnostic #8232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wuxibin89 merged 1 commit into from
Oct 15, 2026
Merged

[ckpt, hardware] fix: make checkpoint CPU offload device-agnostic #8232

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 5 additions & 9 deletions verl/utils/checkpoint/fsdp_checkpoint_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from transformers import GenerationConfig, PreTrainedTokenizer, ProcessorMixin
from transformers.dynamic_module_utils import custom_object_save

from verl.utils.device import is_cuda_available
from verl.utils.device import is_device_available
from verl.utils.fs import copy_to_local, is_non_local, local_mkdir_safe
from verl.utils.fsdp_utils import fsdp_version, get_fsdp_state_ctx, normalize_peft_param_name
from verl.utils.logger import log_with_rank
Expand Down Expand Up @@ -291,14 +291,10 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte

# every rank download its own checkpoint
state_dict_cfg = (
ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
if self.should_load_model
else None
ShardedStateDictConfig(offload_to_cpu=is_device_available()) if self.should_load_model else None
)
optim_cfg = (
ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
if self.should_load_optimizer
else None
ShardedOptimStateDictConfig(offload_to_cpu=is_device_available()) if self.should_load_optimizer else None
)
with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
if self.should_load_model:
Expand Down Expand Up @@ -400,8 +396,8 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
)

# every rank will save its own model and optim shard
state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=is_device_available())
optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=is_device_available())
with warnings.catch_warnings():
warnings.simplefilter("ignore")
with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
Expand Down
Loading