Skip to content

CPU offload toggles feature #1228

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

Draft
ShotaKaji5207 wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

CPU offload toggles feature #1228

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d952d6a
Add CPU offload toggles feature
ShotaKaji5207 Nov 25, 2025
2d33fe7
Merge remote-tracking branch 'upstream/main' into feat/cpu-offload-to…
ShotaKaji5207 Nov 25, 2025
28583a2
Merge remote-tracking branch 'upstream/main' into feat/cpu-offload-to…
ShotaKaji5207 Nov 28, 2025
0d07ba8
Merge remote-tracking branch 'upstream/main' into feat/cpu-offload-to…
ShotaKaji5207 Dec 4, 2025
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
8 changes: 7 additions & 1 deletion open_instruct/grpo_fast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,8 @@ class Args:
"""whether to offload parameters to CPU (reduces GPU memory usage)"""
deepspeed_offload_optimizer: bool = False
"""whether to offload optimizer states to CPU (reduces GPU memory usage)"""
deepspeed_cpu_adam: bool = False
"""Whether to use DeepSpeedCPUAdam optimizer"""
gather_whole_model: bool = True
"""whether to gather the whole model to boardcast (not doable for 70B but can be faster for 8B)"""
enable_queue_dashboard: bool = True
Expand Down Expand Up @@ -708,7 +710,11 @@ def load(self, path: str, map_location=None):
optim_params = get_optimizer_grouped_parameters(self.policy, args.weight_decay)
else:
optim_params = self.policy.parameters()
self.optimizer = torch.optim.AdamW(optim_params, lr=args.learning_rate, fused=args.fused_optimizer)
if args.deepspeed_cpu_adam:
from deepspeed.ops.adam import DeepSpeedCPUAdam
self.optimizer = DeepSpeedCPUAdam(optim_params, lr=args.learning_rate)
else:
self.optimizer = torch.optim.AdamW(optim_params, lr=args.learning_rate, fused=args.fused_optimizer)
num_scheduler_steps = args.num_training_steps * args.num_epochs * args.num_mini_batches
warm_up_steps = args.warm_up_steps
if args.warmup_ratio > 0.0:
Expand Down